论文

Muon 是否需要细粒度谱整形?

Does Muon Need Fine-Grained Spectral Shaping?

模型训练预训练

摘要

Muon 将当前和过去的梯度组合成矩阵动量。对于 $M=UΣV^\top$,理想化的极分解更新 $Q=UV^\top$ 为每个奇异方向赋予相同的权重。我们将此称为平坦谱型。最近的几个优化器用细粒度的谱映射取代了这种平坦的轮廓,这些谱映射为每个方向提供了自己的增益。我们询问 Muon 更新需要多少谱细节。我们的谱诊断表明,测量的奇异模式的大约 $94$--$97\%$ 位于估计的噪声边缘下方,但总体上与参考梯度正向对齐。我们引入了 BulkBoost,一种具有固定秩和噪声校准变体的两频段频谱重新加权框架。后者使用分割小批量梯度差异来校准 Muon Nesterov 输入的 Marchenko-Pastur 参考边缘,将边缘下方的大部分与其上方的尖峰分开。两种变体都通过一个共享增益增加了谱主体的相对权重,同时保留了每个矩阵未重新加权方向的弗罗贝尼乌斯范数。对于固定分区,我们的理论给出了一阶条件,在该条件下,将权重移向谱主体会降低损失。它还量化了两个频段可以捕获的所有每模式重新分配的最大一阶改进率的分数。在涵盖 Pythia-14M 至 410M 和 6 个语料库的 30 个持续预训练设置中,两波段重新加权可与 Freon 的细粒度幂律曲线相媲美,并且优于 Spectra。根据 Muon 的平坦轮廓进行测量,Freon将预适应损失的最终损失平均减少了 $0.022\%$,而两波段变体实现了 $0.073$--$0.147\%$ 的减少。这些观察表明,与平坦轮廓的有用偏离是令人惊讶的低维:单一主体对尖峰增益至少捕获与细粒度谱轮廓一样多的好处。