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.
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
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia