Samba: 効率的な無制限コンテキスト言語モデリングのためのシンプルなハイブリッド状態空間モデル

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

June 11, 2024

著者: Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, Weizhu Chen

cs.AI

要旨

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

English

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.