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

arXiv: https://arxiv.org/abs/2608.22253 · HTML: https://mlsanity.com/paper/2608.22253/

## 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 (most cited first, 22 of 22)

- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [Breaking the Softmax Bottleneck: A High-Rank RNN Language Model](https://mlsanity.com/paper/1711.03953/) (2017)
- [On the Properties of the Softmax Function with Application in Game Theory and Reinforcement Learning](https://mlsanity.com/paper/1704.00805/) (2017)
- [Muon Optimizes Under Spectral Norm Constraints](https://mlsanity.com/paper/2506.15054/) (2025)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [Scalable Optimization in the Modular Norm](https://mlsanity.com/paper/2405.14813/) (2024)
- [Softmax is 1/2-Lipschitz: A tight bound across all ℓp norms](https://mlsanity.com/paper/2510.23012/) (2025)
- [RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization](https://mlsanity.com/paper/2603.20527/) (2026)
- [Memory-Efficient LLM Pretraining via Minimalist Optimizer Design](https://mlsanity.com/paper/2506.16659/) (2025)
- [On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer](https://mlsanity.com/paper/2603.09952/) (2026)
- [Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers](https://mlsanity.com/paper/2605.18106/) (2026)
- [Gradient Multi-Normalization for Stateless and Scalable LLM Training](https://mlsanity.com/paper/2502.06742/) (2025)
- [Lost in Backpropagation: The LM Head is a Gradient Bottleneck](https://mlsanity.com/paper/2603.10145/) (2026)
- [Non-linear Embeddings in Hilbert Simplex Geometry](https://mlsanity.com/paper/2203.11434/) (2022)
- [On a Generalization of Wasserstein Distance and the Beckmann Problem to Connection Graphs](https://mlsanity.com/paper/2312.10295/) (2023)
- [Does the LM Head Create a Harmful Gradient Bottleneck? A Causal Test](https://mlsanity.com/paper/2608.16671/) (2026)

## 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.