Adam améliore Muon : Estimation adaptative des moments avec momentum orthogonalisé
Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum
February 19, 2026
papers.authors: Minxin Zhang, Yuxuan Liu, Hayden Scheaffer
cs.AI
papers.abstract
L'optimisation stochastique efficace intègre typiquement une direction de mise à jour performante en régime déterministe avec un mécanisme s'adaptant aux perturbations stochastiques. Si Adam utilise des estimateurs de moments adaptatifs pour assurer la stabilité, Muon exploite la structure matricielle des couches de poids via un momentum orthogonalisé, démontrant des performances supérieures pour l'entraînement de grands modèles de langage. Nous proposons un nouvel optimiseur et son extension diagonale, NAMO et NAMO-D, offrant la première intégration principlée du momentum orthogonalisé avec l'adaptation au bruit de type Adam basée sur la norme. NAMO met à l'échelle le momentum orthogonalisé en utilisant un unique pas d'apprentissage adaptatif, préservant l'orthogonalité tout en surpassant Muon à un coût additionnel négligeable. NAMO-D multiplie plutôt à droite le momentum orthogonalisé par une matrice diagonale à entrées tronquées. Cette conception permet une adaptation au bruit par neurone et s'aligne avec la structure hessienne commune quasi-bloc-diagonale. Sous des hypothèses standard, nous établissons des taux de convergence optimaux pour les deux algorithmes en cadre déterministe et montrons qu'en cadre stochastique, leurs garanties de convergence s'adaptent au niveau de bruit des gradients stochastiques. Les expériences sur le pré-entraînement de modèles GPT-2 démontrent une amélioration des performances de NAMO et NAMO-D par rapport aux bases de référence AdamW et Muon, avec NAMO-D obtenant des gains supplémentaires grâce à un hyperparamètre de troncature additionnel qui équilibre les objectifs concurrents de maintien d'une direction de mise à jour bien conditionnée et d'exploitation d'une adaptation fine au bruit.
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.