事前学習におけるミューオンの実用的効率性

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におけるすべての誤差源を考慮しながら、リソースのオーバーヘッドを最小限に抑える単純なテレスコーピングアルゴリズムを提示した。これらの知見は、最大40億パラメータのモデルサイズにおける広範な実験と、データ分布およびアーキテクチャに関するアブレーションスタディを通じて検証された。

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.