CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
Ranking
Overall
67
Content
80
Popularity
36
Observed public metrics from 1 member.
Merged summary
TL;DR - CMuon is a modification of the Muon optimizer that splits fused weight tensors in Diffusion Transformers into independent chunks before momentum orthogonalization, removing a convergence bottleneck and cutting DiT training cost by more than half versus AdamW.
- Diagnosis: standard DiT architectures fuse functionally distinct weights (e.g., AdaLN and QKV projections) into single tensors for efficiency; orthogonalizing these fused tensors with vanilla Muon induces implicit subspace coupling that distorts update directions and hurts global optimization.
- Fix: partition the fused matrices into independent sub-components before the orthogonalization step — a simple, architecture-aware change rather than a new optimizer family.
- Results: a 675M-parameter DiT reaches FID 1.18 on ImageNet 256 in 200 epochs, reported as >2x training speedup over AdamW while avoiding vanilla Muon's late-stage convergence plateaus.
- Relevance: targets the practical bottleneck of prohibitively expensive visual generative model pretraining, suggesting optimizer design must account for how weights are packed in the implementation.
Sources (1)
CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
Public signals
Semantic Scholar citations 0 · Semantic Scholar influential citations 0
TL;DR - CMuon is a modification of the Muon optimizer that splits fused weight tensors in Diffusion Transformers into independent chunks before momentum orthogonalization, removing a convergence bottleneck and cutting DiT training cost by more than half versus AdamW.
- Diagnosis: standard DiT architectures fuse functionally distinct weights (e.g., AdaLN and QKV projections) into single tensors for efficiency; orthogonalizing these fused tensors with vanilla Muon induces implicit subspace coupling that distorts update directions and hurts global optimization.
- Fix: partition the fused matrices into independent sub-components before the orthogonalization step — a simple, architecture-aware change rather than a new optimizer family.
- Results: a 675M-parameter DiT reaches FID 1.18 on ImageNet 256 in 200 epochs, reported as >2x training speedup over AdamW while avoiding vanilla Muon's late-stage convergence plateaus.
- Relevance: targets the practical bottleneck of prohibitively expensive visual generative model pretraining, suggesting optimizer design must account for how weights are packed in the implementation.