提出Aurora优化器,解决大矩阵更新不均问题。
Aurora: A Leverage-Aware Spectral Optimizer

- 通过保持行均匀性并保留动量极分解几何结构优化矩阵更新
- 在预训练中优于Muon,且在modded-nanoGPT速度赛中达顶尖性能
- 对超宽MLP层效果更显著,适合大规模模型训练
我们发现,对于高瘦型矩阵参数(如MLP层中的投影矩阵),Muon更新的行范数可能任意不均匀。这会引发自我强化反馈环,导致部分神经元持续接收微小更新,最终对网络输出无贡献。虽然增加行归一化可缓解此问题,但现有方法会改变Muon更新的极分解几何结构,我们发现这是不可取的。为此,我们提出Aurora:一种在保持行均匀性的同时尊重Muon极分解几何的优化器。Aurora在预训练实验中表现优于Muon;与现有方法结合后,在modded-nanoGPT速度赛的优化器赛道上达到当前最佳性能。此外,我们发现Aurora相对于Muon的增益随MLP扩展因子增大而提升,表明其可能支持高效训练极宽的MLP层。
原文摘要 · Abstract (English)
We show that for tall matrix parameters, like projection matrices in the MLP layers, the Muon update can have row norms that are arbitrarily non-uniform. This can lead to a self-reinforcing feedback loop whereby neurons receive persistently small updates and eventually do not contribute meaningfully to network outputs. This problem is effectively mitigated by an additional row normalization step, but current methods do this in a way that moves the Muon update geometry away from the polar factor of the momentum matrix, which we find is undesirable. We propose Aurora, an optimizer that enforces row-uniformity of matrix parameter updates while respecting Muon's polar factor geometry. Aurora outperforms Muon in our pre-training experiments and, when combined with existing methods, achieves state-of-the-art performance among spectral optimizers on the optimizer track of the modded-nanoGPT speedrun. Additionally, we find that Aurora's empirical gains over Muon scale with the MLP expansion factor, suggesting that Aurora may allow for effective training of very wide MLP layers.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。