ステップ・バイ・ステップでスケールする：大規模Mixture-of-Expertsのための計算効率的なハイパーパラメータ転移

要旨

Mixture-of-Experts（MoE）アーキテクチャは、計算コストを比例的に増加させることなくモデル容量を大幅に拡張する。しかし、モデルサイズとトークン予算の両方が極端に大きいスケールにおいて、スイープによるハイパーパラメータ（特に学習率）の最適化は計算上実行不可能である。本論文では、計算効率の高い2段階のハイパーパラメータ転移フレームワークを提案する。このフレームワークは、モデル幅のスケーリングにわたって大規模MoEモデルの訓練に最適な学習率を転移させ、その後、数兆トークンの学習期間へ外挿することで最適学習率を推定する。第一に、Multi-head Latent Attention（MLA）とMuonオプティマイザを活用するMoEアーキテクチャに対する最大更新パラメータ化（μP）の適応を定式化し、最適な学習率が幅スケーリングモデル間で一貫して転移することを実証する。第二に、予測的スケーリング法則を確立することで、この転移可能性をトークン次元に沿って拡張する。限られた予算で小規模プロキシモデルから導出された最適値に線形回帰を適用することにより、大規模な訓練期間（例：10兆トークン）に対する理想的な学習率を高い精度（R²=0.95）で外挿することに成功する。このことから、小規模モデルによるプロキシ訓練が、大規模MoEの大規模訓練に必要な最適学習率を決定するのに十分であることが示される。提案手法を基礎モデル（総パラメータ155B、活性パラメータ17B）のスクラッチからの事前学習に適用し、安定した訓練および評価結果により、フルスケールのターゲットモデルに対する最適設定が最小限のアブレーションコストで正確に予測可能であることを検証する。

English

Mixture-of-Experts (MoE) architectures significantly expand model capacity without a proportional increase in computational cost. However, optimizing their hyperparameters---particularly the learning rate---at extreme scales of both model size and token budget via sweeping remains computationally prohibitive. In this paper, we propose a compute-efficient, two-step hyperparameter transfer framework that estimates optimal learning rates for training large MoE models by transferring them across scaling model widths, and subsequently extrapolating to trillion-token horizons. First, we formulate a Maximal Update Parameterization (μP) adaptation for MoE architectures utilizing Multi-head Latent Attention (MLA) and the Muon optimizer, demonstrating that optimal learning rates transfer consistently across width-scaled models. Second, we extend this transferability along the token dimension by establishing a predictive scaling law. By applying linear regression to the optimal values derived from small proxy models on limited budgets, we successfully extrapolate the ideal learning rate to massive training horizons (e.g., 10 trillion tokens) with high fidelity (R^2=0.95). Consequently, this indicates that proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs. We apply the proposed methodology to pretrain our foundation model (155B total, 17B active parameters) from scratch, and the stable training and evaluation results validate that optimal configurations for full-scale target models can be accurately predicted with minimal ablation costs.