Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean

論文 · arXiv 1701.06538

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean

公開日 2017-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

XYZNOT で読む（全文翻訳と AI 解説）

3行まとめ

🎯 密なニューラルネットワークは、大規模データセットを吸収するためにモデル容量を拡張すると、計算コストが二乗的に爆発する課題を抱える。
🔧 分散GPUシャード全体でNoisy Top-Kゲーティングと補助的な負荷分散損失を用いるSparsely-Gated Mixture-of-Experts（MoE）層を提案。
📈 言語モデルおよび翻訳モデルを最大1370億パラメータまで拡張し、より低い計算コストで既報の最良LM Perplexityを24%〜39%改善。

概要

ニューラルネットワークが情報を吸収する能力は、そのパラメータ数によって制限されます。入力例ごとにネットワークの一部のみを活性化する条件付き計算（Conditional computation）は、計算量を比例して増やすことなくモデルの容量を飛躍的に高める手法として理論上提案されてきました。しかし実際には、アルゴリズム面および性能面で大きな課題が存在します。本研究ではこれらの課題を解決し、条件付き計算の可能性をついに現実のものとすることで、最新のGPUクラスタにおける計算効率の低下をわずかに抑えつつ、モデル容量を1000倍以上に向上させました。我々は、最大で数千のフィードフォワード・サブネットワークから構成されるSparsely-Gated Mixture-of-Experts layer（MoE）を提案します。学習可能なゲーティングネットワークが、入力例ごとに使用するこれらエキスパートのスパースな組み合わせを決定します。我々はこのMoEを、訓練コーパスに存在する膨大な知識を吸収するためにモデル容量が極めて重要となる言語モデリングおよび機械翻訳のタスクに適用しました。最大1370億パラメータを持つMoEを、積み重ねられたLSTM層の間に畳み込み的に適用するモデルアーキテクチャを提示します。大規模な言語モデリングおよび機械翻訳のベンチマークにおいて、これらのモデルはより低い計算コストで当時の最高性能（SOTA）を大幅に上回る結果を達成しました。

英語の原文概要

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

論文解説

問題設定と動機

ディープラーニングモデルは、大規模データセットで学習する際にパラメータ容量から大きな恩恵を受けます。しかし、標準的なアーキテクチャはすべての入力トークンやサンプルに対してネットワーク全体を活性化します。これにより、ネットワークサイズとデータセットの規模が拡大するにつれて、計算コストが現実的ではないレベルまで増加してしまいます。

サンプルごとにモデルのサブセットのみを活性化する条件付き計算（Conditional computation）は理論的には提案されていたものの、従来のハードウェア実装では最新のアクセラレータ上で以下の課題に直面していました：

GPUが分岐処理よりも均一な演算を好む点。
サンプルのサブセットをルーティングすることで、個々のサブネットワークが受け取るバッチサイズが極小かつ非効率になる「シュリンクバッチ問題（shrinking batch problem）」。
デバイス間のネットワーク帯域幅のボトルネック。
少数のサブネットワークに処理が集中し、他のサブネットワークが活用されないルーティングの不均衡。

主要なアイデア

Sparsely-Gated MoE層: フィードフォワード・サブネットワーク（「エキスパート」）を再帰層（スタックされたLSTMなど）の間に配置します。位置ごとにスパースな上位 $k$ 個のエキスパートのみが評価されます。
Noisy Top-K Gating: 上位 $k$ 個の選択とSoftmaxを適用する前のゲーティングロジットに調整可能なガウスノイズを加えることで、連続的で逆伝播可能な勾配を提供し、エキスパート間でのバランスの取れた探索を促進します。
データ並列とモデル並列の併用: 標準的な層とゲーティングはデバイス間でデータ並列に実行されますが、エキスパートはGPU間でモデル並列に分割されます。全デバイスからのバッチが集約されるため、各エキスパートは十分に大きなミニバッチを受け取ることができます。
負荷分散の制約: エキスパートの重要度（Importance）と負荷（Load）に関する明示的な補助損失を導入し、ゲーティングネットワークが一握りのエキスパートを過剰に利用する縮退解に陥るのを防ぎます。

手法（主要な数式を含む）

MoE層は、 $n$ 個のフィードフォワード・エキスパートネットワーク $E_{1}, \dots, E_{n}$ とゲーティングネットワーク $G (x)$ から構成されます。この層の出力は以下のようになります： $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ $G (x)_{i} = 0$ の場合、その計算は完全にスキップされます。

Noisy Top-K Gating

スパース性と微分可能なゲーティングを両立させるため、線形射影にノイズを注入します： $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty if v_{i} is in the top k elements of v otherwise$ $G (x) = Softmax (KeepTopK (H (x), k))$

負荷分散損失

均等な総ゲーティング重みを強制するため、変動係数（ $C V$ ）を用いてバッチ $X$ に対する重要度損失を定義します： $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ エキスパートごとのサンプル数を均等にするため、正規累積分布関数 $Φ$ を用いた滑らかな負荷推定量を使用します： $P (x, i) = Φ (\frac{( x \cdot W _{g} ) _{i} - kth_excluding ( H ( x ) , k , i )}{Softplus (( x \cdot W _{noise} ) _{i} )})$ $Load (X)_{i} = x \in X \sum P (x, i)$ $L_{load} (X) = w_{load} \cdot C V (Load (X))^{2}$

結果（主要な数値を含む）

10億語ベンチマーク（1 Billion Word Benchmark）: 計算量を揃えたMoE-4096-hモデルはテストPerplexity 34.1を達成しました（計算量を揃えたベースラインLSTMと比較して24%の削減）。高予算のMoEモデルは32台のK40 GPUで47時間でテストPerplexity 28.0に到達し、従来のSOTAである34.7（32台のK40 GPUで59時間）を上回りました。
1000億語コーパス: エキスパート数を65,536（680億パラメータ）までスケールさせたところ、密なベースラインと比較してテストPerplexityが39%低下し（28.9 vs. 47.0）、0.72 TFLOPS/GPUを維持しました。
機械翻訳: MoEを組み込んだGNMTは、WMT'14 En $\to$ Frで40.56 BLEU（RLなしのGNMTから+1.34 BLEU）、WMT'14 En $\to$ Deで26.03 BLEU（+1.12 BLEU）を達成しつつ、タイムステップあたりの演算量を半分未満に抑えました。
多言語翻訳: 単一の多言語MoEモデルが、12言語ペア中11言語ペアにおいて多言語GNMTベースラインを上回り（韓国語 $\to$ 英語で最大+5.84 BLEU）、開発セットのPerplexityを19%削減しました。

限界と未解決の課題

極端なスパース性は収穫逓減や性能低下につながる可能性があります。1000億語のデータセットにおいて、エキスパートを65,536から131,072に増やした際にPerplexityの悪化が見られました。
分散通信のオーバーヘッドは依然として隠れ層と入力/出力の比率に依存するため、エキスパートの容量が小さいと帯域幅効率が低下します。
RNNの再帰ステップ内でMoEを反復的に適用すると、シーケンス位置にまたがるバッチ展開の最適化が破綻するため、アクティベーション再計算スキームが必要になります。

重要用語

Sparsely-Gated Mixture-of-Experts Layer (MoE): 複数の並列なエキスパート・サブネットワークを含み、学習可能なゲーティングネットワークが各入力サンプルに対してスパースなサブセットのエキスパートのみを選択するニューラルネットワーク層。
Conditional Computation (条件付き計算): 計算量を比例して増やすことなくパラメータ容量を増加させるために、入力例ごとにニューラルネットワークの一部を動的に活性化または非活性化するアーキテクチャパラダイム。
Noisy Top-K Gating: 上位k個の値を保持してSoftmaxを適用する前にロジットへ調整可能なガウスノイズを注入することで、逆伝播を可能にし、ルーティングの崩壊を防ぐゲーティングメカニズム。
Shrinking Batch Problem (シュリンクバッチ問題): 入力が多数のエキスパートにスパースに分割されることで個々のサブネットワークが受け取る実効バッチサイズが縮小し、ハードウェアの演算効率が低下する問題。
Importance Loss (重要度損失): エキスパートの均等な割り当てを促進するために、バッチ全体におけるエキスパートごとのゲート値の合計の二乗変動係数にペナルティを課す補助損失。
Load Loss (負荷損失): ハードウェアの負荷を均衡させるために、各エキスパートへのトークンの推定確率重み付き割り当て数の分散を最小化する補助目的関数。
Hierarchical Mixture of Experts (階層型Mixture of Experts): 分岐係数と分散通信のオーバーヘッドを削減するために、一次ゲーティングネットワークが入力を二次MoEクラスタにルーティングする2段階のMoE構造。

よくある質問

ゲーティングネットワークはどのようにして離散的なTop-kルーティングを通じた逆伝播を可能にしているのですか？

Top-kの選択がSoftmaxの前に行われるため、選択されたエキスパートに対する非ゼロのゲート値は入力重みに関して連続的な導関数を保持し、勾配がゲーティングネットワークを通り逆伝播することを可能にしています。

分散デバイス間でシュリンクバッチ問題はどのように解決されていますか？

モデルはデータ並列とモデル並列を組み合わせています。標準的な層はデバイス間で独立したミニバッチを処理しますが、異なるGPUに配置された定常エキスパートへルーティングされる前に、すべてのデバイスからのトークンが一括して集約されます。

ゲーティングネットワークになぜノイズが必要なのですか？

注入されたガウスノイズは、学習初期におけるエキスパート間の探索を促進するとともに、負荷分散損失の滑らかな推定量を提供し、エキスパートの利用が一部に偏る自己強化的な崩壊を防ぎます。

モデルを131,072個のエキスパートまでスケールさせると何が起きますか？

Perplexityは28.9（65,536エキスパート時）から29.2（131,072エキスパート時）へとわずかに悪化し、計算効率は0.30 TFLOPS/GPUまで低下しました。これは極端なスパース性と、スケールされずに固定されたバッチサイズが原因と考えられます。

このページの解説は AI が論文から作成したもので、誤りを含む場合があります。詳細は原文をご確認ください。

XYZNOT で読む（全文翻訳と AI 解説）