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

AI 中文解读

AI绘画大模型训练有了新突破！这项叫CMuon的技术，直接把图像生成模型DiT的训练速度提了2倍多。以前训练这类模型用的是AdamW优化器，虽然好用但速度慢；后来有人提出Muon方法，效果更好但到了训练后期容易“卡壳”，进展缓慢。研究人员发现，问题出在模型内部把不同功能的参数“揉”在一起处理，导致更新方向被干扰。他们的解决方案很简单——把这些参数拆开，各自独立优化。实际效果很惊艳：一个拥有6.75亿参数的模型，用CMuon训练200轮，生成图像质量评分达到1.18，比AdamW整整快了一倍。这意味着，未来AI绘画、AI视频生成的训练成本大幅下降，更快迭代出新版本。对普通人来说，不久的将来，你打开AI绘图软件，生成一张高质量图片可能只需几秒，背后的引擎却比现在便宜、高效得多。AI创作工具的门槛降低，内容生产者的效率将大幅提升。

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.