Mamba (状态空间模型)

概述

Mamba 是由 Albert Gu 和 Tri Dao 于 2023 年提出的基于结构化状态空间模型（Structured State Space Model, SSM）的序列建模架构。它通过引入选择性扫描（Selective Scan）机制，使 SSM 的参数能够根据输入动态变化，从而克服了先前 SSM 模型（如 S4）在内容依赖推理上的不足。Mamba 的核心优势在于实现了线性时间复杂度（$O(n)$），同时保持了与 Transformer 相当的建模能力，特别适合超长序列处理。

为什么需要 Mamba

Transformer 的 Self-Attention 计算复杂度为 $O(n^2)$，在处理长序列时面临严峻挑战：

核心原理

状态空间模型 (SSM) 基础

传统 SSM 通过以下方程将一维输入 $u(t)$ 映射到输出 $y(t)$：

$h^{\prime }(t)=Ah(t)+Bu(t)$ $y(t)=Ch(t)+Du(t)$

其中 $h(t)$ 是隐状态，$A,B,C,D$ 是参数矩阵。

离散化后（使用零阶保持 ZOH）：

$h_k=\bar{A}{h}_{k-1}+\bar{B}{u}_k$ $y_k=C{h}_k$

这类似于 RNN 的递推形式，但参数来自连续时间系统的离散化。

Mamba 的关键创新：选择性扫描

Mamba 的核心突破在于让 SSM 参数 $B,C$ 和离散化步长 $\Delta$ 成为输入的线性函数：

$B_t={\text{Linear}}_B(x_t)$ $C_t={\text{Linear}}_C(x_t)$ ${\Delta }_t=\text{Softplus}({\text{Linear}}_{\Delta }(x_t))$

这使得模型能够根据输入内容选择性地保留或遗忘信息，是 Mamba 超越 S4 等固定参数 SSM 的关键。

Mamba 块结构

Input -> [LayerNorm] -> [SSM (Selective Scan)] -> [SiLU] -> [Linear] -> Output
  |                                                                    ^
  |____________________________________________________________________|
                           (残差连接)

Mamba vs Transformer 对比

Mamba 家族演进

选择性扫描的关键作用

2026 年趋势

1. 混合架构：Jamba、Samba 等混合 Mamba + Transformer 的架构兴起
2. Mamba 在视觉中的应用：Vision Mamba (ViM) 在 CV 任务上挑战 ViT
3. Mamba 在基因组学中的应用：长 DNA 序列建模是 Mamba 的天然优势
4. Mamba + MoE：结合稀疏激活进一步降低计算成本
5. 尚不能完全替代 Transformer：Mamba 在上下文学习（In-Context Learning）和复杂推理任务上仍落后于 Transformer

相关笔记

• Transformer 架构 - Transformer 架构与 Mamba 的对比
• KV Cache (Key-Value 缓存) - Transformer 推理瓶颈与 Mamba 的替代方案
• 长上下文 (Long Context) — 深度解析 - 长上下文处理中 Mamba 的优势
• 位置编码 (Position Encoding) - Transformer 需要位置编码，Mamba 不需要

来源

• Gu & Dao (2023) - “Mamba: Linear-Time Sequence Modeling with Selective State Spaces”: https://arxiv.org/abs/2312.00752
• Dao & Gu (2024) - “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality” (Mamba-2): https://arxiv.org/abs/2405.21060
• Gu et al. (2022) - “Efficiently Modeling Long Sequences with Structured State Spaces” (S4): https://arxiv.org/abs/2111.00396
• Lieber et al. (2024) - “Jamba: A Hybrid Transformer-Mamba Language Model”: https://arxiv.org/abs/2403.19887