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

2026-08-03

AI Summary

This work addresses the high training cost of Diffusion Transformers (DiT) and the convergence difficulties encountered by the existing momentum-based orthogonal optimizer, Muon, when applied to fused weight matrices. To mitigate implicit subspace interference arising from functionally heterogeneous fused weights in DiT, the authors propose Chunked Muon (CMuon), which decomposes these fused weight matrices into independent sub-blocks before applying momentum orthogonalization. This decoupling strategy substantially enhances training stability and convergence speed. On ImageNet 256, CMuon achieves a state-of-the-art FID of 1.18 with a 675M-parameter DiT in only 200 epochs—more than twice as fast as AdamW—and successfully overcomes the convergence bottleneck inherent in the original Muon optimizer.

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.

Problem

• Diffusion Transformers
• Momentum Orthogonalization
• weight fusion
• subspace coupling
• late-stage convergence

Innovation

• Chunked Muon
• Diffusion Transformer
• Momentum Orthogonalization
• Subspace Decoupling
• Training Acceleration