トランスフォーマーの覇権を覆せ！ CMU プリンストンは致命的なバグを解決し推論速度を 5 倍に向上させる新しい Mamba アーキテクチャを発表

ディープラーニングは新たな時代を迎えた。Transformerの優位性は覆されるのか？

2017年6月12日に発表され、NLPの世界に突如変化をもたらしました。長年自然言語分野を席巻してきたTransformerは、ついに新しいアーキテクチャによって打ち破られようとしています。

Transformer は強力ですが、致命的なバグがあります。コアアテンションレイヤーが長期的なコンテキストに拡張できないのです。

ちょうど今、CMU とプリンストンの研究者が Mamba をリリースしました。この SSM アーキテクチャは、言語モデリングの Transformer に匹敵しますが、推論スループットが 5 倍になりながら線形に拡張されます。

論文アドレス: https://arxiv.org/abs/2312.00752

論文の筆頭著者であるアルバート・グ氏は、二次的注意は情報集約型モデルには不可欠だが、今ではもはや必要ないと述べています。

この論文が発表されるやいなや、AIコミュニティに旋風を巻き起こしました。

Nvidiaの主任科学者ジム・ファン氏は、誰かがTransformerを打倒してくれるのを待ち望んでいたと述べ、長年に渡ってTransformerのシーケンシャルアーキテクチャを置き換えようとしてきたアルバート・グー氏とトリ・ダオ氏に感謝の意を表した。

「君たちは本当に素晴らしい研究をしていて、何か新しいものを次々と生み出している。しばらく休んでもらえないかな？」

「レイカーズのファンはマンバという名前にとても満足していると言っていました！」

著者は、このアーキテクチャが Mamba と名付けられている理由についても説明しています。

- 高速: (1) シーケンスの長さに比例して増加する単純な再帰、(2) ハードウェアを考慮した設計と実装

- 致命的: シーケンスモデリングの問題に致命的な魅力がある

- 彼らが発する「音」も非常に似ています。そのコアメカニズムは、構造化状態空間シーケンスモデル (S4) の最新の進化です...SSSS

パフォーマンスがトランスフォーマーを圧倒？

Mamba は、Albert Gu の以前の「構造化状態空間モデル」に関する研究から派生したもので、強力なループ演算子として考えることができます。これにより、シーケンス長の線形スケーリングと高速な自己回帰デコードが可能になります。

論文アドレス: https://arxiv.org/abs/2111.00396

しかし、これまでの再帰型モデルの欠点は、状態サイズが固定されているためコンテキストを圧縮することが難しいことです。

Mamba の主な貢献は、「選択的 SSM」の導入です。これは、入力に選択的に注意を払ったり無視したりできる S4 の単純な一般化です。

この小さな変更 (いくつかのパラメータを入力の関数にするだけ) により、以前のモデルでは非常に困難だったタスクを即座に解決できるようになります。

たとえば、重要な「連想想起」タスクの解決策を無限の長さに外挿することができます。 (トレーニング長 256、テスト長 1M)

要点は、この変更には重要な効率性のトレードオフが伴い、S4 の元の設計には特定の理由があったということです。

DNA やオーディオなどの他のモダリティからの実際のデータでは、Mamba の事前トレーニング済みパフォーマンスは、HyenaDNA や SaShiMi などの以前の専門的なベースラインを上回ります。

合成データ、DNA データ、オーディオ データを問わず、シーケンスの長さが 100 万を超えると、Mamba のパフォーマンスが向上し続けることは注目に値します。

もう一人の筆頭著者である Tri Dao 氏は、この課題に対処するためにハードウェアを考慮した設計を使用する方法と、Mamba の強力な言語パフォーマンスを紹介しました。

彼は、アルバートが言ったように、状態空間モデル (SSM) の特徴は固定サイズの再帰状態であると述べました。より良いパフォーマンスを実現したい場合は、この状態をより大きくし、より表現力豊かにする必要があります。

残念ながら、実際に大きな状態を再帰を使用して計算するには遅すぎます。

これまで、S4 ベースの SSM は、構造上の仮定 (つまり、線形時間不変性) を立てることでこの問題に対処し、大きな状態を実装せずに同等の「畳み込みモード」計算を実行できるようにしてきました。

しかし今回、CMU とプリンストンの研究者が使用した方法は選択的 SSM であり、これは周期的にしか計算できません。

この計算上のボトルネックを解決するために、研究者らは FlashAttention などの他の効率的なハードウェア対応アルゴリズムで使用される手法を活用しました。

Mamba (および一般的な SSM) の場合、このアプローチではメイン メモリではなく SRAM にのみ拡張状態が実装されることに注意することが重要です。

さらに、スキャン実装は、基本の PyTorch/JAX 実装よりも 30 倍高速で、シーケンスの長さが長くなると、二次 FlashAttention よりも桁違いに高速になります。

さらに、固定サイズの再帰状態 (KV キャッシュなし) を使用しているため、Mamba の LM 推論速度は Transformer よりも 5 倍高速です。

経験的に、2 人の著者が達成した最も重要な結果は言語モデリングに関するものであり、これは以前の SSM (Mamba の前身である H3 など) が対象としていた分野でもあります。

論文アドレス: https://arxiv.org/abs/2212.14052

しかし、現時点では、私の作品はまだトランスフォーマーほど良くはありません。そして彼は、当時は、微調整されたトランスフォーマーに真に対抗できるモデルは存在しなかったと語った。

しかし、突然サプライズが起こりました！

チンチラのスケーリング法則に従って事前トレーニングすると、Mamba は突然、非常に強力な最新の「Transformer++」モデル (Llama モデルに近い) よりも優れたパフォーマンスを発揮します。

3000億トークンでトレーニングした後、Mamba のパフォーマンスは同様のオープンソース モデルよりもはるかに優れています。

最後に、著者らは、ハードウェアを考慮した考え方が新しいアーキテクチャ設計の道を開く可能性があると結論付けています。

将来的には、この新しいアーキテクチャは、Transformer を中心に構築されたハードウェア/ライブラリを活用できるでしょうか?他の分野 (ゲノミクス、オーディオ、ビデオ) でのシーケンス拡張はどのように変化するでしょうか?

この目的のために、著者らは一連のモデル重み（最大 28 億のパラメータ、3000 億のトークンでトレーニング済み）と高速推論コードもリリースしました。

プロジェクトアドレス: https://github.com/state-spaces/mamba

Transformerに勝るアーキテクチャはいかにして生まれたか

今日の基本モデルは、ほぼすべて、Transformer アーキテクチャとそのコア アテンション モジュールに基づいて構築されています。

長いシーケンスを処理する際の Transformer の計算効率の悪さの問題を解決するために、学術コミュニティでは、線形アテンション、ゲート畳み込みおよび再帰モデル、構造化状態空間モデル (SSM) など、2 次時間計算量を持つ多くのアーキテクチャが開発されてきました。

ただし、これらのアーキテクチャは、言語を処理する際に従来の注意モデルほど優れたパフォーマンスを発揮しません。

研究者たちは、これらのモデルの主な弱点は内容に基づく推論を実行するのが難しいことであると発見し、いくつかの改良を加えました。

まず、SSM パラメータを入力データの関数にすることで、離散データ型を扱う際のこのタイプのモデルの欠点に対処できます。

これにより、モデルは現在のトークンに基づいて、シーケンスの長さの次元で情報を選択的に伝播または忘れることができます。

第二に、このような調整によりモデルは効率的な畳み込みを使用できなくなりますが、研究者らはハードウェアに適応する並列アルゴリズムを設計し、それを再帰モードで実装しました。

研究者たちは、この選択的 SSM を、注意メカニズムや MLP (多層パーセプトロン) モジュールさえも必要としない、簡素化されたエンドツーエンドのニューラル ネットワーク アーキテクチャに統合しました。研究者たちはこれを Mamba と名付けました。

Mamba は高速推論に優れており (Transformers より 5 倍高速)、そのパフォーマンスはシーケンスの長さに比例して向上し、最大 100 万の長さのシーケンスを処理する場合にさらに優れたパフォーマンスを発揮します。

汎用シーケンス処理モデルとして、Mamba は言語、オーディオ、ゲノミクスなど複数の分野で最先端のパフォーマンスを実現しています。

言語モデリングの点では、Mamba-3B モデルは、事前トレーニングとその後の評価でパラメータ数が 2 倍である Transformers モデルよりも優れています。

研究者は実証的研究を通じて、基礎モデル (FM) のコア フレームワークとしての Mamba の大きな可能性を検証しました。

この可能性は、事前トレーニングの品質だけでなく、複数のモダリティと環境をカバーするドメイン固有のタスクのパフォーマンスにも反映されています。

- 合成タスク

複製や誘導などの重要な合成タスクでは、Mamba はそれらを簡単に解決するだけでなく、無限長 (> 100 万トークン) のソリューションを推論することもできます。

- オーディオとゲノミクス

オーディオ波形と DNA シーケンスのモデリングにおいて、Mamba は、事前トレーニング品質とダウンストリーム メトリックの両方の点で、SaShiMi、Hyena、Transformers などの以前の SOTA モデルよりも優れています (例: 困難な音声生成データセットで FID が半分以上削減されます)。

どちらの場合も、コンテキストの長さが数百万の長さのシーケンスまで増加すると、パフォーマンスが向上します。

- 言語モデル

Mamba は、事前トレーニングの複雑さと下流のタスク評価の両方で Transformer レベルのパフォーマンスを実現する最初の線形時系列モデルです。

研究者らは、モデルを10億のパラメータにまでスケールアップした後、MambaがLlamaなどのいくつかのベースラインモデルよりも優れていることを実証した。

Mamba 言語モデルは、同じサイズの Transformer の 5 倍の生成スループットを備えており、Mamba-3B の品質は、その 2 倍のサイズの Transformer に匹敵します (Pythia-3B と比較すると、平均常識推論スコアは 4 ポイント高く、Pythia-7B を上回ります)。

選択的状態空間モデル

研究者たちは、合成タスクからの直感を利用して選択メカニズムを導き出し、このメカニズムを状態空間モデルに組み込む方法を説明しています。結果として得られる時間変動 SSM では畳み込みを使用できないため、畳み込みを効率的に計算する方法に関して技術的な課題が生じます。

研究者たちは、最新のハードウェア上のメモリ階層を活用するハードウェア対応アルゴリズムを使用してこの問題を克服しました。次に研究者らは、アテンション メカニズムや MLP モジュールさえも使用しない単純な SSM アーキテクチャについて説明します。最後に、研究者らは選択メカニズムのいくつかの追加特性について議論します。

動機: 圧縮手段としての選択

研究者たちは、シーケンスモデリングにおける基本的な問題はコンテキストをより小さな状態に圧縮することだと考えています。彼らは、この観点から、一般的なシーケンス モデルのトレードオフを検討します。

たとえば、注意はいくつかの点で非常に効果的ですが、コンテキストをまったく圧縮しないため、他の点では効果的ではありません。この点から、自己回帰推論にはコンテキスト全体の明示的な保存 (つまり、KV キャッシュ) が必要であり、これが直接的に Transformer の線形時間推論と二次時間トレーニングの遅延につながることがわかります。

一方、再帰型モデルは状態数が有限であるため効率的であり、推論時間は一定で、トレーニング時間は線形になります。

しかし、注意の有効性は、そのような状態が文脈を圧縮する程度によって制限されます。

この原理を理解するために、研究者は 2 つの合成タスクの例を実行することに重点を置きました (下の図 2 を参照)。

選択的コピー タスクは、記憶するトークンの場所を変更することで、一般的なコピー タスクを変更します。関連するタグ (色付き) を記憶し、無関係なタグ (白) を除外するには、コンテンツを認識した推論が必要です。

誘導ヘッドタスクは、LLM の文脈学習能力の大部分を説明するために以前の研究で仮説が立てられたよく知られたメカニズムです。適切なコンテキスト (黒) で正しい出力をいつ生成するかを知るには、コンテキスト認識推論が必要です。

これらのタスクは、LTI モデルの障害モードを明らかにします。再帰的な観点から見ると、それらの一定のダイナミクス（例えば、（2）の（A、B）遷移）は、文脈から正しい情報を選択したり、入力に関連して隠れ状態がシーケンスに沿って渡される方法に影響を与えたりすることを可能にしません。

畳み込みの観点から見ると、グローバル畳み込みは時間的認識のみを必要とするため一般的なコピータスクを解決できますが、コンテンツ認識が不足しているため選択的コピータスクでは困難があることはよく知られています (上図を参照)。

具体的には、入力と出力の間隔は変化するため、静的畳み込みカーネルではモデル化できません。

要約すると、シーケンス モデルの効率性と有効性のトレードオフは、状態を圧縮する度合いによって特徴付けられます。効率的なモデルは小さな状態を持つ必要があり、優れたモデルでは、この小さな状態にコンテキスト内のすべての必要な情報が含まれている必要があります。

代わりに、研究者が構築したシーケンス モデルの基本原理は選択性、つまり一連の状態への入力に焦点を合わせたりフィルター処理したりするコンテキスト認識機能です。

特に、選択メカニズムは、シーケンス次元に沿って情報がどのように伝播または相互作用するかを制御します。

選択による SSM の改善 モデルに選択を組み込む 1 つの方法は、シーケンスの相互作用に影響するパラメーター (RNN の再帰ダイナミクスや CNN の畳み込みカーネルなど) を入力に依存させることです。

アルゴリズム 1 と 2 は、研究者が使用する主な選択メカニズムを示しています。

主な違いは、いくつかのパラメータ Δ、B、C を入力の関数にし、それに応じてテンソル全体の形状が変化するという点です。

これらのパラメータに長さの次元が追加され、モデルが時間不変から時間変動に変わったことに注意してください。

これにより畳み込みとの同等性が失われ、効率に影響します。

簡素化されたSSMアーキテクチャ

構造化 SSM と同様に、選択的 SSM は「ニューラル ネットワークに柔軟に組み込むことができる」独立したシーケンス変換です。

H3 アーキテクチャは最も有名な SSM アーキテクチャの基礎であり、通常は線形アテンションにヒントを得たブロックと MLP (多層パーセプトロン) ブロックがインターリーブされて構成されます。研究者たちは、2 セットの同質スタックを 1 つのコンポーネントに組み合わせることでアーキテクチャを簡素化しました (下の図を参照)。

このアプローチは、Gated Attention Unit (GAU) にヒントを得ています。このアーキテクチャでは、制御可能な拡張係数を通じてモデルの寸法を拡張します。各ブロックでは、ほとんどのパラメータ (3ED^2) が線形投影にありますが、内部 SSM の寄与は少なくなります。比較すると、SSM パラメータの数ははるかに少なくなります。

研究者たちは、このブロックを標準的な正規化と残余接続と交互に繰り返し、Mamba アーキテクチャを形成しました。

実験では、常に x を E = 2 に設定し、2 つのブロック スタックを使用して、Transformer インターリーブ MHA (マルチヘッド アテンション) ブロックと MLP ブロックの 122 個のパラメーターを一致させます。

研究者らは、ゲート付き MLP の一般的な「SwiGLU」バリアントに触発された SiLU/Swish 活性化関数を使用しました。最後に、研究者らは、RetNet が同様の位置で正規化レイヤーを使用しているという事実に基づいて、オプションの正規化レイヤーも使用しました。

選択メカニズムは、より伝統的な RNN や CNN、さまざまなパラメーター (アルゴリズム 2 の A など)、さまざまな変換の使用など、さまざまな方法で適用できるより広範な概念です。

実証的評価

合成タスク: 選択的コピー

コピータスクは、シーケンス モデル、特に再帰モデルの記憶能力をテストするために使用される古典的な合成タスクです。

LTI SSM (線形再帰およびグローバル畳み込み) は、推論データではなく時間のみに焦点を当てることで、このタスクを簡単に解決できます。たとえば、正確な長さの畳み込みカーネルを構築します (図 2)。

これに対応して、選択的レプリケーション タスクは、トークン間の間隔をランダムに変更することで、このショートカットを防ぐことができます。

表 1 は、H3 や Mamba などのゲーティング アーキテクチャではパフォーマンスを部分的にしか向上できないのに対し、選択メカニズム (つまり、S4 から S6 への向上) では、特に強力なアーキテクチャと組み合わせると、この問題を簡単に解決できることを示しています。

合成タスク: 誘導ヘッド

帰納的ヘッドは、機械的な解釈可能性の観点からは比較的単純なタスクですが、予想外に大規模言語モデル (LLM) の文脈学習機能を予測します。

このタスクでは、モデルが連想想起と複製アクションを実行する必要があります。たとえば、モデルが以前にシーケンス内で「ハリー ポッター」などのフレーズに遭遇した場合、同じシーケンス内で「ハリー」が再び出現すると、モデルは履歴情報を確認して「ポッター」を予測できる必要があります。

表 2 は、Mamba モデル、より正確にはその選択的 SSM レイヤーが、関連するトークンを選択的に記憶し、その間にある他のトークンを無視する機能により、タスクを完璧に実行できることを示しています。

さらに、これはトレーニング中に遭遇した長さの 4000 倍である 100 万の長さのシーケンスに完全に一般化されます。比較すると、他の方法の一般化能力は 2 倍を超えることはできません。

言語モデル

研究者らは、Mamba を標準の Transformer アーキテクチャ (GPT-3 アーキテクチャ) および現在の最先端の Transformer (Transformer++) と比較しました。

後者は PaLM および LLaMa アーキテクチャに基づいており、回転埋め込み、SwiGLU MLP、LayerNorm の RMSNorm への置き換え、線形バイアスのキャンセル、より高い学習率の使用などの機能を備えています。

図 4 は、Mamba が、パラメータ サイズが約 1 億 2500 万から約 13 億の範囲で最も強力な Transformer アーキテクチャ (Transformer++) と同等のパフォーマンスを発揮する最初の注意不要モデルであることを示しています。

表 3 は、一連のダウンストリーム ゼロ ショット評価タスクにおける Mamba のパフォーマンスを示しています。

このうち、Mamba はトレーニング中に Pythia および RWKV と同じトークナイザー、データセット、トレーニング長 (300B トークン) を使用しました。

トレーニング中のMambaとPythiaのコンテキスト長は2048であるのに対し、RWKVのコンテキスト長は1024であることに注意してください。

DNAモデリング

大規模言語モデルの成功により、基本モデルのパラダイムをゲノミクスに適用する研究が始まりました。

DNA は、特定の語彙を持つ一連の個別の記号で構成されており、モデル化には長距離の依存関係も必要となるため、言語に例えられています。

私たちは、長配列 DNA モデルに関する最近の研究と同じ文脈で、事前トレーニング済みで微調整された FM バックボーンとして Mamba を研究します。

事前トレーニングに関しては、研究者は基本的に標準的な因果言語モデリング（次のトークン予測）設定に従います。

データセットに関しては、基本的に HyenaDNA の設定に従います。事前トレーニングには HG38 データセットを使用します。このデータセットは単一のヒトゲノムで構成され、トレーニング分割には約 45 億のトークン (DNA 塩基対) が含まれています。

結果は図 5 (左) に示されています。Mamba の事前トレーニングの難しさはモデル サイズの増加とともに着実に増加し、Mamba のスケーラビリティは HvenaDNA や Transformer++ よりも優れています。

たとえば、最大モデル サイズ ≈ 40M パラメータでは、曲線は、Mamba が 3 ～ 4 倍少ないパラメータで Transformer++ および HvenaDNA モデルに匹敵できることを示しています。

さらに、図 5 (右) は、Mamba がより長いコンテキスト、最大 100 万の非常に長いシーケンスも利用できること、およびコンテキストの増加に伴って事前トレーニングの難しさが増加することを示しています。

一方、ハイエナ DNA モデルは、配列の長さが長くなるにつれて悪化します。

畳み込みの観点から見ると、非常に長い畳み込みカーネルは長いシーケンスにわたってすべての情報を集約しています。

図 6 は類人猿 DNA の分類であり、同じコンテキスト長の事前トレーニング済みモデルを使用して長さ 2^10 ～ 2^20 のシーケンスを微調整した後の精度を示しています。

オーディオモデリングと生成

オーディオ波形処理の分野では、主な比較対象は SaShiMi アーキテクチャです。モデルには以下が含まれます:

1. 2 段階のプーリングを備えた U-Net バックボーン。各段階で、プーリング係数 p を使用してモデルの次元 D が 2 倍になります。

2. 各ステージでは、S4 モジュールと MLP モジュールが交互に使用されます。

研究者らは、長いコンテキストでの自己回帰事前トレーニングのために、標準的なピアノ音楽データセットである YouTubeMix を使用して評価しました。このデータセットには、サンプリング レート 16000Hz の 4 時間のソロ ピアノ音楽が含まれています。

図 7 は、計算の複雑さを変えずに、トレーニング シーケンスの長さを 8192 (2^13) から約 1000000 (2^20) に増やした場合の効果を示しています。

Mamba と SaShiMi (S4+MLP) のベースライン モデルのパフォーマンスは、コンテキストの長さが長くなるにつれて着実に向上します。その中でも、Mamba はプロセス全体を通じて優れており、シーケンスが長くなるほどその利点は明らかになります。

自己回帰音声生成に関しては、ベンチマーク音声生成データセット SC09 を使用して評価します。これは、16,000 Hz のサンプリング周波数を持ち、「0」から「9」までの数字とさまざまな特徴を含む 1 秒の音声クリップで構成されています。

表 4 は、Mamba-UNet と、WaveNet、SampleRNN、WaveGAN、DiffWave、SaShiMi などの一連のベースライン モデルの自動評価結果を示しています。

ご覧のとおり、小規模の Mamba モデルは、最先端の GAN および拡散ベースの技術を使用する大規模なモデルよりも優れています。同じパラメータスケールを持つ Mamba モデルは、忠実度の点ではるかに優れたパフォーマンスを発揮します。

表 5 では、小規模の Mamba モデルを使用して、外部ステージと中央ステージのさまざまなアーキテクチャの組み合わせによる効果を調べています。

調査の結果、Mamba モデルは外側のブロックと中央のブロックの両方で S4+​​MLP アーキテクチャよりも優れたパフォーマンスを発揮し、中央のブロックでのパフォーマンス順位は Mamba > S4+MLP > MHA+MLP であることがわかりました。

速度とメモリのベンチマーク

図 8 は、スキャン操作 (状態拡張 N = 16) の速度と Mamba エンドツーエンド推論スループット ベンチマークを示しています。

結果は、シーケンスの長さが 2k を超えると、効率的な SSM スキャンが現在の最良のアテンション メカニズムである FlashAttention-2 よりも高速になることを示しています。さらに、PyTorch の標準スキャン実装と比較すると、速度は最大 20 ～ 40 倍向上します。

キー値 (KV) キャッシュがないため、Mamba はより大きなバッチ サイズをサポートでき、同じサイズの Transformer よりも 4 ～ 5 倍高い推論スループットを実現します。

たとえば、69 億のパラメータを持つトレーニングされていない Mamba (Mamba-6.9B) は、推論処理能力の点で 5 倍小さい、わずか 13 億のパラメータを持つ Transformer モデルよりも優れたパフォーマンスを発揮します。

ほとんどのディープシーケンシャルモデルと同様に、メモリ使用量は活性化テンソルのサイズに比例します。表 15 は、Mamba のメモリ要件が最適化された Transformer のメモリ要件と同等であることを示しています。

125M モデルを単一の A100 80GB GPU でトレーニングするためのメモリ要件

論文の最後で、著者らは、さまざまな分野の基本モデルを構築する上での選択的状態空間モデルの幅広い応用性は非常に興味深いと述べています。

さまざまな実験結果から、Mamba は一般的なシーケンス モデルの主流のフレームワークになる可能性が高く、Transformer と競合する可能性さえあることがわかります。