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
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
FD-DINO ↓
FID ↓
MMD-DINO ↓
Precision ↑
Recall ↑
Coverage ↑
Density ↑
HPSv2 ↑
GenEval2 AM ↑
GenEval2 GM ↑
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
0.0268
0.9487
0.9071
0.9492
0.9918
20.84
44.72
9.38
Periodic Row-wise Muon
45.65
6.29
0.0270
0.9517
0.9065
0.9485
0.9756
20.36
45.89
9.75
4B
AdamW
51.39
7.35
0.0271
0.9371
0.8995
0.9388
0.9898
21.40
45.70
9.13
Muon
41.62
6.08
0.0261
0.9385
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
Step time ↓
Opt. time ↓
Opt. share ↓
AG/step ↓
AR/step ↓
Comm. volume ↓ GiB/rank/step
1.3B
Muon
1.873
0.835
44.6%
10.0
0.0
2.38
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
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
γ
Best FD-DINO ↓
Vanilla Muon
Every step
–
–
41.733
RowNorm every step
None
RowNorm
0.15
51.322
Periodic Row-wise Muon ( γ=1 )
K=3
RowNorm
1.0
44.294
Scalar-controlled ( γ for both NS5 and RowNorm)
K=3
Scalar-controlled
0.15
44.713
Periodic Row-wise Muon ( γ=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×10−4
2.0×10−4
1.4×10−4
1.2×10−4
Muon
5.0×10−3
3.0×10−3
2.4×10−3
2.0×10−3
Periodic Row-wise Muon
5.0×10−3
3.0×10−3
2.4×10−3
2.0×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 ϵ=10−8 .
Setting
Value
Image resolution
512×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 γ . (a) Late-stage FD-DINO, averaged over the 20k, 25k, and 30k checkpoints, across different K and γ . (b) The best mean FD-DINO over γ 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 γ=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
FD-DINO ↓
FID ↓
MMD-DINO ↓
Precision ↑
Recall ↑
Coverage ↑
Density ↑
HPSv2 ↑
GenEval2 AM ↑
GenEval2 GM ↑
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
FD-DINO ↓
FID ↓
MMD-DINO ↓
Precision ↑
Recall ↑
Coverage ↑
Density ↑
HPSv2 ↑
GenEval2 AM ↑
GenEval2 GM ↑
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
FD-DINO ↓
FID ↓
MMD-DINO ↓
Precision ↑
Recall ↑
Coverage ↑
Density ↑
HPSv2 ↑
GenEval2 AM ↑
GenEval2 GM ↑
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
FD-DINO ↓
FID ↓
MMD-DINO ↓
Precision ↑
Recall ↑
Coverage ↑
Density ↑
HPSv2 ↑
GenEval2 AM ↑
GenEval2 GM ↑
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.
Figure 8: Simplified profiler traces for the 1.3B, 4B, and 9B models. The corresponding 15B result is shown in Figure 4 .
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.
Chuyan Chen, Peng Sun, Kun Yuan
1Peking University · 2Westlake University · 3Zhejiang University
Matrix-orthogonalization-based optimizers, exemplified by Muon, have demonstrated strong convergence behavior across a wide range of modern deep learning workloads. The matrix-aware updates offer a compelling alternative to conventional element-wise optimization, particularly as model architectures continue to grow in scale and heterogeneity. Yet contemporary distributed training infrastructure built around the assumption of element-wise optimizers is poorly matched to matrix-level optimizers such as Muon, whose updates couple entire weight matrices and require costly Newton-Schulz iterations. Vanilla Muon implementations incur more than 2x the cost of forward and backward passes. To close this gap, we present DMuon, an open-source distributed Muon implementation that integrates into existing training pipelines as a drop-in module, with no framework-level modifications. Across both embodied foundation model and large language model (LLM) training workloads, DMuon achieves a 1.48x-3.01x speedup in end-to-end step time and a 6.85x-163.00x speedup in optimizer-step time, bringing per-step latency to near-AdamW levels and enabling efficient scaling in our model training.
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× across most settings and up to nearly 3× 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.
Ben S. Southworth, Stephen Thomas
Theoretical Division, Los Alamos National Laboratory, USA. · Lehigh University, Computer Science