Search papers, labs, and topics across Lattice.
This paper introduces Chunked Muon (CMuon), an optimization strategy that enhances the training of Diffusion Transformers (DiTs) by addressing the convergence issues associated with the standard application of Momentum Orthogonalization (Muon). By partitioning the fused tensors in DiT architectures into independent sub-components before orthogonalization, CMuon stabilizes update directions and improves global optimization. Experimental results show that a 675M-parameter DiT trained with CMuon achieves a state-of-the-art FID of 1.18 on ImageNet 256 in just 200 epochs, more than doubling the training speed compared to AdamW while overcoming late-stage convergence challenges.
Chunked Muon accelerates Diffusion Transformer training by over 2x while achieving state-of-the-art performance, solving critical convergence issues.
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.
