MAI-Base-2は、2.8Tパラメータ、80B activeのMoEモデルを、16000基のGB200を用いて38日間、50Tトークンで事前学習したと報告している。 前回のMAI-Baseに続き、公開情報が限られている大規模事前学習について、技術的な詳細や実験結果、開発手法がまとめられており、参考になる。 今回も、改善ループが回せるよう、小規模モデルの実験結果から大規模モデルの性能を正確に予測するために、Source Scalingという手法を導入している。 具体的には、あるデータソースについて本番モデルで8 epoch繰り返す場合、小規模実験でも同じ8 epochを繰り返すように、使用するデータソースのサイズを調整する。 通常、小規模実験では本番学習よりも学習トークン数が少ないため、同じデータソースを使うと繰り返し回数が少なくなる。しかし、それではデータの反復による過学習などの影響を再現できず、小規模実験で得られた最適な設定が本番モデルでも最適になるとは限らない。 Source Scalingでは、学習トークン数だけでなく、ユニークなデータ量やデータの反復率も調整することで、大規模学習の挙動を小規模実験から予測しやすくしている。 もう一つの重要な結果は、学習率（LR）とWeight Decay（WD）を独立に探索するのではなく、その積に着目することで、最適なハイパーパラメータについて次の経験則を得ていることである。 LR × WD = 17.04 / S また、最適な学習率は次の関係を満たす。 log LR = -1.51 + 0.175 log(TPP) - 0.573 log S ここで、Sは学習ステップ数、TPPは学習トークン数をActive Parameter数で割った値である。 なお、LRとWDの積に着目する考え方は、既存研究でもいくつか見られ、初出ではない。 Weight Decayによる累積的な重みの縮小効果を考えると、学習率が一定であると単純化した場合、 W_{t+1} ≒ (1 - LR × WD) W_t となる。これをSステップ繰り返すと、 W_S ≒ exp(-S × LR × WD) W_0 となる。 今回の経験則では、S × LR × WD = 17.04なので、 W_S ≒ exp(-17.04) W_0 となり、学習全体を通じたWeight Decayの累積的な強さは、設定によらず一定に保つという意味を持つ。 なお、実際には学習率のスケジュールや勾配による重みの更新が存在するため、重みのスケールがこのようになるわけではない。 また、これらの係数は今回の実験設定で得られた経験値であり、異なるOptimizerやモデル構造でもそのまま成り立つとは限らない。 このほか、Activation Gradient Clippingによって学習の安定性を大きく改善したことも興味深い。 通常のGradient Clippingでは、重みに対する勾配のノルムを制限する。しかし、巨大なバッチのうち1サンプルだけが異常に大きな勾配を発生させた場合でも、バッチ全体の勾配が縮小されてしまう。 そこで、ネットワーク内の特定箇所におけるActivationの勾配をクリップする方法を採用している。 Activationに対する勾配は、Backward Passではバッチやトークンの次元が保持されているため、重み勾配として合算される前に、個別の勾配を制限できる。 これにより、異常なサンプルの影響を局所的に抑えつつ、正常なサンプルの学習信号を残すことができる。 さらに、異常に大きな勾配を発生させたサンプルを調査したところ、データの前処理のバグにより、想定しない位置でトークンが出現していることが判明し、データ前処理の修正にも役にたった。 計算効率についても注目に値する。MAI-Base-2は、総パラメータ数2.8T、活性化率約2.9%という非常に大規模かつ疎なMoEモデルであるにもかかわらず、モデル全体でMFU 21%を達成している。 GB200のピークBF16性能を2.5 PFLOPSとすると、これは1GPUあたり525 TFLOPSの有効学習演算性能に相当する。特に、MoE層単体ではMFUを24%から35%まで改善しており、全層MoEを採用できるようになったことも重要である。 コメント === 前回のMAI-Baseに引き続き、大規模モデルの事前学習について詳細な技術情報を提供している。 前回は、様々な後続タスクに関連するテキストの予測誤差を指標として、同じ損失に到達するための投入計算資源をどこまで減らせるかという、学習効率の最適化に重点が置かれていた。今回は、こうした開発手法をさらに大規模なモデルへと適用し、実際に学習を完了させた成果といえる。 特に、総パラメータ数が約3T、活性化率が3%を下回るモデルを実際に学習している点は重要であり、モデル規模としてはフロンティアモデルに近い。 今回の報告において独自の評価指標ではDeepSeek-V4-Proの事前学習などを上回る結果が得られたと報告している。今後、Post-trainingによって推論能力やエージェント性能がどの程度向上するのかが注目される。 もう一つ注目したいのは、これほど大規模な事前学習を38日間で完了させたことである。 従来、大規模モデルの事前学習には3〜4か月程度をかける例も多いが、今回は16,000基のGB200を投入し、高い疎率にもかかわらずMFU 21%という学習効率を実現していることで、1ヶ月強で終了させている。 このような短期間で終わるようになれば大規模事前学習でも、アーキテクチャ、データ、Optimizer、学習条件などを改善して再学習するサイクルをすすめらえっるようになる。