状态空间模型(State Space Model,SSM)是一类源于控制理论与信号处理的序列建模方法,通过一个随时间演化的隐藏状态来压缩历史输入信息,并据此产生当前输出。与 Transformer 依赖全局两两比对注意力的思路不同,SSM 在每个时间步只更新固定维度的状态,因此计算复杂度随序列长度线性增长。自 2022 年 S4、2023 年 Mamba 相继提出后,SSM 凭借线性复杂度与恒定的推理内存,成为长序列建模领域备受关注的一类架构。
状态空间模型最初用于描述物理与工程系统在连续时间下的演化。给定一维输入信号 x(t)、输出信号 y(t) 和隐藏状态 h(t),连续形式的状态空间方程为:
其中 A 控制隐藏状态随时间的衰减与演化,B 控制输入写入状态的程度,C 将隐藏状态映射为输出,D 提供从输入到输出的直连通路。这四个矩阵共同决定系统如何"记忆"和"读取"信息。
计算机处理的是离散序列,需将连续方程离散化。采用零阶保持(ZOH)等方法、以步长 Δ 离散后,得到递推形式:
其中 Ā、B̄ 由连续参数 A、B 与步长 Δ 计算得到。离散化后,模型在每个时间步接收当前输入 x_t、结合上一时刻状态 h_{t-1},产生新状态 h_t 与输出 y_t,本质上是一个线性循环结构。
SSM 的核心思想是用一个固定维度的隐藏状态 h_t 概括截至当前时刻的全部历史信息。无论序列多长,状态维度始终保持不变,因此内存占用与序列长度无关。这与 RNN 的循环记忆思路相似,但 SSM 的线性结构与特定参数化方式带来了可并行训练和稳定长程依赖的优势。
基础状态空间模型虽已存在数十年,但在真实序列任务上表现不佳。根本问题在于状态矩阵 A 需要捕获长程依赖,而这难以通过梯度下降有效学习——朴素 SSM 在长序列任务上往往落后于 Transformer。
2022 年 Albert Gu 等人提出 S4(Structured State Space for Sequence Modeling),关键创新是用 HiPPO(High-Order Polynomial Projection Operators)框架初始化状态矩阵 A。HiPPO 提供了一种将长序列信息压缩进固定维度状态的最优多项式逼近方法,使状态矩阵从一开始就具备长程记忆能力,从而显著改善长程依赖建模。
S4 对 A 矩阵施加结构化约束(如低秩加对角形式),使计算复杂度从平方级降到线性级。同时,在线性时不变假设下,SSM 可等效为一个卷积,训练时可借助卷积与快速傅里叶变换并行加速。S4 在 Long Range Arena 等长程基准上取得了领先结果,证明了 SSM 在长序列任务上可与 Transformer 竞争。
在循环模式下,SSM 像 RNN 一样逐时间步处理:每一步仅依据当前输入和上一时刻状态更新,产生当前输出。其单步计算成本为 O(1),与序列总长无关,因此特别适合推理阶段的自回归生成。
在线性时不变(LTI)假设下,整个序列的 SSM 处理可视为一次大型卷积运算。由于系统对任意长度序列的响应是固定的,可借助卷积定理和快速傅里叶变换实现并行计算,复杂度约为 O(N log N)。这种模式适合训练阶段的大批量并行处理。
循环与卷积两种模式在数学上等价,可按场景切换:训练时用卷积模式获得并行加速,推理时用循环模式获得恒定单步成本。这种"两全其美"的对偶性是现代 SSM(如 S4)的核心优势。需要注意的是,引入选择性机制的 Mamba 不再严格满足线性时不变假设,因此无法直接使用卷积模式,需改用并行扫描算法。
朴素循环结构逐时间步串行计算,无法利用现代硬件的并行能力,训练效率低下。并行扫描(Parallel Scan)算法通过将递推分解为可结合的前缀和运算,使循环结构能够在硬件上并行执行。
Mamba 引入的选择性机制使参数依赖输入,破坏了线性时不变性,无法直接用卷积。为此,Mamba 采用选择性扫描算法:将序列分块,块内做二次(类注意力)计算以利用硬件效率,块间传递 SSM 状态以维持整体线性复杂度。该算法在保持线性扩展的同时恢复了并行训练能力。
Mamba 的实现进一步做了 IO 优化:不将庞大的状态矩阵写入 GPU 高带宽内存(HBM),而是在片上高速 SRAM 中完成扫描(recurrence)运算,大幅减少内存读写。官方实现报告,该高效扫描内核相比朴素的选择性循环实现可提速数十倍。
Transformer 的自注意力需要对序列中每个 token 与其他所有 token 计算关系,形成长度为 N 的 N×N 注意力矩阵,因此复杂度为 O(N²)。SSM 不做全局比对,而是逐 token 更新固定维度状态,每个 token 的计算成本恒定。
SSM 的隐藏状态维度不随序列长度增长。处理 N 个 token 需要 N 步,每步成本相同,总复杂度因此为线性的 O(N)。这意味着序列长度翻倍时,计算量只翻倍,而非像 Transformer 那样增长为四倍。
线性复杂度在超长序列上优势尤为明显。当序列达到百万 token 级别时,Transformer 的平方复杂度在现有硬件上几乎不可行,而 SSM 的线性扩展使其能够处理基因组、长音频、长文档等超长输入。
由于计算复杂度线性增长且状态维度固定,SSM 天然适合超长序列。Mamba 在实验中展现出对高达百万 token 上下文的有效处理能力,且随着上下文增长持续受益。
Transformer 推理需缓存随序列增长的 KV 矩阵,内存随长度线性增加;SSM 的推理状态尺寸固定,不随序列增长,因此在长上下文推理时内存占用显著更低。
借助 HiPPO 初始化与选择性机制,SSM 能够在数千个时间步的跨度上保持对关键信息的记忆。在 Long Range Arena 等专门测试长程依赖的基准上,SSM 在长序列任务中表现突出。
SSM 与 RNN 都是逐时间步更新隐藏状态的循环结构,都通过固定维度状态概括历史。这是二者表面的相似之处。
关键区别在于 SSM 是线性的(不含 tanh 等非线性激活),这一结构特性使其兼具并行训练能力与稳定的长程依赖。传统 RNN 因非线性与梯度问题,训练难以并行且长程依赖较弱。
现代 SSM(如 S4、Mamba)通过 HiPPO 初始化与选择性机制,使状态矩阵具备优化的长程记忆与内容感知能力。而朴素 RNN/LSTM 的记忆能力依赖门控结构,在超长序列上往往不如经过专门设计的 SSM。
Transformer 通过自注意力让每个 token 与序列中所有其他 token 直接比对,形成全局感受野,代价是 O(N²) 的平方复杂度。SSM 则通过逐 token 更新隐藏状态来压缩历史,不做全局两两比对,复杂度为线性的 O(N)。
Transformer 的注意力是输入相关的,能动态判断哪些 token 更重要。早期 SSM(如 S4)参数固定、对所有输入一视同仁,缺乏内容感知的选择性。Mamba 通过引入输入相关的选择性机制,弥补了这一短板,使模型能根据内容决定记住或遗忘哪些信息。
Transformer 在需要精确跨 token 比对的任务(如多跳问答、结构化检索)上更强;SSM 在超长序列、推理效率敏感的场景下更有优势。当前不少架构选择将二者结合,以兼顾各自长处。
2024 年发表的 Mamba-2 提出了结构化状态空间对偶(Structured State Space Duality,SSD)框架,从数学上揭示了 SSM 与注意力机制的深层联系:二者可视为同一结构化矩阵的不同分解形式。
SSD 框架证明,当状态矩阵为标量乘以单位矩阵(所有对角元素相同)时,SSM 递推在数学上等价于一种带 1-半可分因果掩码的掩码自注意力。这意味着同一次计算既可表达为 SSM 循环,也可表达为类注意力的矩阵乘法。
这一理论统一表明,SSM 与 Transformer 并非截然对立的两种范式,而是序列模型谱系上的不同取值点。当所有状态值取 1 时,注意力形式退化为标准因果线性注意力。该联系为设计和混合两种架构提供了理论基础。
Transformer 在自回归解码时,每生成一个新 token 都需基于缓存重新计算注意力,KV 缓存随序列增长,单步成本随上下文增大。SSM 每步只更新固定状态,单步成本恒定,与序列长度无关。
Transformer 推理需维护随序列线性增长的 KV 缓存,占用大量显存并受内存带宽制约。SSM 的推理状态尺寸固定,无需增长的缓存,因此内存读写更少。
Mamba 通过硬件感知算法优化 GPU 内存 IO,在片上 SRAM 完成扫描运算。在长序列场景下,Mamba 的生成吞吐量可达同规模 Transformer 的数倍。不过需注意,在较短序列(如 32K token 以下)上,由于 Transformer 拥有更成熟的硬件内核,实际耗时可能仍更快。
DNA 序列是极长的离散数据,碱基之间的依赖关系可跨越数千乃至数百万个碱基。Transformer 的平方复杂度使其难以处理这种超长序列,而 SSM 的线性扩展天然契合基因组建模需求。
Mamba 在基因组序列建模实验中展现出对长达百万 token 序列的有效处理,且随上下文增长持续获益。在一项区分五种大型猿类(相似度约 99%)的下游物种分类任务中,Mamba 利用超长上下文的能力表现尤为突出。
除物种分类外,SSM 还被用于基因表达、蛋白质序列等生物学任务的建模。其线性复杂度与长程记忆能力,使其成为处理基因组尺度数据的有力工具。
原始音频波形采样率高、持续时间长,是典型的高频长序列。Transformer 处理这类信号时平方开销巨大,而 SSM 的线性复杂度使长时音频处理变得可行。
SSM 源于信号处理与控制理论,对连续或类信号数据(如音频、传感器流)具有天然适配性。S4 在音频、时间序列等连续信号任务上表现尤为出色。
在语音识别、长音频转录等任务中,SSM 能够高效处理长时依赖,并以恒定状态控制内存占用。这使其在长音频场景下相较 Transformer 具备明显的效率优势。
时间序列预测(如高频交易、气象建模)高度依赖长历史上下文。SSM 能够高效捕获长程依赖,将长历史压缩进固定状态,为预测提供充分的时序信息。
在长窗口预测场景下,Transformer 的平方复杂度成为瓶颈,而 SSM 的线性扩展使其能够处理更长的历史窗口,同时保持较低的内存占用。
时间序列是典型的连续动态过程,与 SSM 的连续时间动力学本源高度契合。这使得 SSM 在时序预测任务中既有理论适配性,也有实际的效率优势。
Transformer 推理的 KV 缓存随序列增长,对内存受限设备(如笔记本、手机)构成压力。SSM 的推理状态尺寸固定且较小,不随序列增长,更适合在低内存设备上运行。
SSM 的线性复杂度意味着处理长序列时的算力需求可控,降低了边缘设备的计算与功耗负担。
得益于上述特性,SSM 被视为在资源受限设备上部署长上下文模型的重要方向之一。Mamba-3 进一步将状态尺寸压缩至 Mamba-2 的一半而保持同等困惑度,提升了端侧与低显存部署的性价比。
混合专家(Mixture of Experts,MoE)通过路由器为每个 token 只激活部分"专家"网络,在扩大参数规模的同时控制计算量,是千亿、万亿参数模型的常用架构。
将 Transformer 注意力层、Mamba(SSM)层与 MoE 层结合,可形成"三重混合"架构。AI21 Labs 的 Jamba 是较早的生产级代表:它在 72 层中交错排列 Transformer 与 Mamba 层,并引入 MoE 路由,支持 256K token 上下文,在长上下文基准上取得领先结果,同时保持接近更小纯注意力模型的推理效率。
NVIDIA 推出的 Nemotron-H 系列(涵盖 8B、47B、56B 参数规模)也是混合 Mamba-Transformer 架构的代表,验证了 SSM 与注意力、MoE 结合在企业级部署中的可行性。这类混合设计让 SSM 的高效长序列处理能力与注意力的精确检索、MoE 的参数效率互补。
SSM 将全部历史压缩进固定维度状态,这使其在需要精确、结构化检索长上下文中任意内容的任务上弱于 Transformer——后者可通过重访完整 KV 缓存实现精确召回。Mamba-3 在真实检索任务上"有竞争力但非全面领先",在半结构化信息抽取上仍弱于 Transformer。
早期 SSM 因实数状态、一阶离散化等限制,在状态跟踪、逻辑推理、精确位置感知等任务上明显弱于 Transformer。Mamba-2 在奇偶检测等状态跟踪任务上近乎随机水平,直到 Mamba-3 引入复值状态才将该任务准确率提升至接近满分。
在较短序列(如 32K token 以下)任务上,Transformer 凭借更成熟的硬件内核和社区生态,实际表现与效率仍占优。SSM 的工具链、训练稳定性与大规模预训练经验仍在持续完善中,当前更现实的路径是混合架构,而非完全替代 Transformer。