MuonBP:基于块周期正交化的更快Muon
机器学习
2025-10-21 v1 最优化与控制
摘要
梯度正交化是一种简单却在加速梯度下降方面显示出巨大效用的策略。Muon优化器(Jordan、Jin等,2024)将梯度正交化与一阶动量相结合,在语言模型训练中相较于Adam/AdamW(Loshchilov和Hutter,2019)在数据效率上实现了显著提升。然而,在使用模型并行时,梯度正交化由于在不同设备上的梯度矩阵碎片进行额外的收集和散播操作,相对于坐标级优化器(如AdamW)会引入额外的开销。这种额外的通信量在吞吐量上可能下降5%-10%。为此,我们提出带块周期正交化的Muon(MuonBP),该方法在每个设备上独立应用正交化,并定期执行完整正交化以保持大规模训练的稳定性。我们展示了如何从基线调整学习率至MuonBP,并为该算法提供收敛保证。关键的是,我们的理论指出,我们需要两种步长:一种用于块状正交化步骤,另一种用于完整正交化步骤。我们的方法简单,几乎不需要调整超参数,能够在与基线Muon相当的迭代复杂度下,实现与坐标级方法(如AdamW)相当的每迭代吞吐量。在使用八路张量并行和ZeRO优化器状态分片训练8B模型时,MuonBP相对于Muon实现了8%的吞吐量提升,且未出现性能下降。
引用
@article{arxiv.2510.16981,
title = {MuonBP: Faster Muon via Block-Periodic Orthogonalization},
author = {Ahmed Khaled and Kaan Ozkara and Tao Yu and Mingyi Hong and Youngsuk Park},
journal= {arXiv preprint arXiv:2510.16981},
year = {2025}
}