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 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 MoMuSAM, which turns Muon's momentum from an update variable into a sharpness-probing signal. By using normalized momentum as the perturbation direction, MoMuSAM removes the extra clean-gradient computation and replaces the single-batch perturbation proxy with a temporally aggregated one. We prove that both MuSAM and MoMuSAM attain the same asymptotic first-order stationarity order, and further show that MoMuSAM yields a more stable stochastic linearization than MuSAM. Experiments on standard benchmarks demonstrate that sharpness awareness consistently improves Muon.
Figures & tables
Aspect
Muon
MuSAM
MoMuSAM
Perturbation source
–
∇ℓ(Wt−1)
Mt−1
Per-iteration gradient cost
1
2
1
Directional stability †
–
O(T−3/4)
O(T−1)
First-order stationarity ⋄
O(T−1/4)
O(T−1/4)
O(T−1/4)
Table 1: Comparison between Muon , MuSAM , and MoMuSAM .
Figure 1: Performance comparison of different baselines across multiple datasets, showing test accuracy on FMNIST, CIFAR-10, and CIFAR-100, as well as validation loss on OpenWebText.
Figure 3
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: 3D loss landscapes for different algorithms.
Dataset
SAM
MSAM
VASSO
Muon
MuSAM
MoMuSAM
FMNIST
93.83 ± 0.11
93.74 ± 0.14
93.86 ± 0.08
94.10 ± 0.06
94.27 ± 0.08
94.35 ± 0.17
CIFAR-10
89.27 ± 0.05
89.40 ± 0.07
89.26 ± 0.25
90.72 ± 0.23
91.11 ± 0.18
91.21 ± 0.19
CIFAR-100
68.98 ± 0.29
68.48 ± 0.14
68.96 ± 0.33
74.18 ± 0.20
74.62 ± 0.17
75.14 ± 0.16
Appendix
Table 2: Test Accuracy under different datasets.
Algorithm
λ1
λ2
λ3
λ4
λ5
Muon
30.21
17.89
9.50
7.59
6.08
MuSAM
8.71
6.42
5.42
3.48
2.71
MoMuSAM
5.55
3.18
1.78
1.41
1.06
Appendix
Table 3: Top 5 Hessian eigenvalues.
Wall-clock time test accuracy (%)
Algorithm
200 sec
500 sec
800 sec
1100 sec
1400 sec
1700 sec
2000 sec
SAM
68.68
74.26
78.27
81.94
84.65
85.84
86.53
MSAM
75.02
77.78
82.50
84.23
85.15
87.20
87.86
VASSO
68.40
76.17
79.17
85.26
85.35
86.33
87.64
Muon
83.74
85.67
88.38
88.96
89.34
89.82
90.02
MuSAM
83.44
86.61
87.57
89.10
89.49
89.60
89.91
Appendix
Table 4: Wall-clock time test accuracy of CIFAR-10 for different algorithms.
Sharpness-Aware Minimization (SAM) aims to improve generalization by encouraging insensitivity to small, worst-case parameter perturbations. However, the notion of a "small" perturbation is inherently geometry-dependent: while existing SAM variants have explored a wide range of choices, a clear perspective on which geometries are most effective in practice remains elusive. Recent work on matrix-aware optimization, particularly the Muon optimizer, suggests that respecting the matrix structure of hidden-layer weights can lead to strong empirical performance. Motivated by this, we study matrix-aware geometry in both stages of SAM: we introduce a layerwise spectral inner perturbation for matrix-valued hidden-layer parameters and combine it with either AdamW/SGDW or Muon in the outer update. Across ImageNet-1K experiments on ViT-Small/16 and ResNet-50, we find that the combination of a spectral inner step with a Muon outer step performs consistently strongly, achieving the best validation accuracy on both models among the evaluated methods.
Wenzhi Zhong, Edward Milsom, Michael Murray
Department of Mathematical Sciences University of Bath · Department of Computer Science University of Bath
Adversarial sharpness penalty per iteration of a base optimizer, known as sharpness-aware minimization (SAM), is an accepted technique to improve generalization in deep learning. Recently, Muon has emerged as a dominant optimization algorithm for training large language models (LLMs). It achieves lower sample complexity for heavy-tailed gradient updates - a phenomenon common in LLM training. However, sharpness penalties with Muon updates are yet largely unexplored in the literature. Moreover, the behavior of SAM under heavy-tailed noise is unknown both theoretically and empirically. A potential reason for this gap is the additional cost of extra gradient computation per iteration for penalty perturbations in SAM, which becomes prohibitive for LLMs. This paper introduces Sharpness-aware Orthogonalized Momentum Algorithm (SOMA) to combine a computationally efficient and theoretically grounded perturbation scheme with the Muon algorithm. The key idea is to use stale (purely) stochastic gradient noise for adversarial perturbations. Theoretically, we show that SOMA's sample complexity is invariant to the perturbation direction. We provide explorations and insights into SOMA's robustness to heavy-tailed gradient noise. We extensively explore SOMA's efficacy in (a) mitigating quality degradation due to post-training quantization, and (b) catastrophic forgetting during fine-tuning of LLMs.
Department of Computer Science and Engineering Indraprastha Institute of Information Technology, Delhi, India. · Department of Mathematics Indraprastha Institute of Information Technology, Delhi, India.
We study the Muon optimizer using sharpness and provide the first empirical evidence that the Flat Minima Hypothesis also applies to Muon. Further, we derive Muon's implicit preconditioner to analyze its potential Edge of Stability behavior. Experiments on CIFAR-10 show sharpness strongly correlates with generalization gap within each optimizer (Pearson r=0.45-0.69, p≤10−4), confirming the hypothesis holds for Muon similarly to Adam and SGD. However, Muon operates above the classical sharpness-based Adaptive Edge of Stability threshold that bounds diagonally-preconditioned optimizers like Adam and RMSprop, yet remains stable. Our findings suggest Muon operates in a fundamentally different regime than diagonally-adaptive methods. Defining the exact stability boundary for Muon remains an open question. All code is available at https://github.com/NikoCovic/DL-project.
Niko Covic, Baraq Lipshitz, Alexander Padula +1
Department of Informatics, ETH Zurich, Switzerland.