Aurora: A Leverage-Aware Spectral Optimizer
This work addresses the issue in the Muon optimizer where highly anisotropic row norms in tall-and-skinny matrices—such as those in MLP projection layers—lead to severely imbalanced neuron updates, causing some neurons to receive negligible or ineffective gradients. To resolve this, the authors propose the Aurora optimizer, which introduces a row-normalization mechanism while preserving the geometric structure of the momentum matrix’s polar factor. Aurora is the first method to achieve uniformly scaled row-wise updates without compromising this intrinsic geometry—a limitation inherent in prior approaches that trade geometric fidelity for update uniformity. By integrating spectral optimization techniques with polar factor constraints, Aurora outperforms Muon in pretraining and achieves state-of-the-art performance among spectral optimizers on the modded-nanoGPT speedrun benchmark, with gains that become increasingly pronounced as the MLP expansion factor grows.