CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
论文提出 Chunked Muon (CMuon) 优化器，用于加速和稳定 Diffusion Transformer (DiT) 训练。标准 DiT 将功能不同的权重（如 AdaLN 和 QKV 层）融合为统一张量，直接应用 Muon 会导致隐式子空间耦合，损害优化。CMuon 在正交化前将这些矩阵划分为独立子组件。675M 参数 DiT 使用 CMuon 在 ImageNet 256 上 200 epoch 达到 FID 1.18，比 AdamW 快 2 倍以上，并克服了 Muon 的后期收敛平台。
发展脉络
- 首次出现CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum OrthogonalizationarXiv cs.AI
- 当前判断训练加速对扩散模型商业化至关重要。CMuon 提供 2 倍加速，可能降低训练成本，但需验证在更大模型和数据集上的泛化性。可验证信号：是否有第三方复现或采用。Agent Pulse · 分析
Diffusion Transformers (DiTs) 在视觉生成中达到 SOTA，但训练计算成本高。近期提出的 Momentum Orthogonalization (Muon) 优化器是 AdamW 的替代方案，但直接应用于 DiTs 时后期收敛不佳。论文指出根因：标准 DiT 架构将功能不同的权重（如 AdaLN 和 QKV 层）融合为统一张量以提高计算效率，对融合张量应用 Muon 会引入隐式子空间耦合，扭曲更新方向并损害全局优化。为此，论文提出 Chunked Muon (CMuon)，在正交化前将矩阵划分为独立子组件。实验表明，675M 参数 DiT 使用 CMuon 在 ImageNet 256 上 200 epoch 达到 FID 1.18，比 AdamW 快 2 倍以上，并有效克服 Muon 的后期收敛平台。
CMuon 的核心在于将融合张量按功能划分后分别正交化，避免子空间耦合。这提示优化器设计需与架构的权重布局对齐。可验证信号：后续工作是否将 CMuon 应用于其他融合权重架构（如多模态模型）并复现加速效果。
训练加速对扩散模型商业化至关重要。CMuon 提供 2 倍加速，可能降低训练成本，但需验证在更大模型和数据集上的泛化性。可验证信号：是否有第三方复现或采用。
对训练扩散模型的企业，CMuon 可降低计算成本和时间，加速模型迭代。但需评估在自有数据上的效果。
CMuon 可能成为 DiT 训练的标准优化器，并扩展到其他架构。未来可能看到与架构搜索结合，或针对融合权重设计更通用的优化方法。