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

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.