従来の MoE よりも推論速度が速く、パフォーマンスが高い新世代アーキテクチャが登場しました。

この汎用アーキテクチャはMoE++と呼ばれ、 Yan Shuicheng氏が率いるKunlun Tech 2050 Research Instituteと北京大学のYuan Li氏が率いるチームによって共同で提案されました。

全体的に、MoE++ の革新性は「ゼロ計算コストのエキスパート」の導入にあり、この設計により新しいアーキテクチャに 3 つの大きな利点がもたらされます。

• 計算コストの削減: MoE++ では、各トークンが可変数の FFN エキスパートを使用でき、現在の MoE レイヤーを完全にスキップすることもできます。
• パフォーマンスの向上: 単純なトークンに必要な FFN エキスパートの数を減らすことで、MoE++ ではより多くのエキスパートが複雑なトークンに集中できるようになり、従来の MoE よりも優れたパフォーマンスの可能性を引き出します。
• 計算コストゼロのエキスパートのパラメータは非常に小さく、すべての計算コストゼロのエキスパートを各 GPU に同時に展開できるため、分散 FFN エキスパートの展開によって発生する通信オーバーヘッドやエキスパートの負荷の不均一な問題を回避できます。

さらに、MoE++ では、エキスパートを選択するときに、各トークンが前のレイヤーのルーティング パスを参照することもできます。

実験結果によると、パラメータスケールが 0.6B ～ 7B の LLM では、MoE++ は同じモデルサイズで従来の MoE よりも優れており、エキスパートのスループットは 1.1 ～ 2.1 倍向上します。

さらに、モデルの重みがオープンソース化されました!

MoE++ がこれを実際にどのように実行するかを見ていきましょう。

MoE++ はこれをどうやって実現するのでしょうか?

既存のハイブリッド エキスパート (MoE) メソッドのほとんどは、通常、すべてのトークンに対して固定数の FFN エキスパートをアクティブ化します。

多くの研究では、各トークンは上位2つのFFNエキスパートを選択し、その出力を重み付けして合成し、次の層の入力としています。しかし、すべてのトークンの予測難易度は同じではありません。

たとえば、コンマなどの単純な記号は、FFN スペシャリスト 1 人だけで処理できる場合があります。

特殊なケースでは、特定のトークンが現在の MoE レイヤーの専門家と一致しない場合、レイヤーをバイパスして上位 2 人の FFN 専門家を選択しない方が効率的な場合があります。

この洞察に基づいて、研究チームは、既存の MoE で使用されている固定されたハイブリッド メカニズムにより、トレーニングと推論の効率が低下し、モデルのパフォーマンスが制限される可能性があると考えています。

速度とパフォーマンスの両方を向上させるために、研究チームは MoE++ と呼ばれる汎用的な異種 MoE フレームワークを提案しました。

具体的には、チームは計算コストがゼロの 3 種類の専門家を導入しました。

• ゼロエキスパート、空のベクトルを出力します。
• コピーの専門家は入力をそのまま出力として使用します。
• 定数エキスパートは、入力の代わりにトレーニング可能なベクトルを出力として使用します。

図 1 に示すように、従来の MoE 方式とは異なり、MoE++ では各トークンが可変数の FFN エキスパートを使用でき、定数ベクトル置換を受け入れ、現在の MoE++ レイヤーを完全にスキップすることもできます。

△図1: MoE++と通常のMoEの比較

この異種構造により、モデルのフィッティング能力が向上し、ネットワークの組み合わせ空間が拡張されるため、計算コストが大幅に削減されます。

さらに、研究チームは、前の層のルーティングスコアを現在の層の専門家の選択に統合し、トークンが専門家を選択する際に以前のルーティングパスを参照できるようにすることで、より安定した専門家の割り当てを実現しました。

研究チームは、新しく設計された MoE アーキテクチャは次の基準を満たす必要があると考えています。

1. 単純なトークンを効率的に処理するには、設計はできるだけ単純にする必要があります。
2. 公平な比較を行うには、追加されるパラメータを無視できる範囲内に抑える必要があります。

これらの原則に従って、研究チームは計算コストがゼロの専門家を導入しました。各専門家は最も基本的な操作のみを実行します。

図 2(a) に示すように、チームは、破棄、スキップ、置換操作にそれぞれ対応する、ゼロ エキスパート、コピー エキスパート、定数エキスパートの 3 種類のゼロ計算エキスパートを設計しました。

△図2: MoE++のコアコンポーネント

ゼロエキスパート

最も単純なゼロ計算エキスパートは、現在の入力を破棄するゼロ エキスパートです。

本質的に、ゼロの専門家の存在は、トップ 2 MoE++ レイヤーをトップ 1 MoE++ レイヤーにダウングレードする可能性があります。

具体的には、Zero エキスパートがアクティブ化されると、Top-2 MoE++ レイヤーの出力は、他のエキスパートの単一の出力と同等になります。

このように、ゼロ エキスパートの導入により、単純なトークンと複雑なトークンの両方を処理するモデルの柔軟性が向上します。

コピーエキスパート

コピー エキスパートは入力を直接出力として使用します。これは直感的には、現在の MoE++ レイヤーをスキップすることと同じです。

具体的には、入力トークンが既存のエキスパートとうまく一致しない場合は、MoE++ レイヤーをバイパスする方が有利な場合があります。

常連の専門家

定数エキスパートは入力トークンをトレーニング可能なベクトルに置き換えます。

ただし、完全な置換を行うと、入力トークンの情報が失われます。

この問題に対処するため、研究チームは、置換率を動的に予測するための学習可能な重み行列を導入しました。Constantエキスパートの計算コストは​​非常に低いため、依然として計算コストゼロのエキスパートとして分類されます。

ルーティングスコア残差

MoE++ には異種のエキスパートが含まれるため、通常の MoE と比較してルーターの設計がより重要になります。

これを解決するために、研究チームは、図2(b)に示すように、適切なエキスパートを選択する際に前の層で使用されたパスを考慮するパス認識ルーターを提案しました。

具体的には、MoE++ は学習可能な変換マトリックスを使用して、前のレイヤーからのルーティング スコアを現在のレイヤーのエキスパート選択に組み込みます。

これらのルーティング スコア残差により、各トークンはエキスパートを選択するときに以前のルーティング パスを考慮できるようになります。

MoE++ が MoE より優れているのはなぜですか?

この問題には主に 3 つの原因があります。

1 つ目は、計算リソースの柔軟な割り当てです。

MoE++ は、より少ない FFN スペシャリストを単純なトークンに割り当てることでコンピューティング リソースの割り当てを最適化し、より多くの FFN スペシャリストがより難しいトークンの処理に集中できるようにします。

図 3 に示されているように、研究チームは、MoE++ では動詞が最も多くの FFN エキスパートを活性化し、次に名詞が活性化したのに対し、分割された単語断片は最も少ない FFN エキスパートを活性化したことを発見しました。

これは、MoE++によって、意味的複雑度の低いトークンではFFNエキスパートの数が少なくなり、より多くのエキスパートが意味的複雑度の高いトークンを処理できるようになることを示しています。したがって、MoE++は計算オーバーヘッドを削減するだけでなく、全体的なパフォーマンスも向上させます。

△図3: MoE++における異なるトークンに必要なFFNエキスパートの平均数

第二に、安定したルーティングが重要です。

MoE++ は、前のレイヤーのルーティング スコアを現在のレイヤーの専門家の選択にマージします。

これらのルーティング スコア残差により、各トークンはエキスパートを選択するときに以前のルーティング パスを考慮できます。

図 4 に示すように、ルーティング スコア残差は、異なる MoE++ レイヤー間の接続を効果的に確立し、ルーティング スコアの分散を削減します。

一方、ルーティングスコア残差はルーティングスコアの平均や範囲を変化させません。そのため、ルーティングスコア残差は、MoE++における異種エキスパートアーキテクチャにおいて安定したルーティングを実現するのに役立ちます。

△図4: ルーティングスコア残差がルーティングスコア分布に与える影響

最後に、計算の複雑さが低くなります。

下の表に示すように、MoE++ は通常の MoE よりも理論的な計算複雑度が低くなります。

実験結果

パラメータ数を 0.6B から 7B まで徐々に拡大していく広範な実験結果により、MoE++ メソッドが通常の MoE メソッドよりも大幅に優れていることが実証されました。

同じサイズの標準的な MoE モデルと比較すると、MoE++ はエキスパート スループットが約 15% ～ 111% 向上し、優れたパフォーマンスも実現します。

既存の LLM モデルは通常、大きなトレーニング予算を必要とします。たとえば、OpenMoE-8B/32E は 1.1T トークンを使用し、TinyLlama-1.1B は 3T トークンを使用します。

研究者らは、MoE++モデルのトレーニング予算を1Tトークンまで拡張しました。

研究者らは、MoE++ モデルのパフォーマンスは、活性化パラメータが 2 ～ 3 倍の高密度モデルのパフォーマンスに匹敵することを発見しました。

MoE++ は、最初からトレーニングされ、より多くのトレーニング トークンを使用する、より大規模な MoE モデルである OpenMoE-8B/32E よりも優れていることは注目に値します。

これらの結果は、MoE++ フレームワークが有望な LLM フレームワーク スキームであることを示しています。

タスクレベルの専門家の作業負荷分布の可視化

研究者らはまた、MoE++モデルを用いて、様々なタスクにおける熟練者の負荷スコアを調査しました。これらの視覚化によって、いくつかの興味深い発見が明らかになりました。

1. エキスパート負荷は異なる層、特に隣接する層間で相関関係にあります。例えば、FFNエキスパートの大部分が層jで活性化されている場合、FFNエキスパートが層j+1でも同様に高い割合で活性化される可能性が非常に高くなります。
2. 中間層と比較して、浅層と最終層における専門家の割り当てパターンは、タスク間でより大きなばらつきを示しています。これは、モデルが様々なタスクに適応する際に、中間層ではなく、主に浅層と最終層が重要な役割を果たしていることを示唆しています。今後の研究では、これらの層内により複雑な構造を設計し、モデルの多様なタスクへの適応性を高めることに焦点を当てることができます。
3. トークンごとに活性化されるFFNエキスパートの数はタスクによって大きく異なりますが、タスクが単純であれば必ずしも活性化されるFFNエキスパートの数が少なくなるわけではありません。例えば、ARCチャレンジタスクはARCイージータスクよりも多くのFFNエキスパートを活性化します。これらの結果は、MoE++モデルが、タスク全体の難易度ではなく、知識内容の複雑さとトークンレベルに基づいてエキスパートを割り当てていることを示唆しています。
4. すべてのエキスパートタイプの中で、ゼロエキスパートの平均活性化回数が最も高く、より単純なタスクではさらに高い平均活性化回数を示しました。例えば、ARCイージータスクでは、ARCチャレンジタスクよりもゼロエキスパートの活性化回数が多く見られました。これは、ゼロエキスパートの活性化レベルが、モデルのタスク難易度の指標となる可能性があることを示唆しています。
5. MoE++ モデルのすべてのレイヤーにおいて、異なるタスク トピックに対する専門家の割り当ては大きく異なります。これは、MoE++ モデルが異なる専門家の割り当てパターンを採用して、さまざまなタスク トピックを処理していることを示しています。

論文の宛先:
https://arxiv.org/abs/2410.07348

GitHub アドレス:
https://github.com/SkyworkAI/...

ハギングフェイス URL:
https://huggingface.co/Chat-U...