Mamba: Linear-Time Sequence Modeling with Selective State Spaces 论文分析

论文概述
动机：Transformer 的局限
核心思想：Selective State Spaces
硬件优化：Selective Scan 算法
- 问题：HBM ↔ SRAM 瓶颈
- 解决方案：Kernel Fusion + Recomputation
Mamba 模块架构
实验结果
Mamba vs Transformer：该用哪个？
后续研究：Mamba-2 与混合架构
结语
小测验

论文概述

标题: Mamba: Linear-Time Sequence Modeling with Selective State Spaces
作者: Albert Gu, Tri Dao (Carnegie Mellon University, Princeton University)
发表: 2023 年 12 月（arXiv: 2312.00752），被 ICML 2024 收录
代码: github.com/state-spaces/mamba

动机：Transformer 的局限

Transformer 稳坐序列建模的王座，但存在几个根本性的问题：

O(N²) 复杂度：Self-attention 相对于序列长度呈二次复杂度
长序列处理的瓶颈：上下文超过 128K 时内存/计算量会爆炸式增长
推理低效：KV Cache 会随序列长度成比例增长

State Space Model（SSM）理论上具有 O(N) 复杂度，但既有的 SSM（S4、H3 等）因为与输入无关的固定参数，在语言建模上落后于 Transformer。

核心思想：Selective State Spaces

既有 SSM 的问题

既有 SSM 是 Linear Time-Invariant（LTI） 系统：

$h'(t) = Ah(t) + Bx(t)$ $y(t) = Ch(t)$

这里的 A、B、C 都是固定矩阵。也就是说，无论输入是 "hello" 还是 "world"，都以同样的方式处理。这意味着content-aware reasoning是不可能的。

Mamba 的解决方案：Selection Mechanism

Mamba 的核心创新，是让B、C、Δ依赖于输入：

# 既有 SSM：固定参数
B = nn.Parameter(torch.randn(N))  # 可学习，但与输入无关

# Mamba：输入依赖参数
B = nn.Linear(D, N)(x)      # B 会随 x 变化
C = nn.Linear(D, N)(x)      # C 会随 x 变化
delta = nn.Linear(D, D)(x)  # step size 也会随 x 变化

这正是Selective State Space 的含义。根据输入，选择要把哪些信息存入状态、哪些信息忽略掉。

直观理解

把 Selection Mechanism 直观地理解一下：

Δ(delta)："这个 token 要看得多重要" — Δ 大就强烈反映当前输入，Δ 小就保留之前的状态
B："这个输入的哪部分要记录进状态"
C："从状态里取出哪部分作为输出"

这与 Transformer 中 attention 所起的作用类似，但以 O(N) 的复杂度完成。

硬件优化：Selective Scan 算法

由于 Selective SSM 依赖输入，无法使用既有 SSM 那种高效的卷积（convolution）技巧。Mamba 用硬件感知算法解决了这个问题。

问题：HBM ↔ SRAM 瓶颈

GPU 的内存层级：

┌────────────────────┐
│   SRAM (20MB)      │  ← 非常快，非常小
├────────────────────┤
│   HBM (80GB)       │  ← 慢，大
└────────────────────┘

朴素的实现会反复把中间状态存入/读出 HBM，从而变成内存受限（memory-bound）。

解决方案：Kernel Fusion + Recomputation

Mamba 的 Selective Scan 算法：

在 SRAM 中一次性计算：把离散化（discretization）、selective scan、输出计算融合进一个 kernel
不保存中间状态：前向传播只把 (N, B, C, Δ) 存入 HBM，反向传播时重新计算中间状态
结果：与 FlashAttention 类似的 IO-aware 优化

# 简化版 Selective Scan（概念代码）
def selective_scan(x, A, B, C, delta):
    """
    x: (batch, length, d_model)
    Returns: (batch, length, d_model)
    """
    batch, length, d = x.shape
    n = A.shape[1]  # state dimension

    # Discretize A, B using delta
    deltaA = torch.exp(delta.unsqueeze(-1) * A)  # (B, L, D, N)
    deltaB_x = delta.unsqueeze(-1) * B.unsqueeze(2) * x.unsqueeze(-1)

    # Sequential scan (实际中通过 CUDA kernel 并行化)
    h = torch.zeros(batch, d, n, device=x.device)
    ys = []
    for i in range(length):
        h = deltaA[:, i] * h + deltaB_x[:, i]
        y = (h * C[:, i].unsqueeze(1)).sum(-1)
        ys.append(y)

    return torch.stack(ys, dim=1)

Mamba 模块架构

Mamba 从 H3 架构获得启发，采用了简化的模块设计：

Input (x)
    │
    ├──── Linear projection (expand) ────┐
    │                                     │
    ▼                                     ▼
  Conv1D                              SiLU (gate)
    │                                     │
    ▼                                     │
  SiLU                                    │
    │                                     │
    ▼                                     │
  SSM (selective scan)                    │
    │                                     │
    ▼                                     │
    ×─────────────────────────────────────┘
    │
    ▼
  Linear projection (reduce)
    │
    ▼
  Output

用 PyTorch 实现如下：

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = d_model * expand

        # Input projection
        self.in_proj = nn.Linear(d_model, d_inner * 2, bias=False)

        # Conv1D
        self.conv1d = nn.Conv1d(
            d_inner, d_inner, d_conv,
            padding=d_conv - 1, groups=d_inner
        )

        # SSM parameters
        self.x_proj = nn.Linear(d_inner, d_state * 2 + 1, bias=False)  # B, C, delta
        self.dt_proj = nn.Linear(1, d_inner, bias=True)

        # A parameter (structured)
        A = torch.arange(1, d_state + 1).float().repeat(d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))

        self.D = nn.Parameter(torch.ones(d_inner))
        self.out_proj = nn.Linear(d_inner, d_model, bias=False)

    def forward(self, x):
        batch, length, _ = x.shape

        # Dual path
        xz = self.in_proj(x)
        x, z = xz.chunk(2, dim=-1)

        # Conv + activation
        x = self.conv1d(x.transpose(1, 2))[:, :, :length].transpose(1, 2)
        x = F.silu(x)

        # SSM
        y = self.ssm(x)

        # Gate and project
        y = y * F.silu(z)
        return self.out_proj(y)

实验结果

语言建模（The Pile）

模型	参数量	Perplexity
Transformer++	1.3B	8.94
H3	1.3B	9.19
Hyena	1.3B	9.07
RWKV-4	1.5B	9.02
Mamba	1.4B	8.63

Mamba 在相近规模下达到了比 Transformer更低的 perplexity。尤其是在 2.8B 规模下，相较 Transformer++，Mamba 展现出约0.3 perplexity的提升。

推理速度

序列长度	Transformer	Mamba
2K	1x	1x
8K	1x	快 3 倍
64K	1x	快 5 倍
1M	OOM	正常运行

序列长度越长，Mamba 相对 Transformer 的速度优势就越明显。

Selective Copying & Induction Heads

论文用两个核心的合成任务验证了 Selection Mechanism 的效果：

Selective Copying：从长序列中选择性地只复制特定 token → 既有 SSM（S4）失败，Mamba 成功
Induction Heads：识别 "A B ... A → B" 模式 → 既有 SSM 失败，Mamba 成功

这两个任务都必须具备content-aware reasoning能力，没有 Selection Mechanism 就无法解决。

Mamba vs Transformer：该用哪个？

标准	Transformer	Mamba
长序列	△ (O(N²))	◎ (O(N))
推理速度	△ (KV Cache 增长)	◎ (固定状态)
In-context learning	◎	○
训练并行化	◎ (完全并行)	○ (scan 并行化)
生态/工具	◎	△

后续研究：Mamba-2 与混合架构

Mamba 之后有了许多进展：

Mamba-2（Dao & Gu, 2024）：提出 State Space Duality（SSD） 框架，揭示了 SSM 与 Attention 之间的理论联系
Jamba（AI21, 2024）：Transformer + Mamba 混合架构，256K 上下文
Zamba（Zyphra, 2024）：在小型混合模型中达到 SOTA

结语

Mamba 用Selection Mechanism这一简单却强大的思想，克服了 SSM 的局限，以 O(N) 复杂度实现了媲美甚至超越 Transformer 的性能。硬件感知算法设计，正是把理论上的效率转化为实际速度的关键。

尽管 Transformer 依然是主流，但 Mamba 与混合架构正在长序列处理至关重要的领域里，逐渐扩大存在感。

小测验

Q1: Mamba 想要解决的 Transformer 核心局限是什么？

Self-attention 的 O(N²) 时间/空间复杂度。序列长度越长，计算量与内存都会呈平方增长，导致长序列处理效率低下。

Q2: Selective State Space 中"Selective"是什么意思？

让 SSM 的参数（B、C、Δ）依赖于输入，从而可以选择（select）哪些信息存入状态、哪些信息忽略。

Q3: 既有 SSM（S4 等）在语言建模上表现不佳的根本原因是什么？

Linear Time-Invariant（LTI）特性。参数与输入无关、固定不变，因此无法实现 content-aware reasoning。

Q4: Mamba 中 Δ(delta) 参数的直观含义是什么？

决定当前输入要被多大程度地反映出来的"step size"。Δ 大就强烈反映当前输入，Δ 小就保留之前的状态。

Q5: 为什么 Mamba 的 Selective Scan 无法使用既有 SSM 的卷积技巧？

由于参数依赖输入，系统不再是 LTI 系统，因此转换成卷积、用 FFT 并行计算的技巧不再适用。

Q6: Mamba 的硬件优化中，与 FlashAttention 类似的核心策略是什么？

IO-aware 优化。在 SRAM 中融合（kernel fusion）所有计算，不把中间状态存入 HBM，而是在反向传播时重新计算（recomputation）。

Q7: Selective Copying 任务想要验证的能力是什么？

从长序列中选择性地记忆并复制特定 token 的 content-aware reasoning 能力。固定参数的 SSM 无法解决这个任务。

Q8: Mamba-2 的核心贡献 State Space Duality（SSD）是什么？

一个揭示 SSM 与 Attention 在数学上是同一运算的不同表现形式的框架。它为结合两种方法的优点提供了理论基础。