Transformers are SSMs: Generalized Models and Efficient Algorithms (Mamba 2)

为什么这篇论文重要

2023 年的 Mamba 被 ICLR 拒了—— 但社区反响巨大。

2024 年作者 Tri Dao + Albert Gu 又发了 Mamba 2—— 这次被 ICML 接收 + 引爆学术界。

更深远的意义：

证明了 SSM 和 Attention 在数学上是等价的——它们是同一种”状态空间序列模型”的两种表达。

这把”Transformer vs Mamba”的二元对立—— 变成了”同一框架下的两种实现”。

论文的核心定理

定义新框架：State Space Duality (SSD)—— 状态空间对偶性。

任何可表达为”线性时不变状态空间方程”的模型—— 都同时有两种表达：

• “循环视角” （如 Mamba / SSM）
• “注意力视角” （如 Transformer）

两者：

• 训练时用注意力视角（并行化好）
• 推理时用循环视角（速度快、显存少）

这是计算图论上的优雅—— 让两个 community 合并。

Mamba 2 的关键改进

1. 简化 SSM 数学

原 Mamba 中 $A, B, C$ 矩阵复杂—— Mamba 2 用 scalar A + diagonal A —— 减少参数 + 更适合硬件。

2. 大幅提速

通过新算法 + 硬件优化：

• Selective scan: 4× 比 Mamba 1 快
• State expansion size: 8× 提升
• 显存使用更优

3. 模型质量提升

同尺寸下击败 Transformer—— 小尺度上的优势更明确。

那个核心思想：序列建模的统一观

论文最深的部分—— 所有”高效序列模型”都是 SSD 的特例：

• Transformer（标准 attention）
• Linear Attention
• Mamba / S4 / S5
• RWKV
• RetNet
• xLSTM

它们的差异：

• 不同的 “结构假设”（什么矩阵被允许）
• 不同的 “计算视角”（attention vs 循环）
• 不同的 “硬件优化”

但底层数学一致。

这是物理学般的”统一理论”—— 减少了”哪种架构更好”的争论。

论文之后

2024-2025 SSM 生态爆发

• Falcon-Mamba （UAE TII）—— 第一个商用 Mamba LLM
• Codestral Mamba（Mistral）—— 编程专用
• Zamba（Zyphra）—— Mamba + Attention 混合
• Jamba（AI21）—— MoE + Mamba 混合
• Granite Mamba（IBM）—— 企业级

工业接受度

2025 年 SSM 已经成为主流选项之一—— 但 Transformer 仍占主导。

实际部署：

• 文本任务：Transformer + Mamba 混合 增多
• 视频 / 长序列：Mamba 优势明显
• 代码 / 推理：Transformer 仍占优

一些有意思的细节

”Tri Dao + Albert Gu”

两人在系统 ML 上几乎”出圈了”：

• Tri Dao：FlashAttention + Mamba
• Albert Gu：S4 + Mamba 系列
• 合作多次——研究方向高度互补

他们代表了算法 + 硬件协同设计的范式。

ICML vs ICLR

Mamba 1 被 ICLR 拒了—— Mamba 2 被 ICML 接收。

12 个月间，社区从怀疑到接受—— 学术圈的”风向” 转变之快。

推荐配套阅读

• HelloAI: L3-04 RNN / LSTM + L3-05 Attention
• Mamba 2 论文 —— 数学硬核
• State Space Models series（Albert Gu）
• FlashAttention 2 / 3 —— 同作者