プリンストン大学とCMUがマンバ配列モデルで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 関連研究

C 選択的 SSM のメカニズム

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

E 実験の詳細と追加結果

抽象的な

現在、ディープラーニングのエキサイティングなアプリケーションのほとんどを支えている基礎モデルは、ほぼ例外なく、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、et al. 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 を上回っています)。

図 1: (概要) 構造化 SSM は、入力 x の各チャネル (例: D = 5) を、より高次元の潜在状態 ℎ (例: N = 4) を介して出力 y に個別にマッピングします。従来の SSM では、時間不変性を必要とする巧妙な代替計算パスによって、この大きな有効状態 (DN、バッチ サイズ B、シーケンス長 L の積) が実現されることを回避していました。つまり、(∆、A、B、C) パラメーターは時間を通じて一定です。当社の選択メカニズムでは、入力依存のダイナミクスが追加され、拡張された状態を GPU メモリ階層のより効率的なレベルでのみ実現するための、慎重なハードウェア対応アルゴリズムも必要になります。