Toward a First-Principles Update Geometry for the Language-Model Head
Aditya Somasundaram ⋅ Charles Guille-Escuret ⋅ Alexander Moreno ⋅ Zhengzhong Liu ⋅ Eric Xing
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_2 \right\rVert \leq H$ is exactly $H D(S)$, where $D(S)=\max_{i
Chat is not available.
Successful Page Load