Mamba状态空间模型原理

Mamba由Gu和Dao于2023-2024年提出，基于结构化状态空间模型（Structured State Space Sequence Model, S4），在长序列建模上展现出了超越Transformer的效率和潜力。

状态空间模型（SSM）核心原理： SSM用一组微分方程描述系统的输入-输出关系： h'(t) = A·h(t) + B·x(t) y(t) = C·h(t) + D·x(t)

• h(t)：隐状态。
• x(t)：输入序列。
• y(t)：输出序列。
• A, B, C, D：可学习的参数矩阵。

离散化：将连续SSM转化为离散形式： h_t = Ā·h_{t-1} + B̄·x_t y_t = C·h_t + D·x_t 其中Ā = exp(ΔA)，B̄ = (ΔA)^{-1}(exp(ΔA)-I)·ΔB，Δ为采样步长。

S4（Structured SSM）的突破：

• 使用HiPPO矩阵参数化A，实现对长程依赖的学习。
• 利用结构化矩阵实现O(N log N)的计算复杂度（卷积模式）。
• 但A, B, C是输入无关的（静态/时不变），限制了表达能力。

Mamba的选择性状态空间（Selective SSM）： 核心创新：让SSM参数(\Delta, B, C)成为输入的函数，实现输入依赖的动态选择。

1. 选择性机制：

   • Δ(xt) = softplus(Linear(xt))：控制每个token关注当前输入的程度。
   • B(xt)：控制输入如何进入状态。
   • C(xt)：控制状态如何映射到输出。
   • 实现了类似注意力的'选择关注'能力，但计算为线性复杂度。

2. 硬件感知并行扫描：

   • 使用并行扫描（Parallel Scan）算法计算SSM的前向传播。
   • 利用GPU的共享内存优化IO开销。
   • 达到与Transformer相当的训练速度，但推理时更高效。

Mamba块结构：

• Mamba Layer = 选择性SSM + 门控MLP + 残差连接 + 层归一化。
• 没有自注意力机制！

与Transformer对比： | 特性 | Transformer | Mamba | |------|-----------|-------| | 计算复杂度 | O(N²·d) | O(N·d·logN) | | 推理速度 | 随序列长度平方增长 | 线性增长 | | 长序列处理 | 受限于注意力窗口 | 理论上无限制 | | 训练稳定性 | 成熟（LayerNorm） | 需要SSM初始化技巧 | | 硬件效率 | 有高效实现（FlashAttention） | 同等或更好 |

潜在应用：

• 长文本建模（100K+ tokens）。
• 多模态序列（视频帧、音频波形）。
• 高效率的语言模型。
• Vision Mamba：将Mamba应用于图像任务。

局限性：

• 在较小规模下不如Transformer灵活。
• 未来发展：Mamba-2（更大的状态维度）+ 集成注意力机制。