论文: 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 的线性动力学，实现了以下突破：
- 首次将并行扫描（Parallel Scan）应用于卷积循环，实现次二次复杂度（subquadratic）并行化训练
- 建立了 ConvSSM 与 SSM 之间的数学等价性，从而利用 HiPPO 初始化等 SSM 核心技术实现长程依赖建模
- 训练比 ConvLSTM 快 3 倍，自回归生成比 Transformer 快 400 倍，同时匹配或超越 SOTA
二、核心动机
ConvRNN 的致命伤：ConvLSTM 的张量更新使用非线性循环（tanh/sigmoid 门控），这种非线性顺序计算无法并行化。训练 300 帧序列在 8×V100 上需要 75 天，而 ConvS5 仅需 25 天。
关键洞察：如果将 ConvRNN 的非线性更新替换为线性卷积动力学，就可以利用并行扫描对序列进行并行化计算，同时保持张量结构。
解决方案：ConvSSM 使用线性卷积状态更新（类似 SSM 的线性 ODE），结合并行扫描实现训练并行化，同时通过 SSM 等价性利用 HiPPO 初始化和连续时间参数化来建模长程依赖。
三、方法
3.1 模块整体设计
图1：ConvRNN（左）、SSM（中）和 ConvSSM（右）对比。ConvSSM 使用张量状态（像 ConvRNN）和线性动力学（像 SSM）
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