mlsanityData to 14 Oct 2026

Toward a First-Principles Update Geometry for the Language-Model Head

arXiv 2608.22253 ↗First posted 23 Aug 2026cs.LGcited by 0

Not on the map. Why?

Abstract (arXiv)

Muon motivates designing optimizer geometry around the function of each parameter block and uses the spectral norm for hidden linear layers. For the language-model head, the spectral norm is not a faithful measure of functional change. Softmax removes shared logit shifts, whereas the spectral norm can assign arbitrarily large size to updates that change no output probability. We therefore treat the LM head and softmax as one module and derive an update geometry for their composition. Hilbert's projective distance respects this invariance as it measures the largest change in pairwise log odds. For an update S with token rows si^top, we show that the largest Hilbert distance over leftlVert hrightrVert2≤ H is exactly H D(S), where D(S)=maxi<jleftlVert si - sjrightrVert2 is the Euclidean row diameter. This diameter replaces the spectral norm in the resulting Muon-style steepest descent problem. An exact solution is possible, but its direct formulation contains one d-dimensional vector variable for every token pair. For a vocabulary size of approximately 50k, this means more than one billion token pairs, making the calculation impractical at every training step. We instead impose a stronger common-ball constraint and derive projected RowNorm as an O(Vd) solution. For the exact RowNorm oracle, we prove that its first-order decrease is at least 1/√2 of the exact diameter-constrained optimum. With Muon on the backbone, experiments across three seeds at 190M, 380M, and 640M parameters show that RowNorm reduces mean final step diameters and empirical Hilbert RMS perturbations by factors of 45--60 and 12--15, respectively, with only a 0.0057--0.0153 increase in mean final validation loss.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 22 of its 45 references that are in the map. First posted on arXiv 23 Aug 2026.

Builds on

22 of 45 references.

References · 22

  1. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
  2. 2020GLU Variants Improve Transformercited by 1,804
  3. 2019Root Mean Square Layer Normalizationcited by 1,480
  4. 2015Optimizing Neural Networks with Kronecker-factored Approximate Curvaturecited by 796
  5. 2025Muon is Scalable for LLM Trainingcited by 428
  6. 2024Old Optimizer, New Norm: An Anthologycited by 193
  7. 2025Training Deep Learning Models with Norm-Constrained LMOscited by 192
  8. 2017Breaking the Softmax Bottleneck: A High-Rank RNN Language Modelcited by 152
  9. 2017On the Properties of the Softmax Function with Application in Game Theory and Reinforcement Learningcited by 142
  10. 2025Muon Optimizes Under Spectral Norm Constraintscited by 86
  11. 2024Modular Duality in Deep Learningcited by 77
  12. 2024Scalable Optimization in the Modular Normcited by 68
10 more
  1. 2025Softmax is 1/2-Lipschitz: A tight bound across all ℓp normscited by 22
  2. 2026RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimizationcited by 15
  3. 2025Memory-Efficient LLM Pretraining via Minimalist Optimizer Designcited by 13
  4. 2026On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfercited by 12
  5. 2026Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routerscited by 11
  6. 2025Gradient Multi-Normalization for Stateless and Scalable LLM Trainingcited by 9
  7. 2026Lost in Backpropagation: The LM Head is a Gradient Bottleneckcited by 9
  8. 2022Non-linear Embeddings in Hilbert Simplex Geometrycited by 6
  9. 2023On a Generalization of Wasserstein Distance and the Beckmann Problem to Connection Graphscited by 5
  10. 2026Does the LM Head Create a Harmful Gradient Bottleneck? A Causal Testcited by 1

Built on by

Not cited yet.