AuON: Alternative Unit-Norm Optimizer
- AuON is a linear-time optimizer that normalizes gradient updates using Frobenius normalization and cosh-RMS scaling to control the spectral norm.
- It preserves the direction and internal correlation of momentum updates while avoiding expensive SVD or iterative Newton–Schulz computations.
- AuON enhances stability and allows larger learning rates compared to SGD and AdamW by effectively reconditioning updates in a spectral-norm trust region.
AuON, short for Alternative Unit-norm momentum updates by Normalized nonlinear scaling, is a linear-time optimizer for deep learning that transforms a gradient or momentum update by first normalizing it, then measuring its tail-heaviness with a hyperbolic-cosine RMS statistic, and finally shrinking the entire update by a single global factor so that its spectral norm is strictly bounded. In the formulation presented by Ryuzan, AuON is intended as a linear-time alternative to semi-orthogonal momentum updates such as Muon: it preserves the direction and internal structure of the update, suppresses spiky updates through tail-aware scaling, and avoids explicit semi-orthogonalization and the overhead of Newton–Schulz iterations (Maity, 29 Sep 2025).
1. Optimization context and motivating problem
AuON is positioned within the broader line of orthogonality-inspired optimizers. The motivating observation is that gradients and momentum updates in deep networks are often ill-conditioned: a small number of directions can dominate, while others remain weak. In this setting, strict orthogonalization replaces an update matrix with its polar factor from the SVD , thereby preserving singular vectors while flattening singular values to $1$. This yields an update with spectral norm on the active subspace and promotes more isotropic exploration (Maity, 29 Sep 2025).
The paper identifies two limitations of strict orthogonalization. First, SVD- or QR-based methods incur cost for square-like matrices, which is prohibitive in modern deep networks. Second, orthogonal-SGDM-style approaches can underperform well-tuned SGD with momentum because they orthogonalize gradients first and apply momentum afterward, thereby destroying the variance-reducing property of momentum, removing magnitude information in , and over-constraining updates to isometries (Maity, 29 Sep 2025).
Muon improves on this by applying momentum before semi-orthogonalization and approximating the polar factor by Newton–Schulz iterations, reducing the dominant cost to per step. AuON is introduced to remove that remaining quadratic bottleneck. Its stated objective is to retain several benefits associated with semi-orthogonal updates—spectral-norm control, better conditioning, and stability with larger learning rates—while using only elementwise operations and simple reductions, hence linear time in the number of parameters (Maity, 29 Sep 2025).
| Method | Defining operation | Reported cost |
|---|---|---|
| Strict orthogonalization | SVD/QR-based orthogonalization | |
| Muon | Newton–Schulz iterations on momentum | 0 per step |
| AuON | Frobenius normalization plus cosh-RMS scaling | 1 |
| Hybrid-AuON | One Newton–Schulz iteration plus AuON scaling | Between AuON and Muon |
This comparison suggests that AuON is less an orthogonalizer than a trust-region reconditioner for momentum updates: it does not construct a semi-orthogonal matrix, but instead seeks to preserve correlation structure while shrinking updates into a bounded spectral region.
2. Core transformation and update semantics
For an update tensor or matrix 2 with 3 elements, AuON defines a three-stage transformation. In the notation used in the paper,
4
The resulting 5 is then used in the parameter update
6
An equivalent presentation computes
7
and then scales by 8. The nonlinear part is therefore confined to the statistic used to determine a global scaling factor; the final update itself is a uniformly scaled version of the normalized input (Maity, 29 Sep 2025).
The paper emphasizes four consequences of this construction. First, AuON is scale-invariant after the initial Frobenius normalization: multiplying the raw update by a positive scalar does not materially change the transformed direction. Second, it is tail-aware because 9 is quadratic near zero and exponential in the tails; large coordinates contribute disproportionately to the scaling statistic and therefore induce stronger shrinkage. Third, it preserves direction in the precise sense that all coordinates are scaled by the same factor, leaving their relative ratios unchanged. Fourth, it preserves structural alignment because it does not decorrelate rows or columns the way full orthogonalization does (Maity, 29 Sep 2025).
In practical use, AuON is generally applied not to the raw gradient but to a momentum-smoothed update:
0
The transformed update 1 is then computed from 2. This ordering mirrors the motivation taken from Muon and from the critique of methods that orthogonalize before momentum (Maity, 29 Sep 2025).
A common misconception is that AuON applies 3 directly to the update sent to parameters. The paper states the opposite: 4 is used only to compute the scalar shrinkage factor. Another misconception is that “Alternative Unit-norm” means exact unit 5 norm or exact unit RMS. The paper explicitly notes that AuON does not force the step to exactly unit L2 or unit RMS; rather, it guarantees a bound on the spectral norm and makes the step smaller when the cosh-RMS statistic is larger (Maity, 29 Sep 2025).
3. Spectral-norm trust region and semi-orthogonal interpretation
The central formal claim is that AuON places each transformed update inside a strict spectral-norm trust region. For
6
the paper derives
7
Using 8 and 9, it follows that 0, so every AuON step lies in a strictly contracted spectral ball 1 (Maity, 29 Sep 2025).
This trust-region interpretation is strengthened by the paper’s analysis of tail sensitivity. If one coordinate of the normalized update satisfies 2, then
3
A single large spike therefore induces exponential contraction of the allowable spectral norm. This is the formal mechanism behind the claim that AuON damps heavy-tailed or harmful updates (Maity, 29 Sep 2025).
The paper also analyzes correlation energy and a form of “near semi-orthogonality.” Defining
4
one has
5
The isotropy residual obeys
6
Thus, both total spectral energy and deviation from isotropy are shrunk by 7. The paper is explicit that this is not true orthogonalization: AuON does not drive off-diagonal correlations to zero. Instead, it uniformly contracts correlation energy and anisotropy, which is presented as an approximation to some semi-orthogonal properties in a single cheap pass (Maity, 29 Sep 2025).
Under standard assumptions—8-smoothness, alignment of the AuON update with the true gradient, and bounded variance—the appendix proves an average-gradient-norm convergence rate of 9 when $1$0. The deterministic trust-region bound yields $1$1, which the paper states tightens the variance term relative to vanilla SGD (Maity, 29 Sep 2025).
4. Hybrid-AuON, implementation, and computational profile
Hybrid-AuON augments the base method with a single Newton–Schulz-like step before the AuON scaling. For a normalized update matrix $1$2, the paper gives
$1$3
with example constants $1$4, $1$5, and $1$6 in the code. This step is described as reducing correlations between columns or rows and moving $1$7 somewhat closer to semi-orthogonal, after which the standard AuON normalization is applied (Maity, 29 Sep 2025).
The computational distinction is straightforward. AuON uses norms, $1$8, RMS computation, and scalar divisions, all of which are reported as $1$9 per tensor. Hybrid-AuON adds one matrix multiplication 0 and one 1, giving cost on the order of 2 for square matrices, but still less than Muon because it performs only one Newton–Schulz iteration rather than several (Maity, 29 Sep 2025).
Implementation details in the reference PyTorch code are concrete. The AuON scaling function zeropower_via_cosh_rms(G, steps, newton_s1) casts to bfloat16, normalizes by L2 norm, optionally performs the Newton–Schulz step for Hybrid-AuON, and then applies the cosh-RMS scaling. The optimizer is integrated as a Muon-like class with a standard momentum buffer. Numerical stability is handled by small constants such as 3 and 4, and the paper reports that 5 is usually safe because it is applied to normalized inputs (Maity, 29 Sep 2025).
The reported hyperparameter behavior is also distinctive. AuON typically uses larger learning rates than AdamW in the provided experiments—for example, 6 versus 7 in the nanoGPT experiment—and momentum values around 8. This suggests that the trust-region effect is intended not merely as a stabilizer but as an enabler of larger nominal step sizes (Maity, 29 Sep 2025).
5. Empirical behavior across language and vision benchmarks
The paper reports experiments on language modeling with nanoGPT on SmolLM-Corpus, an MNIST MLP, and CIFAR-10 CNNs. In the nanoGPT-style transformer setting—hidden size 9, 0 layers, 1 heads, FFN dimension 2, RoPE, RMSNorm, SwiGLU, FlashAttention-2, trained for 3 steps with global batch size 4 on a 5k-token SmolLM-Corpus subset—the reported results are:
- AdamW: loss 6, accuracy 7, perplexity 8
- AuON: loss 9, accuracy 0, perplexity 1
- Hybrid-AuON: loss 2, accuracy 3, perplexity 4
- Muon: loss 5, accuracy 6, perplexity 7
The paper’s interpretation is that all AuON variants and Muon outperform AdamW under the given tuning, with Muon best and Hybrid-AuON close behind (Maity, 29 Sep 2025).
On MNIST with an MLP, the reported final losses are 8 for SGD, 9 for AdamW, and 0 for AuON, with runtimes of 1 s, 2 s, and 3 s respectively. AuON is therefore described as much better than SGD but worse than tuned AdamW in that simple setting, with modest overhead of about 4 versus SGD (Maity, 29 Sep 2025).
On a small CIFAR-10 CNN experiment, the reported results are:
- SGD: final loss 5, test accuracy 6, time 7 s
- AdamW: final loss 8, test accuracy 9, time 0 s
- AuON: final loss 1, test accuracy 2, time 3 s
In this case, AuON achieves the best reported accuracy and lowest loss while having runtime essentially identical to SGD and AdamW. In a larger reduced-scale CIFAR-10 setup with approximately 4M parameters, 5 epochs, batch size 6, AdamW base learning rate 7, AuON learning rate 8, weight decay 9, and momentum 00, the paper reports 01 test accuracy for AdamW and 02 for AuON, a gain of 03 (Maity, 29 Sep 2025).
The overall empirical picture given by the paper is therefore specific rather than universal: AuON consistently improves over SGD, often matches or slightly surpasses AdamW, Hybrid-AuON is somewhat better than AuON and closer to Muon, and Muon remains best on some metrics but at higher computational cost.
6. Practical interpretation, limitations, and research position
The paper recommends AuON when better conditioning and stability than SGD are desired, when a drop-in alternative to AdamW with similar runtime is attractive, and when larger learning rates are useful but Muon’s Newton–Schulz overhead is undesirable. Hybrid-AuON is recommended when one can afford a modest increase in compute to obtain partial decorrelation and performance closer to Muon. Muon remains preferable when strong semi-orthogonality is required and 04 cost is acceptable (Maity, 29 Sep 2025).
Several limitations are stated explicitly. AuON does not explicitly decorrelate directions, so applications requiring stronger semi-orthogonal behavior may favor Muon or Hybrid-AuON. On very large transformer models, the paper reports potential instabilities such as exploding attention logits, similar to issues observed for other orthogonalization-based optimizers. Large-scale evaluations are left to future work, and techniques such as QK-clipping are mentioned as possible mitigations (Maity, 29 Sep 2025).
In the broader literature position given by the paper, AuON belongs to a family that includes SVD-based orthogonal-SGDM, Muon and its variants, and trust-region interpretations of orthogonalized gradients. It is also connected to non-Euclidean trust-region methods, because orthogonalized momentum can be viewed as solving a trust-region problem in the spectral norm, whereas AuON attains a related effect by directly enforcing 05 through cosh-RMS scaling. The cosh-based scaling is further related to cosh gradient systems and tail-sensitive metrics, which the paper invokes as a way of reconditioning updates to be more robust to heavy tails (Maity, 29 Sep 2025).
Future directions proposed in the paper include scaling AuON and Hybrid-AuON to larger architectures and longer runs, combining them with QK-clipping in transformers, designing more sophisticated correlation-reducing steps compatible with linear-time cost, and integrating AuON with adaptive stepsize schemes in the style of AdaGrad-meets-Muon (Maity, 29 Sep 2025).
Taken together, these elements define AuON as a spectral-trust-region optimizer rather than a literal orthogonalizer: it preserves update direction and correlation structure, contracts spectral energy and anisotropy through a tail-aware global normalization, and aims to capture part of the optimization behavior associated with semi-orthogonal momentum updates at 06 cost.