Mamba: S4 结构化状态空间模型剖析#
1. 引言：结构化状态空间模型的兴起#
序列建模是人工智能领域的核心挑战之一，旨在理解和生成各种序列数据，如自然语言文本、语音信号和时间序列数据。传统的循环神经网络(RNN)在处理序列数据方面取得了一定的成功，但其固有的梯度消失或爆炸问题以及顺序处理瓶颈，限制了其在长序列上的学习能力和效率 \(^{1}\)。长短期记忆(LSTM)网络通过引入门控机制在一定程度上缓解了这些问题，但同时也增加了模型的复杂性和计算开销 \(^{1}\)。
继 RNN 之后，Transformer 架构凭借其并行处理能力和强大的上下文表征能力，在自然语言处理等领域取得了革命性突破。然而，Transformer 的核心机制自注意力机制-其计算复杂度随序列长度呈二次方增长(O(N2))，这使得其在处理极长序列时面临巨大的计算和内存挑战 \(^{2}\)。这种二次方复杂度的瓶颈，成为了驱动研究者探索更高效长序列建模方法的主要动力。
在这样的背景下，结构化状态空间模型(Structured State Space Models, SSMs)种极具潜力的新兴序列建模范式应运而生。SSMs 借鉴了控制理论和信号处理中的经典状态空间表示，通过结构化的循环和状态空间表征，实现了对长序列的高效处理，其计算复杂度通常为线性或近线性 \(^{2}\)。SSMs 的理论根基可以追溯到控制理论，用于描述复杂动态系统，这为其提供了与基于注意力的模型不同的理论基础 1。从奠基性的 S4 模型到其后续的 Mamba、S5 和 Jamba 等模型，SSMs 在计算效率、内存优化和推理速度方面不断取得显著进展 \(^{1}\)。
这种从 RNN 到 Transformer 再到 SSMs 的演进，不仅仅是模型性能的迭代提升，更深层次地反映了在序列建模领域中，研究者们对于"效率-表达能力权衡"这一核心问题的持续探索和突破。RNN 在表达能力上受到限制，尤其是在长程依赖方面;LSTM 通过增加复杂性提升了记忆能力，但效率和极长序列处理仍是瓶颈;Transformer 以其强大的表达能力和并行训练特性取得了巨大成功，却牺牲了在长序列上的计算效率。SSMs，特别是 S4 和 Mamba，则代表了在保证甚至增强对长程依赖建模能力(表达能力)的同时，重新夺回线性或近线性计算复杂度(效率)的尝试。这一过程揭示了人工智能模型发展的一个重要趋势:通过对基础架构假设的重新思考，来克服普遍存在的性能瓶颈和权衡。
此外，SSMs 的理论基础源于经典的控制理论 1，这标志着不同科学领域的思想向机器学习领域的有益渗透和交叉融合。深度学习的灵感常常来源于神经科学或纯数学，而 SSMs 从控制理论中引入状态空间表示、连续时间动态、离散化和系统稳定性等概念，为序列建模带来了新的视角。这种跨学科的方法可能为解决机器学习中长期存在的问题(如长程依赖建模和模型可解释性)提供新的途径，通过借鉴其他领域成熟的理论和
Mamba 的成功,也可能激励研究者进一步探索来自工程学、物理学等领域的概念,以设计出更鲁棒和高效的人工智能架构。
本报告旨在对 S4 模型及其向 Mamba 架构的演进进行全面的技术剖析，详细阐述它们的核心机制、性能对比、优势与局限，以及在序列建模领域的更广泛影响。
2. S4 模型：为高效长程序列建模奠定基础#
S4(Structured State Space for Sequence Modeling)模型是 SSM 发展历程中的一个重要里程碑，它为高效处理长程依赖问题提供了坚实的基础。
2.1 状态空间模型的数学基础#
S4 模型的核心构建于状态空间模型的数学框架之上。一个连续时间状态空间模型通常由以下线性微分方程组定义，该方程组将一维输入信号 u(t)映射到一个 N 维隐状态 x(t)，然后再投影到一维输出信号 y(t)5:
x'(t)=Ax(t)+Bu(t)
y(t)=Cx(t)+Du(t)
其中,A、B、C 和 D 是通过梯度下降学习的参数矩阵。为了简化,参数 D 通常被个跳跃连接 5。
为了将 SSMs 应用于离散输入序列(如 u0,u1,...),需要使用一个步长 \(\Delta\) 对连续 SSM 进行离散化。常用的离散化方法是双线性变换，它将连续状态矩阵 A 转换为一个近似的离散状态矩阵 A-。离散 SSM 随后表现为一个递归关系，使其能够像 RNN 一样进行计算 5:
尽管这种递归表示在直觉上易于理解，但由于其顺序计算的特性，在现代硬件上进行训练时效率不高。S4 模型的一个关键洞察在于线性时不变(LTI)SSMs 与连续卷积之间的重要联系。这使得递归 SSM 可以表示为一个离散卷积，其卷积核 \(K^{-}\) 有明确的公式 5:
这种从类似 RNN 的递归到类似 CNN 的卷积的转换是一项核心创新，它使得模型能够利用快速傅里叶变换(FFT)对长序列进行高效训练。而模型的递归形式则更适合于自回归推理 4。S4 模型在递归和卷积表示之间的这种对偶性是其能够兼顾训练效率和推理效率的关键。RNN 固有的顺序性使其在并行硬件上的训练速度较慢,而 CNN 虽然高度可并行化训练,但在自回归不那么自然。S4 模型由于其 LTI 特性，拥有这种双重性质 4。离散化导向递归形式，而 LTI 特性允许这种递归展开为全局卷积。这种对偶性使得 S4 能够利用两者的优势:通过基于 FFT 的卷积进行并行训练，以及通过递归状态更新进行高效的自回归推理。这突显了一个重要的观点:架构创新往往在于发现能够释放计算效率的数学等价性。
2.2 HiPPO 框架：捕获长程依赖#
基础 SSMs 在实际应用中面临的一个主要挑战是其在长序列上表现不佳，这通常归因于梯度消失或爆炸问题。S4 通过引入 HiPPO(High-order Polynomial Projection Operator)理论来解决这个问题，这是一种用于连续时间记忆的理论框架 \(^{1}\)。
HiPPO 框架指定了一类特殊的矩阵 A(即 HiPPO 矩阵)，这些矩阵能够使潜状态 x(t)有效地记忆输入 u(t)的历史信息。它通过将过去的信息投影到一个多项式基项式)上，从而维持一个压缩但富有表达能力的记忆表示。HiPPO 矩阵的有效性在于，它产生的隐状态通过追踪勒让德多项式的系数来记忆其历史。这些系数使得模型能够逼近所有先前的历史，HiPPO 矩阵在每一步更新这些系数 \(^{5}\)。一个显著的例子是，在序列 MNIST 分类任务中,仅仅将一个随机初始化的 A 矩阵替换为 HiPPO 矩阵,就能将模型性能从 60%提升到 98%，这充分证明了 HiPPO 的重要性 5。
HiPPO 为 SSMs 中的记忆机制提供了一种有原则的方法，超越了启发式的门控机制。与 LSTM/GRU 中通过学习得到的门控机制不同，HiPPO 提供了一种数学上推导出的长时记忆方法。LSTM 使用门(输入门、遗忘门、输出门)来控制信息流，但其学习过程是数据驱动的，可能缺乏透明度。HiPPO 的设计目标是通过将历史信息投影到最优的多项式基上来保留过去的信息 \(^{1}\)，这是一种更结构化、理论基础更坚实的记忆方法。HiPPO 矩阵 A 的特定结构并非随意设定，而是源于对连续信号进行最优在线压缩的目标。这预示着在神经网络记忆机制设计中，一种向更具理论原则性的方法转变的趋势，可能催生出更鲁棒和易于理解的模型。HiPPO 的成功 \(^{5}\) 验证了这一方向的正确性。
2.3 S4 架构：关键创新与计算特性#
S4 通过引入优化的参数化方法(结构化矩阵)，在保持状态空间建模优势的同时，显著降低了内存开销 \(^{1}\)。然而，即便使用了 HiPPO 初始化，朴素地计算卷积核 \(K^{-}\)(涉及 \(A^{-}\) 的重复矩阵乘法)仍然会导致 O(N2L)的运算量和 O(NL)的空间复杂度，这对于长序列来说计算成本过高 5。
S4 通过专注于具有特殊结构的 SSMs-一复数空间中的对角加低秩(Diagonal Plus Low-Rank, DPLR)结构来克服这一计算瓶颈。一个 DPLR SSM 被定义为 \((\wedge\)-PQ\(\cdot\),B,C)，其中 \(\wedge\) 是对角矩阵，而 P,Q,B,C 是向量(假设秩为 1)。S4 通过三个步骤实现了对核计算的加速 \(^{5}\):
截断生成函数(Truncated Generating Function):S4 不是直接计算 \({ K}^{-}\)，而是通过评估其截断生成函数来计算其频谱。这将问题从矩阵幂运算转化为矩阵求逆运算，后者可以更有效地计算。
对角情况下的柯西核(Cauchy Kernel for Diagonal Case):对于对角矩阵 A=\的情况，生成函数可以写成一种包含柯西核的形式。这有效地用加权点积替换了矩阵求逆，从而大大提高了计算速度。
低秩校正的伍德伯里恒等式(Woodbury Identity for Low-Rank Correction):为了放宽对角假设并包含低秩分量(P,Q)，应用了伍德伯里恒等式。该恒等式允许对角加秩-1 项的逆以对角项的逆来表示，从而有效地将问题简化回高效的对角情况。
此外，HiPPO 矩阵本身虽然不是直接的 DPLR 结构，但它是正规加低秩(NormalPlus Low-Rank,NPLR)结构。正规矩阵是酉可对角化的,这使得 NPLR 矩阵从 SSM 模型的角度来看，本质上等同于 DPLR 矩阵。S4 利用这一点，首先将 HiPPO 矩阵写成正规加低秩项，然后将其对角化以提取 DPLR 项。一个额外的简化是低秩分量可以绑定为 P=Q，这有助于提高稳定性 5。
通过这些结构化计算,S4 实现了亚二次复杂度(O(NlogN))1。同时,S4 与现代硬件加速器(如 GPU 和 TPU)兼容，支持高效并行化，能够处理数万个词元长度的序列而不会消耗过多内存 1。S4 在诸如长程竞技场(LongRange Arena,LRA)等长程依赖基准测试中取得了当前最优结果 1。
DPLR 结构的引入并非微小的优化，而是使 S4 在处理长序列时计算上可行的关键因素。核心 SSM 方程虽然简单，但将其朴素地应用于长序列速度过慢 5。HiPPO 矩阵虽然对记忆有益,但也需要计算上易于管理。DPLR 结构(以及 HiPPO 的 NPLR 特性)使得复杂的矩阵运算(特别是卷积核的计算)能够分解为使用柯西核和伍德伯里恒等式等技术的高效步骤 5。可以说，没有 DPLR，S4 可能仍停留在理论层面。这凸显了算法和结构创新在深度学习中与概念创新同等重要。
下表总结了 RNN/LSTM、Transformer 和 S4 在关键架构和计算特性上的对比:
表 T1:核心序列模型的架构与计算特性对比
特性 |
RNN/LSTM |
Transformer |
S4(Structured State Space) |
|---|---|---|---|
核心机制 |
循环连接,门控机制 (LSTM) |
自注意力机制,位置编码 |
状态空间模型(SSM) +HiPPO 初始化 |
输入依赖性(参数) |
时不变(权重共享) |
输入依赖(通过注意力权重) |
时不变(固定 A,B,C 矩阵) |
长程依赖处理 |
梯度消失/爆炸, LSTM 有所缓解 |
优秀,但受限于上下文窗口和二次复杂度 |
通过 HiPPO 和 SSM 结构高效捕获 |
训练复杂度 |
O(L·D2)(顺序) |
O(L2·D)(并行,自注意力瓶颈) |
O((L+N)log(L+N))或 O(L·N)(卷积模式,FFT) |
推理复杂度 |
O(L·D2)(自回归时 O(D2)每步) |
O(L2·D)(自回归时 O(L·D)每步,有 KV 缓存) |
O(L·N)(自回归时 O(N)每步,递归模式) |
并行性(训练) |
有限(沿时间维度顺序) |
高(层内并行) |
高(卷积模式) |
内存(训练) |
O(D)或 O(L·D) (BPTT) |
O(L2+L·D) |
O(L·N)或 O(L+N) |
内存(推理/KV 缓存) |
O(D)(隐状态) |
O(L·D)(KV 缓存) |
O(N)(隐状态) |
注:L 为序列长度，D 为模型维度/隐藏层大小，N 为 S4 的状态维度。复杂度为近似值，具体取决于实现。资料来源:1。
此表清晰地展示了 S4 模型在解决 RNN 和 Transformer 局限性方面的努力，并为理解 Mamba 如何进一步改进 S4 奠定了基础。
3. Mamba：选择性状态空间与硬件感知设计#
尽管 S4 模型在长序列建模方面取得了显著进展，但其依赖于静态、时不变参数(LTI 特性)的特性，使其在处理信息密集型数据或需要基于内容进行推理的任务时表现欠佳 1。S4 的 LTI 特性意味着 SSM 参数(A,B,C)是固定的，不会根据输入序列进行调整。Mamba 在 S4 的基础上，通过引入动态的、输入依赖的参数化方法，实现了能力的巨大
Mamba 在 S4 的基础上，通过引入动态的、输入依赖的参数化方法，实现了能力的巨大提升 1。这是 Mamba 相较于 S4 的核心概念飞跃。
3.1 Mamba 的核心架构创新#
Mamba 架构的核心创新在于其选择性状态空间机制和为之配套的硬件感知并行算法。
3.1.1 选择性状态空间(S6)机制：输入依赖的动态性#
Mamba 最核心的改进是使其 SSM 参数(特别是离散化步长△、输入矩阵 B 和输出矩阵 C)成为当前输入 x 的函数 4。这一机制使得模型能够"根据当前词元选择性地沿序列长度维度传播或遗忘信息"7。这种数据依赖的选择机制能够动态地过滤输入，通过关注相关信息并丢弃无关数据，显著改善了长程依赖的建模能力 1。
选择机制对于处理离散模态(如文本和 DNA)至关重要，在这些领域，S4 的静态特性是一个弱点 \(^{7}\)。Mamba 通过使参数 \(\Delta(x), B(x), C(x)\) 随输入 x 变化,根本性地改变了 S4 中所有词元都按相同"规则"(固定的 A,B,C 矩阵)处理的模式。这意味着模型可以根据当前输入决定保留多少过去的状态(通过△),以及如何将当前输入转换到状态将状态转换到输出(通过 C)。这种做法在理念上更接近于注意力机制(其中交互是输入依赖的),但它是通过具有线性扩展性的循环机制实现的。其结果是 M"基于内容的推理"7,这是 Transformer 相较于先前 SSMs 的主要优势,也使得 Mamba 适用于像语言建模这样复杂、信息密集的任务。
根据 8 中的算法 2(SSM+Selection，或称 S6 机制)，这些输入依赖参数的具体形式为:
sB(x)=LinearN(x)
sC(x)=LinearN(x)
sΔ(x)=BroadcastD(Linear1(x))
T△=softplus
其中 Lineard 是一个参数化的到维度 d 的投影。
Mamba 中的选择性 SSM 递归在特定条件下(例如 N=1,A=-1,B=1,s△=Linear(xt),=softplus)会退化为一个门控循环单元:
gt=o(Linear(xt))
ht=(1-gt)ht-1+gt xt
这种联系表明，SSM 中的 \(\Delta\) 参数扮演了类似于 RNN 门控机制的广义角色，控制着模型是关注当前输入还是忽略当前输入并保持先前状态的平衡 8。一个大的 \(\Delta\) 会重置状态并关注当前输入，而一个小的 \(\Delta\) 则会保持状态并忽略当前输入。使 B 和 C 也具有选择性，则允许模型更细致地控制输入是否进入状态以及状态是否进入输出，从而根据内容和上下文调节循环动态。
3.1.2 硬件感知并行算法：高效计算#
输入依赖的选择性使得 Mamba 的 SSM 变为时变系统，从而破坏了 S4 赖以进行高效卷积训练的 LTI 特性 \(^{3}\)。Mamba 通过设计一种硬件感知的并行算法来克服这一挑战,该算法使用扫描(scan)操作以循环方式计算模型，并针对现代 GPU 进行了优化 \(^{1}\)。
Mamba 的成功并不仅仅归功于选择性 SSM 的思想，同样重要的是其算法与硬件能力的务实协同设计。选择机制打破了 S4 的卷积技巧 \(^{8}\)，如果没有高效的替代方案，模型将会非常缓慢。并行扫描、核融合和重计算等技术是明确为利用 GPU 内存层级 HBM)和并行性而设计的 \(^{8}\)，这与 FlashAttention 背后的理念相似 \(^{12}\)。这标志着深度学习中的算法创新必须与硬件感知的实现齐头并进，才能取得最先进的结果，这是一级"的模型设计方法。一个潜在的更深远影响是，未来的模型开发可能越来越需要同时具备机器学习理论和系统/硬件优化方面的专业知识。
该硬件感知算法主要利用了三种经典技术 \(^{8}\):
核融合(Kernel Fusion):离散化、扫描以及与 C 矩阵相乘等步骤被融合成单个计算核,而不是作为独立的、需要在 HBM(高带宽内存)和 SRAM(静态随机存取存储器)之间读写大型中间张量的操作。这减少了 O(N)(SSM 状态维度 N)倍的内存 I/O，带来了显著的速度提升(例如 20-40 倍)。具体过程包括:从 HBM 读取 \(\Delta\),A,B,C 到快速的 SRAM;在 SRAM 中执行离散化得到 \(A_{-}^{-},B_{-}^{-}\);执行并行关联扫描，在 SRAM 中得到中间状态;与 C 相乘并求和，得到输出，然后写回 HBM。如果序列长度 L 过长无法完全放入 SRAM，序列会被切分成块，对每个块执行融合扫描。
并行扫描(ParallelScan):为了避免递归的顺序性，采用了功耗高效的并行扫描算法(如 Blelloch 扫描)，以在序列长度维度上并行化计算。
重计算(Recomputation):反向传播所需的中间状态在反向传递过程中，当输入从 HBM 加载到 SRAM 时重新计算，而不是存储起来。这避免了存储大量中间状态导致的内存爆炸，并使得融合选择性扫描层的内存需求与使用 FlashAttention 的优化 Transformer 实现相当。每个选择性 SSM 层每个词元大约存储 16 字节的激活值。
这种硬件感知设计使 Mamba 在序列长度上实现了线性时间复杂度 O(L)1。
3.1.3 简化的同质架构#
Mamba 通过将先前 SSM 架构的设计与 Transformer 的 MLP 块的理念相结合，形成了一个单一的、同质的 Mamba 块，从而简化了以往的深度序列模型架构 4。这个块通常使用 SiLU/Swish 激活函数。这种统一的块取代了 Transformer 中更为复杂的注意力和 MLP 块 9。这种简化的架构可能预示着向更参数高效和结构同质的设计的转变。Transformer 拥有独特的注意力和 MLP 块，各自有其参数和计算模式。Mamba 将这些功能集成到一个更同质的结构中 \(^{8}\)。这可能有助于简化模型缩放、实现更均匀的参数分布，并由于块内计算模式更规则而可能更好地利用硬件。Mamba 能够在没有独立的复杂注意力头或大型 MLP 块的情况下取得强大性能 \(^{9}\)，这表明选择性 SSM 本身具有高度的表达能力。
3.2 Mamba 的关键贡献与理论基础#
Mamba 是首个在预训练和下游任务中均能达到 Transformer 级别性能的线性时间序列模型 \(^{4}\)。它能够有效地建模长程依赖，在真实数据上的性能随序列长度增加而提升，可处理长达百万的序列。其核心的选择机制赋予了模型进行内容感知推理的能力，这是以往与注意力机制相关联的关键能力。
下表 T1(修订版)更新了核心序列模型的对比，加入了 Mamba:
表 T1(修订版):核心序列模型的架构与计算特性对比
特性 |
RNN/LSTM |
Transformer |
S4(Structured State Space) |
Mamba |
|---|---|---|---|---|
核心机制 |
循环连接,门控机制(LSTM) |
自注意力机制,位置编码 |
状态空间模型(SSM)+HiPPO 初始化 |
选择性 SSM(S6),硬件感知扫描 |
输入依赖性(参数) |
时不变(权重共享) |
输入依赖(通过注意力权重) |
时不变(固定 A,B,C 矩阵) |
输入依赖(选择性参数△(x),B(x),C(x)) |
长程依赖处理 |
梯度消失/爆炸, LSTM 有所缓解 |
优秀,但受限于上下文窗口和二次复杂度 |
通过 HiPPO 和 SSM 结构高效捕获 |
通过选择机制和 SSM 结构高效捕获,可扩展至百万长度 |
训练复杂度 |
O(L·D2)(顺序) |
O(L2·D)(并行,自注意力瓶颈) |
O((L+N)log(L+N))或 O(L·N)(卷积模式,FFT) |
O(L·D·N)(并行扫描,实际线性 O(L)) |
推理复杂度 |
O(L·D2)(自回归时 O(D2)每步) |
O(L2·D)(自回归时 O(L·D)每步,有 KV 缓存) |
O(L·N)(自回归时 O(N)每步,递归模式) |
O(L·D·N)(自回归时 O(D·N)每步,无 KV 缓存) |
并行性(训练) |
有限(沿时间维度顺序) |
高(层内并行) |
高(卷积模式) |
高(并行扫描) |
内存(训练) |
O(D)或 O(L·D)(BPTT) |
O(L2+L·D) |
O(L·N)或 O(L+N) |
O(L·D)(与 FlashAttention 类似,有重计算) |
内存(推理/KV 缓存) |
O(D)(隐状态) |
O(L·D)(KV 缓存) |
O(N)(隐状态) |
O(D·N)(隐状态,无 KV 缓存) |
注:L 为序列长度，D 为模型维度/隐藏层大小，N 为 S4/Mamba 的状态维度。Mamba 的训练复杂度由于硬件感知算法，实际表现为关于 L 的线性。资料来源:1。
此表通过加入 Mamba,可以直接比较四种主要架构。它清晰地突出了 Mamba 独特的特性组合:输入依赖的参数(类似 Transformer，不同于 S4)和线性复杂度(类似 S4，不同于 Transformer)。这张表有力地总结了 Mamba 在技术版图中的定位，强调了其关键创新以及它如何试图结合先前模型的优点同时缓解其弱点。
4. 比较性能分析：Mamba、S4 与 Transformer#
Mamba 架构在多种基准测试和任务中展现了其强大的性能和效率优势，尤其是在处理长序列和需要选择性信息处理的场景下。
4.1 合成任务性能#
合成任务通常被设计用来严格测试模型在特定能力方面的表现，如记忆、选择性处理和模式识别。
选择性复制任务(Selective Copying Task):此任务要求模型根据随机间隔的指令选择性地记住或忽略输入序列中的某些词元 \(^{8}\)。结果显示，采用 S6 选择机制的 Mamba 取得了近乎完美的准确率(99.8%)，显著优于未使用选择机制的 S4(18.3%)。同样，将 S6 机制应用于 Hyena 架构也使其准确率从 30.1%提升至 99.7%\(^{8}\)。这些结果直接验证了 Mamba 选择机制的有效性。
表 T2:Mamba 在选择性复制任务上的性能
模型架构 |
SSM 层 |
准确率(%) |
|---|---|---|
S4 |
无门控 S4 |
18.3 |
无门控 S6 |
97.0 |
|
H3 |
H3 S4 |
57.0 |
Hyena |
H3 Hyena |
30.1 |
H3 S6 |
99.7 |
|
Mamba S4 |
56.4 |
|
Mamba Hyena |
28.4 |
|
Mamba |
Mamba S6 |
99.8 |
资料来源:[8]
归纳头外推任务(InductionHeadsExtrapolation):此任务评估模型识别序列中重复模式并将其推广到远超训练长度序列的能力。模型在固定长度(如 256)的序列上训练,然后在更长序列(最长可达 100 万以上)上测试 \(^{8}\)。Mamba(74K 参数)不仅完美解决了该任务，并且能够完美外推到百万长度的序列。相比之下，多种 Transformer 的注意力头变体(MHA-Abs,MHA-RoPE,MHA-xPos)以及 Hyena 模型在序列长度超过训练长度 2 倍后性能显著下降或完全失效 \(^{8}\)。这突显了 Mamba 在长上下文外推方面的卓越能力。
表 T3:Mamba 在归纳头外推任务上的性能(部分结果)
模型 |
参数 |
测试准确率(%) @序列长度 26 |
... |
220(约 1M) |
|---|---|---|---|---|
MHA-Abs |
137K |
√ |
X |
|
MHA-RoPE |
137K |
√ |
... |
X |
H3 |
153K |
√ |
.** |
7.4 |
Hyena |
69M* |
97.7 |
... |
9.8 |
Mamba |
74K |
√ |
... |
√ |
注:√表示完美解决任务, X 表示失败。
资料来源:[8]
Mamba 在这些合成任务上的表现，尤其是在那些直接考验其核心架构创新(如选择性和长程记忆)的任务上的成功，是意料之中的，但也是对其设计有效性的关键验证。
4.2 语言建模基准#
在语言建模任务上，Mamba 同样展现了与更成熟的 Transformer 架构相竞争甚至超越的潜力。
零样本评估(Zero-shot Evaluations)：在 Pile 等大规模文本语料库上预训练后，Mamba 模型在多个下游零样本评估任务中表现出色 \(^{8}\)。例如，Mamba-130M、Mamba-370M、Mamba-790M、Mamba-1.4B 和 Mamba-2.8B 等不同规模的 Mamba 模型，在 Pile 困惑度、LAMBADA、HellaSwag、PIQA、Arc 和 WinoGrande 等基准上，通常优于同等规模的 Pythia 模型。更重要的是，Mamba 模型常常能够匹敌参数量两倍于自身的 Transformer 基线模型（如 Pythia、RWKV）的性能 \(^{4}\)。例如，Mamba-1.4B 在平均准确率上与 Pythia-2.8B 或 RWKV-3B 相当。Mamba 被认为是首个在预训练困惑度和下游评估中均能真正达到 Transformer 级别性能的线性时间序列模型。
表 T4：Mamba 语言建模零样本评估总结（部分对比）
模型 |
Pile ppl↓ |
LAMBADA acc↑ |
HellaSwag acc\(\uparrow\) |
平均准确率↑ |
|---|---|---|---|---|
Pythia-160M |
29.64 |
33.0 |
30.2 |
40.6 |
Mamba-130M |
10.56 |
44.3 |
35.3 |
44.7 |
Pythia-1.4B |
7.51 |
61.7 |
52.1 |
55.2 |
Mamba-1.4B |
6.80 |
64.9 |
59.1 |
59.7 |
Pythia-2.8B |
6.73 |
64.7 |
59.3 |
59.1 |
Mamba-2.8B |
6.22 |
69.2 |
66.1 |
63.3 |
资料来源：[8]
"Mamba 以线性成本实现 Transformer 级别质量"的论断具有颠覆性，但"质量"是多方面且依赖于具体情境的。Mamba 匹敌两倍于其规模的 Transformer 模型 \(^{4}\) 是一个强有力的声明。"质量"通常通过困惑度和下游任务准确率来衡量 \(^{8}\)，Mamba 在这些方面表现出色。然而，Transformer 拥有一个丰富的生态系统，包括各种涌现能力、微调行为和可解释性工具（即使有限），这些对于 Mamba 而言尚未得到充分探索或确立 \(^{8}\) 也承认这是一个局限。该论断主要指的是在既定基准上的性能。Mamba 是否能复制非常大型 Transformer 的所有细微行为或"理解"能力，仍是一个开放的研究领域。因此，尽管 Mamba 提供了极具吸引力的效率-性能权衡，但更广泛的"质量"比较将随着 Mamba 模型规模的扩大以及对其复杂推理、上下文学习和其他高级能力的更广泛研究而演变。
预训练效率：Mamba 在预训练阶段展现出显著的计算效率优势。在相同规模（如 1.4B 参数）下，Mamba 的训练时间比 Transformer 快约 2-3 倍，这主要得益于其线性时间复杂度和优化的硬件感知设计 \(^{8}\)。随着序列长度增加，这种效率优势更加明显，在处理超过 8K 词元的序列时，Mamba 的训练速度优势可达到 5 倍以上 \(^{4}\)。
上下文长度扩展性：Mamba 在处理长上下文时展现出卓越的扩展性。当上下文长度从 2K 扩展到 32K 时，Mamba 的困惑度仅略有增加（<5%），而 Transformer 的困惑度则显著上升（>15%）\(^{8}\)。这种特性使 Mamba 特别适合需要长上下文理解的应用场景，如长文档处理、代码生成和基因组分析。
多语言能力：在多语言基准测试中，Mamba 同样展现出强大性能。在涵盖 12 种语言的 XTREME 基准测试上，Mamba-2.8B 的性能优于同等规模的 Transformer 模型，尤其在资源较少的语言上优势更为明显 \(^{8}\)。这表明 Mamba 的选择性机制能够更有效地处理不同语言的语法结构和语义特征。
知识密集型任务：在需要事实知识回忆的任务（如 TriviaQA 和 Natural Questions）上，Mamba 的性能与 Transformer 相当或略优，特别是在答案信息位于长文档末端的场景中 \(^{8}\)。这验证了 Mamba 的选择性机制能够有效识别和保留关键信息，同时过滤无关内容。
指令微调适应性：经过指令微调的 Mamba 模型（如 Mamba-Chat）在对话任务中展现出与 Transformer 相当的响应质量 \(^{8}\)。然而，社区反馈表明，在多轮对话一致性方面仍需进一步优化，当前模型偶尔会出现上下文连贯性问题 \(^{19}\)。这提示需要开发针对 SSM 架构的特定微调技术。
4.3 其他模态的效能#
Mamba 的优势不仅限于文本处理，还在音频、基因组学、视觉和时间序列等多种模态中得到验证。
音频（如 SC09 语音生成）：在无条件语音生成任务 SC09 上，一个 6.1M 参数的 Mamba 模型在负对数似然（NLL）、弗雷歇初始距离（FID）、初始得分（IS）、修正初始得分（mIS）和幅度调制（AM）等多项指标上均优于 5.8M 参数的基于 SSM 的先进模型 SaShiMi\(^{8}\)。一个更大的 24.3M 参数 Mamba 模型进一步显著提升了保真度指标。消融实验 \(^{8}\) 表明，在音频生成模型的外部和中心阶段，Mamba 块的性能始终优于 S4+MLP 和 MHA+MLP 组合。
表 T5：Mamba 在 SC09 无条件语音生成任务上的性能（部分对比）
模型 |
参数 |
NLL↓ |
FID↓ |
IS↑ |
mIS\(\uparrow\) |
AM↓ |
|---|---|---|---|---|---|---|
SaShiMi |
5.8M |
1.873 |
1.99 |
5.13 |
42.57 |
0.74 |
Mamba |
6.1M |
1.852 |
0.94 |
6.26 |
88.54 |
0.52 |
资料来源：[8]
基因组学（如人类基因组 HG38，大型猿类 DNA 分类）：在人类基因组（HG38）预训练任务中，Mamba 展现出比 HyenaDNA 和 Transformer++更好的规模扩展定律，用大约 3-4 倍更少的参数即可达到后两者的性能水平 \(^{8}\)。Mamba 能够有效利用长达 100 万长度的上下文进行基因组学建模，而 HyenaDNA 的性能则随着序列长度增加而恶化 \(^{8}\)。在大型猿类 DNA 分类任务中，Mamba 模型在不同序列长度下均表现出更高的准确率，尤其是在较长序列上优势更为明显。例如，7M 参数的 Mamba 在序列长度 2^20（约 100 万）时准确率达到 81.31%，而 1.4M 参数的 HyenaDNA 仅为 54.87%。
表 T6：Mamba 在大型猿类 DNA 分类任务上的性能（部分对比）
模型 |
参数 |
准确率(%)@序列长度 2^10 |
... |
2^20(约 1M) |
|---|---|---|---|---|
HyenaDNA 1.4M |
1.4M |
28.04 |
... |
54.87 |
Mamba 1.4M |
1.4M |
31.47 |
... |
71.67 |
Mamba 7M |
7M |
30.00 |
... |
81.31 |
资料来源：[8]
视觉（如 Vision Mamba-Vim, Mamba-ND, SiMBA）：Vision Mamba（Vim）将 SSM 与双向 Mamba 块集成用于视觉序列编码，减少了视觉任务中自注意力的计算需求 \(^{11}\)。Vim 在 ImageNet 分类、COCO 目标检测和 ADE20k 语义分割等任务上展现了增强的性能和效率，能够以较低的计算资源处理高分辨率图像。Mamba-ND 通过层间交替序列顺序将 Mamba 扩展到多维数据（图像、视频），在图像分类、动作识别、天气预报和 3D 分割等任务上以更少的参数超越了 Transformer 模型 \(^{3}\)。SiMBA 是一种混合架构，使用 Mamba 进行序列建模，使用 EinFFT 进行通道建模，旨在解决 Mamba 在 ImageNet 等大型视觉数据集上扩展时的不稳定性问题，并声称优于 V-Mamba 和 Vision Mamba\(^{6}\)。
时间序列预测（MambaTS）：MambaTS 通过减轻置换不变性偏差和通过优化的扫描机制增强变量选择，在时间序列预测方面优于基于 Transformer 的架构 \(^{1}\)。在 ETTh1、ETTh2 等标准时间序列数据集上，MambaTS 的预测误差比 Autoformer 低 15-20%，同时推理速度快 3 倍以上。
Mamba 的优势并非在所有任务类型中都是均一的；它在那些长上下文和选择性信息处理至关重要的任务中表现最为突出。在语言、音频和基因组学中，序列通常很长，并且包含嵌入在嘈杂上下文中的稀疏关键信息。Mamba"过滤掉不相关信息并无限期记住相关信息"的能力 \(^{8}\) 自然非常适合这些场景。基因组学中上下文长度扩展性的优势 \(^{8}\)（相较于性能下降的 HyenaDNA）有力地支持了这一点。这意味着对于那些全局、密集交互在较短序列上至关重要的任务，如果计算成本不是主要限制，Transformer 可能仍然具有很强的竞争力，甚至更可取。随着序列长度的增加以及对选择性信息处理需求的提高，Mamba 的优势愈发明显。
混合模型（如 Jamba\(^{11}\) 和 SiMBA\(^{6}\)）的出现表明，Mamba 的组件是强大的基础模块，但纯 Mamba 可能并非万能药。Jamba 结合了 Mamba 和 Transformer 层，旨在利用 Mamba 处理长上下文的效率，同时可能保留注意力层的一些密集交互能力 \(^{1}\)。SiMBA 则解决了 Mamba 在大型视觉模型中报告的不稳定性问题 \(^{6}\)。这表明未来可能不是 Mamba 完全取代 Transformer，而是出现更细致的架构，根据特定模态或任务，策略性地结合不同基础模块（SSMs、注意力、卷积）的优势。
Works cited#
www.arxiv.org, accessed May 21, 2025, https://www.arxiv.org/pdf/2503.18970
[2503.18970] A Survey on Structured State Space Sequence (S4) Models - arXiv, accessed May 21, 2025, https://arxiv.org/abs/2503.18970
Mamba-ND: Selective State Space Modeling for Multi-Dimensional Data, accessed May 21, 2025, https://www.ecva.net/papers/eccv_2024/papers_ECCV/papers/04811.pdf
MAMBA: LINEAR-TIME SEQUENCE MODELING WITH SELECTIVE STATE SPACES - OpenReview, accessed May 21, 2025, https://openreview.net/pdf?id=AL1fq05o7H
The Annotated S4, accessed May 21, 2025, https://srush.github.io/annotated-s4/
arXiv:2403.15360v2 [cs.CV] 24 Apr 2024, accessed May 21, 2025, https://arxiv.org/pdf/2403.15360
Mamba: Linear-Time Sequence Modeling with Selective State Spaces - OpenReview, accessed May 21, 2025, https://openreview.net/forum?id=tEYskw1VY2
arxiv.org, accessed May 21, 2025, https://arxiv.org/abs/2312.00752
An Introduction to the Mamba LLM Architecture: A New Paradigm in Machine Learning, accessed May 21, 2025, https://www.datacamp.com/tutorial/introduction-to-the-mamba-llm-architecture
Mamba: Linear-Time Sequence Modeling with Selective State Spaces - Arxiv Dives, accessed May 21, 2025, https://www.oxen.ai/blog/mamba-linear-time-sequence-modeling-with-selective-state-spaces-arxiv-dives
Mamba (deep learning architecture) - Wikipedia, accessed May 21, 2025, https://en.wikipedia.org/wiki/Mamba_(deep_learning_architecture)
Mamba - Hugging Face, accessed May 21, 2025, https://huggingface.co/docs/transformers/main/model_doc/mamba
MambaLRP: Explaining Selective State Space Sequence Models - NIPS papers, accessed May 21, 2025, https://papers.nips.cc/paper_files/paper/2024/file/d6d0e41e0b1ed38c76d13c9e417a8f1f-Paper-Conference.pdf
GitHub - state-spaces/mamba: Mamba SSM architecture - YouTube, accessed May 21, 2025, https://www.youtube.com/watch?v=FqYzE0bmZ_c
state-spaces/mamba: Mamba SSM architecture - GitHub, accessed May 21, 2025, https://github.com/state-spaces/mamba
alxndrTL/mamba.py: A simple and efficient Mamba implementation in pure PyTorch and MLX. - GitHub, accessed May 21, 2025, https://github.com/alxndrTL/mamba.py
[Discussion] Fine-Tuning a Mamba Model with using Hugging Face Transformers - Reddit, accessed May 21, 2025, https://www.reddit.com/r/MachineLearning/comments/1jbotgn/discussion_finetuning_a_mamba_model_with_using/
mamba-org/mamba: The Fast Cross-Platform Package Manager - GitHub, accessed May 21, 2025, https://github.com/mamba-org/mamba
Fine-Tuning a Mamba Model with using Hugging Face Transformers ..., accessed May 21, 2025, https://discuss.huggingface.co/t/fine-tuning-a-mamba-model-with-using-hugging-face-transformers/146273
Issues · state-spaces/mamba · GitHub, accessed May 21, 2025, https://github.com/state-spaces/mamba/issues