PaperlayerView exact revision on arXiv

Source-linked AI summary

Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning

Ziyue Liu, Ruijie Zhang, Zhengyang Wang, Yequan Zhao, Yupeng Su, Zi Yang, Zheng Zhang

arXiv:2604.09967v2cs.LGcs.AI

TL;DR

MUON’s matrix-aware orthogonalization is limited by the cost of Newton–Schulz iterations and the conditioning of its momentum input. MUON2 applies Adam-style adaptive second-moment preconditioning before orthogonalization, improving practical orthogonalization and model-training efficiency; its evidence covers models up to 13B parameters, not frontier-scale settings.

  • Problem

    MUON orthogonalization quality depends on costly Newton–Schulz iterations, whose computation and communication overhead increases the burden of large-model pretraining.

  • Method

    MUON2 applies Adam-style adaptive second-moment scaling to the momentum matrix before MUON’s orthogonalization step.

  • Results

    Across GPT, LLaMA, and Mixture-of-Experts pretraining up to 13B parameters, MUON2 consistently outperforms MUON and its variants with 40% fewer polar iterations and saves up to 25% GPU-hours at the same training loss.

  • Takeaways & Limitations

    MUON2 and MUON2-F provide practically sufficient orthogonalization with lower iteration cost while retaining improved model performance; MUON2-F preserves most gains with lower second-moment memory overhead.

  • Takeaways & Limitations

    MUON2 has not yet been evaluated at frontier-scale settings viable under industrial budgets, and other memory-efficient designs may better balance quality and optimizer-state cost.

Abstract

from arXiv · show

Muon has emerged as a promising optimizer for large-scale foundation model pre-training by exploiting the matrix structure of neural network updates through iterative orthogonalization. However, the orthogonalization quality of Muon hinges on the number of Newton--Schulz (NS) iterations performed, which poses efficiency challenges due to its non-trivial computation and communication cost. We propose Muon$^2$, an extension of Muon, to improve both quality and efficiency by applying Adam-style adaptive second-moment preconditioning before orthogonalization. Our key insight is that the core challenge of polar approximation in Muon lies in the ill-conditioned momentum matrix, of which the spectrum is substantially improved by Muon$^2$, leading to faster convergence toward a practically sufficient orthogonalization. We further characterize the practical orthogonalization quality via directional alignment, under which Muon$^2$ demonstrates dramatic improvement over Muon at each polar step. Across GPT, LLaMA, and Mixture-of-Experts pre-training experiments up to 13B parameters, Muon$^2$ (and its memory-efficient variant Muon$^2$-F that preserves most of its benefits) consistently outperforms Muon and its variants while reducing NS iterations by 40%, and saves up to 1/4 training time over Muon when achieving the same loss.

1 Introduction

MUON2 addresses MUON’s costly orthogonalization by preconditioning its momentum matrix with Adam-style adaptive second-moment scaling. This improves conditioning and orthogonalization while preserving efficiency gains in large-model pretraining.

  • 1 Introduction: MUON explicitly exploits matrix structure through Newton–Schulz polar approximation, unlike adaptive first-order optimizers that overlook this structure.This motivates alternatives to Adam and AdamW for large-model training.
  • 1 Introduction: MUON variants have often modified update rules, leaving the computation and communication burden of repeated Newton–Schulz iterations insufficiently addressed.The section frames simultaneous quality improvement and burden reduction as the central question.
  • 1 Introduction: MUON2 preconditions MUON’s momentum matrix with Adam-style adaptive scaling before orthogonalization.The method is designed to improve both model performance and training efficiency.
  • 1 Introduction: Ill-conditioned momentum spectra limit polar approximation, while MUON2 improves conditioning, directional alignment, and iteration efficiency.These properties support practically sufficient orthogonalization with fewer polar iterations.
  • 1 Introduction: MUON2-F reduces full second-moment memory overhead while preserving most of MUON2’s performance gain.It uses a factorized second-moment preconditioner.
  • 1 Introduction: 40% fewer Newton–Schulz iterations accompany consistent performance gains across GPT, LLaMA, and Mixture-of-Experts models up to 13B parameters.The experiments compare the MUON2 family with MUON and its variants.

2 Related Work

Related work spans coordinate-wise adaptive optimizers, matrix-structured preconditioners, and MUON variants. These approaches target adaptive scaling, matrix conditioning, polar convergence, memory, or distributed scalability through different mechanisms.

  • Coordinate-wise Adaptive Methods: Coordinate-wise methods use historical gradients or first- and second-moment statistics for per-parameter adaptive scaling.Adafactor factorizes second-moment statistics to reduce memory overhead.
  • Matrix-Structured Methods: Matrix-structured methods improve conditioning through Kronecker-factored preconditioning, adaptive stabilization, or direct polar-factor approximation.The cited methods include Shampoo, SOAP, and MUON.
  • Variants of MUON: MUON variants improve polar convergence, incorporate second-moment statistics, or use low-rank orthogonalization for distributed scalability.PolarExpress and Turbo-Muon target convergence; NorMuon and AdaMuon use second-moment statistics; Dion explores low-rank orthogonalization.

3 The MUON2 Optimizer

MUON2 adds Adam-style second-moment preconditioning before MUON’s orthogonalization, improving the spectrum entering Newton–Schulz iterations and enabling stronger polar approximation with fewer steps.

  • 3 The MUON2 Optimizer: MUON2’s only modification is second-moment scaling before orthogonalization, which improves the spectral properties of the matrix entering Newton–Schulz iterations.The resulting orthogonalization converges better and faster than MUON’s.
  • 3 The MUON2 Optimizer: Cosine similarity measures directional alignment with the true orthogonalized update, while exact orthogonality can be misleading under global scaling.Scaling can be absorbed into the learning rate, so directional alignment better reflects the update’s practical role.
  • 3 The MUON2 Optimizer: MUON2 shifts the NS input spectrum toward larger singular values, with its early-training center around 10^-2 versus MUON’s 10^-3.MUON’s distribution spans 10^-4 to 1, whereas most MUON2 values fall into the transition zone.
  • 3 The MUON2 Optimizer: At Ns = 3, MUON2 produces a tighter spectrum than MUON and reduces the density of extremely small singular values by roughly half compared with MUON at Ns = 5.By Ns = 5, MUON2’s singular values fall almost entirely within the target range.
  • 3 The MUON2 Optimizer: MUON2 at Ns = 3 reaches cosine similarity 0.916 versus MUON at Ns = 5 with 0.931, then reaches 0.975 at Ns = 5.Reducing MUON from Ns = 5 to Ns = 3 lowers cosine similarity from 0.931 to 0.808.
  • 3 The MUON2 Optimizer: MUON2 reduces the NS iterations required for practically sufficient orthogonalization by 40%.Three-step MUON2 already achieves better model performance than five-step MUON.
  • 3 The MUON2 Optimizer: MUON2-F stores row- and columnwise second-moment statistics and remains close to exact MUON2 in iteration requirements and model performance.Its memory overhead remains nearly unchanged from MUON.

4 Experiments

Across GPT, LLaMA, and Mixture-of-Experts pretraining, MUON2 consistently improves performance while requiring fewer Newton–Schulz iterations than MUON and related variants.

  • 4 Experiments: MUON2 consistently outperforms MUON across GPT, LLaMA, and Mixture-of-Experts pretraining while using substantially fewer Newton–Schulz iterations.Experiments cover multiple model scales and architectures, including models up to 13B parameters.
  • 4.1 Pre-Training GPT: With Ns = 3, MUON2 outperforms MUON with Ns = 5 at all three GPT scales.The comparison holds even when MUON2 uses fewer orthogonalization iterations.
  • 4.1 Pre-Training GPT: MUON2’s benefits persist across learning-rate choices, suggesting applicability beyond a single GPT-Large hyperparameter setting.The GPT-Large comparison uses a learning-rate sweep.
  • 4.2 Pre-Training LLaMA: MUON2 remains superior on LLaMA models with fewer NS iterations, while reducing NS steps causes much smaller performance drops than for MUON.The gap between MUON2 and MUON2-F is also less pronounced under reduced iteration budgets.
  • 4.3 Pre-Training Mixture-of-Experts: MUON2 outperforms MUON on the evaluated MoE model at both NS settings and reduces NS steps by 40% while improving model performance.The MoE experiment uses a 7B-A1B model with 1B active parameters.
  • 4.4 Comparison with MUON Variants: Existing MUON variants do not match MUON2’s combination of performance and reduced NS iterations.PolarExpress and Turbo-Muon underperform MUON2, while NorMuon and AdaMuon fail to reduce necessary NS iterations and underperform MUON2.

NS Steps Avg ARC-c ARC-e OBQA HellaS. PIQA WinoG. LAMBADA MMLU

The reported evaluations compare MUON2 with MUON and related variants using zero-shot metrics, training loss, runtime, and memory usage.

  • Zero-shot evaluation: Zero-shot evaluations on the trained LLaMA-1B align with validation perplexity, with MUON2 consistently outperforming MUON.The evaluations are summarized in Table 4 across different NS iterations.
  • Training efficiency: MUON2 achieves the same LLaMA-1B training loss as MUON with up to 25% fewer training steps.The comparison is shown in Figure 5.
  • Training efficiency: 25% fewer training time corresponds to 1042 versus 796 GPU-hours when achieving the same loss of 2.36.MUON2’s second-moment scaling leaves per-step time almost unchanged from MUON.
  • Training efficiency: For the same training steps, MUON2 with Ns = 3 reaches loss 2.34 versus 2.36 for MUON while reducing GPU-hours from 1044 to 1016.This comparison reports both loss and compute differences at matched training-step budgets.
  • Memory efficiency: MUON2-F uses practically the same memory as MUON while preserving most of MUON2’s performance gains.The memory comparison is reported across several optimization approaches.

5 Conclusion

Muon2 combines adaptive second-moment preconditioning with MUON orthogonalization to improve optimization behavior while reducing orthogonalization cost. It achieves practically sufficient orthogonalization with fewer polar iterations and lowers GPU-hours at the same training loss.

  • Muon2 applies adaptive second-moment preconditioning before orthogonalization to improve optimization behavior and reduce orthogonalization cost.The method addresses the ill-conditioned input to polar approximation and improves orthogonalization quality.
  • 40% fewer polar iterations are achieved by MUON2(-F) while consistently improving model performance.
  • Up to 25% fewer GPU-hours are required than MUON when achieving the same training loss.

6 Limitation

The evaluation does not yet cover frontier-scale settings that require industrial budgets. The authors also leave broader comparisons of memory-efficient optimizer designs for future work.

  • MUON2 has not been evaluated at frontier-scale settings because the study was conducted under academic budgets.The reported experiments include models up to 13B parameters, exceeding the scales used in most cited studies.
  • Other memory-efficient designs may better balance optimization quality and optimizer-state cost, but they remain unexamined.

7 Ethical Considerations

The work uses public corpora and introduces neither new datasets nor personally identifiable information. Models trained with MUON2 may still inherit harmful behaviors from their training data and require standard safety evaluation.

  • The study uses publicly available corpora and does not introduce new datasets or collect personally identifiable information.
  • MUON2-trained models may inherit biases, toxic content, hallucinations, or other harmful behaviors from their underlying training data.The authors recommend standard safety evaluation and responsible deployment practices.
  • The implementation adopts Adafactor and leaves comprehensive comparison of alternative memory-overhead reductions for future work.

B Discussion on Cosine Similarity

Cosine similarity evaluates directional alignment more faithfully than scale-sensitive exact orthogonality error for MUON’s practically inexact orthogonalization. The loose target reaches sufficient alignment quickly despite a larger exact orthogonality error.

  • Discussion on Cosine Similarity: Cosine similarity measures directional alignment in [0, 1], where 1 indicates perfect alignment and 0 indicates an almost orthogonal direction.It is robust to global scaling, unlike exact orthogonality error.
  • Discussion on Cosine Similarity: Figures 6 and 7 compare singular-value mappings and convergence zones for different NS coefficients and polar methods.
  • Discussion on Cosine Similarity: MUON’s loose target maps singular values roughly from [0, 1] to [0.7, 1.3], while the exact target maps most values to one.
  • Discussion on Cosine Similarity: 5 NS iterations suffice for MUON’s loose target, which has 0.31 exact orthogonality error but 0.98 cosine similarity.The exact target has 0.03 error, illustrating that exact orthogonality error can disagree with practical directional alignment.
  • Discussion on Cosine Similarity: The orthogonalization objective maps singular values toward [1 −ϵ, 1 + ϵ], with convergence-zone boundaries depending on the polar method.PolarExpress changes the zone boundaries relative to MUON’s method.

D Detailed Evaluations

The detailed evaluations cover learning-rate sweeps across GPT and LLaMA model scales, a Mixture-of-Experts configuration, downstream zero-shot benchmarks, and comparisons with Muon variants. Results are organized around validation perplexity, benchmark accuracy, and different Newton–Schulz iteration settings.

  • Learning-rate sweeps: Learning-rate sweeps evaluate GPT-Small, GPT-Base, and GPT-Large using validation perplexity.The corresponding results are reported in Tables 8–10 and Figures 8–9 and 4.
  • Learning-rate sweeps: Learning-rate sweeps likewise cover LLaMA-60M, LLaMA-350M, and LLaMA-1B using validation perplexity.Architecture configurations are listed separately for the LLaMA models.
  • Model configurations: The Mixture-of-Experts evaluation specifies its model configuration separately from the learning-rate sweep results.The configuration is provided in Table 15.
  • Downstream evaluation: Downstream evaluation uses zero-shot accuracy on ARC, OpenBookQA, PIQA, WinoGrande, LAMBADA, and MMLU, plus their unweighted average.The suite spans commonsense reasoning, reading comprehension, knowledge recall, and multitask language understanding.
  • Comparisons: Muon² is compared with Polar-Express, Turbo-Muon, NorMuon, and AdaMuon, with learning rates swept for fairness.Additional evaluations include GPT-Large results across different Newton–Schulz iteration counts.

D.6 Training Speedup

The training-speed evaluation compares wall-time against training loss and reports per-step timing for LLaMA-1B. It also documents the hardware and precision conditions used for memory profiling.

  • Training speedup: Wall-time versus training loss is evaluated for LLaMA-1B, with the figure reporting that MUON2 reaches the same loss with 1/4 fewer GPU-hours than MUON.Per-training-step time is reported separately in Table 20.
  • Evaluation conditions: Memory profiling uses a single 94GB H100 GPU with bfloat16 precision, sequence length 4096, and micro batch size 1.The training framework is adopted from Nanotron.
  • Authorship statement: The paper states that language models were used only for grammar and wording refinement, not for ideas, experiments, analysis, or interpretation.The authors report that they produced and verified the technical content and conclusions.
Loading 2604.09967v2…