用于 LLM 预训练的具有球体约束的曲率条件多尺度动量

预训练占 LLM 训练总计算成本的很大一部分。然而，噪声主导的梯度和高度病态的损失景观带来了严峻的挑战。尽管 AdamW 和 Muon 等现代自适应优化器在大规模预训练中取得了巨大成功，但它们对梯度归一化的依赖对病态曲率的缓解效果有限。沿着平坦方向（小特征值的特征方向）的进展，主导着最终的损失减少，仍然相对缓慢。为了增强沿平面方向的训练动态，我们提出了一种具有球体约束的曲率条件多尺度动量方法，在 LLM 预训练中提供稳定的加速度。这种仅沿平面方向施加的多尺度动量，将用于降噪的慢衰减分量与用于快速曲率适应的快衰减分量配对，利用它们的互补优势。至关重要的是，我们采用球体约束技术来防止参数膨胀和过快的有效学习率衰减，否则这些问题可能会因简单的组合而产生。大量实验表明，所提出的方法显着加速了跨不同架构（密集、MoE）和模型大小（0.12B--2.3B 参数）的 Muon。从理论上讲，我们验证了加速效果，并深入了解了平面方向多尺度动量的设计原理。