什么是状态空间模型(State Space Model)?
状态空间模型(State Space Model,SSM)是一类用隐状态随序列推进的递归结构来描述序列的模型。Mamba 论文指出,此前这类模型在语言等重要模态上不如注意力机制,原因是无法做基于内容的推理;论文让 SSM 参数成为输入的函数,并设计了硬件感知的并行算法,从而在保持线性时间的同时提升效果。
状态空间模型(State Space Model,SSM)是一类把序列建模为「隐状态随时间步演化、并由隐状态产生输出」的递归结构的模型。它要解决的问题是:在处理长序列时,注意力机制的计算代价随序列长度呈二次增长,而状态空间模型希望以随序列长度线性增长的开销完成序列建模。
为什么重要
按照 Mamba 论文(arXiv:2312.00752)的叙述,当前支撑深度学习多数应用的基座模型几乎都建立在 Transformer 架构及其核心的注意力模块之上。为了缓解 Transformer 在长序列上的计算低效,研究者提出了多种次二次时间(subquadratic-time)的架构,包括线性注意力、门控卷积与循环模型,以及结构化状态空间模型(structured state space models,SSMs)。
但该论文同时指出,这些次二次架构在语言这类重要模态上并没有取得与注意力相当的表现。论文把这一差距归因于一个关键弱点:这类模型无法执行基于内容的推理(content-based reasoning)。也就是说,模型在沿序列推进时,无法根据当前 token 的内容决定该保留还是丢弃哪些信息。状态空间模型之所以重要,正是因为它试图在保留线性时间开销的同时,补上这一能力。
工作机制
Mamba 论文给出的改进可以拆成两个要点。
- 让 SSM 参数成为输入的函数。论文指出,仅仅让状态空间模型的参数依赖于输入,就能解决这类模型在离散模态上的弱点。这样一来,模型可以沿序列长度维度,依据当前 token 选择性地传播或遗忘信息。这种机制被称为选择性状态空间模型(selective SSM)。
- 硬件感知的并行算法。论文说明,让参数随输入变化会使得原本高效卷积形式不再可用,因此作者设计了一种在循环模式下运行的、硬件感知的并行算法,以恢复计算效率。
论文把上述选择性 SSM 集成进一个简化的端到端神经网络架构中,该架构不含注意力模块,甚至不含 MLP 模块,作者将其命名为 Mamba。据论文描述,Mamba 具备较快的推理速度(吞吐量比 Transformer 高 5 倍),在序列长度上呈线性扩展,并且在真实数据上可扩展到百万长度的序列。
从实现角度看,官方代码仓库 state-spaces/mamba 提供了名为 mamba-ssm 的软件包。该仓库文档说明,默认安装只安装核心包,不编译 selective_scan_cuda 扩展,也不下载带 CUDA 的预编译 wheel;若需要 CUDA 选择性扫描扩展,需通过环境变量显式开启。仓库列出的核心依赖要求为 Linux、Python 3.10+、PyTorch 1.12+;CUDA 选择性扫描构建与 GPU 执行还要求 NVIDIA GPU 与 CUDA 11.6+。
典型例子
论文中给出的具体结果包括:作为通用序列模型骨干,Mamba 在语言、音频、基因组学等多个模态上取得了当时的最优表现。在语言建模上,论文称 Mamba-3B 模型在同规模对比中优于 Transformer,并追平两倍于其规模的 Transformer,这一结论同时体现在预训练与下游评测上。
官方仓库还列出了与状态空间模型相关的后续工作条目,包括《Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality》以及《Mamba-3: Improved Sequence Modeling using State Space Principles Through Structured State Space Duality》。仓库说明,若要使用最新源码树中的 Mamba-3,需要从源码安装。
边界与常见误解
第一,容易把「状态空间模型」与「Mamba」画等号。按论文的表述,Mamba 是把选择性 SSM 集成进一个不含注意力与 MLP 模块的简化架构后得到的模型,选择性 SSM 是其中的机制,二者不是同一层级的概念。
第二,容易把线性时间开销理解为没有代价。论文明确指出,让 SSM 参数依赖输入会破坏原本可用的高效卷积形式,因此必须改用循环模式下的硬件感知并行算法来弥补;这类实现细节与硬件、编译环境绑定,官方仓库也把 CUDA 扩展列为需要显式开启的可选项。
第三,论文中关于性能的表述(例如吞吐量倍数、与同规模或两倍规模 Transformer 的对比、百万长度序列)都是该论文在其设定下报告的结果,属于论文作者的主张,不应被当作对所有任务、所有硬件与所有实现都成立的普遍结论。
第四,状态空间模型并非只适用于语言。论文将其定位为通用序列模型骨干,并在音频与基因组学等模态上做了评测;但具体到某一任务是否适用,仍取决于该任务的序列特性与实现条件。