Chaos and Order

💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

1. 引言:Transformer 的局限与 SSM 的崛起
2. 背景:Structured State Space Models(S4)
3. Mamba:Selective State Space Models
4. Mamba-2:State Space Duality
- 4.1 SSD(State Space Dual)模型
- 4.2 性能对比
5. 实战:使用 Mamba 模型
- 5.1 安装与推理
- 5.2 使用 Mamba-2
6. Mamba 与 Transformer 对比
7. 局限与展望
- 7.1 当前局限
- 7.2 混合方案
8. 小测验

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) = \mathbf{A}h(t) + \mathbf{B}x(t)

y(t) = \mathbf{C}h(t)

$h(t) \in \mathbb{R}^N$ :hidden state
$x(t) \in \mathbb{R}$ :输入
$y(t) \in \mathbb{R}$ :输出
$\mathbf{A} \in \mathbb{R}^{N \times N}$ , $\mathbf{B} \in \mathbb{R}^{N \times 1}$ , $\mathbf{C} \in \mathbb{R}^{1 \times N}$

2.2 离散化(Zero-Order Hold)

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

\overline{\mathbf{A}} = \exp(\Delta \mathbf{A})

\overline{\mathbf{B}} = (\Delta \mathbf{A})^{-1}(\exp(\Delta \mathbf{A}) - \mathbf{I}) \cdot \Delta \mathbf{B}

离散化后的递归关系:

h_t = \overline{\mathbf{A}} h_{t-1} + \overline{\mathbf{B}} x_t

y_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、Δ)。这带来了两个问题:

无法做 content-based reasoning:在需要根据输入内容做不同处理的任务(例如 selective copying、induction heads)上性能下降
放弃了 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)  ← 依赖输入!

具体来说:

\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 的直观理解

关键在于 $\Delta_t$ (step size):

$\Delta_t$ 较大时: $\overline{\mathbf{A}} \approx \mathbf{I}$ → 保留此前的 state(忽略输入)
$\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.7B	6.7	1.0x	1.0x
Mamba-1	2.8B	6.2	1.3x	5.2x
Mamba-2	2.7B	6.1	2.1x	5.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 对比

特性	Transformer	Mamba
训练复杂度	O(N²)	O(N)
推理复杂度	O(N) per token(KV Cache)	O(1) per token
内存(推理)	O(N) KV Cache	O(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。