AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models

Arash Lagzian; Paniz Halvachi; Junming Zhang; Zhouhan Lin; Dianbo Liu

cs.LGOct 1, 2026

AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models

Authors: Arash Lagzian, Paniz Halvachi, Junming Zhang, Zhouhan Lin, Dianbo Liu

Organizations: National University of Singapore, Singapore · Independent Researcher · LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University, Shanghai, China

Abstract

Muon improves large-scale training by applying a spectral-norm steepest-descent update to matrix parameters, but practical models also contain parameter blocks that do not fit dense-matrix geometry. One important case is the tied vocabulary table, which appears in language models and other token generators and can receive multiple structurally different gradient sources, from sparse input lookups to dense output-classifier updates. In the reference recipe these blocks are handed to an auxiliary AdamW optimizer, which restores second-moment state and updates the aliased table as a generic tensor. We propose AF-Muon, an AdamW-free extension of Muon that keeps the Muon matrix update for hidden weight matrices while using a support-aware finite-cap linear minimization oracle for tied vocabulary tables and an RMS-normalized update for one-dimensional auxiliary parameters. AF-Muon therefore trains every parameter class with a single first-moment buffer and no second-moment state, saving around 20% optimizer-state memory relative to Hybrid Muon in our benchmark. Across nine tied-token settings - decoder-only language models from 124M to 1B parameters, a fully shared T5-style encoder-decoder, and ImageGPT-style image-token, protein, and sparse-MoE variants, spanning text, image, and protein-sequence data - AF-Muon improves mean validation loss and perplexity over both Hybrid Muon and a SCION-style Sign endpoint. Long-horizon runs and hyperparameter sensitivity studies confirm the gain is robust, and identical-momentum diagnostics attribute it to the finite cap, which preserves more within-row magnitude than Sign while bounding the coordinate concentration of row-RMS. These results identify tied vocabulary tables as a distinct optimizer geometry and yield a robust AdamW-free Muon variant across models, modalities, and architectures, with about 1% step-time overhead in matched training.

Figures & tables

Figure 1: AF-Muon tied-table geometry. Top: a hard-tied vocabulary table serves both sparse input lookup and dense output-softmax roles, so their gradients merge on one tensor. AF-Muon assigns one first-moment LMO per parameter class: Muon for hidden matrices, a finite-cap clipped-row LMO for the tied table, and RMS normalization for 1-D vectors. Bottom: the tied-table feasible set is the intersection of a fixed row-

\ell_{2}

ball and a coordinate cap box. Varying

c

interpolates between the Sign endpoint (

c=1

), the finite-cap AF-Muon update (

1<c<\sqrt{d}

), and row-RMS (

c\geq\sqrt{d}

Model / setting	Optimizer	Train loss $\downarrow$	Val. loss $\downarrow$	Val. PPL / bpd $\downarrow$
NanoGPT	Hybrid Muon	3.5150 $\pm$ 0.0028	3.5177 $\pm$ 0.0032	33.71 $\pm$ 0.11
	SCION-style Sign	3.5125 $\pm$ 0.0107	3.5188 $\pm$ 0.0124	33.75 $\pm$ 0.42
	AF-Muon	3.4898 $\pm$ 0.0011	3.4919 $\pm$ 0.0019	32.85 $\pm$ 0.06
SmolLM2-135M	Hybrid Muon	3.5088 $\pm$ 0.0028	3.5065 $\pm$ 0.0021	33.33 $\pm$ 0.07
	SCION-style Sign	3.4727 $\pm$ 0.0069	3.4706 $\pm$ 0.0040	32.16 $\pm$ 0.13
	AF-Muon	3.4453 $\pm$ 0.0006	3.4446 $\pm$ 0.0036	31.33 $\pm$ 0.11

Table 1: Main optimizer comparison across tied-token models and domains. Results are mean

\pm

SD over three paired seeds. All rows compare matched tied-table optimizer arms under the same model, data, token budget, precision, batching, and matrix-update protocol; detailed settings and per-seed results are reported in Appendix J . NanoGPT, SmolLM2, Qwen2.5, Llama, T5, MoE, and protein rows report validation perplexity; ImageGPT reports validation bits/dim.

Figure 2: Decoder-only tied-vocabulary training dynamics on FineWeb for SmolLM2-135M (top) and Llama-3.2-1B (bottom). Each row reports metrics under the matched protocol in Section 3.1 . Curves show the mean over three paired seeds. Llama model long horizon curves reported in Appendix J.4.1

Figure 3: Validation-loss curves beyond decoder-only language modeling. Each panel shows one tied-token setting with its own training-token scale. AF-Muon reaches the lowest final validation loss in all four settings; full training-loss and perplexity curves are in the related appendix.

Figure 4: NanoGPT

2\times 2

factorial ablation over training tokens (billions) for three paired seeds.

Figure 5: NanoGPT batch-size sensitivity at a matched 458.8M-token comparison. The x-axis varies the number of tokens per optimizer update. AF-Muon leads across all tested batch sizes.

Appendix figures & tables52 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Seed	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	43	3.5182	3.5212	33.83
Hybrid Muon	44	3.5133	3.5152	33.62
Hybrid Muon	45	3.5134	3.5165	33.67
SCION-style Sign	43	3.5185	3.5272	34.03
SCION-style Sign	44	3.5001	3.5046	33.27
SCION-style Sign	45	3.5187	3.5247	33.94

Appendix

Table 2: Per-seed NanoGPT-style decoder-only results. All runs use tied input/output embeddings and 2.67B training tokens. The final row reports a paired two-sided

t

-test over validation loss between AF-Muon and the strongest non-AF baseline by mean validation loss. The validation-loss entry is the mean paired gain, computed as baseline minus AF-Muon, so positive values indicate that AF-Muon is better.

Method	Seed	Tied RMS	Tied max	Matrix RMS	Aux WD
Hybrid Muon	43	0.0427	0.2500	0.0463	0.01
Hybrid Muon	44	0.0429	0.2543	0.0465	0.01
Hybrid Muon	45	0.0431	0.2568	0.0464	0.01
SCION-style Sign	43	0.5528	7.8773	0.0473	0.00
SCION-style Sign	44	0.6238	7.8289	0.0466	0.00
SCION-style Sign	45	0.5548	7.8897	0.0471	0.00

Appendix

Table 3: Final NanoGPT tied-vocabulary diagnostics. Hybrid Muon uses auxiliary AdamW weight decay

0.01

; SCION-style Sign and AF-Muon use zero auxiliary/tied weight decay.

Figure 6: NanoGPT-style decoder-only training dynamics on FineWeb. Curves show the mean over three paired seeds, and shaded regions denote

\pm 1

standard error of the mean.

Method	Seed	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	43	3.5118	3.5047	33.27
Hybrid Muon	44	3.5082	3.5058	33.31
Hybrid Muon	45	3.5063	3.5088	33.41
SCION-style Sign	43	3.4777	3.4744	32.28
SCION-style Sign	44	3.4757	3.4663	32.02
SCION-style Sign	45	3.4649	3.4712	32.17

Appendix

Table 4: Per-seed SmolLM2-135M decoder-only results. All runs use tied input/output embeddings and 2.49B training tokens. The final row reports a paired two-sided

t

-test over validation loss between AF-Muon and the strongest non-AF baseline by mean validation loss; the gain is baseline minus AF-Muon, so positive values indicate that AF-Muon is better.

Method	Seed	Tied RMS	Tied max	Matrix RMS	Aux WD
Hybrid Muon	43	0.0549	0.3561	0.0561	0.01
Hybrid Muon	44	0.0552	0.3029	0.0559	0.01
Hybrid Muon	45	0.0551	0.3498	0.0560	0.01
SCION-style Sign	43	0.6611	10.0026	0.0562	0.00
SCION-style Sign	44	0.6524	10.0348	0.0560	0.00
SCION-style Sign	45	0.6668	10.0104	0.0560	0.00

Appendix

Table 5: Final SmolLM2-135M tied-vocabulary diagnostics. Hybrid Muon uses auxiliary/tied AdamW weight decay

0.01

; SCION-style Sign and AF-Muon use zero auxiliary/tied weight decay.

Method	Seed	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	43	3.6186	3.6349	37.90
Hybrid Muon	44	3.6293	3.6406	38.12
Hybrid Muon	45	3.6329	3.6424	38.18
SCION-style Sign	43	3.5935	3.6075	36.87
SCION-style Sign	44	3.5990	3.6133	37.09
SCION-style Sign	45	3.6026	3.6119	37.04

Appendix

Table 6: Per-seed Qwen2.5-0.5B short-horizon results. Values are reported at the last evaluation before the 1B-token budget, corresponding to 0.918B training tokens. The final row reports a paired two-sided

t

-test over validation loss between AF-Muon and the strongest non-AF baseline by mean validation loss; the gain is baseline minus AF-Muon, so positive values indicate that AF-Muon is better.

Method	Seed	Tied RMS	Tied max	Matrix RMS	Aux WD
Hybrid Muon	44	0.0481	0.3676	0.0435	0.01
Hybrid Muon	45	0.0488	0.3256	0.0433	0.01
SCION-style Sign	44	0.2732	2.6204	0.0440	0.00
SCION-style Sign	45	0.2666	2.6204	0.0439	0.00
AF-Muon	44	0.1997	2.6786	0.0445	0.00
AF-Muon	45	0.1934	3.0957	0.0444	0.00

Appendix

Table 7: Final Qwen2.5-0.5B tied-table diagnostics for the 1B short-horizon seed-44 and seed-45 runs. Hybrid Muon uses auxiliary AdamW weight decay

0.01

; SCION-style Sign and AF-Muon use zero auxiliary/tied weight decay. Seed-43 checkpoint diagnostics are reported separately for the long-horizon continuation in Table 9 .

Figure 7: Qwen2.5-0.5B training dynamics on FineWeb. Curves show the mean over three paired seeds, and shaded regions denote

\pm 1

standard error of the mean. The plot shows the short-horizon region used for the main three-seed comparison, omitting the earliest training region for readability. All methods use tied input/output embeddings, FP32 trainable parameters and optimizer state with BF16 autocast, matrix weight decay

0.1

, and the same matched large-batch constant-learning-rate protocol.

Optimizer	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	3.4254	3.4274	30.80
SCION-style Sign	3.4068	3.4096	30.25
AF-Muon	3.3922	3.3957	29.84

Appendix

Table 8: Qwen2.5-0.5B long-horizon seed-43 endpoint results. The run continues the same seed used in the paired short-horizon comparison to 6.40B training tokens; the validation endpoint below is the last scheduled evaluation, at 6.29B tokens. This single-seed continuation is used as a stability check, not as a three-seed estimate.

Optimizer	Tied RMS	Tied max	Matrix RMS	Matrix WD	Aux./tied WD	Precision
Hybrid Muon	0.0553	0.4511	0.0439	0.1	0.01	FP32 + BF16
SCION-style Sign	0.5815	7.4900	0.0450	0.1	0.00	FP32 + BF16
AF-Muon	0.3868	6.5167	0.0452	0.1	0.00	FP32 + BF16

Appendix

Table 9: Final Qwen2.5-0.5B long-horizon seed-43 checkpoint diagnostics. The tied-table and matrix norms are measured from the final model-only checkpoint at 6.40B training tokens. Hybrid Muon uses auxiliary/tied AdamW weight decay

0.01

; SCION-style Sign and AF-Muon use zero auxiliary/tied weight decay. All runs use FP32 trainable parameters and optimizer state with BF16 autocast.

Figure 8: Qwen2.5-0.5B long-horizon seed-43 continuation. The three optimizer arms are trained under the same matched large-batch protocol as the short-horizon runs. The plotted curves omit the earliest training region and begin after 1B tokens to focus on the long-horizon behavior. This figure is a single-seed stability check, not a replacement for the three-seed short-horizon comparison. AF-Muon preserves the best validation loss at the last scheduled evaluation point, 6.29B tokens.

Method	Seed	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	43	3.4927	3.5348	34.29
Hybrid Muon	44	3.4908	3.5400	34.47
Hybrid Muon	45	3.4962	3.5415	34.52
SCION-style Sign	43	3.4755	3.5139	33.58
SCION-style Sign	44	3.4750	3.5280	34.06
SCION-style Sign	45	3.4809	3.5256	33.97

Appendix

Table 10: Per-seed Llama-3.2-1B short-horizon results. Values are reported at the final available evaluation near the 2B-token budget. Seed 43 is read from the long-horizon continuation at the last pre-2B evaluation point, 1.966B tokens; seeds 44 and 45 are read from the 2B short-horizon runs. The final row reports a paired two-sided

t

-test over validation loss between AF-Muon and the strongest non-AF baseline by mean validation loss; the gain is baseline minus AF-Muon, so positive values indicate that AF-Muon is better.

Method	Seed	Tied RMS	Tied max	Matrix RMS	Aux WD
Hybrid Muon	44	0.0311	0.3537	0.0287	0.01
Hybrid Muon	45	0.0310	0.2848	0.0287	0.01
SCION-style Sign	44	0.1690	2.2859	0.0287	0.00
SCION-style Sign	45	0.1718	2.3012	0.0287	0.00
AF-Muon	44	0.1129	2.2891	0.0290	0.00
AF-Muon	45	0.1125	1.3577	0.0290	0.00

Appendix

Table 11: Final Llama-3.2-1B tied-table diagnostics for the 2B short-horizon seed-44 and seed-45 runs. Hybrid Muon uses auxiliary/tied AdamW weight decay

0.01

; SCION-style Sign and AF-Muon use zero auxiliary/tied weight decay. Seed-43 checkpoint diagnostics are reported separately for the long-horizon continuation in Table 13 .

Optimizer	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	3.3848	3.4251	30.73
SCION-style Sign	3.3534	3.3951	29.82
AF-Muon	3.3504	3.3912	29.70

Appendix

Table 12: Llama-3.2-1B long-horizon seed-43 endpoint results.

Optimizer	Grad norm	Tied RMS	Tied max	Matrix WD	Aux. WD	Tied WD	Precision
Hybrid Muon	0.3709	0.0368	0.3828	0.1	0.01	0.01	FP32 + BF16
SCION-style Sign	0.3029	0.3656	6.1341	0.1	0.00	0.00	FP32 + BF16
AF-Muon	0.3084	0.2463	2.8759	0.1	0.00	0.00	FP32 + BF16

Appendix

Table 13: Final Llama-3.2-1B long-horizon seed-43 diagnostics. All runs use FP32 trainable parameters and optimizer state with BF16 autocast.

Figure 9: Llama-3.2-1B long-horizon seed-43 continuation. The curves omit the earliest training region and begin after 1B tokens to focus on long-horizon behavior. This is a single-seed stability check, not a replacement for the three-seed short-horizon comparison. AF-Muon preserves the best endpoint validation loss at 9.98B tokens, while both AdamW-free tied-table arms remain ahead of Hybrid Muon.

Figure 10: Small T5-style training dynamics under shared-vocabulary and untied-vocabulary topologies. Curves show the mean over three seeds. The untied Hybrid Muon control has nearly twice as many parameters, because the encoder embedding, decoder embedding, and output head are separate.

Method	$\rho_{\mathrm{out}}/\rho_{\mathrm{hid}}$	Tied LR	Eval loss $\downarrow$	Eval PPL $\downarrow$	Shared RMS	Shared max
Hybrid Muon	–	–	5.3468	209.93	0.142	3.19
AF-Muon	1	$1.9531\times 10^{-5}$	5.2652	193.48	0.318	7.51
AF-Muon	5	$9.7656\times 10^{-5}$	5.3536	211.36	0.721	8.99
AF-Muon	10	$1.9531\times 10^{-4}$	5.4220	226.34	1.136	20.43
AF-Muon	60	$1.1719\times 10^{-3}$	5.4864	241.38	2.157	37.02

Appendix

Table 14: Single-seed

\rho

sweep for AF-Muon on small T5-style shared-vocabulary encoder–decoder with a fully shared vocabulary table. All runs use 110M target tokens, matrix weight decay

0.1

, and zero auxiliary/tied weight decay. We select

\rho_{\mathrm{out}}/\rho_{\mathrm{hid}}=1

for the main three-seed experiment.

Method	Topology	Seed	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon		43	5.3372	207.94
Hybrid Muon		44	5.3272	205.86
Hybrid Muon		45	5.3476	210.11
SCION-style Sign		43	5.1643	174.91
SCION-style Sign	Fully shared	44	5.1714	176.17
SCION-style Sign		45	5.1930	180.01

Appendix

Table 15: Per-seed small T5-style shared-vocabulary encoder–decoder results for the selected shared-vocabulary setting. All runs use 275M target tokens. Hybrid Muon uses auxiliary AdamW weight decay

0.01

; SCION-style Sign and AF-Muon use zero auxiliary/tied weight decay. The table reports paired two-sided

t

-tests over validation loss comparing AF-Muon with Hybrid Muon and with the strongest non-AF baseline – in the fully shared setting – by mean validation loss.

Method	Seed/Summary	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	43	3.6366	3.6501	38.48
Hybrid Muon	44	3.6455	3.6578	38.77
Hybrid Muon	45	3.6407	3.6566	38.73
Hybrid Muon	Mean $\pm$ SD	3.6409 $\pm$ 0.0044	3.6548 $\pm$ 0.0041	38.66 $\pm$ 0.16
SCION-style Sign	43	3.5705	3.5871	36.13
SCION-style Sign	44	3.5703	3.5867	36.11

Appendix

Table 16: T5Gemma2-style shared-vocabulary encoder–decoder results at 750M target tokens. Runs use the random-init google/t5gemma-2-270m-270m configuration with a shared 262k-row vocabulary table. Per-seed rows report validation metrics; summary rows report mean

\pm

SD over paired seeds. Paired tests compare validation loss against AF-Muon; positive gain means AF-Muon is better.

Figure 11: T5Gemma2-style shared-vocabulary encoder–decoder training dynamics after 100M target tokens. Curves show the mean over three paired seeds for Hybrid Muon, SCION-style Sign, and AF-Muon.

Figure 12: Tied-embedding NanoGPT-MoE training dynamics on FineWeb. Curves show the mean over three paired seeds, and shaded regions denote

\pm 1

standard error of the mean. All methods use four top-1 routed experts per layer and the same router auxiliary load-balancing coefficient.

Optimizer	Train loss $\downarrow$	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	3.4873 $\pm$ 0.0031	3.4902 $\pm$ 0.0019	32.79 $\pm$ 0.06
SCION-style Sign	3.5106 $\pm$ 0.0033	3.5129 $\pm$ 0.0025	33.55 $\pm$ 0.08
AF-Muon	3.4662 $\pm$ 0.0011	3.4677 $\pm$ 0.0007	32.06 $\pm$ 0.02

Appendix

Table 17: Tied-embedding NanoGPT-MoE results on FineWeb for 2.62B tokens across three paired seeds. We report mean

\pm

SD over seeds. The model uses four top-1 routed experts per layer and a router auxiliary load-balancing coefficient of

0.01

Method	Seed	Validation loss $\downarrow$	Validation PPL $\downarrow$
Hybrid Muon	43	3.4883	32.73
Hybrid Muon	44	3.4901	32.79
Hybrid Muon	45	3.4921	32.85
SCION-style Sign	43	3.5112	33.49
SCION-style Sign	44	3.5158	33.64
SCION-style Sign	45	3.5117	33.50

Appendix

Table 18: Per-seed tied-embedding NanoGPT-MoE results. All runs use 2.62B training tokens, four top-1 routed experts per layer, and the same router auxiliary loss. The final row reports a paired two-sided

t

-test over validation loss between AF-Muon and the strongest non-AF baseline by mean validation loss; the gain is baseline minus AF-Muon, so positive values indicate that AF-Muon is better.

Method	Tied RMS	Tied max	Router entropy $\uparrow$	Max expert frac. $\downarrow$	Router CV $\downarrow$	Unused experts
Hybrid Muon	0.0424 $\pm$ 0.0001	0.2656 $\pm$ 0.0120	0.9996 $\pm$ 0.0001	0.2600 $\pm$ 0.0020	0.0290 $\pm$ 0.0056	0.00 $\pm$ 0.00
SCION-style Sign	0.6357 $\pm$ 0.0161	7.8549 $\pm$ 0.0381	0.9978 $\pm$ 0.0011	0.2687 $\pm$ 0.0025	0.0604 $\pm$ 0.0096	0.00 $\pm$ 0.00
AF-Muon	0.3736 $\pm$ 0.0044	6.1092 $\pm$ 2.0978	0.9994 $\pm$ 0.0001	0.2626 $\pm$ 0.0009	0.0396 $\pm$ 0.0043	0.00 $\pm$ 0.00

Appendix

Table 19: Final tied-table and router diagnostics for the tied-embedding NanoGPT-MoE ablation. Values are mean

\pm

SD over three paired seeds.

Optimizer	Train loss $\downarrow$	Val. loss $\downarrow$	Val. bits/dim $\downarrow$	Val. PPL $\downarrow$
Hybrid Muon	2.3236 $\pm$ 0.0067	2.3346 $\pm$ 0.0049	1.1227 $\pm$ 0.0024	10.33 $\pm$ 0.05
SCION-style Sign	2.3568 $\pm$ 0.0049	2.3712 $\pm$ 0.0055	1.1403 $\pm$ 0.0027	10.71 $\pm$ 0.06
AF-Muon	2.3208 $\pm$ 0.0101	2.3243 $\pm$ 0.0055	1.1177 $\pm$ 0.0026	10.22 $\pm$ 0.06

Appendix

Table 20: Tied ImageGPT-style ImageNet-32 results with a deterministic RGB554 16k color vocabulary. Results are mean

\pm

SD over three paired seeds. All methods use the same tied input/output color-token table, FP32 trainable parameters and optimizer state with BF16 autocast, batch size 512, and 1.31B image tokens.

Method	Seed	Train loss $\downarrow$	Val. loss $\downarrow$	Val. bits/dim $\downarrow$	Val. PPL $\downarrow$
Hybrid Muon	43	2.3301	2.3363	1.1235	10.34
Hybrid Muon	44	2.3240	2.3290	1.1200	10.27
Hybrid Muon	45	2.3166	2.3384	1.1245	10.36
SCION-style Sign	43	2.3537	2.3671	1.1383	10.67
SCION-style Sign	44	2.3625	2.3690	1.1392	10.69
SCION-style Sign	45	2.3543	2.3775	1.1433	10.78

Appendix

Table 21: Per-seed tied ImageGPT-style ImageNet-32 results. All runs use 1.31B image tokens and the same deterministic RGB554 16k tokenizer. The final row reports a paired two-sided

t

-test over validation loss between AF-Muon and the strongest non-AF baseline by mean validation loss; the gain is baseline minus AF-Muon, so positive values indicate that AF-Muon is better.

Method	Tied RMS	Tied max	Rows clipped	Tail energy	Cos(AF, Sign)	Obj. ratio(AF/Sign)
Hybrid Muon	0.0327 $\pm$ 0.0000	0.1430 $\pm$ 0.0021	–	–	–	–
SCION-style Sign	0.5332 $\pm$ 0.0037	4.9483 $\pm$ 0.6251	1.000 $\pm$ 0.000	0.8012 $\pm$ 0.0005	–	–
AF-Muon	0.3096 $\pm$ 0.0006	2.2829 $\pm$ 0.3463	0.773 $\pm$ 0.015	0.0350 $\pm$ 0.0013	0.799 $\pm$ 0.001	1.251 $\pm$ 0.000

Appendix

Table 22: Final tied-table diagnostics for the tied ImageGPT-style ImageNet-32 ablation. Values are mean

\pm

SD over three paired seeds. The cosine and objective-ratio columns compare the AF-Muon tied-table update geometry to the Sign endpoint.

Figure 13: Tied ImageGPT-style ImageNet-32 training dynamics with a deterministic RGB554 16k color vocabulary. Curves show the mean over three paired seeds, and shaded regions denote

\pm 1

standard error of the mean. To focus on the stable late-training regime, the plot starts at step 1000; training loss is shown every 250 steps, and validation metrics are shown every 500 steps. AF-Muon achieves the best final validation loss and validation bits/dim, while SCION-style Sign remains worse throughout this late window.

Optimizer	Completion loss $\downarrow$	Bits/token $\downarrow$	Bits/dim $\downarrow$	PPL $\downarrow$
Hybrid Muon	2.0273 $\pm$ 0.0038	2.9248 $\pm$ 0.0054	0.9749 $\pm$ 0.0018	7.59 $\pm$ 0.03
SCION-style Sign	2.0511 $\pm$ 0.0105	2.9591 $\pm$ 0.0151	0.9864 $\pm$ 0.0050	7.78 $\pm$ 0.08
AF-Muon	2.0117 $\pm$ 0.0076	2.9023 $\pm$ 0.0109	0.9674 $\pm$ 0.0036	7.48 $\pm$ 0.06

Appendix

Table 23: Top-half image-completion metrics for the tied ImageGPT-style ImageNet-32 experiment. Metrics are computed by conditioning on the first 512 tokens of fixed validation images and evaluating the likelihood of the held-out bottom half. Results are mean

\pm

SD over three trained seeds; all methods use the same validation images.

Method	Seed	Completion loss $\downarrow$	Bits/token $\downarrow$	Bits/dim $\downarrow$	PPL $\downarrow$
Hybrid Muon	43	2.025113	2.921620	0.973973	7.58
Hybrid Muon	44	2.0316	2.9310	0.9770	7.63
Hybrid Muon	45	2.025116	2.921625	0.973975	7.58
SCION-style Sign	43	2.0456	2.9512	0.9837	7.73
SCION-style Sign	44	2.0631	2.9765	0.9922	7.87
SCION-style Sign	45	2.0444	2.9495	0.9832	7.72

Appendix

Table 24: Per-seed top-half completion results for tied ImageGPT-style ImageNet-32. The same fixed validation images and completion protocol are used for all methods.

Panel of Figure 14 — Figure 14: Qualitative top-half completions from the tied ImageGPT-style ImageNet-32 seed-43 checkpoints. Each panel uses the same fixed validation images and the same sampling seed. Within each panel, columns show the top-half prime, the ground-truth full image, and the sampled bottom-half completion. The samples are illustrative only; quantitative completion likelihoods are reported in Table 23 .

Figure 15: Protein language-modeling training dynamics on UniRef50 / ProtGPT2-BPE with a 363M-parameter tied-vocabulary RoPE decoder. Curves show the mean over three paired seeds under the matched large-batch protocol, after omitting the first 100M training tokens to focus on the stable training regime. AF-Muon maintains the lowest training loss, validation loss, and validation perplexity among the tested optimizers.

AdamW LR	Eval loss $\downarrow$	Eval PPL $\downarrow$
$1\times 10^{-4}$	8.2669	3892.69
$3\times 10^{-4}$	8.2007	3643.67
$6\times 10^{-4}$	8.1369	3418.32
$1\times 10^{-3}$	8.1672	3523.36

Appendix

Table 25: Seed-43 AdamW LR pilot. All runs use the same 363M RoPE decoder and 500 steps / 32.77M tokens. We use

\eta=6\times 10^{-4}

for the longer AdamW comparison in table

Optimizer	Seed	Train step	Eval step	Train loss	Eval loss	Eval ppl	Best loss
AdamW	43	850	875	7.9344	8.0297	3070.73	7.8998
AdamW	44	850	875	7.9773	7.8709	2619.91	7.8709
AdamW	45	850	875	7.9421	7.8737	2627.15	7.8737
Hybrid Muon	43	850	875	7.4835	6.8770	969.68	6.8770
Hybrid Muon	44	850	875	7.4868	6.8806	973.21	6.8806
Hybrid Muon	45	850	875	7.4633	6.8650	958.14	6.8650

Appendix

Table 26: Per-seed protein language-modeling results on UniRef50 / ProtGPT2-BPE. All runs use the matched large-batch protocol and are compared at 458.8M training tokens. The final row reports a paired two-sided

t

-test over eval loss between AF-Muon and the strongest non-AF baseline by mean eval loss; the gain is baseline minus AF-Muon, so positive values indicate that AF-Muon is better.

Model	Optimizer	Opt. state	Saved	Peak alloc.	Step (s)	Tok/s	Rel. thr.	Train loss
NanoGPT	Hybrid Muon	0.63 GiB	–	18.603 GiB	6.306	83.1k	1.000 $\times$	6.5425
NanoGPT	SCION-style Sign	0.49 GiB	22.2%	18.461 GiB	6.331	82.8k	0.996 $\times$	6.2402
NanoGPT	AF-Muon	0.49 GiB	22.2%	18.461 GiB	6.386	82.1k	0.987 $\times$	6.1697
SmolLM2-135M	Hybrid Muon	0.55 GiB	–	25.000 GiB	11.858	44.2k	1.000 $\times$	6.8592
SmolLM2-135M	SCION-style Sign	0.45 GiB	18.2%	24.895 GiB	11.864	44.2k	1.000 $\times$	6.6022
SmolLM2-135M	AF-Muon	0.45 GiB	18.2%	24.895 GiB	11.928	44.0k	0.994 $\times$	6.5485

Appendix

Table 27: Training-overhead benchmark under the matched large-batch protocol. Optimizer state memory is computed analytically from the parameter partition and FP32 optimizer states. “Saved” reports the optimizer-state reduction relative to Hybrid Muon within the same model. Peak allocated memory is measured exactly with torch.cuda.max_memory_allocated() after resetting CUDA peak statistics before the training loop. Step time is the median over 30 training steps after discarding the first five steps as warmup. The short-run training loss is included as a sanity check; long-horizon performance is reported in the main experiment tables.

Tied table	1D aux.	Val. loss $\downarrow$	Val. PPL $\downarrow$
AdamW	AdamW	3.6296 $\pm$ 0.0030	37.70 $\pm$ 0.11
AF-Muon	AdamW	3.6163 $\pm$ 0.0108	37.20 $\pm$ 0.40
AdamW	RMS-LMO	3.6221 $\pm$ 0.0028	37.41 $\pm$ 0.11
AF-Muon	RMS-LMO	3.5883 $\pm$ 0.0016	36.17 $\pm$ 0.06

Appendix

Table 28: NanoGPT

2\times 2

factorial ablation separating the tied vocabulary-table update from the remaining one-dimensional auxiliary update. All runs use 1.05B training tokens and three paired seeds.

Vocab	Optimizer	Params	Train loss $\downarrow$	Val. loss $\downarrow$	Val. PPL $\downarrow$
50k	SCION-style Sign	123.4M	3.9881 $\pm$ 0.0079	3.9260 $\pm$ 0.0014	50.70 $\pm$ 0.07
50k	AF-Muon	123.4M	3.9561 $\pm$ 0.0017	3.8897 $\pm$ 0.0051	48.90 $\pm$ 0.25
100k	SCION-style Sign	161.8M	4.0691 $\pm$ 0.0069	3.9947 $\pm$ 0.0041	54.31 $\pm$ 0.22
100k	AF-Muon	161.8M	4.0275 $\pm$ 0.0059	3.9500 $\pm$ 0.0061	51.94 $\pm$ 0.32
150k	SCION-style Sign	200.2M	4.0881 $\pm$ 0.0108	4.0228 $\pm$ 0.0070	55.86 $\pm$ 0.39
150k	AF-Muon	200.2M	4.0405 $\pm$ 0.0055	3.9706 $\pm$ 0.0041	53.02 $\pm$ 0.22

Appendix

Table 29: NanoGPT real-vocabulary re-tokenization ablation. We train matched ByteLevel BPE tokenizers with fix d=768, and 50k, 100k, and 150k vocabulary sizes on the same FineWeb shard, then train each optimizer for 500M tokens. All runs use FP32 trainable parameters and optimizer state with BF16 autocast, matrix weight decay

0.1

, zero auxiliary/tied weight decay, and effective batch size 512. The 50k runs use microbatch 16; 100k and 150k use microbatch 8 for memory. We report mean

\pm

sample standard deviation over three paired seeds. related figure 16

Panel of Figure 16 — Figure 16: Real-vocabulary re-tokenization ablation on NanoGPT. Each row uses a separately trained ByteLevel BPE tokenizer with vocabulary size 50k, 100k, or 150k and compares AF-Muon with the SCION-style Sign endpoint for 500M training tokens. Curves show the mean over three paired seeds, with shaded regions denoting $\pm 1$ standard error. Since tokenizers differ across rows, absolute validation losses and perplexities should not be compared across vocabulary sizes; the relevant comparison is the within-vocabulary AF-Muon versus Sign gap.

Method	$s$	$d_{\mathrm{eff}}$	Train loss $\downarrow$	Val. loss $\downarrow$	Val. PPL $\downarrow$	Tied RMS
SCION-style Sign	1.0000	–	3.9971	3.9250	50.65	0.2147
AF-Muon	0.5000	768	3.9558	3.8867	48.75	0.1241
AF-Muon	0.4300	893	3.9510	3.8799	48.42	0.1066
AF-Muon	0.3750	1024	3.9409	3.8697	47.93	0.0956
AF-Muon	0.2500	1536	3.9377	3.8687	47.88	0.0692
AF-Muon	0.1875	2048	3.9406	3.8720	48.04	0.0554

Appendix

Table 30: NanoGPT AF-Muon tied-scale sweep at fixed model width and vocabulary size. All runs use a 12-layer NanoGPT model with hidden size

d=768

, a 50k vocabulary, 500M training tokens, FP32 trainable parameters and optimizer state with BF16 autocast, matrix weight decay

0.1

, and zero auxiliary/tied weight decay. AF-Muon uses

c=3

and varies the tied scale

s

. The SCION-style Sign endpoint uses

c=1,s=1

under the same training protocol.

Figure 17: NanoGPT tied-scale sweep training dynamics. Top: training loss, validation loss, and validation perplexity for SCION-style Sign and AF-Muon with different tied scales

s

. Bottom: token-wise AF-Muon minus SCION-style Sign gaps for the same metrics. All runs use the same fixed 50k vocabulary,

d=768

model width, 500M training-token budget, FP32 trainable parameters and optimizer state with BF16 autocast, matrix weight decay

0.1

, and zero auxiliary/tied weight decay.

Width $d$	Optimizer	Params	Train loss $\downarrow$	Val. loss $\downarrow$	Val. PPL $\downarrow$
768	SCION-style Sign	123.4M	3.9881 $\pm$ 0.0079	3.9260 $\pm$ 0.0014	50.70 $\pm$ 0.07
768	AF-Muon	123.4M	3.9561 $\pm$ 0.0017	3.8897 $\pm$ 0.0051	48.90 $\pm$ 0.25
1024	SCION-style Sign	202.2M	3.9205 $\pm$ 0.0035	3.8464 $\pm$ 0.0005	46.82 $\pm$ 0.02
1024	AF-Muon	202.2M	3.8822 $\pm$ 0.0026	3.8068 $\pm$ 0.0026	45.00 $\pm$ 0.12
1536	SCION-style Sign	416.6M	3.8475 $\pm$ 0.0020	3.7573 $\pm$ 0.0041	42.83 $\pm$ 0.18
1536	AF-Muon	416.6M	3.8136 $\pm$ 0.0031	3.7253 $\pm$ 0.0051	41.48 $\pm$ 0.21

Appendix

Table 31: Real-width NanoGPT ablation at fixed 50k vocabulary and 500M training tokens. We report mean

\pm

sample standard deviation over three paired seeds. Rows clipped denotes the fraction of tied-table rows for which the finite-cap constraint is active in the final tied-momentum diagnostic.

Width $d$	Optimizer	Seed	Train loss $\downarrow$	Val. loss $\downarrow$	Val. PPL $\downarrow$
768	SCION-style Sign	43	3.9971	3.9250	50.65
768	SCION-style Sign	44	3.9844	3.9253	50.67
768	SCION-style Sign	45	3.9827	3.9276	50.79
768	AF-Muon	43	3.9558	3.8867	48.75
768	AF-Muon	44	3.9546	3.8956	49.19
768	AF-Muon	45	3.9580	3.8869	48.76

Appendix

Table 32: Per-seed real-width NanoGPT results for the width ablation in Table 31 .

Width $d$	Seeds	$\Delta$ Val. loss	$\Delta$ Val. PPL
768	43,44,45	$-0.0363\pm 0.0058$	$-1.81\pm 0.28$
1024	43,44,45	$-0.0397\pm 0.0023$	$-1.82\pm 0.10$
1536	43,44,45	$-0.0321\pm 0.0085$	$-1.35\pm 0.36$

Appendix

Table 33: Paired AF-Muon minus SCION-style Sign gaps for the real-width NanoGPT ablation. Negative values indicate that AF-Muon is better.

Panel of Figure 18 — Figure 18: Real-width NanoGPT training dynamics at fixed 50k vocabulary and 500M training tokens. Each row compares AF-Muon and the SCION-style Sign endpoint for one hidden width $d\in\{768,1024,1536\}$ . Curves show the mean over three paired seeds, with shaded regions denoting $\pm 1$ standard error.

Batch	Tokens/update	Hybrid Muon	SCION-style Sign	AF-Muon	AF–Sign
32	32,768	4.2078	4.1472	4.1378	-0.0094
64	65,536	3.9918	3.9738	3.9514	-0.0224
128	131,072	3.8571	3.8559	3.8180	-0.0379
256	262,144	3.7950	3.7890	3.7544	-0.0346
512	524,288	3.8200	3.7649	3.7400	-0.0249

Appendix

Table 34: NanoGPT batch-size sensitivity at a matched 458.8M-token comparison point. All runs use seed 43. Tokens/update is the number of training tokens per optimizer step.

Aux/tied LR	Mult.	Hybrid Muon	SCION-style Sign	AF-Muon	AF–Sign	AF–Hybrid
$7.5\times 10^{-5}$	$0.25\times$	4.1498	3.7743	3.7557	-0.0185	-0.3941
$1.5\times 10^{-4}$	$0.5\times$	3.9716	3.7664	3.7472	-0.0192	-0.2244
$3.0\times 10^{-4}$	$1.0\times$	3.8182	3.7726	3.7337	-0.0389	-0.0845
$6.0\times 10^{-4}$	$2.0\times$	3.7641	3.7522	3.7284	-0.0238	-0.0357
$1.2\times 10^{-3}$	$4.0\times$	3.7304	3.7393	3.7220	-0.0173	-0.0084
$2.4\times 10^{-3}$	$8.0\times$	3.7233	3.7283	3.7208	-0.0075	-0.0024

Appendix

Table 35: NanoGPT auxiliary/fallback learning-rate sensitivity at a matched 458.8M-token comparison point. All runs use seed 43. The hidden-matrix Muon learning rate is fixed at

0.02

; the auxiliary/fallback learning-rate scale is varied. For Hybrid Muon this branch includes the tied vocabulary table, while the AdamW-free arms use their tied-table LMO and vary the remaining auxiliary/vector scale.

Cap $c$	Scale $s$	Train loss	Val. loss	Val. PPL	Tied RMS	Tied max
1	0.5	3.7880 $\pm$ 0.0208	3.7831 $\pm$ 0.0033	43.95 $\pm$ 0.15	0.1270 $\pm$ 0.0008	0.6653 $\pm$ 0.0033
2	0.5	3.7885 $\pm$ 0.0250	3.7826 $\pm$ 0.0021	43.93 $\pm$ 0.09	0.1444 $\pm$ 0.0004	1.1725 $\pm$ 0.0324
3	0.5	3.7872 $\pm$ 0.0253	3.7815 $\pm$ 0.0026	43.88 $\pm$ 0.12	0.1418 $\pm$ 0.0029	1.5712 $\pm$ 0.0923
4	0.5	3.7892 $\pm$ 0.0241	3.7856 $\pm$ 0.0062	44.06 $\pm$ 0.27	0.1409 $\pm$ 0.0039	1.7930 $\pm$ 0.1915
6	0.5	3.7910 $\pm$ 0.0232	3.7845 $\pm$ 0.0005	44.01 $\pm$ 0.02	0.1408 $\pm$ 0.0017	2.1051 $\pm$ 0.2569
10	0.5	3.7909 $\pm$ 0.0244	3.7867 $\pm$ 0.0014	44.11 $\pm$ 0.06	0.1399 $\pm$ 0.0018	2.5313 $\pm$ 0.2197

Appendix

Table 36: AF-Muon cap and scale sensitivity on the NanoGPT tied-vocabulary task at the 393.2M-token intermediate comparison point. Results are mean

\pm

SD over three paired seeds. The default setting used in the main experiments is

c=3

s=0.5

Cap $c$	Scale $s$	Matrix RMS	Coord. capped	Rows clipped	Tail energy	Cos(AF, Sign)	Obj. ratio
1	0.5	0.0493 $\pm$ 0.0002	1.0000 $\pm$ 0.0000	1.0000 $\pm$ 0.0000	0.8005 $\pm$ 0.0002	1.000 $\pm$ 0.000	1.000 $\pm$ 0.000
2	0.5	0.0501 $\pm$ 0.0001	0.0579 $\pm$ 0.0000	1.0000 $\pm$ 0.0000	0.2723 $\pm$ 0.0021	0.818 $\pm$ 0.000	1.246 $\pm$ 0.000
3	0.5	0.0500 $\pm$ 0.0001	0.0028 $\pm$ 0.0001	0.8949 $\pm$ 0.0042	0.0303 $\pm$ 0.0014	0.799 $\pm$ 0.000	1.253 $\pm$ 0.000
4	0.5	0.0500 $\pm$ 0.0002	0.0001 $\pm$ 0.0000	0.0484 $\pm$ 0.0050	0.0012 $\pm$ 0.0002	0.798 $\pm$ 0.000	1.253 $\pm$ 0.000
6	0.5	0.0500 $\pm$ 0.0001	0.0000 $\pm$ 0.0000	0.0002 $\pm$ 0.0004	0.0000 $\pm$ 0.0000	0.798 $\pm$ 0.000	1.253 $\pm$ 0.000
10	0.5	0.0499 $\pm$ 0.0001	0.0000 $\pm$ 0.0000	0.0000 $\pm$ 0.0000	0.0000 $\pm$ 0.0000	0.798 $\pm$ 0.000	1.253 $\pm$ 0.000

Appendix

Table 37: Tied-table diagnostics for the NanoGPT cap and scale sweep at 393.2M training tokens. Results are mean

\pm

SD over three matched seeds. Coordinate and row fractions describe cap activity; tail energy is computed from the row-RMS endpoint before applying the finite-cap oracle. Cosine similarity and linear-objective ratio compare the normalized AF-Muon direction with Sign evaluated on the same oracle input within each run.

Tied-table gradient	Train loss	Val. loss	Val. PPL	Input rows	Output rows	Input-output cos.
Sparse input only	4.0353	4.0513	57.47	0.601	–	–
Dense output only	3.7471	3.7634	43.10	–	1.000	–
Sparse input + dense output	3.7217	3.7395	42.08	0.601	1.000	0.003

Appendix

Table 38: NanoGPT tied-table gradient-source ablation at a matched 458.8M-token comparison point. Only the gradient used for the tied input/output vocabulary table is changed; all other model parameters use the same AF-Muon update.

Figure 19: Training dynamics for the NanoGPT tied-table gradient-source ablation. Curves show seed-43 runs after the first 100M training tokens. All variants use the same AF-Muon optimizer for non-tied parameters; only the gradient supplied to the tied input/output vocabulary table is changed. The dense output-only variant learns much faster than the sparse input-only variant, but the full sparse-input plus dense-output tied gradient remains best.

Matrix	Mean pair cosine	Min. pair cosine	Negative pairs	Norm CV	Cancellation $\kappa$	Direction cosine	Objective retained	Split error
Up projection	$0.3396$	$0.0507$	$0.0\%$	$0.3247$	$0.6803$	$0.6785$	$74.08\%$	$5.22\times 10^{-8}$
Down projection	$0.5014$	$0.2251$	$0.0\%$	$0.2800$	$0.7697$	$0.7764$	$83.85\%$	$5.11\times 10^{-8}$

Appendix

Table 39: Homogeneous dense-sharing diagnostic. One up-projection and one down-projection are reused at four depths. Statistics are measured at the final audit after

5.0

M predicted tokens. Objective retention compares the equal-spectral-radius sum of per-depth Muon directions with Muon applied once to the accumulated gradient. The experiment uses one seed and is intended as a boundary diagnostic rather than a performance comparison.

Explore similar work

May 26, 2026cs.LG

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

Muon Momentum Stochastic Gradient Descent Large Language Model Training Curvature-Aware Spectral Framework Orthogonality

Ask PDF

May 11, 2026cs.LG

Muown: Row-Norm Control for Muon Optimization

Muon has emerged as a strong competitor to AdamW for language model pre-training, yet its behavior at scale is sensitive to weight decay. Recent work has observed that, for Muon without decoupled weight decay, the spectral norm of weight matrices drifts upward over training. Through a decomposition of the spectral norm into a row-magnitude factor and a row-coherence factor, we identify the former as the empirical driver of this drift under Muon, while the latter remains well-behaved along the trajectory. Motivated by this diagnosis, we introduce Muown, a drop-in replacement for Muon that treats the row-magnitude vector as an explicit optimizer variable, updating it under the

\ell_\infty

geometry induced by the decomposition, while applying Muon unchanged to the remaining direction component. We prove that Muown attains the optimal non-convex rates in both deterministic and stochastic regimes under a dual norm aligned with the underlying geometries and with a stochastic noise coefficient that empirically remains below that of Muon throughout training. Across GPT-style pre-training on FineWeb-Edu with model sizes from 124M up to 2.7B parameters, Muown improves perplexity over Muon, SOAP, AdamW, and Lion. It also widens the plateau of near-optimal learning rates across model scales, reduces sensitivity to weight decay, and avoids the spectral norm drift at negligible step-time overhead when appropriately sharded.

Kai Lion, Florian Hübler, Bingcong Li +2

Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany

Muon Optimizer Design

Ask PDF

Aug 4, 2026cs.LG

Muon Meets Mamba: Spectral Optimization for State Space Models

Muon is a recent optimizer that orthogonalizes the update to each weight matrix with a Newton-Schulz iteration, which performs steepest descent under the spectral norm. Almost all the evidence for it comes from Transformer models, and its behavior on state-space models is largely unreported. We compare Muon with AdamW on Mamba-2 130M under a controlled protocol that varies only which weight groups are trained with Muon. The benefit is localized. Muon on the output projection alone beats Muon on the input projection or on both. The advantage is mainly one of token efficiency. It holds on two corpora and two token budgets, and persists when training continues well past the compute-optimal point. Conditioning does not explain the gain. Muon lowers the condition number of whichever projection it trains, but the better-conditioned input projection is not the one that helps.

Arslan Battalov, Karim Kramin, Alexander Markotenko +1

HSE University · Moscow, Russia

Muon State Space Models Optimizer Design Model Weights Mamba Transformer Architectures

Ask PDF