Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

论文 · arXiv 2312.00752

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

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

发表于 2023-12-01cs.LGcs.AI arXiv PDF

用 XYZNOT 阅读（全文翻译与 AI 讲解）

三句话总结

🎯 Transformer 随序列长度呈二次方复杂度扩展，自回归推理需要巨大的 KV 缓存；而以往的亚二次模型（如 LTI 状态空间模型）在语言建模等依赖离散内容推理的任务上表现不佳。
🔧 使 SSM 参数依赖于输入（选择性）以动态过滤或传播信息，并结合硬件感知的融合并行结合扫描与重计算算法，避免在 GPU HBM 中具象化展开后的隐藏状态。
📈 Mamba-3B 在下游零样本任务上匹敌或超越两倍于其规模的 Transformer（例如平均准确率达 63.3%，而 Pythia-2.8B 为 59.1%、Pythia-6.9B 为 61.7%），同时提供高达 5 倍的推理吞吐量。

摘要

基础模型如今推动了深度学习领域绝大多数令人振奋的应用，它们几乎普遍基于 Transformer 架构及其核心的注意力模块。许多亚二次时间复杂度的架构（如线性注意力、门控卷积与循环模型，以及结构化状态空间模型 SSMs）已被开发出来，用以解决 Transformer 在长序列上的计算低效问题，但在语言等重要模态上它们的表现仍不及注意力机制。我们发现此类模型的一个关键缺陷在于无法进行基于内容的推理，并据此进行了多项改进。首先，简单地将 SSM 参数设为输入的函数即可解决它们在离散模态上的不足，使模型能够根据当前 token 沿着序列长度维度选择性地传递或遗忘信息。其次，尽管这一改动使得无法再利用高效卷积，我们在循环模式下设计了一种硬件感知的并行算法。我们将这些选择性 SSM 集成到一个简化的端到端神经网络架构中，该架构不包含注意力甚至无需 MLP 模块（Mamba）。Mamba 具有快速的推理速度（吞吐量比 Transformer 高 5 $\times$ ）以及沿序列长度的线性扩展能力，并且其性能在长达百万序列长度的真实数据上持续提升。作为通用的序列模型骨干网络，Mamba 在语言、音频和基因组学等多种模态上均取得了最先进的性能。在语言建模方面，我们的 Mamba-3B 模型在预训练和下游评测中均超越了同等规模的 Transformer，并匹敌两倍于其规模的 Transformer。

英文原文摘要

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5 $\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

论文解读

问题与动机

基础模型极度依赖 Transformer 架构和自注意力机制。尽管稠密注意力能够在序列上下文窗口中高效路由信息，但它带来了两个根本性缺陷：相对于序列长度 $L$ 的二次计算开销 $O (L^{2})$ ，以及必须存储历史键和值（KV cache），这导致每步推理需要线性时间，并引发显存瓶颈。

以往的亚二次复杂度方案——例如线性注意力、循环架构以及结构化状态空间模型（SSMs，如 S4）——在时间上呈线性或拟线性扩展。然而，在离散、信息密集型模态（如文本）上，结构化 SSM 的表现历来难以匹敌 Transformer。作者指出了其根本原因：以往的 SSM 是线性时不变系统（LTI）。由于其转移动态不会根据输入内容发生改变，因此无法选择性地压缩上下文、过滤无关噪声或在可变间距间调节记忆保留。

核心思想

作为上下文压缩的选择机制（Selection Mechanism）： 序列建模需要将历史信息压缩为状态。以往的 LTI SSM 使用固定的转移参数，迫使模型对所有时间步一视同仁。通过基于当前输入 token 对 SSM 矩阵进行参数化，模型获得了根据内容选择性过滤或记忆信息所需的感知能力。
攻克非卷积瓶颈： 使参数依赖于输入打破了线性循环与全局卷积之间的数学等价性。失去卷积后，朴素的循环计算将需要在 GPU 高带宽显存（HBM）中实例化一个形状为 $(B, L, D, N)$ 的膨胀隐状态。Mamba 通过将离散化、并行结合扫描（parallel associative scan）和输出投影融合到 GPU 片上高速 SRAM 中解决了这一问题，将 HBM 读写通信降至最低。
统一架构： Mamba 用单一的同构模块替代了交替堆叠的多头注意力（MHA）和 MLP 模块，该模块结合了 H3 的门控设计与 SwiGLU 风格的门控 MLP。

方法

连续状态空间模型通过隐变量状态 $h (t) \in R^{N}$ 将输入 $x (t) \in R$ 映射到 $y (t) \in R$ ： $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

利用步长为 $Δ$ 的零阶保持（ZOH）离散化，该连续系统转化为离散循环参数： $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

选择机制（S6）

Mamba 不再将 $(Δ, B, C)$ 随时间保持恒定，而是让它们依赖于输入：

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ ，其中 $τ_{Δ} = softplus$ 且 $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

这推广了经典的 RNN 门控机制：当 $N = 1, A = - 1, B = 1$ 时，该循环直接简化为门控更新 $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ ，其中 $g_{t} = σ (Linear (x_{t}))$ 。

硬件感知状态膨胀（Hardware-Aware State Expansion）

为了在不向 HBM 写入 $(B, L, D, N)$ 状态的情况下保持高吞吐量：

参数 $(Δ, A, B, C)$ 直接从 HBM 加载到片上 SRAM 中。
离散化为 $(\overline{A}, \overline{B})$ 以及高效的并行结合扫描均完全在 SRAM 内部计算完成。
隐状态在写回 HBM 之前与 $C$ 收缩收束为形状 $(B, L, D)$ 。
在反向传播过程中，中间状态在 SRAM 中重新计算，而不是保存在 HBM 中。

结果

合成任务外推： 在 Induction Heads 任务中，基于序列长度 256 训练的 Mamba 能够完美外推（100% 准确率）至长度 $2^{20} = 1, 048, 576$ （4000 倍序列长度），而注意力基线在超过 $2^{14}$ 后便失效或发生 OOM。在 Selective Copying 任务中，Mamba 达到 99.8% 的准确率，而 H3 为 57.0%，S4 为 18.3%。
语言建模： 在 Pile 数据集上按照 Chinchilla 定律扩展至 1.3B 参数评测，Mamba 匹敌或超越了 Transformer++（包含旋转位置编码、SwiGLU、RMSNorm）。在下游零样本任务（如 LAMBADA、PIQA、HellaSwag、WinoGrande、ARC）中，Mamba-2.8B 达到 63.3% 的平均准确率，优于 Pythia-2.8B（59.1%）和 Pythia-6.9B（61.7%），并能匹敌或超越 6B–7B 参数规模的模型。
基因组学与音频： 在序列长度为 1M（ $2^{20}$ ）的大型类人猿 DNA 分类任务中，Mamba 取得了 71.67%（1.4M 参数）和 81.31%（7M 参数）的准确率，而 HyenaDNA 仅为 54.87%。在 SC09 语音生成任务上，Mamba 将 FID 从 1.99（SaShiMi）分别降低至 0.94（6.1M 参数）和 0.67（24.3M 参数）。
效率： 融合扫描算子的执行速度比标准 PyTorch 扫描实现快 20–40 倍，在 32K 上下文长度下比 FlashAttention-2 快多达 7 倍。推理吞吐量最高可达同级别 Transformer 的 5 倍。

局限与待探索问题

连续信号上的表现： 尽管输入相关的选择机制提升了离散推理能力（文本、DNA），但在原始音频波形预训练的消融实验中表明，纯线性时不变（LTI）模型（如标准 S4）在均匀采样的连续数据上表现更佳。
模型规模： 评测模型的规模上限约为 3B 参数和 300B 预训练 token；在 7B–70B+ 及更大规模下的经验表现仍需进一步探索。
下游生态： Transformer 拥有围绕提示词工程（prompt engineering）、上下文内操作、KV-cache 量化和适配器微调（adapter tuning）的成熟工具链。SSM 如何更好地适配这些交互模式仍是一个活跃的研究方向。

关键术语

Structured State Space Models (SSMs) (结构化状态空间模型): 一类由连续线性动力系统衍生而来的深度序列模型，通过隐变量状态对一维输入进行映射。
Linear Time Invariance (LTI) (线性时不变性): 一种数学性质，指系统转移参数在所有时间步保持静态不变，这使得模型可以通过快速全局卷积表示，但限制了基于内容的推理能力。
Selection Mechanism (选择机制): 将 SSM 转移动态 (Δ, B, C) 参数化为当前输入 token 的函数，从而允许模型沿着序列长度动态传播、重置或忽略上下文。
S6: 通过结合并行扫描而非卷积计算的选择性结构化状态空间模型（selective structured state space models）的简称。
Parallel Associative Scan (并行结合扫描): 一种高效的前缀和（prefix computation）算法，可在加速器硬件上跨序列长度并行计算线性循环。
Hardware-Aware State Expansion (硬件感知状态膨胀): 一种 GPU 算子融合与重计算策略，在高速片上 SRAM 中执行离散化和并行扫描，避免将庞大的中间状态频繁往返读写至高带宽显存（HBM）。
Induction Heads (感应头): 一种合成关联召回机制，用于评估模型是否能够在出现前序上下文 token 后复现先前观察到的对应 token。
Mamba Block (Mamba 模块): 一种简化的同构序列模型模块，集成了输入投影、一维卷积、选择性 SSM 和乘性门控，无需独立的注意力或 MLP 层。

常见问题

为什么选择性 SSM 不能使用快速 FFT 卷积进行训练？

快速全局 FFT 卷积要求系统动态特性（矩阵 A、B 和 C）在时间维度上保持恒定（线性时不变性）。由于 Mamba 将 B、C 和 Δ 变成了每个时间步输入 token 的函数，卷积核在每个位置都会发生变化，打破了卷积等价性，因此必须改用循环并行扫描。

Mamba 在训练过程中如何避免循环状态膨胀带来的显存爆炸？

状态扩展维度为 N 的朴素循环需要在 GPU HBM 中实例化形状为 (B, L, D, N) 的中间状态。Mamba 采用算子融合技术，仅将低维输入加载至 SRAM 中，在高速 SRAM 内部计算离散化和并行扫描，仅将收束后的 (B, L, D) 输出写回 HBM，并在反向传播时重新计算中间激活值。

Mamba 的 Δ 参数与循环门控机制之间有什么关系？

步长参数 Δ 起到了广义门控机制的作用。在零阶保持（ZOH）且 N=1、A=-1 的离散化设定下，选择性循环在数学上直接退化为标准 RNN 门控：较大的 Δ 会使模型关注当前输入并重置状态，而较小的 Δ 则会持久保留历史状态并忽略当前输入。

Mamba 在生成阶段使用 KV 缓存吗？

不使用。由于 Mamba 在生成过程中作为自回归状态空间循环运行，它仅需维护固定大小的隐藏状态（每层维度为 D × N），而无需缓存所有历史 token 的键和值，从而实现每步常数时间 O(1) 的推理。

本页解读由 AI 根据论文生成，可能有错误；细节请以原文为准。

用 XYZNOT 阅读（全文翻译与 AI 讲解）