Aurora: A Leverage-Aware Spectral Optimizer
For deep learning practitioners training large models with wide MLP layers, Aurora improves optimization stability and performance over existing spectral optimizers.
Aurora addresses the issue of non-uniform row norms in Muon updates for tall matrices, which can cause neurons to become inactive. It enforces row-uniformity while preserving Muon's polar factor geometry, outperforming Muon in pre-training and achieving state-of-the-art results among spectral optimizers on the modded-nanoGPT speedrun.
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.