CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
2026-08-03 • Artificial Intelligence
Artificial Intelligence
AI summaryⓘ
The authors studied a type of AI model called Diffusion Transformers (DiTs), which are very good at creating images but slow to train. They looked at a new training method called Momentum Orthogonalization (Muon), which can speed things up but didn’t work well with DiTs because of how the model combines certain parts for efficiency. To fix this, they developed Chunked Muon (CMuon), which splits these combined parts before training, avoiding problems with updates. This new approach allowed much faster training without losing quality, achieving strong results on a large image dataset.
Diffusion TransformersMomentum OrthogonalizationAdamW optimizerAdaLNQKV layersorthogonalizationFID scoreImageNettraining convergencesubspace coupling
Authors
Chuyan Chen, Peng Sun, Kun Yuan
Abstract
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.