Convolutional State Space Models for Long-Range Spatiotemporal Modeling
Authors: Jimmy T.H. Smith
Shalini De Mello
Jan Kautz
Scott Linderman
Wonmin Byeon
Venue: NeurIPS 2023 poster (2023)
Keywords: spatiotemporal modeling
ConvLSTM
RNN
state spaces
SSM
S4
S5
long-range dependencies
video prediction
TL;DR: Introduces convolutional state space models that overcome weaknesses of both transformers and convolutional RNNs for long-range spatiotemporal modeling.

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 $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.

Convolutional State Space Models for
Long-Range Spatiotemporal Modeling
Jimmy T.H. Smith 224,, Shalini De Mello', Jan Kautz', Scott W. Linderman Wonmin Byeon
'NVIDIA, "Work performed during internship NVIDIA
Institute for Computational and Mathematical Engineering, Stanford University.
DDpartment of Statistics, Stanford University.
*Wu Tsai Neurosciences Institute, Stanford University.
[jsmith14, scott. linderman @stanford. edu (shalinig, jkautz, wbyeon] @nvidia, com.
Abstract
Effectively modeling long spatiotemporal sequences is challenging due the need to model complex spatial correlations and long-range temporal dependencies simul- taneously. ConvLSTMs attempt to address this by updating tensor-valued states with recurrent neural networks, but their sequential computation makes them slow train. In contrast, Transformers can process an entire spatiotemporal sequence, compressed into tokens, in parallel. However, the cost of attention scales quadrat- ically length, limiting their scalability 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 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 ConySSM 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 of-thof-thh-art methods challenging DMLab, Minecraft and Habitat prediction benchmarks and enables new directions for modeling long spatiotemporal sequences.
Moving-MNIST
1 Introduction
Developing methods that efficiently and effectively model long-range spatiotemporal dependencies is 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 the costs 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.
IImplementation available at: https: //ll:////ittt
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

W..F E RIxU
Uus -R. """""
24
Ui
X
Xk-
XL
X
G(A
Axp
Bu
X, + 24
Noniinear function:
h(
[x(+
In(
Linear function
ConVRNN
24
SSM
u'
ConvSSM (this paper) 24
Figure I: ConvRNNs [17. 18] (left) model spatiotemporal sequences using tensor-valued states, Xs. and nonlinear RNN update, GO, that uses convolutions instead of matrix-vector multipiccations. A position-wise nonlinear function, h(), transforms the states into the output sequence. Deep SSMs [19, 20] (center) model vector- valued input sequences using 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 mangeranee dependencies.
Convolutional recurrent networks (ConvRNNs) such ConvLSTM [17] and ConvGRU [18] 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 reflect the spatial structure of the data while simultancously 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/cxploding 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 bounded context window and whose computational complexity scales quadratically in sequence length for training and inference [2]. More efficient Transformer methods improve on the complexity of attention [33-39], but these methods can fail on sequences with long- rongrranee dependencies [40, 133. 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, 431, 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, introduce convolational state space models (ConySSMs), which have a tensor-structured state like ConyRNNs 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 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 connection between

---

the dynamics of SSMs and ConvSSMs (Section 3.3) and leverage this, in Section 3.4, to introduce parameterization and initialization design that can capture long-range spatiotemporal dependencies.
As a result, we introduce ConvS5, a new spatiotemporal layer that an efficient ConvSSM variant. It is parallelizable and overcomes difficulties during training (e.g.. vanishing/explodnng 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 x faster than ConvLSTM on this task and generates samples x00x faster than the Transformer. Finally, we show that ConvS5 matches or exceeds the performance of various state-of ffte-offthhe-rt methods on challenging DMLab, Minecraft, and Habitat long-range video prediction benchmarks [13].
2
Background
This section provides the background necessary for ConvSSMs and ConvS5, introduced in Section 3.
2.1
Convolutional Recurrent Networks
Given sequence of inputs U1:L E RLxU an RNN updates its state, X E Rº, using the state update equation X₁; = F uk), where F() is nonlinear function. For example, vanilla RNN can be represented (ignoring the bias term) as
X =
Buk)
(I)
with state matrix A E PxP, input matrix B E PxU and tanh() applied elementwise. Other RNNs such as LSTM [21] and GRU [22] utilize more intricate formulations of 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 length L sequence of frames, UisL E LL xW with height H', width W and U features, a ConvRNN updates its state, X E HxxPP, with state update equation X G(X 2x),, where G() is nonlinear function. Analogous to (I), we can express the state update equation for a vanilla ConvRNN as
X tanh( X 13 Ux),
where * is a spatial convolution operator with state kernel A R (using an [output features, input features, kernel height, kernel width] convention), input kernel B E xxpx and tanh() is applied elementwise. More complex updates such as ConvLSTM [17] and ConvGRU [18] 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.
Given a continuous input signal u(t) E R a latent state x(t) E R and an output signal yt) E RM, continuous-tima linear SSM is defined using a differential equation:
Linear State Space Models
x x((1= Ax(() + Bu(l),
y(1) Cx(1) Du((),
and is parameterized by state matrix A RPP. an input matrix B PxU output matrix
C e IR MPP and a feedthrough matrix D E IR Mx Given sequence, u1.z E RixU the SSM can be discretized define a discrete-time SSM
Ax + Bu:
Ys === Cx Dug,
(4)
where the discrete-time parameters are a function of the continuous time parameters and a timescale parameter, We define A = DISCRETIZE, (A, A) and B DISCRETIZEg A, B, A) where DISCRETIZE() is discretization method such Euler, bilinear zero-order hold [55].
3

---

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-timm 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 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 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].
stateful
SSM Parameterization and Initialization Parameterization and initialization are crucial aspects that allow deep SSMs capture long-range dependencies more effectively than prior attempts at linear RNNs [59-61]. The general setup includes mitin 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 iraameterzztiios/iiiiiiizziiiiii
in more detail and provides insight into this setup's '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 introduce in Section 3. See Blelloch [63] for a more detailed review. A scan operation, given a binary associative operator (i.e. (a o == a (b c)) and a sequence of L
elements [01, d2,,.., az], yields the sequence: [a1, (a1 02), .... (a1 0g
aL)].
Parallel scans use the fact that associative operators can be computed any order. A parallel sean can be defined for the linear recurrence of the state update in (4) by forming the initial scan tuples Cj = -Ck.o: Ck,b) := (A, BuR) and utilizing binary associative operator that takes two tuples qs. 95 (either the initial tuples C2 Cy or intermediate tuples) and produces new tuple of the same type, qi 9j == (9i.a @ qi.a: Gj.a @ qi + &z,b), where @ is matrix-matrix multiplication and @ -etrix-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].
3 Method
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 ConySSM variant, ConvS5.
3.1 Convolutional State Space Models
Consider a continuous tensor-valued input utt) E H'''''XUU with height H', width W", and number of input features U. We will define continuous-tmme, linear convolutional state space model (ConvSSM) with state (t) E IR HxxWxx xP derivative X'(t) E R xP' and output (t) e
xW XU using a differential equation:
x²( = A (t) + x (t
(5)
((t) C +Dt+
(6)
where denotes the convolution operator, A PPPPxx is the state kernel, B R XxUxkg is the input kernel, E x is the output kernel, and D E x is the feedthrough kernel. For simplicity, we pad the convolution ensure the same spatial resolu- tion, x W, is maintained in the states and outputs. Similarly, given a sequence of L in- we define a discrete-time convolutional state space model as
puts, 41 E RLXH xW'xU
X = A X B U4:
(7)
V: = C + Xx + D U4
(8)
where A E &PxxxkaxkA and B R Px U xka xke denote that these kernels are in discrete-time

---

State evolution of ConvSSM with pointwise state kemel
X (8)
*
X
+
B
*
(4
WXWXP
State pixels evolie according to SSM state update
(1)
AssM
+
BSSM
Px P
PX
Figure 2: The dynamics of a ConvSSM with pointwise state kernel (top) can be equivalently viewed the dynamies an SSM (bottom). See Proposition Each ConvSSM state pixel evolves according to an SSM state update with shared state matrix, AsSM. and input matrix, BSSM 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 binary associative operator. Given that convolutions are associative, we show:
Proposition 1. Consider convolutional recurrence as in (7) and define initial parallel scan elements Ck (Ck,o, Ck.b) == (A, B U. The binary operator &, defined below, is associative.
qs 9j = (gj,a c B.oa 9j.a qs + 95,6),
where o denotes convolution of two kernels, denotes convolution and + elementwise addition.
Proof. See Appendix A.I.
Therefore, in theory, we can use this binary operator with a parallel scan to compute the recurrence in (7). However, the binary operator, @, requires convolving two K æą resolution state kernels together. To maintain equivalence with the sequential scan, the resulting kernel will have resolution 2k. 1 2k. 1. This implies that the state kernel will grow during the parallel scan computations for general kernels with a resolution greater than 1 x 1. This allows the receptive field to grow the time direction, 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 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, choose to use 1 state kernels and perform pointwise state convolutions for the convolutional recurrence of (7). When we stack multiple layers 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 * U11L E R and a pointwise state kernel A E R PxXxxx the computational cost of computing the convolutional recurrence in Equation 7 with a parallel sean is (L + HW))
Proof. See Appendix A.2.
Further, the ConvS5 implementation introduced below admits a diagonalized parameterization that reduces this cost to O(LPHW). See Section 3,4 and Appendix B for more details.
5

---

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 -etrix-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 ceamettrization/iiiillazaaii Section 2.2 for modeling long-range dependencies. We show that each pixel of the state, X (t) E IR,, can be equivalently described as evolving according to a differential equation with a shared state matrix, AsM. and input matrix, BSSM- See Figure 2.
ideas discussed in
Proposition 3. Consider ConvSSM state update as in (5) with pointwise state kernel A Pxxxxxx input kernel B RPPUxxn xkv, and input Ul) E B xU, Let Vimzoal (1) RH be the reshaped result of applying the Image to Column (im2col) [66, 67] operation on the input (tt). Then the dynamics of each state pixel of (5), (() E IR,,, evolve according to the following differential equation
X' (1))s AA) + Bsss
(10)
where the state matrix, AssM RPXP, and input matrix, BSSM E ()UX)))), can be formed by reshaping the state kernel, A. and input kernel, B, respectively.
Proof. See Appendix A.3.
Thus, to endow these SSMs with the same favorable long-range dynamical properties as S4/S5 methods, initialize ASSM with HiPPO [62] inspired matrix and discretize with learnable timescale parameter to obtain ASSM and BASM- Due to the equivalence of Proposition 3, 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 x 1, then the ConvSSM formulation is equivalent to "convolving" 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.
3.4 Efficient ConvSSM for Long- Range Dependencies: ConvS5
Here, we introduce ConvS5, which combines ideas of parallelization of convolutional recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that leverages parallel schemes. Given Proposition 3, we implicitly
seans and deep SSM parameterizaaton//ittiiiiizn parameterize a pointwise state kernel, A E RR Pxpx1x1 and input kernel B E PxUxxx x in (5)
using SSM parameters as used by S5 [20], Ass E PxP and Bss E (UU1))) We discretize these S5 SSM parameters as discussed in Section 2.2 to give
Ass DISCRETIZE (Ass, A)
Bsg DISCRETIZEg (Ass. Bs,, A),
(11)
and then reshape to give the ConvS5 state update kernels:
Ass E PxP reshape AconvS5 E R PxPxix1
(12)
Bss E PX(UA)) reshape,
B CouvvS E R PxU
(13)
We then run the discretized ConvSSM system (7- 8), using parallel scans to compute the recurrence. In practice, this setup allows us to parameterize ConvS5 using a diagonalized parameterization [41, 42, 20 which reduces the cost of applying the parallel scan in Proposition 2 to OLLPHW). See Appendix B for a more detailed discussion of parameterization, initialization and discretization.
We define a ConvS5 layer as the combination of ConvS5 with nonlinear function applied to the ConvS5 outputs. For example, for the experiments in this paper, we use ResNet[68] blocks for the nonlinear activations between layers. However, this is modular, and other choices such as ConvNext [69] or S4ND [45] blocks could easily be used as well. Finally, many ConvS5 layers can be stacked to form deep spatiotemporal sequence model.

---

Table 1: Computational complexity of Transformers, ConvRNNS and ConvS5 with respect to sequence length. Metrics are inference cost (cost per step of autoregressive generation), cost per training step, and parallelization ability. ConvS5 combines the best of Transformers and ConvRNNs.
Transformer
ConvRNNs
ConvS5
Inference
(L)
(1)
O(1)
Training
O(L*)
O(L)
O(L)
Parallclizable
Yes
No
Yes
3.5 ConvS5 Properties
Refer to Table 1 for comparison of computational complexity for Transformers, ConvRNNs and ConvS5. The parallel scans allow ConvS5 to be parallelized across the sequence like a Transformer but with cost only scaling linearly with the sequence length. In addition, ConvS5 is stateful model like Conv allowing fast autoregressive generation, extrapolation to different sequence lengths, and unbounded context window. The connection with SSMs such as S5 derived in Section 3.3 allows for precisely specifying the initial dynamics to enable the modeling of long-range dependencies and to realize benefits from the unbounded context. Finally, the parallel scan of ConvS5 can allow parameterization to process irregularly sampled sequences in parallel.
leveraging the continuous- S5 achieves this by providing suitable spacing to the discretization operation [20]. procedure that ConvS5 can also use.
We have proposed general ConvSSM structure that can be easily adapted to future innovations in the deep SSM literature. ConvS5's parallel sean could be used to endow ConvS5 with time-varying dynamics such the input-dependett dynamics shown to be beneficial in the Liquid-S4 [43] work. Multiple works [51, 52, 50, 53] proposed adding multiplicative gating to allow SSM-based methods to overcome some weaknesses language modeling. Similar ideas could be useful to ConvSSMs for reasoning over long spatiotemporal sequences.
Related Work
This work is most closely related to the ConvRNNs and deep SSMs already discussed. We note here that ConvRNNs have been used in numerous domains, including biomedical, robotics, traffic modeling and weather forecasting [70, 1, 71, 2, 3, 72-75, 4, 76, 77]. In addition, many variants have been proposed [78-84, 64]. SSMs have been considered previously for video classification [46, 45, 85], however none addressed the challenging problem of generating long spatiotemporal sequences. S4ND [45] proposed applying separate S4 layers to each axis of an image or video, similar to a PixeIRNN [86] approach, and then combining the results with outer product. While that work mostly focused on images, they also show results for a short 30-frame video classification task. We note this approach could be complementary to replace some of the ResNet blocks used as activations in our model. Other model architectures explored in the literature for spatiotemporal modeling include 31) convolution approaches [87-89], transformers, [9, 13, 14. 24-27] and standard RNNs [16, 90]. Attempts to address the problem of modeling long spatiotemporal sequences have involved compressed representations [9, 91, 10, 92, 27, 29], training on sparse subsets of frames [15, 93, 8, 94, 95], temporal hierarchies [16], continuous-tin time neural representations [93, 95], training on different length sequences at different resolutions [96] and strided sampling [14, 97]. Of particular interest, Yan et al. [13] introduced the Temporally Consistent (TECO) Video Transformer, which achieved state-of-the -aet-hh-att performance on challenging 3D environment benchmarks designed to contain long-range dependencies
ours and the S42D blocks could potentially be used
[13]. This approach combines vector-quantized (VQ) latent dynamics with MaskGit [98] dynamics prior, and several training tricks to model long videos.
5 Experiments
In Section 5. we present a long-horizon Moving-MNIST experiment to compare ConvRNNs, Transformers and ConvS5 directly. In Section 5.2, we evaluate ConvS5 on the challenging 3D
7

---

Table Quantitative evaluation on the Moving-MNIST dataset [54]. Top: evaluate, we condition on 100 frames, and then show results after generating 800 and 1200 frames. An expanded Table 6 is included in Appendix C with more results, error bars and ablations. Bold scores indicate the best performance and underlined scores indicate the second best performance. Bottom: Computational cost comparison for the 600 frame task. Compare Table
Trained 300 frames
100 + 800
100 + 1200
Method
FVD
PSNR
SSIM
LPIPS
FVD
PSNR
SSIM
LPIPS
Transformer [24]
159
12.6
0.609
0.287
265
12.4
0.591
0.321
Performer [33]
234
13.4
0.652
0.379
275
13.2
0.592
0.393
CW-VAE [16]
104
12.4
0.592
0.277
117
12.3
0.585
0.286
ConVLSTM [17]
128
15.0
0.737
0.169
187
14.1
0.706
0.203
ConvS5
72
16.0
0.761
0.156
187
14.5
0.678
0.230
Trained on 600 frames
Transformer
42
13.7
0.672
0.207
91
13.1
0.631
0.252
Performer
93
12.4
0.616
0.274
243
12.2
0.608
0.312
CW-VAE
94
12.5
0.598
0.269
107
12.3
0.590
0.280
ConVLSTM
15.5
0.757
0.149
137
14.6
0.727
0,180
ConvS5
47
16.4
0.788
0.134
71
15.6
0.763
0.162
GFLOPS
Train Step Time (s)
Train Cost (V100 days)
Sample Throughput (frames/s) f
Transformer
70
0.77(1.0x)
50
0.21 (1.0x)
ConvLSTM
65
3.0(3.9x)
150
1 117 (557x)
ConvS5
97
0193(
50
90 (429x)
environment benchmarks proposed in Yan et al. [13]. Finally, in Section 5.3, discuss ablations that highlight the importance of ConvS5's parameterization.
5.1 Long Horizon Moving-MNIST Generation
There few existing benchmarks for training on and generating long spatiotemporal sequences. We develop a long-horizon Moving-MNIST [54] prediction task that requires training on 300-600 frames and accurately generating up 1200 frames. This allows for direct comparison of ConvS5, ConvRNNs and Transformers as well as an efficient attention alternative (Performer [33]) and CW- VAE [16], temporally hierarchical RNN based method. We first train all models on 300 frames and then evaluate by conditioning on 100 frames before generating 1200. We repeat the evaluation after training on 600 frames. See Appendix for more experiment details. We present the results after generating 800 and 1200 frames in Table 2. See Appendix C for randomly selected sample trajectories. ConvS5 achieves the best overall performance. When only trained on 300 frames, ConvLSTM and ConvS5 perform similarly when generating 1200 frames, and both outperform the Transformer. All methods benefit from training on the longer 600-frame sequence. However, the longer training length allows ConvS5 to significantly outperform the other methods across the metrics when generating 1200 frames.
In Table 2-bottom we revisit the theoretical properties of Table and compare the empirical compu- tational costs of the Transformer, ConvLSTM and ConvS5 on the 600 frame Moving-MNIST task. Although this specific ConvS5 configuration requires few more FLOPs due to the convolution com- putations, ConvS5 is parallelizable during training (unlike ConyLSTVV and has fast autoregressive generation (unlike Transformer) - training 3x faster than ConvLSTM and generating samples 400x faster than Transformers.
5.2 Long-range 3D Environment Benchmarks
Yan et al. [13] introduced challenging video prediction benchmark specifically designed to con- tain long-range dependencies. This one the first comprehensive benchmarks for long-range spatiotemporal modeling and consists of 300 frame videos of agents randomly traversing 3D environ-

---

Table Quantitative evaluation on the DMLab long-range benchmark [13]. Results from Yan et al. [13] are indicated with x, Methods trained using the TECO [13] training framework are at the bottom of the table. TECO methods are slower sample due to the MaskGit [98] procedure. The expanded
Table 8 in Appendix C includes eroor bars and ablations.
DMLab
Method
FVD
PSNR
SSIM 1
LPIPS
Sample Throughpul (frames/s) 1
FitVid [90]
176
12.0
0.356
0.491
CW-VAE" [16]
125
12.6
0.372
0.465
Perceiver AR [39]
96
11.2
0.304
0.487
Latent FDM [15]
181
17.8
0.588
0.222
Transformer [24]
97
19.9
0.619
0.123
9 (1.0x)
Performer [33]
80
17.3
0.513
0.205
7 (0.8x)
S5 [20]
221
19.3
0.641
0.162
28 (3.1x)
ConvS5
66
23.2
0.769
0.079
56 (6.2x)
IECO-Transformes [13]
28
22.4
0.709
0.155
16 (1.8x)
TECO-Transformer (our run)
28
21.6
0.696
0.082
16 (1.8x)
TECO-S5
35
20.1
0.687
0.143
21 (2.3x)
TECO-ConvSS
31
23.8
0.803
0.085
18 (2.0x)
ments in DMLab [99], Minecraft [100], and Habitat [101] environments. See Appendix C for more experimental details and Appendix E for more details on each dataset.
We train models using the same 16 x 16 vector-quantized (VQ) codes from the pretrained VQ- GANs [30] used for TECO and the other baselines in Yan et al. [13]. In addition to ConvS5 and the existing baselines, we also train a Transformer (without the TECO framework), Performer and S5. ablation on ConvS5's convolutional tensor-structured approach. Finally, since TECO is essentially a training framework (specialized for Transformers), we also use ConvS5 and S5 layers original version of TECO as TrcO--ransformer, the ConvS5 version as TECO-ComsST and the S5 version TECO-S5. See Appendix D for detailed information training procedures, architectures, and hyperparameters.
The S5 baseline serves
drop-in replacement for the Transformer in TECO. Therefore, refer to the
The results for DMLab are presented in Table 3. Of the methods trained without the TECO framework in the top section of Table ConvS5 outperforms all baselines, including RNN [90, 16), efficient attention [39, 33] and diffusion [15] approaches. ConvS5 also has much faster autoregressive generation than the Transformer. ConvS5 significantly outperforms S5 all metrics, pointing the value of the convolutional structure of ConvS5.
DMLab
For the models trained with the TECO framework, the same FVD and LPIPS as TrCo-Traasformer, while significantly improving PSNR and SSIM. Note the sample speed comparisons procedure is relatively slow. Still, the sample throughput of TECO-ConvS5 and TECO-S5 remains constant, while SSCO-Transformer's throughput decreases with sequence length.
that TECO-ConvSs achieves essentially
less dramatic in this setting since the MaskGit [98] sampling
Minecraft and Habitat Table 4 presents the results on the Minecraft and Habitat benchmarks. On Minecraft, TECO-ConvS5 achieves the best FVD and performs comparably to TECO-Transformer on the other metries, intperfoTarming all other baselines. On Habitat, TECO-ConvS5 is the only method to achieve comparable FVD to TECO- TrCn-Tranforrer,, while outperforming it on PSNR and SSIM.
5.3 ConvS5 ablations
In Table 5 we present ablations on the convolutional structure of ConvS5. We compare different input and output kernel sizes for the ConvSSM and also compare the default ResNet activations a channel mixing GLU [102] activation. Where possible, when reducing the sizes of the ConvSSM kernels, we redistribute parameters to the ResNet kernels or the GLU sizes to compare similar parameter counts. The results suggest more convolutional structure improves performance.

---

Table Quantitative evaluation on the Minecraft and Habitat long-range benchmarks [13]. Results from Yan et al. [13] are indicated with x. See expanded Table 10 with error bars in Appendix C.
Minecraft
Habitat
Method
FVD +
PSNR
SSIM 1
LPIPS + FVD
PSNR 1 SSIM f LPIPS +
FitVid
956
13.0
0.343
0.519
CW-VAE"
397
13.4
0.338
0,441
Perceiver AR*
76
13.2
0.323
0.441
164
12.8
0.405
0.676
Latent FDM*
167
13.4
0.349
0.429
433
12.5
0.311
0.582
RECO-Transformer
116
15.4
0.381
0.340
76
12.8
0.363
0.604
TECO-ConvS5
71
14.8
0.374
0.355
95
12.9
0.390
0.632
Table 5: Ablations of ConvS5 convolutional structure for DMLaa long-range henchmark dataset [13].
More convolutional structure improves overall performance. See expanded Table 9 in Appendix.
DMLab
conv.
B kernel
C kernel
nonlinearity
FVD 4
PSNR
SSIM t
LPIPS +
GLU
221
19.3
0.641
0.162
1x 1
1 X 1
GLU
187
21.0
0.689
0.112
x 1
5
GLU
89
21.5
0.713
0.106
3 x 3
3 3 3
GLU
96
22.7
0.762
0.088
1x1
1x1
ResNet
81
23.0
0.767
0.083
1 1
x
ResNet
68
22.8
0.756
0.085
3
3
ResNet
67
23.2
0.769
0.079
We also perform ablations to evaluate the importance of ConvS5's deep SSM-inspired parameter- rration/iitiaaizztion. We evaluate the performance kernel on both Moving-Mnist and DMLab. See Table and Table 8 in Appendix C. We observe a degradation in performance in all settings for this ablation. This reflects prior results for deep SSMs [56, 19, 20, 57] and highlights the importance of the connection developed in Section 3.3.
ConvSSM with randomly initialized state
Discussion
This work introduces ConvS5, a new spatiotemporal modeling layer that combines the fast, stateful autoregressive generation of ConvRNNs with the ability be parallelized across the sequence like Transformers. Its computational cost scales linearly with the sequence length, providing better scaling for longer spatiotemporal sequences. ConvS5 also leverages insights from deep SSMs to model long-range dependencies effectively.
We note that despite the ConvS5's sub-quadratic scaling, it did not show significant training speedups over Transformers for sequence lengths of 300-600 frames. (See detailed run- un-timee comparisons in Appendix C.) We expect ConvS5 to excel in training efficiency when applied to much longer spatiotemporal sequences where the quadratic scaling of Transformers dominates. We hope this work inspires the creation of longer spatiotemporal datasets and benchmarks. At the current sequence lengths, future optimizations of the parallel scan implementations in common deep learning frame- works will be helpful. In addition, the ResNet blocks used as the activations between ConvS5 layers could be replaced with efficient activations such as sparse convolutions [103] or S4ND [45].
An interesting future direction is to further utilize ConvS5's continuous-time parameterization. Deep SSMs show strong performance when trained one resolution and evaluated on another [19, 41, 42, 20, 43, 45]. In addition, S5 can leverage its parallel scan to effectively model irregularly sampled sequences [20]. ConvS5 can used for such applications in the spatiotemporal domain [104, 93, 95]. Finally, ConvS5 is modular and flexible. We have demonstrated that ConvS5 works well drop-in replacement in the TECO [13] training framework specifically developed for Transformers. Due to its favorable properties, expect ConvS5 to also serve building block of new approaches for modeling much longer spatiotemporal sequences and multimodal applications.
10

---

References
[1] Chelsea Finn, lan Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. Advances in neural information processing systems, 29, 2016.
[2] Yi Xu, Longwen Gao, Kai Tian, Shuigeng Zhou, and Huyang Sun. Non-local ConvLSTM for video compression artifact reduction. In Proceedings of the IEEE/CVF international conference on computer vision, pages 7043-7052, 2019.
[3] He Huang, Zheni Zeng, Danya Yao, Xin Pei, and Yi Zhang. Spatial-temports ConvLSTM for vehicle driving intention prediction. Tsinghua Science and Technology. 27(3):599 60999660090000 2021.
[4] Xiaoyu Chen, Xingsheng Xie. and Da Teng. Short-term traffic flow prediction based on ConvLSTM model. In 2020 IEEE 5th Information Technology and Mechatronics Engineering Conference (ITOEC), pages 846-850. IEEE, 2020.
[5] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopad- hyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcastnet: A global data-driven high-resolution weather model using adaptive Fourier neural operators. arXiv preprint arXiv:2202. 2022.
[6] Jonathan A Weyn, Dale R Durran, Rich Caruana, and Nathaniel Cresswell-Clay, Sub-seasonal forecasting with a large ensemble of deep-learning weather prediction models. Journal of Advances in Modeling Earth Systems, 13(7)=2002 MS002502, 2021.
Jonathan Ho, Tim Salimans, Alexey A. Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffusion models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[8] Aidan Clark, Jeff Donahue, and Karen Simonyan. Adversarial video generation on complex datasets. arXiv preprint 1907: (9109990077771 2019.
[9] Wilson Yan, Yunzhi Zhang. Pieter Abbeel, and Aravind Srinivas. Videogpt: Video generation using VQ-VAL and Transformers. arXiv preprint arXiv:2104. 10157, 2021.
[10] Guillaume Le Moing, Jean Ponce, and Cordelia Schmid. CCVS: context-aware controllable video synthesis. Advances in Neural Information Processing Systems, 34:14042-14055, 2021.
[11] Yu Tian, Jian Ren, Menglei Chai, Kyle Olszewski, Xi Peng, Dimitris N. Metaxas, and Sergey Tulyakov. A good image generator what you need for high-resolution video synthesis. In International Conference on Learning Representations, 2021.
[12] Pauline Luc, Aidan Clark, Sander Dieleman, Diego de Las Casas, Yotam Doron, Albin Cassirer, and Karen Simonyan. Transformation-basee adversarial video prediction large-scale data. arXiv preprint arXiv:2003.04033: 2020.
[13] Wilson Yan, Danijar Hafner, Stephen James, and Pieter Abbeel. Temporally consistent video Transformer for long-term video prediction. arXiv preprint arXiv:2210.023300 2022.
[14] Songwei Ge, Thomas Hayes, Harry Yang, Xi Yin, Guan Pang. David Jacobs, Jia-Bin Huang, and Devi Parikh. Long video generation with time-agnostic VQGAN and time-sensitive Transformer. Computer ESCon 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XVII, pages 102-118. Springer, 2022.
[15] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Dietrich Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[16] Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. Ad- vances in Neural Information Processing Systems, 1429224-22225 2021.
11

---

[17] Xingjian Shi, Zhourong Chen, Hao Wang. Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional LSTM network: machine learning approach for precipitation nowcast- ing. Advances in neural information processing systems, 28, 2015.
[18] Nicolas Ballas, Li Yao, Chris Pal, and Aaron Courville. Delving deeper into convolutional networks for learning video representations. arXiv preprint arXiv: 111111111.22. 2015.
[19] Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2021.
[20] Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.
[21] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, (8):1735-1780, 1997.
[22] Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using mn encoder- decoder for statistical machine translation. arXiv preprint arxXiv:1406. 1078, 2014.
[23] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrett neural networks. In International Conference on Machine Learning, pages 1310-1318. PMLR, 2013.
[24] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[25] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[26] Anurag Arnb, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. ViVIT: A video vision Transformer. In Proceedings of the IEEE/CVF international conference on computer vision, pages 6836-6846, 2021.
[27] Agrim Gupta, Stephen Tian, Yunzhi Zhang, Jiajun Wu, Roberto Martín-Martin and Li Fei-Fei. Maskvit: Masked visual pre-training for video prediction. arXiv preprint arXiv:2206. / 1894, 2022.
[28] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
[29] Jacob Walker, Ali Razavi, and Aäron van den Oord. Predicting video with VQVAE. arXiv preprint 1103::::00309900000 2021.
[30] Patrick Esser, Robin Rombach, and Bjor Ommer. Taming Transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12873- 12883, 2021.
[31] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling Transformer inference. arXiv preprint arXiv:2211.05 2022.
[32] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with 1O-awareness. Advances in Neural Information Processing Systems, 35:16344-16335 2022.
[33] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with Performers. In International Conference on Learning Representations, 2021.
12

---

[34] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs:: Fast autoregressive Transformers with linear attention. International Conference Machine Learning, pages 5156-5165. PMLR, 2020.
[35] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient Transformer. In International Conference on Learning Representations, 2020.
[36] Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document Transformer. arXiv preprint arXiv:2004.055500 2020.
[37] Ankit Gupta and Jonathan Berant. Gmat: Global memory augmentation for Transformers, 2020.
[38] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0477 2020.
[39] Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, pages 4651-4664. PMLR, 2021.
[40] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long Range Arena: benchmark for efficient Transformers. In International Conference on Learning Representations, 2021.
[41] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In Advances in Neural Information Processing Systems, 2022.
[42] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, 2022.
[43] Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. In International Conference on Learning Representations, 2023.
[44] Karan Goel, Albert Gu, Chris Donahue, and Christopher Re. It's raw! Audio generation with state-space models. In Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 7616-7633, PMLR, 17-23 Jul 2022.
[45] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4ND: Modeling images and videos multidimensional signals with state spaces. In Advances in Neural Information Processing Systems, 2022.
[46] Md Mohaiminul Islam and Gedas Bertasius. Long movie clip classification with state-space video models. In Computer Vision-ECCV 2022: 17th European Conference. Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXXV, pages 87-104, 2022.
[47] Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision S4: Efficient sequence-based RL via state spaces layers. In The Eleventh International Conference on Learning Representations, 2023.
[48] Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. arXiv preprint arXiv:2303. 03982, 2023.
[49] Linqi Zhou, Michael Poli, Winnie Xu, Stefano Massaroli, and Stefano Ermon. Deep latent state space models for time-series generation. arXiv preprint arXiv:22 12749 2022.
[50] Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, 2023.
13

---

[51] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In The Eleventh International Conference on Learning Representations, 2023.
[52] Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint arXiv:2212. 10544, 2022.
[53] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint 2rXiv:2302.10866, 2023.
[54] Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning of video representations using LSTMs. In International conference machine learning, pages 843-852, PMLR, 2015.
[55] Arieh Iserles. A first course in the numerical analysis of differential equations. 44, Cambridge university press, 2009.
[56] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in Neural Information Processing Systems, 34, 2021.
[57] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. arXiv preprint S.XXix:3333.633999 2023.
[58] Antonio Orvieto, Soham De, Caglar Gulcehre, Razvan Pascanu, and Samuel Smith. On the universality of linear recurrences followed by nonlinear projections. arXiv preprint arXiv:2307. 1888, 2023.
[59] Eric Martin and Chris Cundy, Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations, 2018.
[60] James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. In International Conference on Learning Representations, 2017.
[61] Tao Lei, Yu Zhang, Sida Wang, Hui Dai, and Yoav Artzi. Simple recurrent units for highly parallelizable recurrence. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 4470-4481, 2018
[62] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HiPPO: Recurrent memory with optimal polynomial projections. Advances Neural Information Processing Systems, 33:1474-1487 2020.
[63] Guy Blelloch. Prefix sums and their applications. Technical report, Tech. rept. CMU-CS-90- 190. School of Computer Science, Carnegie Mellon, 1990.
[64] Wonmin Byeon, Qin Wang, Rupesh Kumar Srivastava, and Petros Koumoutsakos. ContextVP: Fully context-aware video prediction. In Proceedings of the European Conference on Computer Vision (ECCV), pages 753-769, 2018.
[65] Vincent Dumoulin and Francesco Visin. A guide to convolution arithmetic for deep learning. arXiv preprint arXiv: /603.6728. 2016.
[66] Kumar Chellapilla, Sidd Puri, and Patrice Simard. High performance convolutional neural networks for document processing. In Tenth international workshop on frontiers in handwriting recognition. Suvisoft, 2006.
[67] Yangqing Jia. Learning semantic image representations at a large scale. 2014.
[68] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of IEFE conference on computer vision and pattern recognition, pages 770-778, 2016.
14

---

[69] Zhuang Liu, Hanzi Mao, Chao- Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A Convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11976-11986 2022.
[70] Marijn F Stollenga, Wonmin Byeon, Marcus Liwicki, and Juergen Schmidhuber. Parallel multi-dimensional LSTM, with application to fast biomedical volumetric image segmentation. Advances in neural information processing systems, 28, 2015.
[71] Reza Azad, Maryam Asadi Aaghi-ghbb,,,,hi Mahmood Fathy, and Sergio Escalera. Bi-directional ConvLSTM U-net with densely connected convolutions. In Proceedings of the IEEE/CVF international conference on computer vision workshops, pages 0-0, 2019.
[72] Si Woon Lee and Ha Young Kim. Stock market forecasting with super-high dimensional time-series data using ConvLSTM, trend sampling, and specialized data augmentation. expert systems with applications, 161:113700 2020.
[73] Qingging Wang, Ye Huang, Wenjing Jia, Xiangjian He, Michael Blumenstein, Shujing Lyu, and Yue Lu. FACLSTM: ConvLSTM with focused attention for scene text recognition. Science China Information Sciences, 63:1-14, 2020.
[74] Mohamadreza Bakhtyari and Sayeh Mirzaei. ADHD detection using dynamic connectivity patterns of EEG data and ConvLSTM with attention framework. Biomedical Signal Processing and Control, 103708 2022.
[75] Li Kang. Ziqi Zhou, Jianjun Huang, Wenzhong Han, and IEEE Member. Renal tumors segmentation in abdomen CT images using 3D-CNN and COnVLSTM. Biomedical Signal Processing and Control, 72:103334, 2022.
[76] Tie Liu, Mai Xu, and Zulin Wang. Removing rain in videos: a large-scale database and a two-stream ConvLSTM approach. In 2019 IEEE International Conference on Multimedia and Expo (ICME), pages 664-669. IEEE, 2019.
[77] Xiaofang Xia, Jian Lin, Qiannan Jia, Xiaoluan Wang. Chaofan Ma, Jiangtao Cui, and Wei Liang. ETD-ConVLSTM: A deep learning approach for electricity theft detection in smart grids. IEEE Transactions on Information Forensics and Security, 2023.
[78] Xingjian Shi, Zhihan Gao, Leonard Lausen, Hao Wang. Dit-Yan Yeung, Wai-kin Wong, and Wang-chun Woo. Deep learning for precipitation nowcasting: A benchmark and new model. Advances in neural information processing systems, 30, 2017.
[79] Mohammad Babacizadch, Chelsea
Finn, Dumitru Erhan, Roy H. Campbell, and Sergey Levine. Stochastic variational video prediction. In International Conference Learning Representations, 2018
[80] Yunbo Wang, Haixu Wu, Jianjin Zhang, Zhifeng Gao, Jianmin Wang, S Yu Philip, and Mingsheng Long. PredRNN: recurrent neural network for spatiotemporal predictive learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(2):2208 2225, 2022.
[81] Yunbo Wang, Zhifeng Gao, Mingsheng Long, Jianmin Wang, and S Yu Philip. PredRNN++: Towards a resolution of the deep-in-time dilemma in spatiotemporal predictive learning. In International Conference on Machine Learning, pages 5123-5132. PMLR, 2018.
[82] Yunbo Wang, Lu Jiang, Ming-Hsuaa Yang, Li-Jia Li, Mingsheng Long, and Fei-Fei. Fidetic 3D LSTM: A model for video prediction and beyond. In International conference on learning representations, 2019.
[83] Wei Yu, Yichao Lu, Steve Easterbrook, and Sanja Fidler. Efficient and information-ppessrviin future frame prediction and beyond. In International Conference on Learning Representations, 2020.
[84] Jiahao Su, Wonmin Byeon, Jean Kossaifi, Furong Huang, Jan Kautz, and Anima Anandkumar. Convolutional tensor-train LSTM for spatio-temporal learning. Advances in Neural Information Processing Systems, 177-11166. 2020.
15

---

[85] Jue Wang. Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-fon form video understanding. arXiv preprint arXiv:2303. 14526, 2023.
[86] Aäron Van Den Oord, Nal Kalchbremner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference machine learning, pages 1747-1756. PMLR, 2016.
[87] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, Springer, 2015.
2015, Proceedings, Part III 18, pages
[88] Özgün Çiçek, Ahmed Abdulkadir, Soeren Lienkamp, Thomas Brox, and Olaf Ronneherger. learning dense volumetric segmentation from sparse annotation. In Medical Image Computing and Computer-Assiste Intervention-MICCAN 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19, pages 424 432. Springer, 2016.
3D U-net:
[89] Tobias Höppe, Arash Mehrjou, Stefan Bauer, Didrik Nielsen, and Andrea Dittadi. Diffusion models for video prediction and infilling. Transactions on Machine Learning Research.
[90] Mohammad Babaeizadeh, Mohammad Taghi Saffar, Suraj Nair, Sergey Levine, Chelsea Finn, and Dumitru Erhan. FitVid: Overfitting in pixel-level video prediction. arXiv preprint arXiv:2106.1313 2021.
[91] Ruslan Rakhimov, Denis Volkhonskiy, Alexey Artemov, Denis Zorin, and Evgeny Burnaev. Latent video Transformer. arXiv preprint arXiv:2006. 10704, 2020.
[92] Younggyo Seo, Kimin Lee, Fangchen Liu, Stephen James, and Pieter Abbeel. HARP: Autore- gressive latent video prediction with high-fideliyy image generator. In 2022 IEEE International Conference on Image Processing (ICIP), pages 3943-3947. IEEE, 2022.
[93] Ivan Skorokhodov, Sergey Tulyakov, and Mohamed Elhoseiny. StyleGAN-V: A continuous video generator with the price, image quality and perks of StyleGAN2. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 3626-3636, 2022.
[94] Masaki Saito and Shunta Saito. TGAN22: Efficient training of large models for video generation with multiple subsampling layers. arXiv preprint 1e1iv:1811.0921 2(6). 2018.
[95] Sihyun Yu, Jihoon Tack, Sangwoo Mo, Hyunsu Kim, Junho Kim, Jung- Woo Ha, and Jinwoo Shin. Generating videos with dynamics-aware implicit generative adversarial networks. In International Conference on Learning Representations, 2022.
[96] Tim Brooks, Janne Hellsten, Miika Aittala, Ting-Chun Wang, Timo Aila, Jaakko Lehtinen, Ming-Yu Yu Liu, Alexei Efros, and Tero Karras. Generating long videos of dynamic scenes. Advances in Neural Information Processing Systems, 15:31769-31788 2022.
[97] Wenyi Hong. Ming Ding, Wendi Zheng, Xinghan Liu, and Jie Tang. Cogvideo: Large-scale pretraining for text-to-video generation via Transformers. arXiv preprint 2rXiv::205.15866 2022.
[98] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. MaskGit: Masked generative image Transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11315-11325, 2022.
[99] Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, et al. Deepmind Lab. arXiv preprint arXiv: 2..12.0001, 2016.
[100] William H Guss, Brandon Houghton, Nicholay Topin, Phillip Wang, Cayden Codel, Manuela Veloso, and Ruslan Salakhutdinov. MineRL: A large-scale dataset of Minecraft demonstrations. arXiv preprint 17Xiv:1907.13440, 2019.
16

---

[101] Manolis Savva, Abhishek Kadian, Oleksandr Maksymets. Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, al. Habitat: A platform for embodied Al research. In Proceedings of the IEEF/CVF international conference on computer vision, pages 9339-9347, 2019.
[102] Yann Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International Conference Machine Learning, pages 933-941. PMLR, 2017.
[103] Keyu Tian, Yi Jiang, Qishuai Diao, Chen Lin, Liwei Wang, and Zehuan Yuan. Designing BERT for convolutional networks: Sparse and hierarchical masked modeling. arXiv preprint arXiv:2301.00 003580, 2023.
[104] Sunghyun Park, Kangyeol Kim, Junsoo Lee, Jaegul Choo, Joonseok Lee, Sookyung Kim, and Edward Choi. Vid-ODE: Continuous-time video generation with neural ordinary differential equation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 2412-2422, 2021.
[105] Mark Harris, Shubhabrata Sengupta, and John D Owens. Parallel prefix sum (scan) with CUDA. GPU gems, 3(39):8518876, 2007.
[106) Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Re. How to train your HIPPO: State space models with generalized orthogonal basis projections. In International Conference on Learning Representations, 2023.
[107] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1xXiv:6607.0645555 2016.
[108) Thomas Unterthiner, Sjoerd Van Steenkiste, Karol Kurach, Raphael Marinier, Marcin Michal- ski, and Sylvain Gelly. Towards accurate generative models of video: new metric challenges. arXiv preprint arXiv: /8/2.01777 2018.
[109] Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing, 13(4): 600-612, 2004.
[110] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreason- able effectiveness of deep features as a perceptual metric. In CVPR, 2018.
[111] Yann LeCun. The MNIST database of handwritten digits. http://yann. lecun. com/exdd/mmisst/ 1998.
[112] Santhosh Ramakrishnan, Aaron Gokaslan, Erik Wijmans, Oleksandr Maksymets, Alex Clegg, John Turner, Eric Undersander, Wojciech Galuba, Andrew Westbury, Angel X Chang, et al. Habitat-Matterport 3D dataset (HM3D): 1000 large-scale 3D environments for embodied AL. arXiv preprint arXiv:2109008233 2021.
[113] Angel Chang, Angela Dai, Thomas Funkhouser, Maciej Halber, Matthias Niessner, Manolis Savva, Shuran Song, Andy Zeng, and Yinda Zhang. Matterport3ID: Learning from RGB-I) data in indoor environments. arXiv preprint arXiv: 2017.
14] Fei Xia, Amir Zamir, Zhiyang He, Alexander Sax, Jitendra Malik, and Silvio Savarese. Gib- son env: Real-world perception for embodied agents. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 9068-9079, 2018
17

---

Appendix for: Convolutional State Space Models for Long-range Spatiotemporal Modeling
Contents:
Appendix A: Propositions
Appendix B: ConvS5 Details: Parameterization, Initialization, Discretization
Appendix C: Supplementary Results
Appendix D: Experiment Configurations
Appendix E: Datasets
18

---

A Propositions
A.1 Parallel Scan for Convolutional Recurrences
Proposition 1. Consider a convolutional recurrence as in (7) and define initial parallel scan elements
Ck (Ck,a) Ck.b) == (A, B x U.). The binary operator &, defined below, is associative.
C c Cj == (Cj.a o Ci,ar Cja Ci + Cj,b),
(14)
where denotes convolation of two kernels, denotes convolution between a kernel and input, and is elementwise aldition.
Proof. Using that is associative and the companion operator of i.e. (doc) * d (e Blelloch [63], Section 1.4), we have:
(see
(c: @ Cj) C Ck
(eja o Ci,a Gja GiP + Cj&b (CR,as Ck,b)
(15)
Ckaa (eza Cia), Ck (eja Gibb Cj.b)) + ck.bb)
(16)
(Cka o Czaa) Cz Cke (oja c.b) Ckaa * Cj,b + Ck.b
(17)
(Ck.a o Cjaa) o Cia (CR, o (ja) Gi.b + Ckaa Gj.b + Ck.b
(18)
C (Ck,a o Cjar CKa Cj.b ck.b)
(19)
G @ (c, @ Ck)
(20)
A.2 Computational Cost of Parallel Scan for Convolutional Recurrences
Proposition 2. Given the effective inputs B U1 I and pointwise state kernel A E Pxxxxxx, the computational cost of computing the convolutional recurrence in Equation 7 with parallel scan is O (L(33 + P2 HW)).
Proof. Following Blelloch [63]. given a single processor, the cost of computing the recurrence
sequentially using the binary operator defined in Proposition is O (T..TT T. + T+)) where Te refers to the cost of convolving two kernels, T. is the cost of convolution between a kernel and input (PHHW) and T4 is the cost of elementwise addition. The cost of elementwise addition is T4 For state kernels with resolution AA T. ((p))4) and T. 2 HW). For pointwise convolutions this becomes T. O(Pa) and T. HPP) Thus, the cost of computing the recurrence sequentially using @ is (L(33 + pa HW)). Since there are work-efficient algorithms for parallel scans [105], the overall cost of the parallel scan also (LL(p' + H²))))
Note that ConvS5's diagonalized parameterization discussed in Section 3.4 and Appendix B leads to T. = ((P) and T. O(PHW). Therefore the cost of applying the parallel scan with ConvS5 is OLPHHW).
A.3 Connection Between ConvSSMs and SSMs
Proposition 3. Consider ConvSSM state update as in (5) with pointwise state kernel A PxxxIxI input kernel B P xka and input u(1 E R''''W' xU Let Mimzcal (t) H xUKă, be the reshaped result of applying the Image to Column (im2col) [66, 67] operation on the input (L).. Then the dynamics of each state pixel of (5), X X()) E RR,, evolve according to the following differential equation
X'())s 2 + Bsssu (1)
(21)
where the state matrix, AssM E PxP, and input matrix, BSSM E (X(Uk)))) can be formed by reshaping the state kernel, A. and input kernel, 13, respectively.
19

---

Proof. Let Umm2ool E RUR HW denote the result of performing the im2col operation on the input (t) for convolution with the kernel B. Reshape this matrix into the tensor U (tcool))) E R x Reshape Limerol (t) once more into the tensor (t)) E xxxxxxxxxe x
Now, we can write the evolution for the individual channels of each pixel, X' (().,k in (5)
P-1
U- 8-1kn- 1
(t))ks
Σ (() 1=0
+
Bk..m. (()
(22)
q=0
==0
Let ASSM E RxxP be a matrix with rows, ASSMi R., corresponding flattened version of the output features of A, i.e. A, E Pxix1, Similarly, reshape B into matrix BSSM E Px where the rows, BSSM,; E UM correspond to a flattened version of the output features of B, i.e. B; E gU
Then rewrite (22) equivalently
kg-kkl-k ka-1 Bkg.mm (t)...g..m
X'(t
Σ
A (
(23)
P-1
Ukb -1
= (
BSSM o (t)je
(24)
1=
A
&
(1)s.
& (t)i.
(25)
20

---

B ConvS5 Details: Parameterization, Discretization, Initialization
B.1 Background: S5
S5 Parameterization and Discretization S5 [20] uses a diagonalized parameterization of the general SSM in (3).
Let Ast VAst V E RPxP" where Ass E CPPP is a complex-valued diagonal matrix and
V E CPxP corresponds to the eigenvectors. Defining x(() = V-)x(1), B = V-1B, and Č = CV
can reparameterize the SSM of (3) the diagonalized system:
dx(() (t) + Bu(t),
y(t) Cx(() + Du(t).
(26)
dt
S5 uses learnable timescale parameters IRP and the following zero-order hold (ZOH) disretiza-
tion:
Ass DISCRETIZE (A55,
S5, A) As
(27)
Bss = DISCRETIZEg (Ass B, A) == Ags Ass
I)B
(28)
S5 Initialization
S5 initializes its state matrix by diagonalizing Ass as defined here:
(n + 2 / + 4)2
n > n = k. 12+2) + 1)1²² < k
As...
(29)
2
This matrix the normal part the normal plus low-rank HiPPO-Legs matrix from the HipPO frame- work [62] for online function approximation. S4 originally initialized its single-input, single-output (SISO) SSMs with a representation of the full HiPPO-LegS matrix. This was shown to be approximat- ing long-range dependencies at initialization with respect to infinitely long, exponentially-decat measure [106]. Gupta et al. [41] empirically showed that the low-rank terms could be removed with- out impacting performance. Gu et al. [42] showed that in the limit of infinite state dimension, the linear, single-input ODE with this normal approximation to the HiPPO-LegS matrix produces the same dynamics as the linear, single-inph input ODE with the full HiPPO-LegS matrix. The S5 work extended these findings to the multi-input SSM setting [20].
Importance of SSM Parameterization, Discretization and Initialization Prior research has highlighted the importance of parameterization, discretization and initialization choices of deep SSM methods through ablations and analysis [56, 19, 42, 20, 57]. Concurrent work from Orvieto et al. [57] provides particular insight into the favorable initial eigenvalue distributions provided by initializing with HiPPO-in HuPPO-inspired matrices well as important normalization effeet provided by the explicit discretization procedure. They also introduce a purely discrete-timg time parameterization that can perform similarly the continuous- practically ends up quite similar to the equations of (27-28). We choose to use the continuous-time parameterization of S5 for the implicit parameterization of ConvS5 since it can also be leveraged for zero-shot resolution changes [19, 20, 45] and processing irregularly sampled time-series in parallel [20]. However, due to Proposition 3, other long-range SSM parameterization strategies can also be used, such as in Orvieto et al. [57] or potential future innovations.
discretization of $4 and S5. However, their parameterization
B.2 ConvS5 Diagonalization
We leverage S5's diagonalized parameterization to reduce the cost of the parallel scan of ConvS5, Concretely, we initialize Ass as in (29) and diagonalize as Ass VAgs V-1 To apply ConvS5, we compute Ass and Bst using (27-28), and then form the ConvS5 state and input kernels:
Ass E PxP
reshape,
AConvS5
E ER. 1xPxxx
(30)
Bas E R (
reshape
B CouvSS
R PxUxkxx
(31)
See Listing for example of the core implementation. Note, the state kernel Aconvss is "diagonal- ized" in the sense that all entries in the state kernel are zero except AConv55.d, Ast.a, Vi E [P].
21

---

This means that the pointwise convolutions reduce to channel- multiplications. However, this does not reduce expressivity compared full pointwise convolution. This is because, given the ConvSSM to SSM equivalence of Proposition 3 and the use of complex-valued kernels, the diagonal- ization maintains expressivity since almost all SSMs are diagonalizable [41, 42], which follows from the well-known fact that almost all square matrices diagonalize over the complex plane.
22

---

import jax
2
import jax numpy as
3
from ConVSSM helpers import discretize, Conv2D, Restiet_Blod para_lel_sces jax. lax. awscciative__cce
def Sopplyy_ovvs_____ layer
R_ahape, c_karnel, log_pelta, resmet_parama, xC)
**coompute the outputs of Conves layer given Input sequence. Azgs:
A (complex6d): 55 state matrix 0 [complex61]: 55 input matrix a_shape (tuple): shape of g_kerssi karnnl (complex64) output kernai log_neita (flost32); learnable timescaie parame (P,) resnet_pazams (dict): ReeNet block params (floet32): ingut sequence festures x0 (complex64): initial state
(#,) (E, 022
(U, P, k_ k_t)
14
(L, (baz, H, R1,
16
Returns:
outputs : the ConvsS layer outputs x (compiex64): the last state of the ConvssM (b5z, H, P)
(1, baz,
19
*un
21
Discretize and reshape ConvS5 state and input keznels Pr U, k H_shape A_bar, B_bar discretiza (A, B, np.exp Pelta)
Akkeznl A.bar slready correct shape due to diagonalization R_kernel s_.bar.reshapee
k_B, k_s)
Apply Coav85
28
apply_ConvS3 kernel, B_keznel, C_kernel, us, x0)
Apply Restet activation fanction
outputé vmap (ResMet_3Loce axis- (None, 0}) (reanet_parart ys)
outputa, xa[11]
apply_ kernel, B_kernel, C_kernel, us, x0): ***compute the output sequence of the convolutionsl sa8 given input sequence using peralial scon. Computes x_ y_ - x
[2-1]
t_k
where
for operator.
40
Azgs:
A_keznel (complexod) A_kernel (complex64): input kernei kekeeel (complex64) output (float32); input sequence x2 (complex64): initial state
kernei
(B,) (2, k_ k_ (J, k_G, k_cj
42
(1, UJ De,2,i,,,P3
Returns:
47
(float32); the conv55 outputs x_1 (complexod):
(i, H, 0] (b82,
state
Computie initia! scan elements
As repest (A_keznel [None,
shape(0], ex1s-0)
Bus jex. vmap (Conv2D) (B. (e_kernel, up complexx1 (us))
and expand_dime dine (a_bar, 1, 21)
54
Convolutional recurrence persliel seen
parallal Baralllll_ccaaeeee scan (connannnoonoooooeeoo
(As, Busl)
Compute Coav65 outputs
ys jax. vmap (Conv20) (c_kernel, xs) .eeal ys, xs
conv binary www.rbinnnnnaeaaaeeeeeeeeeeeeeee operstor 9-1): sunginary operator for convolutional recurrence with "dlagonalized" IXI state keznels.
Azgs:
66
L,i I (tuples): sean el elements ii=(i, #t_i) 1) shere Ai (complex61) is kernei (P,)
68
8 (complex64) is effective input (bsz, 8,
Returns:
oucput tuple a. 'oiroteoast a_
ALi, _U_i 9_i
A_i. PU_1 = 9.3
Convolve "diagonal* kernels
AA A_ A_
Convalva "élagona)* A_3 with AI] ABBB_ up.expand_dian (AJ,, (9, 1, 21) B return AA, A_ 30_]
Listing 1: JAX implementation of core code apply a single ConvS5 layer batch of spatiotem- poral input sequences.
23

---

C
Supplementary Results
We include expanded tables and sample trajectories from the experiments in the main paper. Sample videos can be found at:
C.1 Moving-MNIST
Table Full results on the Moving-MNIST dataset [54]. For the top table, all models are trained on 300 frames. For the bottom table, all models are trained on 600 frames. The evaluation task is to condition on 100 frames, and then generate forward 400, 800 1200 frames. ConvSSM (ablation) is performed by randomly initializing the state kernel (see Section 5.3 and Appendix D.4).
Trained 390 frames
100
1200
Method
Params
PVD 4
PSNR
SSIM ?
LPIPS
FVD J
PSNK
SSIM
LPIPS 4
FVD J
PSNK
SSIM 1
LPIPS J
Transformer
154M
73 d:
13.5
0.1
0.669
0.002
0.213
1: 0.003
159
12.6 0.1
0.609
0.002
0.287
0.001
265 1:
12.4 0.1
0.591
0.002
0.321 4: 0.002
Performer
184M
111 9
13.4
0.1
0.653
0.002
0.288
1: 0.001
234 =
13.4
0.652
1 0.006
0.379 0.379 0.0002
275 1.5
13.2
0.1
0.592
0.001
0.393 1: 0.001
CW VAB
20M
0.611
0.001
12.4
0.092
ComVLSTM
16.9
0.004
0.002
15.0
14.1
CoeVSSM (ablation)
20M
67
3
16.5
+
0.1
0.742
0.001
0.368
+ 0.901
287
13.6
0.577
+ 0.001
0.293
0.001
511 5
13.3
0.1
0.515 0.001
0.348
+ 0.001
Conv55
20M
26
1
18.1
0.1
0.830
0.003
0.094
+ 0.002
72 1
16.0 +
0.1
0.761
0.005
0.156
0.003
187 +
14.5
0.1
0.878
0.008
0.230
+ 0.004
Trained 600 Trames
400
100
800
100
1200
FVD4
PSNR
SSIM 1
Performer
164M
27 +
13.1 0.1
0.664
0.004
0.206 + 0.001
93 5
11.4
0.618 + 0.002
0.274
0.001
12.2 0.1
0.808
0.001
0.312 + 0.002
VAR
20M
73 +
12.9 0.1
0.621
0.004
0.242 + 0.001
94
12.5 0.9
0.598 + 0.004
0.269 0.00
107 +
12.3 0.1
0.500
0.004
0.280 = 0.002
CTM
0.767
0.005
0.003
CSeVSSM (ablation
20M
6
15.5
0.743
0.002
0.163 1 0.003
145
14.3
0.698
0.002
0.218
0.002
215
0.1
0.614
0.001
0.287 2 0.220.000 0.001
Conv55
20M
23 13
18.1 2 0.1
0.882
0.003
0.092
+ 0.008
=7
16.4
0. 0.1
0.788
= 0.002
O. 134 0.008
71 99
15.6
2 0.1
0.768
0.002
0.162 = 0.008
Table 7: Model runtime comparison for Moving-MNIST results in Table 6. ConvS5 can be parallelized like Transformer but maintains the constant cost-per-step autoregressive generation of ConvRNNs.
100 + 400
100 - 800
100 - 1200
Method
Parallelizable
Train Step Time (8)
Sample Throughput (frames/s) f
Sample Throughput (frames/s) f Sample Throughput (frames/s)
Transformer
YES
0.77 (1.0x)
1.1 (1.0x)
0.34 (1.0x)
0.21 (1.0x)
ConvLSTM
NO
3.0 (3.9x)
117 (106x)
117 (345x)
117 (557x)
ConvS5
YES
0.93 (1.2x)
(82x)
(265x)
90 (429x)
24

---

Condition
Predicted Frames
Ground Truth
39
3
9
9
3
9
3
9
3
2
9
9
3
9
3
9
3
3
Transformer
39
3
9
9
3
9
3
9
3
9
93
ConvLSTM
39
3
9
9
3
9
3
9
3
4
3
A
3
3
3
ConvS5
39
3
9
9
3
9
3
3
2
9
3
R
3
3
9
3
3
9
9
t=0
50
99
100
250
400
550
700
850
1000
1150
1299
(a) Example
Condition
Predicted Frames
0
00
Ground Truth
0
00
7
Transformer
7
ConvLSTM
ConvS5
0
00
8
8
00
0
0
t=0
50
99
100
250
400
550
700
850
1000
1150
1299
(b) Example 2
Condition
Predicted Frames
66
4
Ground Truth
6
67
76
7
Transformer
4
7
7
7
G
IT'
u
b
ConvLSTM
6
7
76
ConvS5
U
7
7
4
L7
70
7
7
t=0 0
50
99
100
250
400
550
700
850
1000
1150
1299
(c) Example
Figure 3: MNVIS Samples: 1200 frames generated conditioned on 100.
25

---

C.2 3D Environments
Table 8: Full results for DMLab long-range benchmark dataset [13]. Results from Yan et al. [13] are indicated with *. We separate out the methods trained using the TECO [13] training framework in the bottom of the table. TECO-COnVSSM (ablation) refers to the ablation performed by randomly initializing the state kernel (see Section 5.3 and Appendix D.5).
DMLab
Method
Params
FVD 4
PSNR f
SSIM f
LPIPS +
FitVid*
165M
176 + 4.86
12.0 0.013
0.356 + 0.00171
0.491 + 0.00108
CW-VAB*
111M
125 + 7.95
12.6 + 0.059
0.372 + 0.00033
0.465 + 0.00156
Perceiver AR
30M
11.2 0.004
0.304 + 0.00004
+ 0 0.00123
Latent FDM*
31M
18.2
0.588 0.00453
0.222 0.00499
Transformer
152M
9..9
19.9 0.108
0.619 0.00506
0.123 + 0.00191
Performer
152M
3.23
17.3 0.074
0.513 = 0.00492
0.205 + 0.00315
S5
140M
221 13.1
19.3 - 0.128
0.641 = 0.00400
0.162 = 0.04510
ConvS5
100M
4.81
23.2 + 0.053
0.769 0.01020
0.079 0.00073
TECO-Trane Transformer*
173M
1.77
22.4 0.368
0.709 = 0.0119
0.155 = 0,00958
TECO-Transformer (our run)
173M
28.2
21.6 - 0.079
0.696 = 0.02640
0.082 0082=00 0.00119
TECO-S5
180M
0.26
20.1 0.037
0.687 = 0.00132
= 0.143 0.00049
TECO-ConVSSM (ablation)
175M
44.3
21.0 + 0.106
0.691 + 0.00004
0.010 + 0.00267
TECO-ConvS5
175M
+
23.8 + 0.056
0.803 + 0.0020
0.085 + 0.00179
Table 9: Full results for ablation of ConvS5 convolutional ablations for DMLab long-range benchmark dataset [13]. To make parameter counts comparable for different configurations, when possible, we adjust parameters the activation blocks, e.g. increasing the size of the ResNet convolution kernels or increasing the features of the GLU activation. With these adjustments, the models also have similar training speeds. Note the last entry the S5 which serves an additional ablation of the convolutional structure.
DMLab
B kernel
C kernel
Activation
Params
FVD
PSNR 1
SSIM 1
LPIPS
Train Step Time (s)
3
3
ResNet
100M
66.6 + 4.81
23.2 0.053
0.769 + 0.00043
0.079 + 0.00073
2.31
1 x 1
3 x 3
ResNet
85M
67.5 0.73
0.756 #: 0.00039
0.085 + 0.00124
2.52
1x1
ResNet
100M
81.1
0.074
0.767 - 0.00186
0.083 + 0.00139
2.21
3 3x3
3 3 3
GLU
71M
96.1
0.177
0.762 :: 0.00363
0.088 + 0.00287
2.75
1x1
5
GLU
83M
89.1 0.61
2.072
0.713 = 0.00290
0.106 + 0.00289
2.58
1x1
1x1
GLU
30M
187 2.77
0.689 + 0.00007
0.112 + 0.00183
1.73
GLU
140M
221 13.1
0.128
0.641 = 0.00400
0.162 + 0.04510
1.34
26

---

Condition
Predicted Frames
Ground Truth
Transfomer
S5
ConvS5
TECO-Transfomer
TECO-S5
TECO-ConvS5
t=0
73
143
144
163
182
202
221
240
260
279
299
(a) 156 frames generated conditioned on 144 Duction-conditions
Condition
Predicted Frames
Ground Truth
Transfomer
$5
ConvS5
TECO-Transfomer
TECO-S5
TECO-ConvS5
t=0
17
35
36
68
101
134
167
200
233
266
299
(b) 264 frames generated conditioned on
action action conditionin))).
Figure 4: DMLab Samples
27

---

Condition
Predicted Frames
Ground Truth
TECO-Transformer
TECO-ConvS5
t=0
73
143
144
163
182
202
221
240
260
279
299
(a) 156 frames generated conditioned on 144
Condition
Predicted Frames
Ground Truth
TECO-Transfomer
TECO-ConvS5
t=0
17
35
36
68
101
134
167
200
233
266
299
(b) 264 frames generated conditioned on 36 (action-conditioned)
Figure 5: Minecraft Samples
Condition
Predicted Frames
Ground Truth
TECO-Transfomer
TECO-ConvS5
t=
73
143
144
163
182
202
221
240
260
279
299
(a) 156 frames generated conditioned on 144 Daction-conditiona
Condition
Predicted Frames
Ground Truth
TECO-Transfomer
TECO-ConvS5
t=0
17
35
36
68
101
134
167
200
233
266
299
(b) 264 frames generated conditioned on 36 (no action-condttioning)
Figure 6: Habitat Samples
28

---

Table 10: Full results the Minecraft and Habitat long-range benchmark datasets [13]. Results from Yan et al. [13] are indicated with s. Note that Yan et al. [13] did not evaluate FitVid or CW-VAE on Habitat due cost.
Minecraft
Method
Params
FVD 4
PSNR t
SSIM
LPIPS +
F
176M
956 + 15.8
13.0 + 0.0089
0.343 + 0.00380
0.519 + 0.00367
CW-VAES
140M
397 + 15.5
13.4 + 0.0610
0.338 + 0.00274
0.441 + 0.00367
Perceiver AR*
166M
76.3 + 1.72
13.2 + 0.0711
0.323 + 0.00336
0.441 + 0.00207
Latent FDM*
33M
167 6.26
13.49000000
0.349 = 0.00327
0.429 + 0.00284
TECO- Transformer*
274M
116 = 5.08
0.000
0.381 + 0.00192
0.340 + 0.00264
TECO-ConvS5
214M
70.7 = 3.05
14.8 0.0984
0.374 + 0.00414
0.355 + 0.00467
Habitat
Method
Params
FVD 4
PSNR t
SSIM
LPIPS +
Perceiver AR*
200M
164 12.6
12.8 + 0.0423
0.405 0.00248
0.676 + 0.00282
Latent FDM
87M
12.5 0.0121
0.311 + 0.00083
0.582 + 0.00049
TECO- Transformer*
386M
76.3
12.8 0.0139
0.363 = 0.00122
0.604 0.00451
351M
9
0.390 = 0.01238
0.632 + 0.00823
Table 11: Model runtime comparison for 3D Environment results in Tables 8-10. The implementations of the baselines Vir CW-VAE, Perceiver AR and Latent FDM used in the TECO work [13] publicly available in the TECO repository, so we were unable to include direct runtime comparisons for those methods.
DMLab
Method
Train Step Time (s) $
Sampling Speed (frames/s)
Transformer
1.25 (1.0x)
9.1 (1.0x)
Performer
1.25 (L.x))
7.6(0.8x)
S5
1.34 (1.1x)
28 (3.1x)
ConvS5
2.31 (1.8x)
56 (6.2x)
TTCOO-rrasformer
0.75 (0.6x)
16 (1.8x)
TECO-S5
0.81 (0.7x)
21 (2.3x)
TECO-ConvS5
1.17 (0.9x)
18 (2.0x)
Minecraft
Method
Train Step Time (8)
Sampling Speed (frames/s)
FFCO-Trrasformer
1.91 (1.0x)
8.1 (1.0x)
TECO-ConvS5
2.53 (1.3x)
14 (1.7x)
Habitat Sampling Speed (frames/s)
Method
Train Step Time (8) 4
IFCO-Transformes
2.71 (1.0x)
(1.0x)
TECO-ConvS5
3.10 (1.1x)
11 (1.6x)
29

---

D Experiment Configurations
Our codebase modifies the TECO codebase from Yan et al. [13] and we reuse their core Transformer and TECO framework implementations. More architectural details and dataset-specific details are described below.
D.1 Spatiotemporal Sequence Model Architectures
ConvS5, ConvLSTM and S5 models are formed by stacking multiple ConvS5, ConvLSTM or S5 layers, respectively. For each of these models, layer normalization [107] with a post-norm setup used along with residual connections. For the Transformer, we use the Transformer implementation from Yan al. [13] which consists of stack of multi-head attention layers.
ConvS5 and ConvLSTM are applied directly to sequences of frames of shape [sequence length, latent height, latent width, latent features], where the original data has been convolved and latent number of features. Since S5 and Transformer act on vector-valued sequences, these models require an additional downsampling convolution operation to project the latent frames into token and an upsampling transposed convolution operation to project the Transformer backbone output tokens back into latent frames. We use the same sequence of compression operations for this as in Yan et al. [13]. The Encoder and Decoder referred to for all models in the hyperparameter tables below consist of ResNet Blocks with 3 kernels as implemented in Yan et al. [13].
latent resolution
D.2 Evaluation Metrics
We follow Yan al. [13] and evaluate methods by computing Fréchet Video Distance (FVD) [108], peak signal-to-noise ratio (PSNR), structural similarity index measure [109] and Learned Perceptual Image Patch Similarity (LPIPS) [110] between sampled trajectories and ground truth trajectories. See Yan et al. [13] for more in-depth discussion of the use of these metrics for the 3D environment benchmarks.
D.3
Compute
All models were trained with 32GB NVIDIA V100 GPUs. For Moving-MNIST. models were trained with 8 V100s. For all other experiments, models were trained with 16 V100s. We list V100 days in the hyperparameters, which denotes the number of days it would take to train on a single V100.
D.4 Moving-MNIST
All models were trained to minimize L1+L2 loss over the frames directly in pixel space, as in Su et al. [84]. We trained models 300 frames. We then repeated the experiment and trained models 600 frames. For ConvS5 and ConyLSTM, we fixed the hidden dimensions (layer input/output features) and state sizes to be 256, and swept
the following learning rates [1 10 5 x 10 1 x 10-3J and chose the best model. For Transformer, we swept over model size, considering hidden dimensions of [512, 2014] and learning rates [1 10-1, 10-1 10-31 and chose the best model. We also observed better performance for the Transformer by convolving frames down to an 8 x 8 latent resolution (rather than the 16 16 used by ConvS5 and ConvLSTM) before downsampling to a token. All other relevant training parameters were kept the same between the three methods. See Tables 12-14 for detailed experiment configurations.
Each model was evaluated by collecting 1024 trajectories using the following procedure: condition on 100 frames from the ground truth test then generate forward 1200 frames. These samples were compared with the ground truth to compute FVD, PSNR, SSIM and LPIPS.
The ConvSSM ablation was performed using the exact settings as ConvS5, except the state kernel was initialized with a Gaussian and we swept over the following learning rates [1 10-4, 10-1
1 x 10-31.
30

---

Table 12: Experiment Configuration for ConvS5 on Moving-MNIST experiments
Hyperparameters
Moving-MNIST-300
Moving-MNIST-600
V100 Days
25
50
Params
20M
20M
Input Resolution
64
64
64 x 64
Latent Resolution
16 x 16
16 x 16 8
Batch Size
8
Sequence Length
300
600
LR
x 10-3
x 10-3
LR Schedule
cosine
cosine 5k 300K
Warmup Steps
5k 300K
Max Training Steps
Weight Decay
1 x 10-5
x 10
Encoder
Depths
64, 128, 256
64, 128. 256 1
Blocks
Decoder
Depths
64, 128, 256
64, 128, 256
Blocks
Hidden Dim (U)
256
256
State Size (P)
256
256
Kernel Size
x 3
3 x
ConvS5
C Kernel Size
x 3
3 X
Layers
8
8
Dropout
0
0
Activation
ResNet
ResNet
Table 13: Experiment Configuration for ConvLSTM on Moving-MNIST experiments
Hyperparameters
Moving-MNIST-600
V100 Days
75
150
Params
20M
20M
Input Resolution
64 x 64
64 x 64
Latent Resolution
16 16
16 16
Batch Size
8
8
Sequence Length
300
600
LR
5 10-4
5 10-4
LR Schedule
cosine
cosine
Warmup Steps
5k 300K
5k 300K
Max Training Steps
Weight Decay
1 10-5
1 10
Encoder
Depths
64, 128, 256
64, 128, 256 1
Blocks
Decoder
Depths
64, 128, 256
64, 128, 256
Blocks
Hidden Dim
256
256
State Size
256
256
Kernel Size
3 3
3 x 3
ConvLSTM
Layers
8
Dropout
0
31

---

Table 14: Experiment Configuration for Transformer on Moving-MNiin experiments
Hyperparameters
Moving-MNISTIN Moving-MNIST--600
V100 Days
25
50
Params
164M
164M
Input Resolution
64 64
64 64
Latent Resolution
8 x
8 x 8
Batch Size
8
Sequence Length
300
600
LR
x 10-4
1 x 10-4
LR Schedule
cosine
cosine 5k 300K
Warmup Steps
5k
Max Training Steps
300K
Weight Decay
x 10
1 x 10
Encoder
Depths
64, 128, 256, 512 1
64, 128, 256, 512
Blocks
Decoder
Depths
64, 128, 256, 512
64, 128, 256, 512
Blocks
Downsample Factor
8
Hidden Dim
1024
1024
Temporal
Feedforward Dim
4096
4096
Transformer
Heads
16
16
Layers
8
Dropout
0
32

---

All sequence models we trained used the same number of layers as the Transformer used in the TECO-S5 and TECO-ConvS5 models we trained used the exact encoder/decoder configuration and MaskGit D.5 Long-Range 3D Environment Benchmarks
We follow the procedures from Yan et al. [13] and train models the same pre-trained vector- quantized (VQ) 16 x 16 codes used by the basclines evaluated in that work. Models were trained to optimize cross-entropy reconstruction loss between the predictions and true VQ codes. The evaluation of DMLab and Habitat involves both an action-conditione and unconditioned setting. Therefore, as in Yan et al. [13], the actions were randomly dropped out half the time during training on these datasets.
After training, we follow the procedure from Yan et al. [13] for evaluation in two different settings. The first setting involves computing PSNR, SSIM and LPIPS from 1024 samples generated by conditioning on the first 144 frames and then generating the next 156 frames while providing the model with past and future actions. The second setting does not provide actions as input (with the exception of Minecraft, which also provides actions in this setting). It involves computing FVD using 1024 samples generated by conditioning on the first 36 frames and then predicting the remaining 264 frames.
TECO-Transforme trained by Yan et al. [13]. In addition, the
configuration as in Yan et al. [13]. The Transformer, S5 and ConvS5 models we trained without the TECO framework were all trained using the same encoder/decoder configuration. See Tables 15-21 for more detailed experimental configuration details. See dataset-specific paragraphs below for hyperparameter tuning information.
TECO Training Framework Yan et al. [13] proposed the TECO training framework to train Transformers on long video data. For some of our experiments, we use ConvS5 layers and S5 layers as drop-in replacement for the Transformer in this framework. We refer the reader to Yan et al. [13] for full details. Briefly, given the original VQ codes, TECO trains an additional encoder/decoder that compresses the frames to a lower latent resolution (e.g., from 16 x 16 to 8 x 8 by training an additional encoder/decoder with codebook loss, Lvq- In addition, a MaskGit [98] dynamics prior loss, ppiirr. is used for the latent transitions. The sequence model (e.g. Transformer, S5, ConvS5) takes the latent frames (compressed into tokens in the case of Transformer and S5) and produces an output which is used along with the latents by the decoder to produce predictions and reconstruction loss, Lrecon. Models are trained to minimize the following total loss:
LrEco = LvQ + Lrocon + Lprior-
(32)
In addition, TECO includes the use of DropLoss [13], which drops out percentage of random timesteps that are not decoded and therefore do not require computing the expensive Crecon and Cptior terms.
DMLab As mentioned above, the actions were randomly dropped out of sequences half the time (due to the evaluation scenarios, eetion-conditioneed and unconditioned). We observed that for DMLab, when provided past and future actions, models converged faster using the simple masking strategy discussed in Gu et al. [19] that masks the future inputs rather than feeding the predicted inputs (or true inputs during training) autoregressively. Therefore we trained all models (Transformer, Performer, $5, ConvS5, the actions were provided, and using the autoregressive strategy when actions were not provided. We observed this significantly improved the LPIPS of the Transformer baselines. Note, pilot runs for Minecraft and Habitat, we observed this strategy led to lower-quality frames and did not use it for the reported results for those datasets.
TECO-S5, TECO-ConVS5) by using this strategy when
We trained each model, Transformer, Performer, S5, ConvS5, Pico-Transformer, TECO-S5, TECO-
ConvS5, with three different learning rates [1 10 10-4, 5 x 10-4, 4, 1 x 10 3 and selected the best run for each model. See Tables 15-20 for more experiment configuration details.
The TECO-ConvSE ablation used the exact same settings TECO-ConvS except the state kernel was initialized with random Gaussian and a lower learning rate of x 10-5 was required for stable training.
33

---

Table 15: Experiment Configuration for ConvS5 on DMLab
Hyperparameters
DMLab
VI00 Days
150
Params
101M
Input Resolution
64 64
Latent Resolution
16 16
Batch Size
Sequence Length
300
LR
5 x 10-4
LR Schedule
cosine
Warmup Steps
5k
Max Training Steps
500K
Weight Decay
1 x 10
Encoder
Depths
256
Blocks
Decoder
Depths
256
Blocks
Hidden Dim (U)
512
State Size (P)
512
Kernel Size
3 3
ConvS5
C Kernel Size
3 x 3
Layers
Dropout
Activation
ResNet
Table 16: Experiment Configuration for S5 on DMLab
Hyperparameters
DMLab
V100 Days
125
Params
140M
Input Resolution
64 64
Latent Resolution
16 x 16
Batch Size
16
Sequence Length
300
LR
x 10-3
LR Schedule
cosine
Warmup Steps
5k
Max Training Steps
500K
Weight Decay
x 10 5
Encoder
Depths
256
Blocks
Decoder
Depths
256
Blocks
Downsample Factor
16
Hidden Dim (U)
1024
State Size (P)
1024
S5
Layers
Dropout
Activation
GLU (half)
34

---

Table 17: Experiment Configuration for Transformer on DMLab
Hyperparameters
DMLab
V100 Days
125
Params
152M
Input Resolution
64 64
Latent Resolution
16 x 16
Batch Size
Sequence Length
300
LR
5 x 10
LR Schedule
cosine
Warmup Steps
5k
Max Training Steps
500K
Weight Decay
1 x 10 5
Encoder
Depths
256
Blocks
Decoder
Depths
256
Blocks
Downsample Factor
Hidden Dim
512
Temporal
Feedforward Dim
2048
Transformer
Heads
16
Layers
8
Dropout
0
35

---

Table 18: Experiment Configuration for TECO-ConvS5 on DMLab
Hyperparameters
DMLab
V100 Days
110
Params
175M
Input Resolution
64 x 64
Latent Resolution
8 x 8
Batch Size
16
Sequence Length
300
LR
5 10-4
LR Schedule
cosine 5k
Warmup Steps
Max Training Steps
500K
Weight Decay
1 x 10
DropLoss Rate
0.9
Encoder
Depths
256, 512
Blocks
2
Size
1024
Codebook
Embedding Dim
32
Decoder
Depths
256, 512
Blocks
4
Hidden Dim (U)
512
State Size (P)
1024
B Kernel Size
3 x 3
ConvS5
Kernel Size
x
Layers
8
Dropout
0
Activation
ResNet
Mask Schedule
cosine
Hidden Dim
512
Feedforward Dim
2048
MaskGit
Heads
8
Layers
8
Dropout
0
36

---

Table 19: Experiment Configuration for TECO-S5 on DMLab
Hyperparameters
DMLab
V100 Days
80
Params
180M
Input Resolution
64 64
Latent Resolution
8 x 8
Batch Size
Sequence Length
300
LR
1 x 10-3
LR Schedule
cosine 5k 500K
Warmup Steps
Max Training Steps
Weight Decay
1 x 10
DropLoss Rate
0.9
Encoder
Depths
256, 512 2
Blocks
Size
1024
Codebook
Embedding Dim
32
Decoder
Depths
256, 512
Blocks
4
Downsample Factor
8
Hidden Dim (U)
2048
State Size (P)
2048
S5
Layers
Dropout
0
Activation
GLU (half)
Mask Schedule
cosine
Hidden Dim
512
Feedforward Dim
2048
MaskGit
Heads
Layers
8
Dropout
0
37

---

Table 20; Experiment Configuration for TECO-Transforme on DMLab
Hyperparameters
DMLab
V100 Days
80
Params
173M
Input Resolution
64 64
Latent Resolution
8 x 8
Batch Size
Sequence Length
300
LR
x 10-4
LR Schedule
cosine
Warmup Steps
5k 500K
Max Training Steps
Weight Decay
1 x 10-5
DropLoss Rate
0.9
Encoder
Depths
256, 512
Blocks
2
Size
1024
Codebook
Embedding Dim
32
Decoder
Depths
256, 512
Blocks
Downsample Factor
8
Hidden Dim
1024
Temporal
Feedforward Dim
4096
Transformer
Heads
16
Layers
8
Dropout
Mask Schedule
cosine
Hidden Dim
512
Feedforward Dim
2048
MaskGit
Heads
Layers
8
Dropout
0
38

---

Minecraft and Habitat For Minecraft and Habitat, we only trained TECO-Conv's due to the costs of training on these datasets. See dataset details Appendix and reported compute costs in Yan et al. [13]. For Minecraft, we evaluated two different learning rates [1 10-1, 5 x 10=2) and chose the best. For Habitat, we only performed one run with no further tuning. See Table 21 for further experiment configuration details.
Table 21: Experiment Configuration for TECO-ConvS5 on Minecraft and Habitat
Hyperparameters
Minecraft
Habitat
V100 Days
470
575
Params
214M
351M
Input Resolution
128 x 128
128 128
Latent Resolution
8 x 8
8 x 8
Batch Size
16
Sequence Length
300
300
LR
x 10
1 x 10-4
LR Schedule
cosine
cosine
Warmup Steps
5k
5k
Max Training Steps
1M
IM
DropLoss Rate
0.9
0.9
Encoder
Depths
256, 512
256, 512
Blocks
4
Size
1024
1024
Codebook
Embedding Dim
32
32
Decoder
Depths
256, 512 8
256, 512
Blocks
Hidden Dim (U)
512
512
State Size (P)
512
512
B Kernel Size
3
3
ConvS5
C Kernel Size
3 x 3
3 x 3
Layers
12
Dropout
0
0
Activation
ResNet
ResNet
Mask Schedule
cosine
cosine
Hidden Dim
768
1024
Feedforward Dim
3072
4096
MaskGit
Heads
12
16
Layers
6
Dropout
0
39

---

E Datasets
E.1 Moving-MNIST
The MovIng-MNIST [54] dataset is generated by moving two 28 28 size MNIST digits from the MNIST dataset [111] inside a 64 64 black background. The digits begin a random initial location, and move with constant velocity, bouncing when they reach the boundary. For each of the sequence lengths we consider, 300 and 600, we follow Wang al. [81] and Su al. [84] and generate 10,000 sequences for training.
E.2 DMLab
We use the DMLab long-range benchmark designed by Yan al. [13] using the DeepMind Lab (DMLab) [99] simulator. The simulator generates random 3D mazes with random floor and wall textures. The benchmark consists 40K action-conditioned, 300 frame videos 64 x 64 resolution. The videos are of an agent randomly navigating 7 x 7 mazes by choosing random points in the maze and navigating them through the shortest path.
E.3 Minecraft
We use the Minecraft [100] long-range benchmark designed by Yan et al. [13]. The game features 3D worlds that contain complex terrains such as hills, forests, rivers and lakes. The benchmark was constructed by collecting 200K action-conditioned 300 frame videos at a 128 x 128 resolution. The videos Minecraft's marsh biome and the agent iterates walking forward for random number of steps and randomly rotating left or right. This results in parts of the scene going out of view and coming back into view later.
E.4 Habitat
We use the Habitat long-range benchmark designed by Yan et. [13] using the Habitat simulator [101] The simulator renders trajectories using of real 3D scenes. Yan et al. [13] compiled 1400 indoor scans from HM3D 12], Matterport3D [113] and Gibson [114] to generate 200K action-conditioned. 300 frame videos with 128 128 resolution. Yan et al. [13] used Habitat's in-built path traversal algorithm to construct action trajectories that move the agent between randomly sampled locations.
40


Reviews:
[{"review_id": "L0IrfdzQjX", "reviewer": "Reviewer_P1WH", "summary": "This paper introduces a method for long-term sequential modeling. The authors extend prior RNN-based work for sequential modeling, i.e. S5[20],  by substituting the linear operations with convolutions.   The authors show the superiority of their method on the future prediction task with transformers. A favorable property of the RNN methods, in general, compared to transformers is that they scale linearly with respect to the sequence length, while the transformer's complexity is quadratic with respect to time.", "questions": "Please see the weaknesses. ", "limitations": "No, they have not. However, I do not see a particular negative social impact associated for this work.", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1) The proposed method makes sense and is a natural extension of prior work.\n2) The paper is well-written and easy to follow.", "weaknesses": "1) Compared to the transformers, the linear complexity with respect to the sequence length is clearly favorable. However, it would be informative to compare the actual computational cost, e.g. in terms of flops.\n\n2) Despite transformers, RNNs are notorious for being difficult to train. It would be great to compare the training of two architecture as well. This can be a potential blocker for scaling the method for more complex datasets.\n\n3) I am not an expert in this domain to have a proper evaluation on the impact of the experiments, but the contributions of this paper sound marginal with the NeurIPS standards. Also, it looks like the paper is missing some prior work to compare[1].\n\n[1] Gao et al,  Simvp: Simpler yet better video prediction, CVPR 2022. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time they took to review our paper and for their feedback. We are glad the reviewer found the proposed method easy to follow.\n\n-- **Re. flops:** Please see the flops comparison in the computational cost section of the General Response above. Note the appendix of the original submission included the computational cost including training and inference runtimes, parameter counts, etc.\n\n-- **Re. training RNNs:** We agree RNNs are notoriously difficult to train, generally due to the vanishing/exploding gradient problem. We note the SSM line of work (e.g. S4/S5, etc.) that ConvS5 extends to the spatiotemporal domain effectively mitigates the vanishing/exploding gradient problem through its special parameterizations and initialization schemes. One of the major contributions of our work is to ensure the convolutional recurrence formulation retains these same favorable properties (See Section 3.3). We will add further discussion in the paper to make this more clear. \n\nIn practice, we did not observe any stability issues when training ConvS5. However, we found that Transformer and ConvLSTM can be unstable with higher learning rates during the hyperparameter search (discussed in Appendix D), especially with Moving-MNIST trained on 600-length context. In addition, random initialization of the ConvSSM kernel instead of careful design of the kernel parameterization/initialization proposed in our paper (Section 3.3) was also unstable with higher learning rates.  \n\nWe do not think the recurrent structure of ConvS5 is a blocker to scaling the method, and the empirical results on complex, large scale datasets also support this. Please let us know if you have any other suggestions for comparing the training of the two architectures. \n\n-- **Re. contribution:** We respectfully disagree with the reviewer that the contributions are marginal. Please see the Technical Contributions section of the General Response above for a detailed explanation of the contributions.\n\nIn short, we have introduced a method that can be parallelized during training like a transformer, provides fast autoregressive generation like an RNN, and provides high performance on complex video generation tasks requiring long-range reasoning and high-quality frames. Achieving these three aspects required contributions to ensure both computational efficiency and high performance.\n\nWe show how convolutional recurrences can be parallelized both theoretically and practically (Section 3.2). To our knowledge, this has not been previously considered. Also, to capture long-range spatiotemporal dependencies, we developed a connection between the dynamics of SSMs and ConvSSMs (Proposition 3), which informs our parameterization and initialization design as discussed in Sections 3.3 and 3.4. \n\nFinally, our empirical results validate these design choices.\n\n-- **Re. SimVP:** We are happy to include a citation of SimVP in the related works. However, this work did not consider long-term prediction tasks (they only considered up to 40 frames). Since the SimVP model considers the sequence length as an input channel of the convolution performed by its Translator module, the SimVP model size has to scale with the sequence length. This may affect the scaling of this model to long-sequence video modeling. In our work, we compare against numerous state-of-the-art video prediction baselines designed for modeling long video sequences and evaluate on the most challenging existing video prediction baselines (training on hundreds of frames and predicting hundreds to thousands of frames). Nonetheless, we are working to include this method as an additional baseline. There was not time to complete the training and evaluation of this prior to the rebuttal deadline, but we will include it in the final version.\n\n**Thank you again** for taking the time to review our paper. We hope we have addressed the reviewer's questions and the reviewer is willing to raise their score. Please let us know if we can provide any additional clarification or information. "}, {"review_id": "g9GGboAjlF", "reviewer": "Reviewer_Kz49", "summary": "This paper proposes convolutional state space models (ConvSSMs), that combine ConvLSTM with the long sequence modeling approaches like S4/S5. In particular, authors propose ConvS5, that allow parallelizing the stateful autoregression of convRNNs across the sequential direction. Naively applying parallel scan speedup tricks like in S4/S5, leads to convolutional kernels of exploding kernel size. Hence, the authors propose an interesting way of structuring modeling complexity, choosing to have simple linear dynamics in the sequence direction that are learnt with 1x1 conv kernels, and complex non-linear operations in depth that can be learnt with entire resnet blocks. The authors benchmark the proposed ConvS5 model on the Moving-MNIST dataset by training on 300 and 600 frames and predicting 400/800/1200 frames conditioned on the initial 100 frames. Also benchmarking results are reported on the DMLab, Minecraft, and Habitat long-range benchmarks. Across the axes of image quality such as FVD, PSNR, SSIM and LPIPS as well as efficiency axes such as sampling speed.   ", "questions": "- What hyperparameter sweeps and effort went into optimizing the Transformer performance on the proposed benchmarks? Was it commensurate with the proposed convS5? Also see weaknesses. ", "limitations": "Yes ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "+ **Well motivated and principled design** : The design considerations in ConvS5 is very well motivated with both theoretical (most of which are directly inspired from Deep SSMs) and practical arguments that make an interesting and worthwhile contribution to the community. \n\n+ **Broad benchmarking and strong results** : The authors perform benchmarking across a number of datasets and benchmarks across all of which ConvS5 outperforms prior works such as ConvLSTM and a vanilla transformer on image quality as well as sample speed. \n\n+ Paper is also well written with good exposition of the prior work ideas that are necessary to grasp convSSM which is difficult to balance given the large amount of history in S5 model family development. ", "weaknesses": "- **Model Size** such as FLOPs / parameters and **Training Efficiency** such as Training Speed /  Number of Epochs need to reported for both convS5 and also all the baselines that are being compared to. Without these details, quality results on their own are meaningless since it is unclear if the best footing was also provided to the baselines as the proposed method. It is mentioned as an offhand remark that convS5 uses ResNet blocks as intermediate non linearities but that can lead to major slowdown and burdens (as also alluded to by the authors in limitations). If so, this should be made clear in the tables through the metrics mentioned before.  \n\n- **Boarder Ablations** : The paper is quite weak on ablation studies presented. Since the authors propose a new architecture, boarder ablations on more than just initialization are requires to be convincing of choices proposed. \n\n- Experiments on non synthetic real data/video. While the authors do experiment on a bunch of datasets, they're all either toy like Moving-MNIST or, synthetic like DMLab/Minecraft/Habitat. For ConvS5 to be successful we need to make sure that it does well on real world benchmarks as well. \n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the time the reviewer took to review our paper and thank the reviewer for their positive feedback. We are glad the reviewer found the design of our method well-motivated and principled and also recognized the broad benchmarking of our paper and strong results of our approach.\n\n-- **Re. model size and training efficiency**: Please see the discussion of the computational cost and the other comparisons including flops between Transformer, ConvLSTM and ConvS5 in the General Response above. We also note that parameters, training speed, and max number of training steps were included for all experiments in Appendix C and D of the original submission. Parameter counts were included in the expanded Tables 5,7,8 and train step speeds were included in Tables 6,9. Parameters, max number of training steps and V100 days are listed in Tables 10-19. As the reviewer can see, the training footprints for the methods were commensurate. We will include a full summary of the tables in the final version to make this more clear. \n\n-- **Re. ablations:** Please see the discussion of ablations in the General Response above. Even though we did not explicitly label it as an ablation, the S5 baseline was meant to be an ablation on the convolutional structure of ConvS5, since it is the closest possible model without the convolutional structure. We have added a separate ablation table and discussion that makes the ablation of the convolutional structure (using S5) more clear and also added new ablations of the nonlinearity choice between layers. \n\nAs our model sticks to a pretty basic architecture, there were not many other obvious design choices that can be ablated. If the reviewer has suggestions for other ablations we should run to improve the paper, please let us know and we will be happy to include!\n\n-- **Re. datasets:** We agree that real-world long-range spatiotemporal datasets are important to test and improve long-range spatiotemporal models. Unfortunately, strong benchmarks in this area do not currently exist, and the 3D Environment tasks from TECO are the most challenging and well-thought-out long-range video benchmarks we are aware of. We hope these recent works on long-range video prediction help to inspire the creation of better real-world long-range video benchmarks. \n\n--**Re. hyperparameter sweeps:** All of this information is included in Appendix D of the supplement of the original submission. \nHere is the summary: \n\n- Moving-Mnist: For the Transformer, we swept over 2 model sizes (hidden dimension of 512 and 1024) and 3 learning rates and chose the best model. For ConvS5 and ConvLSTM, we chose a single model size (less than the Transformer) and swept over the same learning rates as the Transformer. \n\n- DMLab: For each of the methods we chose model sizes and hyperparameters very close to those used in the TECO paper and swept over 3 learning rates for the Transformer, S5, ConvS5, TECO-Transformer, TECO-S5, and TECO-ConvS5, and chose the best run with the best learning rate for each model.\n\n- Minecraft: TECO-ConvS5 was run with two different learning rates with no other tuning due to the cost of this experiment. \n\n- Habitat: We only did 1 run with no further tuning due to the cost of this experiment. \n\n**Thank you again for your review.** We hope we have addressed the reviewer's questions and they are willing to increase their score. Please let us know if we can provide any additional clarification or information."}, {"review_id": "9wvlXn0H5g", "reviewer": "Reviewer_R5qx", "summary": "The paper builds upon ConvRNNs and proposes the use of SSM (State Space Models) as a replacement for RNNs. This allows for efficient computation using parallel scan. The proposed method is evaluated on the Long Horizon Moving-MNIST Generation and Long-range 3D Environment Benchmarks datasets, where it achieves promising results.", "questions": "1. The motivation for the proposed approach should be made more explicit, whether it is based on intuitive explanations or considerations of efficiency, among other factors.\n\n2. The experiments conducted in the paper are not comprehensive enough. It would be beneficial to compare the proposed method with other efficient attention-based approaches, such as kernel-based linear attention, 1+elu, performer, and cosformer, on long sequence tasks.\n\n3. More ablation studies are needed to validate the rationale behind the design, specifically regarding the convolutional aspect. The existing ablation studies primarily focus on validating the initialization, but it is equally important to investigate and verify the effectiveness of the convolutional components.\n\nReferences:\n[1] Fast Autoregressive Transformers with Linear Attention\n[2] Rethinking Attention with Performers\n[3] cosFormer: Rethinking Softmax in Attention", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The organization of this paper is clear and easy to understand. It starts from ConvRNN and replaces RNN with SSM, leading to the instantiation of ConvSSM known as ConvS5.", "weaknesses": "I think the main issue with this paper is the lack of motivation for the proposed approach, as replacing the matrix multiplication of SSM with convolution seems trivial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our paper and provide feedback.\n\n--**Re. motivation:** Thank you for this feedback. We were a little surprised since several other reviewers found the design well-motivated and clear (e.g. see Reviewer Kz49). Here we will walk through the motivation as presented in the paper. Please let us know where we can improve the presentation or provide further clarification.\n\nSpatiotemporal prediction methods need to scale efficiently with sequence length and effectively capture long-range dependencies in the spatiotemporal data. ConvRNNs have been historically popular due to their spatiotemporal modeling strengths and fast inference abilities. However, they are slow to train and suffer from vanishing/exploding gradients. On the other hand, Transformer methods provide strong performance but poor scaling in sequence length and slow inference. The recent SSM methods show subquadratic scaling in sequence length, ability to capture long-range temporal dependencies, and fast inference. Our method combines the best of ConvRNNs (strong spatiotemporal modeling abilities and fast inference) and SSMs (parallelizable and favorable scaling, long-range dependencies, and fast inference).\n\nWe hope this answer provides clarification, but if the reviewer has concrete suggestions for how we can further improve this exposition of the motivation we would love to incorporate their recommended changes!\n\n--**Re. contributions**:\nWe also note here that the move from SSMs to an effective version of a ConvSSM is not as trivial as simply replacing the SSM matrix-vector multiplications. Please see our detailed response regarding Technical Contributions in the General Response above where we outline the technical contributions required to provide a modern and scalable method.\n\n\n--**Re. experiments**: We were surprised by the feedback that our experiments were not comprehensive enough. We chose the most challenging and well-thought out long-range video prediction benchmarks that exist (the 3D Environment Benchmarks from TECO) and compare to state-of-the-art methods, many of which have been specifically designed for long-range video prediction.  We would like to note that reviewers Kz49 and xT6A complimented our paper for the broad benchmarking, comprehensive experiments and strong results of our paper and method. But we are also happy to include any additional baselines.\n\nAs requested, we add the Performer baseline on DMLAB below. We see that ConvS5 also outperforms the Performer. We also note that one of our original baseline methods, Perceiver AR is a modern efficient attention alternative that was published more recently than linear attention or Performer and concurrently with Cosformer. We also include S5 as a baseline, which has been shown to significantly outperform Linear attention, Performer, CosFormer and a host of other efficient attention alternatives on long-range sequence tasks (see Table 10 in [1] and Table F.1 in [2] and also compare these to the results in Table 4 of the Cosformer paper [3].).  \n\n**New Performer baseline on DMLab:**\n|             | FVD &darr; | PSNR &uarr; | SSIM &uarr; | LPIPS &darr; |\n|-------------|------------|-------------|-------------|--------------|\n| Perceiver-AR |  96     |      11.2      |    0.304     |    0.487       |\n| Performer   |     78     |     17.3   |    0.513   |      0.203    |\n| Transformer |     97     |     19.9    |    0.619    |     0.123    |\n| S5          |     221    |     19.3    |    0.641    |     0.162    |\n| ConvS5      |   **53**   |   **23.6**  |  **0.782**  |   **0.074**  |\n\nWe are also in the process of training Performer on the Moving MNIST experiments and will add this baseline to the final version as well.\n\n--**Re. ablations**: While we did not explicitly label it as an ablation, the S5 baselines also serve as the best ablation of the ConvS5 convolutional approach, as it is the closest possible variant that does not include ConvS5's convolutions. We have also added an ablation of the nolinearity choice that connects the ConvS5 layers. Please see the new ablation table and the discussion in the General Response. We will add this separate Ablation Table and discussion to make the ablations more clear.\n\n\n**We thank the reviewer again** for taking the time to review our paper. We hope we have addressed the reviewer's questions and the reviewer is willing to increase their score. Please let us know if we can provide additional clarification or information.\n\n**References:** [1] Efficiently Modeling Long Sequences with Structured State Spaces. [2] Simplified State Space Layers for Sequence Modeling. [3] CosFormer: Rethinking Softmax in Attention."}, {"review_id": "Wg8QWIpFog", "reviewer": "Reviewer_rhqo", "summary": "This paper proposes a new state space model for spatiotemporal modeling by introducing inductive bias of spatial locality. The core idea is to extend SSMs to ConvSSMs (just like extending FC-LSTM to ConvLSTM), which has a inherent convolutional structure. The new model also establishes an equivalence between the dynamics of particularly structured ConvSSMs and SSMs. Based on recent state space method S5, a model instantiation ConvS5 combines the stateful autoregressive generation with the ability to be parallelized across the sequence. A experimental evaluation shows that the proposed method captures spatiotemporal information better.", "questions": "* The idea of parallel scan for SSMs is not new. This design is considerable overlap with S5, which also leverages parallel scans to efficiently compute the states. The authors are suggested to discuss more about this.\n* In Table 3, ConvS5 shows a significant improvement than S5 without TECO. When trained using the TECO, ConvS5 only achieves a comparable result with S5. It makes me wonder if the proposed method is effective consistently on various settings (e.g. framework, SSMs backbone, etc). \n* As in Weakness 3, I wonder if the proposed method is applicable to other SSMs. The authors are suggested to discuss more about this.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The idea is simple but effective as it combines fast autoregressive generation and parallelized process.\n* The proposed method achieves better performance than Transformer and ConvRNN on Moving-MNIST prediction task. It also strikes a balance on computational complexity between the two methods.", "weaknesses": "* The lack of different ConvSSM variants. The paper proposes ConvSSMs to address long-range spatiotemporal modeling. However, only a variant, ConvS5, is present and evaluated. It is not convincing to support the basic idea of ConvSSMs. More variants are needed. \n* The experiments about comparison with Transformer and ConvLSTM could be more extensive. In Table 2, we see that the ConvS5 outperforms the other methods across the metrics. We notice that ConvLSTM achieves a slightly lower performance result. However, ConvLSTM shows a slightly better sampling speed than ConvS5. It is not sufficiently significant that ConvS5 has a better quality-speed tradeoff than ConvLSTM. \n* There lacks explanation and experiments to the design of convolutional operator upon SSMs. The comparison results of whether convolution or not are missing. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback. We are glad the reviewer appreciated the simplicity and effectiveness of the proposed method.\n\n--**Re. convSSM variants:** We present the general ideas of ConvSSMs (tensor-valued states, linear transitions and continuous-time parameterization) in Section 3.1, similar to how RNNs and ConvRNNs are general model classes with different variants. We note here that the general ConvSSM formulation of Section 3.1 could be run sequentially as is, similar to ConvRNNs. This could be considered a Vanilla ConvSSM variant. The downside of course is that this would be slow to train.\n\nWe then work through the steps required to get a practical, effective, and modern version of a ConvSSM that allows efficient parallelization (Section 3.2), enables modeling long-range dependencies (Sections 3.3 and 3.4), while retaining the fast autoregressive generation abilities of ConvRNNs. The result is the ConvS5 method that we propose.\n\nNote that various recent extensions to SSMs such as data-dependent gating (H3 [1]) and data-dependent dynamics (Liquid SSMs [2]) could be applied to provide different extensions of ConvS5 (and thus different ConvSSM variants).  We can add a brief discussion of these potential extensions and variants in the discussion. However, our main focus in this work is to show the base ConvS5 architecture is effective. \n\nWhile the ConvS5 formulation we develop is simple and effective, we hope our work inspires others to consider the possibilities for this type of linear SSM and ConvRNN-inspired architecture, as well as other alternatives.\n\nIf the reviewer has suggestions to improve the exposition of the general idea of ConvSSMs and the particular formulation of ConvS5, we are happy to incorporate these changes to make these distinctions more clear.\n\n--**Re. more comparison with Transformer and ConvLSTM:** We compare 3 aspects between these models. We believe they illustrate the main benefits of our method. \n\n1. Performance: The experiments on Moving-MNIST show that \nConvS5 is able to significantly outperform ConvLSTM across the metrics (cutting FVD nearly in half while also outperforming by a fair margin on PSNR/SSIM/LPIPS) when trained on the longer context (600 frames). This indicates that ConvS5 is able to take better advantage of the longer context and better capture the long-range dependencies than ConvLSTM. \n\n2. Faster training: Because ConvS5 can be parallelized across the sequence, like Transformer, it trains much faster than the sequential ConvLSTM (>3X faster for these experiments, see Table 6 in the Appendix). \n\n3. Fast autoregressive generation: Like ConvLSTM, ConvS5 has a much faster autoregressive generation speed than the Transformer. It is true that ConvLSTM had a slight edge (557x vs 427x faster than the Transformer) in these experiments, but the main point is that both methods are orders of magnitude faster than the Transformer. \n\nAlso, ConvLSTM and ConvS5 are both stateful recurrent-based methods, so there is nothing that makes ConvLSTM inherently faster at autoregressive generation than ConvS5. Since both ConvS5 and ConvLSTM were much faster than the Transformer, and ConvS5 trains much faster than ConvLSTM, we did not spend time optimizing the speed of the generation process. Nonetheless, given the results from this experiment, there are likely many applications in which halving the FVD for a slightly slower inference speed (though still much faster than Transformer) would be worth it.\n\nWe welcome suggestions from the reviewer for any additional comparisons to add here. \n\n\n--**Re. design and ablation of convolutions**: We discuss in detail the design of the convolution operators of ConvS5 in Section 3.2-3.4. Please let us know if we can provide further clarification on this point. \n\nIn terms of ablating the convolutions, while we did not explicitly label it as an ablation, the S5 baselines in the DMLab experiments serve as this ablation since S5 is the closest possible variant to ConvS5 without the convolution structure. We will make a separate ablations table and discussion of these points to make this more clear. Please see the new ablations table and discussion in the General Response above. \n\n### Questions:\n\n--**Re. the parallel scan design:** We propose a parallel scan for convolutional recurrences which is different from the one in S5. S5's parallel scan is for a linear recurrence while ConvS5's is for a convolutional recurrence.  There were several challenges for applying a parallel scan to a convolutional recurrence. We discuss the details of this in Section 3.2.  Section 2.3 explains the parallel scan used in S5. We would also be happy to incorporate suggestions the reviewer has on this point to help provide further clarification.\n\n--**Re. the effect of the TECO Framework:** We agree with the reviewer that the TECO framework provides S5 a large performance improvement on FVD (though still slightly below ConvS5). However we note that even with the TECO framework, ConvS5 still significantly outperforms S5 on the other long-range consistency metrics (PSNR/SSIM/LPIPS), so we respectfully disagree on the point that TECO-S5 achieves a comparable overall result to TECO-ConvS5. Therefore, ConvS5 performs well with and without the TECO framework, while S5 only performs well with the TECO framework. This would seem to suggest ConvS5 might be expected to generalize well to other settings since its performance on these experiments is not dependent on the framework.\n\n--**Re. other SSMs methods**: Please see the response of ConvSSM variants above.\n\n**We again thank the reviewer** for taking the time to review our paper. We hope we have answered the reviewer's questions and the reviewer is willing to increase their score. Please let us know if we can provide additional clarification or information.\n\n**References:** [1] Hungry Hungry Hippos: Towards Language Modeling with State Space Models [2] Liquid Structural State-Space Models"}, {"review_id": "JY7JK4Ty1r", "reviewer": "Reviewer_xT6A", "summary": "This paper presents ConvS5, a convolutional state-space model that aims to model long-range spatiotemporal dependencies in video data. They extend the prior S5 model to operate in a convolutional state-space, and retain the long-range benefits of state-space models by using a point-wise convolutional kernel for diagonalizable and S5 HiPPO initialization. The resulting model can then be efficiently trained using parallel scans due to linearity in the recurrent mechanism. They show strong results on several complex long-range video benchmarks in Moving MNIST, DMLab, Minecraft, and Habitat.", "questions": "- For Moving MNIST, how would ConvS5 potentially compare against temporally hierarchical methods in modeling long-range dependencies (e.g. Clockwork-VAE)? Or would an analogous hierarchical version of ConvS5 be better than a hierarchical GRU model?\n- How would ConvS5 compare for Minecraft / Habitat? I believe the primary difference is without the MaskGit / iterative decoding on the prior, so would ConvS5 struggle with more stochastic environments?\n- What is the potential cause of the difference  in performance between FVD and PSNR / SSIM / LPIPS in Minecraft? Is FVD lower due to better frame fidelity of TECO-ConvS5, and TECO-transformer has better long-term consistency? Or perhaps other factors?\n- Are generated video samples available anywhere to view? It is hard to evaluate temporal coherence in the rows of frames provided in the appendix.", "limitations": "Yes, the authors discuss limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is generally clear and well written\n- Comprehensive experiments on a range of datasets focused on long-range understanding in videos, showing that their proposed method is better than or competitive with SOTA methods\n- Experiments show that better scaling in sequence length allows for training on longer sequences, which results in better / more stable generation\n- Faster inference speed compared to existing methods while retaining high generation fidelity", "weaknesses": "- Overall technical novelty is a little low, as the model itself is a straightforward extension of the S5 model to a convolutional setting.\n- A primary concern I have about recurrent-based models for long-range understanding in video generation is limited capacity in the recurrent state, as these models are required to store all fine-grained details about the scene in the recurrent state (since potentially every part of the scenes will be visited / need to be reconstructed at a future timestep), whereas transformer-based models can dynamically retrieve from all prior states for each timestep. A rigorous study on this aspect would benefit the paper greatly, as this issue is more present in visually complex videos (minecraft -> habitat -> real videos), and could be a potential bottleneck to the scalability of these models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback. We are glad the reviewer appreciated the comprehensive experiments as well as strong results. \n\n**--Re. technical novelty**: We respectfully disagree that the proposed approach is a straightforward extension of S5. Please see the Technical Contributions section of the General Response for a detailed explanation of the contributions required to ensure the convolutional recurrence could be efficiently parallelized and model long-range spatiotemporal dependencies. \n\n**--Re. recurrent vs attention architectures**: Thank you for the question. This is an interesting discussion. While it is true that recurrent methods are required to store the history in its recurrent states, this capacity can be increased by adding more hidden units or layers. On the other hand, the attention mechanisms have bounded context by their design which limits the amount of history they can retrieve. Therefore, it is not necessarily true that the transformer-based models can retrieve from all prior states for each timestep or that recurrent models are a bottleneck for scalability.\n\nIn addition, recent architecture improvements have been proposed for SSMs in language modeling, such as multiplicative gating/routing operators (e.g. in H3 [1]) to allow for recall operations similar to those you have suggested. These ideas can be applicable to ConvSSMs as well. Another possibility is to apply chunked/bounded context attention mechanisms to the outputs of an SSM/ConvSSM as suggested in MEGA [2].\n\nHowever, the main focus of this paper was to develop the base ConvSSM architecture. Our evaluation on the long-range video benchmarks (the most challenging ones that currently exist) show that ConvS5 performs comparably or better than Transformers on these tasks. We believe this provides some evidence that this type of method is capable of scaling to more complex long-range datasets. \n\nWe would be happy to include these discussions and potential modifications in the final version. Please let us know if the reviewer has suggestions for an additional study they would like to see in this regard.  \n\n### Questions:\n\n**--Re. comparing against temporally hierarchical methods:** Thank you for your suggestion!  We added the CW-VAE comparison to the MNIST table for models trained on 600 sequence length below. We will add a table for models trained on 300 sequence lengths for the final version (we could not finish both before the rebuttal deadline).  We found that ConvS5 outperforms CW-VAE in this setting across the metrics with both models having a similar number of parameters. This result indicates that ConvS5 itself compares well to temporally hierarchical methods in modeling long-range dependencies. However, we also agree that it could be interesting to combine temporally hierarchical methods with ConvSSM methods. We can include a discussion of this in the final version.\n\n**New CW-VAE baseline trained on 600 Frames:**\n|             |            |  100 &rarr; | 800         |              |            |  100 &rarr; | 1200        |              |\n|-------------|------------|------------:|-------------|--------------|------------|------------:|-------------|--------------|\n|             | FVD &darr; | PSNR &uarr; | SSIM &uarr; | LPIPS &darr; | FVD &darr; | PSNR &uarr; | SSIM &uarr; | LPIPS &darr; |\n| Transformer |   **42**   |     13.7    |    0.672    |     0.207    |     91     |     13.1    |    0.631    |     0.252    |\n| ConvLSTM    |     91     |     15.5    |    0.757    |     0.149    |     137    |     14.6    |    0.727    |     0.180    |\n| CW-VAE      |     93     |     12.5    |    0.599    |    0.268     |     109    |     12.4    |    0.590    |     0.280    |\n| ConvS5      |     47     |   **16.4**  |  **0.788**  |   **0.134**  |   **71**   |   **15.6**  |  **0.763**  |   **0.162**  |\n\n**--Re. convS5 comparison for Minecraft/Habitat**: \nWe were not able to add the comparisons of (non-TECO) Transformer, S5, and ConvS5, because they are a lot more expensive to run than the TECO variants. Note that the size of the datasets and the model size used for these datasets are much bigger than DMLAB. However, we believe that the trend would be similar to the DMLAB non-TECO comparisons. \n\nNevertheless, we may be able to try the comparison with smaller model sizes for these datasets if the reviewer would like to see the relative comparisons of the models for the final version. Please let us know. \n\n**--Re. performance difference between metrics on Minecraft**: \nWe are not confident to make any strong conclusion about the difference in performance between FVD vs PSNR/SSIM/LPIPS. However, the gap between TECO-Transformer and TECO-ConvS5 is small for all metrics (especially relative to the baselines), and the visual differences in the prediction between these models were not significant. Our conclusion is that the performance of these models is comparable on this dataset, but our model has a faster sampling speed than TECO-Transformer (1.7x at this sequence generation length; note that the TECO-ConvS5 speed will stay constant for generating longer sequences while the TECO-Transformer speed will decrease).  \n\n**--Re. link to videos**: \nIn the appendix of the original submission (Line 746 in Appendix C), we included a link to an anonymized website that includes the video samples. We will move this link to the main paper to make this clear. Please see the link in the appendix. Note that we cannot include any external link in our response due to NeurIPS rebuttal policy. \n\n**We thank the reviewer again** for the time they have taken to review our paper and read this rebuttal. We hope we have addressed the reviewer's questions and the reviewer is willing to increase their score. Please let us know if we can provide additional clarification or information.\n\n**References:**\n [1] Hungry Hungry Hippos: Towards Language Modeling with State Space Models. [2] Mega: Moving Average Equipped Gated Attention."}]

Author rebuttal:
# General Response
We thank the reviewers for reviewing our submission and providing constructive feedback. We provide a general response here and respond to each reviewer individually. We presented the ConvS5 spatiotemporal sequence model which has parallelizable training, fast autoregressive inference, and effectively captures long-range spatiotemporal dependencies.

We were pleased all reviewers agree that the paper is **1)** easy to follow. Also, reviewers Kz49, rhqo, and xT6A appreciated  **2)** the effective model design for fast autoregressive generation and the parallelization process, and **3)** broad benchmarking and strong empirical results. 

## Technical Contribution
There was some concern from reviewers P1WH and xT6A regarding technical novelty (though some reviewers such as Kz49 praised ConvS5's well motivated and principled design based on theoretical and practical arguments). We agree that the basic idea of going from SSM to ConvSSM by replacing the SSM's matrix-vector multiplications with convolutions is relatively straightforward. However, there are challenges to make this approach scalable and effective for modeling long-range spatiotemporal data. 

**1. Computational efficiency, parallelization across the sequence for fast training and inference.** For the linear recurrence of SSMs, there are different ways to parallelize the model across the sequence (e.g. FFTs or parallel scans). However, a parallel scan for convolutional recurrences has not been studied before. In this paper, we are the first to introduce the parallelization of convolutional recurrences using a binary associative operator. In Section 3.2, we show both theoretical (Proposition 1) and practical results required to make this feasible and efficient.

**2. Capture long-range spatiotemporal dependencies.** We developed a theoretical connection between the dynamics of SSMs and ConvSSMs (Proposition 3, Section 3.3). Based on this result, in Section 3.4, we are able to introduce a parameterization and initialization design that allows ConvS5's convolutional recurrence to capture long-range spatiotemporal dependencies.

**Result:** The result of these contributions is ConvS5 which is parallelizable and overcomes difficulties during training (e.g., vanishing/exploding gradient problems) that traditional ConvRNN approaches experience. It also provides fast (constant time and memory per step) autoregressive generation compared to Transformers. Furthermore, our empirical results validate these design choices for modeling long-range spatiotemporal dependencies.

We argue that these contributions are not trivial, and our paper provides a rigorous framework that ensures both computational efficiency and modeling performance for spatiotemporal sequence modeling.

## Computational Cost
The appendix of the paper includes the experimental details including the number of parameters, training and inference times for all experiments (Tables 5-19). As a couple of reviewers asked us to provide FLOPs, we provide the comparison table below for Moving-MNIST trained on 600 frames. We will include the full summary in the final paper. 

Although ConvS5 requires a few more FLOPs due to the convolution computations and some architecture choices (ResNet blocks for nonlinearity), our model is parallelizable during training (unlike ConvLSTM) and has fast autoregressive generation (unlike Transformer) ---  training 3x faster than ConvLSTM and generating samples 400x faster than Transformers. Note that the number of parameters are comparable.

|             | GFLOPS &darr; | Parallelizable     | Train Step Time (s) &darr; | Train cost (V100 days) &darr; | Sampling Speed (frames/s) &uarr; |
|-------------|:------------:|:------------:|:--------------:|:---------------------------:|:---------------------------------:|
| Transformer |      70.0     |       o      |         0.77 (1.0x)         |    50       |            0.21 (1.0x)            |
| ConvLSTM    |      64.9     |       x       |          3.0 (3.9x)       |    150      |             117 (557x)            |
| ConvS5      |      96.8     |       o      |         0.93 (1.2x)         |    50       |             90 (429x)             |

## Ablations
A common request was the ablation of the convolutional structure of ConvS5. We note that the S5 baseline also serves as an ablation of ConvS5 since it is the closest possible model without ConvS5's convolutional structure, even though we did not explicitly label this as an ablation. We will add a separate ablation table in the final version and make this more clear.

We additionally include the ablations of the nonlinear connections used between the ConvS5 layers. We use ResNet blocks but also consider elementwise GELU and GLU activations (used in the S5 paper). Due to the time limit, we only show the nonlinearity ablation for non-TECO models here. We will include the rest in the final version. 

**DMLAB Ablations**
|                    | conv. | nonlinearity | FVD &darr; | PSNR &uarr; | SSIM &uarr; | LPIPS &darr; |
|--------------------|---|----------|------------|-------------|-------------|--------------|
| S5                 | x |   Glu    |     221    |     19.3    |    0.641    |     0.162    |
| ConvS5             | o | Gelu     |     129   |     21.4    |    0.709    |     0.110    |
| ConvS5             | o | GLU      |     112    |     21.6    |    0.720    |     0.098    |
| ConvS5 (ours)      | o | ResNet   |     53     |     23.6    |    0.782    |     0.074    |
| TECO-S5            | x |  Glu    |     35     |     20.1    |    0.687    |     0.143    |
| TECO-ConvSSM (random init.)      | o |    ResNet      |     44     |     21.0    |    0.691    |     0.010    |
| TECO-ConvS5 (ours) |  o |     ResNet    |     31     |     23.8    |    0.803    |     0.085    |

**We once again thank the reviewers for their time and positive and constructive feedback**. We will now respond to individual comments.

--The ConvS5 authors

Meta-review: This paper proposes a convolutional state space model that aims to overcome the efficiency concerns of previous methods, including ConvLSTM and transformers, for long-range spatiotemporal modeling.

The paper is generally clear and well written. The motivation for this work is presented well and the proposed model is based on a principled design. Strong results from the experiments are reported, showing that the proposed method can achieve significant speedup. Many of the comments given by the five reviewers were related to the experiments, demanding more extensive experimentation. We thank the authors for their effort in conducting further experiments and answering the questions in detail. Their effort has led four of the five reviewers to upgrade their overall ratings.

While we lean towards acceptance of this paper, I highly recommend the authors to consider the comments and suggestions of the reviewers thoroughly in revising their paper to make it appeal better to the research community.
Decision: Accept (poster)