15日前

Mamba:選択的状態空間を用いた線形時間シーケンスモデリング

Albert Gu, Tri Dao
Mamba:選択的状態空間を用いた線形時間シーケンスモデリング
要約

現在、深層学習における多くの革新的な応用を支えている基盤モデル(Foundation Models)は、ほぼすべてがTransformerアーキテクチャおよびその中心となるアテンションモジュールに基づいている。長文系列におけるTransformerの計算効率の低さに対処するため、線形アテンションやゲート付き畳み込み、再帰型モデル、構造化状態空間モデル(SSM)など、サブ二次時間計算量を実現する多くのアーキテクチャが開発されてきたが、言語などの重要なモダリティにおいてはアテンションの性能に及ばないことが続いていた。本研究では、こうしたモデルの根本的な課題が「コンテンツに基づく推論」が行えない点にあると指摘し、複数の改良を提案する。第一に、SSMのパラメータを入力のアドレス関数として定義することで、離散的モダリティにおけるこの欠点を克服し、現在のトークンに応じて系列長方向に情報を選択的に伝搬または忘却する能力を獲得した。第二に、この変更により効率的な畳み込みが利用できなくなるものの、再帰モードにおいてハードウェアに配慮した並列アルゴリズムを設計した。これらの選択的SSMを、アテンションやMLPブロックを一切含まない簡略化されたエンドツーエンドのニューラルネットワークアーキテクチャに統合した(Mamba)。Mambaは高速な推論を実現し、Transformerに比べて5倍のスループットを達成し、系列長に対して線形スケーラビリティを示す。実データ上で100万単位の長さの系列にも対応し、性能が継続的に向上する。汎用的な系列モデルのバックボーンとして、Mambaは言語、音声、ゲノミクスなど複数のモダリティで最先端の性能を達成した。言語モデリングにおいて、Mamba-3Bモデルは同サイズのTransformerを上回り、サイズが2倍のTransformerと同等の性能を、事前学習および下流タスク評価の両方で達成した。

Mamba:選択的状態空間を用いた線形時間シーケンスモデリング | 最新論文 | HyperAI超神経