はじめに
前編では Sparsely-Gated MoE レイヤの数式と、Noisy Top-k Gating やロードバランス損失といった「ゲート側の設計」を中心に見てきました。
論文紹介：Sparsely-Gated Mixture-of-Experts Layer ～その1～
後編となる今回は、論文のもう一つの柱である
- 巨大 MoE を高速に回すための実装トリック
- 実験結果とスケーリング挙動
- その後の MoE 系 LLM への影響
について確認していきます。
参考文献
- N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. V. Le, G. E. Hinton & J. Dean. "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer." ICLR 2017
🐣 MoE を実装した最初の苦労を感じ取りたいと思います
巨大 MoE を高速に回すための実装トリック
MoE の数式そのものは前編で見た通り思ったよりはシンプルでした。論文のページ数的にも「実装と分散学習」の話の方がずっと多いです。ここでは特に重要だと思った点に絞って解説します。
Shrinking Batch 問題【主に学習時】
エキスパートが $n$ 個あって、Top-$k$ でルーティングすると、バッチサイズ $b$ に対して 1 エキスパートあたりの平均バッチサイズは以下のように概算できます。
\frac{k b}{n}
この式は、エキスパート数 $n$ を増やすほど各エキスパートに届くサンプル数が小さくなるという「Shrinking Batch」問題が発生し、GPU の計算効率が下がることを表しています。
論文では、これを緩和するために「データ並列」と「エキスパートのシャーディング」を併用します。ここで論文中の用語は model parallelism ですが、現代の LLM 文脈で言うテンソル並列というより、expert をデバイスに分割配置する方式です。今の呼び方だと expert parallel に寄った概念だと思います。
- 通常のレイヤやゲートはデータ並列
- エキスパート集合はデバイスごとに分割して保持
- 各デバイスの入力バッチを同期して MoE にまとめて流し込む
これにより、デバイス数を $d$ とすると 1 エキスパートあたりの「見かけのバッチサイズ」はほぼ $d$ 倍になります。
\frac{k b d}{n}
さらに論文では、RNN の時間方向の構造も活用します。
- LSTM の各タイムステップに同じ MoE を差し込む構成にする
- いったん前段 LSTM を全タイムステップ分まとめて計算
- その出力をまとめて MoE に投入
こうすると MoE 側から見ると「タイムステップ × バッチサイズ」を束ねた大きなミニバッチとして処理できるため、Shrinking Batch の影響をかなり抑えられます。なお、これは「逐次計算が本質の RNN で、MoE を時間方向に畳み込みっぽく使う」から成立する工夫です。
🐣 Transformer はトークン並列が前提なので悩みどころが少し違いますね
ネットワーク帯域とエキスパートのサイズ【学習時・推論時共通】
もう一つのボトルネックは GPU 間通信です。MoE では
- エキスパートを各デバイス上に常駐させる（シャーディングして置く）
- 送受信するのは MoE の入力と出力
という方針を取ります。このとき重要になるのが 1 入力あたりの「計算量 / 転送量」の比です。エキスパート内部の計算が軽すぎると、ネットワーク越しに入出力を運ぶコストの方が目立ってスループットが落ちます。そこで論文では
- エキスパートの隠れ層次元を数千に設定
- 必要に応じて層も深くする
といった形で、1 トークンあたりの計算を意識的に重くしています。計算量が増えるほど、相対的に通信の割合が下がって効率が良くなる、という発想です。
また、この論文のスケーリングは「エキスパート数を増やすときにデバイスも増やす」前提なので、1 デバイスあたりに抱えるエキスパートの量を増やさずに総パラメータ数を盛れます。計算量だけでなく、メモリもハードウェア増設に寄せて解決していくタイプの設計です。
ロードの偏りはメモリも殴る
前編で触れたロードバランス損失は「全 expert が同じくらい使われる」ための仕掛けでした。ただ、論文中でも強調されているのが
- 重要度（ゲート重みの総和）が揃っても
- 実際に割り当てられるサンプル数が偏ると
- 分散環境でメモリや性能の問題が出る
という点です。そこで論文では、importance を揃える損失に加えて「load（割当数）も揃える」損失を導入しています。前編を読んだ後だと、ここはかなり腑に落ちました。ゲートを賢くするだけでなく、実装が落ちないように注意していることがわかります。
どこに MoE を入れているのか？
論文での代表的な構成は「LSTM + MoE」です。典型的には次のようなブロックになっています。
- 単語埋め込み
- LSTM（1 層目）
- MoE レイヤ
- LSTM（2 層目）
- 出力 softmax
この論文では「既存のモデルの途中に MoE を差し込む」形で効かせています。Transformer 時代の「FFN を MoE に差し替える」流れも、発想としてはここからの自然な延長に見えます。
MoE レイヤ内部の流れ
ゲートとエキスパートのパイプラインは、概念的には次のような三段構成です。
流れを整理すると
- 入力 $x$ からゲートがスコアを計算
- Top-k で使うエキスパートを選択（学習時のみノイズを加える）
- Dispatch で該当デバイスへ入力を転送
- エキスパートで FFN を計算
- Combine で結果を集約して、ゲートの重みで足し合わせて出力
というパイプラインになっています。前編で説明したように、推論時はノイズなしの決定論的ルーティングを用います。
🐣 「ルーター → 専門家クラスタ → 結果を集約」という流れを頭に入れておくと、MoE 論文が読みやすくなります
実験結果のまとめ
この論文の実験は、細部を見るほど面白いです。ただ今回は分量を抑えるため、スケーリング挙動が見えるところだけを 1 つの表に圧縮します。数値は論文本文と Table 1〜5 の代表値から抜粋しています。
| タスク | スケール | 主な結果 |
|---|---|---|
| 1 Billion Word LM（計算量はほぼ固定） | 約 8M ops/timestep 最大 4096 experts（階層型） 各入力で 4 experts を使用 |
最大構成でテスト perplexity が計算量マッチのベースラインより 24% 低い |
| 1 Billion Word LM（計算量も増やす） | MoE 層は約 4B params 規模 高計算量モデルでベースラインと同程度の計算量 |
10 epochs 時点のテスト perplexity が 18% 低い |
| 100B Word Google News（内部コーパス） | 約 8M ops/timestep 最大 131,072 experts MoE 層は最大 137B params |
65,536 experts（68B params）まで順調に改善し、ベースラインより 39% 低い 131,072 experts では悪化 |
| WMT14 En→Fr | GNMT を軽量化し MoE を挿入 2048 experts |
newstest2014 BLEU 40.56（GNMT+RL は 39.92） |
| WMT14 En→De | 同上 | newstest2014 BLEU 26.03（GNMT は 24.91） |
| 多言語翻訳（12 言語ペア） | 1 モデルに 12 ペアを詰め込む | dev perplexity が multilingual GNMT より 19% 低い BLEU は 11/12 で multilingual GNMT を上回り、8/12 で monolingual GNMT も上回る |
表を眺めると、この論文の狙いがかなり素直に見えます。
- 1 トークンあたりの計算量を抑えたまま、容量（パラメータ数）だけを増やす
- ただし増やしすぎると破綻しうる（131,072 experts で悪化）
- 翻訳でもきちんと BLEU 改善が出る
この論文が今の MoE 系 LLM に与えた影響
論文の時点では LSTM ベースの言語モデルや機械翻訳が中心でしたが、ここで出てくる課題設定と解法は、そのまま Transformer 時代の大規模 MoE に受け継がれています。
- Noisy Top-k のようなルーティング設計
- ロードバランス損失で偏りを抑える発想
- 分散前提で expert をシャーディングして回す（expert parallel の原型）
- 「計算量を増やさず容量だけ増やす」というスケーリング思想
この後の流れとしては、たとえば GShard や Switch Transformers が Transformer の FFN を MoE 化して一気に押し広げました。さらに最近だと、ルーターが 2 experts を選ぶ構成がはっきり書かれた Mixtral of Experts のようなモデルも出てきています。
個人的には「アイディアはシンプルなのに、成立させるには分散と通信と負荷分散の現実と殴り合う」点が、今の MoE 系 LLM にもそのまま残っている気がします。
おわりに
この論文では MoE という秀逸かつシンプルなアイディアと、それを支える地道なエンジニアリングの積み重ねが印象に残りました。通信、バッチサイズ、メモリ、ロードバランスといった実装面の努力が縁の下の力持ちとなって、MoE というアイディアを世に送り出していたのですね。ではまた次の記事でお会いしましょう。
🐣 アイディアを思いつくこととそれを実現することの間にはものすごく距離があることを学びました