[论文解读] 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（选择性状态空间模型）与滑动窗口注意力联合集成，实现无限上下文长度，线性时间复杂度，在基准测试上表现强劲，吞吐量显著提升。

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。
• 在多达 3.2T token 上训练 421M、1.3B、1.7B、3.8B 参数模型，并在多样化基准上评估。
• 执行高效的长上下文外推测试（高达 256K 提示和 1M 上下文）以及吞吐比较。

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

实验结果

研究问题

• RQ1 一个简单的层级混合的输入相关 SSM 与滑动窗口注意力是否能实现线性时间复杂度的无限上下文长度？
• RQ2 Samba 与纯注意力和纯 SSM 模型在长上下文和短上下文基准上的比较？
• RQ3 哪些架构选择（如注意力分布 vs 递归的分配）能为长上下文语言建模带来最佳性能与效率？
• RQ4 在外推上下文长度时，记忆回忆与检索在多大程度上可以被保留或增强？
• RQ5 对提示处理与生成的效率与吞吐有何影响？

主要发现

• Samba 在广泛基准集上对比到强大的开源 LLMs（高达 8B 参数）取得了最先进的平均表现。
• 3.8B Samba 模型在 MMLU（71.2）和 GSM8K（69.6）以及 4K 训练长度预训练上的 HumanEval（54.9）获得高分。
• Samba 显示出近乎完美的记忆回忆和在最小微调后对齐到 256K 上下文的有效外推，以及在 Proof-Pile 上对数困惑度提升可达1M 上下文。
• 吞吐方面的收益：在 128K 提示下，Samba 相比 Llama-3 1.6B 提高约 3.73 倍的提示处理吞吐，在 64K 令牌生成并带流式传输时提升约 3.64 倍。
• Samba 保持线性时间复杂度，并在长上下文任务和指令微调后的下游长上下文摘要中优于仅 SWA 或纯注意力的基线。
• 消融表明，将 Mamba 与 SWA 混合在所有任务中都获得最佳性能，而减少注意头数量仍可保留甚至提升性能。

(b) Decoding throughput with a batch size of 16