Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training

Ben S. Southworth; Stephen Thomas

cs.LGMar 18, 2026

Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training

Authors: Ben S. Southworth, Stephen Thomas

Organizations: Theoretical Division, Los Alamos National Laboratory, USA. · Lehigh University, Computer Science

Abstract

Orthogonalized-momentum optimizers such as Muon improve transformer training by approximately whitening/orthogonalizing matrix-valued momentum updates via a short polar-decomposition iteration. However, polar-factor approximations typically require multiple large matrix multiplications, and the resulting overhead can be substantial and hardware-dependent. We introduce MUD (MomentUm Decorrelation), a complementary whitening approach that replaces Muon's polar update with a triangular (Cholesky-like) whitening surrogate inspired by classical Gram--Schmidt and Gauss-Seidel ideas. We show that row-orthonormal matrices are fixed points of the MUD map, relate the inner step to symmetric Gauss-Seidel preconditioning of the Gram matrix, and prove quadratic local convergence near the fixed point. In terms of time-to-perplexity, MUD yields consistent 10-50% wall-clock improvements over tuned AdamW and Muon in time-to-perplexity, typically converging slightly slower per step than Muon but with substantially lower optimizer overhead -- relative to Muon, MUD improves peak tokens/s by roughly $1.3-2.6\times$ across most settings and up to nearly $3\times$ on GPT-2 large on an A100. We also demonstrate training a ESM-2 150M protein language model, where MUD matches Muon-level validation perplexity in significantly less wall-clock time.

Figures & tables

Method	#Grams	#Applies	#TRSM	Leading-order FLOPs
Muon5 (quintic)	$5$	$10$	$0$	$\approx 30\,k^{2}d$
MUD1	$1$	$0$	$1$	$\approx 2.5\,k^{2}d$
MUD2	$2$	$0$	$2$	$\approx 5\,k^{2}d$

Table 1: Dominant dense-linear-algebra cost per parameter update (after transposing so

M\in\mathbb{R}^{k\times d}

k=\min(n,m)

d=\max(n,m)

). “Gram” denotes

XX^{T}

\approx 2k^{2}d

FLOPs; “Apply” denotes multiplying a

k\times k

matrix into a

k\times d

matrix,

\approx 2k^{2}d

FLOPs; TRSM denotes solving a

k\times k

triangular system with

k\times d

right-hand side,

\approx\tfrac{1}{2}k^{2}d

FLOPs (counting multiply-adds). Constants ignore lower-order reductions (row norms) and scalar operations.

Model	Layers	Heads	$d_{\text{model}}$	Params
GPT-2 small	12	12	768	$\approx$ 125M
GPT-2 medium	24	16	1024	$\approx$ 355M
GPT-2 large	36	20	1280	$\approx$ 775M

Table 2: GPT-2 model configurations used in the experiments (decoder-only Transformers).

GPU	$(B,T)$	AdamW	Muon	MUD1	$\frac{\text{MUD}}{\text{Muon}}$	$\frac{\text{AdamW}}{\text{Muon}}$
A100	$(8,1024)$	$131,800$	$46,900$	$87,000$	1.85	2.81
A100	$(12,2048)$	$130,000$	$80,300$	$108,000$	1.35	1.62
A100	$(8,4096)$	$125,700$	$88,000$	$113,700$	1.29	1.43
GH200	$(8,1024)$	$342,400$	$109,000$	$185,800$	1.70	3.14
GH200	$(12,2048)$	$391,100$	$214,600$	$295,300$	1.38	1.82
MI250	$(8,1024)$	$28,200$	$14,900$	$23,000$	1.54	1.89

Table 3: Training throughput (tokens/s) for GPT-2 small on the NanoGPT/OpenWebText benchmark. Absolute throughput and speedups relative to Muon are reported.

GPU	$(B,T)$	AdamW	Muon	MUD1	$\frac{\text{MUD}}{\text{Muon}}$	$\frac{\text{AdamW}}{\text{Muon}}$
A100	$(8,1024)$	$49,300$	$12,600$	$29,100$	2.31	3.91
A100	$(12,2048)$	$55,200$	$26,400$	$43,450$	1.65	2.09
GH200	$(8,1024)$	$142,900$	$33,600$	$68,500$	2.04	4.25
GH200	$(12,2048)$	$159,000$	$72,400$	$114,600$	1.58	2.20
GH200	$(12,4096)$	$134,000$	$89,400$	$116,200$	1.30	1.50
MI250	$(8,1024)$	$28,100$	$14,900$	$23,000$	1.54	1.89

Table 4: Training throughput (tokens/s) for GPT-2 medium on WikiText-103 (GPT-2 BPE tokens). Absolute throughput and speedups relative to Muon are reported.

GPU	$(B,T)$	AdamW	Muon	MUD1	$\frac{\text{MUD}}{\text{Muon}}$	$\frac{\text{AdamW}}{\text{Muon}}$
A100	$(8,1024)$	$25,000$	$4,500$	$13,200$	2.93	5.56
MI250	$(8,1024)$	$13,900$	$6,200$	$10,900$	1.76	2.24
GH200	$(8,1024)$	$70,700$	$12,000$	$30,600$	2.55	5.89
GH200	$(8,2048)$	$71,600$	$20,900$	$44,500$	2.13	3.43
GH200	$(8,4096)$	$66,400$	$31,400$	$52,200$	1.66	2.11

Table 5: Training throughput (tokens/s) for GPT-2 large on FineWeb-Edu (GPT-2 BPE tokens). Absolute throughput and speedups relative to Muon are reported.

Figure 1: Nvidia GH200 ( bs =12, T =2048, GPT2-Medium): validation perplexity vs. step and vs. wall-clock time for AdamW, Muon, and MUD.

Figure 2: Nvidia GH200 ( bs =12, T =2048, GPT2-Medium): (left) speedup relative to AdamW vs. achieved validation perplexity; (right) relative perplexity improvement rate vs. perplexity, showing diminishing relative gains at low perplexity.

Figure 3: Nvidia GH200 ( bs =12, T =2048, GPT2-Small): validation perplexity vs. step and vs. wall-clock time for AdamW, Muon, and MUD.

Figure 4: Nvidia A100 ( bs =12, T =2048, GPT2-Medium): validation perplexity vs. step and vs. wall-clock time for AdamW, Muon, and MUD.

Figure 5: Nvidia A100 ( bs =12, T =2048, GPT2-Medium): (left) speedup relative to AdamW vs. achieved validation perplexity; (right) relative perplexity improvement rate vs. perplexity, showing diminishing relative gains at low perplexity.

Figure 6: Nvidia GH200 ( bs =12, T =4096, GPT2-Medium): (left) speedup relative to AdamW vs. achieved validation perplexity; (right) relative perplexity improvement rate vs. perplexity, again showing diminishing relative gains at low perplexity.

Figure 7: Nvidia GH200 ( bs =8, T =2048, GPT2-Large): (left) Validation loss as a function of step; (right) speedup relative to AdamW vs. achieved validation loss.

Setting	Method	Val. Acc. (mean $\pm$ std)	Time [s] (mean $\pm$ std)
8 epochs (default)	MUD1	$0.9334\pm 0.0014$	$2.9795\pm 0.0101$
	MUD3	$0.9339\pm 0.0013$	$3.2644\pm 0.0111$
	Muon3	$0.9402\pm 0.0011$	$2.9690\pm 0.0071$
	Muon5	$0.9403\pm 0.0010$	$3.0393\pm 0.0085$
13 epochs (+ momentum 0.7)	MUD1	$0.9407\pm 0.0014$	$4.7498\pm 0.0178$
	MUD3	$0.9405\pm 0.0011$	$5.2078\pm 0.0132$

Table 6: CIFAR-10 speed running on AIRBench. Starting from the Muon speed running script, the Muon Newton-Schulz step is exchanged for MUD, using either MUD1 or MUD3. Muon baselines use 3 or 5 NS iterations. Results are mean

\pm

std across 20 runs.

Figure 8: ESM-2 150M training from scratch on omg_prot50 (5,000 steps, AMD MI300X): validation perplexity vs. training step (left) and vs. wall-clock time (right), both on log–log axes, for AdamW (

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

), Muon (

\eta_{\mathrm{muon}}=3\times 10^{-4}

), MUD1, and MUD2 (

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

). Final perplexities at step 4,999 are 17.176 (AdamW, 466s), 16.840 (Muon, 997s), 16.855 (MUD1, 636s), and 16.795 (MUD2, 710s). All four curves are still descending at the end of the training budget.

Target ppl	AdamW	Muon	MUD1	MUD2
17.5	141.2s	127.8s (1.10 $\times$ )	78.0s (1.81 $\times$ )	78.2s (1.81 $\times$ )
17.3	245.1s	188.4s (1.30 $\times$ )	110.0s (2.23 $\times$ )	86.2s (2.84 $\times$ )
17.2	392.8s	236.7s (1.66 $\times$ )	126.0s (3.12 $\times$ )	131.3s (2.99 $\times$ )
Final ppl (step 4,999)	17.176	16.840	16.855	16.795
Total wall time	466s	997s	636s	710s

Table 7: ESM-2 150M training from scratch on omg_prot50 (AMD MI300X, 5,000 steps): wall-clock time to first reach each target validation perplexity, with speedup relative to AdamW in parentheses. MUD1 and MUD2 reach every target faster than both AdamW and Muon.

Deep Learning Optimization

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

See all 199 papers

T-MUON: Tensorizing Muon via t-SVD for Large Language Model Pre-training
Oct 18, 2026Dongyang Li, Mingsong Yan, Zheng Zhang +1Muon Optimizer Language Model Pretraining
University of California, Santa Barbara Santa Barbara, CA 93106, USA
Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates
Dec 4, 2025Thibaut Boissin, Thomas Massena, Franck Mamalet +1Newton-Schulz Iteration Orthogonal Optimization
IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT +2
Spectrally Targeted Muon
Oct 7, 2026Vishrut Goyal, Rohan RamkumarSpectral Optimization Deep Learning Optimization
Department of Computer Science, Stanford University · Department of Electrical Engineering, Stanford University.

Oct 18, 2026J/K move · Enter open · S save

T-MUON: Tensorizing Muon via t-SVD for Large Language Model Pre-training

Dongyang Li, Mingsong Yan, Zheng Zhang, Sui Tang

University of California, Santa Barbara Santa Barbara, CA 93106, USA

Muon Optimizer Language Model Pretraining Transformer Training Dynamics

Open paper AskPDF arXiv

Muon has recently emerged as a strong optimizer for large language model pretraining by exploiting matrix structure through polar orthogonalization. However, it processes each parameter matrix independently, whereas recent tensor-aware extensions of Muon jointly update groups of eligible matrix parameters through matrix unfoldings. In contrast, we propose T-MUON, a t-SVD-based generalization that groups structurally related Transformer gradients into third-order tensors and defines the corresponding Muon-style update directly through the

M

-transform. We show that the exact T-MUON update is a steepest-descent direction under the

\star_M

-spectral norm and establish a convergence guarantee. Across GPT-2-style and LLaMA-style models from 51M to 1.3B parameters trained on FineWeb and FineWeb-Edu, head-wise T-MUON with the identity transform consistently achieves lower validation perplexity than AdamW, Muon, and variants using DFT, DCT, or DST transforms. These gains persist in additional experiments with longer context lengths and larger token budgets, although the preferred tensorization mode varies across settings.

Abstract

Deep Learning Optimization

Explore similar work

T-MUON: Tensorizing Muon via t-SVD for Large Language Model Pre-training

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates

Spectrally Targeted Muon