Efficacité pratique de Muon pour le pré-entraînement

Practical Efficiency of Muon for Pretraining

May 4, 2025

Auteurs: 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

Résumé

Nous démontrons que Muon, l'instanciation la plus simple d'un optimiseur du second ordre, élargit explicitement la frontière de Pareto par rapport à AdamW en termes de compromis entre temps de calcul et performance. Nous constatons que Muon est plus efficace qu'AdamW pour préserver l'efficacité des données avec des tailles de batch importantes, bien au-delà de la taille critique dite "critical batch size", tout en restant efficace sur le plan computationnel, permettant ainsi un entraînement plus économique. Nous étudions la combinaison de Muon avec la paramétrisation par mise à jour maximale (muP) pour un transfert efficace des hyperparamètres et présentons un algorithme de télescopage simple qui prend en compte toutes les sources d'erreur dans muP tout en introduisant seulement une surcharge modeste en ressources. Nous validons nos résultats par des expériences approfondies avec des modèles allant jusqu'à quatre milliards de paramètres et des ablations sur la distribution des données et l'architecture.

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.