Samba：用于高效无限上下文语言建模的简单混合状态空间模型

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

📜 文献卡

Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling
作者: Liliang Ren; Yang Liu; Yadong Lu; Yelong Shen; Chen Liang; Weizhu Chen
DOI: 10.48550/arXiv.2406.07522
摘要: 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.
**github（Pytorch）: **https://github.com/microsoft/Samba

⚙️ 内容

研究提出了一种名为SAMBA的简单混合状态空间模型，旨在高效处理无限上下文长度的语言建模问题。SAMBA结合了Mamba（一种选择性状态空间模型）和滑动窗口注意力（Sliding Window Attention, SWA），旨在通过层间混合这些组件来解决序列建模中的长期依赖捕捉和内存效率问题。

💡 创新点

• 高效无限上下文建模：SAMBA能够以线性时间复杂度处理任意长度的序列，显著提高了处理长序列的效率。

• 混合架构设计：结合Mamba的序列结构捕获能力和SWA的精确记忆检索能力，SAMBA在不牺牲性能的前提下实现了高效压缩和精确回忆。

• 卓越性能与扩展性：模型被扩展至高达3.8亿参数，并在仅训练4K长度序列的基础上，展示了对长达256K甚至1M上下文长度的有效外推能力，同时保持线性解码时间复杂度。

• 实际应用速度提升：在处理用户提示和生成长序列时，SAMBA相比基于Transformer的模型具有显著的速度优势，分别达到3.73倍和3.64倍的吞吐量提升。

• 指令微调下的记忆召回：经过少量步骤的指令微调，SAMBA能够在Passkey Retrieval任务中实现256K上下文长度的完美记忆召回，而基于SWA的模型无法超越4K限制。

🧩 不足

• 对于一些替代的线性循环层（如多尺度保留和全局线性注意力）的测试显示，它们的表现不如Mamba，但具体原因和优化空间尚需进一步探索。
• 添加特定技术（如SC）到模型的不同层会导致性能下降，这表明对于这类混合架构，组件之间的平衡与集成策略需要更精细的调整。

🔁 实验卡

💧 数据

研究使用了规模庞大的训练数据集，包含3.2万亿个标记，这充分展示了模型在大数据集上进行训练和泛化的强大能力。通过对如此海量数据的学习，SAMBA能够学习到丰富的语言规律和模式，从而在各种任务中展现出优越的泛化性能。

👩🏻‍💻 方法

SAMBA的核心创新在于其独特的混合架构设计，该架构融合了三种关键组件：Mamba、滑动窗口注意力（SWA）和多层感知机（MLP）。具体来说：

• Mamba作为选择性状态空间模型，专注于捕捉序列中的非马尔可夫依赖关系，有效维护长期记忆。
• 滑动窗口注意力（SWA）则通过操作一个固定大小的窗口滑过输入序列，直接访问窗口内的内容以获取近期历史的高精度信号，补充Mamba在中短期记忆捕捉上的不足。
• 多层感知机（MLP）则提供了一个非线性转换和事实知识回忆的机制，增强了模型的表达能力。

研究人员探索了这些组件的不同组合策略，以求找到最高效的架构配置，确保模型在语言建模任务中既能高效运行，又能准确处理无限长度的上下文信息。

🔬 实验

为了验证SAMBA的性能，研究团队实施了一系列广泛的实验：

• 基准测试：涵盖了常识推理（如ARC、PIQA）、语言理解（如HellaSwag、BoolQ）、数学和编程任务（如GSM8K、MBPP），结果显示SAMBA在众多任务上超越了当前最佳的纯注意力或状态空间模型。
• 消融实验与架构分析：进行了深入的分析，以确认设计决策的有效性，比如对模型的不同组件进行单独测试，以理解它们各自贡献及相互作用机制，这有助于揭示SAMBA成功的关键因素。

📜 结论

最终，SAMBA不仅在理论性能评估中取得了优异成绩，还在实际应用中展现出了处理长序列数据的高效能，尤其是在那些需要深刻理解大量上下文信息的任务上。该模型的成功不仅为解决无限上下文长度序列建模问题提供了新的思路，还为未来混合模型设计和参数优化策略的研究开辟了新的方向，证明了通过精心设计的混合架构，可以在保持高效率的同时，增强模型的泛化和记忆能力。

🤔 总结卡

研究成功地开发了一种在理论上和实践中都表现出色的混合模型，SAMBA通过巧妙结合不同的建模方法，解决了无限上下文建模的效率与能力问题。其在处理超长序列时的性能提升尤为引人注目，同时也展示了在不同任务上的广泛适用性和优越性。然而，对于替代组件表现不佳的原因以及如何进一步优化混合架构的内部平衡，仍有待深入探究。未来工作可考虑如何更好地整合不同模块，比如理解SC为何在某些情况下导致负面效果，并探索如何最大化利用每个组件的优势。