LionMuon：交替使用谱下降与符号下降实现高效训练

arXiv:2609.35297v3 公告类型：替换 摘要：预训练语言模型需要大量计算资源，而合适的优化器可以节省其中相当一部分。Muon 的谱更新步能提供比符号更新步更强的方向，但代价高昂。每一步都要对完整矩阵执行 Newton-Schulz 迭代，在分布式训练中还需要额外执行一次全归约。Lion 和 Signum 所采用的符号更新步成本低，且始终在各设备本地执行。我们提出 LionMuon，每隔 $P$ 次迭代执行一次 Muon 更新步，其间执行 Lion 更新步，两者共享一个双指数移动平均（EMA）动量缓冲区。Muon 的计算和通信开销每 $P$ 步才产生一次，且优化器状态大小仅为 AdamW 的一半。采用单 EMA 的变体 SignMuon 就已优于 Muon。我们证明了重尾噪声下的复杂度界，其中周期决定了 Muon 与 Lion 的平滑性常数和噪声常数之间的插值，并揭示了 LionMuon 何时比两者都更快。在 FineWeb 上训练的 124M 和 355M 模型中，当处理的词元数量相同时，$P=2$ 和 $P=5$ 的 LionMuon 均取得了比 Muon、AdamW、Lion 和 Signum 更低的损失。在采用 PCIe 互连的 4-GPU 数据并行训练中，它达到 Muon 最终损失所需的实际运行时间减少了三分之一；同时，在不增加未被计算掩盖的通信开销且保留精确梯度的情况下，其损失表现优于通信高效的 Muon 变体 Dion 和 MuonBP。代码：https://github.com/brain-lab-research/lion-muon