PaperlayerView exact revision on arXiv

Source-linked AI summary

AdaMuon: Adaptive Muon Optimizer

Chongjie Si, Debing Zhang, Wei Shen

arXiv:2507.11005v3cs.LG

TL;DR

Large-scale training needs optimizers that handle heterogeneous gradient statistics without sacrificing stable update geometry. AdaMuon combines sign-stabilized orthogonalization, element-wise second-momentum scaling, and RMS-aligned rescaling, with experiments reporting more than 40% training-efficiency improvement over Adam. The method also introduces additional memory and runtime costs.

  • Problem

    Muon’s global orthogonalization does not model individual-parameter variance, while large-scale training involves heterogeneous gradient statistics.

  • Method

    AdaMuon applies element-wise second-momentum estimation to orthogonalized directions, sign-transforms momentum before orthogonalization, and uses RMS-aligned rescaling.

  • Results

    More than 40% training-efficiency improvement over Adam is reported for AdaMuon in large-scale scenarios, alongside stable convergence.

  • Takeaways & Limitations

    AdaMuon consistently improves training efficiency and final performance over Adam and Muon across GPT-2 and Qwen2.5 experiments.

  • Takeaways & Limitations

    AdaMuon requires an additional second-moment buffer and incurs minor per-step overhead from Newton–Schulz and variance EMA.

Abstract

from arXiv · show

We propose AdaMuon, a novel optimizer that combines element-wise adaptivity with orthogonal updates for large-scale neural network training. AdaMuon incorporates two tightly coupled mechanisms: (1) an element-wise second momentum estimator applied to orthogonalized update directions, and (2) a sign-stabilized orthogonal update, where the momentum is first sign-transformed before orthogonalization. These two components jointly enable variance-adaptive scaling while maintaining stable update geometry. In addition, AdaMuon employs an RMS-aligned rescaling strategy to match the root-mean-square update magnitude to Adam, allowing direct reuse of existing learning rate schedules without extra tuning. Experiments demonstrate that AdaMuon not only maintains stability but can surpass Adam by more than 40\% training efficiency in large-scale scenarios.

1 INTRODUCTION

AdaMuon is proposed to combine Muon’s geometry-aware orthogonal updates with element-wise variance adaptation for large-scale training. Its RMS-aligned rescaling is designed to preserve compatibility with existing learning-rate schedules, and experiments report improved efficiency over Adam.

  • Muon applies polar decomposition to momentum matrices, producing spectrally normalized, direction-only updates for two-dimensional parameter blocks.
  • AdaMuon combines element-wise variance adaptation with orthogonal updates to address heterogeneous gradient statistics in large-scale training.The design applies second-momentum estimation to orthogonalized directions and uses sign-stabilized orthogonal updates.
  • AdaMuon uses a sign-stabilized orthogonal update and an element-wise second momentum estimator applied after orthogonalization.
  • RMS-aligned rescaling matches update magnitude to Adam, enabling reuse of existing learning-rate schedules without extra tuning.
  • More than 40% training-efficiency improvement over Adam is reported for AdaMuon in large-scale scenarios.
  • Experiments across multiple models and datasets are described as demonstrating consistent outperformance and effective scaling to large models.

2 PRELIMINARY

Muon uses momentum orthogonalization to produce geometry-aware matrix updates, while scaling and weight-decay adjustments support practical large-scale use. The preliminary discussion identifies a remaining gap: global orthogonalization does not adapt to coordinate-level variance differences, motivating element-wise second-momentum estimation.

  • 2.1 MUON OPTIMIZER: Muon is designed for matrix-shaped parameters and uses Newton–Schulz steps to approximate the polar decomposition of momentum updates.The approximation preserves update direction while removing anisotropic scaling and avoids the cost of a full SVD.
  • 2.1 MUON OPTIMIZER: Newton–Schulz normalization uses iteration coefficients and T steps, with a cited configuration of a = 3.4445, b = −4.7750, c = 2.0315, and T = 5.
  • 2.2 SCALING UP FOR MUON: Muon faces practical scaling challenges because weight RMS can become too large for bf16 and one-dimensional parameters remain optimized with Adam.These issues create separate learning-rate and tuning requirements in existing pipelines.
  • 2.2 SCALING UP FOR MUON: RMS matching with γ = 0.2√max(m, n) enables Muon and Adam to share a unified learning-rate schedule.The adjustment is described as stabilizing training and supporting integration into existing optimization pipelines.
  • 2.3.1 NECESSITY: Muon’s polar decomposition encodes row–column second-order interactions without explicitly storing full covariance matrices.
  • 2.3.1 NECESSITY: Global orthogonalization does not model individual-parameter variance, risking overshooting noisy coordinates and under-updating low-variance informative ones.The passage connects this mismatch with potentially slower convergence, reduced stability, and limited effectiveness under heterogeneous gradient statistics.
  • 2.3.1 NECESSITY: Integrating element-wise second-momentum estimation requires choosing whether to accumulate on Gt, Mt, or Ot and matching Adam-like normalization.

3 ADAMUON

AdaMuon combines element-wise second-momentum estimation with sign-stabilized orthogonal updates to adapt scaling while preserving coherent update geometry. It further rescales updates to match Adam’s RMS magnitude for schedule compatibility.

  • Core design: AdaMuon combines element-wise second-momentum estimation with sign-stabilized orthogonal updates to capture variance while preserving orthogonal geometry.The second momentum is accumulated on orthogonalized directions, while sign transformation stabilizes the input before orthogonalization.
  • Element-wise second momentum: The second-momentum buffer tracks exponential moving averages of element-wise squared orthogonalized updates, analogous to Adam’s variance accumulator.The coefficient β is inherited from Muon, so this mechanism introduces no additional hyper-parameters.
  • Element-wise second momentum: Variance normalization divides each orthogonalized update element by its estimated standard deviation, suppressing noisy coordinates while preserving global update geometry.A small ε prevents division by zero, and the operation uses element-wise division.
  • Sign-stabilized orthogonal updates: Sign transformation is selected as the canonical admissible element-wise transformation because it is scale-invariant, direction-preserving, and bounded.The construction complements global scale invariance by mitigating element-level volatility before orthogonalization.
  • RMS alignment: AdaMuon rescales its variance-normalized update so its RMS magnitude matches Adam’s empirical value of ≈0.2, enabling compatibility with Adam learning-rate schedules.The rescaled update is then applied to the parameters, alongside weight decay.
  • Algorithm and analysis: The method is presented with pseudocode and additional discussions of omitted second-momentum bias correction and convergence analysis.These supporting analyses appear in the algorithm presentation and appendices.

4 EXPERIMENT

Experiments compare AdaMuon with AdamW and Muon across GPT-2 and Qwen2.5 models, training settings, datasets, and benchmarks. AdaMuon generally converges fastest, achieves strong efficiency gains, and benefits from both sign stabilization and second-momentum scaling.

  • GPT-2 Results: AdaMuon consistently achieves the highest training efficiency across four GPT-2 scales and learning-rate settings.The comparison uses AdamW as the baseline and attributes the result to second-moment scaling with RMS-norm alignment.
  • Qwen2.5 Results: 30.8% less wall-clock time than Adam and 2.9% less than Muon are reported for AdaMuon on Qwen2.5 time-to-target reduction.AdaMuon also outperforms Muon and Adam across all evaluation metrics after 100B-token training on 15 benchmarks.
  • Ablation Study: Retaining both sign stabilization and second momentum produces the best overall ablation performance.Sign stabilization improves orthogonal-update stability, while second momentum provides element-wise normalization.
  • Training Behavior: AdaMuon produces smoother gradient-norm trajectories and smaller parameter norms than AdamW and Muon.The reported training-behavior analysis also compares per-layer maximum attention logits across shallow, middle, and deep layers.
  • Sensitivity Analysis: Final GPT-small training losses cluster tightly across reasonable β values, allowing reuse of Muon’s default β = 0.95 without special tuning.The sensitivity analysis varies the first- and second-momentum coefficient at learning rate 6 × 10^-4.

5 CONCLUSION

AdaMuon combines Muon’s geometry-aware stability with coordinate-wise variance adaptivity through sign-stabilized orthogonal updates, second-momentum estimation, and RMS-aligned rescaling. Experiments across GPT-2 and Qwen2.5 show consistent improvements over Adam and Muon, with up to 40% higher training efficiency and superior final performance.

  • 5 CONCLUSION: AdaMuon unifies geometry-aware stability with coordinate-wise adaptivity through three integrated update mechanisms.These are sign-stabilized orthogonal updates, an element-wise second-momentum estimator, and RMS-aligned rescaling.
  • 5 CONCLUSION: Experiments across multiple GPT-2 and Qwen2.5 scales show consistent improvements over Adam and Muon.The reported gains include superior final performance and up to 40% training-efficiency improvement.
  • 5 CONCLUSION: AdaMuon is presented as a versatile algorithm that scales effectively to large-scale models.The paper identifies practical and scalable second-order optimization as an open question for future work.

6 LIMITATION

The paper identifies additional memory and computation as AdaMuon’s main limitations.

  • 6 LIMITATION: AdaMuon requires an additional second-moment buffer, increasing computation and memory usage.The paper also reports minor per-step runtime overhead from Newton–Schulz and variance EMA.
  • 6 LIMITATION: Newton–Schulz and variance EMA introduce minor per-step runtime overhead.The paper leaves addressing these drawbacks to future work.

A PROOF OF THEOREM 1

The proof characterizes orientation-preserving, odd, scale-invariant functions as sign functions up to a finite positive scale, then explains why AdaMuon omits explicit momentum bias correction. RMS alignment removes constant second-momentum bias, while orthogonalization removes constant first-momentum magnitude bias.

  • A PROOF OF THEOREM 1: Orientation preservation and oddness constrain f(x) to h(|x|) · sign(x), while scale invariance makes h constant on positive inputs.The proof concludes f(x) = κ · sign(x) for finite κ.
  • A PROOF OF THEOREM 1: Global scale invariance eliminates the multiplicative factor κ, yielding the canonical representative f(x) = sign(x).The proof applies this result to the normalized update construction.
  • Bias Correction: RMS alignment makes explicit second-momentum bias correction unnecessary by canceling any constant multiplicative bias in V_t.The rescaling targets a fixed RMS magnitude of 0.2, matching Adam’s typical update norm.
  • Bias Correction: Orthogonalization removes constant magnitude bias in the first-momentum buffer, so explicit correction of M_t is unnecessary.The normalized direction is unaffected by a constant multiplicative factor in M_t.

C CONVERGENCE ANALYSIS

The convergence analysis studies AdaMuon under smoothness, unbiased stochastic gradients, and directional alignment assumptions. Its preconditioner remains bounded, while RMS alignment fixes the update norm.

  • Assumptions: AdaMuon is analyzed under smooth nonconvex assumptions with unbiased stochastic gradients and positive directional alignment after normalization.The alignment condition accounts for sign transformation, polar-factor orthogonalization, and element-wise normalization.
  • RMS alignment: RMS-aligned rescaling fixes the step norm at the prescribed constant c_rms.The analysis uses c_rms = 0.2, while noting that any fixed positive constant works.
  • Preconditioner bounds: The normalized second-momentum preconditioner V_t stays between 0 and 1 entrywise throughout the recursion.This follows from initializing V_0 at zero and updating it with squared orthogonalized directions.
  • Preconditioner bounds: The preconditioner bounds control the normalized update magnitude through ε, yielding a bound on ∥bO_t∥_F.The resulting bound is ∥bO_t∥_F ≤ ε^-1∥O_t∥_F = √r/ε.

C.2 ONE-STEP PROGRESS

The one-step analysis combines smoothness with directional alignment and bounded preconditioning to obtain a descent inequality for AdaMuon. This inequality supports the subsequent convergence results.

  • One-step inequality: Under smoothness, AdaMuon's expected one-step loss change is bounded using the gradient, update direction, and update norm.The proposition provides the starting inequality for the convergence analysis.
  • One-step inequality: Directional alignment converts the update-gradient inner product into a negative term proportional to ∥G_t∥_F.The bound uses −α γ η_t∥G_t∥_F after applying the alignment assumption.
  • One-step inequality: The RMS-aligned step norm supplies the remaining quadratic update term in the one-step bound.The argument invokes the fixed-norm result from Lemma 2.
  • Accumulated progress: Summing the expected one-step inequalities over iterations yields a bound on average gradient magnitude.The analysis defines the initial loss gap Δ_0 before taking expectations and summing.

C.3 PL-BASED CONVERGENCE

Under the PL condition and directional alignment, AdaMuon converges to an O(η^2) neighborhood. The analysis also characterizes the tradeoff between step size, hitting time, and asymptotic accuracy.

  • C.3 PL-based convergence: The loss gap eventually remains below the squared threshold determined by the recursion's balance between descent and error.The limit-superior argument shows lim sup_{t→∞} x_t ≤ (b/a)^2.
  • C.3 PL-based convergence: Under PL and cosine-alignment assumptions, AdaMuon converges to an O(η^2) neighborhood.The constant-step recursion is x_{t+1} ≤ x_t − a√x_t + b, with a = α γ η and b = Kη^2 in the stated analysis.
  • Convergence speed: The target-radius analysis gives an O(1/η^2) iteration bound for entering a fixed inflated neighborhood.For δ = 1, the stated hitting-time bound is T_hit(1) ≤ (x_1 − 4(b/a)^2)/b = O(1/η^2).
  • Convergence speed: Larger η reduces hitting time as O(1/η^2) but enlarges the asymptotic neighborhood as O(η^2).This is the explicit step-size tradeoff identified by the analysis.
  • Stronger alignment: With stronger alignment, AdaMuon has linear convergence to an O(η^2) neighborhood under PL.The contraction factor is ρ = 2κ γ µ η, with an additive Kη^2 term.

D.1 MOE ARCHITECTURE

The study evaluates AdaMuon on a 1.3B-parameter DeepSeek V3 mixture-of-experts model trained for 300B tokens, and also broadens GPT2-Small comparisons to Shampoo and Lion.

  • MOE ARCHITECTURE: The DeepSeek V3 evaluation uses a 1.3B-parameter MoE model with 0.11B activated parameters.The experiment uses batch size 1024, sequence length 8192, and a learning rate of 4.2×10^-4.
  • MOE ARCHITECTURE: The DeepSeek V3 model is trained for 300B tokens after 16.8B warmup tokens using the Qwen training dataset.Training-loss curves are reported in Figure 5.
  • MOE ARCHITECTURE: On GPT2-Small, the broader comparison includes AdamW, Muon, Shampoo, Lion, and AdaMuon under the same training recipe.Shampoo and Lion use their recommended default hyperparameters without bespoke tuning.
Loading 2507.11005v3…