EMO: 自然発生するモジュール性のためのエキスパート混合の事前学習
- +8
本日、EMO を公開します。これは、新しい混合専門家（MoE）モデルで、データから直接モジュール構造が立ち上がるように、エンドツーエンドで事前学習されています。人が定義した事前知識（prior）に頼ることなく、モジュール構造がデータそのものから出現します。EMOでは、特定のタスクに対してモデル全体のうち専門家の小さな一部（総数の12.5%）だけを使いながら、ほぼフルモデルと同等の性能を維持できます。また、全専門家を一緒に使った場合でも、強力な汎用モデルとして機能します。
大規模言語モデルは通常、単体のシステムとして学習され、デプロイされます。つまり、1つのモデルを初期化し、事前学習し、微調整し、そして統一された1つの実体として提供します。しかし、アプリケーションによっては、コード生成、数学的推論、あるいはドメイン固有の知識など、必要なのは能力の一部だけであることが多いです。最前線の言語モデルではパラメータが数兆にまで日常的に到達しているため、フルモデルを使いこなしたり適応させたりすることは、多くのユーザーにとって現実的ではなく、必要かどうかも分からないパラメータをホストするための不必要な計算コストとメモリを招きます。
混合専門家（MoE）モデルは、この制約を緩和する自然な方法のように見えます。各層で1つの大きなフィードフォワードネットワークを使うのではなく、MoEは多数の小さなネットワーク（専門家）を含み、入力トークンごとにごく一部だけを活性化します。原理的には、あるタスクが1つの能力だけを必要としているなら、その能力に関連する専門家だけを読み込めばよいはずです。
しかし実際には、既存のMoEは依然として、うまく機能するためにはフルモデルが必要です。同一の入力の中でも、異なるトークンが異なる専門家を活性化することがよくあるため、生成の過程でタスクが結果的にすべての専門家を使ってしまうことがあります。私たちの論文が示すように、これは部分的に、標準的なMoEにおける専門家が、高レベルの領域や能力ではなく、前置詞や句読点のような低レベルの語彙パターンにしばしば特化してしまうために起きます。その結果、小さな専門家のサブセットは、それだけでは確実に利用できません。
私たちが目指すのは、選択的に使ったり組み合わせたりできる、首尾一貫したグループに専門家が整理されるMoEモデルです。
そのために事前学習中に行える1つの方法は、数学、生物学、コードのような事前定義された意味領域に基づいてトークンを専門家へルーティングすることです。BTXや、私たちのFlexOlmoプロジェクトのような先行研究もこれを試みています。ただし、事前定義された領域には重要な制限があります。事前学習コーパス全体でドメインラベルが必要で、曖昧で取得コストが高い可能性があります。また、モデルが自分自身をどう整理することが許されるかに対して、人間のバイアスを過剰に注入してしまうかもしれません。さらに重要なのは、領域を最初に固定してしまうと、モジュール構造そのものも固定されてしまう点です。推論時に新しい領域や能力が現れても、どの専門家を使うべきかが自明ではありません。
そこでEMOが登場します。
私たちは、EMO――1Bアクティブ（活性）、14B総パラメータ（8専門家が活性、128専門家が総数）のMoEで、1兆トークンで学習――が、選択的な専門家の利用を支えることを示します。特定のタスクやドメインに対して、専門家の小さなサブセットだけ（総専門家の12.5%だけ）を使いながら、ほぼフルモデル性能を維持できます。同時に、すべての専門家を一緒に使う場合でも、EMOは依然として強力な汎用モデルであり続けます。対照的に、同じデータで学習された同等のアーキテクチャの標準MoEでは、その専門家サブセットを選択的に使うと、深刻な劣化が見られます。
EMOは、モジュール性を第一級の目的として訓練されたMoEです。特定のドメイン（例：数学、コード、生物医学）に対して、ユーザーは任意のサイズの小さな専門家（エキスパート）のサブセットを選択し、ほぼフルモデルの性能を維持できます。これにより、単一のモデルをコンポーザブル（組み立て可能）なアーキテクチャへと変換し、大規模で疎なMoEに対するメモリと精度のトレードオフを改善しつつ、柔軟なデプロイを可能にします。
モジュール性はどうやって自然に現れるのか？
MoEでは、ルータと呼ばれる小さなネットワークが各トークンをどの専門家（エキスパート）に割り当てるかを決定します。私たちは、類似したドメインのトークンは、同様の専門家サブセットを活性化するようにルータが学習することを望んでいます。我々の重要な観察は、同じドキュメントのトークンは通常、同じドメインに由来するという点です。そこで、ドキュメント境界を弱い教師信号として用います。つまり、訓練中は、あるドキュメント内のすべてのトークンが、共通の専門家プールから活性化する専門家を選ぶよう制限します。
標準のMoEとEMOの訓練の比較（k = 2、n = 10、説明のため共有エキスパートは省略）。（左）標準のMoEでは、各トークンが独立に上位k個の専門家を選択します。トークン間では、すべての専門家が使用されます。（右）EMOでは、ルータがまず各ドキュメントに対して専門家のサブセットを選択し、すべてのトークンがこのサブセット内に収まるよう制約されます。これにより、ドキュメント内で一貫した専門家の使用が強制され、専門家グループがドメイン特化へと形成されることを促します。
例えば、合計10人の専門家があり、トークンごとに活性化される専門家が2人のMoEでは、上の図に示すように、ドキュメント内のすべてのトークンが同じ4人の専門家プール内へルーティングするよう制限されます。このプールはルータ自身によって選ばれます。具体的には、ドキュメント内のすべてのトークンに対するルータの専門家嗜好を平均し、最も使われた専門家をドキュメントの共有プールとして選択します。異なるドキュメントは異なるプールを使うことができ、その結果、専門家グループが訓練データから直接、繰り返し現れるようになります。
システムを実装する際には、いくつかの考慮事項があります：
負荷分散（ロードバランシング）。 1つの技術的な課題は負荷分散です。標準のMoE訓練では、モデルが少数の専門家に崩壊してしまうのを防ぐために、負荷分散の目的関数が用いられます。一見すると、これはEMOの訓練目的と衝突しているように見えます。というのも、私たちは各ドキュメントに対して明示的に専門家のサブセットしか使わせないからです。
衝突が生じるのは、負荷分散が通常適用されるスケールに由来します。多くのMoE実装では、負荷分散は局所的に計算されます。多くの場合、それは少数のドキュメントしか含まないマイクロバッチ内で行われます。この局所的な目的関数は、同じドキュメント内のトークンに対して多くの専門家に分散させる方向に働き、ドキュメント内で専門家の使用を一貫させるというEMOの目的に真っ向から反します。
これを解決するために、私たちは多くのドキュメントにまたがって負荷分散をグローバルに適用します。このより大きなスケールでは、2つの目的は補完的になります。EMOは同じドキュメント内のトークンが首尾一貫した専門家プールを使うことを促し、グローバル負荷分散は異なるドキュメントが集合的にすべての専門家をカバーすることを促します。実際、安定した訓練にはグローバルな負荷分散が重要であることが分かりました。
ドキュメントプールサイズ。 ドキュメントプールサイズは、モジュール性の制約がどれほど厳しいかを制御します。プールが小さいほど、同じドキュメント内のトークンがより狭い範囲の専門家を共有することを強いられ、より強いモジュール性を促進します。一方、プールが大きいほど、モデルはより柔軟になりますが、制約は弱まります。
1つのプールサイズを固定するのではなく、訓練中にランダムにサンプリングします。これにより、EMOが単一のサブセットサイズに過学習するのを防ぎ、推論時に異なる専門家サブセットサイズをサポートできるようになります。
ベンチマーク結果
一般用途のベンチマークにおいて、EMOは標準のMoEモデルと同等の性能を達成し、モジュール性の目的がフルモデルの性能を犠牲にするものではないことを示しています。しかし、より重要な問いは、「専門家（エキスパート）のサブセットだけを残した場合でも、モデルは機能するのか？」です。この設定では、少量のタスク検証データにおけるルーティング使用量のランキングに基づいて専門家を並べ、最も使われた専門家を保持し、それ以外を捨てることで、タスク固有の専門家サブセットを構築します。
下の図は、EMOが選択的な専門家利用に対して頑健であることを示しています。専門家の25%だけを保持する（32専門家サブセット）場合、EMOの性能低下は全ベンチマークで絶対値としてわずか約1%にとどまります。さらに、専門家の12.5%だけを保持する（16専門家サブセット）場合でも、全体の低下は約3%程度にすぎません。これは微調整（ファインチューニング）の前後の両方で成り立ちます。対照的に、マッチさせる標準のMoEは、専門家サブセットが小さくなるにつれて急激に劣化し、最小の専門家サブセット設定では、しばしばランダム性能に近い、またはそれを下回るほどまで落ちます。
さらに、あるタスクに対して適切な専門家を選択することのコストが驚くほど低いことも示します。少数ショットのデモンストレーションを伴う単一の例で、フルの検証セットを使って選択した場合と同等の性能を発揮するモジュールを特定するのに十分です。そしてEMOは、特定の選択手法に結びついているわけではありません。Easy-EPのような既存の専門家プルーニング手法とも良好に動作し、両者は相補的です。
より小さい130Bトークン設定。メモリ予算の異なる条件下で、16のMMLUカテゴリにわたって平均した性能。EMOのエキスパート部分集合は、メモリと精度のトレードオフにおけるパレート境界を押し広げ、標準MoEや、同じく固定予算でゼロから訓練したモデルをも上回ります。
エキスパート部分集合は何に特化しているのか？
EMOが学習後に実際に何を得たのかを見るために、12Kの事前学習ドキュメントにまたがって、最初の100トークンのルータのアクティベーションをクラスタリングしました。標準MoEとの差は明らかです。
EMOのトークンクラスタは、たとえばHealth, Medical & Wellness（健康、医療、ウェルネス）、News Reporting（ニュース報道）、US Politics & Elections（米国の政治と選挙）、Film & Music（映画と音楽）といったものに対応しています。標準MoEでは、前置詞、固有名詞、コピュラ動詞、定冠詞のようなクラスタが生成されます。EMOでは、あるドキュメントからのトークンは主に同じクラスタに着地しますが、標準MoEでは、多くのクラスタに散らばってしまいます。
この対比は、1つの具体例を見るのがいちばん簡単です。健康記事を例に取りましょう。EMOでは、ほぼすべてのトークンがHealth, Medical & Wellnessクラスタにルーティングされます。標準MoEでは、最上位クラスタはPossessives & Definite Articles（所有格と定冠詞）です。このモデルは、その記事が何についてのものかにかかわらず、単語theやyourを使っているあらゆる他のテキストと一緒に記事をまとめてしまいます。
1Tトークンで学習したMoEにおける事前学習データのトークンクラスタ。EMOは意味的に意味のある領域にクラスタを形成し、同じドキュメント由来のトークンは主にひとまとまりにされます。標準MoEの学習は、表層的あるいは構文的な特徴のクラスタを生成し、ドキュメントのトークンは複数のクラスタに分散します。
EMOは表層的な特徴ではなく意味領域に対応するモジュールを形成するため、小さなエキスパート部分集合を選んでも、ちゃんと機能するモデルを得られます。このグループは、実際の能力に対応しているからです。
クラスタリングの結果は、私たちのインタラクティブな可視化で、ご自身でもいじって確かめられます。
私たちが公開するもの
私たちは、EMOを学習した完全なモデル、同じデータで学習した対応する標準MoEのベースライン、そして学習コードを公開します。MoEにおける創発的なモジュール性を研究する他のチームにとって、これらの成果物が役に立つことを願っています。
まだやるべきことはあります。EMOは、大規模な疎モデルをよりモジュール化可能にするための初期の一歩ですが、未解決の問いは数多く残っています。エキスパート部分集合をより良く選択し、組み合わせるにはどうすべきか。全体モデルを壊さずにモジュールを更新するにはどうすべきか。より優れた解釈可能性と制御のために、モジュール構造をどのように活用するのか。これらのモデルを公開することで、コミュニティがこれらの問いを研究し、より展開しやすく、適応しやすく、検査しやすく、組み合わせやすいモジュール型言語モデルの構築へと前進できるはずです。
この記事で言及されたコレクション 1
- +2
