When and How Does Muon Precondition? The Role of Structured Gradient Noise
Authors: Xiaoran Chen, Binghui Li, Liwei Wang, Pinyan Lu
Organizations: School of EECS, Peking University · Center for Machine Learning Research, Peking University · Center for Data Science, Peking University · State Key Laboratory of General Artificial Intelligence, Peking University · Key Laboratory of Interdisciplinary Research of Computation and Economics, Shanghai University of Finance and Economics
Muon has become a popular optimizer for large language model training, yet its theoretical foundations remain poorly understood. In this work, we study how gradient noise geometry shapes Muon's preconditioning benefits. We first show, theoretically and empirically, that under the classical isotropic noise assumption, Muon offers no acceleration beyond a rescaling of the learning rate. This motivates us to investigate structured noise arising from mini-batch sampling in associative memory with heterogeneous task frequencies. We establish an exact balancing threshold for two tasks and extend the analysis to multiple tasks, revealing how task coverage and gradient reliability jointly shape learning. Our analysis shows that Muon can soften the dependence on task frequency from linear to square-root, helping balance the learning of frequent and rare tasks. We further identify rotation-induced negative transfer in SignSGD that Muon avoids, and explain how momentum reduces sampling noise while tracking changing gradients. Controlled numerical experiments validate these mechanisms, and a preliminary LLM pretraining study illustrates the practical promise of fast-slow momentum.
Figures & tables
Figure 1: Isotropic noise removes first-order directional preconditioning at low SNR. Left: for an anisotropic d=8 signal, the response gain along each singular direction, normalized by cd , collapses near one across a wide range of signal singular values. Right: the full conditional-mean matrix remains within a few percent of the scalar-response prediction over the tested low-SNR amplitudes.
Figure 2: Mini-batch sampling controls the transition from frequency-limited to balanced growth. Left: the two-task gap-growth transition occurs at the exact threshold B/Bbal=1 . Right: in the higher-rank hierarchy, late margin-growth rates from different frequency blocks collapse when plotted against the expected task count Bpj and approach the full-batch unit response as that count increases. Error bars summarize 64 trajectories; the dotted curve is the task-presence guide 1−e−Bpj .
Figure 3: Corrupted observations introduce a second sampling scale. Left: full-matrix Muon changes from positive to negative large-margin drift as the expected number of corrupted task- j observations, Bpjρα , crosses an order-one scale; dashed curves show the exact projected counting formula, whose crossover approaches log2 in the Poisson limit. Right: near the population optimum under label noise, the full-categorical Muon response scales empirically as pj0.494 , close to the square-root prediction under the auxiliary covariance and frequency approximations, whereas the SGD coefficient scales linearly in pj .
Figure 4: Sparse support yields slow positive tail learning for Muon but negative transfer for SignSGD. Both panels use the same B=1 , p1=0.95 , p2=0.05 mini-batch streams over 64 runs. Left: Muon preserves support, so the rare-task margin remains nonnegative and grows on its occurrence clock. Right: at θ=1.15>π/4 , rotated SignSGD drives the rare-task margin downward while the frequent task continues to improve. The positive answer-embedding scalar c is absorbed into the SignSGD learning rate. Curves show medians and ribbons show interquartile ranges.
Figure 5: Momentum trades variance reduction against tracking. Left: different physical batch sizes and momentum coefficients nearly collapse when the two-task gap bias is plotted against the variance-equivalent task count Beffpmin . Right: with B and β fixed and a common horizon τmax=50 , the late-window mean-gap error stays small while η/(1−β) is small and grows once the EMA horizon becomes comparable to the state-evolution time scale.
Figure 6: Fast–slow momentum: core update and pretraining results. Left: schematic momentum mixing following Zhu et al. (2026) ; RowScale normalizes rows, and parameter and update scaling follow the cited implementation. Right: validation loss from steps 20k to 30k for a 0.12B dense language model, with a final reduction of 0.0209 . Other training settings are fixed; curves are unsmoothed, with one run per method.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Higher-rank trajectory view underlying Figure 2 (b). At small expected task count ( B=8 ), block margins grow at visibly different rates; at large expected task count ( B=1024 ), all four blocks approach the unit full-batch rate. The main text uses the rate-versus- Bpj collapse because it exposes the controlling expected-count variable more directly.
Diagnostic
measured
prediction
Isotropic anisotropic-signal median matrix error
3.05%
→0 as ε→0
Two-task equilibrium-gap bias exponent
−1.011
−1
Two-task gap-fluctuation exponent
−0.531
−1/2
Hierarchical spatial-coupling exponent
−0.512
−1/2
Hierarchical momentum-coupling exponent
−0.509
−1/2
Projected large-margin drift RMSE
0.00618
0
Appendix
Table 1: Quantitative summary of the controlled numerical checks. Exact and reduced-model validations are kept separate from transfer tests. The isotropic and closure errors are finite-signal/finite-Monte-Carlo quantities whose limiting prediction is zero; the full-categorical label-noise exponent is a stress test rather than a direct validation of the isotropic closure.
Figure 8: Two-task large-batch asymptotics. Both panels use p2=0.05 , η=0.003 , and 64 trajectories. (a) Absolute late-time bias of Δ1−Δ2 from the full-batch target log(p1/p2) versus batch size B on log–log axes; the fitted exponent is −1.011 , compared with the B−1 prediction. (b) Late-time stationary standard deviation of the gap versus B ; the fitted exponent is −0.531 , compared with the B−1/2 diffusion prediction.
Figure 9: Higher-rank perturbation checks at a frozen block-symmetric state. (a) Mean finite-difference derivative norm versus step h : the contrast-to-contrast block decays to numerical zero while contrast-to-average coupling remains O(1) , verifying first-order cancellation. (b) RMS contrast-to-average coupling versus expected task count Bpμ/C ; the fitted log–log slope is −0.512 , matching the predicted −1/2 . (c) The same coupling versus variance-equivalent task count Beffpμ/C under EMA momentum; the fitted slope is −0.509 . Each point uses NMC=8192 draws.
Figure 10: Direct validation of the projected large-margin counting law. Points are Monte Carlo estimates in the longitudinal projected model and dashed curves are the exact finite- B formula. The full-matrix transfer test is shown in Figure 3 (a), where deviations from the projected formula are expected because transverse categorical modes are retained.
Figure 11: Direct check of the moment-matched isotropic covariance closure. The Monte Carlo coefficient follows the predicted p1/2 scaling with fitted exponent 0.502 . This panel validates the stated Gaussian proxy itself; the main text uses the full-categorical simulation as the more demanding stress test.
Muon has recently demonstrated strong empirical performance in large language model training, but the theoretical role of momentum in Muon remains unclear. Existing analyses of Muon either remove momentum to study spectral updates in isolation, or retain momentum without explaining why it improves empirical performance. Our work bridges this gap by showing momentum in Muon acts as a spectral filter. Under a structured signal-plus-perturbation gradient model, we prove that momentum suppresses perturbations while preserving the dominant signal, thereby enlarging the spectral gap between them. This enlarged gap stabilizes the singular subspaces of the matrix passed to Muon's orthogonalization step, making the resulting update more reliable. We further show that applying momentum before orthogonalization achieves provably stronger alignment with the signal component of the gradient than either reversing this order or simply removing momentum. Experiments across diverse tasks, including LLM pretraining, support our theoretical analysis. More broadly, our theory offers a starting point for understanding the benefits of momentum in other matrix-based optimizers.
Xianliang Li, Zihan Zhang, Weiyang Liu +1
1The Institute of Statistical Mathematics · 2The Graduate Institute for Advanced Studies, SOKENDAI · 3National Institute of Informatics +3
Muon is increasingly used for language-model pretraining, yet its large-step dynamics are not captured by the classical edge-of-stability (EoS) picture of gradient descent (GD). In GD, loss neutrality, equal-magnitude update reversal, and marginal stability meet at a single learning-rate-dependent edge. We show that Muon breaks this coupling. For stochastic no-momentum Muon, we derive a coherence-corrected conditional loss-neutral boundary 2ρb/η, while temporal alignment follows a separate geometry. Controlled experiments show that loss balance and temporal alignment respond differently to learning rate and batch size. Across our language model experiments, the 130M Llama-like LLM runs exhibit loss-boundary tracking with weak negative alignment, whereas the studied 1B LLM configuration shows stronger partial cancellation; in both settings, directions remain far from coherent reversal while training continues to improve. These results support a split EoS picture for Muon: a stochastic loss-neutral edge survives, but it is not accompanied by a universal temporal-direction signature. The source code for reproducing the experiments can be found in https://github.com/cyzebra/Muon-Sublates-the-Edge-of-Stability-in-LLM-Pretraining
Yanzhe Chen, Qifang Zhao, Xiaoxiao Xu +1
School of Mathematical Sciences, Shanghai Jiao Tong University, China · Alibaba Inc., Hangzhou, China · School of Mathematical Sciences, Institute of Natural Sciences and MOE-LSC, Shanghai Jiao Tong University, China
Muon orthogonalizes the momentum buffer before each update, replacing its singular values with ones via Newton-Schulz iterations. This simple change lets Muon tolerate far larger learning rates and converge faster than other optimizers, but why? We show that the mechanism is spectral flattening, and develop two results around it. First, we prove that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest, which bottlenecks standard gradient descent. Second, we recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor, with the improvement controlled by the spectrum of the gradient covariance. Extensive experiments validate both results: Muon remains stable at learning rates that cause SGD to diverge within the first few iterations, and reaches accuracy milestones several epochs earlier even at identical step sizes. Taken together, our results offer a principled, geometric explanation for Muon's empirical success.
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia