Skip to content
AI Atlas
PaperActive

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

arxiv.org/abs/2608.22253

quality89

Updated 5 h ago · first seen 11 Sept 2026

paper_01M294FSF7HB06SV0E2BAAS0YY

Published
11 Sept 2026
T1 · 5 h ago
arXiv
2608.22253
T1 · 5 h ago
Category
cs.LG
T1 · 5 h ago

Abstract

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 $s_i^\top$, we show that the largest Hilbert distance over $\left\lVert h\right\rVert_2\leq H$ is exactly $H D(S)$, where $D(S)=\max_{i<j}\left\lVert s_i - s_j\right\rVert_2$ 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 $50$k, 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/\sqrt{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.

Authors 5

Aditya Somasundaram, Charles Guille-Escuret, Alexander Moreno, Zhengzhong Liu, Eric Xing

Specification

Official page

Source:arXiv (Atom API + RSS)T1observed 5 h agohigh

Arxiv announce type
replace

Source:arXiv (Atom API + RSS)T1observed 5 h agohigh

arXiv id
2608.22253

Source:arXiv (Atom API + RSS)T1observed 5 h agohigh

Categories
cs.LG

Source:arXiv (Atom API + RSS)T1observed 5 h agohigh

PDF

Source:arXiv (Atom API + RSS)T1observed 5 h agohigh

Primary category
cs.LG

Source:arXiv (Atom API + RSS)T1observed 5 h agohigh

Published
11 Sept 2026

Source:arXiv (Atom API + RSS)T1observed 5 h agohigh

Each value shows its source, tier and observation time. Conflicting claims are kept side by side and flagged — never averaged. How AI Atlas records facts →

Provenance

Attributed facts

9

Source tiers

T19

Freshest observation

5 h ago

Conflicts

None