Eficiencia Práctica de Muon para el Pretrenamiento

Practical Efficiency of Muon for Pretraining

May 4, 2025

Autores: Essential AI, Ishaan Shah, Anthony M. Polloreno, Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh J Shah, Khoi Nguyen, Kurt Smith, Michael Callahan, Michael Pust, Mohit Parmar, Peter Rushton, Platon Mazarakis, Ritvik Kapila, Saurabh Srivastava, Somanshu Singla, Tim Romanski, Yash Vanjani, Ashish Vaswani

cs.AI

Resumen

Demostramos que Muon, la instanciación más simple de un optimizador de segundo orden, expande explícitamente la frontera de Pareto sobre AdamW en el equilibrio entre tiempo de cómputo y rendimiento. Encontramos que Muon es más efectivo que AdamW para mantener la eficiencia de datos con tamaños de lote grandes, mucho más allá del llamado tamaño crítico de lote, mientras mantiene la eficiencia computacional, permitiendo así un entrenamiento más económico. Estudiamos la combinación de Muon con la parametrización de actualización máxima (muP) para la transferencia eficiente de hiperparámetros y presentamos un algoritmo telescópico simple que contabiliza todas las fuentes de error en muP mientras introduce solo un modesto incremento en el uso de recursos. Validamos nuestros hallazgos mediante experimentos extensos con tamaños de modelo de hasta cuatro mil millones de parámetros y ablaciones sobre la distribución de datos y la arquitectura.

English

We demonstrate that Muon, the simplest instantiation of a second-order optimizer, explicitly expands the Pareto frontier over AdamW on the compute-time tradeoff. We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes, far beyond the so-called critical batch size, while remaining computationally efficient, thus enabling more economical training. We study the combination of Muon and the maximal update parameterization (muP) for efficient hyperparameter transfer and present a simple telescoping algorithm that accounts for all sources of error in muP while introducing only a modest overhead in resources. We validate our findings through extensive experiments with model sizes up to four billion parameters and ablations on the data distribution and architecture.