Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling

要約

無限のコンテキスト長を持つシーケンスを効率的にモデル化することは、長年にわたる問題でした。
過去の研究は、二次計算の複雑さ、または長さの一般化における外挿能力の制限のいずれかに悩まされていました。
この研究では、選択的状態空間モデル (SSM) である Mamba とスライディング ウィンドウ アテンション (SWA) をレイヤーごとに組み合わせたシンプルなハイブリッド アーキテクチャである Samba を紹介します。
Samba は、アテンション メカニズムによって記憶を正確に呼び出す機能を維持しながら、特定のシーケンスを選択的に繰り返し隠れた状態に圧縮します。
3.2T のトレーニング トークンを使用して Samba を 3.8B パラメータまでスケールし、幅広いベンチマークで Samba が純粋な注意や SSM に基づく最先端のモデルよりも大幅に優れていることを示しました。
4K 長のシーケンスでトレーニングすると、Samba は完全なメモリ呼び出しで 256K のコンテキスト長まで効率的に外挿され、最大 1M のコンテキスト長まで改善されたトークン予測を示すことができます。
線形時間シーケンス モデルとして、Samba は 128K の長さのユーザー プロンプトを処理する際に、グループ化されたクエリ アテンションを備えた Transformers と比較して 3.73 倍高いスループットを実現し、無制限のストリーミングで 64K トークンを生成する際には 3.64 倍の高速化を実現します。
Samba のサンプル実装は、https://github.com/microsoft/Samba で公開されています。

要約(オリジナル)

Efficiently modeling sequences with infinite context length has been a long-standing problem. Past works suffer from either the quadratic computation complexity or the limited extrapolation ability on length generalization. In this work, we present Samba, a simple hybrid architecture that layer-wise combines Mamba, a selective State Space Model (SSM), with Sliding Window Attention (SWA). Samba selectively compresses a given sequence into recurrent hidden states while still maintaining the ability to precisely recall memories with the attention mechanism. We scale Samba up to 3.8B parameters with 3.2T training tokens and show that Samba substantially outperforms the state-of-the-art models based on pure attention or SSMs on a wide range of benchmarks. When trained on 4K length sequences, Samba can be efficiently extrapolated to 256K context length with perfect memory recall and show improved token predictions up to 1M context length. As a linear-time sequence model, Samba enjoys a 3.73x higher throughput compared to Transformers with grouped-query attention when processing user prompts of 128K length, and 3.64x speedup when generating 64K tokens with unlimited streaming. A sample implementation of Samba is publicly available in https://github.com/microsoft/Samba.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google