SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\boldsymbol {\mathcal {X}} ^ {\prime} (t) _ {i, j, k} = \sum_ {l = 0} ^ {P - 1} \boldsymbol {\mathcal {A}} _ {k, l, 0, 0} \boldsymbol {\mathcal {X}} (t) _ {i, j, l} + \sum_ {q = 0} ^ {U - 1} \sum_ {m = 0} ^ {k _ {B} - 1} \sum_ {n = 0} ^ {k _ {B} - 1} \boldsymbol {\mathcal {B}} _ {k, q, m, n} \boldsymbol {\mathcal {V}} (t) _ {i, j, q, m, n}. \tag {22}
$$
Let  $\mathbf{A}_{\mathrm{SSM}} \in \mathbb{R}^{P \times P}$  be a matrix with rows,  $\mathbf{A}_{\mathrm{SSM},i} \in \mathbb{R}^P$ , corresponding to a flattened version of the output features of  $\mathcal{A}$ , i.e.  $\mathcal{A}_i \in \mathbb{R}^{P \times 1 \times 1}$ . Similarly, reshape  $\mathcal{B}$  into a matrix  $\mathbf{B}_{\mathrm{SSM}} \in \mathbb{R}^{P \times (Uk_B^2)}$  where the rows,  $\mathbf{B}_{\mathrm{SSM},i} \in \mathbb{R}^{Uk_B^2}$  correspond to a flattened version of the output features of  $\mathcal{B}$ , i.e.  $\mathcal{B}_i \in \mathbb{R}^{U \times k_B \times k_B}$ .
Then we can rewrite (22) equivalently as
$$
\begin{array}{l} \boldsymbol {\mathcal {X}} ^ {\prime} (t) _ {i, j, k} = \sum_ {l = 0} ^ {P - 1} \boldsymbol {\mathcal {A}} _ {k, l, 0, 0} \boldsymbol {\mathcal {X}} (t) _ {i, j, l} + \sum_ {q = 0} ^ {U - 1} \sum_ {m = 0} ^ {k _ {B} - 1} \sum_ {n = 0} ^ {k _ {B} - 1} \boldsymbol {\mathcal {B}} _ {k, q, m, n} \boldsymbol {\mathcal {V}} (t) _ {i, j, q, m, n} (23) \\ = \sum_ {l = 0} ^ {P - 1} \mathbf {A} _ {\text {S S M}, k, l} \boldsymbol {\chi} (t) _ {i, j, l} + \sum_ {v = 0} ^ {U k _ {B} ^ {2} - 1} \mathbf {B} _ {\text {S S M}, k, v} \boldsymbol {\mathcal {U}} _ {\text {i m 2 c o l}} (t) _ {i, j, v} (24) \\ = \mathbf {A} _ {\mathrm {S S M}, k} ^ {T} \boldsymbol {\chi} (t) _ {i, j} + \mathbf {B} _ {\mathrm {S S M}, k} ^ {T} \boldsymbol {\mathcal {U}} _ {\mathrm {i m 2 c o l}} (t) _ {i, j} (25) \\ \end{array}
$$
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  $\mathbf{A}_{\mathrm{S5}} = \mathbf{V}\boldsymbol{\Lambda}_{\mathrm{S5}}\mathbf{V}^{-1} \in \mathbb{R}^{P \times P}$  where  $\boldsymbol{\Lambda}_{\mathrm{S5}} \in \mathbb{C}^{P \times P}$  is a complex-valued diagonal matrix and  $\mathbf{V} \in \mathbb{C}^{P \times P}$  corresponds to the eigenvectors. Defining  $\tilde{\mathbf{x}}(t) = \mathbf{V}^{-1}\mathbf{x}(t)$ ,  $\tilde{\mathbf{B}} = \mathbf{V}^{-1}\mathbf{B}$ , and  $\tilde{\mathbf{C}} = \mathbf{C}\mathbf{V}$  we can reparameterize the SSM of (3) as the diagonalized system:
$$
\frac {\mathrm {d} \tilde {\mathbf {x}} (t)}{\mathrm {d} t} = \boldsymbol {\Lambda} _ {\mathrm {S} 5} \tilde {\mathbf {x}} (t) + \tilde {\mathbf {B}} \mathbf {u} (t), \quad \mathbf {y} (t) = \tilde {\mathbf {C}} \tilde {\mathbf {x}} (t) + \mathbf {D} \mathbf {u} (t). \tag {26}
$$
S5 uses learnable timescale parameters  $\Delta \in \mathbb{R}^P$  and the following zero-order hold (ZOH) discretization:
$$
\overline {{\boldsymbol {\Lambda}}} _ {\mathrm {S} 5} = \operatorname {D I S C R E T I Z E} _ {\mathrm {A}} \left(\boldsymbol {\Lambda} _ {\mathrm {S} 5}, \boldsymbol {\Delta}\right) := e ^ {\boldsymbol {\Lambda} _ {\mathrm {S} 5} \boldsymbol {\Delta}} \tag {27}
$$
$$
\overline {{\mathbf {B}}} _ {\mathrm {S} 5} = \operatorname {D I S C R E T I Z E} _ {\mathrm {B}} \left(\boldsymbol {\Lambda} _ {\mathrm {S} 5}, \tilde {\mathbf {B}}, \boldsymbol {\Delta}\right) := \boldsymbol {\Lambda} _ {\mathrm {S} 5} ^ {- 1} \left(\overline {{\boldsymbol {\Lambda}}} _ {\mathrm {S} 5} - \mathbf {I}\right) \tilde {\mathbf {B}} \tag {28}
$$
S5 Initialization S5 initializes its state matrix by diagonalizing  $\mathbf{A}_{\mathrm{S5}}$  as defined here:
$$
\mathbf {A} _ {\mathrm {S} 5 _ {n k}} = - \left\{ \begin{array}{l l} (n + \frac {1}{2}) ^ {1 / 2} \left(k + \frac {1}{2}\right) ^ {1 / 2}, & n > k \\ \frac {1}{2}, & n = k. \\ (n + \frac {1}{2}) ^ {1 / 2} \left(k + \frac {1}{2}\right) ^ {1 / 2}, & n <   k \end{array} \right. \tag {29}
$$
This matrix is the normal part of the normal plus low-rank HiPPO-LegS matrix from the HiPPO framework [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 approximating long-range dependencies at initialization with respect to an infinitely long, exponentially-decaying measure [106]. Gupta et al. [41] empirically showed that the low-rank terms could be removed without 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-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-inspired matrices as well as an important normalization effect provided by the explicit discretization procedure. They also introduce a purely discrete-time parameterization that can perform similarly to the continuous-time discretization of S4 and S5. However, their parameterization 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.
B.2 ConvS5 Diagonalization
We leverage S5's diagonalized parameterization to reduce the cost of the parallel scan of ConvS5.
Concretely, we initialize  $\mathbf{A}_{\mathrm{S5}}$  as in (29) and diagonalize as  $\mathbf{A}_{\mathrm{S5}} = \mathbf{V}\boldsymbol{\Lambda}_{\mathrm{S5}}\mathbf{V}^{-1}$ . To apply ConvS5, we compute  $\overline{\mathbf{A}}_{\mathrm{S5}}$  and  $\overline{\mathbf{B}}_{\mathrm{S5}}$  using (27-28), and then form the ConvS5 state and input kernels:
$$
\overline {{\boldsymbol {\Lambda}}} _ {\mathrm {S} 5} \in \mathbb {R} ^ {P \times P} \xrightarrow {\text {r e s h a p e}} \overline {{\boldsymbol {\mathcal {A}}}} _ {\text {C o n v S} 5} \in \mathbb {R} ^ {P \times P \times 1 \times 1} \tag {30}
$$
$$
\overline {{\mathbf {B}}} _ {\mathrm {S} 5} \in \mathbb {R} ^ {P \times \left(U k _ {B} ^ {2}\right)} \xrightarrow {\text {r e s h a p e}} \overline {{\boldsymbol {\mathcal {B}}}} _ {\text {C o n v S} 5} \in \mathbb {R} ^ {P \times U \times k _ {B} \times k _ {B}}. \tag {31}
$$
See Listing 1 for an example of the core implementation. Note, the state kernel  $\overline{\mathcal{A}}_{\mathrm{ConvS5}}$  is "diagonalized" in the sense that all entries in the state kernel are zero except  $\overline{\mathcal{A}}_{\mathrm{ConvS5},i,i} = \overline{\Lambda}_{\mathrm{S5},i,i}\forall i\in [P]$
This means that the pointwise convolutions reduce to channel-wise multiplications. However, this does not reduce expressivity compared to a full pointwise convolution. This is because, given the ConvSSM to SSM equivalence of Proposition 3 and the use of complex-valued kernels, the diagonalization 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.
import jax
import jax numpy as np
from ConvSSMHelpers import discretize, Conv2D, ResNet_DBlock
parallel_scan = jax.laxassociative_scan
def apply_ConvS5_layer(A, B, B_shape, C_kernel, log_Delta, resnet_parameters, us, x0):
    """Compute the outputs of ConvS5 layer given input sequence.
Args:
    A(complex64): S5 state matrix (P.)
    B(complex64): S5 input matrix (P,Uk_B^2)
    B_shape (tuple): shape of B_kernel)
    C_kernel(complex64) output kernel (U,P,k_C,k_C)
    log_Delta(float32): learnable timescale params (P)
    resnet_parameters (dict): ResNet block params
    us(float32): input sequence of features (L,bsz,H,W,U)
    x0(complex64): initial state (bsz,H,W,P)
Returns:
    outputs(float32): the ConvS5 layer outputs (L,bsz,H,W,U)
    x_L(complex64): the last state of the ConvSSM (bsz,H,W,P)
    """
    # Discretize and reshape ConvS5 state and input kernels
    P, U, k_B = B_shape
    A_bar, B_bar = discretize(A, B, np.exp(log_Delta))
    A_kernel = A_bar # already correct shape due to diagonalization
    B_kernel = B_bar.reshape(P, U, k_B, k_B)
    # Apply ConvS5
    xs = apply_ConvS5(A_kernel, B_kernel, C_kernel, us, x0)
    # Apply ResNet activation function
    outputs = jax.vmap(ResNet_DBlock, axis=(None,0)) (resnet_parameters, ys)
    return outputs, xs[-1]
def apply_ConvS5(A_kernel, B_kernel, C_kernel, us, x0):
    """Compute the output sequence of the convolutional SSM
given the input sequence using a parallel scan.
Computes x_k = A * x_{<1} + B * u_k
y_k = C * x_k
where * is a convolution operator.
Args:
    A_kernel(complex64): state kernel (P.)
    B_kernel(complex64): input kernel (P,U,k_B,k_B)
    C_kernel(complex64): output kernel (U,P,k_C,k_C)
    us(float32): input sequence (L,bsz,H,W,U)
    x0(complex64): initial state (bsz,H,W,P)
Returns:
    ys(float32): the convS5 outputs (L,bsz,H,W,U)
    x_L(complex64): the last state (bsz,H,W,P)
   """
    # Compute initial scan elements
As = np.reduce(A_kernel[None,...], us.shape[0], axis=0)
Bus = jax.vmap(Conv2D)(B_kernel, npcomplex64(us))
Bus = Bus.at[0].add(npexpand_dims(A_bar,(0,1,2)) * x0)
# Convolutional recurrence with parallel scan _
_, xs = parallel_scan(conv_binary_operator, (As,Bus))
# Compute ConvS5 outputs
ys = jax.vmap(Conv2D)(C_kernel, xs).real
return ys, xs
def conv_binary_operator(q_i, q_j):
    """Binary operator for convolutional recurrence with "diagonalized" 1X1 state kernels.
Args:
    q_i, q_j (tuples): scan elements q_i=(A_i, BU_i) where
        A_i(complex64) is state kernel (P,)
BU_i(complex64) is effective input (bsz,H,W,P)
Returns:
    output tuple q_i \ circledast q_j
   """
    A_i, BU_i = q_i
A_j, BU_j = q_j
# Convolve "diagonal" 1X1 kernels
AA = A_j * A_i
# Convolve "diagonal" A_j with BU_i
A_jBU_i = np-expand_dims(A_j,(0,1,2)) * BU_i
return AA, A_jBU_i + BU_j
Listing 1: JAX implementation of core code to apply a single ConvS5 layer to a batch of spatiotemporal input sequences.
1   
2   
3   
4   
5   
6   
7   
8   
9   
10   
11   
12   
13   
14   
15   
16   
17   
18   
19   
20   
21   
22   
23   
24   
25   
26   
27   
28   
29   
30   
31   
32   
33   
34   
35   
36   
37   
38   
39   
40   
41   
42   
43   
44   
45   
46   
47   
48   
49   
50   
51   
52   
53   
54   
55   
56   
57   
58   
59   
60   
61   
62   
63   
64   
65   
66   
67   
68   
69   
70   
71   
72   
73   
74   
75   
76   
77
C Supplementary Results
We include expanded tables and sample trajectories from the experiments in the main paper. Sample videos can be found at:
https://sites.google.com/view/convssm.
C.1 Moving-MNIST
Table 6: 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 and 1200 frames. ConvSSM (ablation) is performed by randomly initializing the state kernel (see Section 5.3 and Appendix D.4).
Trained on 300 frames
Method	Params	100 → 400	100 → 800	100 → 1200
FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓
Transformer	164M	73 ± 3	13.5 ± 0.1	0.669 ± 0.002	0.213 ± 0.003	159 ± 7	12.6 ± 0.1	0.609 ± 0.002	0.287 ± 0.001	265 ± 8	12.4 ± 0.1	0.591 ± 0.002	0.321 ± 0.002
Performer	164M	111 ± 9	13.4 ± 0.1	0.653 ± 0.002	0.288 ± 0.001	234 ± 1	13.4 ± 0.1	0.652 ± 0.006	0.379 ± 0.002	275 ± 5	13.2 ± 0.1	0.592 ± 0.001	0.393 ± 0.001
CW-VAE	20M	78 ± 1	12.7 ± 0.1	0.611 ± 0.002	0.254 ± 0.001	104 ± 2	12.4 ± 0.1	0.592 ± 0.002	0.277 ± 0.002	117 ± 2	12.3 ± 0.1	0.585 ± 0.002	0.286 ± 0.001
ConvLSTM	20M	57 ± 3	16.9 ± 0.2	0.796 ± 0.004	0.113 ± 0.002	128 ± 4	15.0 ± 0.1	0.737 ± 0.003	0.169 ± 0.001	187 ± 6	14.1 ± 0.1	0.706 ± 0.003	0.203 ± 0.001
ConvSSM (ablation)	20M	67 ± 3	15.5 ± 0.1	0.742 ± 0.001	0.168 ± 0.001	287 ± 5	13.6 ± 0.1	0.577 ± 0.001	0.293 ± 0.001	511 ± 8	13.3 ± 0.1	0.515 ± 0.001	0.348 ± 0.001
ConvS5	20M	26 ± 1	18.1 ± 0.1	0.830 ± 0.003	0.094 ± 0.002	72 ± 3	16.0 ± 0.1	0.761 ± 0.005	0.156 ± 0.003	187 ± 5	14.5 ± 0.1	0.678 ± 0.003	0.230 ± 0.004
Trained on 600 frames
Method	Params	100 → 400	100 → 800	100 → 1200
FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FV D ↓	PSNR ↑	SSIM ↑	LPIPS ↓
Transformer	164M	21 ± 1	15.0 ± 0.1	0.741 ± 0.002	0.138 ± 0.001	42 ± 2	13.7 ± 0.1	0.672 ± 0.002	0.207 ± 0.003	91 ± 6	13.1 ± 0.1	0.631 ± 0.004	0.252 ± 0.002
Performer	164M	27 ± 1	13.1 ± 0.1	0.654 ± 0.004	0.206 ± 0.001	93 ± 5	12.4 ± 0.1	0.616 ± 0.002	0.274 ± 0.001	243 ± 7	12.2 ± 0.1	0.608 ± 0.001	0.312 ± 0.002
CW-VAE	20M	73 ± 2	12.9 ± 0.1	0.621 ± 0.004	0.242 ± 0.001	94 ± 3	12.5 ± 0.9	0.598 ± 0.004	0.269 ± 0.001	107 ± 2	12.3 ± 0.1	0.590 ± 0.004	0.280 ± 0.002
ConvLSTM	20M	39 ± 5	17.3 ± 0.2	0.812 ± 0.005	0.100 ± 0.003	91 ± 7	15.5 ± 0.2	0.757 ± 0.005	0.149 ± 0.003	137 ± 9	14.6 ± 0.1	0.727 ± 0.004	0.180 ± 0.003
ConvSSM (ablation)	20M	81 ± 6	15.5 ± 0.1	0.743 ± 0.002	0.163 ± 0.003	145 ± 8	14.3 ± 0.1	0.696 ± 0.002	0.218 ± 0.002	215 ± 9	13.4 ± 0.1	0.614 ± 0.001	0.287 ± 0.001
ConvS5	20M	23 ± 3	18.1 ± 0.1	0.832 ± 0.003	0.092 ± 0.003	47 ± 7	16.4 ± 0.1	0.788 ± 0.002	0.134 ± 0.003	71 ± 9	15.6 ± 0.1	0.763 ± 0.002	0.162 ± 0.003
Table 7: Model runtime comparison for Moving-MNIST results in Table 6. ConvS5 can be parallelized like a Transformer but maintains the constant cost-per-step autoregressive generation of ConvRNNs.
Method	Parallelizable	100 → 400 Train Step Time (s) ↓	100 → 800 Sample Throughput (frames/s) ↑	100 → 1200 Sample Throughput (frames/s) ↑	Sample Throughput (frames/s) ↑
Transformer	YES	0.77 (1.0×)	1.1 (1.0×)	0.34 (1.0×)	0.21 (1.0×)
ConvLSTM	NO	3.0 (3.9×)	117 (106×)	117 (345×)	117 (557×)
ConvS5	YES	0.93 (1.2×)	90 (82×)	90 (265×)	90 (429×)
(a) Example 1
(b) Example 2
(c) Example 3
Figure 3: Moving-MNIST Samples: 1200 frames generated conditioned on 100.
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).
Method	DMLab
Params	FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓
FitVid*	165M	176 ± 4.86	12.0 ± 0.013	0.356 ± 0.00171	0.491 ± 0.00108
CW-VAE*	111M	125 ± 7.95	12.6 ± 0.059	0.372 ± 0.00033	0.465 ± 0.00156
Perceiver AR*	30M	96.3 ± 3.64	11.2 ± 0.004	0.304 ± 0.00004	0.487 ± 0.00123
Latent FDM*	31M	181 ± 2.20	17.8 ± 0.111	0.588 ± 0.00453	0.222 ± 0.00493
Transformer	152M	97.0 ± 5.98	19.9 ± 0.108	0.619 ± 0.00506	0.123 ± 0.00191
Performer	152M	80.3 ± 3.21	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	66.6 ± 4.81	23.2 ± 0.053	0.769 ± 0.01020	0.079 ± 0.00073
TECO-Transformer*	173M	27.5 ± 1.77	22.4 ± 0.368	0.709 ± 0.0119	0.155 ± 0.00958
TECO-Transformer (our run)	173M	28.2 ± 0.66	21.6 ± 0.079	0.696 ± 0.02640	0.082 ± 0.00119
TECO-S5	180M	34.6 ± 0.26	20.1 ± 0.037	0.687 ± 0.00132	0.143 ± 0.00049
TECO-ConvSSM (ablation)	175M	44.3 ± 2.69	21.0 ± 0.106	0.691 ± 0.00004	0.010 ± 0.00267
TECO-ConvS5	175M	31.2 ± 0.23	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 in 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 is the S5 run which serves as an additional ablation of the convolutional structure.
B kernel	C kernel	Activation	DMLab
Params	FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓	Train Step Time (s) ↓
3 × 3	3 × 3	ResNet	100M	66.6 ± 4.81	23.2 ± 0.053	0.769 ± 0.00043	0.079 ± 0.00073	2.31
1 × 1	3 × 3	ResNet	85M	67.5 ± 0.73	22.8 ± 0.041	0.756 ± 0.00039	0.085 ± 0.00124	2.52
1 × 1	1 × 1	ResNet	100M	81.1 ± 6.06	23.0 ± 0.074	0.767 ± 0.00186	0.083 ± 0.00139	2.21
3 × 3	3 × 3	GLU	71M	96.1 ± 5.20	22.7 ± 0.157	0.762 ± 0.00363	0.088 ± 0.00287	2.75
1 × 1	5 × 5	GLU	83M	89.1 ± 0.61	21.5 ± 0.072	0.713 ± 0.00290	0.106 ± 0.00289	2.58
1 × 1	1 × 1	GLU	30M	187 ± 2.77	21.0 ± 0.064	0.689 ± 0.00007	0.112 ± 0.00183	1.73
-	-	GLU	140M	221 ± 13.1	19.3 ± 0.128	0.641 ± 0.00400	0.162 ± 0.04510	1.34
(a) 156 frames generated conditioned on 144 (action-conditioned).
(b) 264 frames generated conditioned on 36 (no action-conditioning).
Figure 4: DMLab Samples
(a) 156 frames generated conditioned on 144 (action-conditioned).
(b) 264 frames generated conditioned on 36 (action-conditioned).
Figure 5: Minecraft Samples
(a) 156 frames generated conditioned on 144 (action-conditioned).
(b) 264 frames generated conditioned on 36 (no action-conditioning).
Figure 6: Habitat Samples
Table 10: Full results on the Minecraft and Habitat long-range benchmark datasets [13]. Results from Yan et al. [13] are indicated with  $*$ . Note that Yan et al. [13] did not evaluate FitVid or CW-VAE on Habitat due to cost.
Method	Minecraft
Params	FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓
FitVid*	176M	956 ± 15.8	13.0 ± 0.0089	0.343 ± 0.00380	0.519 ± 0.00367
CW-VAE*	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.4 ± 0.0904	0.349 ± 0.00327	0.429 ± 0.00284
TECO-Transformer*	274M	116 ± 5.08	15.4 ± 0.0603	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 ↓	PSNR ↑	SSIM ↑	LPIPS ↓
Perceiver AR*	200M	164 ± 12.6	12.8 ± 0.0423	0.405 ± 0.00248	0.676 ± 0.00282
Latent FDM*	87M	433 ± 2.67	12.5 ± 0.0121	0.311 ± 0.00083	0.582 ± 0.00049
TECO-Transformer*	386M	76.3 ± 1.72	12.8 ± 0.0139	0.363 ± 0.00122	0.604 ± 0.00451
TECO-ConvS5	351M	95.1 ± 3.74	12.9 ± 0.212	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 FitVid, CW-VAE, Perceiver AR and Latent FDM used in the TECO work [13] are not 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.0×)	9.1 (1.0×)
Performer	1.25 (1.0×)	7.6(0.8×)
S5	1.34 (1.1×)	28 (3.1×)
ConvS5	2.31 (1.8×)	56 (6.2×)
TECO-Transformer	0.75 (0.6×)	16 (1.8×)
TECO-S5	0.81 (0.7×)	21 (2.3×)
TECO-ConvS5	1.17 (0.9×)	18 (2.0×)
Minecraft
Method	Train Step Time (s) ↓	Sampling Speed (frames/s)
TECO-Transformer	1.91 (1.0×)	8.1 (1.0×)
TECO-ConvS5	2.53 (1.3×)	14 (1.7×)
Habitat
Method	Train Step Time (s) ↓	Sampling Speed (frames/s)
TECO-Transformer	2.71 (1.0×)	6.8 (1.0×)
TECO-ConvS5	3.10 (1.1×)	11 (1.6×)
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 is used along with residual connections. For the Transformer, we use the Transformer implementation from Yan et al. [13] which consists of a 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 to a latent resolution 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 a 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 \times 3$  kernels as implemented in Yan et al. [13].
We follow Yan et 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 a 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 on 300 frames. We then repeated the experiment and trained models on 600 frames. For ConvS5 and ConvLSTM, we fixed the hidden dimensions (layer input/output features) and state sizes to be 256, and we swept over the following learning rates  $[1 \times 10^{-4}, 5 \times 10^{-4}, 1 \times 10^{-3}]$  and chose the best model. For Transformer, we swept over model size, considering hidden dimensions of [512, 2014] and learning rates  $[1 \times 10^{-4}, 5 \times 10^{-4}, 1 \times 10^{-3}]$  and chose the best model. We also observed better performance for the Transformer by convolving frames down to an  $8 \times 8$  latent resolution (rather than the  $16 \times 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 set, 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 \times 10^{-4}, 5 \times 10^{-4}, 1 \times 10^{-3}]$ .
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 × 64
Latent Resolution	16 × 16	16 × 16
Batch Size	8	8
Sequence Length	300	600
LR	1 × 10-3	1 × 10-3
LR Schedule	cosine	cosine
Warmup Steps	5k	5k
Max Training Steps	300K	300K
Weight Decay	1 × 10-5	1 × 10-5
Encoder	Depths	64, 128, 256	64, 128, 256
Blocks	1	1
Decoder	Depths	64, 128, 256	64, 128, 256
Blocks	1	1
ConvS5	Hidden Dim (U)	256	256
State Size (P)	256	256
B Kernel Size	3 × 3	3 × 3
C Kernel Size	3 × 3	3 × 3
Layers	8	8
Dropout	0	0
Activation	ResNet	ResNet