提出CMuon方法，让扩散Transformer训练更快更稳
CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
- 将权重分块后分别正交化，避免更新方向扭曲
- 675M参数模型200轮达FID 1.18，速度超AdamW两倍
- 适合追求高效稳定训练的视觉生成研究者
扩散Transformer（DiT）在视觉生成建模中达到顶尖性能，但训练成本仍极高。尽管近期提出的动量正交化（Muon）优化器为AdamW提供了有前景的替代方案，但直接应用于DiT时会出现后期收敛不佳的问题。本文发现根本原因在于：标准DiT架构为提升效率，将功能不同的权重（如AdaLN和QKV层）融合成统一张量。对这些融合张量应用Muon会引发隐式子空间耦合，扭曲更新方向并损害全局优化。为此，我们提出分块动量正交化（CMuon），在正交化前将矩阵划分为独立子组件。大量实验表明，使用CMuon训练的675M参数DiT仅需200个周期即可在ImageNet 256上实现FID 1.18，相比AdamW提速超2倍，且有效克服了原始Muon的后期收敛瓶颈。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。