SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jimmy T.H. Smith*,2,4, Shalini De Mello1, Jan Kautz1, Scott W. Linderman3,4, Wonmin Byeon1
$^{1}$ NVIDIA, *Work performed during internship at NVIDIA
$^{2}$ Institute for Computational and Mathematical Engineering, Stanford University.
$^{3}$ Department of Statistics, Stanford University.
$^{4}$ Wu Tsai Neurosciences Institute, Stanford University.
{jsmith14,scott.linderman}@stanford.edu
{shalinig, jkautz, wbyeon}@nvidia.com.
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  $(\mathrm{ConvSSM})^1$  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  $3\times$  faster than ConvLSTM and generating samples  $400\times$  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.
Developing methods that efficiently and effectively model long-range spatiotemporal dependencies is a challenging problem in machine learning. Whether predicting future video frames [1, 2], modeling traffic patterns [3, 4], or forecasting weather [5, 6], deep spatiotemporal modeling requires simultaneously capturing local spatial structure and long-range temporal dependencies. Although there has been progress in deep generative modeling of complex spatiotemporal data [7-12], most prior work has only considered short sequences of 20-50 timesteps due to the costs of processing long spatiotemporal sequences. Recent work has begun considering sequences of hundreds to thousands of timesteps [13-16]. As hardware and data collection of long spatiotemporal sequences continue to improve, new modeling approaches are required that scale efficiently with sequence length and effectively capture long-range dependencies.
<sup>1</sup>Implementation available at: https://github.com/NVlabs/ConvSSM.
Figure 1: ConvRNNs [17, 18] (left) model spatiotemporal sequences using tensor-valued states,  $\mathcal{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 states into the output sequence. Deep SSMs [19, 20] (center) model vector-valued input sequences using a discretized linear SSM. The linear dynamics can be exploited to parallelize computations across the sequence and capture long-range dependencies. We introduce ConvSSMs (right) that model spatiotemporal data using tensor states, like ConvRNNs, and linear dynamics, like SSMs. We also introduce an efficient ConvSSM variant, ConvS5, that can be parallelized across the sequence with parallel scans, has fast autoregressive generation, and captures long-range dependencies.
Convolutional recurrent networks (ConvRNNs) such as ConvLSTM [17] and ConvGRU [18] are common approaches for spatiotemporal modeling. These methods encode the spatial information using tensor-structured states. The states are updated with recurrent neural network (RNN) equations that use convolutions instead of the matrix-vector multiplications in standard RNNs (e.g., LSTM/GRUs [21, 22]). This approach allows the RNN states to reflect the spatial structure of the data while simultaneously capturing temporal dynamics. ConvRNNs inherit both the benefits and the weaknesses of RNNs: they allow fast, stateful autoregressive generation and an unbounded context window, but they are slow to train due to their inherently sequential structure and can suffer from the vanishing/exploding gradient problem [23].
Transformer-based methods [9, 13, 14, 24-27] operate on an entire sequence in parallel, avoiding these training challenges. Transformers typically require sophisticated compression schemes [28-30] to reduce the spatiotemporal sequence into tokens. Moreover, Transformers use an attention mechanism that has a bounded context window and whose computational complexity scales quadratically in sequence length for training and inference [31, 32]. More efficient Transformer methods improve on the complexity of attention [33-39], but these methods can fail on sequences with long-range dependencies [40, 13]. Some approaches combine Transformers with specialized training frameworks to address the attention costs [13]. However, recent work in deep state space models (SSMs) [19, 41, 42, 20, 43], like S4 [19] and S5 [20], has sought to overcome attention's quadratic complexity while maintaining the parallelizability and performance of attention and the statefulness of RNNs. These SSM layers have proven to be effective in various domains such as speech [44], images [45] and video classification [45, 46]; reinforcement learning [47, 48]; forecasting [49] and language modeling [50-53].
Inspired by modeling ideas from ConvRNNs and SSMs, we introduce convolutional state space models (ConvSSMs), which have a tensor-structured state like ConvRNNs but a continuous-time parameterization and linear state updates like SSM layers. See Figure 1. However, there are challenges to make this approach scalable and effective for modeling long-range spatiotemporal data. In this paper, we address these challenges and provide a rigorous framework that ensures both computational efficiency and modeling performance for spatiotemporal sequence modeling. First, we discuss computational efficiency and parallelization of ConvSSMs across the sequence for scalable training and fast inference. We show how to parallelize linear convolutional recurrences using a binary associative operator and demonstrate how this can be exploited to use parallel scans for subquadratic parallelization across the spatiotemporal sequence. We discuss both theoretical and practical considerations (Section 3.2) required to make this feasible and efficient. Next, we address how to capture long-range spatiotemporal dependencies. We develop a connection between
the dynamics of SSMs and ConvSSMs (Section 3.3) and leverage this, in Section 3.4, to introduce a parameterization and initialization design that can capture long-range spatiotemporal dependencies.
As a result, we introduce ConvS5, a new spatiotemporal layer that is an efficient ConvSSM variant. It is parallelizable and overcomes difficulties during training (e.g., vanishing/exploding gradient problems) that traditional ConvRNN approaches experience. ConvS5 does not require compressing frames into tokens and provides an unbounded context. It also provides fast (constant time and memory per step) autoregressive generation compared to Transformers. ConvS5 significantly outperforms Transformers and ConvLSTM on a challenging long horizon Moving-MNIST [54] experiment requiring methods to train on 600 frames and generate up to 1,200 frames. In addition, ConvS5 trains  $3 \times$  faster than ConvLSTM on this task and generates samples  $400 \times$  faster than the Transformer. Finally, we show that ConvS5 matches or exceeds the performance of various state-of-the-art methods on challenging DMLab, Minecraft, and Habitat long-range video prediction benchmarks [13].
This section provides the background necessary for ConvSSMs and ConvS5, introduced in Section 3.
2.1 Convolutional Recurrent Networks
Given a sequence of inputs  $\mathbf{u}_{1:L} \in \mathbb{R}^{L \times U}$ , an RNN updates its state,  $\mathbf{x}_k \in \mathbb{R}^P$ , using the state update equation  $\mathbf{x}_k = \mathbf{F}(\mathbf{x}_{k-1}, \mathbf{u_k})$ , where  $\mathbf{F}()$  is a nonlinear function. For example, a vanilla RNN can be represented (ignoring the bias term) as
$$
\mathbf {x} _ {k} = \tanh  \left(\mathbf {A} \mathbf {x} _ {k - 1} + \mathbf {B} \mathbf {u} _ {\mathbf {k}}\right) \tag {1}
$$
with state matrix  $\mathbf{A} \in \mathbb{R}^{P \times P}$ , input matrix  $\mathbf{B} \in \mathbb{R}^{P \times U}$  and  $\tanh()$  applied elementwise. Other RNNs such as LSTM [21] and GRU [22] utilize more intricate formulations of  $\mathbf{F}()$ .
Convolutional recurrent neural networks [17, 18] (ConvRNNs) are designed to model spatiotemporal sequences by replacing the vector-valued states and inputs of traditional RNNs with tensors and substituting matrix-vector multiplications with convolutions. Given a length  $L$  sequence of frames,  $\mathcal{U}_{1:L} \in \mathbb{R}^{L \times H' \times W' \times U}$ , with height  $H'$ , width  $W'$  and  $U$  features, a ConvRNN updates its state,  $\pmb{\mathcal{X}}_k \in \mathbb{R}^{H \times W \times P}$ , with a state update equation  $\pmb{\mathcal{X}}_k = \mathbf{G}(\pmb{\mathcal{X}}_{k-1}, \pmb{\mathcal{U}}_k)$ , where  $\mathbf{G}()$  is a nonlinear function. Analogous to (1), we can express the state update equation for a vanilla ConvRNN as
$$
\boldsymbol {\mathcal {X}} _ {k} = \tanh  (\boldsymbol {\mathcal {A}} * \boldsymbol {\mathcal {X}} _ {k - 1} + \boldsymbol {\mathcal {B}} * \boldsymbol {\mathcal {U}} _ {k}), \tag {2}
$$
where  $*$  is a spatial convolution operator with state kernel  $\mathcal{A} \in \mathbb{R}^{P \times P \times k_A \times k_A}$  (using an [output features, input features, kernel height, kernel width] convention), input kernel  $\mathcal{B} \in \mathbb{R}^{P \times U \times k_B \times k_B}$  and  $\tanh()$  is applied elementwise. More complex updates such as ConvLSTM [17] and ConvGRU [18] are commonly used by making similar changes to the LSTM and GRU equations, respectively.
2.2 Deep State Space Models
This section briefly introduces deep SSMs such as S4 [19] and S5 [20] designed for modeling long sequences. The ConvS5 approach we introduce in Section 3 extends these ideas to the spatiotemporal domain.
Linear State Space Models Given a continuous input signal  $\mathbf{u}(t) \in \mathbb{R}^U$ , a latent state  $\mathbf{x}(t) \in \mathbb{R}^P$  and an output signal  $\mathbf{y}(t) \in \mathbb{R}^M$ , a continuous-time, linear SSM is defined using a differential equation:
$$
\mathbf {x} ^ {\prime} (t) = \mathbf {A} \mathbf {x} (t) + \mathbf {B} \mathbf {u} (t), \quad \mathbf {y} (t) = \mathbf {C} \mathbf {x} (t) + \mathbf {D} \mathbf {u} (t), \tag {3}
$$
and is parameterized by a state matrix  $\mathbf{A} \in \mathbb{R}^{P \times P}$ , an input matrix  $\mathbf{B} \in \mathbb{R}^{P \times U}$ , an output matrix  $\mathbf{C} \in \mathbb{R}^{M \times P}$  and a feedthrough matrix  $\mathbf{D} \in \mathbb{R}^{M \times U}$ . Given a sequence,  $\mathbf{u}_{1:L} \in \mathbb{R}^{L \times U}$ , the SSM can be discretized to define a discrete-time SSM
$$
\mathbf {x} _ {k} = \overline {{\mathbf {A}}} \mathbf {x} _ {k - 1} + \overline {{\mathbf {B}}} \mathbf {u} _ {k}, \quad \mathbf {y} _ {k} = \mathbf {C} \mathbf {x} _ {k} + \mathbf {D} \mathbf {u} _ {k}, \tag {4}
$$
where the discrete-time parameters are a function of the continuous-time parameters and a timescale parameter,  $\Delta$ . We define  $\overline{\mathbf{A}} = \mathrm{DISCRETIZE}_{\mathrm{A}}(\mathbf{A}, \Delta)$  and  $\overline{\mathbf{B}} = \mathrm{DISCRETIZE}_{\mathrm{B}}(\mathbf{A}, \mathbf{B}, \Delta)$  where  $\mathrm{DISCRETIZE}()$  is a discretization method such as Euler, bilinear or zero-order hold [55].
S4 and S5 Gu et al. [19] introduced the structured state space sequence (S4) layer to efficiently model long sequences. An S4 layer uses many continuous-time linear SSMs, an explicit discretization step with learnable timescale parameters, and position-wise nonlinear activation functions applied to the SSM outputs. Smith et al. [20] showed that with several architecture changes, the approach could be simplified and made more flexible by just using one SSM as in (3) and utilizing parallel scans. SSM layers, such as S4 and S5, take advantage of the fact that linear dynamics can be parallelized with subquadratic complexity in the sequence length. They can also be run sequentially as stateful RNNs for fast autoregressive generation. While a single SSM layer such as S4 or S5 has only linear dynamics, the nonlinear activations applied to the SSM outputs allow representing nonlinear systems by stacking multiple SSM layers [56-58].
SSM Parameterization and Initialization Parameterization and initialization are crucial aspects that allow deep SSMs to capture long-range dependencies more effectively than prior attempts at linear RNNs [59-61]. The general setup includes continuous-time SSM parameters, explicit discretization with learnable timescale parameters, and state matrix initialization using structured matrices inspired by the HiPPO framework [62]. Prior research emphasizes the significance of these choices in achieving high performance on challenging long-range tasks [19, 20, 56, 57]. Recent work [57] has studied these parameterizations/initializations in more detail and provides insight into this setup's favorable initial eigenvalue distributions and normalization effects.
2.3 Parallel Scans
We briefly introduce parallel scans, as used by S5, since they are important for parallelizing the ConvS5 method we introduce in Section 3. See Blelloch [63] for a more detailed review. A scan operation, given a binary associative operator  $\bullet$  (i.e.  $(a\bullet b)\bullet c = a\bullet (b\bullet c))$  and a sequence of  $L$  elements  $[a_1,a_2,\dots ,a_L]$ , yields the sequence:  $[a_{1}$ $(a_{1}\bullet a_{2})$  ...  $(a_{1}\bullet a_{2}\bullet \ldots \bullet a_{L})]$ .
Parallel scans use the fact that associative operators can be computed in any order. A parallel scan can be defined for the linear recurrence of the state update in (4) by forming the initial scan tuples  $c_{k} = (c_{k,a},c_{k,b})\coloneqq (\overline{\mathbf{A}},\overline{\mathbf{B}}\mathbf{u}_{k})$  and utilizing a binary associative operator that takes two tuples  $q_{i},q_{j}$  (either the initial tuples  $c_{i},c_{j}$  or intermediate tuples) and produces a new tuple of the same type,  $q_{i}\bullet q_{j}\coloneqq (q_{j,a}\odot q_{i,a},q_{j,a}\otimes q_{i,b} + q_{j,b})$ , where  $\odot$  is matrix-matrix multiplication and  $\otimes$  is matrix-vector multiplication. Given sufficient processors, the parallel scan computes the linear recurrence of (4) in  $O(\log L)$  sequential steps (i.e., depth or span) [63].
This section introduces convolutional state space models (ConvSSMs). We show how ConvSSMs can be parallelized with parallel scans. We then connect the dynamics of ConvSSMs to SSMs to motivate parameterization. Finally, we use these insights to introduce an efficient ConvSSM variant, ConvS5.
3.1 Convolutional State Space Models
Consider a continuous tensor-valued input  $\mathcal{U}(t) \in \mathbb{R}^{H' \times W' \times U}$  with height  $H'$ , width  $W'$ , and number of input features  $U$ . We will define a continuous-time, linear convolutional state space model (ConvSSM) with state  $\mathcal{X}(t) \in \mathbb{R}^{H \times W \times P}$ , derivative  $\mathcal{X}'(t) \in \mathbb{R}^{H \times W \times P}$  and output  $\mathcal{Y}(t) \in \mathbb{R}^{H \times W \times U}$ , using a differential equation:
$$
\boldsymbol {\mathcal {X}} ^ {\prime} (t) = \boldsymbol {\mathcal {A}} * \boldsymbol {\mathcal {X}} (t) + \boldsymbol {\mathcal {B}} * \boldsymbol {\mathcal {U}} (t) \tag {5}
$$
$$
\boldsymbol {\mathcal {Y}} (t) = \boldsymbol {\mathcal {C}} * \boldsymbol {\mathcal {X}} (t) + \boldsymbol {\mathcal {D}} * \boldsymbol {\mathcal {U}} (t) \tag {6}
$$
where  $*$  denotes the convolution operator,  $\mathcal{A} \in \mathbb{R}^{P \times P \times k_A \times k_A}$  is the state kernel,  $\mathcal{B} \in \mathbb{R}^{P \times U \times k_B \times k_B}$  is the input kernel,  $\mathcal{C} \in \mathbb{R}^{U \times P \times k_C \times k_C}$  is the output kernel, and  $\mathcal{D} \in \mathbb{R}^{U \times U \times k_D \times k_D}$  is the feedthrough kernel. For simplicity, we pad the convolution to ensure the same spatial resolution,  $H \times W$ , is maintained in the states and outputs. Similarly, given a sequence of  $L$  inputs,  $\mathcal{U}_{1:L} \in \mathbb{R}^{L \times H' \times W' \times U}$ , we define a discrete-time convolutional state space model as
$$
\boldsymbol {\mathcal {X}} _ {k} = \overline {{\boldsymbol {\mathcal {A}}}} * \boldsymbol {\mathcal {X}} _ {k - 1} + \overline {{\boldsymbol {\mathcal {B}}}} * \boldsymbol {\mathcal {U}} _ {k} \tag {7}
$$
$$
\mathcal {Y} _ {k} = \mathcal {C} * \mathcal {X} _ {k} + \mathcal {D} * \mathcal {U} _ {k} \tag {8}
$$
where  $\overline{\mathcal{A}}\in \mathbb{R}^{P\times P\times k_A\times k_A}$  and  $\overline{\pmb{B}}\in \mathbb{R}^{P\times U\times k_B\times k_B}$  denote that these kernels are in discrete-time.
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,  $\mathbf{B}_{\mathrm{SSM}}$ , that can be formed by reshaping the ConvSSM's state kernel and input kernel. This allows leveraging parameterization insights from deep SSMs [19, 41, 42, 20, 57] to equip ConvS5 to model long-range dependencies.
3.2 Parallelizing Convolutional Recurrences
ConvS5 leverages parallel scans to efficiently parallelize the recurrence in (7). As discussed in Section 2.3, this requires a binary associative operator. Given that convolutions are associative, we show:
Proposition 1. Consider a convolutional recurrence as in (7) and define initial parallel scan elements  $c_k = (c_{k,a},c_{k,b})\coloneqq (\overline{\mathcal{A}},\overline{\mathcal{B}} *\mathcal{U}_k)$ . The binary operator  $\text{串}$ , defined below, is associative.
$$
q _ {i} * q _ {j} := \left(q _ {j, a} \circ q _ {i, a}, q _ {j, a} * q _ {i, b} + q _ {j, b}\right), \tag {9}
$$
where  $\circ$  denotes convolution of two kernels,  $^*$  denotes convolution and  $+$  is elementwise addition.
Therefore, in theory, we can use this binary operator with a parallel scan to compute the recurrence in (7). However, the binary operator,  $\otimes$ , requires convolving two  $k_{A} \times k_{A}$  resolution state kernels together. To maintain equivalence with the sequential scan, the resulting kernel will have resolution  $2k_{a} - 1 \times 2k_{a} - 1$ . This implies that the state kernel will grow during the parallel scan computations for general kernels with a resolution greater than  $1 \times 1$ . This allows the receptive field to grow in the time direction, a useful feature for capturing spatiotemporal context. However, this kernel growth is computationally infeasible for long sequences.
We address this challenge by taking further inspiration from deep SSMs. These methods opt for simple but computationally advantageous operations in the time direction (linear dynamics) and utilize more complex operations (nonlinear activations) in the depth direction of the model. These nonlinear activations allow a stack of SSM layers with linear dynamics to represent nonlinear systems. Analogously, we choose to use  $1 \times 1$  state kernels and perform pointwise state convolutions for the convolutional recurrence of (7). When we stack multiple layers of these ConvSSMs, the receptive field grows in the depth direction of the network and allows the stack of layers to capture the spatiotemporal context [64]. Computationally, we now have a construction that can be parallelized with subquadratic complexity with respect to the sequence length.
Proposition 2. Given the effective inputs  $\overline{\mathcal{B}} * \mathcal{U}_{1:L} \in \mathbb{R}^{L \times H \times W \times P}$  and a pointwise state kernel  $\mathcal{A} \in \mathbb{R}^{P \times P \times 1 \times 1}$ , the computational cost of computing the convolutional recurrence in Equation 7 with a parallel scan is  $\mathcal{O}\big(L(P^3 + P^2 HW)\big)$ .
Proof. See Appendix A.2.
Further, the ConvS5 implementation introduced below admits a diagonalized parameterization that reduces this cost to  $\mathcal{O}(LPHW)$ . See Section 3.4 and Appendix B for more details.
3.3 Connection to State Space Models
Since the convolutions in (5-6) and (7-8) are linear operations, they can be described equivalently as matrix-vector multiplications by flattening the input and state tensors into vectors and using large, circulant matrices consisting of the kernel elements [65]. Thus, any ConvSSM can be described as a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise state kernels, as described in the previous section, provides an alternative SSM equivalence, which lends a special structure that can leverage the deep SSM parameterization/initialization ideas discussed in Section 2.2 for modeling long-range dependencies. We show that each pixel of the state,  $\mathbf{x}(t)_{i,j} \in \mathbb{R}^P$ , can be equivalently described as evolving according to a differential equation with a shared state matrix,  $\mathbf{A}_{\mathrm{SSM}}$ , and input matrix,  $\mathbf{B}_{\mathrm{SSM}}$ . See Figure 2.
Proposition 3. Consider a ConvSSM state update as in (5) with pointwise state kernel  $\mathcal{A} \in \mathbb{R}^{P \times P \times 1 \times 1}$ , input kernel  $\mathcal{B} \in \mathbb{R}^{P \times U \times k_B \times k_B}$ , and input  $\mathcal{U}(t) \in \mathbb{R}^{H' \times W' \times U}$ . Let  $\mathcal{U}_{\mathrm{im2col}}(t) \in \mathbb{R}^{H \times W \times Uk_B^2}$  be the reshaped result of applying the Image to Column (im2col) [66, 67] operation on the input  $\mathcal{U}(t)$ . Then the dynamics of each state pixel of (5),  $\mathcal{X}(t)_{i,j} \in \mathbb{R}^P$ , evolve according to the following differential equation
$$
\boldsymbol {\mathcal {X}} ^ {\prime} (t) _ {i, j} = \mathbf {A} _ {\text {S S M}} \boldsymbol {\mathcal {X}} (t) _ {i, j} + \mathbf {B} _ {\text {S S M}} \boldsymbol {\mathcal {U}} _ {i m 2 c o l} (t) _ {i, j} \tag {10}
$$
where the state matrix,  $\mathbf{A}_{\mathrm{SSM}} \in \mathbb{R}^{P \times P}$ , and input matrix,  $\mathbf{B}_{\mathrm{SSM}} \in \mathbb{R}^{P \times (Uk_B^2)}$ , can be formed by reshaping the state kernel,  $\mathcal{A}$ , and input kernel,  $\mathcal{B}$ , respectively.
Proof. See Appendix A.3.
Thus, to endow these SSMs with the same favorable long-range dynamical properties as S4/S5 methods, we initialize  $\mathbf{A}_{\mathrm{SSM}}$  with a HiPPO [62] inspired matrix and discretize with a learnable timescale parameter to obtain  $\overline{\mathbf{A}}_{\mathrm{SSM}}$  and  $\overline{\mathbf{B}}_{\mathrm{SSM}}$ . Due to the equivalence of Proposition 3, we then reshape these matrices into the discrete ConvSSM state and input kernels of (7) to give the convolutional recurrence the same advantageous dynamical properties. We note that if the input, output and dynamics kernel widths are set to  $1\times 1$ , then the ConvSSM formulation is equivalent to "convolving" an SSM across each individual sequence of pixels in the spatiotemporal sequence (this also has connections to the temporal component of S4ND [45] when applied to videos). However, inspired by ConvRNNs, we observed improved performance when leveraging the more general convolutional structure the ConvSSM allows and increasing the input/output kernel sizes to allow local spatial features to be mixed in the dynamical system. See ablations discussed in Section 5.3.