Mamba

[2023]

Gu A, Dao T. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752, 2023.

Mamba

Mamba 是一种面向长序列的序列模型。它想做的事情可以概括成一句话：尽量保留 RNN / SSM 那种沿序列线性扩展的计算方式，同时补上模型对长距离依赖和内容选择的建模能力。论文把这个核心改动概括为 selective state space：状态不再按固定规则机械更新，而是会根据当前输入决定哪些信息该写入、该遗忘、该继续保留。

动机

Transformer 之所以强，很大程度上是因为自注意力允许信息在上下文窗口内密集路由，所以它特别擅长处理复杂依赖。但这种做法也带来一个根本代价：随着窗口长度增长，计算和存储成本会上升，而且自回归推理还要显式保存过去上下文。相对地，递推式模型只有一个持续更新的状态，训练和推理都更容易做到线性扩展，可问题在于，这类模型往往不够擅长“按内容做选择”。Mamba 的出发点正是这里：真正的问题不只是能不能省算力，而是线性模型能不能像强模型那样，根据内容决定什么该记、什么该忘。

状态空间模型（SSM）

经典状态空间模型可以先写成连续形式：

$$\frac{dh(t)}{dt}=Ah(t)+Bx(t)$$$$y(t)=Ch(t)$$

这里的 $x(t)$ 是输入，$h(t)$ 是隐藏状态，$y(t)$ 是输出。直觉上，它表示模型用一个隐状态去压缩过去信息，再让输入不断推动这个状态演化。离散化到序列上以后，可以写成

$$h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$$$$y_t=C{h}_t$$

所以从感觉上看，它很像一种结构更规整、数学形式更强的 RNN。论文也强调，structured SSM 既可以按 recurrence 来理解，也能在特定条件下改写成 convolution，这正是它能在线性或近线性复杂度下处理长序列的重要原因。

为什么传统 SSM 还不够

早期 structured SSM 已经说明，长序列不一定非得依赖标准注意力来建模，递推和卷积式路径也能取得不错效果。但它们有一个明显限制：参数通常是固定的，或者至少不够强地依赖输入内容。这样一来，模型面对不同 token 时，更新状态的规则几乎不变，更像在用一套统一模板处理所有位置，而不是根据当前内容动态决定“现在最重要的是什么”。论文把这点直接概括为 prior models 缺少 content-based reasoning，这也是它们在文本这类离散、信息密集模态上往往不如 Transformer 的关键原因。

Selective State Space

Mamba 不是把所有参数都直接改成输入相关。更准确地说，论文和官方实现里，连续参数 $A$ 仍然是共享的；真正由当前输入生成的是 $\mathrm{\Delta }$、$B$、$C$。只是在离散化之后，真正参与递推的 ${\overline{A}}_t$ 和 ${\overline{B}}_t$ 才因此变成随时间步变化的量。换句话说，Mamba 的 selectivity 不是“每一步都换一整个系统”，而是让系统根据当前内容动态调节记忆更新的方式。

如果把它写成更直观的离散形式，可以把 Mamba 理解为

$$h_t=\overline{A}\ast th\ast t-1+{\overline{B}}_t{x}_t$$$$y_t=C_t{h}_t+D{x}_t$$

这里最关键的不是公式外形，而是这些带 $t$ 的量意味着：模型面对不同输入时，会决定旧状态该保留多少、新信息该写入多少，以及最后该从状态里读出什么。实现里的 $D{x}_t$ 还对应一个 skip 路径，所以读出也不是单纯只看状态。

Mamba block 的结构

从实现上看，一个 Mamba block 是一个相当紧凑的统一模块。输入先经过线性投影，然后被分成两支。主支路先过一个局部的一维卷积，用来补足短程模式建模，接着再通过投影生成 $dt$、$B$、$C$，送进 selective SSM；另一支路则保留下来，在末端经过激活后对主支路输出做门控，最后再投影回原维度。论文也明确把 Mamba 描述成一种不含 attention、也不再单独保留传统 MLP block 的简化架构。

这里加局部卷积是很自然的，因为很多序列任务同时需要两种能力：一类是邻近位置的局部模式，另一类是跨很长上下文的依赖。Mamba 没有把所有事情都丢给 SSM，而是让卷积先处理更短程的结构，再让 selective SSM 负责更长程的状态建模。这样一来，整个 block 既统一，又有比较明确的分工。

为什么它能做到线性扩展

Mamba 不显式构造 token 两两之间的全连接关系，而是沿序列递推更新状态，所以它对序列长度的扩展是线性的。难点在于，一旦参数依赖输入，原来那种依赖时间不变性的高效卷积路径就不能直接用了。为了解决这个问题，论文专门设计了 hardware-aware parallel scan：理论对象仍然是 recurrence，但工程实现不等于老式 RNN 那种按时间步死板串行地跑。

Mamba 的确不像 Transformer 那样维护一个会随着上下文不断变大的 KV cache，但这不等于它“完全没有缓存”。官方实现里会维护固定大小的 recurrent states，包括卷积状态 conv_state 和状态空间的 ssm_state。

Mamba 到底在“记忆”什么

如果把 Transformer 和 Mamba 放在一起看，可以把它们理解成两种不同的记忆哲学。

Transformer 更像显式记忆：历史 token 基本都被保留下来，当前 token 需要时再回头查。

Mamba 则更像隐式记忆：它不会把所有历史原样摊开，而是不断把历史压缩进一个持续更新的状态里。这样做的好处是效率高、长度扩展更自然，代价是模型必须学会如何做足够聪明的压缩。

论文在动机部分其实把这一点讲得很直接：序列建模的核心问题，本来就是如何把上下文压缩进更小的状态，而真正决定压缩质量的关键，就是有没有选择性。

一个通俗类比

如果把长文本建模看成“做读书笔记”，Transformer 更像把整本书一直摊在桌面上，查阅时很直接，但书越厚、桌子越满，成本也越高。Mamba 更像一边读一边写动态摘要：不重要的内容可以跳过，关键线索会被保留下来，遇到新信息时还会重写已有摘要的组织方式。

从这个角度看，Mamba 的 selective mechanism，本质上就是在让这个“动态摘要系统”更聪明。它不是机械压缩，而是按内容决定怎么压缩。