Aurora: A Leverage-Aware Optimizer for Rectangular Matrices | Tilde

Loading...

1: Introduction

The Muon optimizer [1] first attracted attention due to its success in the nanoGPT speedrun competition [2], where it outperformed AdamW [3] in wall-clock time to convergence despite requiring more computation per step. Since then, substantial effort has gone into scaling Muon through improved distributed implementations [4], more accurate orthogonalization routines [5] [6], and methods to reduce compute and communication overhead [4]. The result is that Muon (including clipped and per-head versions) has become an increasingly popular choice for training frontier-scale models [7] [8] [18].

Recently, there has been a surge in Muon variants pushing the state-of-the-art on open research benchmarks [9] [10]. One example is NorMuon [9], which is currently SoTA on the modded-nanoGPT speedrun. NorMuon augments Muon with an additional step that scales each row by its inverse RMS norm, analogously to Adam's per-parameter normalization. Notably, this row normalization step can move the Muon update quite far from the true orthogonal gradient. The fact that NorMuon still succeeds suggests that there may be a gap in the Muon formulation that is being addressed by row normalization.

We study the effects of row normalization and find that Muon can result in neuron death in MLP layers, whereby some neurons receive persistently small updates early in training and fail to recover. We show that this failure mode can be avoided by redistributing mass equally across rows for updates to the up and gate projections in the MLP layers. Motivated by this observation, we propose Aurora, which uses this mechanism to prevent neuron death without sacrificing precision of the gradient orthogonalization. We show that Aurora achieves large gains over both Muon and NorMuon at 1B scale and outperforms the current SoTA on the modded-nanoGPT optimization track. Furthermore, we find that Aurora's performance gains scale with MLP width leading us to believe that Aurora is particularly effective for networks with large MLP expansion factors.

2: Tension Due to Row Normalization

The core algorithmic component in Muon is an iterative algorithm to compute the polar factor of a matrix. For $G \in \mathbb{R}^{m \times n}$ with thin Singular Value Decomposition (SVD) $G = U\Sigma V^\top$, the polar factor $\operatorname{polar}(G) = UV^\top$, is the closest semi-orthogonal matrix to $G$ in Frobenius norm. If $G$ is the gradient of a parameter $W$, then Muon forms $\operatorname{polar}(G) = UV^\top$ and updates $W \gets W - \eta\, UV^\top$ for a learning rate $\eta$.

The existence of matmul-only iterative algorithms for computing $\operatorname{polar}(G)$ [5] [6] [1] is largely what makes Muon feasible at scale. We focus on two such algorithms here: quintic iteration from modded-nanoGPT[21] and Polar Express with 8 steps (PE-8). [5] We give a more thorough overview of different algorithms for computing the polar factor in Appendix A.

We find that in our training settings (as described in Appendix B) Muon achieves monotonically better downstream loss under more precise orthogonalization; i.e. under iterative algorithms that more accurately approximate $\operatorname{polar}(G)$. We plot Muon runs using the least precise iteration we test (quintic) and the most precise (PE-8) in Figure 1 to illustrate this gap. This result confirms that polar factor precision is an important component in strong Muon implementations.

NorMuon augments Muon with an additional step that scales each row of the polar factor to have unit RMS norm.

We show that for tall matrices, this row normalization can significantly reduce polar factor precision, which is undesirable both theoretically and in light of our previous results (Figure 1).

Therefore, row normalization in NorMuon necessarily introduces a precision defect into the orthogonalization routine. We find that this polar precision defect can be quite large for matrices with non-uniform row norms. In particular, the magnitude of the introduced error can be much larger than the precision improvement of PE-8 over quintic iteration in the same setting (see Appendix A for a full comparison of different Newton-Schulz iterations).

Indeed, the performance gap between PE-8 and quintic iteration disappears when these algorithms are used with NorMuon in the same setting. However, both runs outperform our Muon+PE-8 baseline, suggesting that row normalization can be independently useful.

To mitigate the defect introduced by row normalization, we can simply normalize tall matrices to have row norms $\sqrt{n/m}$ instead of one. We call this variant U-NorMuon and find that it outperforms NorMuon in our 340M setting (Figure 4; setting described in Appendix B).

Now we turn our analysis to wide and square matrices, which also receive unit row norm updates under NorMuon. A left-orthogonal wide or square matrix necessarily has all its row norms equal to one, so row-norm uniformity is implied by orthogonality in this case.

Thus, we expect that row-normalization is unnecessary or perhaps even harmful under precise orthogonalization routines like PE-8. Iterative polar factor algorithms can fail to converge precisely early in training when raw gradients tend to be ill-conditioned, allowing for the accumulation of row norm statistics that quickly become stale. In this case, row normalization will make rows less uniform until the statistics are corrected. We find evidence of this occurring at 340M scale where row-normalizing only tall matrices outperforms NorMuon and U-NorMuon by a small but non-trivial margin (Figure 4).

We will show that Muon can allow a large subset of neurons in the MLP layers to effectively die, but that this pathology is mitigated by (U-)NorMuon. Building on ideas in the literature, we propose this as an explanation for the performance gap between Muon and NorMuon that we find in all our settings. We then derive Aurora which effectively row normalizes updates to tall parameters without sacrificing precision of the polar factor.

3: Normalization Prevents Neuron Death

We define a dead model component as a subset of model parameters receiving persistently small learning signal after the earliest phases of training. We identify dead model components with the following three criteria:

1. Low effective gradient norm. Model components which, for a particular batch, receive low effective gradient norm must correspond to relatively flat directions in the loss landscape. Components with persistently small gradient norms will thus receive little learning signal. There is a strong precedent in the attribution and pruning literatures for using gradient norms as a proxy for feature importance [11] [12] [13].

2. Low effective update norm. For networks trained with optimizers like Muon, small grad norms are not themselves sufficient for identifying dead network components. In particular, the orthogonalization step can drastically alter the gradient before it is applied, including amplifying small magnitude directions. We therefore require that effective update norm, which measures the actual parameter change, also be small.

3. Persistence over training. It's possible for parameters to receive small update norms at specific points in training but still contribute non-trivially to the network output. For example, some components may learn useful features early and receive smaller updates as those features stabilize. To distinguish these static but useful components from genuinely dead ones, we require low effective gradient and update norms to be persistent and predictable over training.

We will show that neuron death, as we've defined it, can and does occur in networks trained with Muon because tall matrices in MLP layers are allowed to receive updates with very non-uniform row-norms. In particular, we will show that the Muon update privileges network components that happened to get large updates early in training, allowing others to die. We find that U-NorMuon avoids this problem completely.

The Mechanism of Neuron Death

For a matrix $M \in \mathbb{R}^{m \times n}$ with thin SVD $M = U_r \Sigma V^\top$, define the leverage score of row $i$ as $\ell_i(M) = \|(U_r)_{i,\cdot}\|_2^2$. The leverage score of a row is directly proportional to the scale of its update so that rows with consistently low leverage scores in the update matrix receive a consistently small portion of the update mass. We can show that row leverage is mostly preserved under orthogonalization of reasonably well-conditioned matrices so that under Muon, rows with low leverage in the raw gradient tend to also have low leverage in the update.

Thus, the leverage score (and hence update) for a neuron is bounded by the squared norm of the corresponding row in the momentum buffer, normalized by the square of the smallest singular value. In particular, when $M$ is reasonably well-conditioned (so that $\sigma_n(M)$ is not too small) Muon preserves row leverage anisotropy. This leads to the following self-reinforcing feedback loop in networks trained with Muon:

1. Dead neurons receive near-zero gradient rows
2. Momentum accumulates into a buffer $M$ with highly anisotropic row norms.
3. The updates inherit this anisotropy, resulting in dead neurons not being updated.
4. Hence, dead neurons remain dead.

We find significant empirical evidence that this loop occurs in our standard 1B and 340M transformer settings.

Empirical Evidence of Neuron Death

We train vanilla transformer models with 340 million non-embedding parameters following the recipe given in Appendix B using both Muon and U-NorMuon. For each, we extract the momentum buffer every 50 steps and compute row-wise gradient norms and leverage scores. The U-NorMuon update has approximately uniform row norms (and hence leverage scores) by construction, but leverage scores tend to be very anisotropic for tall matrices under Muon. Our models use SwiGLU MLPs and we verify later that our results transfer to models that use ReLU$^2$ (see Appendix G).

Neuron death is visible in activation space. We first examine the MLP latent statistics directly. For each neuron, we track the mean and standard deviation of activations across tokens at each training step. Under Muon, the distribution of gate means diverges sharply: as training progresses we find that a band of neurons drifts toward zero mean with collapsing variance; we interpret this as activation space collapse. By contrast, gate means and variances remain consistently high under U-NorMuon.

Muon kills neurons early and permanently. We visualize leverage scores for the tall MLP projections of a middle layer over training (Figure 9). Under Muon, neurons are initially alive with uniformly high leverage, but a large fraction of neurons die during learning rate warmup and never recover. By step 500, more than one in four neurons are effectively dead, producing a sharply bimodal distribution of leverage scores; one mass of neurons receives near-zero updates, and the other receives disproportionately large ones.

This neuron death is self-reinforcing and leads to a fundamental collapse in learning dynamics (Figure 10). On a log-log plot, row gradient norms and leverage scores are tightly correlated, indicating that the momentum buffer tends to be in the well-conditioned domain where our bound from Claim 2 is tight.

To verify that neuron death persists throughout training, we split neurons into quartiles by their leverage score at step 100 (out of 10 thousand) and track each cohort as training progresses (Figure 11). The bottom quartile continues to decay, while the remaining cohorts quickly stabilize at high leverage under a "rich get richer" dynamic. If anisotropy were i.i.d. across rows, early leverage would not predict late leverage. This strong stratification confirms that, in our experiments, neuron death is a fixed point of learning rather than just training noise.

U-NorMuon prevents death, and the prevention propagates.

Under U-NorMuon, leverage scores are isotropic and approximately distributed according to a Gaussian throughout training with no dead-neuron tail (Figures 9–10). One might expect this to be a trivial consequence of the normalization step forcing uniformity post-hoc. Instead, we find that the momentum buffer itself converges to a regime with isotropic row norms.

The normalization intervenes most aggressively early in training when anisotropy is developing, and as the buffer stabilizes, the correction shrinks. More surprisingly, this benefit extends to parameters that U-NorMuon does not directly normalize. For example, in Figure 12 we plot column leverage scores for the down projection matrix, which is a wide matrix and thus receives no benefit from row normalization.

By contrast, column leverage becomes anisotropic under Muon. Dead rows in the up/gate projections starve the corresponding columns of the down projection of gradient signal. Under U-NorMuon, keeping the up/gate rows alive ensures isotropic gradient flow into the down projection, stabilizing its column leverage without any direct intervention.

Takeaway. U-NorMuon's benefit is not simply normalization. U-NorMuon breaks a pathological self-reinforcing collapse that wastes a quarter or more of the network's MLP parameters. However, it does so by forcefully overriding the polar factor with uniform row norms, sacrificing precision of the polar factor which is both empirically and theoretically desirable in the Muon framework. We will now present Aurora, which achieves both goals simultaneously. We show that Aurora has better downstream performance than both Muon and (U-)NorMuon at 1B scale and on the modded-nanoGPT optimizers track.

4: Augmenting the Muon Objective

Recall that the Muon update direction can be motivated as the solution to the following constrained optimization problem:

Where $G = \nabla_W \mathcal{L}$ is the gradient with respect to a matrix parameter $W$, and $\|U\|_2$ is the spectral norm (ignoring the learning rate scalar for simplicity). There is no constraint on the uniformity of the row norms of the update $U\in \mathbb{R}^{m \times n}$, which we found can be quite far from uniform for tall matrices (Sections 2 & 3). Motivated by our results in Section 3, we augment the Muon objective for tall matrices ($m \ge n$) with the additional requirement that all the update row norms be equal to $\sqrt{n/m}$:

The constant row norm constraint fixes the squared Frobenius norm to $\|U\|_F^2 = m \cdot (n/m) = n$. Because the squared Frobenius norm also equals the sum of the squared singular values and the spectral norm bound implies each $\sigma_j \le 1$, all $n$ singular values must exactly equal 1. This forces $U$ to be left semi-orthogonal. Therefore, the problem reduces to:

The solution is steepest descent under two constraints: a left semi-orthogonality constraint to make the operator RMS-to-RMS norm preserving, and a constant row norm constraint to enforce uniform updates. Solving this problem yields Aurora.

A Riemannian Solution

We first consider a Riemannian formulation of the constrained update-selection problem. We restrict the parameter update $U$ to lie on the joint Stiefel/equal-row-leverage manifold

Given a matrix gradient $G$, we seek the feasible update most aligned with $G$:

A Riemannian approach approximately solves this problem by projecting the Euclidean direction $G$ onto the tangent space of $\mathcal{M}$, taking a step along the projected direction, and retracting the result back toward $\mathcal{M}$ (Figure 13). We refer to this reference solver as Riemannian-Aurora; the full derivation is given in Appendix D.

Unlike post-hoc row normalization, which first computes a Muon update and then rescales its rows, Riemannian-Aurora treats orthogonality and equal row leverage as a single joint constraint. Both the initialization of $U_0$ and retraction are done approximately via alternately computing the polar factor and row normalizing the update in our implementation.

Some of the steps in Riemannian-Aurora are particularly expensive, making it infeasible for use on moderately-sized networks as written. Namely:

1. The tangent space projection involves solving an $m$-dimensional linear system involving $rI - P \odot P$, where $P = UU^\top$. For an $m \times n$ matrix, explicitly forming $P$ costs $O(m^2 n)$ memory/computation and storing it costs $O(m^2)$, which is prohibitive at moderate scales.
2. Solving for the tangent constraint with conjugate gradient requires repeated projections of the form $v \mapsto (P \odot P)v$, i.e. dense row interactions.

We present it mostly as a faithful solution to our dual-constraint problem and as a comparison point for our practical instantiation of Aurora. We include a small-scale verification of Riemannian-Aurora in Appendix E.

An Iterative Approximation: Aurora

Though the joint Stiefel/equal-row-norm projection has no simple closed form, it admits a natural iterative approximation which alternates between the two constraint sets, starting with row normalization, then using Newton-Schultz iteration to approximate the polar factor.

In practice, we use a damped iteration with an EMA of row normalization factors to improve convergence and prevent oscillation. This requires just one additional damping parameter $\beta$.

These steps alternate because the polar factor computation can disturb row-norm uniformity and vice versa. We show that, under mild conditions, this alternating projection converges to a point on or near the intersection of the Stiefel and row-oblique manifolds (proof in Appendix F). Notably, the final iterate is always on the Stiefel manifold as the last step is always orthogonalization. Aurora is roughly $K$ times more compute expensive than Muon but in practice, we can effectively overlap much of this overhead with gradient communication in distributed settings. We present an interactive simulation of vanilla (iterative) Aurora below.

5: Results

nanoGPT Speedrun

We evaluate Aurora on the modded-nanoGPT speedrun benchmark, comparing against NorMuon and the previous state-of-the-art run. Pure Aurora (2 damping iterations, $\beta = 0.5$) improves on NorMuon by 25 steps, reaching the 3.28 target loss in 3225 steps. When combined with Contra-Muon and update/weight flooring, Aurora achieves a new SoTA of 3175 steps.

Aurora 1.1B Pretraining

We train 1.1B-parameter transformers on ~100B tokens (see Appendix B for full details) and compare Aurora against Muon and NorMuon, each using PE-8. Aurora achieves the lowest final loss of all methods, reaching a smoothed loss of 2.26 at step 24k, which is a clear improvement over Muon (2.31) and NorMuon (2.33).

Downstream evaluations. Aurora's loss improvement translates to consistent gains on standard benchmarks. On Hellaswag, Aurora achieves 67.6% - a 2.5 point improvement over both Muon and NorMuon. Gains are similarly consistent on ARC-Challenge and Winogrande. Strikingly, Aurora improves MMLU scores by 10 points over Muon. We hypothesize that since MLPs are predominantly responsible for memorization, Aurora's gains are most visible on memorization-intensive benchmarks like MMLU.

Remarkably, Aurora-1.1B matches or approaches Qwen3-1.7B (36T tokens) on Hellaswag and Winogrande despite training on two orders of magnitude fewer tokens, having 400M fewer parameters, and an uncurated data mixture. The scatter plot reveals the striking token efficiency: Aurora sits far to the left of frontier small language models on the log-token axis while achieving competitive scores on benchmarks that primarily test language understanding rather than domain-specific knowledge. Note that our 1.1B model is trained on ~100B tokens of the Nemotron CC dataset, which contains no explicit math or science data, explaining the lower MMLU scores compared to models trained on curated multi-domain corpora.

Notably, despite having 25% fewer parameters, 2 orders of magnitude fewer training tokens, and using fully open-source internet-only data, we match Qwen3-1.7B on several benchmarks.

Update Quality

We verify that Aurora's damped iteration actually solves the joint Stiefel–oblique problem for real momentum buffers, and compare it against the Riemannian reference solver from Section 4.

Gradient alignment vs. leverage uniformity

Figure 18 shows the tradeoff between gradient alignment ($\langle G, U \rangle / \|G\|$) and row-norm coefficient of variation (CV) as a function of iteration count, computed on a real momentum buffer from early in training. Standard Muon achieves high alignment but poor row-norm uniformity. Aurora and Riemannian-Aurora converge to low CV with high alignment, but from opposite directions. Whereas Aurora starts orthogonal and corrects row norms, Riemannian-Aurora starts from the tangent projection and retracts onto the manifold. Riemannian-Aurora achieves 0.6% higher gradient alignment at convergence, suggesting that the practical Aurora iteration sacrifices very little relative to the reference solver.

Leverage score distributions

Figure 19 compares leverage score distributions across optimizers. Muon produces a heavy tail of dead rows spanning several orders of magnitude. Aurora with 1 damping iteration (equivalent to Muon-Eq) improves on Muon but retains visible spread. Aurora with 2 iterations produces by far the most isotropic distribution, tightly clustered around the uniform value $\frac{n}{m}$.

Weight spectra

We examine the singular value spectra of the trained weight matrices. Under NorMuon, the down projection spectra are inflated by ~100x ($\sigma_{max} \approx 500$ to $1500$ vs Aurora's $11$ to $21$), while gate projection and up projection spectra remain comparable. This is consistent with the weight matrix absorbing scale information that the row-normalized update discards.

6: Inflection

Four things excite us about this work:

1. Strong empirical gains. We stress that Aurora 1.1B was trained to be a research artifact: we did not do any tricks with data, architecture, or optimizer beyond simply using Aurora, to isolate its effect. Despite this, it achieved 100x data efficiency on several key benchmarks, and we see this as very promising.

2. Architecture-Optimizer Codesign. Aurora is an optimizer specifically for up and gate projections, which we show corrects a fundamental pathology in model training dynamics under Muon. We believe this is a nice example of optimizer-architecture co-design at a level of granularity that is very uncommon in the literature. Crucially, Aurora's advantage over Muon grows monotonically with MLP expansion factor (Figure 21), suggesting that wider MLPs, which increase the row-to-column ratio of the up and gate projections, amplify exactly the pathology Aurora corrects.

3. Compression / Dense Solutions. Aurora's row uniformity constraint is complementary to Muon's motivating constraint. We propose thinking about the row-uniformity constraint as encouraging compression. Under Aurora, networks tend to more effectively utilize the parameters in their MLPs; that is, models tend to denser solutions. The empirical success of Aurora indicates that dense solutions to language modelling can be feasible.

4. Bottom-Up Optimizer Design. Aurora was derived by mechanistically studying a specific pathology in Muon and designing the minimal objective that corrects it. This contrasts with the top-down approach that dominates optimizer research, where one posits a mathematical formulation and hopes it yields empirical gains; an approach that offers little guidance on which of many sound framings to choose, and no explanation for why one works better than another (Figure 22). We strongly believe the field would benefit from more bottom-up research as a whole.

Understanding how Muon fails on tall matrices led us to a better objective, a better optimizer, and "denser" models. We are interested in faster and more precise methods of solving the dual-constraint steepest descent objective, and extending our results to large networks with wide MLPs.

@article{dewulf2026aurora,
  title   = {Aurora: A Leverage-Aware Optimizer for Rectangular Matrices},
  author  = {Dewulf, Alec and Pai, Dhruv and Yang, Li and Zhang, Ashley
             and Keigwin, Ben},
  year    = {2026},
  url     = {https://tilderesearch.com/blog/aurora}
}

Appendix

A. Comparing Newton-Schulz Iterations

We present a cost analysis of different polynomials for iterative orthogonalization. As discussed in Sections 2 & 5, we find that more precise polynomials are always helpful for Muon and Aurora, which justifies the additional computation in many settings.

CANS-12 and Polar Express-8 are both much more precise than quintic iteration, which has a persistent spectral error of 0.3, even in the limit. Both CANS-12 and PE-8 achieve nearly perfect accuracy for reasonably well-conditioned gradients.

B. Training Setting

We report exact configurations and training settings for our 340M and 1.1B runs to maximize the reproducibility of our results. We arrived at these settings via hyperparameter sweeps; we spent roughly the same compute sweeping each optimizer and setting.

We used an internal tokenizer with 128k vocab size for all our experiments, which was released in the model repository. For reproducibility, we trained on fully open-source internet data from NVIDIA Nemotron CC v2 [19].

Transformer 340M:

Transformer 1.1B:

C. Literature on Parameter Death

In Transformer FFNs, prior work has shown that intermediate hidden dimensions are not uniformly utilized: FFN activations are often highly sparse, and some intermediate neurons can remain inactive over large corpora, often referred to as dead or dormant neurons [14] [15] [16]. These neurons correspond to the hidden FFN dimensions between the up/gate projections and the down projection. When a hidden neuron remains inactive or near-zero for many tokens/steps across training, the gradients to its associated weights are greatly reduced. As a result, only a subset of the FFN intermediate dimensions and their corresponding weight matrices effectively participate in computation, reducing the model's effective capacity and parameter utilization [17].

D. Solving for the Riemannian Gradient

We want to solve

Let $r := \frac{n}{m}$. The feasible set is the equal-row-norm Stiefel manifold

Tangent space. Let $Z \in \mathbb{R}^{m \times n}$ be a tangent vector at $U \in \mathcal{M}$. The Stiefel constraint requires $(U + \epsilon Z)^\top (U + \epsilon Z) = I_n + O(\epsilon^2)$ for small $\epsilon > 0$. Keeping only first-order terms gives

The row-norm constraint requires $\|U_{i:} + \epsilon Z_{i:}\|_2^2 = r + O(\epsilon^2)$ for all $i$. Expanding to first order gives $\langle U_{i:}, Z_{i:} \rangle = 0$ for all $i$, or equivalently $\operatorname{diag}(ZU^\top) = 0$. Therefore,

Riemannian gradient. We compute the Riemannian gradient by projecting $G$ onto the tangent space:

The projected vector has the form $Z = G - US - DU$, where $S = S^\top \in \mathbb{R}^{n \times n}$ and $D = \operatorname{diag}(\lambda) \in \mathbb{R}^{m \times m}$. The $US$ term enforces the Stiefel tangent constraint, while the $DU$ term enforces the row-norm tangent constraint. Define

We impose the Stiefel tangent constraint $U^\top Z + Z^\top U = 0$. Substituting $Z = G - US - DU$, we get

Using $U^\top U = I_n$, $S = S^\top$, and $D = D^\top$, this becomes

Therefore,

Now impose the row-norm tangent constraint $\operatorname{diag}(ZU^\top) = 0$. Using $Z = G - US - DU$, this gives

Since $\operatorname{diag}(UU^\top) = r\mathbf{1}$ and $D = \operatorname{diag}(\lambda)$, we have $\operatorname{diag}(DUU^\top) = r\lambda$. Substituting $S = B - U^\top DU$ into $\operatorname{diag}(USU^\top)$ gives

With $P = UU^\top$, this becomes $\operatorname{diag}(USU^\top) = \operatorname{diag}(UBU^\top) - \operatorname{diag}(PDP)$. Since $D = \operatorname{diag}(\lambda)$,

Therefore $\operatorname{diag}(USU^\top) = \operatorname{diag}(UBU^\top) - (P \circ P)\lambda$. Plugging this into the row constraint gives

Rearranging,

After solving this linear system for $\lambda$, define $D = \operatorname{diag}(\lambda)$, $S = B - U^\top DU$, and the tangent projection is $Z = G - US - DU$.

E. CIFAR-10 Results

We train a 2-layer MLP with ReLU for 3 epochs on CIFAR-10 as a toy verification of Riemannian-Aurora. We use $K=20$ conjugate gradient iterations, Nesterov Momentum, lr=4e-3, 2 retraction steps, PE-8, a Riemannian learning rate of $\eta = 0.1$ and an AdamW learning rate of 4e-3 with betas $(0.9, 0.95)$. We found these parameters to be approximately optimal, after a brief sweep.

We find that Riemannian-Aurora outperforms the AdamW baseline by a significant amount, and performs very similarly to Muon in terms of train accuracy, but has slightly higher test accuracy. Unlike Muon, it leads to much more uniform row leverage scores.

F. Alternating Iteration Converges

Let $m > n$, $r = n/m$, and define

Let $\mathcal{M} = \mathcal{S} \cap \mathcal{O}_r$. The alternating projection iteration is

where $\Pi_{\mathcal{O}_r}$ is row-wise rescaling and $\Pi_{\mathcal{S}}$ is the polar factor.

Lemma F.1: Nonemptiness

$\mathcal{M} \neq \emptyset$.

Proof. By the Schur–Horn theorem, there exists a rank-$n$ orthogonal projector $P \in \mathbb{R}^{m \times m}$ with constant diagonal $P_{ii} = n/m$. For example, the diagonal vector $(n/m, \ldots, n/m)$ is majorized by the eigenvalue vector $(1, \ldots, 1, 0, \ldots, 0)$ with $n$ ones. Write $P = UU^\top$ with $U^\top U = I_n$. Then $\|U_{i,:}\|_2^2 = (UU^\top)_{ii} = P_{ii} = n/m$, so $U \in \mathcal{S} \cap \mathcal{O}_r$. Hence $\mathcal{M} \neq \emptyset$. $\square$

Regularity Condition

For $U \in \mathcal{M}$, define $P = UU^\top$ and $L_U = rI_m - P \circ P$, where $\circ$ denotes the Hadamard product. Since $P^2 = P$ and $P_{ii} = r$,

Therefore $L_U$ is a graph Laplacian:

Call $U$ regular if $\ker L_U = \operatorname{span}\{\mathbf{1}\}$.

Lemma F.2: Clean Intersection

If $U \in \mathcal{M}$ is regular, then $\mathcal{S}$ and $\mathcal{O}_r$ intersect cleanly at $U$.

Proof. The normal spaces are $N_U \mathcal{S} = \{UA : A = A^\top\}$ and $N_U \mathcal{O}_r = \{DU : D = \operatorname{diag}(\lambda)\}$. Suppose a vector lies in both normal spaces, so $UA = DU$. Since $P = UU^\top$ is the projector onto $\operatorname{col}(U)$, this implies $(I - P)DU = 0$. But

Using $\|U_{i,:}\|_2^2 = r$, we get $\|DU\|_F^2 = r\|\lambda\|_2^2$ and $\|PDU\|_F^2 = \lambda^\top (P \circ P)\lambda$, so

By regularity, this forces $\lambda = \alpha \mathbf{1}$, hence $D = \alpha I_m$, $A = \alpha I_n$, and $DU = UA = \alpha U$. Therefore $N_U \mathcal{S} \cap N_U \mathcal{O}_r = \operatorname{span}\{U\}$. This one-dimensional common normal direction is unavoidable since both constraint sets imply the same Frobenius norm constraint $\|U\|_F^2 = n$. Thus the combined constraint map has maximal rank modulo this single redundancy. By the constant-rank theorem, $\mathcal{M}$ is locally a smooth manifold and $T_U \mathcal{M} = T_U \mathcal{S} \cap T_U \mathcal{O}_r$, so the intersection is clean. $\square$

Theorem F.3: Local Linear Convergence

Let $U_\star \in \mathcal{M}$ be regular. Then for every initialization $U_0$ sufficiently close to $U_\star$, the alternating projection sequence $U_{k+1/2} = \Pi_{\mathcal{O}_r}(U_k)$, $U_{k+1} = \Pi_{\mathcal{S}}(U_{k+1/2})$ is well-defined and converges linearly to some $\widehat{U} \in \mathcal{M}$.

Proof. Both $\mathcal{S}$ and $\mathcal{O}_r$ are $C^\infty$ embedded submanifolds. Near $U_\star$, projection onto $\mathcal{S}$ is uniquely given by the polar factor, since nearby matrices have full column rank. Projection onto $\mathcal{O}_r$ is uniquely given by row-wise normalization, since all rows of $U_\star$ have norm $\sqrt{r} > 0$. By Lemma F.2, $\mathcal{S}$ and $\mathcal{O}_r$ intersect cleanly at $U_\star$. Clean intersection of smooth manifolds implies local linear convergence of alternating projections. Therefore the iterates converge linearly to a point $\widehat{U} \in \mathcal{S} \cap \mathcal{O}_r = \mathcal{M}$. $\square$

G. Neuron Death in ReLU² MLPs

The neuron death analysis in Section 3 focuses on SwiGLU MLPs, but, as we show below, this pathology is not specific to the SiLU gating function. For any activation $\phi$ with $\phi(0) = 0$ and $\phi'(0) \approx 0$, the same vanishing-gradient feedback loop occurs: neurons with small pre-activations receive small gradients, accumulate small momentum, and produce small leverage scores under Muon. Common activation functions including ReLU² (defined by $\phi(x) = \max(0,x)^2$), GELU, and SiLU all share this property.

We verify this by repeating our leverage score analysis on a 400M transformer trained with ReLU² MLPs (same setting as Appendix B, with SwiGLU replaced by ReLU²). The results closely mirror those of the SwiGLU experiments.

Neuron death grid (ReLU²). Under Muon, a large fraction of neurons die early in training. Under U-NorMuon, row leverage scores remain isotropic throughout training.

Leverage persistence (ReLU²). We sort neurons into cohorts by their leverage score at step 500 and track each cohort over training. The pattern matches our SwiGLU results: under Muon, the lowest-leverage cohorts collapse and never recover, while under U-NorMuon, all cohorts converge toward uniform leverage.

H. A Note on MoE Architectures

Aurora's benefits are most pronounced for tall matrices, where the ratio $m/n$ is large. This is the typical regime for up and gate projections in dense transformers, which commonly use MLP expansion factors of 4× or more. In MoE architectures, capacity is distributed across many smaller experts, each with a proportionally smaller hidden dimension and a lower $m/n$ ratio. We therefore expect the neuron death pathology to be less severe in MoE models of moderate size, though not entirely absent for experts that remain meaningfully tall.

I. Code Release

Our implementation of Aurora is open-sourced at github.com/tilde-research/aurora-release. The repository includes both the practical damped-iteration Aurora optimizer and the Riemannian-Aurora reference solver. Our modded-nanoGPT speedrun submission (Aurora + Contra-Muon + u/w-floor, 3175 steps) is available as a pull request at KellerJordan/modded-nanogpt#284. We also release the Aurora 1.1B pretrained model at tilde-research/aurora-1.1B.