NIPS 2023 | ConvSSM：卷积状态空间模型，线性复杂度并行化让长视频预测快400倍！

论文: Convolutional State Space Models for Long-Range Spatiotemporal Modeling
作者: Jimmy T.H. Smith, Shalini De Mello, Jan Kautz, Scott W. Linderman, Wonmin Byeon
发表: NeurIPS 2023
论文链接: https://arxiv.org/abs/2310.19694
代码链接: https://github.com/NVlabs/ConvSSM

一、引言

长程时空序列建模是机器学习的核心挑战之一。无论是视频预测、交通流建模还是天气预报，都需要同时捕捉局部空间结构和长程时间依赖。

现有方法各有瓶颈：

• ConvRNN（ConvLSTM/ConvGRU）：使用张量状态建模时空数据，但顺序计算导致训练极慢，且面临梯度消失/爆炸问题
• Transformer：可并行处理整个序列，但注意力复杂度随序列长度二次增长，推理时自回归生成极慢（每步需重新计算整个序列）
• S4/S5 等状态空间模型（SSM）：擅长长序列建模，但作用于向量化序列，缺乏对张量结构空间信息的建模能力

本文提出的ConvSSM（Convolutional State Space Model）结合了 ConvRNN 的张量状态建模和 SSM 的线性动力学，实现了以下突破：

1. 首次将并行扫描（Parallel Scan）应用于卷积循环，实现次二次复杂度（subquadratic）并行化训练
2. 建立了 ConvSSM 与 SSM 之间的数学等价性，从而利用 HiPPO 初始化等 SSM 核心技术实现长程依赖建模
3. 训练比 ConvLSTM 快 3 倍，自回归生成比 Transformer 快 400 倍，同时匹配或超越 SOTA

二、核心动机

ConvRNN 的致命伤：ConvLSTM 的张量更新使用非线性循环（tanh/sigmoid 门控），这种非线性顺序计算无法并行化。训练 300 帧序列在 8×V100 上需要 75 天，而 ConvS5 仅需 25 天。

关键洞察：如果将 ConvRNN 的非线性更新替换为线性卷积动力学，就可以利用并行扫描对序列进行并行化计算，同时保持张量结构。

解决方案：ConvSSM 使用线性卷积状态更新（类似 SSM 的线性 ODE），结合并行扫描实现训练并行化，同时通过 SSM 等价性利用 HiPPO 初始化和连续时间参数化来建模长程依赖。

三、方法

3.1 模块整体设计

ConvSSM 的核心是一个连续的、线性的卷积状态空间模型。给定连续时间输入U ( t ) ∈ R H ′ × W ′ × U \mathcal{U}(t) \in \mathbb{R}^{H' \times W' \times U}U(t)∈RH′×W′×U，状态X ( t ) ∈ R H × W × P \mathcal{X}(t) \in \mathbb{R}^{H \times W \times P}X(t)∈RH×W×P和输出Y ( t ) ∈ R H × W × U \mathcal{Y}(t) \in \mathbb{R}^{H \times W \times U}Y(t)∈RH×W×U，其微分方程为：

X ′ ( t ) = A ∗ X ( t ) + B ∗ U ( t ) (1) \mathcal{X}'(t) = \mathcal{A} * \mathcal{X}(t) + \mathcal{B} * \mathcal{U}(t) \tag{1}X′(t)=A∗X(t)+B∗U(t)(1)

Y ( t ) = C ∗ X ( t ) + D ∗ U ( t ) (2) \mathcal{Y}(t) = \mathcal{C} * \mathcal{X}(t) + \mathcal{D} * \mathcal{U}(t) \tag{2}Y(t)=C∗X(t)+D∗U(t)(2)

其中∗ *∗是卷积算子，A ∈ R P × P × k A × k A \mathcal{A} \in \mathbb{R}^{P \times P \times k_A \times k_A}A∈RP×P×kA​×kA​是状态核，B ∈ R P × U × k B × k B \mathcal{B} \in \mathbb{R}^{P \times U \times k_B \times k_B}B∈RP×U×kB​×kB​是输入核。

离散化后（零阶保持法 ZOH），得到离散时间 ConvSSM：

X k = A ˉ ∗ X k − 1 + B ˉ ∗ U k (3) \mathcal{X}_k = \bar{\mathcal{A}} * \mathcal{X}_{k-1} + \bar{\mathcal{B}} * \mathcal{U}_k \tag{3}Xk​=Aˉ∗Xk−1​+Bˉ∗Uk​(3)

Y k = C ∗ X k + D ∗ U k (4) \mathcal{Y}_k = \mathcal{C} * \mathcal{X}_k + \mathcal{D} * \mathcal{U}_k \tag{4}Yk​=C∗Xk​+D∗Uk​(4)

3.2 关键操作：并行扫描与 ConvSSM-SSM 等价性

并行化卷积循环

卷积操作具有结合律，因此可以定义二元结合算子⊛ \circledast⊛：

c i ⊛ c j : = ( c j , a ∘ c i , a , c j , a ∗ c i , b + c j , b ) (5) c_i \circledast c_j := (c_{j,a} \circ c_{i,a},\; c_{j,a} * c_{i,b} + c_{j,b}) \tag{5}ci​