[論文レビュー] Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling

Liliang Ren, Yang Liu+4 more|arXiv (Cornell University)|Jun 11, 2024

Topic Modeling被引用数 7

ひとこと要約

Samba は Mamba（a selective State Space Model）と Sliding Window Attention を共同で統合し、線形時間計算で無限の文脈長を実現し、ベンチマーク全体で強力な性能と大幅なスループットの向上を提供します。

ABSTRACT

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.

研究の動機と目的

• トレーニング時の文脈長を超えて外挿する、効率的な言語モデリングを動機づける。
• SSMsとアテンションを層ごとに組み合わせた、シンプルなハイブリッドアーキテクチャを提案する。
• 長い文脈能力を持つ数十億パラメータ超のモデルへのスケーラビリティを実証する。
• 推論、数学、コーディング、長文脈検索など多様なタスクにわたる性能を評価する。
• ハイブリッド設計を理解するための訓練および推論の効率性とアブレーション研究を分析する。

提案手法

• Samba を導入する、Mamba (SSM) と Sliding Window Attention (SWA) の層別ハイブリッド。
• 入力依存の選択的状態空間を介して時間依存的意味論を捕捉するために Mamba レイヤを使用する。
• 2048 ウィンドウの SWA を組み込み、線形時間フレームワーク内で正確なメモリ検索を実現する。
• 非線形変換には SwiGLU MLP を活用し、異なる情報ストリームには別々の MLP を用いる。
• 421M、1.3B、1.7B、3.8B パラメータのモデルを最大 3.2T トークンで訓練し、多様なベンチマークで評価する。
• 効率的な長文脈の外挿テスト（最大 256K プロンプトと 1M コンテキスト）およびスループット比較を実施する。

(a) Perplexity on the test set of Proof-Pile

実験結果

リサーチクエスチョン

• RQ1 入力依存の SSM と sliding window attention のシンプルな層ごとのハイブリッドで、線形時間計算で無限の文脈長を達成できるか？
• RQ2 Samba は長文脈および短文脈のベンチマークで、純粋なアテンションモデルおよび純粋な SSM モデルとどのように比較されるか？
• RQ3 長文脈言語モデリングにおいて、どのようなアーキテクチャの選択（例：アテンションと再帰の分布）が最良の性能と効率をもたらすのか？
• RQ4 文脈長を外挿する際、記憶のリコールと検索をどの程度保持または強化できるか？
• RQ5 プロンプト処理と生成のための効率性とスループットへの影響はどのようなものか？

主な発見

• Samba は 8B パラメータまでの強力なオープンソース LLM と比較して、広範なベンチマーク群で平均的に最先端のパフォーマンスを達成する。
• 3.8B Samba モデルは MMLU (71.2) と GSM8K (69.6) および 4K-training-length pretraining での HumanEval (54.9) において高得点を達成。
• Samba はほぼ完璧な记憶リコールと、最小限の微調整後の最大 256K context までの有効な外挿、そして Proof-Pile の困惑度改善のため 1M context までを示す。
• スループットの利点: Samba は 128K prompts で Llama-3 1.6B より約 3.73x のプロンプト処理スループットを提供し、ストリーミングを伴う 64K トークン生成で 3.64x の速度向上を実現。
• Samba は線形時間計算を維持し、長文脈タスクおよび instruction tuning 後の長文文脈要約において、SWA のみあるいは純粋なアテンションベースのベースラインを上回る。
• アブレーションから、Mamba と SWA のハイブリッド化が全タスクで最良の性能を生み出す一方、注意ヘッドを削減しても性能を維持・向上させられることが示される。

(b) Decoding throughput with a batch size of 16