PaperlayerView exact revision on arXiv

Source-linked AI summary

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

Minxin Zhang, Yuxuan Liu, Hayden Schaeffer

arXiv:2602.17080v2cs.LGmath.OC

TL;DR

The paper addresses the challenge of combining structured update directions with noise adaptation in stochastic optimization. It proposes NAMO and NAMO-D, integrating norm-based Adam-type scaling with orthogonalized momentum. Both methods have convergence guarantees under standard assumptions and outperform AdamW and Muon in GPT-2 pretraining, with NAMO-D improving further through neuron-wise clamped adaptation.

  • Problem

    Stochastic optimization needs update directions that work well without noise and adaptive stepsizes that stabilize updates under gradient uncertainty.

  • Method

    NAMO uses a single norm-based adaptive scalar for orthogonalized momentum, while NAMO-D uses clamped column-wise diagonal scaling for neuron-wise noise adaptation.

  • Results

    NAMO and NAMO-D achieve theoretical convergence guarantees and outperform AdamW and Muon in GPT-2 pretraining, with NAMO-D providing further gains over NAMO.

  • Takeaways & Limitations

    Orthogonalized momentum can be combined with norm-based noise adaptation, while diagonal clamping trades strict orthogonality for finer-grained adaptation and improved reported performance.

  • Takeaways & Limitations

    The theory assumes smooth loss functions and unbiased stochastic gradients with uniformly bounded variance.

Abstract

from arXiv · show

Efficient stochastic optimization typically integrates an update direction that performs well in the deterministic regime with a mechanism adapting to stochastic perturbations. While Adam uses adaptive moment estimates to promote stability, Muon utilizes the weight layers' matrix structure via orthogonalized momentum, showing superior performance in large language model training. We propose a new optimizer and a diagonal extension, NAMO and NAMO-D, providing the first principled integration of orthogonalized momentum with norm-based Adam-type noise adaptation. NAMO scales orthogonalized momentum using a single adaptive stepsize, preserving orthogonality while improving upon Muon at negligible additional cost. NAMO-D instead right-multiplies orthogonalized momentum by a diagonal matrix with clamped entries. This design enables neuron-wise noise adaptation and aligns with the common near block-diagonal Hessian structure. Under standard assumptions, we establish optimal convergence rates for both algorithms in the deterministic setting and show that, in the stochastic setting, their convergence guarantees adapt to the noise level of stochastic gradients. Experiments on pretraining GPT-2 models demonstrate improved performance of both NAMO and NAMO-D compared to the AdamW and Muon baselines, with NAMO-D achieving further gains over NAMO via an additional clamping hyperparameter that balances the competing goals of maintaining a well-conditioned update direction and leveraging fine-grained noise adaptation.

1 Introduction

The paper frames efficient stochastic optimization as combining a strong deterministic update direction with adaptive noise control. It motivates integrating Muon’s orthogonalized momentum with Adam-type adaptation, yielding NAMO and NAMO-D.

  • Motivation: Stochastic optimizers must balance effective noise-free update directions with stepsize mechanisms that stabilize iterates under gradient uncertainty.The paper presents direction selection and adaptive stepsizes as complementary ingredients.
  • Related methods: Adam uses coordinate-wise adaptive stepsizes, whereas Muon exploits matrix structure through orthogonalized momentum.Muon orthogonalizes momentum matrices using an approximate polar factor computed by Newton–Schulz iterations.
  • Motivation: Orthogonalization can amplify stochastic noise because it is unbounded, increasing training instability and sensitivity to hyperparameter choices.This motivates pairing Muon’s update direction with explicit noise adaptation.
  • Contributions: NAMO combines norm-based moment estimation with orthogonalized momentum, using one adaptive scalar that preserves update-direction orthogonality.The method targets noise adaptation while retaining Muon’s structured direction.
  • Contributions: NAMO-D right-multiplies orthogonalized momentum by a diagonal adaptive matrix, enabling neuron-wise scaling but relaxing strict orthogonality.Its clamping parameter balances well-conditioned directions against finer-grained adaptation.
  • Related work: The proposed methods are presented as theoretically principled adaptive Muon variants, complementing related approaches that lack theoretical guarantees or LLM-training evaluation.The introduction positions NAMO and NAMO-D relative to AdaMuon, NorMuon, and AdaGO.

2 New Optimization Algorithms: NAMO and NAMO-D

NAMO uses a norm-based scalar adaptation for orthogonalized momentum, while NAMO-D uses clamped column-wise adaptation for neuron-wise noise control. The algorithms are analyzed under smoothness and unbiased bounded-variance assumptions, with approximate orthogonalization used in experiments.

  • Assumptions: The analysis assumes a Lipschitz-continuous gradient and stochastic gradients that are unbiased with uniformly bounded variance.These assumptions support the deterministic and stochastic convergence analyses.
  • Implementation: Exact orthogonalization is assumed for theory, whereas experiments use Newton–Schulz iterations to compute an approximate orthogonalization.This separates the theoretical algorithms from their practical implementation.
  • NAMO: NAMO maintains a momentum matrix and a second raw-moment estimate based on the squared Frobenius norm of stochastic gradients.Bias correction produces the quantities used to adapt the scalar multiplier on orthogonalized momentum.
  • NAMO: NAMO scales orthogonalized momentum with a single adaptive scalar that becomes small under noisy gradients or near-stationary iterates.This scalar is intended to promote stable convergence.
  • NAMO-D: NAMO-D estimates squared norms separately for gradient columns and assigns an individual adaptive stepsize to each neuron.Column-wise scaling provides finer-grained noise adaptation and reflects near block-diagonal Hessian structure.
  • NAMO-D: NAMO-D clamps column-wise stepsizes toward their average, preserving a well-conditioned scaled direction while sacrificing strict orthogonality.The clamp uses a parameter c in (0, 1].

3 Convergence Analysis

The analysis establishes optimal-order convergence guarantees for NAMO and NAMO-D in deterministic and stochastic settings. In the stochastic regime, both methods adapt to gradient-noise levels, with sufficiently large batches recovering the optimal O(T^-1/4) rate.

  • NAMO in the deterministic case: NAMO achieves the optimal deterministic O(T^-1/2) convergence rate under the stated smoothness assumptions.This matches the known lower complexity bound for deterministic first-order methods.
  • NAMO in the stochastic case: NAMO's stochastic convergence adapts to gradient noise and recovers the optimal O(T^-1/4) rate when b = Ω(σ^2T^1/2).The bound contains a variance-dependent term that becomes asymptotically dominated at the specified batch size.
  • NAMO-D in the deterministic case: NAMO-D uses column-wise diagonal scaling instead of NAMO's scalar adaptive stepsize.The diagonal scaling is clamped, supporting control of the scaling matrix's conditioning.
  • NAMO-D in the deterministic case: NAMO-D achieves optimal-order deterministic convergence under Assumption 1 with appropriately chosen learning-rate, moment, regularization, and clamping parameters.The theorem uses full-batch gradients and sets c = Θ(1).
  • NAMO-D in the stochastic case: NAMO-D's stochastic convergence adapts to gradient noise and reaches the optimal O(T^-1/4) rate for sufficiently large batches.The result is stated under Assumptions 1–2 with the theorem's prescribed parameter choices.
  • Overall convergence guarantees: Both methods preserve optimal order in deterministic and stochastic settings, while their bounds explicitly quantify the effects of gradient noise and batch size.The unified analysis combines orthogonalized descent with bias-corrected moments and controlled scalar or diagonal adaptive scaling.

4 Experiments

Experiments compare NAMO and NAMO-D with AdamW and Muon on GPT-2 models under matched training settings. Both proposed methods perform better, while NAMO-D gains further from neuron-wise adaptation and clamping.

  • Experimental setup: Experiments compare NAMO and NAMO-D against AdamW and Muon on GPT-2 models with 124M and 355M parameters.All experiments use OpenWebText and four NVIDIA H100 GPUs.
  • GPT-2 (124M): NAMO and NAMO-D achieve lower training and validation losses across a wider learning-rate range than Muon and AdamW after 10K steps.The sweep evaluates losses at 10K steps and indicates accelerated convergence and improved tuning robustness.
  • GPT-2 (124M): NAMO-D attains lower training and validation losses than NAMO when GPT-2 (124M) training extends to 50K steps using the selected learning rates.The reported advantage is attributed to finer-grained neuron-wise adaptive scaling.
  • GPT-2 (355M): For GPT-2 (355M), the experiments sweep learning rates for all optimizers and additionally sweep NAMO-D’s clamping hyperparameter c.The optimal learning rate is selected by lowest validation loss after 10K steps.
  • GPT-2 (355M): NAMO and NAMO-D outperform Muon and AdamW on GPT-2 (355M), with NAMO-D providing further gains over NAMO through c.The clamping parameter balances preserving a well-conditioned update direction with finer-grained noise adaptation.

5 Conclusions and Future Work

The paper concludes that NAMO and NAMO-D integrate orthogonalized updates with Adam-type noise adaptation while retaining theoretical convergence guarantees. NAMO preserves orthogonality, whereas NAMO-D trades strict orthogonality for finer-grained adaptation controlled by clamping.

  • Conclusions: NAMO and NAMO-D provide a theoretically principled integration of orthogonalized update directions with norm-based adaptive moment estimation.The methods target noise adaptation for matrix-structured parameters.
  • Conclusions: NAMO uses a single adaptive stepsize and preserves the orthogonality of Muon’s update direction at negligible additional computational cost.The additional cost is O(mn), with no additional memory overhead.
  • Conclusions: NAMO-D right-multiplies orthogonalized momentum by a diagonal scaling, enabling neuron-wise noise adaptation without strictly preserving orthogonality.A clamping parameter balances update conditioning against finer-grained adaptation.
  • Theory: Under standard smoothness and unbiased bounded-variance noise assumptions, both algorithms achieve optimal deterministic convergence rates.Their stochastic guarantees adapt to the noise level of stochastic gradients.
  • Theory: In stochastic settings, the convergence guarantees attain the optimal rate when the batch size is sufficiently large.

Appendix B. Proof of Theorem 3

This appendix section presents the detailed proof of NAMO’s deterministic convergence theorem.

  • Appendix B. Proof of Theorem 3: The proof establishes convergence of NAMO in the deterministic setting.

Proof

The proof derives NAMO’s deterministic convergence result through stepsize control, descent inequalities, momentum and adaptive-term bounds, and a final rate argument.

  • Proof: The proof first establishes a uniform upper bound on the adaptive stepsize.
  • Proof: It then derives a descent inequality and averages the resulting bounds.The initial objective gap is represented by Δ := L(Θ0) − minΘ L(Θ).
  • Proof: The analysis bounds the distance between bias-corrected momentum and the true gradient, then bounds the adaptive second-moment term.
  • Proof: The proof relates the adaptive stepsize to the gradient norm using auxiliary quantities and constants depending on µ1 and µ2.
  • Proof: The convergence rate is derived using Lemma 8, completing the proof for sufficiently large T.

Appendix C. Proof of Theorem 4

This section presents the detailed proof of Theorem 4 for NAMO convergence in the stochastic setting.

  • The proof establishes convergence of NAMO in the stochastic setting.

Proof

The proof bounds stepsizes, controls bias-corrected momentum and adaptive estimates, derives an average gradient bound, and obtains a convergence rate under parameter choices.

  • Step 1: A uniform upper bound on stepsize: The proof begins by deriving a uniform upper bound on the stepsize.
  • Step 2: Expected descent and an average bound: Conditional expected descent is combined with bounds on α_t∥M̂_t∥_* to control optimization progress.
  • Step 3: Bounding the distance: The analysis bounds the distance between bias-corrected momentum and the true gradient.
  • Step 4: Bounding E[√v̂_t]: Minkowski, Jensen, and related inequalities bound E[√v̂_t] and connect the adaptive estimate to the gradient norm.
  • Step 6: Deriving the convergence rate: The final steps rearrange the bounds and derive a convergence rate for the expected gradient norm.
  • For large T, the proof considers η = O(T^-3/4), 1 − μ1 = Θ(T^-1/2), 1 − μ2 = Θ(T^-1/2), and ε = O(T^-1/2).

Appendix D. Proof of Theorem 5

Theorem 5 is proved for NAMO-D in the deterministic setting by controlling its diagonal scaling, descent, gradient estimates, and resulting convergence rate.

  • Step 1: A uniform upper bound on stepsize: The proof derives a uniform upper bound on the NAMO-D stepsize.
  • Bias-corrected momentum and adaptive quantities are defined columnwise before analyzing the diagonal update matrix.
  • The diagonal scaling uses a clamping hyperparameter c ∈ (0, 1], and its condition number is bounded through κ_t.
  • Step 2: Descent inequality and averaging: A descent inequality and averaging argument control progress, while the scaled momentum is compared with the true gradient.
  • Steps 5–6: Gradient bounds and rate: The proof lower-bounds d_t,max∥∇L(Θ_t−1)∥_* and then derives an upper bound on the gradient norm.

Appendix E. Proof of Theorem 6

Theorem 6 analyzes NAMO-D convergence in the stochastic setting through stepsize, diagonal-scaling, descent, momentum-error, and rate bounds.

  • Step 1: A uniform upper bound on stepsize: The proof begins with a uniform upper bound on the NAMO-D stepsize.
  • The diagonal update uses a fixed clamping hyperparameter c ∈ (0, 1], with its condition number characterized through κ(diag(d_t)).
  • Step 2: Expected descent inequality and averaging: Conditional expected descent and averaging are established using the bias-corrected momentum error and the diagonal scaling bounds.
  • Step 3: Bounding the distance: The stochastic analysis bounds the distance between bias-corrected momentum and the true gradient using auxiliary inequalities.
  • Step 6: Deriving the convergence rate: The final steps combine gradient-norm bounds to derive the stochastic convergence rate.
Loading 2602.17080v2…