Skip to content

docs: ICNN/KeyNet docstring fixes (architecture, math, citation)#703

Merged
marcocuturi merged 2 commits into
mainfrom
docs/icnn-architecture-docstring
Jun 14, 2026
Merged

docs: ICNN/KeyNet docstring fixes (architecture, math, citation)#703
marcocuturi merged 2 commits into
mainfrom
docs/icnn-architecture-docstring

Conversation

@marcocuturi

@marcocuturi marcocuturi commented Jun 14, 2026

Copy link
Copy Markdown
Contributor

What

Follow-up to #702 (merged). Documentation-only cleanup of ICNN/KeyNet in src/ott/neural/networks/icnn.py.

Changes (no behavior change)

ICNN

  • Architecture block now shows the final layer as linear (no activation, per Remove activation on final layer of KeyNet and ICNN #702) and notes convexity is still preserved.
  • Adds the optional per-layer bias b_i (gated by use_bias); clarifies the W_x input-injection terms are always bias-free (built with use_bias=False), so there is no redundant double bias.

KeyNet

  • Adds an Architecture block in the same style as ICNN (linear final layer, optional bias, residual output).
  • Cites the KeyNet paper — Olausson, Monteiro, Klein, Cuturi, Amortizing Maximum Inner Product Search with Learned Support Functions (arXiv:2603.08001, 2026) — added to docs/references.bib as olausson:26.

Formatting unification

  • Inline formulas (function signature, inner-product potential, residual output) now consistently use :math: LaTeX instead of a mix of plain text and inline code. The affected docstrings are raw strings so LaTeX backslashes are literal.

🤖 Generated with Claude Code

marcocuturi and others added 2 commits June 14, 2026 19:50
Follow-up to #702. The ICNN docstring "Architecture" block still showed
the activation applied at every layer and omitted bias terms entirely:

- show the final layer is linear (no activation, per #702) and note that
  convexity is still preserved;
- include the optional per-layer bias (gated by `use_bias`), and clarify
  that the W_x input-injection terms are always bias-free.

Docstring-only change; no behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
- Give KeyNet an `Architecture::` block in the same style as ICNN, showing
  the linear final layer and optional bias.
- Unify formula rendering across both docstrings: inline expressions (the
  function signature, the inner-product potential, the residual output) now
  use `:math:` LaTeX consistently instead of a mix of plain text and inline
  code. Docstrings are now raw strings so LaTeX backslashes are literal.
- Add the KeyNet reference (Olausson et al., 2026, "Amortizing Maximum
  Inner Product Search with Learned Support Functions") to references.bib
  and cite it from the KeyNet docstring.

Docstring/bib-only change; no behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
@marcocuturi marcocuturi changed the title docs: fix ICNN architecture docstring (linear final layer + optional bias) docs: ICNN/KeyNet docstring fixes (architecture, math, citation) Jun 14, 2026
@marcocuturi marcocuturi merged commit e241ec3 into main Jun 14, 2026
8 of 9 checks passed
@marcocuturi marcocuturi deleted the docs/icnn-architecture-docstring branch June 14, 2026 18:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant