From Geometry to Generalization: Why Row Normalization Can Beat Adam and Muon

Jihwan Kim; Dogyoon Song; Chulhee Yun

cs.LGOct 8, 2026

From Geometry to Generalization: Why Row Normalization Can Beat Adam and Muon

Authors: Jihwan Kim, Dogyoon Song, Chulhee Yun

Organizations: Seoul National University · KAIST InnoCORE LLM · University of California, Davis · KAIST

Abstract

Different optimizers can fit the same training data while selecting classifiers with substantially different geometries, but whether this difference provably affects population performance remains unclear. We show that row-wise normalization can achieve strictly higher population accuracy than full-batch Adam, a proxy for random-reshuffling Adam, and exact-SVD Muon in high-dimensional multiclass classification. Under an isotropic Gaussian-cloud data model, this advantage arises because row normalization's class-wise Euclidean geometry asymptotically preserves the population decision-boundary directions, whereas Adam's coordinate-wise geometry and Muon's spectral geometry introduce nonvanishing distortions. Beyond isotropy, the advantage persists for full-batch training on class means with independently oriented class-mean and test-noise covariances. It holds for power-law spectra with class-mean exponent below one, even under heavily anisotropic test noise. When both covariances are diagonal and sufficiently close, the advantage over Adam can reverse, while applying the same random rotation to both restores it by changing only their alignment with Adam's coordinate axes. Synthetic and last-layer language-model experiments support the predicted advantage.

Figures & tables

Figure 1: Geometry of the learned classifiers. Class means are compared with classifier rows learned by row-normalized GD, Adam, and Muon for 60 classes in 20 dimensions: a 2D projection of two classes, with mean angles over all classes and their proportional-limit predictions from Proposition 3.2 (left); coordinate magnitudes at ranks $1,3,\ldots,19$ after sorting by class-mean magnitude (middle); and singular values at indices $1,3,\ldots,19$ of the full matrices (right).

Figure 2: Synthetic and language-model experiments. Left: Final evaluation accuracy of six deterministic optimizers on 30 draws of Gaussian-cloud data. Right: Validation perplexity during last-layer-only LLaMA training with RowSGDM, Muon, and Adam. Lower perplexity is better.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Additional optimizer comparisons on the baseline synthetic model. Left: The deterministic comparison in Figure 2 (left), augmented with Signum. Right: Random-reshuffling comparison of RR-RowSGDM, RR-Muon, RR-Signum, and RR-Adam with batch size

b=1

. Points are paired seeds and gray lines connect methods evaluated on the same dataset.

Figure 4: Auxiliary optimization geometries. RowGD, FrobGD, ColumnGD, ordinary GD, SpecGD, and SignGD on the same 30 paired baseline datasets, where FrobGD denotes Frobenius-normalized GD.

Panel of Figure 5 — Figure 5: Robustness to training and test noise. Top: training-noise coefficient $0.0447$ (left) and $0.4472$ (right), compared with $0.2236$ in Figure 2 (left). Bottom: test-noise coefficient $1.3416$ (left) and $3.1305$ (right), compared with $2.2361$ in the main-text setting. In all cases, the residual vectors are drawn from $\mathcal{N}({\mathbf{0}},I_{d}/d)$ .

Panel of Figure 6 — Figure 6: Robustness to class proportions and training sample count. Top: balanced test classes (left) and $\operatorname{Dirichlet}(0.3\mathbf{1}_{k})$ test proportions (right), compared with $\operatorname{Dirichlet}(\mathbf{1}_{k})$ in Figure 2 (left). Bottom: exactly one (left) and five (right) training examples per class, compared with independent counts drawn uniformly from $\{1,\ldots,5\}$ in the main-text setting.

Panel of Figure 7 — Figure 7: Robustness to problem size and class-mean distribution. Top: $(k,d)=(30,20)$ (left) and $(120,40)$ (right), compared with $(60,20)$ in Figure 2 (left). Bottom: class means drawn independently and uniformly from the unit sphere instead of from $\mathcal{N}({\mathbf{0}},I_{d}/d)$ .

Figure 8: Independent-orientation anisotropic experiment. Representative setting with

s_{\mu}=0.5

and

s_{\Sigma}=2

, where RowGD and RowGDM outperform SpecGD, Muon, SignGD, and Adam on all

30

paired seeds.

Figure 9: Coordinate-aligned reversal and recovery. RowGDM outperforms Adam at

\lambda=0.25

, Adam outperforms RowGDM at

\lambda=1

, and the RowGDM-over-Adam ordering is restored after applying the same Haar rotation.

Optimizer	NLL $\downarrow$	PPL $\downarrow$	$\widehat{D}_{\mathrm{avg\text{-}bdry}}\downarrow$	$\widehat{C}_{\mathrm{bdry}}\uparrow$
RowSGDM	$\mathbf{3.35503\pm 0.00005}$	$\mathbf{28.6466\pm 0.0013}$	$\mathbf{1.13194\pm 0.00004}$	$\mathbf{0.35652\pm 0.00005}$
Muon	$3.35699\pm 0.00004$	$28.7027\pm 0.0012$	$1.13807\pm 0.00006$	$0.35018\pm 0.00007$
Adam	$3.35759\pm 0.00002$	$28.7199\pm 0.0006$	$1.13350\pm 0.00008$	$0.35481\pm 0.00009$

Appendix

Table 1: Held-out performance and empirical boundary alignment for last-layer-only LLaMA training. Values are mean

\pm

sample standard deviation across five paired training seeds.

Model	$K$	$\widehat{s}_{\mu}$	$K_{\Sigma}$	Mean $\widehat{s}_{\Sigma,i}$	Median	Max
OLMoE-1B-7B-0125	$47{,}190$	$0.957735$	$218$	$0.987339$	$0.973537$	$1.374366$
Qwen3-30B-A3B	$72{,}577$	$0.932780$	$120$	$1.057924$	$1.045483$	$1.258163$
AMUSE-720M	$49{,}348$	$0.985366$	$211$	$1.054793$	$1.045965$	$1.402045$

Appendix

Table 2: Full-rank pointwise-OLS spectral exponents at the three final checkpoints.

K

is the number of observed classes used for the class-mean covariance, and

K_{\Sigma}

is the number of reliable classes used for the within-class covariance summary.

Explore similar work

Muown: Row-Norm Control for Muon Optimization
May 11, 2026Kai Lion, Florian Hübler, Bingcong Li +2Language Model Pretraining Muon Optimizer
Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany
The Row Normalization Puzzle in Muon
Sep 30, 2026Jiayu Zhang, Tianyi LinStochastic Optimization Optimization Convergence Analysis
Department of Industrial Engineering and Operations Research Columbia University
On the Two Faces of Adam in Separable Linear Classification
Sep 27, 2026Chen Fan, Csaba SzepesváriImplicit Bias Classification
University of Alberta

May 11, 2026, cs.LGJ/K move · Enter open · S save

Muown: Row-Norm Control for Muon Optimization

Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto+1

Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany

Language Model Pretraining Muon Optimizer

Open paper AskPDF arXiv

Muon has emerged as a strong competitor to AdamW for language model pre-training, yet its behavior at scale is sensitive to weight decay. Recent work has observed that, for Muon without decoupled weight decay, the spectral norm of weight matrices drifts upward over training. Through a decomposition of the spectral norm into a row-magnitude factor and a row-coherence factor, we identify the former as the empirical driver of this drift under Muon, while the latter remains well-behaved along the trajectory. Motivated by this diagnosis, we introduce Muown, a drop-in replacement for Muon that treats the row-magnitude vector as an explicit optimizer variable, updating it under the

\ell_\infty

geometry induced by the decomposition, while applying Muon unchanged to the remaining direction component. We prove that Muown attains the optimal non-convex rates in both deterministic and stochastic regimes under a dual norm aligned with the underlying geometries and with a stochastic noise coefficient that empirically remains below that of Muon throughout training. Across GPT-style pre-training on FineWeb-Edu with model sizes from 124M up to 2.7B parameters, Muown improves perplexity over Muon, SOAP, AdamW, and Lion. It also widens the plateau of near-optimal learning rates across model scales, reduces sensitivity to weight decay, and avoids the spectral norm drift at negligible step-time overhead when appropriately sharded.

Abstract

Explore similar work

Muown: Row-Norm Control for Muon Optimization

The Row Normalization Puzzle in Muon

On the Two Faces of Adam in Separable Linear Classification