Mamba 架构

概述

Mamba 是 Albert Gu 与 Tri Dao 于 2023 年提出的序列建模架构,核心是选择性状态空间模型(Selective SSM,即 S6)。它不使用注意力机制,以随输入内容动态变化的状态方程递推处理序列。

原理

传统 SSM 的参数与输入无关,难以做内容相关的信息筛选;Mamba 让状态转移和输入门控参数由当前输入决定,使模型能「选择性」记住或遗忘信息。配合面向 GPU 的并行扫描与核融合实现,训练随序列长度线性扩展,推理只需维护固定大小的状态,无 KV 缓存。

应用与影响

Mamba 在语言、音频、基因序列等任务上表现出与同规模 Transformer 相当的能力,长序列效率优势明显。后续的 Mamba-2、Jamba(Mamba 与注意力混合)等工作推动了「后 Transformer 架构」的探索,混合架构被多家大模型公司用于生产。