Sharpness-Aware Muon: Matrix-Aware Optimization with Momentum-Guided Perturbations
Zhiwei Zhai, Wenjing Yan, Ying-Jun Angela Zhang
原文摘要
Muon is a matrix-aware optimizer that combines momentum with orthogonalization to produce structured update directions with strong empirical behavior. However, it only shapes the geometry of the update itself and does not explicitly control neighborhood worst-case loss or bias training toward flat regions, which are often linked to stronger generalization. To address this limitation, we introduce the sharpness-aware formulation of Muon. We first develop \textsc{MuSAM}, which incorporates Sharpness-Aware Minimization (SAM)-style perturbation into Muon's orthogonalized update and establishes a principled bridge between sharpness-aware training and matrix-aware optimization. We then show that this formulation faces an efficiency--stability tension: constructing the perturbation requires an extra gradient evaluation, while the perturbation itself is generated from a single minibatch gradient and is therefore noisy. To resolve this issue, we propose \textsc{MoMuSAM}, which turns Muon's momentum from an update variable into a sharpness-probing signal. By using normalized momentum as the perturbation direction, \textsc{MoMuSAM} removes the extra clean-gradient computation and replaces the single-batch perturbation proxy with a temporally aggregated one. We prove that both \textsc{MuSAM} and \textsc{MoMuSAM} attain the same asymptotic first-order stationarity order, and further show that \textsc{MoMuSAM} yields a more stable stochastic linearization than \textsc{MuSAM}. Experiments on standard benchmarks demonstrate that sharpness awareness consistently improves Muon.
cs.LG