Mamba：基于选择性状态空间的线性时间序列建模

阿尔伯特・顾（Albert Gu）*¹ 和特里・陶（Tri Dao）*²1 卡内基梅隆大学机器学习系 2 普林斯顿大学计算机科学系电子邮箱：agu@cs.cmu.edu，tri@tridao.me

摘要

如今支撑深度学习领域多数热门应用的基础模型，几乎普遍基于 Transformer 架构及其核心的注意力机制。为解决 Transformer 在长序列上的计算低效问题，研究人员已开发出众多亚二次时间复杂度架构，例如线性注意力、门控卷积、循环模型以及结构化状态空间模型（SSMs）等。然而，这些模型在语言等重要模态上的表现始终不及注意力机制。我们发现这类模型的关键缺陷在于缺乏基于内容的推理能力，并为此做出多项改进。首先，让状态空间模型的参数成为输入的函数，这一简单改动解决了它们在离散模态上的短板，使模型能够根据当前令牌，选择性地沿序列长度维度传播或遗忘信息。其次，尽管这一改动导致无法使用高效卷积，但我们设计了一种硬件感知的循环模式并行算法。我们将这些选择性状态空间模型整合到一个简化的端到端神经网络架构中，该架构既无注意力机制，也无多层感知机（MLP）块，即 Mamba。Mamba 兼具快速推理能力（吞吐量是 Transformer 的 5 倍）和序列长度线性缩放特性，在长达百万长度序列的真实数据上性能持续提升。作为通用序列模型骨干网络，Mamba 在语言、音频和基因组学等多种模态上均实现了最先进的性能。在语言建模任务中，我们的 Mamba-3B 模型不仅优于同规模的 Transformer，还能与两倍规模的 Transformer 相媲美，无论是预训练阶段还是下游任务评估均是如此。

1 引言

基础模型（FMs），即先在海量数据上预训练，再适配下游任务的大型模型，已成为现代机器学习中的高效范式。这些基础模型的骨干网络通常是序列模型，能够处理来自语言、图像、语音、音频、时间序列和基因组学等多个领域的任意输入序列。尽管这一概念并不局限于特定模型架构，但现代基础模型主要基于一种序列模型 ——Transformer 及其核心的注意力层。自注意力机制的有效性源于其能够在上下文窗口内密集路由信息，从而建模复杂数据。然而，这一特性也带来了根本性缺陷：无法建模有限窗口之外的信息，且计算复杂度随窗口长度呈二次方缩放。已有大量研究致力于开发更高效的注意力变体以克服这些缺陷，但往往会牺牲使其高效的核心特性。迄今为止，尚未有任何一种变体被证明能在跨领域大规模场景下取得实证有效性。

近年来，结构化状态空间序列模型（SSMs）逐渐成为序列建模领域的潜力架构。这类模型可被解读为循环神经网络（RNNs）与卷积神经网络（CNNs）的结合体，同时借鉴了经典状态空间模型的思想。它们能以循环或卷积两种方式高效计算，计算复杂度随序列长度呈线性或近线性缩放。此外，在特定数据模态下，它们具备建模长程依赖的严谨机制，并在长序列基准测试中占据主导地位。多种状态空间模型变体已在音频、视觉等连续信号数据领域取得成功。但在处理文本等离散且信息密集型数据时，它们的表现却不尽如人意。

我们提出一类新的选择性状态空间模型，通过多方面改进现有工作，在实现 Transformer 级建模能力的同时，保持序列长度线性缩放特性。

选择机制

首先，我们发现现有模型的关键局限：难以以输入依赖的方式高效筛选数据（即聚焦或忽略特定输入）。基于选择性复制和归纳头这类重要合成任务的直觉，我们设计了一种简单的选择机制 —— 将状态空间模型的参数以输入为条件进行参数化。这使得模型能够过滤无关信息，并无限期记忆相关信息。

硬件感知算法

这一简单改动给模型计算带来了技术挑战：事实上，所有现有状态空间模型都必须满足时间不变性和输入不变性，才能保证计算效率。我们通过一种硬件感知算法克服了这一难题 —— 该算法采用扫描而非卷积的方式循环计算模型，但不实例化扩展状态，从而避免了 GPU 存储层次结构不同级别之间的 IO 访问。最终实现无论在理论上（序列长度线性缩放，而所有基于卷积的状态空间模型均为伪线性缩放），还是在现代硬件上（A100 GPU 上速度提升高达 3 倍），都优于以往方法。

架构设计

我们简化了现有的深度序列模型架构，将现有状态空间模型架构设计与 Transformer 的 MLP 块整合为单个块，形成了一种简洁同质的架构设计（Mamba），该架构融合了选择性状态空间。

选择性状态空间模型及其延伸的 Mamba 架构，是完全循环模型，具备成为通用基础模型骨干网络的关键特性：（1）高质量：选择性赋予模型在语言、基因组学等密集模态上的强劲性能；（2）快速训练与推理：训练阶段计算和内存随序列长度线性缩放，推理阶段自回归展开时无需缓存先前元素，每步仅需恒定时间；（3）长上下文：高质量与高效率的结合，使其在长达百万长度的序列上性能持续提升。

我们从预训练质量和特定领域任务性能两方面，在多种模态和场景下实证验证了 Mamba 作为通用序列基础模型骨干网络的潜力：

合成任务：在选择性复制、归纳头这类被认为是大型语言模型核心能力的合成任务上，Mamba 不仅能轻松解决，还能外推至无限长序列（超过 100 万个令牌）。

音频与基因组学：在音频波形和 DNA 序列建模任务中，Mamba 在预训练质量和下游指标上均优于 SaShiMi、Hyena 和 Transformer 等现有最先进模型（例如，在一个具有挑战性的语音生成数据集上，将 FID 值降低一半以上）。在这两种场景下，随着上下文长度增加至百万级，其性能持续提升。

语言建模：Mamba 是首个真正实现 Transformer 级性能的线性时间序列模型，无论是预训练困惑度还是下游任务评估均表现优异。在参数规模高达 10 亿的缩放定律实验中，我们证明 Mamba 的性能超过了众多基准模型，包括基于 LLaMa 的强现代 Transformer 训练方案。Mamba 语言模型的生成吞吐量是同规模 Transformer 的 5 倍，且 Mamba-3B 的质量与两倍规模的 Transformer 相当（例如，在常识推理任务上，平均得分比 Pythia-3B 高 4 分，甚至超过 Pythia-7B）。

模型代码和预训练检查点已开源至https://github.com/state-spaces/mamba。

2 状态空间模型

结构化状态空间序列模型（S4）是深度学习领域一类新兴的序列模型，与循环神经网络、卷积神经网络以及经典状态空间模型密切相关。它们的设计灵感来源于特定的连续系统，该系统通过隐含的 latent 状态将一维函数或序列映射为输出。

具体而言，S4 模型由四个参数（Δ、A、B、C）定义，通过两个阶段实现序列到序列的转换：\(h'(t)=A h(t)+B x(t) \quad (1a)\)\(y(t) = C h(t) \quad (1b)\)\(h_{t}=\overline{A} h_{t-1}+\overline{B} x_{t} \quad (2a)\)\(y_t= C h_t \quad (2b)\)\(\overline{K}=(C \overline{B}, C \overline{A B}, ..., C \overline{A}^{k} \overline{B}, ...) \quad (3a)\)\(y= x*K \quad (3b)\)

离散化

第一阶段通过固定公式将 “连续参数”（Δ、A、B）转换为 “离散参数”（\(\overline{A}\)、\(\overline{B}\)），即\(\overline{A}=f_{A}(\Delta, A)\)、\(\overline{B}=f_{B}(\Delta, A, B)\)，其中（\(f_{A}\)、\(f_{B}\)）被称为离散化规则。可采用多种规则，例如公式（4）定义的零阶保持（ZOH）规则：\(\overline{A}=exp (\Delta A)\)\(\overline{B}=(\Delta A)^{-1}(exp (\Delta A)-I) \cdot \Delta B \quad (4)\)

离散化与连续时间系统密切相关，能赋予模型分辨率不变性等额外特性，并自动确保模型适当归一化。它还与循环神经网络的门控机制存在关联，这一点我们将在 3.5 节进一步探讨。但从计算机制来看，离散化可简单视为状态空间模型前向传播计算图的第一步。部分状态空间模型变体可绕过离散化步骤，直接对（\(\overline{A}\)、\(\overline{B}\)）进行参数化，这种方式可能更易于推理。

计算方式

参数从（Δ、A、B、C）转换为（\(\overline{A}\)、\(\overline{B}\)、C）后，模型可通过两种方式计算：线性循环（公式 2）或全局卷积（公式 3）。通常，模型在训练阶段采用卷积模式（公式 3）以实现高效并行计算（可提前获取整个输入序列），在自回归推理阶段切换为循环模式（公式 2）以实现高效计算（逐时间步处理输入）。

线性时间不变性（LTI）

公式（1）至（3）的一个重要特性是模型动态特性不随时间变化。换句话说，（Δ、A、B、C）以及由此推导的（\(\overline{A}\)、\(\overline{B}\)）在所有时间步均保持固定。这一特性被称为线性时间不变性（LTI），与循环和卷积密切相关。通俗地说，我们认为线性时间不变状态空间模型等价于任何线性循环（公式 2a）或卷积（公式 3b），并将线性时间不变性作为这类模型的统称。

迄今为止，所有结构化状态空间模型均为线性时间不变模型（例如，以卷积形式计算），这主要受限于基本的效率约束（详见 3.3 节）。然而，本研究的核心见解是：线性时间不变模型在建模特定类型数据时存在根本性局限，而我们的技术贡献在于移除线性时间不变性约束的同时，克服效率瓶颈。

结构与维度

结构化状态空间模型之所以得名，是因为高效计算它们还需要对 A 矩阵施加结构约束。最常用的结构是对角结构，我们的研究也采用了这一结构。在这种情况下，A∈ℝ^N×N、B∈ℝ^N×1、C∈ℝ^1×N 矩阵均可通过 N 个数值表示。为处理批量大小为 B、长度为 L、通道数为 D 的输入序列 x，状态空间模型需独立应用于每个通道。此时，每个输入的总隐藏状态维度为 D×N，在整个序列长度上计算需消耗 O (BLDN) 的时间和内存，这也是 3.3 节要解决的核心效率瓶颈。

通用状态空间模型

需要说明的是，“状态空间模型” 这一术语含义广泛，泛指任何包含 latent 状态的循环过程。在不同学科中，它被用于指代多种不同概念，包括强化学习中的马尔可夫决策过程（MDP）、计算神经科学中的动态因果建模（DCM）、控制领域的卡尔曼滤波器、机器学习中的隐马尔可夫模型（HMM）和线性动态系统（LDS），以及深度学习中的循环模型（有时也包括卷积模型）。

在本文中，“SSM” 专指结构化状态空间模型或 S4 模型，且这两个术语可互换使用。为方便表述，我们也可能将这类模型的变体（例如，侧重于线性循环或全局卷积视角的模型）纳入其中，并在必要时说明差异。

状态空间模型架构

状态空间模型是独立的序列转换模块，可灵活整合到端到端神经网络架构中（我们有时也将状态空间模型架构称为 SSNNs，其与状态空间模型层的关系类似于卷积神经网络与线性卷积层的关系）。以下介绍几种最著名的状态空间模型架构，它们也将作为我们的主要基准模型：

线性注意力：是自注意力的一种近似，包含循环结构，可视为退化的线性状态空间模型。

H3：将这一循环结构推广至 S4；其架构可视为在两个门控连接之间嵌入一个状态空间模型，并在主状态空间模型层之前插入一个标准局部卷积（称为移位状态空间模型）。

Hyena：采用与 H3 相同的架构，但将 S4 层替换为 MLP 参数化的全局卷积。

RetNet：在架构中添加了额外的门控，并使用更简单的状态空间模型，允许通过一种多注意力头（MHA）变体而非卷积实现并行计算。

RWKV：是一种最新的循环神经网络，专为语言建模设计，基于另一种线性注意力近似 —— 无注意力 Transformer。其核心的 “WKV” 机制涉及线性时间不变循环，可视为两个状态空间模型的比值。

其他密切相关的状态空间模型和架构将在扩展相关工作（附录 B）中进一步讨论。特别需要强调的是 S5、QRNN 和 SRU，我们认为它们与我们的核心选择性状态空间模型最为相关。

3 选择性状态空间模型

我们首先基于合成任务的直觉，阐述选择机制的设计动机（3.1 节），然后说明如何将该机制整合到状态空间模型中（3.2 节）。由此产生的时变状态空间模型无法使用卷积，这就带来了如何高效计算的技术挑战。我们通过一种利用现代硬件存储层次结构的硬件感知算法克服了这一挑战（3.3 节）。接着，我们描述了一种无注意力机制、甚至无 MLP 块的简化状态空间模型架构（3.4 节）。最后，讨论选择机制的其他特性（3.5 节）。

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

我们认为序列建模的核心问题是将上下文压缩到更小的状态中。事实上，我们可以从这一角度审视主流序列模型的权衡取舍。例如，注意力机制之所以高效且低效，是因为它完全不压缩上下文。这一点从自回归推理需显式存储整个上下文（即键值缓存）便可看出，这直接导致 Transformer 推理的线性时间复杂度和训练的二次时间复杂度。另一方面，循环模型之所以高效，是因为它们具有有限状态，这意味着推理的恒定时间复杂度和训练的线性时间复杂度。然而，它们的有效性受限于状态对上下文的压缩效果。

为理解这一原理，我们聚焦于两个合成任务示例：

选择性复制任务：对经典复制任务进行修改，通过改变需记忆令牌的位置，要求模型具备基于内容的推理能力，以记忆相关令牌（有色）并过滤无关令牌（白色）。

归纳头任务：这是一个著名机制，被认为是解释大型语言模型上下文学习能力的核心，要求模型具备上下文感知推理能力，在适当的上下文中生成正确输出（黑色）。

这些任务揭示了线性时间不变模型的失效模式。从循环视角来看，它们的恒定动态特性（例如，公式（2）中的（\(\overline{A}\)、\(\overline{B}\)）转换）无法让模型从上下文中筛选正确信息，也无法以输入依赖的方式影响沿序列传递的隐藏状态。从卷积视角来看，已知全局卷积能够解决标准复制任务（仅需时间感知），但由于缺乏内容感知能力，难以处理选择性复制任务。具体而言，输入与输出之间的间距是变化的，无法通过静态卷积核建模。

总之，序列模型的效率与有效性权衡取决于状态的压缩效果：高效模型必须具备小状态，而有效模型的状态必须包含上下文中的所有必要信息。由此，我们提出构建序列模型的核心原则 —— 选择性，即基于上下文的能力，能够聚焦或过滤输入到序列状态中的信息。具体而言，选择机制控制信息沿序列维度的传播或交互方式（详见 3.5 节）。

3.2 基于选择机制改进状态空间模型

将选择机制整合到模型中的一种方法，是让影响序列交互的参数（例如，循环神经网络的循环动态或卷积神经网络的卷积核）成为输入依赖的函数。

算法 1 和算法 2 展示了我们使用的主要选择机制。核心差异在于将 A、B、C 等多个参数设为输入的函数，并相应调整张量形状。特别需要强调的是，这些参数现在包含长度维度 L，这意味着模型从时间不变转变为时变。（形状标注详见 2 节）。这一变化导致模型不再等价于卷积（公式 3），其效率影响将在下文讨论。

我们具体选择\(s_{B}(x)=Linear_{N}(x)\)、\(s_{C}(x)=Linear_{N}(x)\)、\(s_{\Delta}(x)=Broadcast_{D}(Linear_{1}(x))\)，且\(\tau_{\Delta}\)=softplus，其中\(Linear _{d}\)表示参数化投影至 d 维。\(s_{\Delta}\)和\(\tau_{\Delta}\)的选择源于与循环神经网络门控机制的关联（详见 3.5 节）。

算法 1 状态空间模型（S4）	算法 2 状态空间模型 + 选择机制（S6）
输入：x : (B, L, D)	输入：x : (B, L, D)
输出：y : (B, L, D)	输出：y : (B, L, D)
1: A : (D, N) ← 参数（表示结构化 N×N 矩阵）	1: A : (D, N) ← 参数（表示结构化 N×N 矩阵）
2: B : (D, N) ← 参数	2: B : (B, L, N) ← \(s_{B}(x)\)
3: C : (D, N) ← 参数	3: C : (B, L, N) ← \(s_{C}(x)\)
4: Δ : (D) ← \(\tau_{\Delta}\)(参数)	4: Δ : (B, L, D) ← \(\tau_{\Delta}\)(参数 + \(s_{\Delta}(x)\))
5: A, B : (D, N) ← 离散化 (Δ, A, B)	5: A, B : (B, L, D, N) ← 离散化 (Δ, A, B)
6: y ← 状态空间模型 (A, B, C)(x)（时间不变：循环或卷积）	6: y ← 状态空间模型 (A, B, C)(x)（时变：仅循环（扫描））
7: 返回 y	7: 返回 y

3.3 选择性状态空间模型的高效实现

卷积和注意力这类硬件友好型原语已得到广泛应用。本文旨在让选择性状态空间模型在现代硬件（GPU）上也能高效运行。选择机制的设计非常自然，早期研究已尝试整合特殊形式的选择，例如让循环状态空间模型中的 Δ 随时间变化。然而，如前所述，状态空间模型的核心局限在于计算效率，这也是 S4 及所有变体采用线性时间不变（非选择性）模型（通常以全局卷积形式）的原因。

3.3.1 现有模型的动机

我们首先回顾这一动机，并概述克服现有方法局限的思路：

总体而言，循环模型（如状态空间模型）始终需要在表达能力和速度之间权衡：如 3.1 节所述，隐藏状态维度越大的模型，理论上效果越好，但速度越慢。因此，我们希望在不增加速度和内存成本的前提下，最大化隐藏状态维度。

循环模式比卷积模式更灵活，因为后者（公式 3）是通过扩展前者（公式 2）推导而来。然而，这需要计算并实例化形状为（B, L, D, N）的 latent 状态 h，其规模（乘以状态空间模型状态维度 N）远大于形状为（B, L, D）的输入 x 和输出 y。因此，研究人员提出了更高效的卷积模式，无需计算状态，仅实例化大小为（B, L, D）的卷积核（公式 3a）。

现有线性时间不变状态空间模型利用循环 - 卷积双重形式，将有效状态维度提升 N（约 10-100）倍，远超传统循环神经网络，且无效率损失。

3.3.2 选择性扫描概述：硬件感知状态扩展

选择机制的设计旨在克服线性时间不变模型的局限；同时，我们需要重新审视状态空间模型的计算问题。我们通过三种经典技术解决这一问题：核融合、并行扫描和重计算，并得出两个核心观察：

朴素循环计算的浮点运算量（FLOPs）为 O (BLDN)，而卷积计算的浮点运算量为 O (BLD log (L))，且前者的常数因子更低。因此，对于长序列和不算太大的状态维度 N，循环模式的浮点运算量实际上更少。

面临的两大挑战是循环的顺序性和巨大的内存占用。为解决后者，与卷积模式类似，我们可以尝试不实例化完整状态 h。

核心思路是利用现代加速器（GPU）的特性，仅在存储层次结构的高效级别实例化状态 h。具体而言，大多数操作（除矩阵乘法外）均受内存带宽限制，包括扫描操作。我们通过核融合减少内存 IO，相比标准实现实现显著提速。

具体来说，我们不将大小为（B, L, D, N）的扫描输入（\(\overline{A}\)、\(\overline{B}\)）准备在 GPU 高带宽内存（HBM）中，而是将状态空间模型参数（Δ、A、B、C）直接从慢速高带宽内存加载到快速静态随机存取存储器（SRAM），在静态随机存取存储器中执行离散化和循环计算，最后将大小为（B, L, D）的最终输出写回高带宽内存。

为避免顺序循环，我们观察到尽管模型不再是线性的，但仍可通过高效并行扫描算法实现并行化。

最后，我们还需避免保存反向传播所需的中间状态。我们巧妙地应用经典的重计算技术以降低内存需求：不存储中间状态，而是在反向传播时，将输入从高带宽内存加载到静态随机存取存储器时重新计算。最终，融合选择性扫描层的内存需求与采用 FlashAttention 优化的 Transformer 实现相当。

融合核和重计算的细节详见附录 D。完整的选择性状态空间模型层和算法如图 1 所示。

3.4 简化的状态空间模型架构

与结构化状态空间模型类似，选择性状态空间模型是独立的序列转换模块，可灵活整合到神经网络中。H3 架构是大多数著名状态空间模型架构的基础，其设计通常包含一个受线性注意力启发的块，与 MLP 块交错排列。我们通过将这两个组件整合为一个块，并进行同质堆叠，简化了该架构。这一设计灵感来源于门控注意力单元（GAU），后者对注意力机制采用了类似的整合思路。

该架构通过可控的扩展因子 E 扩展模型维度 D。对于每个块，大部分参数（3ED²）集中在线性投影（输入投影 2ED²，输出投影 ED²），而内部状态空间模型的参数占比较小。状态空间模型的参数（Δ、B、C 的投影以及矩阵 A）相对更少。我们将该块与标准归一化和残差连接交错堆叠，形成 Mamba 架构。在实验中，我们固定 E=2，并使用两个块堆叠，以匹配 Transformer 交错的多注意力头（MHA）和 MLP 块的 12D² 参数规模。我们采用 SiLU/Swish 激活函数，使门控 MLP 成为流行的 “SwiGLU” 变体。最后，我们还添加了一个可选的归一化层（选择 LayerNorm），这一设计灵感来源于 RetNet 在类似位置使用归一化层的做法。

3.5 选择机制的特性

选择机制是一个更广泛的概念，可通过多种方式应用，例如应用于传统循环神经网络或卷积神经网络、不同参数（如算法 2 中的 A），或使用不同的转换函数 s (x)。

3.5.1 与门控机制的关联

我们强调最重要的关联：循环神经网络的经典门控机制是状态空间模型选择机制的一个实例。循环神经网络门控与连续时间系统离散化之间的关联已得到充分证实。事实上，定理 1 是对相关研究的扩展，将其推广至零阶保持离散化和输入依赖门控（证明详见附录 C）。更广泛地说，状态空间模型中的 Δ 可被视为循环神经网络门控机制的广义形式。与现有研究一致，我们认为状态空间模型的离散化是启发式门控机制的严谨基础。

定理 1：当 N=1、A=-1、B=1、\(s_{\Delta}=Linear(x)\)且\(\tau_{\Delta}\)=softplus 时，选择性状态空间模型循环（算法 2）可表示为：\(g_{t}=\sigma (Linear(x_{t}))\)\(h_{t}=(1-g_{t})h_{t-1}+g_{t}x_{t}\)

如 3.2 节所述，我们对\(s_{\Delta}\)和\(\tau_{\Delta}\)的具体选择正是基于这一关联。特别需要说明的是，如果某个输入\(x_{t}\)应被完全忽略（如合成任务所需），则所有 D 通道都应忽略它，因此我们先将输入投影至 1 维，再通过 Δ 进行重复 / 广播。

3.5.2 选择机制的解读

我们详细阐述选择机制的三个具体机制性作用：

可变间距

选择性允许模型过滤掉感兴趣输入之间的无关噪声令牌。选择性复制任务便是典型示例，但这一特性在常见数据模态中普遍存在，尤其对于离散数据 —— 例如语言中的填充词 “um”。这一特性的实现机制是，模型能够选择性过滤任何特定输入\(x_{t}\)，例如在门控循环神经网络中（定理 1），当\(g_{t} \to 0\)时。

上下文过滤

实证观察表明，许多序列模型的性能并不会随上下文长度增加而提升，尽管理论上更多的上下文应带来更优性能。一个可能的解释是，许多序列模型在必要时无法有效忽略无关上下文；全局卷积（及一般线性时间不变模型）便是直观示例。相反，选择性模型可在任何时候重置状态，去除多余历史，因此其性能理论上随上下文长度单调提升。

边界重置

在多个独立序列拼接的场景中，Transformer 可通过特定的注意力掩码将它们分开，而线性时间不变模型会在序列之间泄露信息。选择性状态空间模型也可在边界处重置状态（例如，\(\Delta_{t} \to \infty\)，或定理 1 中\(g_{t} \to 1\)）。这类场景可能是人为的（例如，为提高硬件利用率而打包文档），也可能是自然的（例如，强化学习中的 episode 边界）。

此外，我们详细阐述每个选择性参数的作用：

Δ 的解读：一般而言，Δ 控制对当前输入\(x_{t}\)的聚焦或忽略程度。它是循环神经网络门控的广义形式（例如，定理 1 中的\(g_{t}\)）：机制上，大 Δ 会重置状态 h 并聚焦当前输入 x，而小 Δ 会保留状态并忽略当前输入。状态空间模型（1）-（2）可被解读为通过时间步 Δ 离散化的连续系统，在这一背景下，大 Δ→∞表示系统长时间聚焦当前输入（因此 “选择” 它并遗忘当前状态），而小 Δ→0 表示输入是瞬时的，将被忽略。

A 的解读：尽管 A 参数也可设计为选择性的，但它最终通过\(\overline{A}=\exp (\Delta A)\)（离散化公式 4）与 Δ 交互，影响模型。因此，Δ 的选择性足以确保（\(\overline{A}\)、\(\overline{B}\)）的选择性，这也是性能提升的主要来源。我们推测，除 Δ 外（或替代 Δ）让 A 具备选择性，可能会获得类似性能，但为简洁起见，本研究未采用这一设计。

B 和 C 的解读：如 3.1 节所述，选择性的核心特性是过滤无关信息，使序列模型的上下文能够被压缩到高效状态中。在状态空间模型中，将 B 和 C 设计为选择性的，可更精细地控制是否允许输入\(x_{t}\)进入状态\(h_{t}\)，或状态进入输出\(y_{t}\)。这可分别解读为允许模型基于内容（输入）和上下文（隐藏状态）调节循环动态。

3.6 模型补充细节

实值与复值：大多数现有状态空间模型在状态 h 中使用复数值，这对于感知模态的许多任务至关重要。然而，实证观察表明，纯实值状态空间模型在某些场景下表现良好，甚至更优。我们默认使用实值，除一个任务外，实值在所有任务中均表现良好；我们推测，复值 - 实值的权衡与数据模态的连续 - 离散谱相关，复数值对连续模态（如音频、视频）有帮助，但对离散模态（如文本、DNA）无明显优势。

初始化：大多数现有状态空间模型建议特殊初始化，尤其是在复值情况下，这在低数据场景等多种设置中可能有帮助。复值情况下的默认初始化采用 S4D-Lin，实值情况下采用 S4D-Real，这两种初始化均基于 HIPPO 理论。它们分别将 A 的第 n 个元素定义为 - 1/2 + ni 和 -(n+1)。然而，我们认为许多初始化方法都能正常工作，尤其是在大数据和实值状态空间模型场景下；4.6 节将讨论部分消融实验。

Δ 的参数化：我们将 Δ 的选择性调整定义为\(s_{\Delta}(x)=Broadcast _{D}(Linear_{1}(x))\)，这一设计源于 Δ 的机制（3.5 节）。我们观察到，可将其从 1 维推广至更大的维度 R。我们将 R 设置为 D 的一小部分，与块中的主要线性投影相比，其参数数量可忽略不计。此外，广播操作也可视为另一种线性投影，初始化为特定的 1 和 0 模式；如果该投影可训练，则可得到替代形式\(s_{\Delta}(x)=Linear_{D}(Linear_{R}(x))\)，可视为低秩投影。

在实验中，Δ 参数（可视为偏置项）的初始化遵循现有状态空间模型的做法，设置为\(\tau_{\Delta}^{-1}( Uniform ([0.001,0.1]))\)。

注 3.1：为简洁起见，在实验结果中，我们有时将选择性状态空间模型简称为 S6 模型，因为它们是带有选择机制的 S4 模型，通过扫描计算。

4 实证评估

4.1 节测试 Mamba 解决 3.1 节提出的两个合成任务的能力。随后，我们在三个领域进行评估，每个领域均包含自回归预训练和下游任务：

4.2 节：语言模型预训练（缩放定律）和零样本下游评估。

4.3 节：DNA 序列预训练和长序列分类任务微调。

4.4 节：音频波形预训练和自回归生成语音片段的质量。

最后，4.5 节展示 Mamba 在训练和推理时的计算效率，4.6 节对架构和选择性状态空间模型的各个组件进行消融实验。

4.1 合成任务

这些任务的完整实验细节（包括任务描述和训练方案）详见附录 E.1。

4.1.1 选择性复制

复制任务是序列建模领域研究最广泛的合成任务之一，最初用于测试循环模型的记忆能力。如 3.1 节所述，线性时间不变状态空间模型（线性循环和全局卷积）可轻松解决该任务，只需跟踪时间而非对数据进行推理；例如，通过构建长度完全匹配的卷积核。这一点已在早期全局卷积研究中得到明确验证。选择性复制任务通过随机化令牌之间的间距，避免了这种捷径。需要说明的是，该任务此前已作为去噪任务被提出。

许多现有研究认为，添加架构门控（乘法交互）可赋予模型 “数据依赖性”，并解决相关任务。然而，我们直观上认为这一解释并不充分，因为此类门控不沿序列轴交互，无法影响令牌之间的间距。具体而言，架构门控并非选择机制的实例（附录 A）。

表 1 证实，H3 和 Mamba 等门控架构仅能部分提升性能，而选择机制（将 S4 改为 S6）可轻松解决该任务，尤其是与这些更强大的架构结合时。

4.1.2 归纳头

归纳头是一个简单任务，源于机制可解释性视角，意外地能预测大型语言模型的上下文学习能力。它要求模型执行关联回忆和复制：例如，如果模型在序列中看到 “哈利・波特” 这样的二元组，那么下次在同一序列中看到 “哈利” 时，应能通过从历史中复制，预测出 “波特”。

数据集：我们在序列长度为 256、词汇表大小为 16 的归纳头任务上训练一个 2 层模型，与现有研究相当，但序列更长。此外，我们通过在测试时使用从\(2^{6}=64\)到\(2^{20}=1048576\)的一系列序列长度，评估模型的泛化和外推能力。

模型：遵循归纳头相关研究的做法，我们使用 2 层模型，这使得注意力机制能够通过机制性方式解决归纳头任务。我们测试了多注意力头（8 个头，采用多种位置编码）和状态空间模型变体。Mamba 的模型维度 D 为 64，其他模型的模型维度 D 为 128。

结果：表 2 显示，Mamba（更具体地说，其选择性状态空间模型层）能够完美解决该任务，因为它能选择性记忆相关令牌，同时忽略中间的所有其他信息。它能完美泛化至百万长度序列，比训练时的序列长度长 4000 倍，而其他方法的泛化能力不超过 2 倍。

在注意力模型的位置编码变体中，专为长度外推设计的 xPos 略优于其他变体；此外，由于内存限制，所有注意力模型的测试序列长度仅达到\(2^{14}=16384\)。在其他状态空间模型中，H3 和 Hyena 的表现相近，这与相关研究的发现相反。

模型	架构	层	准确率
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

4.2 语言建模

我们在标准自回归语言建模任务上，从预训练指标（困惑度）和零样本评估两方面，将 Mamba 架构与其他架构进行比较。我们设置的模型规模（深度和宽度）与 GPT3 规格一致。我们使用 Pile 数据集，并遵循相关研究描述的训练方案。所有训练细节详见附录 E.2。

4.2.1 缩放定律

基准模型包括标准 Transformer 架构（GPT3 架构），以及我们所知的最强 Transformer 方案（称为 Transformer++），该方案基于 PaLM 和 LLaMa 架构（例如，旋转嵌入、SwiGLU MLP、RMSNorm 替代 LayerNorm、无线性偏置、更高学习率）。我们还与其他近期亚二次架构进行了比较。所有模型细节详见附录 E.2。

图 4 展示了在标准 Chinchilla 协议下，参数规模从约 1.25 亿到约 13 亿的模型的缩放定律。Mamba 是首个与强 Transformer 方案（Transformer++）性能相当的无注意力模型，尤其是随着序列长度增加。（需要说明的是，由于缺乏高效实现，导致内存不足或计算需求不切实际，RWKV 和 RetNet 这两个此前强大的循环模型（也可解读为状态空间模型）在上下文长度为 8k 的完整结果缺失。）

4.2.2 下游评估

表 3 展示了 Mamba 在一系列流行的零样本评估任务上的表现。我们与这些规模下最著名的开源模型进行了比较，尤其是 Pythia 和 RWKV，它们与我们的模型使用相同的分词器、数据集和训练长度（3000 亿令牌）。（注：Mamba 和 Pythia 的训练上下文长度为 2048，而 RWKV 的训练上下文长度为 1024。）

表 3 显示，对于每个模型规模，Mamba 在所有评估结果中均为同类最佳，且总体上与两倍规模的基准模型相当。

4.3 DNA 建模

受大型语言模型成功的启发，近期研究开始探索将基础模型范式应用于基因组学。DNA 可被类比为语言，由有限词汇表的离散令牌序列组成。已知建模 DNA 需要长程依赖。我们在与近期长序列 DNA 模型研究相同的设置下，探索 Mamba 作为基础模型骨干网络的预训练和微调效果。具体而言，我们聚焦于模型规模和序列长度的缩放定律（图 5），以及一个需要长上下文的具有挑战性的下游合成分类任务（图 6）。

预训练方面，我们主要遵循标准的因果语言建模（下一个令牌预测）设置，训练和模型细节详见附录 E.2。数据集方面，我们主要遵循 HyenaDNA 的设置，使用 HG38 数据集进行预训练，该数据集包含人类基因组，训练集约有 45 亿个令牌（DNA 碱基对）。

4.3.1 缩放：模型规模

在本实验中，我们研究了不同模型骨干网络的基因组学基础模型的缩放特性（图 5 左）。

训练：为使基准模型更具优势，我们在短序列长度（1024）上训练；如 4.3.2 节所示，我们预计在更长序列长度下，Mamba 的优势将更加明显。我们固定全局批量大小为 1024，每批约 100 万个令牌。模型训练 10000 个梯度步骤，总令牌数为 100 亿。

结果：图 5（左）显示，Mamba 的预训练困惑度随模型规模平稳提升，且缩放效果优于 HyenaDNA 和 Transformer++。例如，在最大模型规模（约 4000 万参数）下，曲线显示 Mamba 仅需约 1/3 至 1/4 的参数，就能达到 Transformer++ 和 HyenaDNA 模型的性能。

4.3.2 缩放：上下文长度

在接下来的 DNA 实验中，我们研究模型随序列长度的缩放特性。我们仅比较 HyenaDNA 和 Mamba 模型，因为二次复杂度的注意力机制在长序列长度下的计算成本过高。我们在序列长度\(2^{10}=1024\)、\(2^{12}=4096\)、\(2^{14}=16384\)、\(2^{16}=65536\)、\(2^{18}=262144\)、\(2^{20}=1048576\)上预训练模型。我们固定模型规模为 6 层 × 宽度 128（约 130-140 万参数）。模型训练 20000 个梯度步骤，总令牌数约 3300 亿。更长的序列长度采用了类似相关研究的序列长度预热策略。

结果：图 5（右）显示，Mamba 能够利用更长的上下文，即使序列长度高达 100 万，其预训练困惑度随上下文长度增加而降低。相反，HyenaDNA 模型的性能随序列长度增加而下降。这与 3.5 节中关于选择机制特性的讨论一致。具体而言，线性时间不变模型无法选择性忽略信息；从卷积视角来看，极长的卷积核会聚合长序列中的所有信息，可能包含大量噪声。需要说明的是，尽管 HyenaDNA 声称性能随上下文长度增加而提升，但它们的结果未控制计算时间。

4.3.3 合成物种分类

我们在下游任务中评估模型：通过随机采样不同物种的连续 DNA 片段，对 5 个不同物种进行分类。该任务改编自 HyenaDNA，后者使用的物种包括人类、狐猴、小鼠、猪、河马。我们对任务进行了修改，使其更具挑战性：对五个类人猿物种（人类、黑猩猩、大猩猩、猩猩、倭黑猩猩）进行分类，这些物种的 DNA 相似度高达 99%。

图 6 显示，使用相同上下文长度的预训练模型进行微调后，随着序列长度从\(2^{10}=1024\)增加到\(2^{20}=1048576\)，模型的准确率变化情况。数值结果详见表 13。

4.4 音频建模与生成

对于音频波形模态，我们主要与 SaShiMi 架构和训练方案进行比较。该模型包括：

U-Net 骨干网络，包含两个池化阶段，每个阶段的池化因子为 P，模型维度 D 翻倍；

每个阶段交替使用 S4 和 MLP 块。

我们考虑用 Mamba 块替代 S4+MLP 块。实验细节详见附录 E.4。

4.4.1 长上下文自回归预训练

我们在 YouTubeMix（DeepSound 2017）上评估预训练质量（自回归下一个样本预测），这是一个标准钢琴音乐数据集，包含 4 小时的独奏钢琴音乐，采样率为 16000Hz。预训练细节主要遵循标准语言建模设置（4.2 节）。图 7 评估了在计算量固定的情况下，训练序列长度从\(2^{13}=8192\)增加到\(2^{20}≈10^6\)的效果。（数据整理存在一些特殊情况，可能导致缩放曲线出现拐点。例如，仅提供分钟级片段，因此最大序列长度实际上受限于\(60s·16000Hz=960000\)。）

Mamba 和 SaShiMi（S4+MLP）基准模型均随上下文长度增加而持续提升；Mamba 整体表现更优，且在更长长度下优势更明显。主要指标是每字节比特数（BPB），它是其他模态预训练标准负对数似然（NLL）损失的常数因子 log (2) 倍。

需要说明的一个重要细节：这是本文中唯一将参数化从实值切换为复值的实验。附录 E.4 提供了更多消融实验。

4.4.2 自回归语音生成

SC09 是一个语音生成基准数据集，包含 1 秒长的片段，采样率为 16000Hz，涵盖数字 “0” 到 “9”，且特征高度可变。我们主要遵循相关研究的自回归训练设置和生成方案。

表 4 显示了 Mamba-UNet 模型与多种基准模型的自动指标比较：WaveNet、SampleRNN、WaveGAN、DiffWave 和 SaShiMi。一个小型 Mamba 模型的性能优于最先进的（且规模大得多的）生成对抗网络（GAN）和扩散模型。一个与基准模型参数规模相当的大型 Mamba 模型，在保真度指标上进一步显著提升。

表 5 使用小型 Mamba 模型，探索了外部阶段和中心阶段不同架构的组合。结果显示，Mamba 在外部块中始终优于 S4+MLP，且在中心块中，Mamba > S4+MLP > MHA+MLP。

4.5 速度和内存基准测试

我们在图 8 中基准测试了状态空间模型扫描操作（状态扩展 N=16）的速度，以及 Mamba 的端到端推理吞吐量。我们的高效状态空间模型扫描在序列长度超过 2K 时，比已知最快的注意力实现（FlashAttention-2）更快，且比 PyTorch 中的标准扫描实现快 20-40 倍。Mamba 的推理吞吐量是同规模 Transformer 的 4-5 倍，因为无需键值缓存，可使用更大的批量大小。例如，一个未训练的 Mamba-6.9B 的推理吞吐量，高于规模小 5 倍的 Transformer-1.3B。细节详见附录 E.5，其中还包含内存消耗基准测试。

4.6 模型消融实验

我们对模型的各个组件进行了一系列详细的消融实验，聚焦于语言建模任务，使用约 3.5 亿参数的模型，遵循 Chinchilla 令牌计数设置（与图 4 相同）。

4.6.1 架构

表 6 探索了架构（块）及其内部状态空间模型层的影响（图 3）。我们发现：

在现有非选择性（线性时间不变）状态空间模型中，等效于全局卷积的模型，性能非常相似。

将现有研究中的复值 S4 变体替换为实值变体，对性能影响不大，这表明（至少对于语言建模）考虑到硬件效率，实值状态空间模型可能是更好的选择。

将这些模型替换为选择性状态空间模型（S6），性能显著提升，验证了 3 节的动机。

Mamba 架构的性能与 H3 架构相当（使用选择性层时似乎略好）。

我们还在附录 E.2.2 中探索了将 Mamba 块与其他块（如 MLP（传统架构）、MHA（混合注意力架构））交错排列的效果。

4.6.2 选择性状态空间模型

表 7 通过考虑选择性 Δ、B、C 参数的不同组合（算法 2），对选择性状态空间模型层进行了消融实验，结果表明，由于与循环神经网络门控的关联（定理 1），Δ 是最重要的参数。

表 8 考虑了状态空间模型的不同初始化方法，现有研究表明，在某些数据模态和设置中，初始化方法的影响很大。在语言建模任务中，我们发现，与更标准的复值参数化（S4D-Lin，第 1 行）相比，更简单的实值对角初始化（S4D-Real，第 3 行）表现更优。随机初始化也能取得良好效果，这与现有研究的发现一致。

表 9 和表 10 分别考虑了 Δ 和（B、C）投影维度的变化。从静态变为选择性带来的收益最大，而进一步增加维度，性能通常会适度提升，且参数数量仅小幅增加。

特别值得注意的是，当状态大小 N 增加时，选择性状态空间模型的性能显著提升，仅增加 1% 的参数，困惑度就降低了超过 1.0。这验证了我们在 3.1 节和 3.3 节中的核心动机。

5 讨论

我们讨论了相关工作、局限性和未来方向。

相关工作：附录 A 讨论了选择机制与类似概念的关系。附录 B 提供了状态空间模型和其他相关模型的扩展相关工作。

无免费午餐：连续 - 离散谱。结构化状态空间模型最初被定义为连续系统的离散化，对音频、视频等连续时间数据模态具有很强的归纳偏置。如 3.1 节和 3.5 节所述，选择机制克服了它们在文本和 DNA 等离散模态上的短板；但相反，这可能会影响它们在 linear time invariant 状态空间模型擅长的连续数据上的性能。我们在音频波形上的消融实验更详细地探讨了这一权衡。

下游适用性：基于 Transformer 的基础模型（尤其是大型语言模型）拥有丰富的生态系统，包括微调、适配、提示、上下文学习、指令调优、基于人类反馈的强化学习（RLHF）、量化等多种特性和交互模式。我们特别感兴趣的是，状态空间模型等 Transformer 替代方案是否具有类似的特性和适用性。

缩放：我们的实证评估仅限于小型模型规模，低于大多数强开源大型语言模型（如 Llama）以及 RWKV 和 RetNet 等其他循环模型的规模（这些模型已在 70 亿参数规模及以上进行了评估）。Mamba 在这些更大规模下是否仍具有优势，仍需进一步评估。此外，状态空间模型的缩放可能涉及本文未讨论的额外工程挑战和模型调整。

6 结论

我们将选择机制引入结构化状态空间模型，使它们能够执行上下文依赖推理，同时保持序列长度线性缩放。当整合到一个简单的无注意力架构中时，Mamba 在多个不同领域取得了最先进的结果，性能与强 Transformer 模型相当或更优。我们对选择性状态空间模型在不同领域构建基础模型的广泛应用感到兴奋，尤其是在基因组学、音频和视频等需要长上下文的新兴模态中。我们的结果表明，Mamba 是通用序列模型骨干网络的有力候选者。