Dion3让优化器更快，提速最高达6倍。
Dion3: Full-Stack Orthogonal Updates
- 用新算法降低正交化计算量，减少浮点操作数
- 通过分块和稀疏更新，使训练速度提升6倍且性能不降
- 适合需要高效优化的深度学习训练场景
Muon优化器因立方时间复杂度的Newton-Schulz正交化步骤而存在显著开销。当权重分片时，通信开销进一步加剧计算成本，削弱了Muon的优势。我们提出Dion3，从全栈层面优化该问题：新提出的Gram Newton-Schulz算法降低正交化计算量；CuteDSL内核利用对称性加速运算；巨批次策略减少通信开销。此外，我们改进更新规则，仅对动量矩阵的部分行进行正交化，相比Dion在速度和性能上均有提升。整体上，Dion3在损失表现上与Muon相当或更优，但优化器单步耗时最多降低6倍。Dion3可通过dion包（https://github.com/microsoft/dion）作为Muon的即插即用替代品获取。
原文摘要 · Abstract (English)
The Muon optimizer incurs a significant overhead cost due to its cubic-time Newton-Schulz orthogonalization step. When weights are sharded, communication overhead compounds this computational cost, eroding the benefits of Muon in many settings. We present Dion3, a revision of Muon that targets this overhead at every level of the stack. Our Gram Newton-Schulz algorithm reduces the FLOP cost of orthogonalization, our CuteDSL kernels accelerate it by exploiting symmetry, and our megabatching strategy reduces communication overhead. Moreover, we propose a simple change to the update rule that cuts costs even further: selecting only a fraction of the momentum matrix's rows to orthogonalize at each step. This update rule improves on Dion (another "compressed" version of Muon), in both speed and performance. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to 6x. Dion3 is available via the dion package (https://github.com/microsoft/dion) as a drop-in replacement for Muon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。