Skip to content
Published on

Mamba 论文精读：用 Selective State Space Models 超越 Transformer

🇰🇷KO🇺🇸EN🇯🇵JA🇨🇳ZH
필사 모드
分享
Authors

1. 引言:Transformer 的局限与 SSM 的崛起

Transformer 架构自 2017 年问世以来,几乎主宰了 NLP、Vision、Audio 等所有序列建模领域。但它存在根本性的局限:

  • O(N²) 复杂度:Self-Attention 的时间/空间复杂度随序列长度呈二次增长
  • 推理效率低:每生成一个 token,都要参照整个 KV Cache
  • 长序列处理困难:在 128K+ 上下文中出现内存瓶颈

State Space Model(SSM)是为克服这些局限而提出的替代方案。尤其是 Mamba(Gu & Dao, ICLR 2024),它在 SSM 中引入了选择性(Selective)机制,以 O(N) 复杂度达成了可与 Transformer 匹敌的性能。

2. 背景:Structured State Space Models(S4)

2.1 连续时间 SSM

SSM 是把连续时间系统离散化后的产物:

h′(t)=Ah(t)+Bx(t)h'(t) = \mathbf{A}h(t) + \mathbf{B}x(t) y(t)=Ch(t)y(t) = \mathbf{C}h(t)
  • h(t)∈RNh(t) \in \mathbb{R}^N:hidden state
  • x(t)∈Rx(t) \in \mathbb{R}:输入
  • y(t)∈Ry(t) \in \mathbb{R}:输出
  • A∈RN×N\mathbf{A} \in \mathbb{R}^{N \times N}, B∈RN×1\mathbf{B} \in \mathbb{R}^{N \times 1}, C∈R1×N\mathbf{C} \in \mathbb{R}^{1 \times N}

2.2 离散化(Zero-Order Hold)

把连续系统转换为离散时间:

A‾=exp⁡(ΔA)\overline{\mathbf{A}} = \exp(\Delta \mathbf{A}) B‾=(ΔA)−1(exp⁡(ΔA)−I)⋅ΔB\overline{\mathbf{B}} = (\Delta \mathbf{A})^{-1}(\exp(\Delta \mathbf{A}) - \mathbf{I}) \cdot \Delta \mathbf{B}

离散化后的递归关系:

ht=A‾ht−1+B‾xth_t = \overline{\mathbf{A}} h_{t-1} + \overline{\mathbf{B}} x_t yt=Chty_t = \mathbf{C} h_t

2.3 S4 的核心:HiPPO 初始化

S4(Structured State Spaces for Sequence Modeling)的核心贡献,是用 HiPPO(High-order Polynomial Projection Operator)矩阵来初始化 A:

import torch
import numpy as np

def make_hippo(N):
    """生成 HiPPO-LegS 矩阵"""
    P = np.sqrt(1 + 2 * np.arange(N))
    A = np.zeros((N, N))
    for i in range(N):
        for j in range(N):
            if i > j:
                A[i, j] = P[i] * P[j]
            elif i == j:
                A[i, j] = i + 1
            # i < j: 0
    return -A  # 取负以保证稳定性

2.4 S4 的局限

S4 使用与输入无关的固定参数(A、B、C、Δ)。这带来了两个问题:

  1. 无法做 content-based reasoning:在需要根据输入内容做不同处理的任务(例如 selective copying、induction heads)上性能下降
  2. 放弃了 Attention 的优势:失去了 Transformer 那种 content-aware 的匹配能力

3. Mamba:Selective State Space Models

3.1 核心思路:Selection Mechanism

Mamba 的核心创新,是让 SSM 参数依赖于输入:

S4:    (A, B, C, Δ) = 固定参数
Mamba: (B, C, Δ) = f(input)  ← 依赖输入!

具体来说:

Bt=LinearB(xt),Ct=LinearC(xt),Δt=softplus(LinearΔ(xt))\mathbf{B}_t = \text{Linear}_B(x_t), \quad \mathbf{C}_t = \text{Linear}_C(x_t), \quad \Delta_t = \text{softplus}(\text{Linear}_\Delta(x_t))

这样一来,模型就能根据输入动态决定该记住哪些信息、遗忘哪些信息。

3.2 Selection 的直观理解

关键在于 Δt\Delta_t(step size):

  • Δt\Delta_t 较大时:A‾≈I\overline{\mathbf{A}} \approx \mathbf{I} → 保留此前的 state(忽略输入)
  • Δt\Delta_t 较小时:更专注于当前输入(更新 state)

这与门控机制(gating mechanism)类似:

LSTM 的 forget gate ≈ Mamba 的 Δ

3.3 Hardware-Aware 算法

一旦引入 Selection,参数就变成依赖输入,于是无法再使用 convolution trick。这就有在训练时退回到 O(N²) 复杂度的风险。

Mamba 用 kernel fusion + recomputation 来解决这个问题:

# Mamba 的 Selective Scan(简化版)
def selective_scan(x, delta, A, B, C):
    """
    x: (B, L, D)  - 输入
    delta: (B, L, D) - step size(依赖输入)
    A: (D, N)      - state matrix
    B: (B, L, N)   - input matrix(依赖输入)
    C: (B, L, N)   - output matrix(依赖输入)
    """
    B_batch, L, D = x.shape
    N = A.shape[1]

    # 离散化
    deltaA = torch.exp(delta.unsqueeze(-1) * A)  # (B, L, D, N)
    deltaB = delta.unsqueeze(-1) * B.unsqueeze(2)  # (B, L, D, N)

    # Sequential scan(训练时使用 parallel scan)
    h = torch.zeros(B_batch, D, N, device=x.device)
    ys = []
    for t in range(L):
        h = deltaA[:, t] * h + deltaB[:, t] * x[:, t, :, None]
        y = (h * C[:, t, None, :]).sum(-1)  # (B, D)
        ys.append(y)

    return torch.stack(ys, dim=1)  # (B, L, D)

在实际实现中,CUDA 核函数会把所有中间状态保留在 SRAM 里,以最小化 HBM 访问(这是一种与 FlashAttention 类似的 IO-aware 做法)。

3.4 Mamba Block 架构

Input
  │
  ├──── Linear(D → 2ED) ──── SiLU ──── Conv1d ──── SiLU ──── SSM ────┐
  │                                                                     │
  └──── Linear(D → ED) ──── SiLU ──────────────────────────── × ───────┘
                                                               │
                                                          Linear(ED → D)
                                                               │
                                                            Output

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

        self.in_proj = nn.Linear(d_model, d_inner * 2, bias=False)
        self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv,
                                padding=d_conv-1, groups=d_inner)

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

        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):
        # x: (B, L, D)
        xz = self.in_proj(x)  # (B, L, 2*ED)
        x, z = xz.chunk(2, dim=-1)

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

        # SSM
        A = -torch.exp(self.A_log)
        B_C_dt = self.x_proj(x)
        # ... selective scan ...

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

4. Mamba-2:State Space Duality

4.1 SSD(State Space Dual)模型

Mamba-2(Dao & Gu, ICML 2024)发现了 SSM 与 Attention 之间的数学二重性(Duality):

  • SSM 视角:递归式的 state update → O(N) 推理
  • Attention 视角:特殊结构的矩阵乘法 → 并行训练
SSM recurrence: h_t = A_t h_{t-1} + B_t x_t
                 y_t = C_t h_t

↕ Dual ↕

Structured Attention: Y = (T ⊙ M) X
  where T = lower-triangular causal mask
        M = structured (semi-separable) matrix

4.2 性能对比

模型参数量Pile (ppl)训练 FLOPS/s推理 (tok/s)
Transformer++2.7B6.71.0x1.0x
Mamba-12.8B6.21.3x5.2x
Mamba-22.7B6.12.1x5.5x

核心改进点:

  • 训练速度提升 2 倍(得益于 structured matrix multiplication)
  • 支持更大的 state size(N=64 → N=256)
  • 引入 multi-head 结构

5. 实战:使用 Mamba 模型

5.1 安装与推理

pip install mamba-ssm causal-conv1d>=1.4.0

from mamba_ssm import MambaLMHeadModel
from transformers import AutoTokenizer

# 加载 Mamba-2.8B
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b")
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")

model = model.cuda().half()

# 推理
input_ids = tokenizer("The future of AI is", return_tensors="pt").input_ids.cuda()
output = model.generate(input_ids, max_length=100, temperature=0.7)
print(tokenizer.decode(output[0]))

5.2 使用 Mamba-2

from mamba_ssm import Mamba2

# 单独使用 Mamba-2 层
layer = Mamba2(
    d_model=2048,
    d_state=128,    # Mamba-2 支持更大的 state
    d_conv=4,
    expand=2,
    headdim=64,     # Multi-head
).cuda()

x = torch.randn(1, 1024, 2048).cuda()  # (batch, seq_len, d_model)
y = layer(x)  # (1, 1024, 2048)

6. Mamba 与 Transformer 对比

特性TransformerMamba
训练复杂度O(N²)O(N)
推理复杂度O(N) per token(KV Cache)O(1) per token
内存(推理)O(N) KV CacheO(1) 固定 state
In-context learning强较弱(改进中)
长序列瓶颈高效
硬件利用率并行化最优递归瓶颈(Mamba-2 有所改进)

7. 局限与展望

7.1 当前局限

  • In-context learning:把无限的上下文信息压缩进固定大小的 state 时会有信息损失
  • Retrieval 任务:需要精确检索 token 时,Attention 更有优势
  • 训练稳定性:长序列上存在 gradient 不稳定的问题

7.2 混合方案

近期的趋势是 Mamba + Attention 混合:

  • Jamba(AI21):Mamba 层 + Attention 层混合
  • Zamba(Zyphra):以 Mamba 为主 + 少量 Shared Attention 层

8. 小测验

Q1. Mamba 相对于 S4 最核心的改进是什么?

把 SSM 参数(B、C、Δ)变成依赖输入,从而使 content-based reasoning 成为可能。S4 使用固定参数,因此无法根据输入内容做选择性处理。

Q2. Mamba 中 Δ(delta)扮演什么角色?

作为 step size,起到门控机制的作用。Δ 较大时保留此前的 state(忽略输入),较小时则专注于当前输入(更新 state)。与 LSTM 的 forget gate 类似。

Q3. 引入 Selection 之后为什么无法使用 convolution trick?

Convolution trick 只在时不变(time-invariant)系统中才成立。一旦 Selection 让参数依赖输入,系统就变成时变(time-varying),全局卷积也就无法进行。

Q4. Mamba 的 Hardware-Aware 算法的核心是什么?

在 CUDA 核函数中把中间 state 保留在 SRAM 里,以最小化 HBM 访问。这是一种与 FlashAttention 类似的 IO-aware 做法,用来优化 memory-bound 运算。

Q5. Mamba-2 的 State Space Duality 指的是什么?

SSM 的递归式 state update 与 structured attention 的矩阵乘法在数学上是等价的。这使得训练时可以利用并行化的矩阵乘法,推理时又能利用高效的递归。

Q6. 相对于 Mamba,Transformer 在哪些任务上更有优势?

In-context learning,以及需要精确检索 token 的 retrieval 任务。Attention 可以直接引用序列内任意位置,而 Mamba 必须把信息压缩进固定大小的 state。

Q7. Jamba、Zamba 这类混合模型的设计原理是什么?

把 Mamba 层的 O(N) 效率,与少量 Attention 层的精确检索能力结合起来。大部分层用 Mamba 构成,只在关键位置放置 Attention。

Discuss on Twitter • View on GitHub