Praktische Effizienz von Muon für das Vortraining

Practical Efficiency of Muon for Pretraining

May 4, 2025

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

Zusammenfassung

Wir zeigen, dass Muon, die einfachste Instanziierung eines Optimierers zweiter Ordnung, die Pareto-Front gegenüber AdamW im Hinblick auf den Kompromiss zwischen Rechenzeit und Leistung explizit erweitert. Wir stellen fest, dass Muon effektiver als AdamW darin ist, die Dateneffizienz bei großen Batch-Größen weit über die sogenannte kritische Batch-Größe hinaus zu bewahren, während es gleichzeitig recheneffizient bleibt und somit eine wirtschaftlichere Trainings ermöglicht. Wir untersuchen die Kombination von Muon mit der maximalen Update-Parametrisierung (muP) für einen effizienten Hyperparameter-Transfer und präsentieren einen einfachen Teleskop-Algorithmus, der alle Fehlerquellen in muP berücksichtigt, während nur ein moderater Ressourcenaufwand entsteht. Unsere Erkenntnisse validieren wir durch umfangreiche Experimente mit Modellgrößen von bis zu vier Milliarden Parametern sowie Ablationen zur Datenverteilung und Architektur.

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.