Scaling Muon for Diffusion Transformers

Chenghao Li; Xiao Han; Xinxin Huang; Wei Liu; Boyang Li; Bing Xiao; Heran Zhang; Juanma Perez Rua

cs.LGAug 21, 2026

Scaling Muon for Diffusion Transformers

Authors: Chenghao Li, Xiao Han, Xinxin Huang, Wei Liu, Boyang Li, Bing Xiao, Heran Zhang, Juanma Perez Rua,

Organizations: University of Southern California · Meta

Abstract

The matrix-aware optimizer Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Diffusion Transformers (DiTs) remain unclear. We first establish Muon's scaling behavior on DiTs from 1.3B to 15B parameters, showing that its optimization and generative quality advantages over AdamW persist across model scales. However, at scale, the 5-step Newton--Schulz iteration (NS5) performed at every optimization step, together with full-momentum materialization, introduces substantial computation and communication overhead that can offset Muon's step-efficiency advantage. We introduce Periodic Row-wise Muon, which performs a full NS5 spectral update once every $K$ steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps. We further co-design a distributed implementation that operates directly on sharded momentum during non-refresh steps and accelerates spectral refreshes through bucketed all-gather and communication--computation overlap. Across all scales, Muon improves the best observed generative quality over AdamW by 12.9--19.1%. Compared with vanilla Muon, Periodic Row-wise Muon remains within 0.5% in best generative quality on the 1.3B--4B models and improves it by 4.5% at 9B. It reduces optimizer time by 46.9--54.3%, end-to-end step time by 15.7--24.3%, and logical communication volume by 66.7%, while reaching its respective best generative quality with 33.7--64.8% less active training time. These results show that Periodic Row-wise Muon preserves Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.

Figures & tables

Panel of Figure 1 — Figure 1: Scaling behavior of Muon for DiTs training. Muon consistently improves step efficiency and generation quality, while its higher per-step overhead can offset this advantage in wall-clock time.

Figure 2: Generative quality across model scales. We report FD-DINO improvement over AdamW at both the final checkpoint and the best checkpoint observed during training.

Figure 3: Generative quality frontier under optional checkpoint selection.

Model	Optimizer	Discrepancy			Fidelity & Diversity				Alignment	Compositionality
Model	Optimizer	FD-DINO $\downarrow$	FID $\downarrow$	MMD-DINO $\downarrow$	Precision $\uparrow$	Recall $\uparrow$	Coverage $\uparrow$	Density $\uparrow$	HPSv2 $\uparrow$	GenEval2 AM $\uparrow$	GenEval2 GM $\uparrow$
1.3B	AdamW	$53.93$	$6.71$	$0.0312$	$0.9405$	$0.8874$	$0.9401$	$0.9665$	$20.53$	$38.66$	$7.10$
	Muon	$46.08$	$6.54$	$\mathbf{0.0268}$	$0.9487$	$\mathbf{0.9071}$	$\mathbf{0.9492}$	$\mathbf{0.9918}$	$\mathbf{20.84}$	$44.72$	$9.38$
	Periodic Row-wise Muon	$\mathbf{45.65}$	$\mathbf{6.29}$	$0.0270$	$\mathbf{0.9517}$	$0.9065$	$0.9485$	$0.9756$	$20.36$	$\mathbf{45.89}$	$\mathbf{9.75}$
4B	AdamW	$51.39$	$7.35$	$0.0271$	$0.9371$	$0.8995$	$0.9388$	$\mathbf{0.9898}$	$\mathbf{21.40}$	$45.70$	$9.13$
	Muon	$\mathbf{41.62}$	$\mathbf{6.08}$	$0.0261$	$0.9385$	$\mathbf{0.9136}$	$0.9431$	$0.9639$	$20.66$	$46.58$	$9.06$

Table 1: Final 60k-step comparison. Bold denotes the best result among optimizers at the same model size.

Model	Method	Computation					Communication
Model	Method	Step time $\downarrow$		Opt. time $\downarrow$		Opt. share $\downarrow$	AG/step $\downarrow$	AR/step $\downarrow$	Comm. volume $\downarrow$ GiB/rank/step
1.3B	Muon	1.873		0.835		44.6%	10.0	0.0	2.38
1.3B	Periodic Row-wise Muon	1.419	$(-24.3\%)$	0.443	$(-46.9\%)$	31.3%	3.3	0.7	0.79	$(-66.7\%)$
4B	Muon	2.352		1.009		42.9%	32.0	0.0	7.58
4B	Periodic Row-wise Muon	1.878	$(-20.2\%)$	0.536	$(-46.9\%)$	28.6%	10.7	0.7	2.53	$(-66.7\%)$
9B	Muon	4.153		1.689		40.7%	80.0	0.0	17.30

Table 2: System-efficiency comparison on 32 H100 nodes.

Figure 4: Simplified profiler traces for the 15B model.

(a) Variant	NS5 Schedule	Off-refresh Update	$\gamma$	Best FD-DINO $\downarrow$
Vanilla Muon	Every step	–	–	41.733
RowNorm every step	None	RowNorm	0.15	51.322
Periodic Row-wise Muon ( $\gamma=1$ )	$K=3$	RowNorm	1.0	44.294
Scalar-controlled ( $\gamma$ for both NS5 and RowNorm)	$K=3$	Scalar-controlled	0.15	44.713
Periodic Row-wise Muon ( $\gamma=0.15$ )	$K=3$	RowNorm	0.15	41.913

Table 3: Ablation studies. (a) Algorithm-level ablation of Periodic Row-wise Muon on the 1.3B model. (b) System-level ablation of the distributed implementation on the 15B model.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Total params.	Trainable params.	Hidden dim.	Blocks	Heads	Muon-routed
1.3B	1,255,586,099	1,255,585,795	1,216	19	19	98.69%
4B	3,950,236,099	3,950,235,715	1,920	24	30	99.20%
9B	9,374,837,315	9,374,836,803	2,560	32	40	99.48%
15B	14,848,237,251	14,848,236,611	2,880	40	45	99.60%

Appendix

Table 4: Model configurations. Muon-routed parameters denotes the fraction of trainable parameters updated by Muon in the Muon-based runs.

	Main learning rate
Method	1.3B	4B	9B	15B
AdamW	$3.0{\times}10^{-4}$	$2.0{\times}10^{-4}$	$1.4{\times}10^{-4}$	$1.2{\times}10^{-4}$
Muon	$5.0{\times}10^{-3}$	$3.0{\times}10^{-3}$	$2.4{\times}10^{-3}$	$2.0{\times}10^{-3}$
Periodic Row-wise Muon	$5.0{\times}10^{-3}$	$3.0{\times}10^{-3}$	$2.4{\times}10^{-3}$	$2.0{\times}10^{-3}$

Appendix

Table 5: Optimizer hyperparameters. “Aux. AdamW” denotes the auxiliary AdamW parameter group used in the Muon-based runs. All methods use weight decay

0.01

and AdamW-updated parameter groups use

\epsilon=10^{-8}

Setting	Value
Image resolution	$512\times 512$
Training length	60,000 steps
Local / global batch size	16 / 4,096
Learning-rate schedule	2k warmup, then cosine decay
Hardware	32 nodes, 256 NVIDIA H100 GPUs
Distributed training	Pytorch FSDP2 with activation checkpointing

Appendix

Table 6: Training settings shared by all methods.

Figure 5: Selection of

K

and

\gamma

. (a) Late-stage FD-DINO, averaged over the 20k, 25k, and 30k checkpoints, across different

K

and

\gamma

. (b) The best mean FD-DINO over

\gamma

at each refresh period. Increasing

K

from 2 to 3 reduces the number of NS5 refreshes by

33\%

with only a

2.8\%

increase in FD-DINO, whereas increasing

K

from 3 to 4 saves a further

25\%

but increases FD-DINO by

8.2\%

. We therefore select

K=3

and

\gamma=0.15

Figure 6: Training loss curves across model scales over 60,000 optimizer steps.

Figure 7: Generation quality for AdamW and Periodic Row-wise Muon.

Optimizer	Step	Fidelity			Diversity				Alignment	Compositionality
Optimizer	Step	FD-DINO $\downarrow$	FID $\downarrow$	MMD-DINO $\downarrow$	Precision $\uparrow$	Recall $\uparrow$	Coverage $\uparrow$	Density $\uparrow$	HPSv2 $\uparrow$	GenEval2 AM $\uparrow$	GenEval2 GM $\uparrow$
AdamW	5k	$266.40$	$9.94$	$0.258$	$0.861$	$0.361$	$0.657$	$0.900$	$15.79$	$17.19$	$1.35$
	10k	$125.51$	$7.57$	$0.094$	$0.908$	$0.681$	$0.842$	$0.947$	$15.93$	$22.24$	$2.77$
	15k	$82.93$	$6.62$	$0.051$	$0.915$	$0.789$	$0.880$	$0.976$	$18.09$	$28.53$	$4.00$
	20k	$86.84$	$7.94$	$0.055$	$0.925$	$0.812$	$0.890$	$1.107$	$17.71$	$29.83$	$4.19$
	25k	$66.30$	$6.58$	$0.037$	$0.929$	$0.849$	$0.921$	$0.999$	$17.55$	$30.96$	$4.75$

Appendix

Table 7: All checkpoint evaluation metrics for the 1.3B model.

Optimizer	Step	Fidelity			Diversity				Alignment	Compositionality
Optimizer	Step	FD-DINO $\downarrow$	FID $\downarrow$	MMD-DINO $\downarrow$	Precision $\uparrow$	Recall $\uparrow$	Coverage $\uparrow$	Density $\uparrow$	HPSv2 $\uparrow$	GenEval2 AM $\uparrow$	GenEval2 GM $\uparrow$
AdamW	5k	$223.73$	$9.17$	$0.199$	$0.870$	$0.453$	$0.700$	$0.901$	$15.90$	$18.62$	$1.78$
	10k	$101.76$	$8.42$	$0.066$	$0.908$	$0.746$	$0.857$	$0.950$	$18.17$	$22.01$	$2.26$
	15k	$72.16$	$7.67$	$0.036$	$0.922$	$0.834$	$0.885$	$0.958$	$18.31$	$34.18$	$6.51$
	20k	$51.87$	$6.87$	$0.023$	$0.941$	$0.866$	$0.926$	$0.998$	$20.73$	$37.09$	$5.76$
	25k	$50.94$	$6.39$	$0.023$	$0.938$	$0.879$	$0.930$	$1.007$	$19.36$	$32.85$	$5.71$

Appendix

Table 8: All checkpoint evaluation metrics for the 4B model.

Optimizer	Step	Fidelity			Diversity				Alignment	Compositionality
Optimizer	Step	FD-DINO $\downarrow$	FID $\downarrow$	MMD-DINO $\downarrow$	Precision $\uparrow$	Recall $\uparrow$	Coverage $\uparrow$	Density $\uparrow$	HPSv2 $\uparrow$	GenEval2 AM $\uparrow$	GenEval2 GM $\uparrow$
AdamW	5k	$184.83$	$8.37$	$0.150$	$0.883$	$0.531$	$0.757$	$0.943$	$16.08$	$19.18$	$1.75$
	10k	$78.05$	$6.05$	$0.043$	$0.924$	$0.794$	$0.890$	$0.991$	$19.32$	$27.01$	$3.41$
	15k	$59.30$	$7.46$	$0.028$	$0.940$	$0.857$	$0.916$	$0.966$	$20.65$	$35.76$	$6.14$
	20k	$43.20$	$6.27$	$0.016$	$0.944$	$0.881$	$0.933$	$0.967$	$20.58$	$41.28$	$8.23$
	25k	$47.67$	$6.43$	$0.019$	$0.944$	$0.879$	$0.940$	$1.020$	$20.21$	$40.96$	$8.46$

Appendix

Table 9: All checkpoint evaluation metrics for the 9B model.

Optimizer	Step	Fidelity			Diversity				Alignment	Compositionality
Optimizer	Step	FD-DINO $\downarrow$	FID $\downarrow$	MMD-DINO $\downarrow$	Precision $\uparrow$	Recall $\uparrow$	Coverage $\uparrow$	Density $\uparrow$	HPSv2 $\uparrow$	GenEval2 AM $\uparrow$	GenEval2 GM $\uparrow$
AdamW	5k	$175.31$	$8.92$	$0.136$	$0.879$	$0.561$	$0.753$	$0.903$	$16.03$	$16.81$	$1.50$
	10k	$74.84$	$7.70$	$0.035$	$0.927$	$0.805$	$0.895$	$0.970$	$19.01$	$25.75$	$3.08$
	15k	$53.21$	$6.42$	$0.023$	$0.939$	$0.856$	$0.920$	$0.970$	$20.82$	$32.86$	$5.35$
	20k	$56.21$	$7.03$	$0.028$	$0.935$	$0.870$	$0.925$	$0.987$	$18.66$	$35.58$	$5.70$
	25k	$40.75$	$5.86$	$0.017$	$0.943$	$0.894$	$0.942$	$0.985$	$21.99$	$40.27$	$7.44$

Appendix

Table 10: All checkpoint evaluation metrics for the 15B model.

Panel of Figure 8 — Figure 8: Simplified profiler traces for the 1.3B, 4B, and 9B models. The corresponding 15B result is shown in Figure 4 .

Muon Optimizer

Muon optimizer uses momentum and matrix orthogonalization to shape neural-network updates, especially in large language-model training.

See all 151 papers

CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
Aug 3, 2026Chuyan Chen, Peng Sun, Kun YuanDiffusion Transformer Image Generation
1Peking University · 2Westlake University · 3Zhejiang University
DMuon: Efficient Distributed Muon Training with Near-Adam Overhead
Jun 25, 2026Vincent Chen, Starrick Liu, Regis Cheng +8Deep Learning Optimization Distributed Optimization
X SQUARE ROBOT TEAM
Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training
Mar 18, 2026Ben S. Southworth, Stephen ThomasDeep Learning Optimization Muon Optimizer
Theoretical Division, Los Alamos National Laboratory, USA. · Lehigh University, Computer Science

Aug 3, 2026J/K move · Enter open · S save

CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization

Chuyan Chen, Peng Sun, Kun Yuan

1Peking University · 2Westlake University · 3Zhejiang University

Diffusion Transformer Image Generation Diffusion Model Acceleration

Open paper AskPDF arXiv

Diffusion Transformers (DiTs) have achieved state-of-the-art (SOTA) performance in visual generative modeling, yet their training remains computationally prohibitive. While the recently proposed Momentum Orthogonalization (Muon) optimizer offers a promising alternative to AdamW, its direct application to DiTs yields suboptimal late-stage convergence. In this paper, we identify the root cause of this bottleneck: standard DiT architectures fuse functionally distinct weights (e.g., within AdaLN and QKV layers) into unified tensors for computational efficiency. Applying Muon to these fused tensors inadvertently induces implicit subspace coupling, which distorts update directions and degrades global optimization. To address this, we introduce Chunked Muon (CMuon), a simple yet highly effective strategy that partitions these matrices into independent sub-components prior to orthogonalization. Extensive experiments demonstrate that a 675M-parameter DiT trained with CMuon achieves a FID of 1.18 on ImageNet 256 in just 200 epochs. This represents more than a 2x training speedup over AdamW, while effectively overcoming the late-stage convergence plateaus of vanilla Muon.

Abstract

Muon Optimizer

Explore similar work

CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

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