[论文解读] Convolutional State Space Models for Long-Range Spatiotemporal Modeling

Jimmy T. H. Smith, Shalini De Mello+3 more|arXiv (Cornell University)|Oct 30, 2023

Geophysics and Gravity Measurements被引用 13

一句话总结

引入 ConvSSMs 和 ConvS5，一种卷积状态空间方法，能够实现并行化的长程时空建模，具有线性时间的每步推断和高效自回归生成，在长期任务上优于 Transformers 和 ConvLSTM。

ABSTRACT

Effectively modeling long spatiotemporal sequences is challenging due to the need to model complex spatial correlations and long-range temporal dependencies simultaneously. ConvLSTMs attempt to address this by updating tensor-valued states with recurrent neural networks, but their sequential computation makes them slow to train. In contrast, Transformers can process an entire spatiotemporal sequence, compressed into tokens, in parallel. However, the cost of attention scales quadratically in length, limiting their scalability to longer sequences. Here, we address the challenges of prior methods and introduce convolutional state space models (ConvSSM) that combine the tensor modeling ideas of ConvLSTM with the long sequence modeling approaches of state space methods such as S4 and S5. First, we demonstrate how parallel scans can be applied to convolutional recurrences to achieve subquadratic parallelization and fast autoregressive generation. We then establish an equivalence between the dynamics of ConvSSMs and SSMs, which motivates parameterization and initialization strategies for modeling long-range dependencies. The result is ConvS5, an efficient ConvSSM variant for long-range spatiotemporal modeling. ConvS5 significantly outperforms Transformers and ConvLSTM on a long horizon Moving-MNIST experiment while training 3X faster than ConvLSTM and generating samples 400X faster than Transformers. In addition, ConvS5 matches or exceeds the performance of state-of-the-art methods on challenging DMLab, Minecraft and Habitat prediction benchmarks and enables new directions for modeling long spatiotemporal sequences.

研究动机与目标

• 激发对可扩展建模长程时空序列的需求，这些序列能够捕捉局部空间结构和长时序依赖。
• 开发将张量状态与线性、连续时间动力学相结合的 ConvSSMs，以实现高效并行化。
• 提出 ConvS5，一种高效的 ConvSSM 变体，具备并行扫描和受 SSM 启发的初始化，以建模长程依赖。
• 在长时间序列数据集和基准测试上，展示相对于 Transformers 和 ConvLSTM 的实证性能提升。

提出的方法

• 定义带张量值状态和线性卷积动力学的卷积状态空间模型（ConvSSMs）。
• 展示如何使用二元结合运算符和并行扫描将 ConvSSM 的递推在序列上并行化。
• 建立 ConvSSMs 与状态空间模型 (SSMs) 之间的等价性，以利用 HiPPO 启发的初始化和离散化。
• 通过对 ConvSSMs 参数化为对角化的类似 SSM 的动力学并应用并行扫描来实现高效的长程建模，提出 ConvS5。
• 堆叠多个带有非线性激活的 ConvS5 层，以捕捉非线性的时空动力学，同时保持可控的训练难度。
• 在长时间跨度的 Moving-MNIST、DMLab、Minecraft 和 Habitat 基准上提供实证证据，且具备快速的自回归生成。

Figure 1 : ConvRNNs [ 17 , 18 ] (left) model spatiotemporal sequences using tensor-valued states, $\mathbfcal{X}_{k}$ , and a nonlinear RNN update, $\mathbf{G}()$ , that uses convolutions instead of matrix-vector multiplications. A position-wise nonlinear function, $\mathbf{h}()$ , transforms the st

实验结果

研究问题

• RQ1 ConvSSMs 是否能够提供对长程时空序列的亚二次级可并行训练与推断？
• RQ2 如何对 ConvSSMs 进行参数化和初始化，以有效捕获长程依赖？
• RQ3 在长时域视频预测基准上，ConvS5 是否优于 Transformer 和 ConvLSTM 基线？
• RQ4 与 Transformers 和 ConvRNNs 相比，ConvS5 的计算权衡（速度、内存）是什么？
• RQ5 ConvS5 是否能够扩展到复杂的三维环境基准（DMLab、Minecraft、Habitat）并保持质量？

主要发现

• ConvS5 实现了具有亚二次复杂度的并行序列处理和快速自回归生成。
• 在长期 Moving-MNIST 上，ConvS5 的训练速度快于 ConvLSTM，样本生成速度也远快于 Transformers。
• ConvS5 在 DMLab、Minecraft 和 Habitat 的长程视频预测基准上达到或超过最先进行的水平。
• ConvS5 受益于具备 SSM 启发参数化和 HiPPO 基于初始化的 ConvSSM 结构。
• 该方法在不依赖基于 token 的帧压缩的前提下实现了强性能。

Figure 2 : The dynamics of a ConvSSM with pointwise state kernel (top) can be equivalently viewed as the dynamics of an SSM (bottom). See Proposition 3 . Each ConvSSM state pixel evolves according to an SSM state update with shared state matrix, $\mathbf{A}_{\mathrm{SSM}}$ , and input matrix, $\math