Variance-Adaptive Muon: Pre-Orthogonalization Variance Modulation for Efficient Language Model Pretraining

AI summary

Overview

• Research area: Optimization algorithms for large language model pretraining (machine learning / deep learning systems).
• Technical level: Advanced. The paper assumes familiarity with Adam/AdamW, momentum EMAs, sign-based normalization, the orthogonal polar factor, singular value decomposition, and Newton–Schulz iteration.
• Scope in one sentence: The paper proposes two variants of the Muon optimizer, Muon-NSR and Muon-VS, that inject Adam-style gradient-variance information into Muon's momentum update before orthogonalization, and evaluates them on Llama-style and GPT-2 pretraining from 125M to 1.2B parameters.

What This Paper Is About

Muon is a strong pretraining optimizer that orthogonalizes momentum updates — a matrix version of sign-based normalization — but unlike Adam it carries no running estimate of gradient variance, so it cannot explicitly damp update directions that are dominated by noise. This paper asks whether Adam-style variance adaptivity can be grafted onto Muon's matrix-valued sign updates while keeping Muon's spectral normalization intact, and answers with two methods, Muon-NSR and Muon-VS, that modulate the update proxy with an online centered variance estimate before the Newton–Schulz step.

Key Contributions

1. Variance-adaptive Muon updates. Muon-NSR attenuates high-uncertainty update entries via a noise-to-signal ratio (NSR) gate, while Muon-VS applies variance scaling (VS). Both modulate the update proxy before Newton–Schulz orthogonalization, preserving Muon's spectral update structure and requiring only one additional variance buffer.
2. Centered-variance formulation instead of uncentered second moments. The variance surrogate is driven by the squared deviation between the current gradient and the previous mean estimate, rather than raw second moments or post-orthogonalization statistics.
3. Matched pretraining evaluation. Both variants are evaluated across Llama-style (Suite A and Suite B) and GPT-2 pretraining regimes spanning scales from 125M to 1.2B parameters, under fixed pretraining budgets and matched recipes, against AdamW and Muon-family baselines.
4. Ablations on the variance coefficient and on modulation ordering. The paper sweeps the Muon-NSR coefficient γ on a logarithmic grid and tests whether NSR gating should be applied before or after orthogonalization.

Main Findings

• Llama-1.2B step-to-target speedup: Muon-VS achieves a 1.33× step-to-target speedup over a well-tuned Muon baseline, using Muon's final validation loss as the target, and outperforms NorMuon at that scale.
• Consistent gains across Suite A scales: Under the tuned Suite A protocol, Muon-NSR and Muon-VS consistently improve over Muon at Llama-130M, Llama-300M, and Llama-1.2B; at the 130M and 300M scales both variants reach lower final validation loss than AdaMuon and NorMuon.
• Suite B speedups: Using the final Muon validation loss as the target, Muon-NSR achieves step-to-target speedups of 1.09× on Llama-210M and 1.11× on Llama-720M.
• GPT-2 speedups: Relative to the Muon baseline, Muon-NSR achieves time-to-target speedups of 1.24× and 1.15× on training-loss curves, and step-to-target speedups of 1.26× and 1.22× on validation-loss curves, for GPT-2 Small and GPT-2 Medium respectively. It recorded the lowest final validation loss among the evaluated optimizers at both scales.
• Muon-VS also beats baselines on GPT-2: It attains lower final validation loss than AdamW, Muon, AdaMuon, and NorMuon at both GPT-2 scales.
• Next-token accuracy tracks loss: Muon-NSR and Muon-VS rank among the leading methods on validation-set next-token top-1 accuracy under Suite B, suggesting the gains are not specific to the loss metric.
• γ sweep sensitivity (Table 1, final validation loss / reduction vs. AdaMuon): On GPT-2 Small, AdaMuon baseline 2.9572, Muon-NSR γ=1 → 2.9531 (+0.0041), γ=10 → 2.9508 (+0.0064), γ=100 → 2.9536 (+0.0036), γ=1000 → 2.9545 (+0.0027), Muon-VS → 2.9543 (+0.0029). On Llama-130M, AdaMuon baseline 3.2126, γ=1 → 3.2104 (+0.0022), γ=10 → 3.2067 (+0.0059), γ=100 → 3.2072 (+0.0054), γ=1000 → 3.2067 (+0.0059), γ=10000 → 3.2077 (+0.0049), Muon-VS → 3.2066 (+0.0060).
• γ=1 is suboptimal: The unit-coefficient case γ=1, the most direct Adam–Muon combination under Adam's equal-beta variance decomposition, improves on AdaMuon in both settings but remains suboptimal; relaxing the coefficient yields lower validation loss.
• The method is robust to γ: All evaluated logarithmic-grid values of γ improve over AdaMuon on GPT-2 Small and Llama-130M, with mild variation, suggesting a coarse log-scale choice is sufficient rather than fine-grained tuning.
• Ordering matters: Applying NSR modulation before orthogonalization yields a smoother optimization trajectory and lower validation loss than applying it afterward on Llama-130M.

Methodology in Plain English

The authors start from a known reinterpretation of Adam: when the first-moment and second-moment exponential moving averages share the same decay rate, Adam's per-coordinate update splits into a sign term and a damping factor of the form 1/sqrt(1 + NSR²), where NSR is the estimated noise-to-signal ratio. Muon already generalizes the sign part to matrices by approximating the orthogonal polar factor of a momentum matrix with Newton–Schulz iterations, but it has no variance estimate, so the damping term is missing.

To fill the gap, they maintain two coupled exponential moving statistics per parameter matrix. One is the momentum (mean estimate), parameterized as an Adam-style EMA and then rescaled by 1/(1−β) so that it reproduces Muon's unnormalized momentum recursion. The other is a centered variance surrogate, updated from the element-wise squared difference between the current gradient and the previous mean estimate.

Muon-NSR divides the Nesterov-style update proxy by sqrt(proxy² + γ·variance) + ε, which acts as a soft gate that damps entries whose estimated uncertainty is large relative to their signal. Muon-VS divides by sqrt(variance) + ε instead, dropping the proxy-squared term and the γ coefficient; the authors describe it as a large-γ limiting case of Muon-NSR up to a scalar absorbed by the subsequent spectral normalization. In both cases the modulated proxy then goes through the same Newton–Schulz orthogonalization and the same decoupled weight-decay update as Muon, with the same dimension-dependent scaling rule as the corresponding Muon baseline.

The experiments reuse Muon's tuned hyperparameters rather than re-tuning from scratch. Suite A follows a published Llama pretraining optimizer protocol (Wen et al., 2026) at a 1× Chinchilla token budget on DCLM, validated on C4-EN with the Llama-2 tokenizer; Suite B follows a published optimizer benchmark protocol (Semenov et al., 2025) on the FineWeb sample-100BT subset; the GPT-2 experiments use the AdaMuon implementation built on nanoGPT, pretraining GPT-2 Small and Medium on OpenWebText. All runs were done in PyTorch on a single node with 8× NVIDIA RTX 5090 GPUs using Distributed Data Parallel.

Speedups are defined as the ratio of the baseline's first step (or wall-clock time) reaching a target loss to the method's, where the target is normally the baseline's own final loss. For Llama-style multi-target analysis, Muon's validation losses at 20% intervals over 20%–100% of its training budget serve as five shared targets.

Why This Matters

Impact on research. The paper shows a principled way to carry Adam's variance-adaptivity into matrix-structured optimizers without discarding Muon's orthogonalized geometry, and it isolates a design choice — modulating before rather than after orthogonalization — that the ablation suggests matters. It also places itself against other adaptive Muon variants (AdaMuon, NorMuon, Muon 2) by using centered variance estimates before orthogonalization rather than uncentered second moments or post-orthogonalization statistics. The added cost is modest: an O(mn) element-wise step per parameter matrix, dominated in practice by the O(Kmn·min(m,n)) Newton–Schulz iterations, plus one extra element-wise variance buffer.

Real-world applications:

• Reducing the number of optimization steps and wall-clock time needed to hit a target validation loss under a fixed pretraining compute budget.
• Lowering the cost of repeated pretraining or domain-adaptation runs where budget, not model scale, is the binding constraint.
• Making optimizer-level efficiency gains available as drop-in replacements in existing Muon-based training stacks, since Muon-VS adds no optimizer-specific hyperparameter beyond Muon's.
• Improving the accuracy-vs-compute trade-off reported in the paper's next-token top-1 accuracy results under Suite B.

Industry relevance. Pretraining cost scales with model size and token count, and industrial-scale reports associate Muon with faster convergence and improved stability. A variant that reduces the steps needed to reach a given loss, while keeping the same optimizer-state footprint class as Adam-type optimizers and reusing existing Muon-tuned hyperparameters, is directly relevant to teams choosing optimizers for large pretraining runs.

Future Directions

• Scaling beyond 1.2B parameters. The authors explicitly list 3B and 7B regimes as beyond the scope of this study, requiring substantially larger training budgets and matched tuning for all baselines.
• Broader and deeper hyperparameter analysis of γ. Although the paper reports robustness across the tested logarithmic grids, the preferred γ setting "can vary across training recipes"; how best to select it in a new regime is left open.
• Reconciling the variance-modulation designs. The paper contrasts centered variance before orthogonalization with Muon 2's uncentered second moments and with AdaMuon's and NorMuon's post-orthogonalization statistics, and reports additional Muon 2 comparisons in Appendix B.6, leaving room for a systematic study of which modulation family wins under which conditions.
• Downstream validation at larger scale. The paper provides only a "compact zero-shot downstream sanity check" on Suite A Llama-130M and Llama-300M checkpoints, reporting consistent macro-average gains over the evaluated Muon-family baselines; downstream behavior at larger scales is not reported.

Target Audience

Optimizer researchers and machine learning engineers working on large-scale language model pretraining who already understand adaptive gradient methods and matrix-structured updates. It is also relevant to practitioners who run Muon-based training pipelines and want a low-overhead variant with limited extra tuning, and to graduate-level readers studying the theory connecting Adam's variance-adaptive interpretation to matrix-valued sign methods. Readers looking for an introductory treatment of optimizers, or for results above 1.2B parameters, will not find them here — the paper states that larger-scale settings remain future work.

Authors’ abstract

Optimizer design plays a central role in efficient language model pretraining, directly affecting optimization dynamics, convergence speed, and compute cost under fixed training budgets. Muon has emerged as a strong optimizer by orthogonalizing momentum updates, yielding a matrix-valued analogue of sign-based normalization. However, unlike Adam-style methods, Muon does not explicitly incorporate gradient-variance information into its updates. Motivated by Adam's variance-adaptive interpretation, we propose Muon-NSR and Muon-VS, two variance-adaptive Muon variants for language model pretraining. Muon-NSR applies noise-to-signal ratio (NSR) modulation before Newton--Schulz orthogonalization, whereas Muon-VS uses variance scaling (VS) without introducing any additional hyperparameters beyond those of Muon. Both methods preserve Muon's spectral normalization structure while requiring only one additional variance buffer. Experiments on Llama-style and GPT-2 pretraining across model scales from 125M to 1.2B parameters show that our methods improve over well-tuned Muon baselines and remain competitive with representative adaptive Muon-family baselines. On Llama-1.2B, Muon-VS achieves a 1.33$\times$ step-to-target speedup over a well-tuned Muon baseline, with Muon's final validation loss as the target. These results indicate that variance-adaptive modulation is a simple and effective mechanism for improving Muon-style optimizers in language model pretraining.