AI界に新たな寵児が誕生—それはトランスフォーマーではない

リンク一覧

要約と1. はじめに

2 状態空間モデル

3 選択的状態空間モデルと3.1 動機: 圧縮手段としての選択

3.2 選択によるSSMの改善

3.3 選択的SSMの効率的な実装

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

3.5 選択メカニズムの特性

3.6 追加モデルの詳細

4 実証的評価と4.1 合成タスク

4.2 言語モデル

4.3 DNAモデリング

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

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

4.6 モデルアブレーション

5 議論

6 結論、謝辞および参考文献

議論：選択メカニズム

B 関連研究と B.1 S4 の派生型と派生

B.2 SSMアーキテクチャ

B.3 RNNとの関係

B.4 線形注意と B.5 長いコンテキストモデル

C 選択的 SSM のメカニズム

D 選択的 SSM のためのハードウェア対応アルゴリズム

E 実験の詳細と追加結果およびE.1 合成タスク

E.2 言語モデル

E.3 DNAモデリング

E.4 オーディオの詳細

E.5 効率ベンチマーク

抽象的な

現在、ディープラーニングのエキサイティングなアプリケーションのほとんどを支えている基礎モデルは、ほぼ例外なく、Transformer アーキテクチャとそのコアとなるアテンション モジュールに基づいています。長いシーケンスでの Transformer の計算効率の悪さに対処するために、線形アテンション、ゲート付き畳み込みおよび再帰モデル、構造化状態空間モデル (SSM) などの多くのサブ 2 次時間アーキテクチャが開発されてきましたが、言語などの重要なモダリティではアテンションほどのパフォーマンスを発揮していません。私たちは、このようなモデルの主な弱点はコンテンツ ベースの推論を実行できないことだと特定し、いくつかの改善を行いました。まず、SSM パラメータを入力の関数にするだけで、離散モダリティに関する弱点に対処し、モデルが現在のトークンに応じてシーケンス長の次元に沿って情報を選択的に伝播または忘れることができるようにします。次に、この変更によって効率的な畳み込みが使用できなくなりますが、再帰モードでハードウェア対応の並列アルゴリズムを設計します。これらの選択的 SSM を、アテンションや MLP ブロック (Mamba) なしで、簡素化されたエンドツーエンドのニューラル ネットワーク アーキテクチャに統合します。 Mamba は高速な推論 (Transformer の 5 倍のスループット) とシーケンス長の線形スケーリングを誇り、最大 100 万長のシーケンスまでの実際のデータでパフォーマンスが向上します。一般的なシーケンス モデルのバックボーンとして、Mamba は言語、オーディオ、ゲノミクスなどの複数のモダリティにわたって最先端のパフォーマンスを実現します。言語モデリングでは、当社の Mamba-3B モデルは、事前トレーニングとダウンストリーム評価の両方で、同じサイズの Transformer よりも優れ、2 倍のサイズの Transformer に匹敵します。

1 はじめに

基礎モデル (FM) は、大規模なデータで事前トレーニングされ、その後下流のタスクに適応された大規模モデルであり、現代の機械学習における効果的なパラダイムとして登場しました。これらの FM のバックボーンは、多くの場合、言語、画像、音声、オーディオ、時系列、ゲノミクスなど、さまざまなドメインからの任意の入力シーケンスで動作するシーケンス モデルです (Brown ら 2020 年、Dosovitskiy ら 2020 年、Ismail Fawaz ら 2019 年、Oord ら 2016 年、Poli ら 2023 年、Sutskever、Vinyals、Quoc V Le 2014 年)。この概念はモデル アーキテクチャの特定の選択に依存しませんが、最新の FM は主に 1 種類のシーケンス モデル、つまり Transformer (Vaswani et al. 2017) とそのコアとなる注意層 (Bahdanau、Cho、および Bengio 2015) に基づいています。自己注意の有効性は、コンテキスト ウィンドウ内で情報を密にルーティングし、複雑なデータをモデル化できることに起因しています。ただし、この特性には根本的な欠点があります。有限のウィンドウの外側をモデル化できないことと、ウィンドウの長さに関して 2 次スケーリングすることです。これらの欠点を克服するために、注意のより効率的な変種に関する膨大な研究が行われてきましたが (Tay、Dehghani、Bahri など 2022)、多くの場合、注意の効果を生み出す特性そのものが犠牲になっています。今のところ、これらの変種のいずれも、ドメイン間で大規模に実験的に有効であることが示されていません。最近、構造化状態空間シーケンス モデル (SSM) (Gu、Goel、および Ré 2022; Gu、Johnson、Goel 他 2021) が、シーケンス モデリングの有望なアーキテクチャのクラスとして登場しました。これらのモデルは、古典的な状態空間モデル (Kalman 1960) からインスピレーションを得た、再帰型ニューラル ネットワーク (RNN) と畳み込みニューラル ネットワーク (CNN) の組み合わせとして解釈できます。このクラスのモデルは、シーケンス長が線形またはほぼ線形にスケーリングされた、再帰または畳み込みのいずれかとして非常に効率的に計算できます。さらに、特定のデータ モダリティで長距離依存関係をモデル化するための原理的なメカニズム (Gu、Dao 他 2020) を備えており、Long Range Arena (Tay、Dehghani、Abnar 他 2021) などのベンチマークで優位に立っています。 SSM の多くの種類 (Gu、Goel、Ré 2022; Gu、Gupta、et al. 2022; Gupta、Gu、Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith、Warrington、Linderman 2023) は、オーディオやビジョンなどの連続信号データを含む領域で成功を収めています (Goel et al. 2022; Nguyen、Goel、et al. 2022; Saon、Gupta、Cui 2023)。ただし、テキストなどの離散的で情報密度の高いデータのモデル化にはあまり効果的ではありません。

私たちは、シーケンスの長さに比例してスケーリングしながらトランスフォーマーのモデリング能力を実現するために、いくつかの軸で以前の研究を改良した、新しいクラスの選択的状態空間モデルを提案します。

選択メカニズム。まず、従来のモデルの主な制限を特定します。それは、入力に依存した方法でデータを効率的に選択する機能です (つまり、特定の入力に焦点を合わせるか無視するか)。選択的コピーや誘導ヘッドなどの重要な合成タスクに基づく直感に基づいて、入力に基づいて SSM パラメーターをパラメーター化することで、シンプルな選択メカニズムを設計します。これにより、モデルは無関係な情報をフィルターし、関連する情報を無期限に記憶できます。

ハードウェア対応アルゴリズム。この単純な変更は、モデルの計算に技術的な課題をもたらします。実際、以前のすべての SSM モデルは、計算効率を上げるために、時間と入力に対して不変でなければなりません。私たちは、畳み込みではなくスキャンを使用してモデルを再帰的に計算するハードウェア対応アルゴリズムでこれを克服しますが、GPU メモリ階層の異なるレベル間の IO アクセスを回避するために、拡張された状態は実現しません。結果として得られる実装は、理論上 (すべての畳み込みベースの SSM の疑似線形と比較して、シーケンス長が線形にスケーリング) と最新のハードウェア (A100 GPU で最大 3 倍高速) の両方で以前の方法よりも高速です。

アーキテクチャ。従来の SSM アーキテクチャ (Dao、Fu、Saab 他 2023) の設計と Transformer の MLP ブロックを 1 つのブロックに組み合わせることで、従来のディープ シーケンス モデル アーキテクチャを簡素化し、選択的状態空間を組み込んだシンプルで均質なアーキテクチャ設計 (Mamba) を実現します。選択的 SSM、およびその拡張である Mamba アーキテクチャは、シーケンスで動作する一般的な基礎モデルのバックボーンとして適した重要な特性を備えた完全再帰モデルです。(i) 高品質: 選択性により、言語やゲノミクスなどの高密度モダリティで優れたパフォーマンスが得られます。(ii) 高速なトレーニングと推論: トレーニング中は計算とメモリがシーケンス長に比例して増加し、推論中にモデルを自己回帰的に展開するには、以前の要素のキャッシュを必要としないため、ステップごとに一定の時間しかかかりません。(iii) 長いコンテキスト: 品質と効率性が相まって、シーケンス長 1M までの実際のデータでパフォーマンスが向上します。

私たちは、いくつかの種類のモダリティと設定において、事前トレーニングの品質とドメイン固有のタスクパフォーマンスの両方において、Mamba が一般的なシーケンス FM バックボーンとして潜在的可能性を持っていることを経験的に検証しました。

• 合成。大規模言語モデルの鍵となると提案されているコピーや帰納ヘッドなどの重要な合成タスクでは、Mamba はそれらを簡単に解決できるだけでなく、無限に長い (>100 万トークン) ソリューションを外挿できます。

• オーディオとゲノミクス。Mamba は、オーディオ波形と DNA シーケンスのモデリングにおいて、事前トレーニング品質とダウンストリーム メトリックの両方において、SaShiMi、Hyena、Transformers などの従来の最先端モデルよりも優れています (例: 難しい音声生成データセットの FID を半分以下に削減)。どちらの設定でも、最大 100 万の長さのシーケンスまでの長いコンテキストでパフォーマンスが向上します。

• 言語モデリング。Mamba は、事前トレーニングの難しさと下流の評価の両方で、Transformer 品質のパフォーマンスを真に達成した最初の線形時間シーケンス モデルです。10 億パラメーターまでのスケーリング則により、Mamba は LLaMa (Touvron ら、2023 年) に基づく非常に強力な最新の Transformer トレーニング レシピを含む、広範囲のベースラインのパフォーマンスを上回ることが示されています。当社の Mamba 言語モデルは、同様のサイズの Transformer と比較して 5 倍の生成スループットを備えており、Mamba-3B の品質は 2 倍のサイズの Transformer の品質に匹敵します (たとえば、常識的推論の平均が Pythia-3B と比較して 4 ポイント高く、Pythia-7B を上回っています)。

モデル コードと事前トレーニング済みのチェックポイントは、https://github.com/state-spaces/mamba でオープンソース化されています。

2 状態空間モデル

一般的な状態空間モデル。状態空間モデルという用語は、潜在状態を伴うあらゆる再帰プロセスの概念を表す非常に広い意味を持っていることに注意してください。これは、マルコフ決定プロセス (MDP) (強化学習 (Hafner 他 2020))、動的因果モデリング (DCM) (計算神経科学 (Friston、Harrison、および Penny 2003))、カルマン フィルター (制御 (Kalman 1960))、隠れマルコフ モデル (HMM) と線形動的システム (LDS) (機械学習)、および再帰型 (場合によっては畳み込み型) モデル全般 (ディープラーニング) など、さまざまな分野のさまざまな概念を指すために使用されてきました。

本論文では、全体を通じて「SSM」という用語を構造化SSMまたはS4モデルのクラスのみを指すために使用し（Gu、Goel、およびRé 2022; Gu、Gupta、et al. 2022; Gupta、Gu、およびBerant 2022; Hasani et al. 2023; Ma et al. 2023; Smith、Warrington、およびLinderman 2023）、これらの用語を互換的に使用します。便宜上、線形再帰またはグローバル畳み込みの観点に焦点を当てたモデル（Y. Li et al. 2023; Orvieto et al. 2023; Poli et al. 2023）など、このようなモデルの派生モデルも含め、必要に応じてニュアンスを明確にすることもあります。

SSM アーキテクチャ。SSMは、エンドツーエンドのニューラル ネットワーク アーキテクチャに組み込むことができるスタンドアロンのシーケンス変換です。(SSM アーキテクチャを SSNN と呼ぶこともあります。SSNN は、CNN が線形畳み込み層に相当するのと同様に、SSM 層に相当します。) ここでは、最もよく知られている SSM アーキテクチャのいくつかについて説明します。その多くは、主要なベースラインとしても機能します。

• 線形注意 (Katharopoulos et al. 2020) は、退化した線形 SSM と見なすことができる再帰を含む自己注意の近似です。

• H3 (Dao, Fu, Saab, et al. 2023) はこの再帰を一般化して S4 を使用しました。これは、2 つのゲート接続に挟まれた SSM を持つアーキテクチャとして見ることができます (図 3)。H3 はまた、メインの SSM レイヤーの前に、シフト SSM としてフレーム化された標準的なローカル畳み込みを挿入します。

• Hyena（Poli et al. 2023）はH3と同じアーキテクチャを使用していますが、S4層をMLPパラメータ化されたグローバル畳み込みに置き換えています（Romero et al. 2021）。

• RetNet (Y. Sun et al. 2023) は、アーキテクチャに追加のゲートを追加し、よりシンプルな SSM を使用することで、畳み込みの代わりにマルチヘッドアテンション (MHA) のバリアントを使用して、代替の並列化可能な計算パスを可能にします。

• RWKV (B. Peng et al. 2023) は、別の線形注意近似 (注意のないトランスフォーマー (S. Zhai et al. 2021)) に基づいて言語モデリング用に設計された最近の RNN です。その主な「WKV」メカニズムは LTI 再帰を含み、2 つの SSM の比率として見ることができます。

その他の密接に関連する SSM とアーキテクチャについては、拡張された関連研究 (付録 B) でさらに詳しく説明します。特に、S5 (Smith、Warrington、Linderman 2023)、QRNN (Bradbury 他 2016)、および SRU (Lei 他 2017) を強調します。これらは、私たちのコアとなる選択的 SSM に最も密接に関連する手法であると考えられます。