Adam Migliora Muon: Stima Adattiva dei Momenti con Momento Ortogonalizzato

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

February 19, 2026

Autori: Minxin Zhang, Yuxuan Liu, Hayden Scheaffer

cs.AI

Abstract

L'ottimizzazione stocastica efficiente integra tipicamente una direzione di aggiornamento che performa bene in regime deterministico con un meccanismo di adattamento alle perturbazioni stocastiche. Mentre Adam utilizza stime adattative dei momenti per promuovere la stabilità, Muon sfrutta la struttura matriciale degli strati di pesi attraverso un momento ortogonalizzato, dimostrando prestazioni superiori nell'addestramento di grandi modelli linguistici. Proponiamo un nuovo ottimizzatore e una sua estensione diagonale, NAMO e NAMO-D, che forniscono la prima integrazione principiata del momento ortogonalizzato con l'adattamento al rumore di tipo Adam basato sulla norma. NAMO scala il momento ortogonalizzato utilizzando un singolo passo di apprendimento adattativo, preservando l'ortogonalità mentre migliora le prestazioni di Muon a un costo aggiuntivo trascurabile. NAMO-D, invece, moltiplica a destra il momento ortogonalizzato per una matrice diagonale con elementi clampati. Questo design consente un adattamento al rumore a livello di neurone e si allinea con la comune struttura quasi a blocchi diagonali dell'Hessiano. Sotto ipotesi standard, stabiliamo tassi di convergenza ottimali per entrambi gli algoritmi in ambito deterministico e dimostriamo che, in ambito stocastico, le loro garanzie di convergenza si adattano al livello di rumore dei gradienti stocastici. Esperimenti sul pre-addestramento di modelli GPT-2 dimostrano prestazioni migliorate sia di NAMO che di NAMO-D rispetto ai baseline AdamW e Muon, con NAMO-D che ottiene ulteriori vantaggi rispetto a NAMO attraverso un iperparametro di clamping aggiuntivo che bilancia gli obiettivi contrastanti di mantenere una direzione di aggiornamento ben condizionata e sfruttare un adattamento al rumore a grana fine.

English

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.