精读笔记 #12:Mamba: Linear-Time Sequence Modeling with Selective State Spaces
论文基本信息
| 项目 | 内容 |
|---|---|
| 标题 | Mamba: Linear-Time Sequence Modeling with Selective State Spaces |
| 作者 | Albert Gu(CMU)、Tri Dao(Princeton) |
| 机构 | Carnegie Mellon University / Princeton University |
| 发表时间 | 2023年12月1日(arXiv 2312.00752) |
| 修订版本 | 2024年5月31日(v2) |
| 核心主题 | 用"选择性状态空间模型"替代 Transformer 注意力机制,实现线性时间复杂度序列建模 |
阅读地图
本文解决的核心问题:Transformer 的注意力机制计算量随序列长度平方增长(O(n²)),处理长序列时代价极高。 Mamba 提出了一种替代方案,让模型在处理序列时只需 O(n) 的计算量,推理时还能以常数大小的内存运行。
阅读建议:
1. 先看 Abstract,理解"三件事"(选择性机制、硬件算法、简化架构)
2. 再看 Introduction,理解为什么 Transformer 有问题、为什么之前的 SSM 也不够好
3. 重点看 Section 3,这是核心方法:选择性机制 → 硬件感知算法 → Mamba 模块
4. 最后看实验,感受 Mamba 与 Transformer 的对比
核心类比(贯穿全文):
把序列建模想象成"边读书边做笔记"。
- Transformer(注意力):读完整本书,随时回头翻看任何一页。效果好,但随着书变厚,翻页成本急剧上升(O(n²))。
- 早期 SSM(固定状态):只有一个固定大小的"笔记本",每读一个字就更新笔记,但更新规则是死的,不管读到什么内容,更新方式都一样。
- Mamba(选择性 SSM):也是固定大小的笔记本,但更新规则会根据读到的内容动态调整——读到重要内容就认真记,读到无关内容就跳过。笔记本大小不变(常数内存),但记的质量大幅提升。
Abstract:摘要精译
原文
"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."
翻译:
驱动当前大多数深度学习应用的基础模型,几乎无一例外地建立在 Transformer 架构及其核心注意力模块之上。研究界已开发了许多次二次方时间复杂度的架构,包括线性注意力、门控卷积与循环模型、结构化状态空间模型(SSM)等,以解决 Transformer 在长序列上的计算效率问题。然而,这些替代方案在语言等重要模态上的表现始终不如注意力机制。我们识别出这类模型的关键弱点在于无法执行基于内容的推理,并据此做出了若干改进。
新手讲解:
这段话在说:很多人试图造一个"便宜版 Transformer",但效果都差一截。作者发现关键问题在于这些替代方案是"内容盲的"——它们处理每个输入时用的方法是固定的,不管这个词重不重要都一视同仁。这就是 Mamba 要修复的核心缺陷。
"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."
翻译:
首先,只需让 SSM 的参数成为输入的函数,就能解决离散模态(如文本)上的弱点——这使模型能够根据当前词元,有选择地沿序列长度方向传播或遗忘信息。
新手讲解:
这是全文最核心的一句话。"参数成为输入的函数"意思是:模型更新笔记的方式,会根据当前读到的内容动态改变。 读到"重要会议时间:下午3点",就把时间记下来;读到"嗯……呃……然后呢……"这类废话,就选择忽略,不更新笔记。这就是"选择性(Selective)"的含义。
"Mamba enjoys fast inference (5× higher throughput than Transformers) and linear scaling in sequence length, and its performance improves monotonically with context length. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size."
翻译:
Mamba 具有快速推理能力(吞吐量比 Transformer 高 5 倍),序列长度呈线性扩展,且性能随上下文长度单调提升。在语言建模任务中,我们的 Mamba-3B 模型超越了相同规模的 Transformer,并与两倍规模的 Transformer 相当。
新手讲解:
结果很漂亮:
- 速度:推理快 5 倍(Transformer 需要存储所有历史的"键值缓存",越长越慢;Mamba 只需固定大小的状态,速度恒定)
- 规模效率:3B 参数的 Mamba 能打赢 6B 参数的 Transformer,相当于用一半的模型大小达到同等效果
- 长度外推:Transformer 在超出训练长度时性能下降,Mamba 则随上下文越长越好
Introduction:引言核心段落精译
段落一:基础模型与 Transformer 的统治地位
原文:
"Foundation models (FMs), or large models pretrained on massive data then adapted for downstream tasks, have emerged as an effective paradigm in modern machine learning. The backbone of these FMs are often sequence models, operating on arbitrary sequences of inputs from a wide variety of domains such as language, images, speech, audio, time series, and genomics."
翻译:
基础模型(FM)——即在海量数据上预训练后适配下游任务的大型模型——已成为现代机器学习的有效范式。这些基础模型的骨干通常是序列模型,处理来自语言、图像、语音、音频、时间序列、基因组学等广泛领域的任意输入序列。
新手讲解:
"序列模型"就是处理"有顺序的数据"的模型。文字是序列(一个词接一个词),音频是序列(一个采样点接一个采样点),DNA 也是序列(一个碱基接一个碱基)。Mamba 的目标是成为一个通用的序列处理骨干,不局限于某一领域。
段落二:Transformer 的根本缺陷
原文:
"The efficacy of self-attention is attributed to its ability to route information densely within a context window, allowing it to model complex data. However, this property brings fundamental drawbacks: an inability to model anything outside of a finite window, and quadratic scaling with respect to the window length."
翻译:
自注意力机制的高效性源于其在上下文窗口内密集路由信息的能力,这使其能够建模复杂数据。然而,这一特性带来了根本性的缺陷:无法建模有限窗口之外的内容,以及随窗口长度呈二次方增长的计算开销。
新手讲解:
这里的"二次方增长"(O(n²))非常关键:
- 序列长度 1000:计算量约 100 万次
- 序列长度 10000:计算量约 1 亿次(增加了 100 倍!)
- 序列长度 100000:计算量约 100 亿次(又增加了 100 倍!)
注意力机制需要让序列中每个位置都"关注"其他所有位置——这就是 O(n²) 的来源。对于长文档、长基因序列、长音频,这个代价是灾难性的。
段落三:早期 SSM 的不足与本文贡献
原文:
"Recently, structured state space sequence models (SSMs) have emerged as a promising class of architectures for sequence modeling. We identify that a key weakness of such models is their inability to perform content-based reasoning. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. Finally, we wrap the selective SSM into a simplified end-to-end neural network architecture without attention or MLP blocks."
翻译:
近年来,结构化状态空间序列模型(SSM)作为一类有前景的序列建模架构涌现出来。我们识别出这类模型的关键弱点在于无法执行基于内容的推理。为此,本文做出三项贡献:第一,让 SSM 参数成为输入的函数,解决其在离散模态上的弱点;第二,尽管此改变使高效卷积不再可用,我们设计了一种在递归模式下运行的硬件感知并行算法;第三,我们将选择性 SSM 嵌入一个简化的端到端神经网络架构,该架构不含注意力层或 MLP 块。
新手讲解:
三个贡献对应三个问题:
1. 早期 SSM 是"内容盲"的 → 解决方案:让参数随输入变化(选择性机制)
2. 加了选择性之后,原来的高效计算方法失效了 → 解决方案:重新设计硬件感知算法
3. 需要一个完整可用的网络结构 → 解决方案:设计 Mamba 模块
这三步是递进关系,后面章节会逐一展开。
Section 2:状态空间模型背景(State Space Models)
2.1 什么是 SSM?
术语解释:
- SSM(State Space Model,状态空间模型):一类用"隐藏状态"来压缩历史信息的模型。最早来自控制论,用于描述动态系统(如火箭轨迹、电路信号)。
- State(状态):模型的"工作记忆",一个固定大小的向量,存储了对过去信息的压缩摘要。
- Recurrence(递归):每一步用上一步的状态加上当前输入,计算新状态。
核心公式(连续形式):
h'(t) = A·h(t) + B·x(t) ← 状态更新:新状态 = A×旧状态 + B×当前输入
y(t) = C·h(t) ← 输出:从状态中读出结果
离散形式(实际用于计算机):
h_t = Ā·h_{t-1} + B̄·x_t ← 每一步更新状态(递归)
y_t = C·h_t ← 每一步产生输出
新手讲解:
想象你在读一本书,h_t 就是你的"读书笔记"(固定大小的向量)。每读一个新词 x_t:
- 用矩阵 Ā 决定"保留多少旧笔记"
- 用矩阵 B̄ 决定"从新词中吸收多少信息"
- 用矩阵 C 决定"从笔记中提取什么作为答案"
关键:笔记本大小固定(N维),不管序列多长,推理时的内存都是常数! 这就是 SSM 在推理阶段比 Transformer 省内存的原因。
2.2 离散化(Discretization)
原文概念:
"Discretization converts continuous parameters (Δ, A, B) to discrete ones (Ā, B̄) through fixed formulas. The zero-order hold rule: Ā = exp(ΔA), B̄ = (ΔA)⁻¹(exp(ΔA) − I)·ΔB."
翻译:
离散化通过固定的转换规则将连续参数(Δ、A、B)转换为离散参数(Ā、B̄)。零阶保持规则:Ā = exp(ΔA),B̄ = (ΔA)⁻¹(exp(ΔA) − I)·ΔB。
新手讲解:
物理世界是连续的(时间是连续流动的),但计算机处理的是离散数据(一个词接一个词)。Δ(Delta)是"时间步长",控制将连续系统"采样"的频率。可以理解为:Δ 越大,相邻两个词之间的"虚拟时间间隔"越长,当前词对状态的影响也越大。
2.3 卷积模式 vs 递归模式
SSM 可以用两种等价方式计算:
| 模式 | 公式 | 优势 | 劣势 |
|---|---|---|---|
| 递归模式(Recurrent) | 逐步更新 h_t | 推理高效(常数内存) | 训练时无法并行 |
| 卷积模式(Convolutional) | y = x * K̄(K̄ 为卷积核) | 训练时可完全并行 | 推理需存整个历史 |
新手讲解:
训练时希望并行(用卷积模式,像 CNN 一样高效);推理时希望省内存(用递归模式,每步只需一个固定大小的状态向量)。经典 SSM(如 S4)两者都能用,因为它的参数是固定的(时不变的)。但后面会看到,Mamba 的"选择性"机制让参数随输入变化,卷积模式就不再适用——这正是需要设计新算法的原因。
2.4 线性时不变性(LTI,Linear Time Invariance)
原文:
"A critical property: 'the model's dynamics are constant through time.' Parameters remain fixed across all timesteps, enabling efficient convolution-based computation but limiting expressiveness."
翻译:
一个关键特性:模型的动态在时间上是恒定的。参数在所有时间步保持固定,这使得基于卷积的高效计算成为可能,但同时限制了模型的表达能力。
术语解释:
- LTI(Linear Time Invariance,线性时不变性):系统的行为规则不随时间改变。无论你在第 1 步还是第 1000 步,处理输入的方式完全相同。
新手讲解:
想象用同一把固定规格的"筛子"来过滤每一个输入词——无论是在第一句话里,还是最后一句话里,筛子的孔径都一样大,无法根据上下文调整。
这就是早期 SSM 的根本局限:处理"天气"这个词和处理"啊"这个语气词,用的是完全相同的更新规则。 它不知道哪个词重要,哪个词可以忽略。
Section 3:选择性状态空间模型(Selective SSMs)——核心方法
3.1 动机:选择性作为压缩手段
原文:
"A fundamental problem of sequence modeling is compressing context into a smaller state. The efficiency vs. effectiveness tradeoff is characterized by how well models compress state: efficient models need small state, effective models need state containing all necessary information."
翻译:
序列建模的根本问题是将上下文压缩进一个更小的状态。效率与效果的权衡取决于模型压缩状态的质量:高效的模型需要小状态,有效的模型需要状态包含所有必要信息。
新手讲解:
这里作者用一个很精辟的框架来看序列建模:
| 架构 | 状态大小 | 效率 | 效果 |
|---|---|---|---|
| Transformer | 随序列增长(KV 缓存) | 差(O(n²)) | 好(什么都记得) |
| 固定 RNN/SSM | 固定大小 | 好(O(n)) | 差(记不住重要内容) |
| Mamba | 固定大小,但智能压缩 | 好(O(n)) | 好(选择性记忆) |
关键洞察:固定大小的状态不意味着效果差,关键在于能否智能地选择记什么、忘什么。
两个动机任务:
任务一:选择性复制(Selective Copying)
普通的复制任务:序列 [A, B, C, D, ...] → 要求模型记住并输出 [A, B, C, D, ...]。
选择性复制任务:序列 [A, 噪, B, 噪, 噪, C, ...] → 要求模型只记住有用的 token,输出 [A, B, C, ...],过滤掉噪声。
原文:
"The Selective Copying task modifies the popular Copying task by varying the position of the tokens to memorize. It requires content-aware reasoning to be able to memorize the relevant tokens and filter out the irrelevant ones."
翻译:
选择性复制任务在传统复制任务基础上,通过改变待记忆词元的位置来增加难度。该任务要求基于内容的推理,以便记住相关词元并过滤掉不相关的词元。
新手讲解:
这就像在听一段演讲,演讲者中间会说很多"嗯……呃……那个……"。一个好的听众会自动过滤这些废话,只记住核心内容。LTI(时不变)模型做不到这一点,因为它对"嗯"和"重要结论"的处理方式完全相同。
任务二:归纳头(Induction Heads)
原文:
"The Induction Heads task is a well-known mechanism hypothesized to explain the majority of in-context learning abilities of LLMs. These tasks reveal the failure mode of LTI models. From the recurrent view, their constant dynamics cannot let them select the correct information from their context, or affect the hidden state in input-dependent ways."
翻译:
归纳头任务是一种公认的机制,被认为是解释大语言模型大部分上下文学习能力的关键。这些任务揭示了 LTI 模型的失败模式:从递归视角看,恒定的动态无法让模型从上下文中选出正确信息,也无法以依赖输入的方式影响隐藏状态。
新手讲解:
"归纳头"是指这样的推理能力:看到 [...A, B, ..., A],就能预测接下来是 B(因为之前 A 后面跟着 B)。这需要模型能"关联"前文中的模式。固定规则的 SSM 做不到这种"找特定上下文模式"的能力,因为它的更新规则不因内容而异。
3.2 选择性机制:让参数随输入变化
这是全文最核心的创新。
原文:
"One method of incorporating a selection mechanism into models is by letting their parameters that affect interactions along the sequence be input-dependent. The main difference is simply making several parameters Δ, B, C functions of the input."
翻译:
在模型中引入选择性机制的一种方法是:让影响序列交互的参数成为输入的函数。 核心区别仅仅是将参数 Δ、B、C 变为输入的函数。
具体实现(算法 2):
s_B(x) = Linear_N(x) ← B 由输入 x 决定
s_C(x) = Linear_N(x) ← C 由输入 x 决定
s_Δ(x) = Softplus(Linear_1(x)) ← Δ 由输入 x 决定(用 softplus 保证正值)
对比经典 SSM 与 Mamba:
| 参数 | 经典 SSM(如 S4) | Mamba(S6) |
|---|---|---|
| A | 固定(训练后不变) | 固定(通常是对角矩阵) |
| B | 固定 | 随每个输入词变化 |
| C | 固定 | 随每个输入词变化 |
| Δ | 固定 | 随每个输入词变化 |
新手讲解(类比):
- 经典 SSM:用同一个公式 笔记 = 0.9×旧笔记 + 0.1×新词,不管新词是什么。
- Mamba:根据当前词动态决定:
- 读到"重要"→ 笔记 = 0.3×旧笔记 + 0.7×新词(多吸收)
- 读到"的"(助词)→ 笔记 = 0.99×旧笔记 + 0.01×新词(几乎忽略)
- 读到"但是"(转折)→ 笔记 = 0.1×旧笔记 + 0.9×新词(几乎重置,开始关注新方向)
这个"动态调整比例"的能力,就是选择性的精髓。
为什么这个改变会带来计算难题?
经典 SSM(参数固定):可以把递归展开成卷积 y = x * K̄,然后用 FFT(快速傅里叶变换)一次性并行处理整个序列,复杂度 O(n log n)。
Mamba(参数随输入变化):每一步的矩阵 Ā 和 B̄ 都不同,无法展开成统一的卷积核,必须逐步递归:
h_1 = Ā_1·h_0 + B̄_1·x_1
h_2 = Ā_2·h_1 + B̄_2·x_2 ← Ā_2 ≠ Ā_1
...
这看起来是串行的(每步依赖上一步),无法直接并行——这就是为什么需要特殊的硬件感知算法。
3.3 硬件感知的并行扫描算法(Hardware-Aware Algorithm)
术语解释:
- HBM(High Bandwidth Memory,高带宽内存):GPU 上的"大容量主存",读写较慢,容量大(数十 GB)。
- SRAM(Static RAM,静态随机存储器):GPU 芯片上的"寄存器/缓存",读写极快,但容量极小(几十 MB)。
- Kernel Fusion(核融合):将多个 GPU 计算步骤合并成一个,减少在 HBM 和 SRAM 之间反复搬运数据的开销。
- Parallel Scan(并行扫描):一种让看似串行的递归在 GPU 上并行执行的算法技巧。
- Linear-time(线性时间):计算量随序列长度线性增长 O(n),而非 O(n²)。
为什么朴素实现会很慢?
原文:
"The selection mechanism is designed to overcome the limitations of LTI models; at the same time, we therefore need to revisit the computation problem of SSMs. When parameters vary over time, convolution equivalence breaks; recurrent mode materializes the expanded state."
翻译:
选择性机制旨在克服 LTI 模型的局限性;与此同时,我们因此需要重新审视 SSM 的计算问题。当参数随时间变化时,卷积等价性不再成立;递归模式需要实体化展开的状态。
新手讲解:
状态 h 的维度是 (B, L, D, N)——批大小 × 序列长度 × 模型维度 × 状态维度。对于长序列,这个张量非常大。如果朴素地在 GPU 主存(HBM)上逐步运算,每一步都要在 HBM 和计算单元之间搬运数据,内存带宽成为瓶颈,速度极慢。
三项核心技术:
技术一:核融合(Kernel Fusion)
原文:
"The main idea is to leverage properties of modern accelerators (GPUs) to materialize the state only in more efficient levels of the memory hierarchy. Load SSM parameters directly from slow HBM to fast SRAM, perform discretization and recurrence in SRAM, then write final outputs back to HBM."
翻译:
核心思路是利用现代加速器(GPU)的特性,只在内存层次结构中更高效的层级(SRAM)上实体化状态。将 SSM 参数直接从慢速 HBM 加载到快速 SRAM,在 SRAM 中执行离散化和递归,然后仅将最终输出写回 HBM。
新手讲解(工厂类比):
- HBM 是"远处的大仓库",存货多但取货慢
- SRAM 是"车间旁的小货架",存货少但取货快
- 朴素做法:每处理一步,去大仓库取材料、做完放回去,再去大仓库取下一步材料……来回跑很耗时
- 核融合做法:一次性从大仓库搬来足够多的原料,在小货架上把所有步骤一口气做完,再一次性运回大仓库。这样大幅减少了"跑腿"次数。
技术二:并行扫描(Parallel Scan)
原文:
"Despite recurrence being inherently sequential, a work-efficient parallel scan algorithm enables parallelization—allowing computation in O(log L) parallel steps while maintaining linear total work."
翻译:
尽管递归在本质上是串行的,但一种工作效率优化的并行扫描算法能够实现并行化——允许在 O(log L) 个并行步骤内完成计算,同时保持线性的总工作量。
新手讲解:
递归 h_t = f(h_{t-1}, x_t) 看起来非得一步一步做,但如果操作 f 满足结合律(associativity),就可以用类似"树形求和"的方法并行。
比如求 8 个数的前缀和 [a1, a1+a2, a1+a2+a3, ...]:
- 朴素做法:7 步,串行。
- 并行树形做法:第一轮同时算 (a1+a2), (a3+a4), (a5+a6), (a7+a8);第二轮同时算 (a1+a2+a3+a4), (a5+a6+a7+a8)……只需 3 步(log₂8)。
SSM 的递归恰好满足这个结合律条件,所以可以用并行扫描在 GPU 上高效执行。这就是把"看似串行"的递归变成"实际并行"的关键技巧。
术语补充:
- Scan(扫描):计算机科学中,"扫描"指计算前缀结果的操作(类似前缀和)。
- Parallel Scan(并行扫描):利用结合律将扫描操作并行化的算法,GPU 硬件天然支持。
技术三:重新计算(Recomputation)
原文:
"Intermediate states are not stored but recomputed in the backward pass when inputs are loaded from HBM to SRAM."
翻译:
中间状态不被存储,而是在反向传播时,当输入从 HBM 加载到 SRAM 时重新计算。
新手讲解:
训练神经网络需要"反向传播"来计算梯度,这通常需要保存前向传播的中间结果(如 h_1, h_2, h_3...)。对于长序列,这些中间状态会占用大量 GPU 内存(O(BLDN))。
Mamba 的解决方案:不存中间状态,反向传播时重算一遍。空间换时间——虽然多了一次前向计算,但内存开销大幅降低,整体仍然比 Transformer(需要存 KV 缓存)更省内存。
这个思路和 FlashAttention 的"不存注意力矩阵,重算"如出一辙。
综合效果:
原文:
"Fused selective scan layer has same memory requirements as optimized transformer with FlashAttention while achieving linear sequence scaling."
翻译:
融合后的选择性扫描层与使用 FlashAttention 优化的 Transformer 具有相同的内存需求,同时实现序列长度的线性扩展。
新手讲解:
Mamba 扫描操作比标准 PyTorch 实现快 40 倍,推理吞吐量比同等规模 Transformer 高 4-5 倍。这不是理论上的,是在实际 GPU(A100)上测量的结果。
3.4 Mamba 架构模块(Mamba Block)
原文:
"We simplify prior deep sequence model architectures by combining the design of prior SSM architectures with the MLP block of Transformers into a single block. This architecture involves expanding the model dimension D by a controllable expansion factor E. We repeat this block, interleaved with standard normalization and residual connections, to form the Mamba architecture."
翻译:
我们通过将先前 SSM 架构的设计与 Transformer 的 MLP 块合并为单一模块来简化之前的深度序列模型架构。该架构将模型维度 D 扩展一个可控的扩展因子 E(固定为 2)。我们重复堆叠此模块,与标准归一化层和残差连接交替使用,构成完整的 Mamba 架构。
Mamba 模块结构图(文字版):
输入 x (维度 D)
↓ ↓
Linear(2D) Linear(2D)
↓ ↓
SSM SiLU激活
(选择性扫描) ↓
↓ × ← 逐元素相乘(门控)
↓___________↓
↓
Linear(D)
↓
输出 y (维度 D)
参数分布:
- 线性投影层占模型参数的绝大部分:约 3ED² 个参数
- SSM 本身的参数(Δ、B、C、A)相对较少
- 扩展因子 E = 2(模型维度先扩大为 2D,最后再压回 D)
激活函数:
"We use the SiLU/Swish activation function, motivated so that the Gated MLP becomes the popular SwiGLU variant."
使用 SiLU(Sigmoid Linear Unit)激活,使门控 MLP 变为 SwiGLU 的变体——这也是 LLaMA 等现代 Transformer 使用的激活函数。
新手讲解:
Mamba 模块的设计非常简洁:
1. 把输入扩展到 2 倍维度(类似 Transformer 的 FFN 层先扩后缩)
2. 一路走 SSM(选择性状态空间,负责序列建模)
3. 另一路走门控(SiLU,负责非线性变换)
4. 两路相乘(门控机制,决定 SSM 的输出有多少被"放行")
5. 压缩回原始维度输出
整个模块没有注意力、没有独立的 MLP 层,用一个统一的模块同时完成了序列建模和特征变换。
3.5 选择性机制的性质分析
性质一:与 RNN 门控的等价性
原文(定理 1):
"When N=1, A=−1, B=1 with selective Δ and softplus, the SSM recurrence becomes: g_t = σ(Linear(x_t)); h_t = (1−g_t)h_{t-1} + g_t x_t — precisely standard RNN gating."
翻译:
当 N=1、A=−1、B=1,并使用选择性 Δ 和 softplus 激活时,SSM 递归变为:g_t = σ(Linear(x_t));h_t = (1−g_t)h_{t-1} + g_t x_t,这正是标准 RNN 的门控机制。
新手讲解:
这是一个非常优美的理论结果。它说明:Mamba 的选择性机制是 LSTM/GRU 门控机制的一种自然推广。
- LSTM 的遗忘门:h_t = (1−f_t)·h_{t-1} + f_t·x_t(f_t 由输入决定)
- Mamba 的 Δ:本质上扮演了同样的角色,只是推导更有原则性(来自连续时间系统的离散化),而不是启发式地"拍脑袋"设计的。
论文的说法是:"SSM 的离散化是启发式门控机制的原则性基础。"
性质二:三个参数的直觉含义
Δ(Delta,时间步长)的作用:
原文:
"Δ controls the balance between focusing on or ignoring the current input x_t. Mechanically, large Δ resets state focusing on current input; small Δ persists state ignoring input."
翻译:
Δ 控制着对当前输入 x_t 的关注程度与忽略程度之间的平衡。在机制层面,大 Δ 会重置状态以聚焦于当前输入;小 Δ 则维持状态而忽略当前输入。
新手讲解:
- Δ 大(如 Δ → ∞):Ā = exp(ΔA) → 0(旧状态衰减为零),B̄ → 大(充分吸收新输入)。效果:几乎完全重置,只看当前词。适合"但是……"这类语义转折处。
- Δ 小(如 Δ → 0):Ā ≈ 1(完全保留旧状态),B̄ ≈ 0(几乎不吸收新输入)。效果:完全保留历史,忽略当前词。适合处理"的、了、啊"等虚词。
B(输入到状态的门)的作用:
原文:
"Making B selective allows finer-grained control over whether input enters state. B selectivity lets models modulate recurrent dynamics based on content (input)."
翻译:
使 B 具有选择性,可以对输入是否进入状态实现更细粒度的控制。B 的选择性让模型能够基于内容(输入)调节递归动态。
新手讲解:
B 决定"新词有多少信息被写入笔记"。B 随输入变化意味着:这个词本身的内容决定了它能在笔记里留下多少痕迹。重要的词留下深痕,无关的词几乎不留痕。
C(状态到输出的门)的作用:
原文:
"C selectivity controls dynamics based on context (hidden states)."
翻译:
C 的选择性根据上下文(隐藏状态)控制动态。
新手讲解:
C 决定"从笔记中读出什么来回答当前问题"。C 随输入变化意味着:当前问题决定了去笔记本的哪个部分找答案。同样的笔记本,针对不同问题可以提取出不同的信息。
性质三:三种选择性带来的能力
原文:
"Selectivity allows filtering out irrelevant noise tokens that may occur between inputs of interest. Selective models can simply reset their state at any time to remove extraneous history. Performance improves monotonically with context length."
三种核心能力:
-
可变间距(Variable Spacing):过滤掉有用 token 之间的无关噪声 token(如口头禅、填充词)。
类比:听演讲时自动过滤"嗯……呃……那个……" -
过滤上下文(Filtering Context):随时重置状态,去除无关历史信息。
类比:读完一本书开始读下一本时,你会先清除对上一本书情节的记忆,以免混淆。 -
边界重置(Boundary Resetting):防止拼接在一起的独立序列之间发生信息污染。
类比:批量处理多篇不相关文章时,每篇文章的信息不应影响下一篇的处理。
Section 4:关键实验结果
4.1 合成任务:验证选择性的必要性
结果:
- 选择性复制任务:LTI 模型(H3)部分改善(97% 准确率),完整选择性机制(S6,即 Mamba 的核心)轻松解决(97-99%)
- 归纳头任务:Mamba 能泛化到比训练长度长 4000 倍的序列(超过 100 万 token),而其他方法最多只能泛化到 2 倍长度
新手讲解:
归纳头任务的外推能力是一个惊人结果。模型在长度 256 的序列上训练,却能在长度超过 100 万的序列上正确执行任务。这证明了 Mamba 学到的不是"背答案",而是真正理解了序列中的关联模式,且这个能力不受序列长度限制。
4.2 语言建模:与 Transformer 的直接对比
设置:
在 Pile 数据集上,从 125M 到 1.3B 参数进行 scaling law(扩展规律)实验,与 Transformer++(含 FlashAttention、SwiGLU 等现代改进的强 Transformer 基线)对比。
关键结论:
原文:
"Mamba is the first attention-free model to match the performance of a very strong Transformer++ recipe. Mamba-3B's quality matches that of Transformers twice its size."
翻译:
Mamba 是首个无注意力模型,能够与非常强大的 Transformer++ 方案性能相当。Mamba-3B 的质量与两倍规模的 Transformer 相当。
零样本下游评估(HellaSwag、PIQA、ARC、WinoGrande 等基准):
- Mamba-2.8B 的零样本表现匹配 Pythia-6.9B(同为在 Pile 上训练)
- 在多个基准上超越同等大小的 Transformer
新手讲解:
这是里程碑式的结果。在此之前,无论 SSM 怎么改进,在语言建模上总是比 Transformer 差一截。Mamba 第一次打破了这个规律,证明"无注意力架构"也可以在语言上媲美 Transformer。
4.3 推理效率对比
| 指标 | Mamba | Transformer |
|---|---|---|
| 推理吞吐量 | 4-5× 更高 | 基准 |
| 内存占用(推理) | 常数(固定状态) | 随序列长度线性增长(KV 缓存) |
| 训练计算量 | O(n) 线性 | O(n²) 二次 |
| 扫描操作速度 | 比 PyTorch 快 40× | N/A |
新手讲解:
Transformer 在推理时需要缓存所有历史 token 的键值对(KV 缓存),序列越长,内存越多。处理 1 万 token 的序列比 1000 token 需要多 10 倍内存,处理 10 万 token 需要多 100 倍内存。
Mamba 不存在这个问题——无论序列多长,推理时的内存永远只是那个固定大小的状态向量。这对部署长上下文模型(如处理长文档、长代码库)有重大价值。
4.4 DNA 建模
设置:
DNA 序列建模(基因组学),序列极长(基因组可达数百万碱基对)。
结论:
- 40M 参数的 Mamba 用 3-4 倍更少的参数匹配基线性能
- 与 HyenaDNA 不同,性能随上下文长度增加而单调提升,直至 100 万 token
新手讲解:
这展示了 Mamba 最重要的优势之一:超长序列下的性能优势。DNA 建模需要处理数百万碱基的超长序列,Transformer 根本不实用(O(n²) 的代价无法承受)。Mamba 在这里展现出了 Transformer 无法企及的能力。
4.5 音频建模
结论:
- 在音高和音乐建模任务上超越 SaShiMi 基线
- 在 SC09 语音生成任务上:小型 Mamba 模型超越了最先进的(且规模更大的)基于 GAN 和扩散模型的方法
新手讲解:
这进一步证明 Mamba 不是只适合文本的专用架构,而是真正的通用序列建模骨干。音频序列往往比文本更长(采样率 16kHz 意味着每秒 16000 个时间步),Mamba 的线性复杂度在这里尤为重要。
总结:Mamba 的核心价值
创新点总结
| 创新 | 解决的问题 | 关键思路 |
|---|---|---|
| 选择性机制 | LTI 模型"内容盲",无法按内容过滤信息 | 让 Δ、B、C 成为输入的函数,动态调整状态更新规则 |
| 硬件感知并行扫描 | 选择性递归在 GPU 上看似串行,朴素实现极慢 | 核融合(SRAM 内完成计算)+ 并行扫描(O(log n)并行步)+ 重算梯度(省内存) |
| 简化 Mamba 模块 | 需要端到端可用的完整架构 | SSM + 门控 MLP 合并为单模块,无注意力层 |
与 Transformer 对比一览
| 维度 | Transformer | Mamba |
|---|---|---|
| 训练复杂度 | O(n²) | O(n)(线性) |
| 推理内存 | O(n)(随序列增长) | O(1)(常数,固定状态) |
| 长序列能力 | 受窗口限制,性能在窗口外下降 | 无长度限制,性能随长度单调提升 |
| 内容感知 | 强(注意力机制天生内容感知) | 强(选择性机制实现内容感知) |
| 语言建模性能 | 强基线 | 3B ≈ Transformer 6B 效果 |
| 推理吞吐量 | 基准 | 4-5× 更高 |
核心类比(最终版)
Transformer 像一个有无限记忆的图书馆员:你问他任何问题,他都能回到书架上翻出每一页来比对。效果极好,但随着书越来越多,找书的时间呈平方级增长。
早期 SSM 像一个只有固定大小笔记本的助手:他会把读过的每一页内容以相同方式摘录进笔记本。笔记本大小恒定(效率好),但因为所有内容一视同仁,重要内容和废话都被同等对待,笔记质量堪忧(效果差)。
Mamba 像一个有固定大小笔记本、但会根据内容判断价值的智能助手:读到关键信息,大篇幅摘录;读到无关内容,直接跳过;读到转折词"但是",先清空旧笔记再重新开始。笔记本永远固定大小(效率好),但记录的内容经过了智能筛选(效果也好)。
延伸阅读
- S4(Gu et al., 2021):Mamba 的前身,第一个实用的深度 SSM,首次展示 SSM 能在长程依赖任务(如 Long Range Arena)上超越 Transformer。
- H3(Fu et al., 2023):将 S4 用于语言建模的尝试,引入了 SSM+门控的组合,是 Mamba 架构的直接前身。
- HiPPO(Gu et al., 2020):提出初始化矩阵 A 的理论框架(高阶多项式投影算子),帮助 SSM 在长序列上保持稳定的记忆。
- FlashAttention(Dao et al., 2022):Tri Dao 同一作者的工作,同样用硬件感知的分块计算优化了 Transformer,与 Mamba 的算法设计思路有异曲同工之妙。
- Mamba-2(Dao & Gu, 2024):进一步统一了 SSM 和注意力机制的理论框架(State Space Duality),证明 Mamba 可以被视为线性注意力的一种受限形式。
精读笔记覆盖章节:Abstract、Introduction(核心段落)、Section 2(SSM 背景,全部)、Section 3(选择性机制、硬件算法、Mamba 模块、性质分析,全部段落)、Section 4(关键实验:语言建模、DNA、音频、效率对比)。
总字数:约 7500 字。