Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

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 LLM 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.

• The paper introduces NAMO and Diagonal NAMO, which combine orthogonalized momentum with norm-based moment estimation to address Muon’s sensitivity to noise.
• Empirical evaluations on GPT-2 models show improved convergence and reduced learning rate sensitivity compared to AdamW and Muon.
• Theoretical analysis confirms optimal convergence rates and highlights the benefits of structured, adaptive scaling for training large-scale neural networks.

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

Motivation and Context

The paper "Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum" (2602.17080) addresses fundamental limitations in stochastic optimization for deep learning, particularly in the context of LLM training. Standard adaptive optimizers, notably Adam and AdamW, utilize coordinate-wise adaptive moment estimates to stabilize updates. While Muon leverages weight matrix structure via orthogonalized momentum (i.e., polar factor based updates), it lacks inherent noise adaptation, resulting in instability and sensitivity to hyperparameters in noisy stochastic gradients. Recent empirical evidence demonstrates Muon's efficacy in accelerating LLM training, but its lack of integrated noise-adaptivity motivates the development of optimizers with both structural awareness and robust stochastic dynamics.

Proposed Algorithms: NAMO and Diagonal NAMO

The paper proposes two new optimization algorithms for matrix-structured parameters: NAMO (Norm-Based Adaptive Moment Estimation with Orthogonalized Momentum) and Diagonal NAMO. NAMO employs a norm-adaptive scalar stepsize for orthogonalized momentum, while Diagonal NAMO introduces neuron-wise noise-adaptive scaling via right-multiplication with a diagonal matrix subject to clamping constraints.

NAMO maintains bias-corrected estimates of the stochastic gradient (first moment) and its squared Frobenius norm (second moment) and updates parameters using:

$\Theta_{t} = \Theta_{t-1} - \eta \alpha_t O_t$

where $O_t$ is the orthogonalized momentum, $\alpha_t$ is a norm-based adaptive scalar, and $\eta$ is the learning rate.

Diagonal NAMO extends this by tracking column-wise second moment statistics and applying a clamped diagonal preconditioner:

$\Theta_t = \Theta_{t-1} - \eta O_t D_t$

where $D_t$ is a diagonal matrix of neuron-wise adaptive stepsizes, clamped to preserve conditioning and ensure robust scale-adaptivity.

Theoretical Guarantees

Both NAMO and Diagonal NAMO are analyzed under standard smoothness and bounded-variance noise assumptions. The convergence analysis yields:

• Optimal deterministic rate: $O(T^{-1/2})$ for both scalar and diagonal versions, matching lower complexity bounds for first-order nonconvex optimization.
• Adaptive stochastic convergence: $O(T^{-1/4} + \sqrt{\sigma}b^{-1/4}T^{-1/8})$ for batch size $b$, with optimal $O(T^{-1/4})$ scaling attained as $b = \Omega(\sigma^2 \sqrt{T})$.
• Unified proof strategy: Orthogonalized descent inequalities combined with matrix-aware adaptive scaling yield tight bounds robust to both drift and noise.

The analysis confirms that structured, noise-adaptive scaling atop orthogonalized momentum preserves Muon's spectral-norm steepest descent properties while stabilizing iterates under stochastic perturbations. Notably, Diagonal NAMO achieves further robustness and faster empirical convergence as it exploits block-diagonal Hessian structure common in neural networks.

Empirical Evaluation on GPT-2 Pretraining

Extensive GPT-2 pretraining experiments are presented, comparing AdamW, Muon, NAMO, and Diagonal NAMO on both small (124M) and medium (355M) models, using the OpenWebText dataset. Optimizers are tuned using grid search for learning rate (and clamping hyperparameters for Diagonal NAMO).

Learning rate sensitivity and stability: NAMO and Diagonal NAMO demonstrate accelerated convergence and reduced sensitivity to learning rate selection relative to AdamW and Muon.

Figure 1: Hyperparameter sweeping results for GPT-2 (124M). Training and validation losses at 10K steps versus learning rate.

Long-run convergence: When training for extended steps, particularly at the optimal learning rates, Diagonal NAMO achieves the lowest training and validation losses, outperforming Muon and NAMO, with the clamping parameter $c$ allowing fine-tuning of conditioning versus adaptivity.

Figure 2: Pretraining losses for GPT-2 (124M) over 50K steps under the optimal learning rates.

Scaling to larger models: On GPT-2 (355M), both NAMO and Diagonal NAMO outperform AdamW and Muon; Diagonal NAMO again provides further gains due to its neuron-wise adaptivity and conditioning control.

Figure 3: Pretraining losses for GPT-2 (355M) over 10K steps, optimal LR and clamping parameter sweep.

Final reported losses show consistent improvements for both proposed methods, with Diagonal NAMO delivering the strongest numerical performance.

Implications and Future Directions

The integration of norm-based moment estimation and matrix-aware orthogonalized momentum yields optimizers with provable and empirical superiority in both deterministic and stochastic regimes. This principled coupling addresses Muon's stochastic sensitivity and enables more robust, transfer-friendly hyperparameter protocols. The diagonal scaling aligns with the now-well-characterized block-diagonal structure of neural network Hessians, making it especially suitable for deep architectures.

Practically, the negligible additional computational overhead (for NAMO) and modest per-layer cost (for Diagonal NAMO) make these optimizers viable replacements in large-scale training pipelines, especially in LLM scenarios requiring robust adaptivity.

Theoretically, this work reinforces the utility of spectral-norm based descent in deep learning and motivates further investigation into structured adaptivity—both layerwise and neuronwise—for high-dimensional matrix optimization. The results also suggest possible avenues for developing tuning-light and further memory-efficient variants, as well as extending the framework to other architectures or distributed regimes.

Conclusion

This work provides a rigorous, principled advancement in adaptive optimization for matrix-structured models. NAMO and Diagonal NAMO combine Muon's orthogonalization with norm-based moment estimation, achieving theoretically optimal convergence rates and empirically superior performance in LLM pretraining. The diagonal extension, through neuron-wise adaptivity and conditioning, further enhances generalization and robustness. Future directions include broader empirical benchmarks, tuning-light variants, and deeper theoretical exploration of structured adaptive scaling in matrix-valued optimization.