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

这项工作解决了长期以来的一个问题，即如何高效地对具有无限上下文长度的序列进行建模。以往的研究要么受到二次计算复杂度的限制，要么在长度推广方面具有有限的外推能力。本文提出了Samba，这是一种简单的混合架构，它将选择性状态空间模型（SSM）Mamba与滑动窗口注意力（SWA）层次化地结合起来。Samba将给定序列有选择地压缩成循环隐藏状态，同时仍然保持了注意机制精确回忆记忆的能力。我们将Samba扩展到了3.8B参数，使用3.2T的训练令牌，并展示了Samba在广泛的基准测试中明显优于基于纯注意力或SSM的最先进模型。当训练4K长度的序列时，Samba可以高效地外推到256K上下文长度，并展示了在最多1M上下文长度下提高的令牌预测能力。作为一种线性时间序列模型，Samba在处理128K长度的用户提示时具有3.73倍的吞吐量，比具有分组查询注意力的Transformer快3.64倍，当生成64K令牌时具有无限流式处理的速度提升。Samba的一个样例实现可在https://github.com/microsoft/Samba上公开获取。

论文旨在解决序列建模中的无限上下文长度问题，即如何在不牺牲计算效率和泛化能力的情况下，对具有无限上下文长度的序列进行建模。

论文提出了一种名为Samba的混合架构，将选择性状态空间模型（SSM）Mamba与滑动窗口注意力（SWA）层级组合。Samba将给定的序列压缩成循环隐藏状态，同时通过注意机制保持精确的记忆回溯能力。

论文在3.2T的训练标记和3.8B的参数规模下，将Samba模型应用于多个基准测试，并在很大程度上优于基于纯注意力或SSMs的现有模型。当在4K长度序列上进行训练时，Samba可以有效地推广到256K上下文长度，并在1M上下文长度的标记预测方面显示出改进。Samba作为一种线性时间序列模型，与分组查询注意力的Transformer相比，处理128K长度的用户提示时吞吐量提高了3.73倍，在生成64K标记时，无限流处理速度提高了3.64倍。Samba的示例实现在https://github.com/microsoft/Samba上公开。

最近在这个领域中，还有一些相关的研究，例如《Attention is All You Need》、《Longformer: The Long-Document Transformer》、《Reformer: The Efficient Transformer》等。