Mamba
Transformer 的 self-attention 核心优势是全局感受野和内容选择性——每个 token 可以直接关注任意其他 token。但在长序列上，这两个优势变成了瓶颈：
- O(L²) 注意力计算 → 序列长度翻倍，计算量翻四倍
- O(L) KV cache → 长上下文推理时显存爆炸
- 内容选择通过可学习的 Query-Key-Value 投影实现，需要足够的参数量支撑
问题的根源在于：attention 是"全对交互"机制。每个时间步必须与所有历史 token 计算注意力分数，这对语言的长程依赖建模是有效的，但对计算资源是奢侈的。
State Space Models（SSM）是应对这一问题的经典方案——它们通过状态压缩将 O(L²) 的交互转化为 O(L) 的递归。但此前 LTI（线性时不变）SSM 在"选择性"任务上完全失效。
标准 SSM（S4）是线性时不变系统。数学形式上，给定输入 $x_t$，状态 $h_t$ 和输出 $y_t$：
$h_t = \bar{A}h_{t-1} + \bar{B}x_t$
$y_t = \bar{C}h_t$
— 其中 $\bar{A}, \bar{B}, \bar{C}$ 是时间步固定的矩阵
时不变性的关键优势：SSM 可以等价于全局卷积——整个序列可以一次性用 FFT 并行计算。这使得训练非常高效。
但致命弱点：参数与输入无关。无论输入是什么，卷积核的权重都是固定的。这意味着模型无法根据内容动态决定"该记住什么、该忘记什么"。
选择性任务失败案例：Selective Copying
任务：输入一串 token，输出在特定标记位置之后复制对应的 token。LTI SSM 的卷积核无法学会"在标记处记住什么"——因为标记的位置是输入决定的，而卷积核的权重是固定的。
这不是调参问题，是结构性的：LTI 系统不具备根据内容调整行为的能力。
在语言建模等真实任务上，LTI SSM（S4）已经输给了 Transformer。核心原因：语言需要内容选择能力。
2.1 选择性机制的形式化
Mamba 的解决思路极为直接：让 SSM 参数变成输入依赖的。
$\Delta_t = \text{softplus}(\text{Linear}_1(x_t))$
$B_t = \text{Linear}_N(x_t)$ # 输入矩阵，输入依赖
$C_t = \text{Linear}_N(x_t)$ # 输出矩阵，输入依赖
$A$ 保持固定（diagonal structure）
— 核心洞察：输入 $x_t$ 决定当前时间步的参数配置
直观理解：每个时间步的 SSM 参数都是该时间步输入的函数。这等价于给模型装上了"选择性滤波器"——能根据内容动态调整状态更新的策略。
2.2 直观解释：SSM 作为记忆压缩
可以把 SSM 递归看作一种记忆压缩机制：
- Attention：显式存储所有历史 token 的表示，O(L) 空间
- SSM：压缩到 N 维状态向量 $h_t$，O(N) 空间（N 是 state dimension，通常远小于 L）
内容选择能力来自哪里？来自让压缩函数（Δ, B, C）变成输入依赖的。如果压缩函数是固定的，模型无法区分重要信息和平凡信息；如果是输入依赖的，模型学会了"什么时候该更新状态、什么时候该保持状态"。
2.3 选择性的代价：失去卷积等价性
选择机制带来的代价：Mamba 不再等价于全局卷积。这意味着无法用 FFT 高效训练，必须回到递归计算——而递归是 O(L)串行的，理论训练效率反而不如并行方法。
Mamba 的解法：hardware-aware parallel scan。
Hardware-aware Parallel Scan
核心思想：不在 HBM（显存）中 materialize 完整状态，而是只在使用时在高速 SRAM 中计算。
并行扫描把 O(L) 的串行递归转化为可并行的扫描模式，但只在 SRAM 中执行——既获得递归的 O(L) 内存，又避免串行计算的耗时。IO 开销最小化。
结果：A100 GPU 上训练速度比 S4 快 3×。
Mamba Block 的设计遵循现代 Transformer 的惯例（类比 RMSNorm + 残差），同时融入了 SSM 的选择性机制：
3.1 三个关键设计
- 选择性 SSM 分支：核心创新所在，SSM 参数由输入决定
- 门控机制：主分支的 Linear + SiLU 提供非线性门控，控制 SSM 输出与主分支信息的融合比例
- 残差连接：标准现代网络设计，稳定训练
3.2 SSM 参数初始化
矩阵 A 的初始化：
$A = \text{diag}(a_1, a_2, ..., a_N) + \text{sub-diagonal}(b_1, ..., b_{N-1})$
其中 $a_i = \lambda_i \cdot \sigma(\text{init})$
对角线元素初始化为负实数，保证特征值在复平面左半平面 → 连续时间系统稳定。sub-diagonal 用于增加 hidden state 之间的信息流动。
3.3 与标准 SSM 的对比
| 维度 | 标准 SSM（S4） | Mamba（选择性 SSM） |
|---|---|---|
| 参数 | 固定（时间步不变） | 输入依赖 |
| 训练方式 | FFT 卷积（高效并行） | 硬件感知并行扫描 |
| 选择性任务 | 失效 | 有效 |
| 推理复杂度 | O(L)（递归） | O(L)（递归，常数 cache） |
| 语言建模 | 差 | 接近 Transformer |
4.1 选择性复制任务的验证
Selective Copying 任务（需要内容选择）和 Induction Head 任务（需要追踪序列中的关联）是检验选择性机制的标准任务。Mamba 在这两个任务上都接近最优，而 LTI SSM（S4）完全失败。
4.2 语言建模 benchmark
| 模型 | 困惑度（PPL） | 吞吐量 | 备注 |
|---|---|---|---|
| Mamba（1.3B） | 竞争水准 | 5× vs Transformer | 线性复杂度带来推理优势 |
| S4（S4D） | 差 | — | LTI，无法做内容选择 |
| Transformer（1.3B） | 基准 | 基准 | O(L²) 推理成本 |
4.3 硬件效率
关键数据：Mamba-1.3B 在 A100 GPU 上的实际推理速度是 Transformer 的 5 倍，显存占用大幅降低。原因在于常数 KV cache 和 O(L) 的递归性质。
4.4 Effective Receptive Field（ERF）分析
ERF 分析显示：Mamba 的状态压缩并没有损失全局感受野，反而比之前的 SSM 方法更接近 Transformer 的全局覆盖能力。这是选择机制带来的间接收益——模型学会了什么信息值得压缩到状态里。
Mamba 的三合一
Mamba 第一次同时实现了：
- O(L) 推理：递归形式，常数 KV cache
- 内容选择能力：选择性 SSM 根据输入动态调整
- 硬件感知实现：并行扫描避免串行计算瓶颈
这三者此前在 SSM 框架下是互相矛盾的——Mamba 通过选择机制和系统设计将其统一起来。
对后续研究的影响
- Mamba-2 / SSD：揭示了 SSM 和 attention 的统一框架，半可分矩阵理论
- 视觉 Mamba：Vim、VMamba 将选择性 SSM 引入 2D/3D 视觉任务
- Jamba / Striide：混合 Transformer + SSM 的工程探索