光谱分配：为什么 Muon 优于 Adam，以及如何改进 Muon

相对于 Adam，诸如 Muon 之类的正交优化器可以显着加速 大语言模型 预训练，但其机制仍不完全清楚。我们通过对 Transformer 损失景观进行样本外光谱探测分析来对此进行研究。在沿着真实训练轨迹的检查点，我们将每个动量缓冲区分解为其奇异方向，并根据保留的数据估计沿着每个方向的损失最佳步长。得到的光谱分布是各向异性的，但在批次和训练阶段是稳定的，并且在优化器和模型尺度上是一致的：在稳定边缘运行的易失性头支持比容错体小得多的步长，这允许更大的步长。该剖面提供了统一的光谱分配说明，说明为什么 Muon 优于 Adam，而 Adam 又优于 SGD。它还暴露了 Muon 均匀缩放的局限性：它仍然没有充分利用体积。在这一发现的指导下，我们引入了光谱感知 Muon (SAMuon)，它将头部保持在 Muon 尺度，并使用静态光谱先验放大体积。我们提供两种变体：完整的 SAMuon 使用低秩随机 SVD 遵循测量轮廓，而简化的 SAMuon-lite 通过一级幂迭代使用两级近似。这两种方法都没有在规模上添加持久优化器状态或超出 Muon 的显着额外 FLOP，并且两者的理想化精确白化版本都保留了标准假设下 Muon 的渐近收敛率。在从 124M 到 1B 参数的“modded-nanogpt”模型中，两种变体在所有评估的模型规模和批量大小配置中均优于调整的 AdamW 和 Muon（Scion 实现）基线。 SAMuon 需要减少 13.3% 到 24.0% 的训练词元才能达到与 Muon 相同的验证损失，而 SAMuon-lite 保留了大部分增益，且挂钟开销接近于零。