Mamba: Linear-Time Sequence Modeling with Selective State Spaces

要約

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

要約(オリジナル)

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers’ computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5$\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google