LoginSignup
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@hitomatagi

ST-MoE論文

Posted at

論文

  • タイトル：ST-MOE: DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS
  • 原文：arXiv:2202.08906
  • 著者：Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, William Fedus

この論文を読んだきっかけ

業務の実プロジェクトで、ロングコンテキストLLMを利用する機会があり、モデルの性能を把握するための調査。

免責事項

この投稿は、個人的なメモです。内容理解や翻訳には、生成AIを利用しています。より詳細な内容や正確な内容を知りたい方は、論文の原文を参照してください。

Abstract

ST-MoE論文のABSTRACTは、大規模言語モデルにおけるスケールと効率性のトレードオフ、そしてその解決策としてのMixture-of-Experts (MoE) の可能性と課題について述べています。

まず、大規模化が自然言語処理の新たなフロンティアを開拓した一方で、高コストという問題が生じていることを指摘しています。その解決策として、エネルギー効率の高い大規模言語モデルを実現する手法としてMoEとSwitch Transformerが提案されてきました。

しかし、これらの手法は、学習の不安定性とファインチューニング時の品質の不確実性という課題を抱えており、幅広い自然言語タスクにおいて最先端技術を進歩させる妨げになってきました。

そこで、本論文ではこれらの課題に焦点を当て、解決策を探求しています。具体的には、スパースモデルを2690億パラメータまでスケールさせ、320億のDense Encoder-Decoder Transformerに匹敵する計算コストでStable and Transferable Mixture-of-Experts (ST-MoE-32B) を実現しました。

そして、スパースモデルとして初めて、推論（SuperGLUE、ARC Easy、ARC Challenge）、要約（XSum、CNN-DM）、クローズドブック質問応答（WebQA、Natural Questions）、敵対的に構築されたタスク（Winogrande、ANLI R3）など、多様なタスクにおいて転移学習で最先端の性能を達成したことを報告しています。

つまり、本論文は、MoEの大規模言語モデルにおける実用性と信頼性を高めることを目的とし、学習の安定化、ファインチューニングの改善、モデル設計の指針、そして具体的な成果としてST-MoE-32Bの性能について述べています。

1 Introduction

ST-MoE論文の1 Introductionでは、スパースニューラルネットワーク、特にMixture-of-Experts (MoE) の利点と課題について概説し、本研究の目的と貢献を明確にしています。

まず、従来の密なニューラルネットワークアーキテクチャと比較して、MoEを含むスパースニューラルネットワークが、大規模化と計算効率の両立を実現する有望な手法であることを説明しています。

MoEの利点：

  • スケーラビリティ: パラメータ数を大幅に増やしつつ、トークンあたりのFLOPsをほぼ一定に保つことができるため、巨大なモデルの構築が可能になります。
  • 効率性: 入力に応じて動的にパラメータを選択するため、従来の密なモデルに比べて計算コストとエネルギー消費を削減できます。
  • 環境負荷の低減: 学習に必要な計算リソースを削減することで、大規模ニューラルネットワークの学習に伴う炭素排出量を大幅に削減できます。

MoEの課題:

  • 学習の不安定性: 従来のTransformerよりも学習が不安定になりやすく、特に大規模モデルでは顕著です。
  • ファインチューニング時の品質: 事前学習では優れた性能を示す一方で、ファインチューニング時に小規模なモデルに劣ることがあります。

これらの課題を踏まえ、本研究ではMoEの実用性と信頼性を向上させることを目標としています。具体的には、以下の貢献を挙げています。

  1. 安定化手法のトレードオフ分析: モデルの安定性と品質のトレードオフについて、大規模な実験を通して分析し、それぞれの技術の影響を明らかにしています。
  2. router z-lossの導入: 学習の不安定性を解消し、モデルの品質を向上させる新しい補助損失関数router z-lossを提案しています。
  3. スパースモデルと密なモデルのファインチューニング比較: バッチサイズや学習率に対する感度が異なることを示し、スパースモデルに適したファインチューニングプロトコルを提案しています。
  4. 分散環境における効率的なスパースモデル設計: 分散環境における計算効率を最大化する、アーキテクチャ、ルーティング、モデル設計の原則を提案しています。
  5. トークンルーティングの可視化分析: トークンがエキスパート層にどのようにルーティングされるかを可視化し、エキスパートの専門化について分析しています。
  6. 2690億パラメータのスパースモデルST-MoE-32B: 多様な自然言語タスクにおいて最先端の性能を達成した、安定性と転移性に優れた2690億パラメータのスパースモデルST-MoE-32Bを開発しています。

本論文は、MoEの大規模言語モデルにおける実用化に向けた重要な課題に取り組み、具体的な解決策と設計指針を提示することで、今後の研究開発を促進することを目指しています。

2 Background

ST-MoE論文の2 Backgroundでは、スパースエキスパートモデル、特に自然言語処理におけるMixture-of-Experts (MoE) の基本的な概念と仕組みについて解説しています。

まず、スパースエキスパートモデルは、ニューラルネットワークの一部の層を複数のエキスパートに置き換え、入力に応じて適切なエキスパートを選択して処理を行うことで、計算効率を向上させる手法であると説明しています。

エキスパート: それぞれ固有の重みを持つ、独立して学習されたニューラルネットワークです。
ルーター (ゲーティングネットワーク): 各トークンがどのエキスパートに送られるかを決定するネットワークです。
トップnルーティング: 各トークンが上位n個のエキスパートにルーティングされるアルゴリズムです。
負荷分散損失: トークンがエキスパート間で均等に分散されるように促す補助的な損失関数です。
グループサイズ: グローバルバッチサイズを分割した、より小さなグループのサイズです。各グループは、エキスパート間での負荷分散のために独立して処理されます。
容量係数 (CF): 各エキスパートが処理できるトークンの最大数を調整するハイパーパラメータです。
FFN: Feed Forward Networkの略で、Transformerの層の一つです。
エンコーダ・デコーダ: 本論文で使用されるTransformerのアーキテクチャです。エンコーダは入力全体を処理し、デコーダはエンコーダの出力を参照しながら自己回帰的に出力を生成します。

MoEでは、入力トークン表現xに対して、ルーターが各エキスパートへのルーティング確率を計算します。この確率は、通常softmax関数によって正規化されます。トークンは、ルーティング確率に基づいて選択されたエキスパートに送られ、各エキスパートの出力が重み付け加算されて最終的な出力となります。

本論文では、ハードウェアの利用効率を高めるために、各エキスパートのバッチサイズを固定しています。エキスパートの容量を超えたトークンは処理されずに次の層に渡されます。負荷分散損失は、すべてのアクセラレータが並列に十分な量のデータを処理できるように、トークンがエキスパート間で均等に分散されるように促す役割を果たします。

2 Backgroundでは、これらの基本的な概念を説明することで、読者がMoEの仕組みを理解し、以降のセクションで説明される技術や分析をスムーズに理解できるようにしています。

3 STABILIZING TRAINING OF SPARSE MODELS

このセクションでは、スパースモデルの学習における不安定性の問題を取り上げています。スパースモデル、特にMoEは、通常のTransformerよりも学習が不安定になりやすく、その不安定性は図1に示されるように、学習損失の急激な発散として現れます。

論文では、この不安定性を解消するための様々なアプローチが検討されていますが、多くの場合、安定性の向上と引き換えにモデルの品質が低下するというトレードオフが存在することが指摘されています。例えば、学習率を極端に小さくしたり、勾配クリッピングを厳しくすると安定性は向上しますが、モデルの性能は低下します。

そこで、論文では安定性の向上と品質の維持を両立させるための手法を探求しています。具体的には、以下の3つのカテゴリーに分類される安定化手法を検証しています。

  1. 乗算的な相互作用の除去: Transformerにおける特定の乗算的な演算が不安定性を招く可能性があるため、それらの演算を除去または変更することで安定性を向上させることを試みます。
  2. モデルノイズの注入: モデルにノイズを注入することで、過学習を防ぎ、安定性を向上させることを試みます。
  3. 活性化と勾配の制約: 活性化関数や勾配の値を制限することで、学習の不安定性を抑制することを試みます。

重要なポイント:

  • スパースモデルの学習の不安定性は、大規模言語モデルの開発における大きな課題です。
  • 安定化手法には、モデルの品質とのトレードオフが存在することが多く、両立させることが重要です。
  • 論文では、乗算的な相互作用の除去、モデルノイズの注入、活性化と勾配の制約という3つのカテゴリーに分類される安定化手法を検証しています。

3.1 Stability and Quality Tradeoffs when Removing Multiplicative Interactions

このセクションでは、Transformer モデルのアーキテクチャにおける乗算的な相互作用に着目し、それらがスパースモデルの学習安定性に与える影響を分析しています。具体的には、GEGLU と RMS正規化のスケールパラメータ という2つの要素を取り上げ、それらを削除することで安定性が向上する一方で、モデルの品質が低下するというトレードオフについて議論しています。

1. GEGLU (Gated Linear Unit)

  • GEGLUは、2つの線形変換の出力を要素ごとに掛け合わせる活性化関数です。一方の線形変換出力はsigmoid関数を通過させることでゲートの役割を果たします。
  • 通常のReLUを用いたFFN層をGEGLU-Linear FFN層に置き換えることで、モデルの品質が向上することが報告されています (Narang et al., 2021)。
  • しかし、GEGLUは加算よりも乗算の要素が多く、これが学習の不安定性を招く可能性があります。

2. RMS正規化のスケールパラメータ

  • RMS正規化は、入力ベクトルをその要素ごとの二乗平均平方根でスケールする手法です。
  • RMS正規化の後には、学習可能なスケールパラメータgによる要素ごとのスケーリングが行われます。
  • このスケールパラメータgによる乗算が、学習の不安定性を悪化させる可能性があります。

実験結果:

  • GEGLU層を削除するか、RMS正規化のスケールパラメータを削除すると、安定性が向上することが確認されました。
  • しかし、これらの変更はモデルの品質を大幅に低下させることも分かりました (表2参照)。
  • 特に、スケールパラメータgは、他のパラメータに比べてモデルの品質に与える影響が大きいことが分かりました。

考察:

  • 乗算的な相互作用は、モデルの表現力を高め、品質向上に貢献する可能性があります。
  • しかし、同時に学習の不安定性を招く要因にもなり得ます。
  • スパースモデルの学習安定性と品質のバランスを取るためには、乗算的な相互作用の利用を慎重に検討する必要があります。

結論:

このセクションでは、スパースモデルの学習安定性と品質のトレードオフについて、乗算的な相互作用に着目して分析を行いました。GEGLUやRMS正規化のスケールパラメータなどの乗算的な要素は、モデルの品質向上に貢献する一方で、学習の不安定性を招く可能性があります。スパースモデルの設計においては、これらのトレードオフを考慮し、適切なアーキテクチャを選択することが重要です。

3.2 Stability and Quality Tradeoffs when Adding Noise

このセクションでは、スパースモデルの学習安定性を向上させるために、ノイズを注入するというアプローチの効果と、それがモデルの品質に与える影響について検証しています。

背景:

  • ノイズの注入は、ニューラルネットワークの学習安定性を向上させる一般的な手法として知られています (Neelakantan et al., 2015)。
  • 特に、ドロップアウトはファインチューニング時にノイズを注入する効果的な手法であり、スパースモデルの安定性向上にも効果が期待されます。

検証内容:

  • 入力ジッター: ルーターへの入力logitsに、[1-10^-2, 1+10^-2] の範囲の一様乱数をかける手法 (Fedus et al., 2021)。
  • ドロップアウト: Transformer全体にドロップアウトを適用する手法。

実験結果:

  • 入力ジッターとドロップアウトの両方とも、ベースラインと比較して安定性が向上することが確認されました (表3参照)。
  • しかし、どちらの手法もモデルの品質を低下させることが分かりました。
  • 特に、XL規模のモデルでは、入力ジッターによる品質低下が顕著でした。

考察:

  • ノイズの注入は、過学習を防ぎ、学習安定性を向上させる効果がありますが、同時にモデルの表現力を低下させる可能性があります。
  • 小規模なモデルでは有効な手法であっても、大規模なモデルでは効果が薄れたり、逆効果になる可能性があります (Kaplan et al., 2020)。
  • スパースモデルの学習安定性と品質のバランスを取るためには、ノイズの注入量を慎重に調整する必要があります。

結論:

このセクションでは、スパースモデルの学習安定性と品質のトレードオフについて、ノイズ注入に着目して分析を行いました。入力ジッターやドロップアウトなどのノイズ注入は、安定性向上に貢献する一方で、モデルの品質を低下させる可能性があります。スパースモデルの設計においては、これらのトレードオフを考慮し、適切なノイズ注入手法と注入量を選択することが重要です。

付記:

  • 本論文では、ドロップアウトをTransformer全体に適用していますが、エキスパートの隠れ状態にのみ適用するエキスパートドロップアウト (Fedus et al., 2021) という手法も存在します。
  • エキスパートドロップアウトは、品質低下を抑えつつ安定性を向上させる効果的な手法であることが報告されていますが、本論文ではXL規模のモデルでは効果が限定的でした。

3.3 Stability and Quality Tradeoffs when Constraining Activations and Gradients

このセクションでは、スパースモデルの学習安定性を向上させるために、活性化関数と勾配を制約するアプローチの効果と、それがモデルの品質に与える影響について検証しています。

背景:

  • 活性化関数や勾配の値を制限することは、ニューラルネットワークの学習安定性を向上させるための一般的な手法として知られています (Pascanu et al., 2013; Ioffe and Szegedy, 2015; Salimans and Kingma, 2016; Ba et al., 2016)。
  • 特に、勾配爆発を防ぐために勾配ノルムをクリッピングする手法は広く用いられています。

検証内容:

  • 更新クリッピング: Adafactorオプティマイザにおける更新のクリッピング値を小さくする手法。Adafactorはメモリ効率の良さから採用されていますが、勾配クリッピングの代わりに更新クリッピングを用いています。
  • router z-loss: ルーターへの入力logitsに制約を加える新しい補助損失関数。これは、Mesh Tensorflowコードベースで使用されているz-loss (Shazeer et al., 2018) を応用したものです。router z-lossは、logitsの値が大きくなりすぎるのを防ぎ、数値精度を向上させる効果があります。

router z-loss:

router z-lossは、以下のように定義されます。

Lz(x) = (1/B) * Σ_{i=1}^{B} log Σ_{j=1}^{N} (p_j(x))^2

ここで、Bはトークン数、Nはエキスパート数、xはルーターへの入力logitsです。

実験結果:

  • 更新クリッピングとrouter z-lossの両方とも、3回の実行すべてにおいてモデルを安定化させることに成功しました (表4参照)。
  • しかし、更新クリッピングはモデルの品質を大幅に低下させることが分かりました。
  • 一方、router z-lossは品質を低下させることなく、むしろわずかに向上させることが分かりました。

考察:

  • 更新クリッピングは勾配の値を直接制限するため、学習の進捗を阻害する可能性があります。
  • router z-lossは、logitsの値を制限することで数値精度を向上させ、学習の安定化に貢献します。
  • router z-lossは、品質を低下させることなく安定性を向上させる効果的な手法であると言えます。

結論:

このセクションでは、スパースモデルの学習安定性と品質のトレードオフについて、活性化関数と勾配の制約に着目して分析を行いました。更新クリッピングは安定性向上に貢献する一方で、モデルの品質を低下させる可能性があります。一方、router z-lossは品質を維持したまま安定性を向上させる効果的な手法であることが分かりました。

付記:

  • router z-lossは、新しいハイパーパラメータczを導入します。czは、router z-lossの重みを調整する係数です。
  • 本論文では、cz = 0.001が最適な値として採用されています。

3.4 Selecting a Precision Format: Trading Efficiency and Stability

このセクションでは、スパースモデルの学習における数値精度の選択が、モデルの効率性と安定性に与える影響について議論しています。

背景:

  • 近年の大規模Transformerモデルの学習では、混合精度演算 (Micikevicius et al., 2017) が広く用いられています。
  • 混合精度演算では、重みはfloat32で保存され、勾配更新に使用されますが、行列乗算などの演算はbfloat16で行われます。
  • bfloat16はfloat32よりも精度が低いため、メモリ使用量と計算コストを削減できますが、丸め誤差が大きくなる可能性があります。

数値精度と丸め誤差:

  • bfloat16はfloat32よりも仮数部が短いため、表現できる数値の精度が低くなります。
  • その結果、bfloat16ではfloat32よりも丸め誤差が大きくなり、学習の不安定性を招く可能性があります。
  • 特に、指数関数や三角関数などの演算では、丸め誤差の影響が大きくなります。

スパースモデルにおける数値精度:

  • スパースモデルは、ルーターによってトークンがエキスパートに振り分けられるため、指数関数を用いたsoftmax演算が頻繁に行われます。
  • そのため、スパースモデルは丸め誤差の影響を受けやすく、数値精度の選択が重要になります。

router z-lossと数値精度:

  • router z-lossは、ルーターへの入力logitsの値を小さくすることで、丸め誤差の影響を軽減する効果があります。
  • router z-lossを用いることで、bfloat16のような低精度演算を用いても安定した学習が可能になります。

実験結果:

  • 本論文では、ST-MoE-32Bモデルの学習において、allreduce通信の精度をbfloat16にすると高速化できることが分かりました。
  • しかし、bfloat16では学習が不安定になる場合もあるため、最終的にはfloat32を採用しています。

考察:

  • スパースモデルの学習においては、数値精度の選択が効率性と安定性の両方に影響を与える重要な要素です。
  • router z-lossを用いることで、低精度演算を用いても安定した学習が可能になります。
  • 今後の研究では、スパースモデルに適した新しい数値精度フォーマットの開発が期待されます。

結論:

このセクションでは、スパースモデルの学習における数値精度の選択について、効率性と安定性のトレードオフに着目して議論しました。bfloat16のような低精度演算は効率的ですが、丸め誤差の影響を受けやすいという問題があります。router z-lossを用いることで、低精度演算を用いても安定した学習が可能になります。

4 FINE-TUNING PERFORMANCE OF SPARSE MODELS

高性能な言語モデルは、一般的に以下の2段階の手順で開発されます。

  1. 事前学習 (Pre-training): インターネット上の膨大なデータを用いてモデルを学習させます。
  2. ファインチューニング (Fine-tuning): 特定のタスク (例: SuperGLUE) に合わせてモデルを調整します。

近年、Few-shot推論 (Brown et al., 2020)、Prefix Tuning (Li and Liang, 2021)、Prompt Tuning (Lester et al., 2021)、アダプターモジュール (Houlsby et al., 2019) などの新しい手法が登場し、ファインチューニングの代替手段として注目されています。

しかし、これらの手法は、まだファインチューニングほどの高い性能を実現できていません。

そのため、本研究ではファインチューニングに焦点を当てています。ただし、スパースモデルがFew-shot設定で成功を収めている事例 (Du et al., 2021; Artetxe et al., 2021) も紹介します。

また、強化学習を用いて大規模言語モデルを適応させる手法 (Ouyang et al., 2022) は今後の課題として残されています。

要約:

  • ファインチューニングは、現在でも高性能な言語モデルを開発するための主要な手法です。
  • 新しい手法は有望ですが、まだファインチューニングに匹敵する性能は達成できていません。
  • スパースモデルはFew-shot設定でも有望な結果を示しています。
  • 強化学習を用いた言語モデルの適応は、今後の研究課題です。

4.1 Hypothesis: A Generalization Problem

このセクションでは、スパースモデルが過学習しやすいという仮説を立て、その根拠となる実験結果を示しています。

背景:

  • スパースモデルは、大規模なデータセットを用いた学習では優れた性能を示すことが知られています。
  • しかし、ファインチューニング時には、小規模なモデルに劣ることが報告されています (Fedus et al., 2021; Artetxe et al., 2021)。

仮説:

  • スパースモデルは、その構造上、過学習しやすい傾向があります。

検証内容:

  • SuperGLUEベンチマーク (Wang et al., 2019) の2つのタスク、Commitment Bank (CB) とReCORDを用いて、スパースモデルと密なモデルのファインチューニング特性を比較しました。
  • CBは学習データが250サンプルと少なく、ReCORDは10万サンプル以上と多いという特徴があります。
  • これらのタスクを用いることで、データセットのサイズが過学習に与える影響を分析できます。

実験結果:

  • 図3に示すように、スパースモデルは学習データに対する精度が急速に100%に達し、データ分布のシフトにうまく適応していることが分かります。
  • ReCORDのように学習データが多いタスクでは、スパースモデルは密なモデルよりも高い汎化性能を示しました。
  • しかし、CBのように学習データが少ないタスクでは、スパースモデルは密なモデルに比べて汎化性能が劣りました。

考察:

  • スパースモデルは、学習データに過剰に適合しやすく、未知のデータに対する汎化性能が低下する可能性があります。
  • データセットのサイズが小さいほど、過学習の影響が大きくなります。

結論:

このセクションでは、スパースモデルが過学習しやすいという仮説を立て、その根拠となる実験結果を示しました。スパースモデルの設計においては、過学習を防ぐための対策を講じることが重要です。

付記:

  • 過学習を防ぐための対策としては、正則化 (regularization) が有効です。
  • 正則化には、ドロップアウト、重み減衰、早期終了など、様々な手法があります。
  • スパースモデルに適した正則化手法の開発は、今後の研究課題です。

4.2 Fine-Tuning a Subset of Model Parameters to Improve Generalization

このセクションでは、スパースモデルの汎化性能を向上させるために、ファインチューニング時にモデルパラメータの一部のみを更新するアプローチの効果を検証しています。

背景:

  • 4.1で示されたように、スパースモデルは過学習しやすい傾向があります。
  • 過学習を防ぐためには、モデルの複雑さを抑制する必要があります。
  • ファインチューニング時に更新するパラメータを制限することで、モデルの複雑さを抑制し、汎化性能を向上させることが期待されます。

検証内容:

  • ST-MoE-Lモデルを用いて、ファインチューニング時に更新するパラメータを以下の5つのパターンに分け、SuperGLUEベンチマークにおける性能を比較しました。

    1. All: すべてのパラメータを更新
    2. Non MoE: MoE層以外のパラメータを更新
    3. MoE: MoE層のパラメータのみを更新
    4. Attention: Self-attentionとEncoder-Decoder attentionのパラメータのみを更新
    5. FFN: MoE層以外のFFN層のパラメータのみを更新

実験結果:

  • 図5に示すように、MoE層のパラメータのみを更新した場合 (MoE)、他のパターンと比べて性能が大幅に低下しました。
  • それ以外の4つのパターン (All, Non MoE, Attention, FFN) では、ほぼ同等の性能が得られました。

考察:

  • MoE層はモデルパラメータの大部分を占めていますが、各トークンは1層あたり最大2つのエキスパートしか通過しないため、MoE層のパラメータのみを更新しても、モデル全体への影響が限定的になります。
  • Non MoEパラメータのみを更新することで、ファインチューニングの高速化とメモリ使用量の削減が可能になります。

結論:

このセクションでは、スパースモデルの汎化性能を向上させるために、ファインチューニング時に更新するパラメータを制限するアプローチの効果を検証しました。MoE層以外のパラメータのみを更新することで、性能を維持したまま、ファインチューニングの効率化が可能になることが分かりました。

付記:

  • 本論文では、エキスパート全体を確率的にマスクするドロップアウトの変種も試みていますが、汎化性能の向上には至りませんでした。
  • 他の実験結果については、付録Jを参照してください。

4.3 SPARSE AND DENSE MODELS REQUIRE DIFFERENT FINE-TUNING PROTOCOLS

このセクションでは、スパースモデルと密なモデルでは、最適なファインチューニングプロトコルが異なることを示し、具体的にバッチサイズと学習率の影響について分析しています。

背景:

  • ファインチューニングは、事前学習済みモデルを特定のタスクに適応させるための重要なプロセスです。
  • ファインチューニングのプロトコル (バッチサイズ、学習率など) は、モデルの性能に大きな影響を与えます。

検証内容:

  • Dense-LモデルとST-MoE-Lモデルを5000億トークンのC4データセットで事前学習し、SuperGLUEベンチマークでファインチューニングを行いました。
  • バッチサイズと学習率を変化させて、それぞれのモデルの性能を比較しました。

実験結果:

  • 図6に示すように、スパースモデルと密なモデルでは、最適なバッチサイズと学習率が異なることが分かりました。
  • スパースモデル: 小さなバッチサイズと高い学習率で良い性能を示しました。
  • 密なモデル: 大きなバッチサイズと低い学習率で良い性能を示しました。

考察:

  • 4.1で示されたように、スパースモデルは過学習しやすい傾向があります。
  • 小さなバッチサイズと高い学習率は、ファインチューニングプロセスにノイズを注入し、過学習を防ぐ効果があります。
  • 密なモデルは、スパースモデルよりも過学習しにくい傾向があるため、大きなバッチサイズと低い学習率で安定した学習が可能になります。

結論:

このセクションでは、スパースモデルと密なモデルでは、最適なファインチューニングプロトコルが異なることを示しました。スパースモデルのファインチューニングでは、小さなバッチサイズと高い学習率を用いることが重要です。

重要なポイント:

  • スパースモデルと密なモデルのファインチューニングには、異なるアプローチが必要です。
  • スパースモデルのファインチューニングでは、バッチサイズと学習率を適切に調整することが重要です。
  • 適切なファインチューニングプロトコルを用いることで、スパースモデルの潜在能力を引き出すことができます。

付記:

  • 表20 (付録F) には、すべての実験結果の詳細が記載されています。

4.4 SPARSE MODELS ARE ROBUST TO DROPPED TOKENS DURING FINE-TUNING

このセクションでは、スパースモデルがファインチューニング中にトークンがドロップされても堅牢であるという、驚くべき特性について検証しています。

背景:

  • スパースモデルでは、各層でトークンが1つ以上のエキスパートにルーティングされます。
  • ハードウェアの効率性を高めるために、各エキスパートが処理できるトークン数 (エキスパート容量) は事前に固定されています。
  • エキスパート容量を超えたトークンはドロップされ、計算処理が行われません。

検証内容:

  • ドロップされるトークン数を制御するために、以下の2つの要素を調整しました。

    1. 容量係数 (CF): エキスパート容量を調整するハイパーパラメータ。CFを大きくすると、ドロップされるトークン数が減少します。
    2. 補助損失: トークンがエキスパート間で均等に分散されるように促す損失関数。ファインチューニング中にこの損失関数を無効化することで、ドロップされるトークン数を増加させました。

  • ST-MoE-Lモデルを用いて、SuperGLUEベンチマークにおける性能を測定しました。

実験結果:

  • 表5に示すように、最大10-15%のトークンがドロップされても、ファインチューニングの品質に大きな影響はありませんでした。
  • 驚くべきことに、1%未満のトークンしかドロップされないモデルと、10-15%のトークンがドロップされるモデルで、ほぼ同等の性能が得られました。
  • また、補助損失を有効化すると、ファインチューニングの品質が向上することが分かりました。

考察:

  • スパースモデルは、トークンのドロップに対して非常に堅牢であることが分かりました。
  • これは、スパースモデルが冗長な表現を学習しているためと考えられます。
  • つまり、一部のトークンがドロップされても、他のトークンから必要な情報を得ることができるということです。
  • 補助損失は、トークンをエキスパート間で均等に分散させることで、冗長な表現の学習を促進する効果があります。

結論:

このセクションでは、スパースモデルがファインチューニング中にトークンがドロップされても堅牢であることを示しました。この特性は、スパースモデルの実用性を高める上で重要な要素となります。

付記:

  • トークンのドロップは、一種の正則化として機能している可能性があります。
  • ドロップされるトークンの割合を変化させることで、モデルの汎化性能を制御できる可能性があります。
  • この点については、今後の研究でさらに詳しく調査する必要があります。

4.5 INSERTING SENTINELS TOKENS DURING FINE-TUNING

このセクションでは、ファインチューニング時にsentinelトークンを挿入することで、事前学習とファインチューニングのドメインギャップを縮小し、性能向上を目指すアプローチについて検証しています。

背景:

  • 事前学習では、多くの場合、span corruption objective (Fedus et al., 2018; Devlin et al., 2018) が用いられます。
  • span corruption objectiveでは、入力文から連続したトークンをマスクし、その部分を予測するようにモデルを学習させます。
  • 一方、ファインチューニングでは、様々なタスク (例: 文書分類、質問応答) に合わせてモデルを調整します。
  • 事前学習とファインチューニングでタスクが異なるため、モデルが入力として受け取るデータの分布が異なり、ドメインギャップが生じます。

sentinelトークン:

  • span corruption objectiveでは、マスクされた部分をsentinelトークンで置き換えます。
  • sentinelトークンは、モデルにとって特別な意味を持つトークンであり、「この部分を予測する」という指示を表します。

検証内容:

  • ファインチューニング時に、入力文にsentinelトークンを挿入することで、事前学習とファインチューニングのドメインギャップを縮小できるのではないかと考えました。
  • 具体的には、表6に示すように、ファインチューニングの入力文にsentinelトークンを追加しました。
  • SuperGLUEベンチマークとGrammar Error Correction (GEC) タスク (Rothe et al., 2021) を用いて、sentinelトークンの効果を検証しました。

実験結果:

  • 表7に示すように、sentinelトークンを挿入することで、GECタスクでは性能が向上しましたが、SuperGLUEタスクでは明確な改善は見られませんでした。
  • ただし、sentinelトークンを挿入することで、密なモデルとスパースモデルの両方で、学習の収束が速くなることが分かりました。

考察:

  • sentinelトークンは、事前学習とファインチューニングのドメインギャップを完全に解消するわけではありません。
  • しかし、学習の初期段階において、モデルがタスクを理解するのを助ける効果があると考えられます。

結論:

このセクションでは、ファインチューニング時にsentinelトークンを挿入するアプローチについて検証しました。sentinelトークンは、GECタスクでは性能向上に貢献しましたが、SuperGLUEタスクでは明確な効果は見られませんでした。ただし、学習の収束を速める効果があることが分かりました。

付記:

  • sentinelトークンを複数挿入する実験も行いましたが、明確な改善は見られませんでした。
  • Lester et al. (2021) でも、同様の概念が用いられています。

5 DESIGNING SPARSE MODELS

このセクションでは、スパースモデル設計における重要な要素と、それらに関する推奨事項を提示しています。

背景:

  • 密なモデルの設計は、Kaplan et al. (2020) の研究によって確立された原則に基づいています。

  • しかし、スパースモデルは、以下のような新たな課題を提起します。

    1. エキスパート数: いくつのエキスパートを使用すべきか？
    2. ルーティングアルゴリズム: どのルーティングアルゴリズムを使用すべきか？
    3. 容量係数: どの値を使用すべきか？
    4. ハードウェア: ハードウェアはこれらの決定にどのように影響するか？

推奨事項:

本論文では、パレート効率的なスパースモデルを構築するための推奨事項として、以下の3点を挙げています。

  1. トップ2ルーティング、容量係数1.25、コアあたり最大1エキスパート:

    • トップ2ルーティングは、各トークンを上位2つのエキスパートにルーティングするアルゴリズムです。
    • 容量係数1.25は、各エキスパートが処理できるトークン数を調整するハイパーパラメータです。
    • コアあたり最大1エキスパートとすることで、メモリ転送量を抑制し、計算効率を高めます。

  2. 評価時の容量係数調整:

    • メモリや計算リソースの制約に合わせて、評価時に容量係数を調整することができます。

  3. 密な層のスタッキングと乗算バイアス:

    • スパース層の前後に密な層をスタッキングすることで、モデルの品質を向上させることができます (付録C)。
    • エキスパート層に乗算バイアスを追加することも、品質向上に効果的です (付録C)。

他の研究:

  • Clark et al. (2022) は、より高い層頻度とトップ1ルーティング (Fedus et al., 2021) を推奨しています。

まとめ:

スパースモデルの設計には、エキスパート数、ルーティングアルゴリズム、容量係数、ハードウェアなど、考慮すべき要素が多数あります。本論文で提示された推奨事項は、パレート効率的なスパースモデルを構築するための指針となります。

注記:

  • パレート効率性とは、ある指標を改善しようとすると、他の指標が悪化してしまう状態を指します。
  • パレート効率的なモデルとは、どの指標も犠牲にすることなく、全体的な性能を最大化しているモデルです。

5.1 SETTING THE NUMBER OF EXPERTS

このセクションでは、スパースモデル設計における重要な要素の一つであるエキスパート数の決定方法について解説しています。

背景:

  • エキスパート数は、スパースモデルの性能と効率性に大きな影響を与えます。
  • エキスパート数が多いほど、モデルの表現力が高まりますが、計算コストも増加します。

先行研究:

  • Fedus et al. (2021) は、Switch Transformerにおいて、C4データセットを用いた事前学習で、最大512エキスパートまで性能が向上することを示しました。
  • Kim et al. (2021) は、最大64エキスパートまで性能が向上することを報告しています。
  • Clark et al. (2022) は、最大512エキスパートまで調査を行っています。
  • しかし、256エキスパートを超えると、性能向上の度合いは小さくなる傾向があります。

ハードウェアの影響:

  • エキスパート数の決定には、使用するハードウェアの特性も考慮する必要があります。
  • 演算強度 (Operational Intensity) は、計算とメモリアクセスの比率を表す指標であり、ハードウェアの効率性を評価する上で重要です。
  • 現代のGPUやTPUでは、演算強度が高いほど効率が良くなります。
  • スパースモデルでは、コアあたり複数のエキスパートを使用すると、メモリ転送量が増加し、演算強度が低下する可能性があります。

推奨事項:

  • 本論文では、TPUシステムを使用する場合、コアあたり1エキスパート以下にすることを推奨しています。
  • これにより、演算強度を高く保ち、評価や推論に必要なコア数を削減できます。
  • また、エキスパート数を減らすことで、モデル並列化に割り当てるコア数を増やし、モデル全体のFLOPsを増やすことができます。

データ並列化とモデル並列化:

  • 本論文では、データ並列化とモデル並列化を組み合わせて使用しています。
  • データ並列化は、データセットを複数のデバイスに分割して並列処理を行う手法です。
  • モデル並列化は、モデルのパラメータを複数のデバイスに分割して並列処理を行う手法です。
  • データ並列化は「行」方向に、モデル並列化は「列」方向に適用されます。
  • コアあたりのエキスパート数をデータ並列化の行数よりも少なくすることで、演算強度を高く保ち、モデル並列化に割り当てるコア数を増やすことができます。

まとめ:

スパースモデルのエキスパート数は、モデルの性能と効率性に大きな影響を与える重要な要素です。エキスパート数の決定には、先行研究の結果だけでなく、使用するハードウェアの特性も考慮する必要があります。本論文では、TPUシステムを使用する場合、コアあたり1エキスパート以下にすることを推奨しています。

5.2 CHOOSING THE CAPACITY FACTOR AND ROUTING ALGORITHM の詳細説明

このセクションでは、スパースモデル設計における容量係数 (CF) と ルーティングアルゴリズム の選択方法について、表8と表9の実験結果を交えながら解説しています。

容量係数 (CF): 各エキスパートが処理できるトークンの最大数を調整するハイパーパラメータです。

ルーティングアルゴリズム: トークンをどのエキスパートに割り当てるかを決定するアルゴリズムです。

実験設定:

  • すべてのモデルは、100kステップ、バッチサイズ1Mトークンで事前学習されています。
  • スパースモデルは32エキスパートを持ち、T5-Large (Raffel et al., 2019) と同等のFLOPsになるように設計されています。

表8: 容量係数とルーティングアルゴリズムの比較

  • 学習CFと評価CF: 学習時と評価時のCFをそれぞれ調整することができます。
  • トップnルーティング: 各トークンを上位n個のエキスパートにルーティングするアルゴリズムです。
  • 結果:
    • 学習CFと評価CFの両方を大きくすると、性能が向上する傾向があります。
    • トップ(n+1)ルーティングは、トップnルーティングよりもわずかに性能が向上する傾向がありますが、その差はわずかです。

表9: TPUにおけるスパースモデルのプロファイリング

  • 学習CF: 学習時のCFを変化させて、ステップ時間 (1ステップにかかる時間) を測定しました。
  • モデル: ST-MoE-L (10億パラメータ) とST-MoE-32B (320億パラメータ)
  • 結果:
    • 学習CFを大きくすると、ステップ時間が長くなる傾向があります。
    • モデルが大きくなるほど、学習CFの影響が大きくなります。

考察:

  • 容量係数:
    • CFを大きくすると、性能は向上しますが、計算コストも増加します。
    • 最適なCFは、使用するハードウェアとソフトウェアの組み合わせによって異なります。
    • all2all通信とallreduce通信が高速なシステムでは、大きなCFが適しています。
    • all2all通信とallreduce通信が遅いシステムでは、小さなCFが適しています。
  • ルーティングアルゴリズム:
    • トップnルーティングでは、nを増やすことによる性能向上はわずかです。
    • 計算コストを考慮すると、トップ2ルーティングが適切な選択肢となります。

推奨事項:

  • 本論文では、TPUとMesh Tensorflowを使用する場合、学習CFは1.25、評価CFは2.0、ルーティングアルゴリズムはトップ2ルーティングを推奨しています。
  • これは、性能と計算コストのバランスを考慮した結果です。

付記:

  • バッチ優先ルーティング (Batch Prioritized Routing, BPR) (Riquelme et al., 2021) は、CFが1.0未満の場合に有効なルーティングアルゴリズムです (付録D)。
  • BPRは、エンコーダ側のみに適用できます。デコーダ側に適用すると、モデルが不正をしてしまう可能性があります。

結論:

容量係数とルーティングアルゴリズムは、スパースモデルの性能と効率性に大きな影響を与える重要な要素です。最適な選択は、使用するハードウェアとソフトウェア、そして目標とする性能と計算コストのバランスによって異なります。

6 EXPERIMENTAL RESULTS

このセクションでは、これまで説明してきた学習の安定化、ファインチューニング、モデル設計に関する改善点を踏まえ、具体的な実験結果 を示し、提案手法の有効性を検証しています。

実験の概要:

  1. T5-Largeと同等のスパースモデル (ST-MoE-L) の検証:
    • まず、T5-Large (Raffel et al., 2019) と同程度のFLOPsを持つスパースモデル (ST-MoE-L) を構築し、その性能を検証しました。
  2. 2690億パラメータのスパースモデル (ST-MoE-32B) の学習:
    • 次に、320億パラメータの密なモデルと同等のFLOPsを持つ、2690億パラメータのスパースモデル (ST-MoE-32B) を設計・学習しました。
    • このモデルは、様々なNLPタスクにおいて最先端の性能を達成しました。

評価に用いたベンチマーク:

  • SuperGLUE (Wang et al., 2019):
    • 感情分析、語義曖昧性解消、文類似性、自然言語推論、質問応答、共参照解決、文完成、文 acceptability などの様々なタスクを含むベンチマークです。
    • SuperGLUEでの高性能は、多くのNLPタスクでの高性能と相関することが知られています (ただし、保証されるわけではありません)。
  • その他のベンチマーク:
    • 文書要約: CNN-DM (Hermann et al., 2015) とBBC XSum (Narayan et al., 2018)
    • 質問応答: SQuAD (Rajpurkar et al., 2016)、ARC Easy、ARC Reasoning Challenge (Clark et al., 2018)
    • クローズドブック質問応答: Natural Questions (Kwiatkowski et al., 2019)、Web Questions (Berant et al., 2013)、Trivia QA (Joshi et al., 2017)
      • クローズドブック質問応答とは、外部知識を参照せずに回答する質問応答です。
    • 常識推論: Winogrande Schema Challenge (Sakaguchi et al., 2020)
    • 自然言語推論: Adversarial NLI Benchmark (Nie et al., 2019)

実験結果:

  • ST-MoE-L:
    • 表10に示すように、ST-MoE-Lは、密なモデルと同程度の計算量で、様々なタスクにおいて優れた性能を示しました。
    • 特に、クローズドブック質問応答タスクでは、大幅な性能向上を達成しました。
  • ST-MoE-32B:
    • 表12に示すように、ST-MoE-32Bは、様々なタスクにおいて最先端の性能を達成しました。
    • SuperGLUEでは、推定される人間の能力を超えるスコアを達成しました。
    • 文書要約、ARC Easy、ARC Challenge、クローズドブック質問応答 (WebQA、NatQA) でも、最先端の性能を更新しました。
    • 敵対的に構築されたデータセット (ANLI R3、Winogrande XL) でも、大幅な性能向上を達成しました。

結論:

本論文で提案された学習の安定化、ファインチューニング、モデル設計に関する改善点は、スパースモデルの性能を大幅に向上させることが実証されました。特に、ST-MoE-32Bは、様々なNLPタスクにおいて最先端の性能を達成し、スパースモデルの潜在能力の高さを示しました。

6.1 ST-MoE-L

このセクションでは、T5-Largeと同等の計算量を持つスパースモデル ST-MoE-L の性能を検証した結果について解説しています。

ST-MoE-Lの構成:

  • パラメータ数: 7.7億
  • エキスパート数: 32
  • エキスパート層の頻度: 1/4 (4つごとに1つのFFN層をMoE層に置き換え)
  • 事前学習とファインチューニングの学習CF: 1.25
  • 評価CF: 2.0

学習データ:

  • C4データセット (Raffel et al., 2019) を5000億トークン使用
  • バッチサイズ: 1Mトークン
  • シーケンス長: エンコーダ512、デコーダ114

評価方法:

  • 複数のタスクを混合して学習し、各タスクの検証データセットで性能を評価しました。
  • タスクの混合比率は、各タスクの学習データサンプル数に比例するように設定しました (最大65536サンプル)。

結果:

  • 表10に示すように、ST-MoE-Lは、密なモデル (Dense-L) と同程度の計算量で、様々なタスクにおいて優れた性能を示しました。
  • 特に、クローズドブック質問応答タスク (Roberts et al., 2020) では、大幅な性能向上を達成しました。
  • 一方、学習データが少ないタスク (CB, WSC) では、密なモデルに比べて性能が劣る結果となりました。

考察:

  • ST-MoE-Lは、密なモデルと同程度の計算量で、多くのタスクにおいて同等以上の性能を達成できることが示されました。
  • 特に、クローズドブック質問応答タスクでの性能向上は、スパースモデルが知識の効率的な表現に優れていることを示唆しています。
  • 学習データが少ないタスクでの性能低下は、4.1で議論したように、スパースモデルが過学習しやすいという特性によるものと考えられます。

結論:

ST-MoE-Lは、計算効率の高いスパースモデルでありながら、多くのNLPタスクにおいて優れた性能を発揮することが実証されました。特に、クローズドブック質問応答タスクでの性能向上は注目すべき成果です。

付記:

  • 本論文では、タスクごとに個別にファインチューニングを行うのではなく、複数のタスクを混合して学習しています。
  • これは、スパースモデルの汎化性能を向上させるための工夫の一つです。

6.2 ST-MoE-32B

このセクションでは、本論文の主要な成果である 2690億パラメータのスパースモデル ST-MoE-32B について、その設計、学習、評価結果を詳しく解説します。

ST-MoE-32Bの設計目標:

  • FLOPsとパラメータ数のバランス: 過去の研究 (Fedus et al., 2021) では、FLOPsが大きすぎるスパースモデルは学習が不安定になる傾向がありました。そこで、ST-MoE-32Bでは、FLOPsとパラメータ数のバランスを重視した設計を行いました。
  • サービス提供とファインチューニングの容易さ: パラメータ数が少ないほど、モデルのサービス提供とファインチューニングが容易になります。ST-MoE-32Bは、Switch-CやSwitch-XXLよりもパラメータ数が少ないため、これらの点で優れています。
  • 計算効率: エキスパートの隠れ層サイズ (dff) を大きくすることで、計算効率を高めました。
  • ハードウェア性能: dkvを128に増やすことで、TPUでの性能を向上させました。

ST-MoE-32Bの構成:

  • パラメータ数: 2690億
  • FLOPs: 320億パラメータの密なモデルと同等
  • エキスパート数: 64
  • エキスパート層の頻度: 1/4
  • Sparse-Dense stacking: 付録Cで説明されているアーキテクチャを採用

学習データ:

  • 英語のC4データセット (Raffel et al., 2019) とGLaM (Du et al., 2021) のデータセットを混合 (付録E参照)
  • 学習トークン数: 1.5兆
  • バッチサイズ: 1Mトークン
  • オプティマイザ: Adafactor (デフォルトのハイパーパラメータを使用)
  • 学習率: 10kステップのウォームアップ後、逆平方根減衰

評価:

  • 表10: FLOPsが一致する密なモデル (Dense-L) とスパースモデル (ST-MoE-L) のファインチューニング性能を比較
  • 表11: 過去のスパースモデル (Switch-XXL, Switch-C) とST-MoE-L, ST-MoE-32Bの構成を比較
  • 表12: ST-MoE-32Bと、ゼロショット、ワンショット、ファインチューニング設定における過去の最先端モデルの性能を比較

結果:

  • ST-MoE-32Bは、様々なNLPタスクにおいて最先端の性能を達成しました。
    • SuperGLUE: テストサーバーで91.2のスコアを達成 (人間の能力を超える)
    • 文書要約: XSumとCNN-DMの両方で最先端の性能
    • ARC: EasyとChallengeの両方で最先端の性能
    • クローズドブック質問応答: WebQAとNatQAで最先端の性能
    • 敵対的に構築されたデータセット: ANLI R3とWinogrande XLで大幅な性能向上

考察:

  • ST-MoE-32Bは、FLOPsとパラメータ数のバランスを重視した設計、学習の安定化、適切なファインチューニングプロトコルによって、高い性能を達成することができました。
  • スパースモデルは、大規模言語モデルのスケーリングにおいて、非常に有効なアプローチであることが示されました。

課題:

  • 小規模なデータセット (SQuAD) では、過去の最先端モデルに及ばない結果となりました。
  • SuperGLUEの一部タスク (CB, WSC) では、性能が向上しませんでした。
  • クローズドブック質問応答 (TriviaQA) では、GPT-3やGLAMに及ばない結果となりました。

結論:

ST-MoE-32Bは、スパースモデルの設計、学習、評価に関する重要な知見を提供し、大規模言語モデルのスケーリングにおける新たな可能性を示しました。今後の研究では、小規模データセットへの対応、さらなる性能向上、効率的な推論方法などが課題となります。

7 TRACING TOKENS THROUGH THE MODEL

このセクションでは、トークンがエキスパート間でどのようにルーティングされるか を可視化することで、スパースモデルの動作を定性的に分析しています。

目的:

  • これまでのセクションでは、モデルの性能を定量的に評価してきました。
  • このセクションでは、モデル内部の動作をより深く理解するために、トークンのルーティングを可視化し、定性的な分析を行います。

方法:

  • ST-MoE-Lモデルにトークンのバッチを入力し、各層におけるトークンの割り当てを手動で調査しました。
  • 事前学習データとして、単一言語のC4コーパス (Raffel et al., 2019) と多言語のmC4コーパス (Xue et al., 2020) の2種類を使用しました。
  • モデルは、エンコーダとデコーダの両方に、それぞれ32エキスパートを持つ6つのスパース層を持っています。

事前知識:

  • Span Corruption Objective: 入力文から可変長の連続したセグメントをマスクし、それを予測するようにモデルを学習させるタスクです。
    • 例:
      • 入力: I went to to buy
      • ターゲット: the store milk
  • エンコーダ・デコーダアーキテクチャ: 入力文はエンコーダに、ターゲット文はデコーダに入力されます。
  • グループ単位のルーティング: トークンは、グループ単位でエキスパートにルーティングされます。
    • 各グループは、Shazeer et al. (2017) で提案された補助損失によって、エキスパート間で負荷分散されるように促されます (付録A参照)。
    • トークンは、バッチ全体ではなく、自分のグループ内の他のトークンとエキスパートの割り当てを競います。
    • エキスパートの専門化は、各グループ内のトークンの分布に大きく影響されます。
    • グループという概念は、正しいトークンを正しいエキスパートに送受信するコストを削減するために導入されました。

分析のポイント:

  • エンコーダにおけるエキスパートの専門化
  • デコーダにおけるエキスパートの専門化
  • 多言語モデルにおけるエキスパートの専門化

まとめ:

このセクションでは、トークンのルーティングを可視化することで、スパースモデルの動作を定性的に分析するための準備を行いました。具体的な分析結果については、7.1以降のセクションで詳しく解説されています。

7.1 ENCODER EXPERTS EXHIBIT SPECIALIZATION

このセクションでは、エンコーダにおけるエキスパートの専門化について、具体的な例を挙げながら解説しています。

専門化とは: 特定のエキスパートが、特定の種類のトークン (例: 句読点、動詞、固有名詞) を処理することに特化している現象を指します。

観察結果:

  • sentinelトークンへの特化: 各層において、少なくとも1つのエキスパートがsentinelトークン (マスクされた部分を予測するための特殊なトークン) を処理することに特化していました。
  • その他の専門化:
    • 句読点
    • 動詞
    • 固有名詞
    • 数値
    • 色、空間的位置
    • 視覚的な描写

表13: エンコーダエキスパートの専門化の例

  • 表13には、エンコーダエキスパートの専門化の具体的な例が示されています。
  • 例えば、あるエキスパートは句読点のみを処理し、別のエキスパートは動詞のみを処理しています。

考察:

  • エンコーダにおけるエキスパートの専門化は、span corruption objective の性質に起因すると考えられます。
  • span corruption objectiveでは、入力文から連続したトークンがマスクされるため、各エキスパートは、マスクされた部分の前後の文脈に基づいて、特定の種類のトークンを予測することに特化することができます。

結論:

エンコーダのエキスパートは、特定の種類のトークンを処理することに特化しており、これはスパースモデルが効率的に学習を行うための重要なメカニズムであると考えられます。

付記:

  • 本論文では、明確な意味的または構文的な専門化が見られないエキスパートも多数存在することが指摘されています。
  • SentencePieceモデル (Kudo and Richardson, 2018) は、語彙に存在しないトークンを分割するため、専門化のパターンに影響を与える可能性があります。

7.2 DECODER EXPERTS LACK SPECIALIZATION

このセクションでは、エンコーダとは対照的に、デコーダのエキスパートには明確な専門化が見られないことを示し、その理由について考察しています。

観察結果:

  • sentinelトークンの均等な分散: デコーダのエキスパートは、sentinelトークンをほぼ均等に処理しており、特定のエキスパートに集中しているわけではありません (表14参照)。
  • 意味的・構文的な専門化の欠如: エンコーダのように、特定の種類のトークンを処理することに特化したエキスパートは見当たりませんでした。

表13と表14:

  • 表13: エンコーダのエキスパートの専門化の例を示しています。
  • 表14: エンコーダとデコーダの各層における、ルーティングされたsentinelトークンのエントロピーを計算した結果を示しています。
    • エントロピーが高いほど、トークンの分散が均等であることを意味します。
    • デコーダのエントロピーは、エンコーダよりも高く、ほぼ均一な分布に近い値を示しています。

考察:

  • デコーダのエキスパートに専門化が見られない理由は、span corruption objective とトークンルーティングのメカニズムに起因すると考えられます。
    • Span Corruption Objective: デコーダは、エンコーダよりも長いシーケンスを処理するため、各グループに含まれるトークン数が少なくなります。
    • トークンルーティング: デコーダでは、各グループに含まれるトークンの多くがsentinelトークンであるため、意味的な空間が狭くなり、エキスパートが専門化しにくくなります。

結論:

デコーダのエキスパートには、エンコーダのような明確な専門化は見られませんでした。これは、デコーダが処理するタスクの性質と、トークンルーティングのメカニズムによる影響と考えられます。

今後の課題:

  • デコーダにおけるスパース性とエキスパートの専門化をより効果的に活用する方法を検討する必要があります。
  • 例えば、デコーダ層のエキスパートを削除することで、自己回帰的なデコードを高速化できる可能性があります (Kudugunta et al., 2021a)。

付記:

  • 本論文では、学習済みルーティングではなく、ランダムなルーティングを用いたアーキテクチャ (Roller et al., 2021; Zuo et al., 2021) も高い性能を示すことが指摘されています。
  • 学習済みルーティングとランダムルーティングのスケーリング特性を比較分析することは、今後の研究課題として興味深いテーマです。

7.3 MULTILINGUAL EXPERTS SPECIALIZE, BUT NOT BY LANGUAGE

このセクションでは、多言語の事前学習データを用いて学習したスパースモデルにおける、エンコーダのエキスパートの専門化について分析しています。

観察結果:

  • 言語に依存しない専門化:
    • 単一言語モデルと同様に、多言語モデルでもエンコーダのエキスパートは専門化を示しました。
    • しかし、エキスパートは特定の言語に特化するのではなく、複数の言語を横断的に処理していました。
    • 表15は、多言語エキスパートの専門化の例を示しています。
      • 例えば、あるエキスパートは、英語の "for" とフランス語の "pour" の両方を処理しています。

考察:

  • 言語に依存しない専門化は、トークンルーティングのメカニズムによって説明できます。
    • 各グループには、通常、1つから数個の言語のトークンしか含まれていません。
    • そのため、すべてのエキスパートは、あらゆる言語のトークンを処理できるように学習されます。

グローバルな負荷分散損失:

  • 本論文では、すべてのエキスパートがすべての言語を処理できるようにするために、グローバルな負荷分散損失を用いた実験も行いました。
  • しかし、この方法では、負荷分散がうまくいかず、モデルの性能も低下しました。

結論:

多言語スパースモデルにおいても、エンコーダのエキスパートは専門化を示しますが、その専門化は特定の言語に限定されません。これは、トークンルーティングのメカニズムと、各グループに含まれる言語の多様性による影響と考えられます。

今後の課題:

  • 多言語スパースモデルにおけるエキスパートの専門化をさらに促進する方法を検討する必要があります。
  • 特に、グローバルな負荷分散損失を改善することで、エキスパートがより効果的に言語を横断的に処理できるようになる可能性があります。

付記:

  • 表15は、多言語エキスパートの専門化の例をいくつか示していますが、明確な意味的または構文的な専門化が見られないエキスパートも多数存在します。
  • 多言語スパースモデルの学習は、単一言語モデルよりも複雑であり、さらなる研究が必要です。

8 RELATED WORK

このセクションでは、本研究に関連する先行研究を以下の項目に分けて解説しています。

1. Mixture-of-Experts (MoE) の歴史:

  • MoEは、少なくとも30年前のJacobs et al. (1991) や Jordan and Jacobs (1994) の研究に遡ります。
  • 当初のMoEは、アンサンブル学習のように、ニューラルネットワーク全体を定義していました。
  • Eigen et al. (2013) は、MoEをより深いネットワークの構成要素として導入することを提案しました。
  • Shazeer et al. (2017) は、1370億パラメータのMoEモデルを構築し、機械翻訳において当時の最先端の性能を達成しました。
  • それ以降、多くの研究がMoEを構成要素として用いるアプローチを採用しています。

2. 自然言語処理におけるスケール:

  • 近年、自然言語処理において、モデルの規模を大きくすることで性能が向上することが明らかになってきました (Kaplan et al., 2020; Brown et al., 2020)。
  • スパースエキスパートモデルは、大規模な密なモデルと同等の性能を、より効率的に達成する方法として注目されています。
  • Fedus et al. (2021) は、T5-XXL (Raffel et al., 2019) よりも4倍高速な事前学習を実現しました。
  • Du et al. (2021) は、GPT-3 (Brown et al., 2020) と同等の品質を、3分の1のエネルギー消費で達成しました。
  • 1兆パラメータを超える深層ニューラルネットワークの効率的な学習も、複数の研究グループによって達成されています (Fedus et al., 2021; Yang et al., 2021; Du et al., 2021; Lin et al., 2021)。
  • スパースエキスパートモデルは、言語モデリング、機械翻訳など、データ量が多く、分布シフトがないタスクで特に成功を収めています。
  • 一方で、事前学習では優れた性能を示す一方で、ファインチューニング時に小規模なモデルに劣ることが報告されています (Fedus et al., 2021; Narang et al., 2021; Artetxe et al., 2021)。
  • 今後の研究では、正則化技術の進歩によって、下流タスクの品質が向上することが期待されます。

3. より良いルーティングアルゴリズムに向けて:

  • BASE layers (Lewis et al., 2021) は、トークンルーティングを線形割り当て問題として再定義し、負荷分散のための補助損失関数を不要にしました。
  • Clark et al. (2022) は、様々なルーティングアルゴリズムのスケーリング特性を詳細に分析し、最適輸送定式化を用いたBASE layersの変種を提案しました。
  • Yang et al. (2021) は、M6-Tアーキテクチャとエキスパートプロトタイピングを導入し、エキスパートを異なるグループに分割して、複数のトップ1ルーティング手順を適用しました。
  • Hazimeh et al. (2021) は、従来のトップkゲーティングよりも優れた性能を示す、連続微分可能なスパースゲートを提案しました。
  • Bengio et al. (2016) は、ルーティング選択を強化学習問題として扱うことを検討しました。
  • Hash layers (Roller et al., 2021) は、学習済みルーティングではなく、ハッシュ関数に基づくランダムな固定ルーティングが、学習済みルーティングに匹敵する性能を達成することを示しました。
  • Zuo et al. (2021) は、学習と推論の両方でエキスパートをランダムに選択するアルゴリズムを提案し、Switch Transformerよりも2 BLEUポイント高いスコアを達成しました。
  • Fan et al. (2021) は、明示的な言語固有のサブ層を持つアーキテクチャを設計し、+1 BLEUの性能向上を達成しました。

4. 他のモダリティにおけるスパースエキスパートモデル:

  • MoEとスパースエキスパートモデルは、自然言語処理以外の分野でも成果を上げています。
  • Riquelme et al. (2021) は、150億パラメータのV-MoEを設計し、ImageNet (Deng et al., 2009) において、より少ない計算リソースで最先端の性能を達成しました。
  • Lou et al. (2021) は、画像パッチとチャネルの両方の次元でMoE層を使用することで、密なビジョンモデルよりも優れた性能を示しました。
  • 音声認識においても、SpeechMoE (You et al., 2021a;b) や MoEを用いたSequence-to-Sequence Transformer、Transformer Transducer (Kumatani et al., 2021) によって性能が向上しています。

5. スパースモデルのデプロイメントの改善:

  • 初期のエキスパート設計 (本研究を含む) では、各トークンを個別にエキスパートにルーティングしていました。
  • このようなアーキテクチャは、パラメータを保存するためのメモリ量が多いため、サービス提供が困難になる可能性があります。
  • 蒸留 (Fedus et al., 2021) はある程度の効果がありますが、最近の研究では、文全体またはタスク全体をルーティングする手法が提案されています (Kudugunta et al., 2021b; Zuo et al., 2021)。
  • これにより、サービス提供時に必要なサブネットワークのみを抽出することができます。
  • Kim et al. (2021) は、タスクに必要のないエキスパートを直接プルーニングする手法を検討しています。

6. MoEを用いたマルチタスク学習:

  • Ma et al. (2018) は、タスクごとに個別のゲーティングネットワークまたはルーターネットワークを使用することを推奨しました。
  • Gururangan et al. (2021) は、言語モデルのモジュール化をさらに進め、ドメイン/タスクラベルまたは推論されたラベルに基づいてエキスパートを条件付きでアクティブにすることを提案しました。

まとめ:

このセクションでは、MoEの歴史から最新の研究動向まで、幅広い先行研究を紹介しました。これらの研究は、スパースエキスパートモデルが自然言語処理や他の分野において、大規模モデルのスケーリングと効率的な学習を実現するための重要なアプローチであることを示しています。

9 DISCUSSION

このセクションでは、スパースモデルに関する本研究を通して得られた知見や考察、今後の研究課題などを以下の項目に分けて議論しています。

1. 多言語データにおける予測不可能な学習ダイナミクス:

  • 多言語データで事前学習を行うと、単一言語データの場合よりも学習のスピードアップが小さくなり、不安定になる傾向があることが観察されました。
  • これは、各グループに含まれる言語の多様性と、バッチごとのシーケンスの分散が原因と考えられます。
  • 多言語データを用いた学習では、グループ内の言語数が限られるため、エキスパート間で言語のバランスが取りにくく、ルーティングが不安定になる可能性があります。

2. スパースモデルの堅牢性:

  • スパースモデルは、ハイパーパラメータやアーキテクチャの変更に対して堅牢であることが分かりました。
  • ルーティングアルゴリズム、トークンのドロップ率、バッチサイズ、学習率などを変化させても、比較的安定した性能を示しました。
  • これは、スパースモデルが冗長な表現を学習しているためと考えられます。
  • しかし、ファインチューニングのプロトコル (特にバッチサイズと学習率) は、スパースモデルの性能に大きな影響を与えるため、適切に調整することが重要です。

3. 適応計算:

  • スパースモデルは、入力に応じて異なる計算を行う適応計算モデルの一種とみなすことができます。
  • 容量係数が1未満の場合、モデルは一部のトークンに対して計算を行わないことを学習します。
  • これは、コンピュータビジョン (Riquelme et al., 2021) や自然言語処理 (付録D) で有望な結果を示しています。
  • 今後は、異なる計算を行う異種エキスパート (heterogeneous experts) を用いたモデルの開発が期待されます。

4. 小規模から大規模へのスケーリング:

  • 小規模な実験で得られた知見が、大規模な実験でも必ずしも成り立つとは限りません。
  • 例えば、Fedus et al. (2021) のトップnルーティングに関する知見は、本論文の大規模な実験では逆転しました。
  • モデルの規模を大きくすると、学習ダイナミクスが変化し、新たな課題が生じる可能性があります。

5. より低い精度での学習:

  • router z-lossを用いることで、bfloat16のような低精度演算を用いても安定した学習が可能になることが分かりました。
  • router z-lossは、logitsの値を小さくすることで、丸め誤差の影響を軽減します。
  • 今後は、スパースモデルに適した新しい数値精度フォーマットの開発が期待されます。

6. より多くの乗算的相互作用を持つ新しい演算の設計:

  • 加算よりも乗算の要素が多い演算や、多くの数値を累積しない演算は、モデルの性能を向上させる可能性があります。
  • 付録Cでは、エキスパート層に乗算的相互作用を追加することで、事前学習のスピードアップを達成した例が示されています。
  • 今後は、モデルの安定性を維持しつつ、乗算的相互作用を効果的に活用するアーキテクチャの開発が期待されます。

7. 望ましくないモデルスケーリングダイナミクスを軽減するための活性化の制約:

  • エンコーダ・デコーダモデルは、デコーダのみのモデルよりも学習が不安定になる傾向があります。
  • また、深いモデルは、浅いモデルよりも学習が不安定になる傾向があります。
  • これは、指数関数の数が増加することが原因と考えられます。
  • 今後は、非スパースモデルのattention softmaxにz-lossペナルティを追加することで、これらの問題を解決できる可能性があります。

8. 密なモデルとスパースモデルのハイパーパラメータ依存性の違い:

  • 密なモデルとスパースモデルでは、最適なファインチューニングプロトコルが異なることが分かりました。
  • スパースモデルは、小さなバッチサイズと高い学習率で良い性能を示す傾向があります。
  • 新しいモデルクラスを評価する際には、主要なハイパーパラメータを徹底的にテストすることが重要です。

まとめ:

このセクションでは、スパースモデルに関する様々な議論を行いました。スパースモデルは、大規模言語モデルのスケーリングにおいて非常に有望なアプローチですが、まだ解決すべき課題も残されています。今後の研究によって、これらの課題が克服され、スパースモデルがさらに広く活用されるようになることが期待されます。

10 CONCLUSION

このセクションでは、本研究の成果をまとめ、スパースモデルの将来展望について述べています。

本研究の成果:

  • スケールに対する過度な期待の抑制: Fedus et al. (2021) で示されたスケールに対する過度な期待を抑制し、パラメータ数が5分の1で計算量 (FLOPs) とパラメータ数のバランスが良いモデルが、より効果的なスパース学習者であることを示しました。
  • 実用性の向上: パラメータ数を削減することで、メモリオーバーヘッドを減らし、スパースモデルのデプロイを容易にしました。
  • 幅広いNLPタスクでSOTAを達成: 開発したスパースモデルを用いて、様々な公開ベンチマークにおいて最先端の性能を達成しました。

スパースモデルの将来展望:

  • 本研究は、モデルのスパース性によって、大規模言語モデルの性能を大幅に向上できることを示しました。
  • この成果は、スパースモデルの採用を促進し、大規模言語モデルの開発に新たな可能性をもたらすことが期待されます。

要約:

本研究は、スパースモデルの学習安定性、ファインチューニング、モデル設計に関する重要な知見を提供し、大規模言語モデルのスケーリングにおける新たな方向性を示しました。スパースモデルは、高性能な言語モデルを構築するための有効な手段であり、今後の発展が期待されます。

特に強調されている点:

  • スパースモデルは、計算効率と性能の両立を実現できる。
  • 本研究の成果は、スパースモデルの実用性を高め、普及を促進する。

今後の研究課題:

  • 本論文では明示的に言及されていませんが、スパースモデルのさらなる性能向上、効率的な学習・推論方法の開発などが、今後の研究課題として考えられます。

J NEGATIVE RESULTS

このセクションでは、本研究において効果が得られなかった、つまり性能向上につながらなかったアイデアを列挙し、その詳細を説明しています。

これらの「失敗例」を共有することで、今後の研究における参考情報とし、同じ轍を踏まないようにすることを意図しています。

1. トークンがドロップされた場合の情報をルーターに追加:

  • 前の層でトークンがドロップされたかどうかをルーターに伝えることで、ルーティングの精度が向上する可能性を検証しました。
  • 具体的には、過去のエキスパート層でトークンがルーティングされた回数をカウントし、その値に対応する埋め込み表現をルーターの入力に追加しました。
  • しかし、この変更は性能に影響を与えませんでした。

2. 明示的なエキスパート位置情報の追加:

  • 各トークンがどのエキスパートに送られたかを明示的に表現することで、性能向上または学習の収束速度向上が期待されました。
  • 具体的には、各トークンが送られたエキスパートに対応する埋め込み表現を追加しました。
  • しかし、この変更は性能に影響を与えませんでした。

3. 事前学習ノイズによる事前学習とファインチューニングの差異の解消:

  • 事前学習とファインチューニングのドメインギャップを解消するために、事前学習時に様々なノイズを追加しました。
  • 具体的には、ドロップアウト、エキスパート全体のドロップアウト、ルーターへのエントロピー最大化補助損失などを試しました。
  • しかし、これらのノイズは事前学習の品質を低下させるか、ファインチューニングの性能向上に効果がありませんでした。

4. 下位n-1エキスパートに対するトップnルーティングの負荷分散:

  • 標準的なトップnルーティングでは、トークンが送られる上位1つのエキスパートの負荷分散のみを考慮します。
  • 下位n-1エキスパートに対しても負荷分散を行うことで、性能が向上する可能性を検証しました。
  • しかし、この変更による効果はわずかでした。

5. 過学習を防ぐための事前学習データとファインチューニングデータの混合:

  • ファインチューニング時の過学習を防ぐために、事前学習データとファインチューニングデータを様々な割合で混合しました。
  • しかし、この方法はファインチューニングの性能向上にはつながらず、学習損失が増加するだけでした。

まとめ:

このセクションでは、本研究で試みたものの効果が得られなかったアイデアを列挙し、その詳細を説明しました。これらの「失敗例」は、今後のスパースモデルの研究において、新たなアイデアを検討する際の参考情報となるでしょう。

重要なポイント:

  • 常にうまくいくアイデアは存在しない。
  • 失敗から学ぶことは、成功と同じくらい重要である。
  • 他の研究者と失敗例を共有することは、研究全体の進歩に貢献する。
0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?