Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

Minxin Zhang, Yuxuan Liu, Hayden Schaeffer

TL;DR

This paper introduces NAMO and NAMO-D, two optimizers that fuse orthogonalized momentum with norm-based Adam-style noise adaptation for matrix-structured parameters. The authors prove optimal convergence rates in both deterministic ($O(T^{-1/2})$) and stochastic ($O(T^{-1/4})$ with batch-size dependent terms) settings, with NAMO-D offering finer-grained, diagonal adaptation. Empirically, NAMO and NAMO-D outperform AdamW and Muon in GPT-2 pretraining on OpenWebText, with NAMO-D delivering additional gains via a clamping hyperparameter. These results show scalable, robust optimization benefits for large-language-model training and motivate further exploration of noise-adaptive matrix updates in deep learning.

Abstract

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.

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

TL;DR

This paper introduces NAMO and NAMO-D, two optimizers that fuse orthogonalized momentum with norm-based Adam-style noise adaptation for matrix-structured parameters. The authors prove optimal convergence rates in both deterministic () and stochastic ( with batch-size dependent terms) settings, with NAMO-D offering finer-grained, diagonal adaptation. Empirically, NAMO and NAMO-D outperform AdamW and Muon in GPT-2 pretraining on OpenWebText, with NAMO-D delivering additional gains via a clamping hyperparameter. These results show scalable, robust optimization benefits for large-language-model training and motivate further exploration of noise-adaptive matrix updates in deep learning.

Abstract

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.
Paper Structure (24 sections, 9 theorems, 152 equations, 3 figures, 2 tables, 2 algorithms)

This paper contains 24 sections, 9 theorems, 152 equations, 3 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related work
  3. Contributions and Organization
  4. New Optimization Algorithms: NAMO and NAMO-D
  5. Convergence Analysis
  6. Analysis of NAMO
  7. Proof sketch.
  8. Proof sketch.
  9. Analysis of NAMO-D
  10. Proof sketch.
  11. Proof sketch.
  12. Experiments
  13. Experimental setup
  14. Model Architectures.
  15. Dataset.
  16. ...and 9 more sections

Key Result

Theorem 3

Suppose Assumptions assum:func holds. Let $\left\{\Theta_t\right\}\subset\mathbb{R}^{m\times n}$ be the sequence of iterates generated by Algorithm alg:namo with full-batch gradients and $0\le \mu_1\le \mu_2<1$. If choosing $\eta = \mathcal{O}\left(T^{-\frac{1}{2}}\right)$, $\epsilon=\mathcal{O}(T^{

Figures (3)

  • Figure 1: Hyperparameter sweeping results for GPT-2 (124M). The training and validation losses at step 10K are reported, where the x-axis is the learning rate.
  • Figure 2: Pretraining GPT-2 (124M) for 50K steps. The optimal LR from sweeping for 10K steps is used.
  • Figure 3: Pretraining GPT-2 (355M) for 10K steps. The optimal LR (and optimal $c$ for NAMO-D) from sweeping for 10K steps are used.

Theorems & Definitions (9)

  • Theorem 3: NAMO in the deterministic case
  • Theorem 4: NAMO in the stochastic case
  • Theorem 5: NAMO-D in the deterministic case
  • Theorem 6: NAMO-D in the stochastic case
  • Lemma 7
  • Lemma 8
  • Lemma 9
  • Lemma 10
  • Lemma 11