[Paper Review] 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 Measurements13 citations

TL;DR

Introduces ConvSSMs and ConvS5, a convolutional state-space approach that enables parallelized, long-range spatiotemporal modeling with linear-time per-step inference and fast autoregressive generation, outperforming Transformers and ConvLSTM on long-horizon tasks.

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.

Motivation & Objective

• Motivate the need for scalable modeling of long-range spatiotemporal sequences that capture local spatial structure and long temporal dependencies.
• Develop ConvSSMs that combine tensor states with linear, continuous-time dynamics to enable efficient parallelization.
• Propose ConvS5, an efficient ConvSSM variant with parallel scans and SSM-inspired initialization to model long-range dependencies.
• Demonstrate empirical performance gains over Transformers and ConvLSTM on long-horizon datasets and benchmarks.

Proposed method

• Define convolutional state space models (ConvSSMs) with tensor-valued states and linear convolutional dynamics.
• Show how to parallelize the ConvSSM recurrence across the sequence using a binary associative operator and parallel scans.
• Establish an equivalence between ConvSSMs and state-space models (SSMs) to leverage HiPPO-inspired initialization and discretization.
• Propose ConvS5 by parameterizing ConvSSMs with diagonalized SSM-like dynamics and applying parallel scans for efficient long-range modeling.
• Stack multiple ConvS5 layers with nonlinear activations to capture nonlinear spatiotemporal dynamics while maintaining tractable training.
• Provide empirical evidence on long-horizon Moving-MNIST, DMLab, Minecraft, and Habitat benchmarks, with fast autoregressive generation.

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

Experimental results

Research questions

• RQ1 Can ConvSSMs provide subquadratic parallelizable training and inference for long spatiotemporal sequences?
• RQ2 How can ConvSSMs be parameterized and initialized to capture long-range dependencies effectively?
• RQ3 Does ConvS5 outperform Transformer and ConvLSTM baselines on long-horizon video prediction benchmarks?
• RQ4 What are the computational trade-offs (speed, memory) of ConvS5 compared to Transformers and ConvRNNs?
• RQ5 Can ConvS5 scale to complex 3D environment benchmarks (DMLab, Minecraft, Habitat) while preserving quality?

Key findings

• ConvS5 enables parallel sequence processing with subquadratic complexity and fast autoregressive generation.
• ConvS5 trains faster than ConvLSTM and generates samples much faster than Transformers on long-horizon Moving-MNIST.
• ConvS5 matches or exceeds state-of-the-art on DMLab, Minecraft, and Habitat long-range video prediction benchmarks.
• ConvS5 benefits from a ConvSSM structure with an SSM-inspired parameterization and HiPPO-based initialization.
• The approach achieves strong performance without relying on token-based compression of frames.

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