微软、哈佛开源创新优化器：全面超越Muon，提升大模型训练效率

64 0 0

文章摘要

大模型训练所需计算资源随功能增强呈爆炸式增长，优化器可节省训练资源。此前的Muon优化器虽有优势，但在大规模训练中应用效率低。为此，微软和哈佛大学研究团队联合开源了优化器Dion，其开源地址为https://github.com/microsoft/dion/?tab=readme-ov-file。

Dion的技术核心是通过创新的低秩近似和分布式计算策略，在大规模分布式训练中实现高效正交归一化更新，避开传统方法的计算与通信瓶颈。在集中式场景下，Dion围绕矩阵参数优化展开工作流程，构建缓冲区累加动量信息与当前梯度，采用低秩近似技术，用两个较小矩阵近似缓冲区，通过幂迭代实现且仅需一次迭代达高精度，避免复杂矩阵分解。同时引入误差反馈机制，计算近似结果与原始缓冲区误差并注入下一次动量更新，维持优化稳定性。之后对近似矩阵正交化处理，用缩放因子调节更新强度，提升超参数迁移能力。

在分布式场景中，面对参数分片存储挑战，Dion采用分布式幂迭代与正交化技术，各设备处理本地参数分片并高效通信聚合结果，避免全局矩阵重构。其分布式正交化采用两步迭代，将全局正交化转化为本地计算与小规模通信结合，降低通信开销。解耦动量机制解决各设备动量不一致问题，保证分布式计算准确性。

效率方面，Dion通过低秩近似降低计算和通信复杂度，即使在全秩设置下计算量也比Muon低30%以上。对非矩阵参数采用兼容策略，减轻超参数调优负担，实现端到端高效训练。

研究团队在120M到3B参数的语言模型上测试，Dion表现出色。在3B参数模型上，不同低秩设置验证损失与全秩Muon接近，且模型越大差距越小。全秩Dion收敛速度超Muon和AdamW，在3B模型上比Muon快30%。随批处理量增大优势更明显，在更新密度实验中性能平稳下降，速度测试中达到目标损失所需时间仅为AdamW的1/3，比Muon快25%。