CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization· CMuon：通过分块动量正交化加速和稳定扩散 Transformer 训练
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
通过分块动量正交化优化扩散 Transformer 训练，提高稳定性和收敛速度。
- 核心方法
- 引入 Chunked Muon (CMuon)，将功能不同的权重分块处理，避免隐式子空间耦合，改进优化方向。
- 适合谁读
- 研究者、工程师
- 要解决的问题
- 扩散 Transformers (DiTs) 训练计算成本高，使用 Momentum Orthogonalization (Muon) 优化器时晚期收敛效果不佳。
- 关键实验
- 通过多个视觉生成任务实验验证了 CMuon 的性能，显示了其在加速训练和提高稳定性方面的优势。
- 主要贡献
- 提出了 CMuon 优化器，显著提高了 DiTs 的训练速度和稳定性，改善了晚期收敛问题。
- 意义与局限
- CMuon 为 DiTs 训练提供了有效的优化方案，降低了计算成本，提升了模型性能，但其在其他类型模型中的适用性仍需进一步验证。