ジェフ・ディーン: 「スパースモデル設計ガイド」を作成しましたので、ぜひご覧ください。

スパースモデルは、ディープラーニングの分野でますます重要な役割を果たしています。特定のトークンまたはサンプルに対して、モデルのごく一部だけをアクティブ化できるため、多数のパラメータを持ちながらも計算が容易になります。しかし、そのようなモデルを確実にトレーニングする方法は、解決すべき問題として残っています。この記事では、Google の Barret Zoph、Irwan Bello、William Fedus、Jeff Dean などの研究者が「効率的なスパース エキスパート モデルの設計ガイド」を提供しています。

スパース エキスパート ニューラル ネットワークは、純粋なスケールの利点を示し、現在一般的に使用されている静的ニューラル ネットワーク アーキテクチャの効果的な代替手段を提供します。スパース エキスパート ネットワークは、すべての入力に同じパラメータを適用するのではなく、各入力に使用するパラメータを動的に選択します。これにより、トークンあたりの FLOP をほぼ一定に保ちながら、ネットワークのパラメータ数を大幅に拡張できるようになります。これらの方法を採用した結果、SOTA 変換モデル、4 ～ 7 倍の事前トレーニングの高速化、および GPT-3 レベルのワンショット パフォーマンスが、トレーニング コストのわずか 3 分の 1 で実現されました。パラメータの数が膨大であるにもかかわらず、スパース モデルは大規模なニューラル ネットワークのトレーニングにかかる​​二酸化炭素排出量を 1 桁削減します。しかし、依然として困難は残っています。

Fedus ら (2021) は、スパース 1.6T パラメータ モデルが、以前の SOTA 手法 (Raffel ら、2019) と比較して 4 倍の事前トレーニングの高速化を達成したものの、SuperGLUE などの一般的なベンチマークで微調整すると、より小規模なモデルに遅れをとることを観察しました。 Artetxe et al. (2021)では、研究者らが領域外データでMoE言語モデルを微調整し、同様のギャップを観察しました。

この問題に対処するために、パラメータは少ないものの、計算フットプリントが元の 8 倍に増加し (FLOP は最大の T5 モデルとほぼ同等)、自然言語理解タスクのパフォーマンスが向上する Switch-XXL モデルが提案されました。しかし、必要な事前トレーニングは、小規模な研究ではこれまで検出されていなかったトレーニングの不安定性によって妨げられています。これらの不安定性は後に他のスパース モデルでも特定されました。これらの結果は、パラメータと計算の必要なバランスを明らかにしていますが、そのようなモデルをどのようにして確実にトレーニングするかについては未解決の問題が残っています。

この論文の目的は、スパースモデルの実用性と信頼性を向上させることです。彼らは両方の問題を研究し、設計ガイドラインを提供しました。最終的に、スパース モデルのパラメーターを 269B にスケールし、計算コストは​​ 32B の高密度エンコーダー/デコーダー トランスフォーマー (Stable and Transferable Mixture-of-Experts、ST-MoE-32B) の計算コストに匹敵しました。スパースモデルが、推論 (SuperGLUE、ARC Easy、ARC Challenge)、要約 (XSum、CNN-DM)、クローズドブック質問応答 (WebQA、Natural Questions)、敵対的構築タスク (Winogrande、ANLI R3) など、さまざまなタスクにわたって転移学習の最先端のパフォーマンスを達成したのはこれが初めてです。

この論文の貢献は次のように要約できます。

1. 安定性技術の品質と安定性のトレードオフに関する大規模な調査を実施しました。2. ルーターの z 損失を導入して、モデルの品質をわずかに向上させながら安定性の問題に対処しました。3. スパース モデルと密なモデルの微調整分析を行い、バッチ サイズと学習率に対するハイパーパラメータの感度の違いを明らかにしました。事前トレーニングによる大幅な加速にもかかわらず、ハイパーパラメータが不十分だと密なモデルの微調整によるメリットはほとんどないことがわかりました。4. 分散環境でパレート効率的なスパース モデルを設計するためのアーキテクチャ、ルーティング、およびモデル設計の原則を示しました。5. エキスパート レイヤー全体でトークン ルーティングの決定を追跡する定性分析を行いました。6. 269B のスパース モデルをトレーニングし、一連のさまざまな自然言語ベンチマークで SOTA パフォーマンスを達成しました。

ルーターのZ損失

ニューラル ネットワークを安定化するための最も効果的なアプローチの 1 つは、アクティベーションと勾配に制約を課すことです。一般的なアプローチは、爆発する勾配を補正するために、ディープ ネットワークを介してバックプロパゲーションを行うときに勾配のノルムをクリップすることです。

この論文では、研究者はメモリ効率の理由から Adafactor オプティマイザーを使用しています (ただし、最近導入された 8 ビット オプティマイザー (Dettmers ら、2021 年) の方がより良いトレードオフを提供できる可能性があります)。 Adafactor は、重みの変更が特定の基準以下にクランプされる勾配クリッピングではなく更新クリッピングを使用します。更新クリッピングをより小さな値に引き締めようとします。

次に、ルータに入るロジットの制約を調べました。ルーターはエキスパートの確率分布を float32 で計算します。しかし、研究者たちは、最大規模では、信頼できるトレーニング結果を生み出すにはこれでは不十分であることを発見しました。この問題を解決するために、ルーターの Z ロスが導入されました。

ここで、B はトークンの数、N はエキスパートの数、x ∈ RB×N はルーターに入るロジットです。

以下の表 4 は、3 回の実行で、更新クリッピングとルーター Z 損失の両方がモデルを安定化しますが、更新クリッピングはモデルの品質に重大な影響を与えることを示しています。そのため、研究者はモデルの安定性を修正するために Z 損失法を使用しました。

ルーターの z 損失は、最適化される全体的な損失の一部としての重み係数である別のハイパーパラメータ (c_z) を導入します。総損失は、クロスエントロピー損失 (L_CE)、補助負荷バランス損失 (L_B)、およびルーター Z 損失 (L_Z) の線形加重組み合わせです。

研究者は、ハイパーパラメータスイープによる事前トレーニング後の最高のモデル品質に基づいて、c_z = 0.001 の値を選択しました。付録 B には事前トレーニング中の損失が記録されています。

スパースモデルの設計

高密度モデルの設計は、Kaplan et al. (2020) の基礎研究に基づいています。しかし、スパースモデルに関しては、解決すべき問題が無数にあります。例えば、(1) 何人の専門家を使用すべきか? （２）どのルーティングアルゴリズムを使用するか？ （３）容量係数の値はいくらですか？ （4）ハードウェアはこれらの決定をどのように変えるのでしょうか？この記事で研究者らは次のように提案している。

1. 彼らの設定では、容量係数 1.25 でコアあたり最大 1 人のエキスパートによるトップ 2 ルーティングが推奨されています。2. 評価中に、容量係数を変更して新しいメモリ/コンピューティング要件に適応できます。3. 高密度レイヤー スタッキングと乗法バイアスにより品質が向上します。

詳細については原文論文を参照してください。