Практическая эффективность Muon для предварительного обучения

Practical Efficiency of Muon for Pretraining

May 4, 2025

Авторы: 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

Аннотация

Мы демонстрируем, что Muon, простейшая реализация оптимизатора второго порядка, явно расширяет границу Парето по сравнению с AdamW в компромиссе между вычислительными затратами и временем. Мы обнаруживаем, что Muon более эффективен, чем AdamW, в сохранении эффективности данных при больших размерах пакетов, значительно превышающих так называемый критический размер пакета, оставаясь при этом вычислительно эффективным, что позволяет проводить более экономичное обучение. Мы исследуем комбинацию Muon с параметризацией максимального обновления (muP) для эффективного переноса гиперпараметров и представляем простой телескопический алгоритм, который учитывает все источники ошибок в muP, вводя лишь умеренные дополнительные затраты ресурсов. Мы подтверждаем наши выводы с помощью обширных экспериментов с моделями размером до четырех миллиардов параметров и анализами распределения данных и архитектуры.

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.