Mamba: Linear-Time Sequence Modeling with Selective State Spaces

公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）

目录

0. 摘要

1. 简介

2. 状态空间模型  

3. 选择性状态空间模型

3.1 动机：选择作为压缩的手段

3.2 通过选择改进 SSM

3.4 一个简化的 SSM 架构

3.5 选择机制的性质

3.5.1 与门控机制的连接

3.5.2 选择机制的解释

4. 实验

---

0. 摘要

目前基础模型驱动着深度学习中大部分令人兴奋的应用，几乎普遍基于 Transformer 架构及其核心注意力模块。许多次二次时间（subquadratic-time）架构，如线性注意力、门控卷积和循环模型，以及结构化状态空间模型（structured state space models，SSM），已经被开发出来以解决 Transformer 在长序列上的计算效率低下的问题，但它们在重要的模态，如语言上，表现不佳。我们确定这些模型的一个关键弱点是它们无法进行基于内容的推理，并进行了几项改进。

• 首先，简单地让 SSM 参数成为输入的函数，可以解决它们在离散模态上的弱点，允许模型根据当前 token 有选择地沿序列长度维度传播或遗忘信息。
• 其次，即使这种改变阻止了有效卷积的使用，我们设计了一种硬件感知（hardware-aware）的并行算法，以递归模式运行。
• 我们将这些选择性的 SSM 集成到一个简化的端到端神经网络架构中，没有注意力甚至 MLP 块（Mamba）。

Mamba 享有快速推理（比 Transformers 高 5 倍的吞吐量）和序列长度的线性扩展，并且在真实数据上的性能提高了，直至百万长度的序列。作为一种通用的序列模型骨干，Mamba 在语言、音频和基因组等多个模态上实现了最先进的性能。在语言建模方面，我们的 Mamba-3B 模型在预训练和下游评估中均优于相同规模的 Transformers，并且与其两倍大小的 Transformers 相匹配。

项目页面：https://github.com/state-spaces/mamba

1. 简介

最近，结构化状态空间序列模型（SSM）（Gu，Goel和Ré 2022; Gu，Johnson，Goel等人2021）已经成为一种有前途的序列建模架构类别。这些模型可以被解释为循环神经网络（RNN）和卷积神经网络（CNN）的组合，灵感来自经典状态空间模型（Kalman 1960）。这类模型可以非常高效地计算，可以作为递归或卷积进行，序列长度呈线性或近线性扩展。此外，它们在某些数据模态中具有建模长程依赖性的原则性机制（Gu，Dao等人2020年），并且在 Long Range Arena（Tay，Dehghani，Abnar等人2021年）等基准测试中占据主导地位。许多种类的 SSM（Gu，Goel和Ré 2022年; Gu，Gupta等人2022年; Gupta，Gu和Berant 2022年; Y. Li等人2023年; Ma等人2023年; Orvieto等人2023年; Smith，Warrington和Linderman 2023年）在涉及连续信号数据的领域取得了成功，如音频和视觉（Goel等人2022年; Nguyen，Goel等人2022年; Saon，Gupta和Cui 2023年）。然而，它们在建模文本等离散和信息密集型数据方面效果较差。 

我们提出了一种新型的选择性状态空间模型（selective state space models），通过在多个方面改进之前的工作，实现了与 Transformer 相当的建模能力，同时在序列长度上呈线性扩展。

2. 状态空间模型  

（2024，ViM，双向 SSM 骨干，序列建模）利用双向状态空间模型进行高效视觉表示学习

（2024，VMamba，交叉扫描，线性复杂度，全局感受野，动态权重）视觉状态空间模型 

（2024，DiS，扩散，状态空间主干，Mamba）具有状态空间主干的可扩展扩散模型 

（2024，Mamba，DiT，之字扫描与空间连续性）ZigMa：之字形 Mamba 扩散模型 

3. 选择性状态空间模型

3.1 动机：选择作为压缩的手段

我们认为序列建模的一个基本问题是将上下文压缩到一个较小的状态中。实际上，我们可以从这个角度来看待流行序列模型的权衡。例如，注意力既有效又低效，因为它明确地不对上下文进行任何压缩。这可以从自回归推理需要明确存储整个上下文（即 KV 缓存）这一事实中看出，这直接导致了 Transformer 的缓慢的线性时间推理和二次时间训练。另一方面，循环模型是高效的，因为它们具有有限的状态，意味着推理时间是恒定的，而训练时间是线性的。然而，它们的有效性受到状态如何压缩上下文的限制。

为了理解这一原理，我们专注于两个合成任务的运行示例（图2）。

• 选择性复制（Selective Copying）任务修改了流行的复制任务（Arjovsky，Shah和Bengio 2016），通过改变要记忆的 token 的位置。它需要内容感知的推理能力，以便能够记住相关的 token（着色）并过滤掉无关的标记（白色）。
• 归纳头部（Induction Heads）任务是一个已知的机制，被假设可以解释 LLM 的大部分上下文学习能力（Olsson等人2022）。它需要上下文感知的推理能力，以了解何时在适当的上下文中产生正确的输出（黑色）。

这些任务揭示了线性时不变模型（Linear Time Invariance，LTI）的失败模式。从循环视角来看，它们的恒定动态（例如公式 2 中的（¯A，¯B）过渡）不能让它们从上下文中选择正确的信息，或以一种与输入相关的方式影响沿着序列传递的隐藏状态。从卷积视角来看，众所周知，全局卷积可以解决普通复制任务（Romero等人2021）因为它只需要时间感知，但是它们在选择性复制任务方面有困难，因为缺乏内容感知（图2）。 更具体地说，输入到输出之间的间隔是变化的，不能通过静态卷积核进行建模。

总之，序列模型的效率与有效性的权衡在于它们如何压缩其状态：高效的模型必须具有较小的状态，而有效的模型必须具有一个包含来自上下文的所有必要信息的状态。反过来，我们提出序列模型构建的一个基本原则是选择性：或者说，关注或过滤顺序状态输入的上下文感知能力。特别是，选择机制控制信息沿着序列维度如何传播或相互作用（有关更多讨论，请参见第 3.5 节）。 

3.2 通过选择改进 SSM

将选择机制纳入模型的一种方法是让影响序列上相互作用的参数（例如 RNN 的循环动态或 CNN 的卷积核）依赖于输入。

算法 1 和 2 说明了我们使用的主要选择机制。主要区别仅在于使几个参数 Δ，B，C 成为输入的函数，以及随之在整个张量形状上的相关变化。特别地，我们强调这些参数现在具有长度维度 L，这意味着模型已从时不变变为时变。这使得它失去了与卷积（公式 3）的等价性，对其效率产生了影响，接下来进行讨论。

特别地，我们选择 s_B(x) = Linear_N(x)，s_C(x) = Linear_N(x)，s_Δ(x) = Broadcast_D(Linear_1(x))，和 τ_Δ = softplus，其中 Linear_d 是一个参数化的到维度 d 的投影。选择 s_Δ 和 τ_Δ 是因为与在第 3.5 节中解释的 RNN 门控机制的连接。

3.4 一个简化的 SSM 架构

与结构化的 SSM 一样，选择性的 SSM 是独立的序列转换，可以灵活地整合到神经网络中。H3 架构是最知名的 SSM 架构（第 2 节）的基础，通常由一个受线性注意力启发的块与一个 MLP 块交错组成。我们通过将这两个组件合并成一个，并且均匀堆叠（图 3）来简化这种架构。这受到了门控注意力单元（GAU）（Hua等人2022年）的启发，后者对注意力做了类似的处理。

这种架构涉及通过可控的扩展因子 E 扩展模型维度 D。对于每个块，大部分参数（3ED^2）都在线性投影中（输入投影为 2ED^2，输出投影为 ED^2），而内部的 SSM 贡献较少。与之相比， SSM 参数的数量（用于 Δ，B，C 的投影和矩阵 A）要小得多。我们重复这个块，交错使用标准的归一化和残差连接，形成 Mamba 架构。在我们的实验中，我们总是将 E 固定为 2，并使用两个块的堆叠来匹配 Transformer 的交错 MHA（多头注意力）和 MLP 块的 12D^2 参数。我们使用 SiLU / Swish 激活函数（Hendrycks和Gimpel 2016; Ramachandran，Zoph和Quoc V Le 2017），这样做的动机是使门控 MLP 成为流行的 “SwiGLU” 变体（Chowdhery等人2023年; Shazeer 2020; Touvron等人2023年）。最后，我们还额外使用一个可选的归一化层（我们选择 LayerNorm（J. L. Ba，Kiros和Hinton 2016）），这是受 RetNet 在类似位置使用归一化层的启发（Y. Sun等人2023年）。

3.5 选择机制的性质

选择机制是一个更广泛的概念，可以以不同的方式应用，例如应用到传统的 RNN 或 CNN 中，应用到不同的参数（例如 Algorithm 2 中的 A），或者使用不同的变换 s(x)。

3.5.1 与门控机制的连接

我们强调最重要的连接：RNNs 的经典门控机制是我们 SSM 的选择机制的一个实例。我们注意到，RNN 门控与连续时间系统的离散化之间的联系是已经确立的（Funahashi和Nakamura 1993; Tallec和Ollivier 2018）。事实上，定理 1 是对 Gu，Johnson，Goel等人（2021年，引理 3.1）的改进，推广到了 ZOH 离散化和依赖于输入的门（附录 C 中的证明）。在更广泛的范围内，SSM 中的 Δ 可以被看作是 RNN 门控机制的广义角色。与之前的工作一致，我们采用了这样的观点，即 SSM 的离散化是启发式门控机制的原则性基础。

定理 1。当 N = 1，A = -1，B = 1，s_Δ = Linear(x)，而 τ_Δ = softplus 时，选择性 SSM 递归（Algorithm 2）的形式如下： 

如第 3.2 节所述，我们对 s_Δ 和 τ_Δ 的具体选择是基于这种连接。特别地，请注意，如果一个给定的输入 xt 应该完全被忽略（在合成任务中是必要的），所有的 D 个通道都应该忽略它，因此我们在将其重复/广播给 Δ 之前将输入投影到 1 维。 

3.5.2 选择机制的解释

我们详细阐述了选择的两个特定机制效应。

可变空间（Variable Spacing）。选择性允许过滤掉可能出现在感兴趣的输入之间的不相关噪声标记。这在选择性复制任务中得到了示范，但在常见的数据模态中普遍存在，特别是对于离散数据 - 例如语言填充词 “um” 的存在。这个属性的出现是因为模型可以从机制上过滤掉任何特定的输入 xt，例如在门控 RNN 情况下（定理 1）当 gt → 0 时。

过滤上下文。已经经验性地观察到许多序列模型即使在更长的上下文中也不能提高性能（F. Shi等人2023年），尽管原则上更多的上下文应该导致严格更好的性能。一个解释是，许多序列模型在必要时不能有效地忽略无关的上下文; 一个直观的例子是全局卷积（以及一般的 LTI 模型）。另一方面，选择性模型可以在任何时候简单地重置它们的状态以删除多余的历史记录，因此它们的性能原则上随着上下文长度的增加而单调提高（例如第 4.3.2 节）。

边界重置。在将多个独立序列拼接在一起的设置中，Transformer 可以通过实例化特定的注意力掩码将它们保持分开，而 LTI 模型将在序列之间泄漏信息。选择性 SSM 也可以在边界处重置它们的状态（例如 Δ_t → ∞ 或定理 1 中当 gt → 1时）。这些设置可能是人工的（例如将文档打包在一起以提高硬件利用率）或自然的（例如在强化学习中的情节边界（Lu等人2023））。

此外，我们详细阐述了每个选择性参数的效果。

Δ 的解释。通常，Δ 控制着在多大程度上关注或忽略当前输入 xt。它将 RNN 门控（例如定理 1 中的 gt）进行了泛化，从机械上讲，大的 Δ 会重置状态 h 并专注于当前输入x，而小的 Δ 会保持状态并忽略当前输入。SSM（公式 1 和 2）可以被解释为通过时间步长 Δ 离散化的连续系统，在这种情况下，直觉是大的 Δ → ∞ 表示系统更长时间地专注于当前输入（因此 “选择” 它并忘记了当前状态），而小的 Δ → 0 表示被忽略的瞬态输入。

A 的解释。我们注意到，虽然 A 参数也可以是选择性的，但它最终通过与 Δ 的交互来影响模型，通过 ¯A = exp(ΔA)（离散化公式 4）。因此，在 Δ 中的选择性足以确保（¯A，¯B）中的选择性，并且是改进的主要来源。我们假设，除了（或者代替）Δ，使 A 也具有选择性， 将具有类似的性能，并且为简单起见而将其排除。

B 和 C 的解释。如第 3.1 节所讨论的，选择性的最重要特性是过滤掉不相关的信息，以便将序列模型的上下文压缩成一个高效的状态。在 SSM 中，将 B 和 C 修改为选择性允许更精细地控制是否将输入 xt 引入状态 ht 或状态输出 yt。这可以解释为允许模型根据内容（输入）和上下文（隐藏状态）调节循环动态。

4. 实验