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

要約

無限の文脈長を持つ配列を効率的にモデル化することは、長い間困難な問題であった。これまでのアプローチは、2次関数的な計算複雑性に悩まされるか、長さの汎化における外挿能力が限られていた。本研究では、選択的状態空間モデル（SSM）であるMambaと、スライディングウィンドウアテンション（SWA）をレイヤーごとに組み合わせた、シンプルなハイブリッドアーキテクチャであるSambaを提案する。Sambaは、注意メカニズムによって最近の記憶を正確に呼び出す能力を維持しながら、与えられたシーケンスを選択的にリカレント隠れ状態に圧縮する。我々は、Sambaを3.2Tの訓練トークンで3.8Bのパラメータまで拡張し、様々なベンチマークにおいて最先端のモデルを大幅に上回ることを実証する。Sambaは4Kの長さのシーケンスで事前に訓練され、ゼロショットでは1Mまでのコンテキストの長さでパープレキシティの向上を示す。Sambaは、4K長のシーケンスで微調整された場合、パスキー検索タスクにおいて、256Kのコンテキスト長に効率的に外挿し、完璧な記憶想起を行う。また、難易度の高い電話帳タスクにおいて、フルアテンションモデルと比較して優れた検索外挿を示す。線形時間シーケンスモデルとして、Sambaは128K長のユーザープロンプトに対してグループ化されたクエリーアテンションを持つTransformersと比較して3.73倍のスループットを達成し、無制限のストリーミングで64Kトークンを生成する場合には3.64倍のスピードアップを達成する。オープンソースデータでトレーニングするためのコードは、https://github.com/microsoft/Samba で公開されている。

要約(オリジナル)

Efficiently modeling sequences with infinite context length has long been a challenging problem. Previous approaches have either suffered from quadratic computational complexity or limited extrapolation ability in 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 recent memories with the attention mechanism. We scale Samba up to 3.8B parameters with 3.2T training tokens and demonstrate that it significantly outperforms state-of-the-art models across a variety of benchmarks. Pretrained on sequences of 4K length, Samba shows improved perplexity in context lengths of up to 1M in zero-shot. When finetuned on 4K-length sequences, Samba efficiently extrapolates to a 256K context length with perfect memory recall on the Passkey Retrieval task, and exhibits superior retrieval extrapolation on the challenging Phonebook task compared to full-attention models. As a linear-time sequence model, Samba achieves a 3.73x higher throughput compared to Transformers with grouped-query attention for user prompts of 128K length, and a 3.64x speedup when generating 64K tokens with unlimited streaming. Our code for training on open source data is publicly available at https://github.com/microsoft/Samba.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL