Mamba: Linear-Time Sequence Modeling with Selective State Spaces

为什么这篇论文重要

2017 年至今，Transformer 是大模型的唯一选择。 但它有一个根本问题：注意力是 O(n²) 复杂度 —— 序列翻倍，计算和显存翻 4 倍。

2023 年这篇 Mamba 论文挑战了这个独家地位—— 用”选择性状态空间模型（Selective SSM）” + 线性复杂度， 在多项任务上接近甚至超过 Transformer。

“Transformer is all you need” 可能不再是真理。

状态空间模型（SSM）是什么

灵感来自经典控制理论的 状态空间方程：

$$h_t = A h_{t-1} + B x_t \quad \text{(state update)}$$ $$y_t = C h_t + D x_t \quad \text{(output)}$$

• $x_t$：当前输入
• $h_t$：当前状态（固定大小，比如 64 维）
• $y_t$：输出
• $A, B, C, D$：参数

核心特点：

• 像 RNN —— 状态串行更新
• 像 CNN —— 可以”展开”成卷积
• 像 Attention —— 能学长距离依赖
• 但复杂度是线性！

Mamba 的关键贡献：选择性

之前的 SSM（S4 等）参数是固定的—— 对所有 token 都用同样的 $A, B, C$。

Mamba 让参数输入相关（input-selective）：

$$A(x), B(x), C(x)$$

——根据当前输入决定 $A, B, C$ 的具体值。

这让模型能”选择性地记或忘” —— 类似 LSTM 的门控机制，但用了完全不同的数学结构。

训练速度的工程突破

经典 SSM 算起来慢—— Mamba 的另一作者 Tri Dao（FlashAttention 作者）实现了硬件友好的并行扫描算法：

• 利用 GPU 的并行性
• 减少显存读写
• 实现 O(n) 训练复杂度

关键代码就 200 行 Triton kernel——但极快。

性能对比

语言建模（Pile 数据集）

模型	参数	Perplexity
Transformer	130M	11.13
H3-Hyena	130M	11.49
Mamba	130M	10.76

同尺寸下击败 Transformer——在小尺度。

长序列任务

序列长度	Transformer	Mamba
1k	1x baseline	1.5x baseline
16k	8x slow	1.5x baseline
64k	OOM	1.6x baseline
1M	impossible	works

这是 Mamba 真正的杀手锏——超长序列性能不退化。

DNA 序列任务

DNA 一段就上百万碱基——Transformer 处理不了。 Mamba 能轻松处理 100 万长度 的 DNA 序列。

DNA / 长视频 / 长代码——这些场景 Mamba 远超 Transformer。

论文的争议

”ICLR 拒收”

最讽刺的事—— Mamba 论文被 ICLR 2024 拒了（评审觉得”对 Transformer 替代性不足”）。 但社区影响巨大——发布后 6 个月被引上千次，工业界跟进。

这成了一个学术界标准的失败案例——评审跟不上社区直觉。

“Mamba 真的能取代 Transformer 吗”

支持者：

• 复杂度优势是真的
• 长序列性能远超 Transformer
• 工程实现极快

怀疑者：

• 大尺度（70B+）的性能还没充分验证
• In-context learning 能力差于 Transformer
• 工具调用、推理任务上仍弱

2026 年答案：两者并存 —— Mamba 在长上下文场景吃下市场，Transformer 仍是通用首选。

Mamba 之后的发展

Mamba 2（2024）

一作 Albert Gu 在 ICML 2024 上发了 Mamba 2 —— 理论上统一了 SSM 和 Attention，并且更快。

Jamba（AI21 2024）

混合架构：大部分层用 Mamba，少数层用 Attention。 性能 + 效率最优。已经有商用模型。

Falcon Mamba 7B（2024）

第一个开源大 Mamba 模型（7B）—— 社区验证它在 7B 量级是真的能跑。

Tegridy / Zamba 等

各种 Mamba 变种和混合模型—— SSM 已经成为可选架构，不只是研究玩具。

用 PyTorch 跑

最简版 Mamba block：

import torch
import torch.nn as nn

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

        # 输入投影
        self.in_proj = nn.Linear(d_model, self.d_inner * 2)

        # 1D 卷积（提供局部上下文）
        self.conv1d = nn.Conv1d(
            self.d_inner, self.d_inner,
            kernel_size=d_conv, groups=self.d_inner, padding=d_conv-1
        )

        # 输入相关的 A, B, C 参数
        self.x_proj = nn.Linear(self.d_inner, d_state * 2 + d_state)
        self.dt_proj = nn.Linear(d_state, self.d_inner)

        # 输出投影
        self.out_proj = nn.Linear(self.d_inner, d_model)

    def forward(self, x):
        # 1. 投影
        x_and_res = self.in_proj(x)
        x, res = x_and_res.chunk(2, dim=-1)

        # 2. 卷积（提供 local context）
        x = self.conv1d(x.transpose(-1, -2))[:, :, :x.size(-2)].transpose(-1, -2)
        x = nn.SiLU()(x)

        # 3. 选择性 SSM（核心）
        delta, B, C = self.x_proj(x).split([self.d_state, self.d_state, self.d_state], dim=-1)
        delta = nn.Softplus()(self.dt_proj(delta))
        # ... 这里有个并行扫描算法 ...

        # 4. Gating
        x = x * nn.SiLU()(res)

        return self.out_proj(x)

完整实现复杂——推荐用 Mamba 官方 GitHub。

一些有意思的事

Albert Gu

Mamba 第一作者，Carnegie Mellon 助理教授。 此前发表过 S4、HiPPO 等 SSM 系列工作—— 他坚持 SSM 方向 5 年，终于在 Mamba 看到回报。

这是”坚持非主流方向”的成功案例。Transformer 火 7 年时，他选择继续做 SSM。

Tri Dao

二作 —— 也是 FlashAttention 的发明人。 他同时既改进 Transformer 又试图替代 Transformer。

真正强的研究者从不固守一个范式。

“选择性”和注意力的相似

Mamba 的”选择性”机制—— 让参数随输入变化—— 本质上是另一种形式的”注意力”。

学界后来证明：SSM 和 Attention 在数学上几乎等价（Mamba 2 的洞察）。

不同架构殊途同归。这暗示更深层的”序列建模本质”。

推荐配套阅读

• HelloAI: L3-04 RNN / LSTM + L3-05 Attention
• Mamba 论文 原文
• Tri Dao 的 YouTube 演讲
• Falcon Mamba 7B GitHub —— 实战
• The Annotated Mamba（社区注释版）