Aurora: A Leverage-Aware Optimizer for Rectangular Matrices

TL;DR

Aurora is a Muon-family optimizer that adds a row-norm-uniformity constraint to Muon’s update objective for tall matrices (MLP up/gate projections). The motivation: Muon’s polar-factor step preserves row-leverage anisotropy of the momentum buffer on tall matrices, so neurons that get small early gradients enter a self-reinforcing “death spiral” and never recover — over 25% of MLP neurons end up effectively dead by step 500 of a 340M run. Aurora is derived as steepest descent under joint left-semi-orthogonality and uniform row-norm constraints, implemented as an alternating projection (row-normalize, then Newton–Schulz polar factor, with EMA damping) at ~6% overhead over Muon. A 1.1B Aurora-trained model on 100B NemotronCC tokens reaches Qwen3-1.7B-level scores on HellaSwag and Winogrande (claimed 100× data efficiency), and Aurora sets a new modded-nanoGPT speedrun SoTA at 3175 steps when combined with Contra-Muon + update/weight flooring.

Key claims

For tall matrices, the polar factor preserves row-leverage anisotropy of the input: rows with small momentum norms get small updates after orthogonalization, so neurons that fall behind early stay behind permanently [§3, Claim 2 formal].
At 340M scale, more than one in four MLP neurons are effectively dead by step 500 under Muon — visible as a sharply bimodal leverage-score distribution and as zero-mean / collapsing-variance gate activations [§3, Figs. 9–11].
The “rich get richer” dynamic is self-reinforcing: bottom-quartile-by-leverage neurons at step 100 stay in the dead tail throughout training, while the other three quartiles stabilize at high leverage [§3, Fig. 11].
A tall matrix cannot be simultaneously column-orthogonal and have uniform unit row norms — a tall column-orthogonal matrix’s average squared row norm is $n/m < 1$ [§2, Claim 1].
NorMuon’s post-hoc row-normalization step therefore necessarily introduces a polar-factor precision defect; on tall matrices, U-NorMuon (target row norm $\sqrt{n/m}$ instead of 1) outperforms NorMuon at 340M [§2, Fig. 4].
U-NorMuon’s benefit propagates to wide matrices it doesn’t directly touch: alive up/gate rows induce isotropic column leverage in the down projection, even though the down projection is never row-normalized [§3, Fig. 12].
Aurora is derived as the solution to a joint constrained optimization — maximize alignment with the gradient subject to (a) left-semi-orthogonality and (b) constant unit row norms — and falls out as alternating row-normalization + polar-factor projection, with EMA damping for convergence [§4].
For square matrices Aurora reduces exactly to standard Muon; the divergence is only for tall matrices [GitHub README].
On the modded-nanoGPT speedrun, pure Aurora reaches the 3.28 target loss in 3225 steps (vs NorMuon’s 3250); combined with Contra-Muon + update/weight flooring, Aurora sets a new SoTA of 3175 steps [§5, speedrun table].
At 1.1B / 100B-NemotronCC tokens, Aurora reaches smoothed val loss 2.26 vs Muon 2.31 / NorMuon 2.33; HellaSwag 67.6% (+2.5 over both baselines), and MMLU +10 points over Muon [§5].
Aurora-1.1B matches Qwen3-1.7B on HellaSwag and Winogrande despite 25% fewer parameters, ~360× fewer training tokens (100B vs 36T), and an uncurated data mixture [§5].
Aurora’s gradient alignment at convergence is within 0.6% of the Riemannian reference solver (true joint Stiefel/oblique projection) while being practically implementable, suggesting the alternating approximation is near-optimal [§5, Fig. 18].
Untuned Aurora costs ~6% more wall-clock per step than Muon and is a drop-in replacement [TL;DR].
The Aurora–Muon gap grows monotonically with MLP expansion factor — i.e. wider MLPs amplify exactly the pathology Aurora corrects [§6, Fig. 21].

Method

Aurora reformulates Muon’s update-selection problem for tall matrices ( $n > m$ ) by adding a constant-row-norm constraint to the standard spectral-norm constraint. Combined with the spectral bound this forces $\Delta \in \mathbb{R}^{n \times m}$ to be left semi-orthogonal and have all row norms equal to $\sqrt{m/n}$ . The resulting steepest-descent problem has no closed form, but admits an alternating-projection solution: start from the polar factor of the momentum buffer (orthogonality), row-normalize (uniformity), re-project to the Stiefel manifold via Newton–Schulz, repeat. Convergence is improved with an EMA of the row-normalization factors plus a damping parameter $\beta$ . The paper proves the alternating iteration converges to a point on (or arbitrarily close to) the intersection of the Stiefel and row-oblique manifolds (Appendix F). A Riemannian reference solver — gradient ascent on the joint manifold via tangent-space projection and retraction — is given for verification but is too expensive ( $O(n^3)$ memory for the projection’s linear system) for production use.

Aurora is applied only to MLP up/gate projections (tall). Wide and square matrices (down projections, attention QKV, head) keep the standard Muon update; the down projection’s column-leverage isotropy comes “for free” through the up/gate fix. Untuned overhead is ~6% per step in PE-8 (Polar Express, 8 steps) implementations, and the additional iterations can be overlapped with gradient communication in distributed settings.

Results

modded-nanoGPT speedrun: pure Aurora (lr=0.05, 2 projection iterations, $\beta=0.5$ ) hits target loss in 3225 steps vs NorMuon’s 3250 and Muon’s prior best. With Contra-Muon + update/weight flooring layered on top, 3175 steps — a new SoTA [§5, speedrun table].
1.1B pretraining on ~100B NemotronCC tokens: Aurora final smoothed loss 2.26 vs Muon 2.31 vs NorMuon 2.33; HellaSwag 67.6% (Muon/NorMuon both ~65%); ARC-Challenge and Winogrande similarly improved; MMLU +10 points over Muon [§5].
Versus Qwen3-1.7B (36T tokens): Aurora-1.1B matches on HellaSwag and Winogrande despite 25% fewer parameters and ~360× fewer training tokens; MMLU is lower (NemotronCC has no math/science) [§5].
Update quality: Aurora-2 produces by far the most isotropic leverage-score distribution (tightly clustered around $1/n$ , no dead tail), vs Muon’s heavy tail across orders of magnitude [§5, Fig. 19].
Weight spectra: under NorMuon, the down-projection singular spectrum is inflated ~100× because the weight matrix absorbs the scale information the row-normalized update discarded; Aurora’s spectra are well-conditioned across all three projections [§5, Fig. 20].
MLP-width scaling: Aurora’s advantage over Muon grows monotonically with expansion factor — wider MLPs amplify the row-leverage anisotropy that Aurora fixes [§6, Fig. 21].
Independent third-party replication: per a follow-up Import AI item, Pleias’s Alexander Doria reproduced Aurora > Muon > AdamW at 600M scale.

Why it’s interesting

Aurora is the most mechanistically-derived entry yet on the optimizer-side fork of Training stability at scale. The optimizer-side cluster — Controlled LLM Training on Spectral Sphere (SSO: constrain weight + update spectra), The Newton-Muon Optimizer (Newton-Muon: right-precondition by $(ZZ^\top)^{-1}$ ), Preconditioned Inexact Stochastic ADMM for Deep Models (PISA / SISA / NSISA) (PISA/NSISA: ADMM with NS-orthogonalized momentum), MuonH (hyperball) sets new modded-NanoGPT speedrun record at 3325 steps — Keller Jordan announcing @wen_kaiyue's result (MuonH: replace WD with Frobenius-sphere) — has each member arguing Muon is “missing something.” Aurora’s diagnosis is the most surgical: Muon is missing a row-norm constraint specifically on tall matrices, because the polar factor cannot equalize row norms there (the tall-Stiefel constraint forbids it, [§2, Claim 1]). The fact that NorMuon’s empirically-known wins on the speedrun were not from “normalization” but from breaking the neuron-death feedback loop is the kind of pathology-first explanation that this cluster has been missing.

The architecture-optimizer codesign framing in the conclusion is the second reason to care. Aurora is designed for SwiGLU MLPs with large expansion factors, and its advantage grows with the row-to-column ratio of up/gate projections. That’s a direct prescription for any from-scratch DiT or DiT-style pretraining where the FFN expansion factor is a free hyperparameter — and it’s the cleanest example yet on the wiki of an optimizer change that interacts non-trivially with the FFN width choice. It also raises a sharp open question for the cluster: do Aurora’s gains compose with SSO’s spectral-sphere constraint, or do both saturate the same underlying degree of freedom (bounded row-leverage variance)? No filed paper has run the stacked ablation.

Operationally: untuned Aurora is a 6% drop-in, the aurora-release repo is on GitHub, and the modded-nanoGPT PR #284 lands the SoTA. The bar for trying this on a real run is low — and the dead-neuron diagnostic (per-step leverage scores on the up/gate projection of a middle MLP layer) is a one-screen monitoring widget any training run can add today.