Author by: 张嘉瑶
S4 及其变体在多个长序列基准测试（如 Long Range Arena）上取得了当时最先进的性能，首次证明了结构化 SSM 在这一领域能够超越 Transformer。它成功地确立了 SSM 作为一种兼具卷积和循环双重优势的强大序列建模新范式。
这种双重性并非偶然的数学技巧，而是由 DPLR 结构所赋予的根本属性。DPLR 是连接 SSM 两种计算视图的桥梁：它使得模型在训练时可以被看作一个可并行的卷积层，在推理时又可以被看作一个轻量级的循环单元。这优雅地解决了长期以来困扰序列模型设计的核心矛盾——RNN 的推理效率与 CNN/Transformer 的训练效率之间的权衡。DPLR 是 S4 最深刻的架构贡献，也为 Mamba 的诞生奠定了基础。
内在局限：
- 输入无关性： 这是 S4 最根本的局限。系统的核心矩阵
$(\overline{A}, \overline{B}, \overline{C})$ 以及由此产生的卷积核$\overline{K}$ ，在模型训练完成后就被固定下来。它们的值不随输入序列的内容而改变。这意味着 S4 本质上是一个线性时不变 Linear Time-Invariant系统。这种特性使其难以处理需要根据内容进行判断的任务，例如，从一段文本中有选择地复制某个特定信息，或者在处理信息时主动忽略无关的干扰项。 - 硬件效率问题： S4 对 FFT 和特殊的柯西矩阵快速算法的依赖，在实践中可能成为性能瓶颈。这些算法在 GPU 上的实现并非总是最优化的，可能导致硬件利用率不高。
如下图的选择性复制任务中，我们希望复制输入的部分内容并按顺序输出：
然而，由于（循环/卷积）SSM 具有线性时不变性，因此它在此任务中表现不佳。正如我们之前所见，对于 SSM 生成的每个 token，矩阵 A 、B 和 C 都是相同的。
因此，SSM 无法进行内容感知推理，因为它会根据固定的 A、B 和 C 矩阵，平等地对待每个 token。无法做到选择性复制提示词。
SSM 表现不佳的第二个任务是感应头，其目标是重现输入中发现的模式：
在上面的例子中，我们本质上是在进行一次性提示，试图“教”模型在每个“问： ”之后给出“答：”的答案。然而，由于 SSM 具有时不变性，它无法从历史记录中选择要调用的先前标记。
让我们通过关注矩阵 B 来说明这一点。无论输入 x 是什么，矩阵 B 都保持不变，因此与 x 无关：
同样，无论输入如何， A 和 C 也保持不变。这证明了我们迄今为止所见的 SSM 的静态特性。
相比之下，这些任务对于 Transformer 来说相对容易，因为它们会根据输入序列动态地改变注意力。它们可以选择性地“查看”或“关注”序列的不同部分。SSM 在这些任务上的糟糕表现说明了时间不变 SSM 的根本问题，矩阵 A 、B 和 C 的静态性质导致了内容感知问题。
S4 的这些局限性并非孤立的缺陷，它们之间存在着紧密的因果联系，并共同指向了下一代模型（即 Mamba）所必须进行的创新方向。首先，最核心的概念缺陷是输入无关性。一个无法根据其处理内容来调整自身行为的模型，其能力必然是受限的。这是 Mamba 必须解决的“为什么”的问题。解决方案显而易见：让系统矩阵（如
Mamba 架构的出现，可以被视为对 S4 局限性的一次直接而精彩的回应。它通过引入选择性机制、硬件感知的并行算法和简化的整体架构，将 SSM 的能力提升到了一个新的高度。
如上图所示，SSM 通过将历史记录压缩成一个较小的状态，从而实现了高效的计算。这个方法能够减少计算开销，因为它不需要保存整个历史数据。然而，与 Transformer 通过注意力机制处理整个历史记录的方式相比，SSM 的表达能力相对较弱。 Mamba 的目标则是结合两者的优点。它通过保持较小的状态，同时也能实现与 Transformer 类似的强大功能，既保持了高效性，又能处理复杂的序列数据。
S4 是一个线性时不变（LTI）系统，而 Mamba 是一个线性时变（LTV）系统。这一转变是两者最核心的区别。Mamba 引入了选择性的核心思想，即模型应具备根据输入内容动态地增强或忽略信息的能力。为了实现这一点，Mamba 将 SSM 的关键参数——步长
为了选择性地压缩信息，我们需要参数依赖于输入。为此，我们首先在训练期间探索 SSM 中输入和输出的维度：
在结构化状态空间模型 (S4) 中，矩阵 A 、B 和 C 与输入无关，因为它们的维度 N 和 D 是静态的且不会改变。
相反，Mamba 通过合并输入的序列长度和批量大小，使矩阵 B 和 C，甚至步长∆依赖于输入：
这意味着对于每个输入标记，我们现在有不同的 B 和 C 矩阵，从而解决了内容感知问题！
注意：矩阵 A 保持不变，因为我们希望状态本身保持静态，但其受影响的方式（通过 B 和 C ）是动态的。
由于它们现在依赖于输入，因此它们会选择性地选择在隐藏状态下保留什么以及忽略什么。
较小的步长∆会导致忽略特定的单词，而是更多地使用先前的上下文，而较大的步长∆会更多地关注输入单词而不是上下文：
综上所述为了选择性保留信息，Mamba 做了以下改动：
-
$\Delta$ （步长）：它控制着模型在“记忆”和“遗忘”之间的平衡。一个较大的$\Delta$ 会使得离散化的状态矩阵$\overline{A}$ 的范数变小，从而更多地“遗忘”之前的状态$x_{t-1}$ ，并更多地关注当前输入$u_t$ 。相反，一个较小的$\Delta$ 会让模型更多地保持和沿用旧的状态。通过动态调整$\Delta$ ，模型可以决定在遇到关键信息时“重置”状态，或在处理连续信息时“维持”状态。 -
$B$ 和$C$ （输入/输出投影）： 这两个矩阵由输入$x_t$ 动态生成。这使得模型可以有选择性地决定将输入的哪些信息写入状态$x_t$ （由$B$ 控制），以及从状态中读取哪些信息作为输出$y_t$ （由$C$ 控制）。 -
$A$ （状态矩阵）： 为了保持系统的稳定性和结构性，核心的状态矩阵$A$ 在 Mamba 中仍然是固定的（非输入依赖的），但其在离散化后的有效作用会受到输入依赖的$\Delta$ 的动态调制。
Mamba 的这种选择机制，实际上是一种以连续动态系统语言实现的“注意力”。传统注意力机制通过计算离散的 softmax 权重来聚合信息，而 Mamba 则通过由
Mamba 的整体架构块（Mamba Block）在设计上追求简洁与高效。它将核心的选择性 SSM（SSM）与一个带有门控机制的类 MLP 结构相结合。输入信号
在机器学习的骨干网络中，通常包含两个基本操作：token 间的信息通信(Communication)和 token 内的计算(Computation)$^{6}$。Mamba 采用 SSM 进行通信(取代了注意力机制)，并保留了类似 MLP 的投影进行计算
Mamba 模型具有多个接口级别，但其核心是封装了选择性 SSM 的 Mamba
下图展示了一个 Mamba 模块的概念框图：
使 Mamba 的动态选择性在实践中可行的，是其在算法层面的关键创新。由于参数是输入依赖的，S4 所依赖的静态卷积技巧不再适用，模型必须依赖循环计算。然而，朴素的循环计算在 GPU 这类并行硬件上效率极低，因为它无法利用大规模并行处理单元。
如图所示，每个状态都是前一个状态（乘以 A ）加上当前输入（乘以 B ）之和。这称为扫描操作，可以使用 for 循环轻松计算。相比之下，并行化似乎是不可能的，因为每个状态的计算都只能在已知前一个状态的前提下进行。
Mamba 的解决方案是设计了一种并行扫描算法。该算法在数学上与顺序的循环计算完全等价，但其计算结构被重塑，使其能够高效地在现代硬件上并行执行。它将一个看似顺序的计算链分解为多个可以并行处理的块，然后在块之间进行少量的顺序聚合。
但 GPU 的一个缺点是其容量较小但效率极高的 SRAM 与容量较大但效率略低的 DRAM 之间的传输 (IO) 速度有限。频繁在 SRAM 和 DRAM 之间复制信息会成为瓶颈。
而 Mamba 算法是“硬件感知”的。Mamba 试图将离散化、循环更新和输出投影等多个计算步骤融合 fuse到单个 GPU 内核中。这样做可以最大限度地减少数据在 GPU 不同层级内存（高带宽内存 HBM 与片上高速 SRAM）之间的读写次数。通过避免在每一步都将完整的中间状态
我们可以通过可视化 Mamba 的基础架构来查看 DRAM 和 SRAM 分配的具体实例：
这里，以下内容融合成一个内核：
-
离散化步骤，步长为
$∆$ -
选择性扫描算法
-
用 C 进行乘法
硬件感知算法的最后一部分是重新计算。
中间状态不会被保存，但对于反向传播计算梯度来说却是必需的。因此，作者在反向传播过程中重新计算了这些中间状态。
尽管这看起来效率不高，但比从相对较慢的 DRAM 读取所有中间状态的成本要低得多。
现在，我们已经介绍了其架构的所有组件，该架构使用文章中的下图进行描述：
这种算法与硬件的协同设计是 Mamba 成功的核心。选择性 SSM 的理论模型若没有并行扫描算法的支持，将因无法高效训练而变得不切实际；而并行扫描算法若没有选择性 SSM 这一应用场景，也只是一个通用的计算技巧。这标志着 AI 研究的一个新前沿：算法创新与系统级工程优化必须紧密结合、同步进行。未来的架构突破可能越来越需要跨越机器学习理论、计算机体系结构和系统编程等多个领域的深度专业知识。
凭借选择性机制和硬件感知算法，Mamba 成功地解决了序列建模领域长期存在的“循环-并行”权衡：
- 训练时： Mamba 利用并行扫描算法，使其本质上是循环的结构能够以接近线性的时间复杂度（$O(L \cdot D \cdot N)$）和极高的硬件利用率进行训练，效率媲美 Transformer。
- 推理时： Mamba 可以切换回其原始的、高效的顺序循环模式。每生成一个新词元，只需进行一次
$O(D \cdot N)$ 的状态更新，计算和内存成本均为常数$O(1)$ ，使其在自回归生成任务中速度极快且内存占用极低。
S4 可以被视为一个里程碑式的“概念验证”，它首次证明了结构化 SSM 作为高效长序列模型的可行性，并建立了其卷积/循环的双重计算模式。而 Mamba 则是其“精炼的继任者”，它精准地解决了 S4 的根本缺陷（输入无关性）和实践瓶颈（硬件利用率）。
从 S4 到 Mamba 的进化，核心是从一个线性时不变（LTI）系统到一个线性时变（LTV）系统的飞跃。这一飞跃之所以能够成功，关键在于 Mamba 通过算法与硬件的协同设计（并行扫描），使得这个更强大、更灵活的时变模型在计算上依然是可行的。
Mamba 与 Transformer 在多个维度上展现了根本性的差异，这些差异可以通过下表进行系统性地总结和分析。
| 特征 | RNN/LSTM | Transformer (Attention) | S4 (结构化 SSM) | Mamba (选择性 SSM) |
|---|---|---|---|---|
| 训练复杂度 |
|
|
|
|
| 自回归推理 |
|
|
|
|
| 训练并行性 | 否 (顺序依赖) | 是 (完全并行) | 是 (通过卷积/FFT) | 是 (通过并行扫描) |
| 长距离依赖处理 | 差 (梯度消失) | 优 (直接路径) | 优 (HiPPO 结构) | 优 (HiPPO+选择性) |
| 输入依赖动态 | 是 (非线性门控) | 是 (自注意力) | 否 (线性时不变) | 是 (选择性参数) |
| 推理内存占用 | 恒定 |
|
恒定 | 恒定 |
从上表中可以清晰地看到，Mamba 在多个关键指标上实现了对以往架构的超越。复杂度： Mamba 在训练时实现了与序列长度
在实际的基准测试中，Mamba 的表现印证了其架构设计的优越性。在语言建模任务上，Mamba 不仅在性能上能够匹配甚至超越同等规模的 Transformer 模型，而且实现了更高的训练和推理吞吐量。在处理基因组学、音频和时间序列等典型的长序列数据模态时，Mamba 展现出了压倒性的优势，成为这些领域的首选架构之一。
从在 HG38 人类基因组数据集上的预训练结果来看，Mamba 在处理基因组学任务时的核心优势可归纳为以下两点，分别对应左图和右图的实验结论：
- 参数规模扩展性更优：小参数量下仍能高效提升性能
- 长序列处理能力更强：上下文长度增加时性能不衰减（甚至提升）
这对基因组学任务至关重要：人类基因组序列极长（总长约 30 亿碱基对），且许多关键生物学信息（如长程基因调控、跨片段突变关联）隐藏在长序列中。传统基线模型（如 Transformer 类）因自注意力机制的计算复杂度随序列长度平方增长，在处理长基因组片段时易出现性能瓶颈（如信息稀释、计算过载）；而 Mamba 的选择机制能高效聚焦长序列中的关键信息（如特定碱基模式、基因片段关联），更适配基因组学对 “长序列语义理解” 的核心需求。
-
Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." arXiv preprint arXiv:2312.00752 (2023).
-
Gu, Albert, et al. "Combining recurrent, convolutional, and continuous-time models with linear state space layers." Advances in neural information processing systems 34 (2021): 572-585.
-
Gu, Albert, et al. "Hippo: Recurrent memory with optimal polynomial projections." Advances in neural information processing systems 33 (2020): 1474-1487.
-
Voelker, Aaron, Ivana Kajić, and Chris Eliasmith. "Legendre memory units: Continuous-time representation in recurrent neural networks." Advances in neural information processing systems 32 (2019).
-
Gu, Albert, Karan Goel, and Christopher Ré. "Efficiently modeling long sequences with structured state spaces." arXiv preprint arXiv:2111.00396 (2021).
-
Grootendorst, M. (2023, October 29). A visual guide to Mamba and state space models. Maarten Grootendorst's Newsletter. https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mamba-and-state