🤖 machine learning

Aurora: A Leverage-Aware Spectral Optimizer

The paper introduces Aurora, a leverage-aware spectral optimizer that enforces row-uniformity in matrix parameter updates while preserving Muon's polar factor geometry, thereby preventing neuron stagnation and achieving state-of-the-art performance, particularly in training very wide MLP layers.

Original authors: Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

Published 2026-06-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a massive team of workers (a neural network) to solve complex puzzles. In this team, there are specific departments called MLP layers where the work is divided into individual stations. Each station is a "neuron" responsible for a tiny part of the final answer.

The paper introduces a new manager named Aurora who fixes a major flaw in how the previous manager, Muon, was assigning work.

The Problem: The "Lazy Worker" Loop

Under the old system (Muon), the manager looked at the whole team's performance and tried to balance the work. However, for certain types of tasks (represented mathematically as "tall matrices"), Muon accidentally created a vicious cycle:

The Uneven Split: Some workers (neurons) got huge, exciting updates to their skills, while others got tiny, almost invisible updates.
The Death Spiral: The workers who got tiny updates started to fade away. Because they weren't learning, they stopped contributing to the team's output. Eventually, they became "dead neurons"—they were still on the payroll but did nothing.
The Feedback Loop: As these workers stopped contributing, the manager (Muon) saw less need to update them, so they got even smaller updates. It was a self-reinforcing loop of neglect.

The paper shows that in large models trained with Muon, a significant number of these workers essentially "quit" on the job, especially in the early layers of the network.

The Old Fix: The "Blunt Force" Approach

Some researchers tried to fix this by simply forcing every worker to have the same amount of energy (row normalization). They would look at the team, see someone was weak, and just boost them up.

However, the paper argues this blunt approach was like trying to fix a delicate watch with a sledgehammer. While it stopped the workers from dying, it distorted the shape of the work. It forced the team into a geometry that didn't fit the natural flow of the puzzle, making the whole team less efficient. It was a trade-off: you saved the workers, but you broke the rhythm of the work.

The New Solution: Aurora

Aurora is a new optimizer that solves this problem without breaking the rhythm. Think of Aurora as a master choreographer who understands two rules at once:

The Geometry Rule: The team must move in a specific, elegant, orthogonal pattern (like a perfectly synchronized dance) to solve the puzzle efficiently. This is the "polar factor" that Muon was good at.
The Fairness Rule: Every single worker must get an equal amount of energy and attention.

Aurora manages to do both simultaneously. Instead of just boosting the weak workers after the fact, it calculates the perfect move that satisfies both the dance pattern and the fairness requirement in one go.

How it works in practice:

It looks at the "tall" parts of the network (the up and gate projections in the MLP layers).
It runs a quick, iterative calculation (like a few rounds of mental math) to find the perfect update that keeps the team's dance steps perfect while ensuring no one is left in the dark.
It prevents the "dead neuron" problem entirely.

The Results

The authors tested Aurora on large language models (340 million and 1.1 billion parameters) and found:

No Dead Neurons: Unlike Muon, Aurora kept every single neuron active and contributing.
Better Performance: The models trained with Aurora learned faster and got higher scores on reasoning and knowledge tests (like MMLU) than those trained with Muon or the "blunt force" fixers.
The Wider, The Better: The paper found that the wider the team (the more neurons in the MLP layers), the bigger the advantage Aurora had. This suggests Aurora is the key to training extremely wide, powerful networks without wasting resources on dead workers.

The Bottom Line

Aurora is a smarter way to train AI models. It fixes a hidden flaw where AI "neurons" were quietly dying off, ensuring that every part of the network is utilized effectively, all while keeping the mathematical structure of the learning process perfect. It's like upgrading from a manager who accidentally ignores half the staff to one who ensures everyone is dancing in perfect sync.

Technical Summary: Aurora – A Leverage-Aware Spectral Optimizer

Problem Statement

The paper identifies a critical pathology in the training of large language models using the Muon optimizer, a popular spectral optimizer that approximates weight updates using the polar factor of the gradient momentum matrix. Specifically, for tall matrix parameters (where the number of rows $m$ exceeds columns $n$ , such as the up and gate projection matrices in MLP layers), Muon updates can exhibit highly non-uniform row norms.

This non-uniformity creates a self-reinforcing feedback loop:

Neurons corresponding to rows with small norms receive persistently small updates.
These neurons fail to contribute meaningfully to network outputs, effectively becoming "dead."
The paper provides mechanistic evidence linking this phenomenon to leverage scores, showing that low row norms in the momentum buffer correlate with low leverage scores, which Muon preserves rather than corrects.

The authors observe that while existing row-normalization methods (e.g., NorMuon) mitigate this issue by rescaling rows, they do so in a way that distorts the geometry of the Muon update, moving it away from the true polar factor. This trade-off between enforcing uniform row norms and preserving the precise polar factor geometry is the central problem addressed.

Methodology

1. Analysis of Normalization and Polar Precision

The authors first analyze the tension between polar factor precision and row normalization.

Polar Approximation Error (PAE): They demonstrate that higher precision in computing the polar factor (via Newton-Schulz iterations like CANS-12) correlates with lower validation loss.
Geometric Conflict: They prove that for tall matrices, it is mathematically impossible to simultaneously satisfy column-orthogonality (the polar factor property) and unit row norms. Post-hoc normalization forces the update away from the true polar factor, introducing significant geometric error.
U-NorMuon: As an intermediate step, they propose U-NorMuon, which applies row normalization only to tall matrices and removes stateful momentum buffers for normalization. This isolates the benefit of row normalization without the overhead of stale statistics, showing it outperforms standard Muon and NorMuon.

2. The Aurora Optimizer

To resolve the conflict between preserving polar geometry and enforcing uniform row norms, the authors propose Aurora.

Formulation: Aurora formulates the update as a constrained optimization problem that maximizes the trace of the gradient dot product subject to two constraints:
1. Left semi-orthogonality: $U^\top U = \eta^2 I$ (preserving the polar factor structure).
2. Equal row norms: $\|U_{i,:}\|^2 = \eta^2 \frac{n}{m}$ for all rows $i$ .
Riemannian-Aurora: They derive a theoretical reference solver, Riemannian-Aurora, which treats these constraints as a joint manifold (the intersection of the Stiefel manifold and the equal-row-norm manifold). While theoretically precise, it is computationally prohibitive due to the need to solve large linear systems involving leverage constraints.
Practical Aurora: They propose an iterative, practical instantiation that alternates between:
1. Row Normalization: Rescaling rows to equalize norms.
2. Polar Factor Approximation: Applying a Newton-Schulz iteration to re-orthogonalize the matrix.
  This process is damped with a parameter $\beta$ to ensure convergence. The authors prove that under local regularity conditions, this alternating projection converges linearly to a solution satisfying both constraints.

Key Contributions

Analysis of Normalized Spectral Optimizers: The paper demonstrates that post-hoc unit row normalization degrades polar factor precision for tall matrices. It introduces U-NorMuon, a stateless variant that restricts normalization to tall matrices, proving that row uniformity and precise orthogonalization are concurrently desirable.
Study of Neuron Death: The authors provide a mechanistic explanation for neuron death in Muon-trained models, linking it to low leverage scores in tall matrix parameters. They verify this empirically in their own training runs and in open-weight models (e.g., Rnj-1, Laguna-XS.2), finding significant cohorts of dead neurons in dense MLP layers.
Aurora and Riemannian-Aurora: They derive Aurora as a method to enforce uniform row norms while strictly adhering to the Muon theoretical framework. They also provide Riemannian-Aurora as a high-cost reference solver for the dual-constraint problem.
Extensive Pre-training Verification: The authors validate Aurora across 340M and 1.1B parameter scales, demonstrating superior downstream performance compared to Muon, NorMuon, and U-NorMuon.

Results

Downstream Performance

Benchmark Scores: At the 1.1B scale, Aurora outperforms the Muon baseline significantly, achieving a 9.1 point higher score on MMLU. It also achieves the highest average score across a suite of reasoning and knowledge benchmarks (ARC, HellaSwag, etc.) at both 340M and 1.1B scales.
Speedrun: On the modded-nanoGPT speedrun (optimizer track), Aurora combined with existing methods (Contra-Muon, flooring) set a new state-of-the-art, reaching the target validation loss 25 steps faster than the NorMuon baseline.

Neuron Utilization

Leverage Scores: Models trained with Muon exhibit a bimodal distribution of leverage scores, with a significant mode near zero (dead neurons). In contrast, Aurora and U-NorMuon produce isotropic leverage score distributions, effectively eliminating dead neurons.
Open-Weight Verification: Analysis of open-weight models trained with Muon (Rnj-1 and Laguna-XS.2) confirms the presence of dead neurons in early dense layers (up to ~28% in the first layer of Laguna-XS.2), validating the generalizability of the problem.

Efficiency and Scaling

Throughput: In a distributed setting, Aurora's additional computational overhead is minimal because the extra matrix multiplications can be overlapped with gradient communication. On a single GPU, throughput decreases slightly, but the wall-clock overhead in distributed training is modest.
MLP Width Scaling: The performance gains of Aurora over Muon scale monotonically with the MLP expansion factor. As the MLP layers become wider, the advantage of Aurora increases, suggesting it is particularly effective for training very wide networks.

Significance and Claims

The paper claims that Aurora resolves a specific, empirically observed pathology in language model training—neuron death in tall matrix parameters—by deriving an optimizer that respects both the theoretical geometry of the Muon update and the practical necessity of uniform row norms.

The authors emphasize a bottom-up design philosophy: deriving algorithms based on observed training dynamics (neuron death, leverage score imbalances) rather than purely theoretical models. They argue that architecture-optimizer codesign, where optimizers are tailored to specific parameter classes (e.g., tall vs. wide matrices), is a promising direction that outperforms standard "one-size-fits-all" strategies.

Finally, the paper suggests that Aurora enables the effective training of very wide MLP layers, a capability that may be limited in current Muon-based training due to the accumulation of dead neurons. The work positions Aurora not just as a performance booster, but as a necessary correction to the geometric assumptions underlying current spectral optimization in large-scale language models.