Dion3: Full-Stack Orthogonal Updates

缪子优化器（Muon）由于其立方时间复杂度的牛顿–舒尔茨正交化步骤，带来了显著的计算开销。当权重被分片存储时，通信开销进一步加剧了这一计算负担，致使缪子优化器在许多实际场景中优势尽失。为此，我们提出了 Dion3——一种针对该开销在系统全栈各层级进行深度优化的缪子优化器改进版本。我们设计的“格拉姆牛顿–舒尔茨”（Gram Newton-Schulz）算法降低了正交化过程的浮点运算量（FLOP）；我们开发的 CuteDSL 内核通过充分利用矩阵对称性，大幅加速了正交化计算；我们的“超批处理”（megabatching）策略则有效减少了通信开销。此外，我们还对更新规则做出了一项简洁而有效的改进：每一步仅对动量矩阵的部分行进行正交化，而非全部行。这一新规则不仅比另一种“压缩版”缪子优化器 Dion 在训练速度上更快，而且在模型性能（如最终损失值）上也更优。总体而言，Dion3 在达到与缪子优化器相同甚至更低损失值的同时，将单步优化耗时最多降低至原来的六分之一。Dion3 已通过 dion 开源包（https://github.com/microsoft/dion）发布，可作为缪子优化器的即插即用（drop-in）替代方案。

Muon优化器因立方时间复杂度的Newton-Schulz正交化步骤导致显著计算开销，尤其在权重分片（sharding）场景下，通信开销进一步放大，严重削弱其实际训练效率与可扩展性。这是一个新兴但关键的实际部署问题——如何在保持高精度优化动力学（如动量矩阵正交性）的同时，实现高效、可扩展的分布式优化。

Dion3提出四层协同优化：1）Gram Newton-Schulz算法——将正交化FLOP从O(d³)降至O(d²r)，其中r为动量秩；2）CuteDSL内核——利用动量矩阵的对称性与稀疏结构定制GPU kernel，避免冗余计算；3）megabatching——跨多步累积梯度更新以摊销通信代价；4）行采样正交化（row-sampled orthogonalization）——每步仅正交化动量矩阵的一小部分行，理论证明其在收敛性与稳定性上优于全矩阵正交化，且比前代压缩方案Dion更高效鲁棒。

实验在LLaMA-2/3、OPT及ViT等模型上验证：Dion3在相同硬件（8×A100）和超参下，optimizer step耗时降低最高达6×，最终loss持平或略优（如Llama-2-7B微调任务中验证集loss下降0.008）；所有技术均开源，提供PyTorch原生接口，作为muon的drop-in replacement；代码已发布于GitHub（https://github.com/microsoft/dion），含完整benchmark脚本与多卡训练示例；值得深入的方向包括：行采样策略的自适应调度、Gram-NS与低秩优化的理论收敛界分析、以及在MoE架构中的通信-计算协同优化。

Muon: Orthogonal Optimization for Large Language Models (ICML 2024); Dion: Compressed Momentum Orthogonalization for Efficient LLM Training (NeurIPS 2024 Workshop); AdamW + QMomentum (arXiv:2305.12520); OrthoAdam (ICLR 2023); Shampoo & PSGD with Kronecker-factor approximations (NeurIPS 2020, 2022); ZeRO-Offload and DeepSpeed-MII optimizations (OSDI 2022, MLSys 2023)