Sharpness-aware Orthogonalized Momentum Algorithms

Geetika; Shivani Singh Chauhan; Gauranshi Gupta; Anant Jain; Kaushik Kalyanaraman; Bapi Chatterjee

cs.LGOct 23, 2026

Sharpness-aware Orthogonalized Momentum Algorithms

Authors: Geetika, Shivani Singh Chauhan, Gauranshi Gupta, Anant Jain, Kaushik Kalyanaraman, Bapi Chatterjee

Organizations: Department of Computer Science and Engineering Indraprastha Institute of Information Technology, Delhi, India. · Department of Mathematics Indraprastha Institute of Information Technology, Delhi, India.

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Performance under anisotropic heavy-tailed noise (Condition no., $\kappa_{\varepsilon}=10$ ). Top: validation accuracy. Bottom: training loss (log scale). Red , Blue and Green lines indicate AdamW -based, SGD -based and Muon -based methods, respectively. Lines show the median over seeds, and shaded bands show the interquartile range.

Panel of Figure 2 — Figure 2: Training curves of the fifteen optimizers (mean over 3 seeds) on FineWeb-Edu NanoGPT, FineWeb-Edu Pythia and the CIFAR-10 Airbench CNN: training loss, and validation perplexity (language models) or validation accuracy (CNN), against the optimizer step (100 steps per epoch). Red , Blue and Green lines indicate AdamW -based, SGD -based and Muon -based methods, respectively.

	FineWeb-Edu NanoGPT			FineWeb-Edu Pythia
Optimizer	FP	INT4	$\Delta\%$	FP	INT4	$\Delta\%$
Muon	22.96 $\pm$ 0.01	23.87 $\pm$ 0.01	+3.98 $\pm$ 0.06	33.60 $\pm$ 0.01	37.28 $\pm$ 0.10	+10.96 $\pm$ 0.27
Ssm	23.09 $\pm$ 0.00	23.87 $\pm$ 0.03	+3.38 $\pm$ 0.11	33.65 $\pm$ 0.02	37.22 $\pm$ 0.01	+10.61 $\pm$ 0.08
FP-Soma	23.07 $\pm$ 0.00	23.86 $\pm$ 0.03	+3.42 $\pm$ 0.09	33.65 $\pm$ 0.02	37.08 $\pm$ 0.10	+10.20 $\pm$ 0.26
Soma	22.95 $\pm$ 0.01	23.86 $\pm$ 0.03	+3.96 $\pm$ 0.12	33.66 $\pm$ 0.01	37.21 $\pm$ 0.08	+10.55 $\pm$ 0.19
Rsm	22.96 $\pm$ 0.01	23.87 $\pm$ 0.01	+3.98 $\pm$ 0.02	33.61 $\pm$ 0.02	37.34 $\pm$ 0.10	+11.13 $\pm$ 0.26

Table 1: INT4 post-training quantization results on held-out FineWeb-Edu perplexity (lower is better), comparing full precision (FP) and 4-bit AWQ (INT4). Results show mean

\pm

std over 3 pretraining seeds (INT4 averaged over 3 calibration seeds); bold indicates the lowest value per column. Row colour marks the base optimizer family: Muon -based ( Green ), AdamW -based ( Red ), SGD -based ( Blue ).

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Algorithm	$\alpha=1.1$		$\alpha=1.6$		$\alpha=2.0$		$\alpha=3.0$
	best acc.	$\rho^{\star}$	best acc.	$\rho^{\star}$	best acc.	$\rho^{\star}$	best acc.	$\rho^{\star}$
Muon	38.1	–	42.9	–	44	–	44.3	–
Ssm	38.6	0.5	43.7	0.5	46.3	0.5	47	0.5
FP-Soma	39.8	1	44.2	0.5	46.9	0.5	46.7	0.5
Soma	39.1	1	45	1	46.3	1	46.5	0.5
Rsm	38.3	0.01	43.3	0.002	44.5	0.1	44.8	1

Appendix

Table 2: Validation accuracy (%): best value over the perturbation radius

\rho

and the radius

\rho^{\star}

where it is attained (median of 8 seeds; Frobenius-ball methods pool the global and per-layer radius modes). Row colour marks the base optimizer family: Muon -based ( Green ), AdamW -based ( Red ), SGD -based ( Blue ); baselines have no

\rho^{\star}

. Best value per

\alpha

in bold.

Algorithm	$\alpha=1.1$		$\alpha=1.6$		$\alpha=2.0$		$\alpha=3.0$
	best loss	$\rho^{\star}$	best loss	$\rho^{\star}$	best loss	$\rho^{\star}$	best loss	$\rho^{\star}$
Muon	1.45	–	0.372	–	0.184	–	0.078	–
Ssm	1.45	0.02	0.372	0.002	0.183	0.001	0.0779	0.001
FP-Soma	1.44	0.2	0.372	0.01	0.183	0.001	0.0768	0.02
Soma	1.43	0.2	0.371	0.02	0.184	0.001	0.0783	0.01
Rsm	1.45	0.1	0.371	0.2	0.184	0.001	0.0779	0.05

Appendix

Table 3: Training loss: best value over the perturbation radius

\rho

and the radius

\rho^{\star}

\rho^{\star}

. Best value per

\alpha

in bold.

Panel of Figure 3 — Figure 3: Performance under anisotropic heavy-tailed noise (Condition no., $\kappa_{\varepsilon}=10$ ). Optimality gap (log scale) vs $\rho$ . Red , Blue and Green lines indicate AdamW -based, SGD -based and Muon -based methods, respectively.

Algorithm	$\alpha=1.1$		$\alpha=1.6$		$\alpha=2.0$		$\alpha=3.0$
	gap	$\rho^{\star}$	gap	$\rho^{\star}$	gap	$\rho^{\star}$	gap	$\rho^{\star}$
Muon	0.413	–	0.062	–	0.0316	–	0.0165	–
Ssm	0.265	1	0.0473	0.5	0.0253	0.5	0.0141	0.5
FP-Soma	0.252	1	0.0457	0.5	0.0247	0.5	0.0143	0.2
Soma	0.263	1	0.0551	0.5	0.0299	0.2	0.0163	0.1
Rsm	0.408	1	0.0609	1	0.0311	1	0.0162	1

Appendix

Table 4: Optimality gap: best value over the perturbation radius

\rho

and the radius

\rho^{\star}

\rho^{\star}

. Best value per

\alpha

in bold.

Panel of Figure 4 — Figure 4: Performance under isotropic heavy-tailed noise (nonlinear setting)(Condition no., $\kappa_{\epsilon}=1$ ). Top: validation accuracy. Bottom: training loss (log scale). Red , Blue and Green lines indicate AdamW -based, SGD -based and Muon -based methods, respectively.

Algorithm	$\alpha=1.1$		$\alpha=1.6$		$\alpha=2.0$		$\alpha=3.0$
	best acc.	$\rho^{\star}$	best acc.	$\rho^{\star}$	best acc.	$\rho^{\star}$	best acc.	$\rho^{\star}$
Muon	36.1	–	43.3	–	45.2	–	43.8	–
Ssm	36.6	0.2	44.2	0.5	45.5	0.005	46.2	0.5
FP-Soma	37.8	1	44.7	1	46.6	0.5	46.1	0.5
Soma	38.2	1	44.8	1	46.6	0.5	46.6	0.5
Rsm	36.3	0.05	43.6	0.1	45.2	0.2	44.8	0.2

Appendix

Table 5: Validation accuracy (%): best value over the perturbation radius

\rho

and the radius

\rho^{\star}

\rho^{\star}

. Best value per

\alpha

in bold.

Algorithm	$\alpha=1.1$		$\alpha=1.6$		$\alpha=2.0$		$\alpha=3.0$
	best loss	$\rho^{\star}$	best loss	$\rho^{\star}$	best loss	$\rho^{\star}$	best loss	$\rho^{\star}$
Muon	1.58	–	0.482	–	0.233	–	0.114	–
Ssm	1.58	0.02	0.482	0.001	0.236	0.001	0.113	0.01
FP-Soma	1.57	0.2	0.481	0.001	0.234	0.02	0.112	0.02
Soma	1.56	0.5	0.483	0.002	0.234	0.001	0.113	0.01
Rsm	1.58	0.005	0.482	0.05	0.233	0.5	0.113	0.2

Appendix

Table 6: Training loss: best value over the perturbation radius

\rho

and the radius

\rho^{\star}

\rho^{\star}

. Best value per

\alpha

in bold.

Panel of Figure 5 — Figure 5: Performance under isotropic heavy-tailed noise (linear setting)(Condition no., $\kappa_{\epsilon}=1$ ). Optimality gap (log scale) vs $\rho$ . Red , Blue and Green lines indicate AdamW -based, SGD -based and Muon -based methods, respectively.

Algorithm	$\alpha=1.1$		$\alpha=1.6$		$\alpha=2.0$		$\alpha=3.0$
	gap	$\rho^{\star}$	gap	$\rho^{\star}$	gap	$\rho^{\star}$	gap	$\rho^{\star}$
Muon	1.01	–	0.231	–	0.114	–	0.0552	–
Ssm	0.648	1	0.141	1	0.0733	1	0.0412	0.5
FP-Soma	0.654	1	0.147	1	0.0823	0.5	0.0413	0.5
Soma	0.63	1	0.176	0.5	0.086	0.5	0.0486	0.5
Rsm	1	1	0.224	1	0.108	1	0.0527	1

Appendix

Table 7: Optimality gap: best value over the perturbation radius

\rho

and the radius

\rho^{\star}

\rho^{\star}

. Best value per

\alpha

in bold.

		Tail index $\alpha$
Direction	Geometry	$1.1$	$1.6$	$2.0$	$3.0$
Random	Frobenius	38.3	43.4	44.5	44.7
Random	Spectral	38.5	44.7	46.5	47.2
Sam - Muon	Frobenius	38.3	43.2	45.0	45.6
Sam - Muon	Spectral	38.6	43.7	46.3	47.0
Soma	Frobenius	38.2	43.3	45.0	45.6

Appendix

Table 8: Nonlinear anisotropic setting: best validation accuracy (%,

\uparrow

) with a Frobenius or a spectral perturbation, for each tail index

\alpha

. The better geometry for each method and

\alpha

is in bold; ties are not marked.

		Tail index $\alpha$
Direction	Geometry	$1.1$	$1.6$	$2.0$	$3.0$
Random	Frobenius	0.410	0.0612	0.0312	0.0163
Random	Spectral	0.347	0.0576	0.0301	0.0163
Sam - Muon	Frobenius	0.404	0.0581	0.0294	0.0149
Sam - Muon	Spectral	0.265	0.0473	0.0253	0.0141
Soma	Frobenius	0.399	0.0592	0.0303	0.0163

Appendix

Table 9: Linear anisotropic setting: best optimality gap (

\downarrow

) with a Frobenius or a spectral perturbation, for each tail index

\alpha

. The better geometry for each method and

\alpha

is in bold; ties are not marked.

			Tail index $\alpha$
Geometry	Direction	Method	$1.1$	$1.6$	$2.0$	$3.0$
Frobenius	Fresh	Sam - Muon	0.3828	0.4316	0.4502	0.4561
Frobenius	Stale	S 2 -Sam - Muon	0.3818	0.4326	0.4502	0.4492
Spectral	Fresh	Ssm	0.3857	0.4365	0.4629	0.4697
Spectral	Stale	S 2 -Sam - Muon	0.3887	0.4443	0.4678	0.4570
Spectral	Fresh	FP-Soma	0.3984	0.4424	0.4688	0.4668

Appendix

Table 10: Nonlinear anisotropic setting: best validation accuracy (%,

\uparrow

) with fresh or stale perturbation direction, for each tail index

\alpha

. The better direction for each geometry and

\alpha

is in bold; ties are not marked.

			Tail index $\alpha$
Geometry	Direction	Method	$1.1$	$1.6$	$2.0$	$3.0$
Frobenius	Fresh	Sam - Muon	0.4037	0.0581	0.0294	0.0149
Frobenius	Stale	S 2 -Sam - Muon	0.4068	0.0601	0.0312	0.0165
Spectral	Fresh	Ssm	0.2648	0.0473	0.0253	0.0141
Spectral	Stale	S 2 -Sam - Muon	0.2846	0.0588	0.0309	0.0165
Spectral	Fresh	FP-Soma	0.2522	0.0457	0.0247	0.0143

Appendix

Table 11: Linear anisotropic setting: best optimality gap (

\downarrow

) with a fresh or a stale perturbation direction, for each tail index

\alpha

. The better direction for each geometry and

\alpha

is in bold; ties are not marked.

	FineWeb-Edu NanoGPT		FineWeb-Edu Pythia		CIFAR-10 Airbench CNN
Optimizer	Train loss	Val ppl	Train loss	Val ppl	Train loss	Val acc (%)
Muon	3.303 $\pm$ 0.004	24.41 $\pm$ 0.02	3.558 $\pm$ 0.003	35.72 $\pm$ 0.03	0.266 $\pm$ 0.002	87.27 $\pm$ 0.08
Ssm	3.305 $\pm$ 0.004	24.55 $\pm$ 0.03	3.559 $\pm$ 0.003	35.78 $\pm$ 0.03	0.271 $\pm$ 0.003	87.63 $\pm$ 0.34
FP-Soma	3.305 $\pm$ 0.004	24.54 $\pm$ 0.03	3.560 $\pm$ 0.003	35.79 $\pm$ 0.03	0.270 $\pm$ 0.001	87.37 $\pm$ 0.44
Soma	3.303 $\pm$ 0.004	24.40 $\pm$ 0.01	3.562 $\pm$ 0.003	35.79 $\pm$ 0.03	0.269 $\pm$ 0.002	87.51 $\pm$ 0.27
Rsm	3.303 $\pm$ 0.004	24.41 $\pm$ 0.02	3.558 $\pm$ 0.003	35.74 $\pm$ 0.03	0.266 $\pm$ 0.004	87.22 $\pm$ 0.07

Appendix

Table 12: Continued-pretraining results of the fifteen optimizers after the last epoch (FineWeb-Edu NanoGPT 15, FineWeb-Edu Pythia 14, CIFAR-10 Airbench CNN 7 epochs). Lower is better for train loss and validation perplexity (Val ppl), higher for validation accuracy (Val acc); the best validation value of each dataset is in bold. Mean

\pm

std over 3 seeds. Row colour marks the base optimizer family: Muon -based ( Green ), AdamW -based ( Red ), SGD -based ( Blue ).

Panel of Figure 6 — Figure 6: Catastrophic forgetting of nanoGPT (GPT-2 small) models pretrained with each optimizer, after fine-tuning on four datasets. Bars: rise $\Delta$ PT of the held-out FineWeb-Edu loss at a common fine-tuning loss (lower is better); diamond: average over datasets; error bars: standard deviation over 3 pretraining seeds. † starting model clearly weaker than AdamW’s (not comparable).

Panel of Figure 7 — Figure 7: Catastrophic forgetting of Pythia-70M models pretrained with each optimizer, after fine-tuning on four datasets. Bars: rise $\Delta$ PT of the held-out FineWeb-Edu loss at a common fine-tuning loss (lower is better); diamond: average over datasets; error bars: standard deviation over 3 pretraining seeds. † starting model clearly weaker than AdamW’s (not comparable).

Optimizer	$L_{\mathrm{before}}$	codeparrot	stackmathqa	musicpile	tulu3	Average	$\Delta\%_{\mathrm{adj}}$
		$L_{\mathrm{after}}$ by fine-tuning dataset
Muon	3.144 $\pm$ 0.000	3.516 $\pm$ 0.074	3.526 $\pm$ 0.004	3.259 $\pm$ 0.000	3.332 $\pm$ 0.001	3.409 $\pm$ 0.020	4.20 $\pm$ 0.31
Ssm	3.147 $\pm$ 0.000	3.386 $\pm$ 0.003	3.561 $\pm$ 0.006	3.251 $\pm$ 0.001	3.331 $\pm$ 0.040	3.382 $\pm$ 0.009	3.73 $\pm$ 0.15
FP-Soma	3.146 $\pm$ 0.000	3.393 $\pm$ 0.004	3.515 $\pm$ 0.084	3.251 $\pm$ 0.001	3.308 $\pm$ 0.002	3.367 $\pm$ 0.022	3.50 $\pm$ 0.35
Soma	3.144 $\pm$ 0.001	3.491 $\pm$ 0.086	3.511 $\pm$ 0.005	3.258 $\pm$ 0.002	3.330 $\pm$ 0.001	3.397 $\pm$ 0.023	4.03 $\pm$ 0.35
Rsm	3.144 $\pm$ 0.001	3.558 $\pm$ 0.009	3.526 $\pm$ 0.007	3.260 $\pm$ 0.000	3.333 $\pm$ 0.001	3.419 $\pm$ 0.004	4.36 $\pm$ 0.06

Appendix

Table 13: Catastrophic Forgetting of nanoGPT (GPT-2 small) models pretrained with different optimizers: held-out FineWeb-Edu loss before (

L_{\mathrm{before}}

) and after (

L_{\mathrm{after}}

) fine-tuning on four datasets (lower is better).

\Delta\%_{\mathrm{adj}}

is the smoothed relative increase of the average

L_{\mathrm{after}}

. Mean

\pm

std over 3 seeds. Row colour marks the base optimizer family: Muon -based ( Green ), AdamW -based ( Red ), SGD -based ( Blue ).

Optimizer	$L_{\mathrm{before}}$	codeparrot	stackmathqa	musicpile	tulu3	Average	$\Delta\%_{\mathrm{adj}}$
		$L_{\mathrm{after}}$ by fine-tuning dataset
Muon	3.515 $\pm$ 0.000	4.361 $\pm$ 0.005	4.464 $\pm$ 0.015	3.732 $\pm$ 0.001	3.888 $\pm$ 0.056	4.111 $\pm$ 0.013	8.45 $\pm$ 0.19
Ssm	3.516 $\pm$ 0.001	4.360 $\pm$ 0.002	4.478 $\pm$ 0.026	3.733 $\pm$ 0.001	3.856 $\pm$ 0.004	4.107 $\pm$ 0.006	8.36 $\pm$ 0.09
FP-Soma	3.516 $\pm$ 0.001	4.361 $\pm$ 0.002	4.480 $\pm$ 0.016	3.733 $\pm$ 0.001	3.852 $\pm$ 0.001	4.107 $\pm$ 0.004	8.36 $\pm$ 0.07
Soma	3.517 $\pm$ 0.000	4.360 $\pm$ 0.004	4.481 $\pm$ 0.018	3.734 $\pm$ 0.001	3.855 $\pm$ 0.004	4.107 $\pm$ 0.003	8.36 $\pm$ 0.05
Rsm	3.514 $\pm$ 0.000	4.358 $\pm$ 0.007	4.468 $\pm$ 0.027	3.732 $\pm$ 0.001	3.857 $\pm$ 0.002	4.104 $\pm$ 0.008	8.35 $\pm$ 0.12

Appendix

Table 14: Catastrophic Forgetting of Pythia-70M models pretrained with different optimizers: held-out FineWeb-Edu loss before (

L_{\mathrm{before}}

) and after (

L_{\mathrm{after}}

) fine-tuning on four datasets (lower is better).

\Delta\%_{\mathrm{adj}}

is the smoothed relative increase of the average

L_{\mathrm{after}}

. Mean

\pm

std over 3 seeds. Row colour marks the base optimizer family: Muon -based ( Green ), AdamW -based ( Red ), SGD -based ( Blue ).

Optimizer	FineWeb-Edu NanoGPT	FineWeb-Edu Pythia	CIFAR-10 Airbench CNN
	s / epoch	s / epoch	s / epoch
Muon	265.2 $\pm$ 3.6	287.2 $\pm$ 2.0	23.6 $\pm$ 0.3
Ssm	772.5 $\pm$ 6.7	236.8 $\pm$ 1.9	23.9 $\pm$ 0.4
FP-Soma	641.8 $\pm$ 129.0	240.8 $\pm$ 4.2	23.9 $\pm$ 0.6
Soma	272.1 $\pm$ 1.8	294.4 $\pm$ 4.8	23.2 $\pm$ 2.1
Rsm	267.2 $\pm$ 1.1	298.2 $\pm$ 2.6	23.1 $\pm$ 0.5

Appendix

Table 15: Pretraining runtime of the optimizers: wall-clock seconds per epoch (100 optimizer steps). Row colour marks the base optimizer family: Muon -based ( Green ), AdamW -based ( Red ), SGD -based ( Blue ).

Optimizer	FineWeb-Edu NanoGPT	FineWeb-Edu Pythia
	s / run	s / run
Muon	48.6 $\pm$ 18.2	25.4 $\pm$ 7.0
Ssm	82.0 $\pm$ 33.7	58.6 $\pm$ 21.4
FP-Soma	84.2 $\pm$ 49.9	50.9 $\pm$ 20.2
Soma	36.0 $\pm$ 2.5	27.8 $\pm$ 6.8
Rsm	31.2 $\pm$ 0.1	25.7 $\pm$ 5.8

Appendix

Table 16: Fine-tuning runtime of the optimizers: wall-clock seconds per fine-tuning run, taken from the fixed-base fine-tuning study (per job the median over its successful runs, then mean

\pm

std over 3 pretrained bases

\times

3 seeds). Runs shared the GPUs with other jobs, so only large differences are meaningful. Row colour marks the base optimizer family: Muon -based ( Green ), AdamW -based ( Red ), SGD -based ( Blue ).

	nanoGPT			Pythia-70M			CIFAR-10 CNN
Optimizer	lr	wd	$\rho$	lr	wd	$\rho$	lr	wd	$\rho$
Muon	$4.24\!\times\!10^{-3}$	$0.01$	–	$8.48\!\times\!10^{-4}$	$0.3$	–	$0.03$	$0$	–
Ssm	$4.24\!\times\!10^{-3}$	$0.01$	$0.01$	$8.48\!\times\!10^{-4}$	$0.3$	$3\!\times\!10^{-4}$	$0.03$	$0$	$0.01$
FP-Soma	$4.24\!\times\!10^{-3}$	$0.01$	$0.01$	$8.48\!\times\!10^{-4}$	$0.3$	$3\!\times\!10^{-4}$	$0.03$	$0$	$0.01$
Soma	$4.24\!\times\!10^{-3}$	$0.01$	$3\!\times\!10^{-4}$	$8.48\!\times\!10^{-4}$	$0.3$	$3\!\times\!10^{-4}$	$0.03$	$0$	$3\!\times\!10^{-3}$
Rsm	$4.24\!\times\!10^{-3}$	$0.01$	$0.01$	$8.48\!\times\!10^{-4}$	$0.3$	$9.5\!\times\!10^{-4}$	$0.03$	$0$	$0.1$

Appendix

Table 17: Selected learning rate, weight decay and perturbation radius

\rho

of the continued-pretraining experiments.

	nanoGPT		Pythia-70M
Optimizer	lr	$\rho$	lr	$\rho$
Muon	$10^{-2}$	–	$10^{-3}$	–
Ssm	$10^{-2}$	$10^{-2}$	$10^{-3}$	$3\!\times\!10^{-4}$
FP-Soma	$10^{-2}$	$10^{-2}$	$10^{-3}$	$3\!\times\!10^{-4}$
Soma	$10^{-2}$	$3\!\times\!10^{-4}$	$10^{-3}$	$3\!\times\!10^{-4}$
Rsm	$10^{-2}$	$10^{-2}$	$10^{-3}$	$9.5\!\times\!10^{-4}$

Appendix

Table 18: Selected learning rate and perturbation radius

\rho

for fine-tuning.

Deep Learning Optimization

Deep learning optimization develops algorithms and training strategies that help neural networks converge, generalize, and use compute efficiently.

See all 206 papers

Oct 27, 2026

Sharpness-Aware Muon: Matrix-Aware Optimization with Momentum-Guided Perturbations

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.

Zhiwei Zhai, Wenjing Yan, Ying-Jun Angela Zhang

Department of Information Engineering, The Chinese University of Hong Kong.

Muon Optimizer Sharpness-Aware Minimization Nonconvex Stochastic Optimization

Ask PDF

May 26, 2026

MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training

The Muon optimizer has recently offered a promising alternative to AdamW for large language model training, leveraging matrix orthogonalization to produce geometry-aware updates. However, like all first-order methods, Muon can become trapped in sharp local minima. In this work, we present MONA, an optimizer that bridges Muon's orthogonalization framework with curvature-aware acceleration. MONA adds an acceleration term directly into Muon's gradient processing pipeline. This term is calculated from the exponential moving average of gradient differences. We provide a detailed convergence analysis for MONA, showing that the acceleration term introduces curvature-sensitive corrections while preserving Muon's spectral-norm regularization. Empirically, MONA achieves better convergence and downstream task performance compared to both Muon and AdamW across three scales of Mixture-of-Experts pretraining, spanning from 1B to 68B parameters, with the largest model trained on 1 trillion tokens. Furthermore, we conduct supervised fine-tuning on the MOE-68B-A3B model and evaluate it on general capability, mathematical reasoning, and code generation benchmarks, where MONA achieves SOTA performance.

Jiacheng Li, Jianchao Tan, Hongtao Xu +5

Meituan, Beijing, China

Language Model Pretraining Gradient Descent Muon Optimizer Accelerated Gradient Methods

Ask PDF

Jul 28, 2026

Sharpness-Aware Minimization and Muon: Robustness under the Spectral Norm

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

Sharpness-Aware Minimization Muon Optimizer

Ask PDF

Abstract

Deep Learning Optimization

Explore similar work

Sharpness-Aware Muon: Matrix-Aware Optimization with Momentum-Guided Perturbations

MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training

Sharpness-Aware Minimization and Muon: Robustness under the Spectral Norm