MoE と Mamba は協力して状態空間モデルを数十億のパラメータに拡張します

State Space Model (SSM) は大きな注目を集めているテクノロジーであり、Transformer の代替として考えられています。 Transformer と比較して、SSM は長いコンテキスト タスクを処理するときに線形時間推論を実現でき、並列トレーニングと優れたパフォーマンスを備えています。特に、選択的 SSM とハードウェアを意識した設計に基づく Mamba は、優れたパフォーマンスを示しており、アテンションベースの Transformer アーキテクチャの強力な代替手段の 1 つとなっています。

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

最近、ポーランドの研究チームは、SSM をハイブリッド エキスパート システム (MoE/専門家の混合) と組み合わせると、SSM の大規模な拡張が期待できることを発見しました。 MoE は、Transformer を拡張するために一般的に使用されるテクノロジです。たとえば、最近の Mixtral モデルではこのテクノロジが使用されています。Heart of the Machine の記事 を参照してください。

このポーランドの研究チームによる研​​究結果は、Mamba とハイブリッド エキスパート レイヤーを組み合わせたモデルである MoE-Mamba です。

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

MoE -Mamba は SSM と MoE の効率を同時に向上させることができます。またチームは、専門家の数が異なる場合でも MoE-Mamba が予測どおりに動作することも発見しました。

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

#関連研究

##状態空間モデル

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

Mamba

Mamba は SSM に基づいて構築されたモデルであり、線形時間推論速度を達成できます (コンテキストのため)また、ハードウェアを意識した設計により効率的なトレーニング プロセスも実現します。 Mamba は、ループの逐次性の影響を軽減する作業効率の高い並列スキャン アプローチを使用し、融合 GPU 操作により拡張状態を実装する必要がなくなります。バックプロパゲーションに必要な中間状態は保存されませんが、バックワード パス中に再計算されるため、メモリ要件が軽減されます。 Mamba のアテンション メカニズムに対する利点は、計算の複雑さが軽減されるだけでなく、メモリ使用量がコンテキストの長さに依存しないため、推論段階で特に重要です。

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

研究によると、Mamba にはより長いコンテキスト (最大 100 万のトークン) を効率的に利用する能力があり、コンテキストの長さが増加するにつれて、トレーニング前の混乱も改善されます。 Mamba モデルは、Mamba ブロックを積み重ねて構成されており、NLP、ゲノミクス、オーディオなどのさまざまな分野で非常に優れた成果を上げています。そのパフォーマンスは、既存の Transformer モデルに匹敵し、それを超えています。したがって、Mamba は一般的なシーケンス モデリング バックボーン モデルの有力な候補モデルとなっています。詳しくは、「5 倍のスループット、パフォーマンスが Transformer を完全に取り囲む: 新アーキテクチャ Mamba が AI の輪を爆発させる##」を参照してください。 #》。

ハイブリッド エキスパート

混合エキスパート (MoE) テクノロジは、モデルのパラメーターの数を大幅に増やすことができます。同時に、モデルの推論とトレーニングに必要な FLOP には影響しません。 MoE は 1991 年に Jacobs らによって初めて提案され、2017 年に Shazeer らによって NLP タスクに使用され始めました。

MoE には利点があります。アクティベーションが非常にまばらであり、処理される各トークンについて、モデルのパラメーターのごく一部のみが使用されます。 Transformer の前方層は、その計算要件により、いくつかの MoE 技術の標準ターゲットとなっています。

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

Fedus らは、2022 年の論文「スイッチ トランスフォーマー: シンプルで効率的なスパース性を備えた兆パラメータ モデルへの拡張」でスイッチを提案しました。これは、各トークンがルーティングされることを組み合わせたトークン選択アーキテクチャです。 1 人の専門家 (K=1) に報告し、彼らはこの方法を使用して、Transformer パラメータのサイズを 1.6 兆まで拡張することに成功しました。ポーランドのこのチームも、実験でこの MoE 設計を使用しました。

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

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

特に言及する価値があるのは、Mistral のオープンソース Mixtral 8× 7B です。そのパフォーマンスは LLaMa 2 70B に匹敵しますが、必要な推論コンピューティングの予算は LLaMa 2 70B の約 6 分の 1 にすぎません。

モデル アーキテクチャ

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

MoE-Mamba は、両方のアーキテクチャの互換性を利用します。図 2 に示すように、MoE-Mamba では、すべての間隔の Mamba 層がスイッチベースの MoE フィードフォワード層に置き換えられます。

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

MoE-Mamba は、Mamba 層によって実行される各トークンの無条件処理と MoE 層によって実行される条件付き処理を分離します。無条件処理は、シーケンスのコンテキスト全体を効率的に統合できます。内部表現では、条件付き処理では各トークンに最も関連性の高い専門家を使用できます。条件付き処理と無条件処理を交互に行うというこのアイデアは、一部の MoE ベースのモデルに適用されていますが、通常は基本層と MoE フィードフォワード層を交互に使用します。

主な結果

トレーニング設定

チームは 5 つを比較しましたさまざまな設定: Basic Transformer、Mamba、Mamba-MLP、MoE、MoE-Mamba。

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

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

#結果

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

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

アブレーション研究

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

図 3 は、さまざまな数のエキスパートを使用した場合のトレーニングの実行手順を示しています。

#表 2 は、100k ステップ後の結果を示しています。

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

以上がMoE と Mamba は協力して状態空間モデルを数十億のパラメータに拡張しますの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。