Phases of Muon: When Muon Eclipses SignSGD
Elliot Paquette, Noah Marshall, Lucas Benigni, Guangyuan Wang, Atish Agarwala, Courtney Paquette
Muon and related spectral optimizers have recently demonstrated strong empirical performance as scalable stochastic methods, often outperforming Adam. Yet their behavior remains poorly understood. We analyze stochastic Muon, which orthogonalizes an exponential moving average of minibatch gradients, on a high-dimensional matrix-valued least squares problem, and we compare it with stochastic SignSGD, a simple baseline comparable to Adam. We derive explicit deterministic dynamics that provide a tractable framework for studying learning behavior. We show that Muon has non-linear sensitivity to the effective batch $B_{\rm eff}=B(1+β)/(1-β)$. With exact whitening and small $B_{\rm eff}$ Muon fails to resolve smaller eigenmodes, slowing down convergence. Momentum and finite Newton--Shulz iterations can act as a ridge which lessen this effect. This behavior is in contrast with SignSGD which has no transition, leading to different optimal learning rates, but with matching convergence rate up to a constant factor in the isotropic setting. For anisotropic data, large batch setting Muon performs a square-root preconditioning with respect to the data covariance spectrum while SignSGD provides no preconditioning for generic covariances. We analyze time to target loss in a power law covariance model and derive a phase diagram in the data exponent $α$ and target exponent $ζ$. Muon is favored for steep covariance spectra with rich targets ($α>2$, $ζ<α+1$), SignSGD is favored for flat spectra ($α<1$), and in other regions the winner depends on how the target loss scales with dimension. Our results shed light on why replacing Adam with Muon is not trivial even from a pure optimization perspective.
尚无 AI 解读，展示原文摘要