面向 $\mu\mathsf{P}$ 的 Muon: 在训练期间确保谱条件

$\mu$-参数化 ($\mu$P$) 为大型语言模型 (LLM) 训练提供了原则性基础, 通过规定与宽度无关的学习动力学, 从而实现可预测的规模行为和跨模型大小的鲁棒超参数迁移。$\mu$P$ 的核心要求是满足若干谱条件, 这些条件确保在模型宽度增长时保持一致的特征学习和优化行为。虽然这些条件在理论上已充分理解, 但在实际训练中保证其有效性仍是未充分研究的课题。现有研究要么未确保谱条件在整个训练期间都成立, 要么需要对权重和更新都施行重复的谱归一化 (或 Newton-Schulz 迭代), 导致显著的计算开销和实用性下降。本文我们证明, 在模型宽度适中时,仅在优化器更新层面维持谱控制足以保持 $\mu$P$-兼容的规模, 无需对权重进行显式谱归一化。基于此原理, 我们发展了 Muon 的一种变体, 即 Muon++, 它在整个训练过程中满足谱条件。我们的结果弥合了$\mu$P$ 的理论承诺与基于矩阵的优化器在长期训练中的实际部署之间的鸿沟。我们还首次通过纳入数据依赖效应来实现自适应谱条件, 以更好地适应长期 LLM 训练。