MoEとMambaが協力し、状態空間モデルを数百億のパラメータに拡張

状態空間モデル (SSM) は、最近注目を集めている Transformer の代替手段です。その利点は、線形時間推論、並列トレーニング、および長いコンテキストのタスクでの強力なパフォーマンスです。選択的 SSM とハードウェア対応設計に基づく Mamba は、パフォーマンスがさらに向上し、アテンションベースの Transformer アーキテクチャの強力な代替手段になります。

最近、一部の研究者は、SSM と Mamba を他の方法と組み合わせて、より強力なアーキテクチャを作成することを検討しています。たとえば、Machine Heart はかつて、「Mamba は Transformer の代わりに使用できますが、組み合わせて使用​​することもできます」と報告しました。

最近、ポーランドの研究チームは、SSM をハイブリッド エキスパート システム (MoE/Mixture of Experts) と組み合わせると、SSM を大規模に拡張できることを発見しました。 MoEは現在Transformerの拡張によく使われている技術です。例えば最近のMixtralモデルはこの技術を使っています。Machine Heartの記事を参照してください。

このポーランドの研究チームによって発表された研究成果は、Mamba と混合エキスパート レイヤーを組み合わせたモデルである MoE-Mamba です。

論文アドレス: https://arxiv.org/pdf/2401.04081.pdf

MoE-Mamba は SSM と MoE の両方の効率を向上させることができます。また研究チームは、専門家の数が変わってもMoE-Mambaが予測通りに動作することを発見した。

チームは、図 1 に示すように、実験的なデモンストレーションも実施しました。結果によると、Mamba と比較して、MoE-Mamba では同じパフォーマンスを達成するために必要なトレーニング ステップが 2.2 倍少なく、Transformer および Transformer-MoE に対する新しい方法の潜在的な利点が強調されています。これらの予備的な結果は、有望な研究の方向性も示しています。SSM は数百億のパラメータにまで拡張できる可能性があります。

関連研究

状態空間モデル

状態空間モデル (SSM) は、シーケンス モデリング用のアーキテクチャのクラスです。これらのモデルのアイデアはサイバネティクスの分野に由来しており、RNN と CNN の組み合わせとして考えることができます。かなりの利点があるにもかかわらず、言語モデリングタスクの主要なアーキテクチャになるのを妨げるいくつかの問題もあります。しかし、最近の研究の進歩により、Deep SSM は計算効率と強力なパフォーマンスを維持しながら、数十億のパラメータに拡張できるようになりました。

マンバ

Mamba は SSM 上に構築されたモデルであり、線形時間の推論速度 (コンテキストの長さに関して) を実現できるほか、ハードウェアを考慮した設計により効率的なトレーニング プロセスも実装します。 Mamba は、ループの順次的な性質の影響を軽減する作業効率の高い並列スキャン アプローチを使用し、融合された GPU 操作により拡張状態を実装する必要がなくなります。バックプロパゲーションに必要な中間状態は保存されず、逆方向パス中に再計算されるため、メモリ要件が削減されます。 Mamba がアテンション メカニズムよりも優れている点は、計算の複雑さが軽減されるだけでなく、メモリ使用量がコンテキストの長さに依存しないことから、推論段階で特に顕著です。

Mamba は、シーケンス モデルの効率性と有効性の間の基本的なトレードオフに対処し、状態圧縮の重要性を強調します。効率的なモデルには小さな状態が必要であり、効果的なモデルにはコンテキストのすべての重要な情報が含まれる状態が必要です。時間的および入力の不変性を必要とする他の SSM とは異なり、Mamba は、シーケンス次元に沿って情報が伝播する方法を制御するための選択メカニズムを導入します。この設計の選択は、選択的コピーや帰納的ヘッド合成などの合成タスクからの直感に触発されたもので、モデルが重要な情報を識別して保持しながら、無関係な情報を除外できるようにします。

調査の結果、Mamba はより長いコンテキスト (最大 100 万トークン) を効率的に利用でき、コンテキストの長さが長くなるにつれて事前トレーニングの難しさが改善することがわかりました。 Mamba モデルは積み重ねられた Mamba ブロックで構成されており、NLP、ゲノミクス、オーディオなど、さまざまな分野で非常に優れた結果を達成しています。そのパフォーマンスは、既存の Transformer モデルに匹敵し、それを上回っています。そのため、Mamba は汎用シーケンス モデリング バックボーン モデルの有力な候補モデルとなっています。詳細は、「スループットが 5 倍、パフォーマンスが Transformer を完全に凌駕: 新アーキテクチャ Mamba が AI 界を席巻」を参照してください。

ミキシングエキスパート

Mixture of Experts (MoE) などの手法を使用すると、モデルの推論とトレーニングに必要な FLOP に影響を与えることなく、モデル パラメータの数を大幅に増やすことができます。 MoE は 1991 年に Jacobs らによって初めて提案され、2017 年に Shazeer らによって NLP タスクに使用されました。

MoE には、アクティベーションがまばらであるという利点があります。つまり、処理されるトークンごとに、モデルのパラメータのごく一部だけが使用されます。計算要件により、Transformer のフォワード レイヤーはさまざまな MoE テクニックの標準ターゲットになっています。

研究コミュニティは、ルーティング プロセスとも呼ばれる、トークンを専門家に割り当てるプロセスである MoE の中心的な問題に対処するために、さまざまな方法を提案してきました。現在、基本的なルーティング アルゴリズムには、Token Choice と Expert Choice の 2 つがあります。前者は各トークンを一定数 (K) の専門家にルーティングし、後者は各専門家に固定数のトークンをルーティングします。

Fedusらが2022年の論文「Switchトランスフォーマー：シンプルで効率的なスパース性を備えた兆パラメータモデルへのスケーリング」で提案したスイッチは、各トークンを単一のエキスパート（K=1）にルーティングするトークン選択アーキテクチャであり、この方法を使用してTransformerのパラメータスケールを1.6兆まで拡張することに成功しました。ポーランドのチームもこの MoE デザインを実験に使用しました。

最近、MoE も OpenMoE などのオープンソース コミュニティに参入し始めています。

プロジェクトアドレス: https://github.com/XueFuzhao/OpenMoE

特に注目すべきは、Mistral のオープンソース Mixtral 8×7B です。これは、LLaMa 2 70B に匹敵するパフォーマンスを持ちながら、後者の推論コンピューティング バジェットの約 6 分の 1 しか必要としません。

モデルアーキテクチャ

Mamba の主な基盤となるメカニズムは Transformer で使用されるアテンション メカニズムとは大きく異なりますが、Mamba は Transformer モデルの高レベルのモジュール ベースの構造を保持しています。このパラダイムを使用すると、同じモジュールで構成される 1 つ以上のレイヤーが互いに積み重ねられ、各レイヤーの出力が残差ストリームに追加されます (図 2 を参照)。その後、この残差ストリームの最終値は、言語モデリング タスクの次のトークンを予測するために使用されます。

MoE-Mamba は、これら 2 つのアーキテクチャの互換性を活用します。図 2 に示すように、MoE-Mamba では、1 つおきの Mamba レイヤーが Switch ベースの MoE フィードフォワード レイヤーに置き換えられています。

しかし、チームはこの設計が「Mamba: 選択的状態空間による線形時間シーケンス モデリング」の設計といくぶん似ていることにも気付きました。後者は Mamba レイヤーとフィードフォワード レイヤーを交互に積み重ねたものですが、結果として得られるモデルは純粋な Mamba よりもわずかに劣っていました。この設計は、図 1 では Mamba-MLP として示されています。

MoE-Mamba は、シーケンスのコンテキスト全体を効率的に内部表現に組み込む Mamba レイヤーによって実行される各トークンの無条件処理と、各トークンに最も関連性の高いエキスパートを使用する MoE レイヤーによって実行される条件付き処理を分離します。条件付き処理と無条件処理を交互に行うというこの考え方は、いくつかの MoE ベースのモデルで使用されていますが、通常は基本レイヤーと MoE フィードフォワード レイヤーが交互に使用されます。

主な結果

トレーニング設定

チームは、基本的な Transformer、Mamba、Mamba-MLP、MoE、MoE-Mamba の 5 つの異なる設定を比較しました。

ほとんどの Transformer では、フィードフォワード層に 8dm² のパラメータが含まれていますが、Mamba の論文では Mamba が小さく (約 6dm²) なっているため、2 つの Mamba 層のパラメータ数は、1 つのフィードフォワード層と 1 つのアテンション層の合計とほぼ同じになります。 Mamba と新しいモデル間でトークンあたりのアクティブ パラメーターの数をほぼ同じに保つために、チームは各エキスパート フォワード レイヤーのサイズを 6dm² に縮小しました。埋め込みレイヤーと埋め込み解除レイヤーを除き、すべてのモデルはトークンごとに約 2,600 万個のパラメータを使用します。トレーニング プロセスでは 65 億トークンが使用され、トレーニング ステップの数は 10 万でした。

トレーニングに使用されるデータセットは英語の C4 データセットであり、タスクは次のトークンを予測することです。テキストは、GPT2 トークナイザーを使用してトークン化されます。表 3 にハイパーパラメータの完全なリストを示します。

結果

表1にトレーニング結果を示します。 MoE-Mamba は、通常の Mamba モデルよりも大幅に優れたパフォーマンスを発揮します。

注目すべきは、MoE-Mamba はトレーニング ステップのわずか 46% で、通常の Mamba と同じレベルの結果を達成することです。学習率は通常の Mamba に合わせて調整されているため、トレーニング プロセスを MoE-Mamba 用に最適化すると、MoE-Mamba のパフォーマンスがさらに向上することが期待できます。

アブレーション研究

Mamba が専門家の数に応じて適切に拡張されるかどうかを評価するために、研究者はさまざまな専門家の数を使用したモデルを比較しました。

図 3 は、異なる数の専門家を使用した場合のトレーニング実行の手順を示しています。

表 2 は 10 万ステップ後の結果を示しています。

これらの結果は、提案された方法が専門家の数に応じて適切に拡張されることを示しています。エキスパートの数が 8 人以上の場合、新しいモデルの最終的なパフォーマンスは通常の Mamba よりも優れています。 Mamba-MLP は通常の Mamba よりも劣っているため、少数の専門家を使用する MoE-Mamba のパフォーマンスは Mamba よりも劣ることが予想されます。専門家の数が 32 人の場合、新しい方法は最良の結果をもたらします。