PaperlayerView exact revision on arXiv

Source-linked AI summary

The Row Normalization Puzzle in Muon

Jiayu Zhang, Tianyi Lin

arXiv:2609.39114v1cs.LG

TL;DR

NorMuon’s growing practical use contrasts with limited understanding of its worst-case convergence guarantees. This paper analyzes row normalization under operator-norm geometry and establishes dimension-dependent bounds, while experiments show slower synthetic performance but better LLM-pretraining performance than Muon.

  • Problem

    Whether row normalization improves provable convergence, including through interactions with approximate polar computation and exponential moving-average momentum, remains an open question.

    Source: §1 · p. 2

  • Method

    The paper analyzes NorMuon’s gradient-update alignment and proves lower and upper convergence bounds under deterministic and stochastic settings.

    Source: §1 · p. 2, §3 · p. 5

  • Results

    NorMuon has matching deterministic worst-case bounds of Θ(mLϵ−2), an extra factor of m over Muon, while synthetic experiments show slowdown and LLM pretraining shows improved performance over Muon.

    Source: §1 · p. 2, §5 · p. 12

  • Takeaways & Limitations

    Balanced neuron-wise updates alone do not guarantee faster optimization in the worst case, despite NorMuon’s improved LLM-pretraining performance.

    Source: §1 · p. 2, §5 · p. 12

    The analysis assumes a lower-bounded objective and, for approximate polar computation, an output satisfying the stated singular-value conditions.

    Source: §2.2 · pp. 4–5

Abstract

from arXiv · show

This paper examines how row-wise renormalization affects Muon, focusing on the gap between NorMuon's worst-case guarantees and its practical performance (Li et al.). Despite its growing adoption and promising performance in large language model (LLM) pretraining, NorMuon's worst-case guarantees remain poorly understood. One fundamental question is: Does row normalization yield provable convergence gains, potentially through its interaction with approximate polar computation and exponential moving-average momentum? Our results show that row normalization introduces a dimension-dependent factor in the worst-case iteration complexity under the operator-norm geometry, which persists even with exact polar computation and any fixed momentum parameters. Indeed, we establish an algorithm-dependent lower bound and a matching upper bound in deterministic settings, and extend our upper bound analysis to stochastic settings. Both upper-bound analyses allow approximate polar computation. Experiments show that NorMuon is slower than Muon on synthetic problems inspired by our worst-case construction, yet outperforms Muon in LLM pretraining. These findings sharpen the puzzle of why row normalization helps in practice and complement the recent findings of Dewulf et al.

1 Introduction

The introduction asks whether NorMuon’s practical gains from row normalization translate into better worst-case convergence guarantees. It presents a dimension-dependent theoretical cost alongside contrasting synthetic and LLM-pretraining results.

1 Introduction

  • Muon updates weight matrices with the polar factor of gradient momentum, while NorMuon adds row-wise second-moment normalization to balance update magnitudes across neurons [Jordan et al., 2024, Li et al., 2026].Related matrix-aware methods include Scion [Pethick et al., 2025].Source: §1 · pp. 1, 3
  • NorMuon is slower than Muon on synthetic problems but outperforms it in LLM pretraining, leaving practical benefits unexplained by worst-case guarantees.The paper frames this contrast as sharpening the row normalization puzzle.Source: §1 · p. 2
  • The motivating question is whether row normalization yields provable convergence gains through its interaction with approximate polar computation and exponential moving-average momentum.The introduction positions this as a worst-case counterpart to empirical findings on balanced neuron updates [Dewulf et al., 2026].Source: §1 · p. 2
  • The paper establishes a NorMuon lower bound of Ω(mLϵ−2), versus Muon’s O(Lϵ−2), and upper bounds of O(mLϵ−2) deterministically and O(mLϵ−2 + m2Lσ2ϵ−4) stochastically.The upper-bound analyses allow approximate polar computation.Source: §1 · p. 2
  • The lower bound holds with exact polar computation, fixed momentum parameters, and deterministic adaptive step sizes, while the analysis studies gradient-update alignment in operator-norm geometry.The construction uses a rank-one gradient with unequal row magnitudes, whose row-balanced polar factor weakens gradient alignment.Source: §1 · pp. 2–3

2 Preliminaries and Technical Background

The preliminaries specify the matrix-optimization setting, Muon and NorMuon update rules, and assumptions used to analyze their convergence. They measure convergence by gradient norms and state the resulting dimension-dependent bounds.

2.1 LLM pretraining and matrix optimization

  • LLM pretraining minimizes expected next-token prediction loss over sequences, and the analysis treats each attention or feed-forward weight matrix as a separate variable.The preliminaries introduce an autoregressive model assigning probabilities to token sequences.Source: §2.1 · p. 3
  • Muon applies a polar factor to normalized gradient momentum; NorMuon additionally rescales each row using an exponential moving average of its mean squared entries.Row normalization changes relative row magnitudes, while final rescaling fixes the Frobenius norm of each nonzero update.Source: §2.1 · p. 4

2.2 Assumptions and basic properties

  • Under operator-norm geometry, the analysis assumes nuclear-norm smoothness and measures convergence by the minimum or expected nuclear norm of the gradient.It establishes matching deterministic bounds with an extra factor of m for row normalization, and extends the upper bound to stochastic settings with approximate polar computation.Source: §2.2 · p. 5

3 Main Results

In deterministic operator-norm geometry, NorMuon’s matching lower and upper bounds show an extra factor of m in iteration complexity. The paper extends its upper-bound analysis to stochastic gradients and sketches how row normalization affects alignment and descent.

3.1 Deterministic setting

  • NorMuon can incur an extra factor of m over Muon’s dimension-independent O(∆Lϵ−2) rate, even with fixed momentum parameters and adaptive step sizes [Shen et al., 2026].The lower-bound construction uses exact gradients and exact polar computation.Source: §3.1 · p. 6
  • The deterministic upper bound matches the lower bound at O(m∆Lϵ−2), and remains valid with approximate polar computation when ℓ < c min{m,n}.The stated bound is attained with β1 = 0.Source: §3.1 · p. 6

3.2 Stochastic setting

  • In the stochastic setting, NorMuon’s upper bound has an additional dimension-dependent factor of m max{1, m/n}; a matching lower bound remains an open direction.Muon’s cited upper bound is O(min{m, n}∆Lσ2ϵ−4) for fixed σ > 0 and sufficiently small ϵ [Zhang and Lin, 2026].Source: §3.2 · p. 6

3.3 Proof sketches

  • The lower-bound construction keeps the gradient constant while row normalization equalizes rows, fixing the update direction for any fixed momentum parameters.The proof sketch gives Dt+1 = 0.2√mnD⋆ and ∇F(Xt) = 2ϵH⋆ along the constructed iterates.Source: §3.3 · p. 7
  • For upper-bound sketches, deterministic descent uses β1 = 0; stochastic analysis decomposes momentum into gradient, drift, and noise terms, controlling drift by Assumption 2.1 and noise via gradient-noise martingale structure and Assumption 2.2.Source: §3.3 · p. 7

4 Experiments

Experiments compare Muon and practical NorMuon on synthetic problems, image classification, and LLM pretraining. The findings vary by setting, with ablations probing step sizes, matrix groups, and normalized optimizer variants.

4.1 Synthetic experiment

  • As m grows from 8 to 2048, NorMuon’s loss reduction drops while Muon’s stays nearly unchanged; an equal-norm NorMuon variant points to row normalization as the source of the gap.Results average loss reduction over 100 random seeds; larger values indicate better performance.Source: §4.1 · p. 8, Table 1 · p. 8

4.2 Image classification

  • All Muon-family optimizers outperform AdamW; column-normalized NorMuon has the highest mean test accuracy, while Muon has the lowest mean test loss.The comparison is for CIFARNET, with mean test loss and accuracy reported over 50 seeds.Source: §4.2 · p. 9, Table 2 · p. 9

4.3 LLM pretraining

  • NorMuon lowers validation loss than Muon for most training across all three model sizes and improves test loss and average downstream accuracy at 1.38B and 6.44B.The 6.44B model uses a different tokenizer, so its loss is not directly comparable with the other models.Source: Figure 1 · p. 10, Table 3 · p. 10, §4.3 · p. 10

4.4 Ablation studies

  • NorMuon has lower final validation loss than Muon at every tested step size; both reach their lowest loss at 0.03 for 286M and 0.01 for 1.38B.The step-size comparisons are shown in Figure 2(a).Source: Figure 2 · p. 11, §4.4 · p. 11
  • Normalizing only QKVO matrices gives nearly the same losses as Muon, while normalizing only MLP input or output projections recovers most of NorMuon’s gain.Row-wise normalization of MLP matrices nearly matches NorMuon applied to all transformer matrices.Source: §4.4 · p. 11
  • Aurora and NorMuon have nearly identical validation and test losses, with Aurora slightly lower; Muon+ follows, while MuonEq-R is nearly identical to Muon.With the same step size and five PolarExpress iterations, Aurora, NorMuon, and Muon+ improve on Muon.Source: §4.4 · p. 11

5 Conclusion

The conclusion contrasts NorMuon’s extra worst-case cost with its improved LLM pretraining performance over Muon. It also identifies explaining when row normalization benefits LLM training as an open direction.

5 Conclusion

  • Matching deterministic lower and upper bounds give NorMuon worst-case complexity Θ(mLϵ−2), an extra factor of m compared with Muon.The lower bound holds with exact polar computation, any fixed momentum parameters, and deterministic adaptive step sizes.Source: §5 · p. 12
  • The upper-bound analyses extend to stochastic settings and allow approximate polar computation.The exact-polar and fixed-momentum conditions apply to the lower bound, not this upper-bound statement.Source: §5 · p. 12
  • Synthetic experiments show NorMuon slowing down, whereas LLM pretraining experiments show improved performance over Muon.The conclusion cautions that balanced neuron-wise updates alone do not guarantee faster worst-case optimization.Source: §5 · p. 12
  • Future work should identify LLM training structures associated with row-normalization benefits and develop theory that captures those benefits.Source: §5 · p. 12

A Further Related Works

The related-work discussion places the paper among matrix-aware optimizers and algorithm-dependent convergence analyses. It distinguishes NorMuon’s dimension-dependent worst-case cost from prior convergence and optimizer results.

Matrix-aware optimization methods

  • Muon and Scion construct updates by dualizing gradients under the operator norm, while several variants alter update scaling or preconditioning.The cited variants include AdaMuon, NAMO, Muon2, and Muown, among others.Source: A Further Related Works · p. 20

Analysis of matrix-aware optimization methods

  • Prior operator-norm analyses establish a dimension-free rate for Muon with exact polar computation and study finite polar approximation, momentum, and noise.The passage also surveys work on implicit biases, structured problems, and nonsmooth settings.Source: A Further Related Works · p. 20

Lower bound analysis

  • Unlike general first-order lower bounds, algorithm-dependent analysis fixes an update rule to identify its best attainable convergence result.Examples include lower bounds for SGD, cyclic coordinate descent, random-shuffling schemes, gradient descent, and sign-based methods.Source: A Further Related Works · pp. 20–21

Other optimizers and analysis

  • NorMuon applies normalization to row-wise norms of the orthogonalized update, unlike coordinate-wise adaptive methods that rescale using accumulated or averaged squared gradients.Source: A Further Related Works · p. 21
  • NorMuon’s worst-case cost is an extra dimension-dependent iteration factor, not non-convergence, and persists under any fixed momentum parameters.The passage contrasts this result with prior counterexamples concerning adaptive methods.Source: A Further Related Works · p. 21

B Missing Proofs

The supplied passage identifies e_i as the ith basis vector of R^n.

B Missing Proofs

B.1 Proof of Theorem 3.1

The proof constructs an objective on which Algorithm 1 follows a prescribed trajectory with constant gradients, then verifies the construction satisfies the theorem’s conditions.

B.1 Proof of Theorem 3.1

  • The constructed trajectory satisfies X_t = X_0 + q_tD⋆, F(X_t) = f_t, and ∇F(X_t) = 2ϵH⋆ through T_0.This gives the fixed gradient history used in the lower-bound construction.Source: §B.1 · p. 22
  • The deterministic adaptive step-size rule observes the same history and therefore selects the same step sizes as in the construction.The proof uses ∥2ϵH⋆∥_nuc = 2ϵ > ϵ to obtain the desired theorem result.Source: §B.1 · pp. 21–22
  • An induction using the update rule maintains the prescribed path, with the polar direction and normalized update aligned with D⋆.The argument sets O_{s+1} = Polar(H⋆) = H⋆ and D_{t+1} = 0.2√(mn)D⋆.Source: §B.1 · p. 22
  • The proof separately checks smoothness and the initial objective-gap condition required for the constructed F.It establishes the Lipschitz-gradient condition and that F is bounded below with finite F⋆.Source: §B.1 · p. 23

B.2 Proof of Theorems 3.2 and 3.3

The proofs derive deterministic and stochastic upper bounds by combining descent with control of the update’s alignment and approximate-polar errors.

B.2 Proof of Theorems 3.2 and 3.3

  • The update-alignment lemma gives ⟨M_{t+1},D_{t+1}⟩ ≥ 0.2κ√(mn)∥M_{t+1}∥_nuc and ∥D_{t+1}∥_op ≤ 0.2√(mn).These bounds hold for any α ≥ 0 and t ≥ 0.Source: §B.2 · p. 23
  • The approximate-polar analysis bounds the polar output’s operator norm between 1 − δ and 1 + δ.The proof relates this control to the singular values of the momentum matrix.Source: §B.2 · p. 24
  • The deterministic analysis gives T = O(m∆L/ϵ^2) iterations to reach an iterate with nuclear-norm gradient at most ϵ.The proof also bounds the total gradient-oracle calls by the iteration count.Source: §B.2 · p. 25
  • For the stochastic bound, the proof chooses T using ∆, L, κ, β1, and ϵ, then bounds stochastic-gradient-oracle calls.The analysis controls the accumulated noise using the stated variance assumption.Source: §B.2 · pp. 25–26

C Additional Experiments

Appendix C presents detailed experimental setups and additional results.

C Additional Experiments

C.1 Setup and additional results for the synthetic experiment

The synthetic experiments compare Muon and NorMuon under fixed and cosine step-size schedules, while the image-classification subsection specifies the training setup.

C.1 Setup and additional results for the synthetic experiment

  • As m grows, NorMuon selects larger steps, lowering loss faster initially before it stalls; the cosine-schedule experiment reports a similar dimension-dependent slowdown.Figure 3 contains synthetic loss curves; the cosine experiment uses T = 25 and reports its loss-reduction evaluation in Table 7.Source: §C.1 · p. 26, Figure 3 · p. 27
  • The fixed-step experiment tunes Muon and NorMuon on 50 of 150 seeds, then evaluates selected settings on the remaining 100 seeds over 1,000 iterations.Settings are selected by minimum loss; the cosine-schedule experiment instead tunes peak step sizes on 50 seeds and evaluates on 100.Source: §C.1 · p. 26

C.2 Setup for image classification

  • Image classification uses 5% linear warmup followed by cosine decay, with method-specific step-size and momentum sweeps.All methods sweep decoupled weight decay; gradients are clipped to Euclidean norm at most 1.Source: §C.2 · p. 28

C.3 Setup and additional results for LLM pretraining

NorMuon reports lower LLM losses than Muon across the tested ablations. The section also describes the pretraining setup, random-seed results, and optimizer-alignment patterns.

Setup for ablation studies

  • LLM runs use decoder-only transformers at depths 12, 24, and 32; ablations fix step size 0.02, five PolarExpress iterations, weight decay 0.02, and β1 = 0.95.Muon and NorMuon update transformer weight matrices, while AdamW updates embeddings, the language-model head, and other trainable parameters.Source: §C.3 · pp. 28–29

Random seeds

  • For the 1.38B model, NorMuon’s average improvement over Muon exceeds one standard deviation across three random seeds.Table 9 reports mean ± standard deviation; Table 3 reports only the mean.Source: §C.3 · p. 28, Table 9 · p. 29

Additional ablation results

  • NorMuon attains lower losses than Muon across tested step sizes, PolarExpress counts, second-moment factors, and weight-decay factors.The comparisons are shown in Figure 5.Source: §C.3 · pp. 29–31, Figure 5 · p. 30

Step-size sensitivity

  • NorMuon has lower validation and test losses at every tested base step size, even when each optimizer uses its own best step size.The tested values are 0.01, 0.02, 0.03, and 0.04; 0.02 is used as the common ablation step size.Source: §C.3 · p. 29, Table 10 · p. 30, Figure 5 · p. 30

Accuracy of polar approximation

  • NorMuon’s loss advantage persists across 3, 5, and 7 PolarExpress steps; increasing from 3 to 5 substantially lowers both optimizers’ test loss, with little further effect at 7.The gap between Muon and NorMuon is largest with 3 PolarExpress steps.Source: §C.3 · p. 30, Figure 5 · p. 30, Table 11 · p. 31

Second-moment factor

Weight decay

Alignment

  • On MLP matrices, NorMuon tends to have slightly higher alignment after the first few hundred steps and larger momentum nuclear and update operator norms; QKVO alignment is similar, and update Frobenius norms are similar.Figure 6 shows relative differences across parameter groups; its error bars span the 10th–90th percentile across matrices in each group.Source: §C.3 · p. 31, Figure 6 · p. 32

Related papers

Loading 2609.39114v1…