Source-linked AI summary
Toward a First-Principles Update Geometry for the Language-Model Head
Aditya Somasundaram
TL;DR
The paper asks how to define update geometry for a language-model head whose logits are immediately transformed by softmax. It derives Hilbert-distance geometry from the composed module, showing that update size is governed by token-row separation, and proposes maximizing minimum separation under a diameter constraint. Whether this yields an efficient optimizer and whether language requires semantic or data-dependent geometry remain open questions.
Problem
The LM head’s token rows are not directly characterized by singular-value conditioning because the head is extremely tall and its logits are passed through softmax.
Method
The paper uses Hilbert’s projective distance to derive the functional geometry of LM-head updates from softmax outputs and logit differences.
Results
Under bounded hidden states, the functional size of an LM-head update is the Euclidean diameter of its token-indexed rows and is independent of the current head U.
Takeaways & Limitations
The proposed design objective maximizes the smallest pairwise row distance while constraining the largest, yielding an approximately equidistant packing problem when V ≫ d.
Takeaways & Limitations
An efficient and effective LM-head optimizer based on this geometry remains to be determined, and its implementation should preserve momentum alignment and large-vocabulary efficiency.
Abstract
from arXiv · showhide
We study the language-model head and softmax as a single module, deriving an update geometry from their composition rather than from the weight matrix in isolation. Under Hilbert's projective distance, the maximum change caused by an update $S$ over $\left|\left|{h}\right|\right|_2\le H$ is $H\max_{i<j}\left|\left|{s_i-s_j}\right|\right|_2$, which is $H$ times the Euclidean diameter of its token rows. Motivated by Muon's singular-value conditioning, we propose maximizing the smallest row separation while constraining this diameter, producing an approximate-equidistance problem when $V\gg d$.
1 Introduction
The paper treats the LM head and softmax as one functional module, asking how updates change output probabilities rather than logits alone. It derives a row-based geometry and motivates approximately equidistant token-row updates.
- An LM head maps a hidden state h ∈ R^d to a probability distribution over V tokens.
- The relevant update effect is the change in the probability distribution after logits pass through softmax.
- Hilbert’s projective distance is chosen because it measures pairwise log-odds changes and is invariant to common logit shifts.
- Applying singular-value conditioning directly does not describe the geometry among token-indexed rows of the extremely tall LM head.
- The proposed row geometry treats update size as row-cloud diameter and seeks to make token rows as evenly separated as possible.
2 Hilbert distance
Hilbert distance measures the largest change in pairwise log odds between probability distributions. For softmax outputs, it depends only on logit differences, linking probability movement and cross-entropy-loss changes.
- Hilbert distance compares probability coordinates through their largest pairwise log-ratio difference.
- 2.1 Softmax and Hilbert distance: For softmax outputs, Hilbert distance depends only on the difference between logits, not the original logits or normalization constants.
- 2.1 Softmax and Hilbert distance: The distance ignores common logit shifts, matching softmax’s invariance to adding the same constant to every logit.
- 2.2 Why choose Hilbert distance?: Hilbert distance equals the largest difference between the cross-entropy-loss changes for two possible target tokens.
- 2.2 Why choose Hilbert distance?: Controlling Hilbert distance simultaneously bounds cross-entropy-loss changes for every possible target token.
3 Geometry of an LM Head Update
The paper derives LM-head update size from Hilbert distance on the composed head–softmax map, showing that it is governed by the Euclidean diameter of the update’s token-row cloud. This geometry constrains both forward probability movement and backward perturbation, enabling a row-diameter steepest-descent formulation.
- 3.1 Forward probability movement: For a hidden state h, the LM-head logit change under update S is Sh, and each token pair’s relative-log-odds change depends on (s_i−s_j)^⊤h.The realized change depends on row-pair alignment with h.
- 3.1 Forward probability movement: D(S) = max_i<j ||s_i−s_j||_2 is the functional size of an LM-head update under Euclidean-bounded hidden states.The maximum Hilbert-distance change is H times this row diameter.
- 3.1 Forward probability movement: The update-size result depends only on S, because Hilbert distance removes the current logits and softmax normalization from the calculation.Common shifts of every row are functionally invisible and can be removed by centering S.
- 3.2 Backward map: The same row-diameter constraint controls the LM-head backward perturbation, including ||S^⊤(p−e_c)|| ≤ D(S) for cross-entropy loss.The argument uses the zero-sum structure of gradients propagated through softmax.
- 3.3 Steepest Descent for the LM Head: The proposed steepest-descent problem maximizes first-order loss reduction subject to D(S) ≤ η, making it the LM-head analogue of spectral descent.The constraint is induced by the Hilbert-distance bound rather than by a spectral norm.
4 A Proposal for an LM Head Optimizer
The optimizer proposal augments the row-diameter constraint with a max–min objective so token-row separations are distributed more evenly. Because tall LM heads cannot realize exact equidistance, the design targets approximate equidistance while acknowledging that equal treatment of token pairs may be too strong.
- 4 A Proposal for an LM Head Optimizer: D(S) constrains the largest row separation but allows some token pairs to remain nearly indistinguishable.The diameter constraint alone does not determine the arrangement of the remaining rows.
- 4 A Proposal for an LM Head Optimizer: The ratio ρ(S) = δ(S)/D(S), with δ(S) the smallest row separation, measures how evenly pairwise token distinctions are distributed.Values near one indicate nearly equal pairwise distances.
- 4 A Proposal for an LM Head Optimizer: A Muon-like LM-head optimizer should use centered raw momentum to expand short row differences and shrink large ones, making pair distances more uniform.The centered momentum retains information about token-pair attraction and repulsion.
- 4 A Proposal for an LM Head Optimizer: V ≤ d + 1 permits a regular simplex with δ(S) = D(S), whereas V > d + 1 makes exact equidistance impossible.Equidistance requires V−1 ≤ d because the corresponding Gram matrix has rank V−1 while rank(SS^⊤) ≤ d.
- 4 A Proposal for an LM Head Optimizer: The proposal treats every token pair equally, which may be too strong for synonyms, spelling variants, or related subwords that should respond similarly.This is an author-stated caveat about the pairwise objective.
5 Future Work
Future work asks whether the derived row-diameter geometry can support a scalable practical LM-head optimizer and how the first-principles analysis should extend across modern neural modules.
- 5 Future Work: A practical LM-head optimizer should improve δ(S)/D(S), preserve alignment with raw momentum, and remain efficient for large vocabularies.These are stated design requirements for the proposed optimizer.
- 5 Future Work: The analysis should extend to other modern-ML modules whose transformations and activation functions may require distinct functional optimizer geometries.Examples include attention projections, SwiGLU layers, embeddings, and routers.
6 Related work
The paper situates its contribution among Hilbert geometry, softmax perturbation analyses, LM-head studies, packing-based classifier geometry, and function-informed optimization. Its distinction is deriving an exact row-diameter seminorm by pulling Hilbert distance through the complete head–softmax module.
- 6 Related work: Hilbert geometry connects probability-simplex distances to log probabilities, contraction theory, positive operators, and Sinkhorn scaling.The paper uses this geometric lineage to motivate its probability-distribution distance.
- 6 Related work: Prior softmax work studies how logit perturbations propagate to probabilities, including Lipschitz and co-coercivity properties and sharp Euclidean and ℓp constants.These analyses provide perturbation results without being the paper’s composed-head update geometry.
- 6 Related work: LM-head research has addressed representational rank constraints, multimodal next-token distributions, and the computational cost of large-vocabulary normalization.The paper addresses a complementary optimization-geometry question.
- 6 Related work: The pullback converts the V LM-head update rows into a point cloud in R^d, linking the proposal to last-layer, neural-collapse, and packing geometry.Its max–min ratio is related to minimal V-point diameter constructions.
- 6 Related work: Unlike spectral approaches that act on hidden matrix weights, the paper treats the head and softmax as a complete module and induces an exact row-diameter seminorm on U-updates.This follows the broader first-principles philosophy of choosing geometry from a parameter block’s function.
7 Conclusion
The paper derives LM-head update geometry from the post-softmax probability distribution using Hilbert’s projective distance. It identifies row separation as the functional geometry and proposes an approximately equidistant row-packing objective for optimization.
- Hilbert’s projective distance measures LM-head updates through their effect on the post-softmax probability distribution.It measures the largest change in pairwise log odds.
- The functional size of an LM-head update is the Euclidean diameter of its token-indexed rows.
- The proposed geometry maximizes the smallest pairwise row distance while constraining the largest row distance.This parallels Muon’s conditioning principle while operating on token-indexed row geometry.
- When V ≫ d, exact pairwise equality is impossible, yielding an approximately equidistant packing problem in the hidden dimension.
- The paper proposes maximizing ρ(S) = δ(S)/D(S) to construct an optimizer update, while leaving semantic or data-dependent geometry as an open modeling question.