Adam Mejora Muon: Estimación Adaptativa del Momento con Momento Ortogonalizado

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

February 19, 2026

Autores: Minxin Zhang, Yuxuan Liu, Hayden Scheaffer

cs.AI

Resumen

La optimización estocástica eficiente generalmente integra una dirección de actualización que funciona bien en el régimen determinista con un mecanismo que se adapta a las perturbaciones estocásticas. Mientras que Adam utiliza estimaciones adaptativas de momentos para promover la estabilidad, Muon aprovecha la estructura matricial de las capas de pesos mediante momentum ortogonalizado, mostrando un rendimiento superior en el entrenamiento de modelos de lenguaje grandes. Proponemos un nuevo optimizador y una extensión diagonal, NAMO y NAMO-D, que proporcionan la primera integración fundamentada del momentum ortogonalizado con la adaptación al ruido de tipo Adam basada en normas. NAMO escala el momentum ortogonalizado utilizando un único tamaño de paso adaptativo, preservando la ortogonalidad mientras mejora a Muon con un coste adicional insignificante. En cambio, NAMO-D multiplica por la derecha el momentum ortogonalizado por una matriz diagonal con entradas limitadas. Este diseño permite una adaptación al ruido a nivel de neurona y se alinea con la estructura común del Hessiano casi diagonal por bloques. Bajo supuestos estándar, establecemos tasas de convergencia óptimas para ambos algoritmos en el entorno determinista y demostramos que, en el entorno estocástico, sus garantías de convergencia se adaptan al nivel de ruido de los gradientes estocásticos. Los experimentos de preentrenamiento con modelos GPT-2 demuestran un mejor rendimiento tanto de NAMO como de NAMO-D en comparación con los baselines AdamW y Muon, logrando NAMO-D ganancias adicionales sobre NAMO mediante un hiperparámetro de limitación adicional que equilibra los objetivos contrapuestos de mantener una dirección de actualización bien condicionada y aprovechar la adaptación al ruido a grano fino.

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.