SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 scans and deep SSM parameterization/initialization schemes. Given Proposition 3, we implicitly parameterize a pointwise state kernel,  $\mathcal{A} \in \mathbb{R}^{P \times P \times 1 \times 1}$  and input kernel  $\mathcal{B} \in \mathbb{R}^{P \times U \times k_B \times k_B}$  in (5) using SSM parameters as used by S5 [20],  $\mathbf{A}_{\mathrm{S5}} \in \mathbb{R}^{P \times P}$  and  $\mathbf{B}_{\mathrm{S5}} \in \mathbb{R}^{P \times (Uk_B^2)}$ . We discretize these S5 SSM parameters as discussed in Section 2.2 to give
$$
\overline {{\mathbf {A}}} _ {\mathrm {S} 5} = \text {D I S C R E T I Z E} _ {\mathrm {A}} (\mathbf {A} _ {\mathrm {S} 5}, \boldsymbol {\Delta}), \quad \overline {{\mathbf {B}}} _ {\mathrm {S} 5} = \text {D I S C R E T I Z E} _ {\mathrm {B}} (\mathbf {A} _ {\mathrm {S} 5}, \mathbf {B} _ {\mathrm {S} 5}, \boldsymbol {\Delta}), \tag {11}
$$
and then reshape to give the ConvS5 state update kernels:
$$
\overline {{\mathbf {A}}} _ {\mathrm {S} 5} \in \mathbb {R} ^ {P \times P} \xrightarrow {\text {r e s h a p e}} \overline {{\mathcal {A}}} _ {\text {C o n v S} 5} \in \mathbb {R} ^ {P \times P \times 1 \times 1} \tag {12}
$$
$$
\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 {{\mathcal {B}}} _ {\text {C o n v S} 5} \in \mathbb {R} ^ {P \times U \times k _ {B} \times k _ {B}}. \tag {13}
$$
We then run the discretized ConvSSM system of (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  $\mathcal{O}(LPHW)$ . See Appendix B for a more detailed discussion of parameterization, initialization and discretization.
We define a ConvS5 layer as the combination of ConvS5 with a 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 a 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	O(L)	O(1)	O(1)
Training	O(L2)	O(L)	O(L)
Parallelizable	Yes	No	Yes
3.5 ConvS5 Properties
Refer to Table 1 for a 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 a stateful model like ConvRNNs, allowing fast autoregressive generation, extrapolation to different sequence lengths, and an 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 leveraging the continuous-time parameterization to process irregularly sampled sequences in parallel. S5 achieves this by providing suitable spacing to the discretization operation [20], a procedure that ConvS5 can also use.
We have proposed a general ConvSSM structure that can be easily adapted to future innovations in the deep SSM literature. ConvS5's parallel scan could be used to endow ConvS5 with time-varying dynamics such as the input-dependent 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 on language modeling. Similar ideas could be useful to ConvSSMs for reasoning over long spatiotemporal sequences.
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 PixelRNN [86] approach, and then combining the results with an 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 ours and the S42D blocks could potentially be used to replace some of the ResNet blocks used as activations in our model. Other model architectures explored in the literature for spatiotemporal modeling include 3D 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-time neural representations [93, 95], training on different length sequences at different resolutions [96] and stripped sampling [14, 97]. Of particular interest, Yan et al. [13] introduced the Temporally Consistent (TECO) Video Transformer, which achieved state-of-the-art performance on challenging 3D environment benchmarks designed to contain long-range dependencies [13]. This approach combines vector-quantized (VQ) latent dynamics with a MaskGit [98] dynamics prior, and several training tricks to model long videos.
In Section 5.1, 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
Table 2: Quantitative evaluation on the Moving-MNIST dataset [54]. Top: To 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 to Table 1.
Trained on 300 frames
Method	100 → 800	100 → 1200
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	91	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) ↑
Transformer 70	0.77(1.0×)	50	0.21 (1.0x)
ConvLSTM 65	3.0(3.9×)	150	117 (557×)
ConvS5 97	0.93(1.2×)	50	90 (429×)
environment benchmarks proposed in Yan et al. [13]. Finally, in Section 5.3, we discuss ablations that highlight the importance of ConvS5's parameterization.
5.1 Long Horizon Moving-MNIST Generation
There are 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 to 1200 frames. This allows for a direct comparison of ConvS5, ConvRNNs and Transformers as well as an efficient attention alternative (Performer [33]) and CWVAE [16], a 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 D 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 1 and compare the empirical computational costs of the Transformer, ConvLSTM and ConvS5 on the 600 frame Moving-MNIST task. Although this specific ConvS5 configuration requires a few more FLOPs due to the convolution computations, ConvS5 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.
5.2 Long-range 3D Environment Benchmarks
Yan et al. [13] introduced a challenging video prediction benchmark specifically designed to contain long-range dependencies. This is one of the first comprehensive benchmarks for long-range spatiotemporal modeling and consists of 300 frame videos of agents randomly traversing 3D environ-
Table 3: Quantitative evaluation on the DMLab long-range benchmark [13]. Results from Yan et al. [13] are indicated with  $\ast$ . Methods trained using the TECO [13] training framework are at the bottom of the table. TECO methods are slower to sample due to the MaskGit [98] procedure. The expanded Table 8 in Appendix C includes error bars and ablations.
Method	FVD ↓	DMLab	LPIPS ↓	Sample Throughput (frames/s) ↑
PSNR ↑	SSIM ↑	LPIPS ↓
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.0×)
Performer [33]	80	17.3	0.513	0.205	7 (0.8×)
S5 [20]	221	19.3	0.641	0.162	28 (3.1×)
ConvS5	66	23.2	0.769	0.079	56 (6.2×)
TECO-Transformer* [13]	28	22.4	0.709	0.155	16 (1.8×)
TECO-Transformer (our run)	28	21.6	0.696	0.082	16 (1.8×)
TECO-S5	35	20.1	0.687	0.143	21 (2.3×)
TECO-ConvS5	31	23.8	0.803	0.085	18 (2.0×)
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 \times 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. The S5 baseline serves as an 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 as a drop-in replacement for the Transformer in TECO. Therefore, we refer to the original version of TECO as TECO-Transformer, the ConvS5 version as TECO-ConvS5 and the S5 version as TECO-S5. See Appendix D for detailed information on training procedures, architectures, and hyperparameters.
DMLab The results for DMLab are presented in Table 3. Of the methods trained without the TECO framework in the top section of Table 3, 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 on all metrics, pointing to the value of the convolutional structure of ConvS5.
For the models trained with the TECO framework, we see that TECO-ConvS5 achieves essentially the same FVD and LPIPS as TECO-Transformer, while significantly improving PSNR and SSIM. Note the sample speed comparisons are less dramatic in this setting since the MaskGit [98] sampling procedure is relatively slow. Still, the sample throughput of TECO-ConvS5 and TECO-S5 remains constant, while TECO-Transformer's throughput decreases with sequence length.
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 metrics, outperfoTarming all other baselines. On Habitat, TECO-ConvS5 is the only method to achieve a comparable FVD to TECO-Transformer, 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 to 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 4: Quantitative evaluation on the Minecraft and Habitat long-range benchmarks [13]. Results from Yan et al. [13] are indicated with  $*$ . See expanded Table 10 with error bars in Appendix C.
Method	Minecraft	Habitat
FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FVD ↓	PSNR ↑	SSIM ↑	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
TECO-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 DMLab long-range benchmark dataset [13]. More convolutional structure improves overall performance. See expanded Table 9 in Appendix.
conv.	B kernel	C kernel	DMLab
nonlinearity	FVD ↓	PSNR ↑	SSIM ↑	LPIPS ↓
x	-	-	GLU	221	19.3	0.641	0.162
o	1 × 1	1 × 1	GLU	187	21.0	0.689	0.112
o	1 × 1	5 × 5	GLU	89	21.5	0.713	0.106
o	3 × 3	3 × 3	GLU	96	22.7	0.762	0.088
o	1 × 1	1 × 1	ResNet	81	23.0	0.767	0.083
o	1 × 1	3 × 3	ResNet	68	22.8	0.756	0.085
o	3 × 3	3 × 3	ResNet	67	23.2	0.769	0.079
We also perform ablations to evaluate the importance of ConvS5's deep SSM-inspired parameterization/initialization. We evaluate the performance of a ConvSSM with randomly initialized state kernel on both Moving-Mnist and DMLab. See Table 6 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.
This work introduces ConvS5, a new spatiotemporal modeling layer that combines the fast, stateful autoregressive generation of ConvRNNs with the ability to 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-time 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 frameworks 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 at 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 be used for such applications in the spatiotemporal domain [104, 93, 95]. Finally, ConvS5 is modular and flexible. We have demonstrated that ConvS5 works well as a drop-in replacement in the TECO [13] training framework specifically developed for Transformers. Due to its favorable properties, we expect ConvS5 to also serve as a building block of new approaches for modeling much longer spatiotemporal sequences and multimodal applications.
[101] Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, et al. Habitat: A platform for embodied AI research. In Proceedings of the IEEE/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 on 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.03580, 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):851-876, 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 arXiv:1607.06450, 2016.
[108] Thomas Unterthiner, Sjoerd Van Steenkiste, Karol Kurach, Raphael Marinier, Marcin Michalski, and Sylvain Gelly. Towards accurate generative models of video: A new metric & challenges. arXiv preprint arXiv:1812.01717, 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 unreasonable effectiveness of deep features as a perceptual metric. In CVPR, 2018.
[111] Yann LeCun. The MNIST database of handwritten digits. http://yann.lecun.com/exdb/mnist/, 1998.
[112] Santhosh K 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 AI. arXiv preprint arXiv:2109.08238, 2021.
[113] Angel Chang, Angela Dai, Thomas Funkhouser, Maciej Halber, Matthias Niessner, Manolis Savva, Shuran Song, Andy Zeng, and Yinda Zhang. Matterport3D: Learning from RGB-D data in indoor environments. arXiv preprint arXiv:1709.06158, 2017.
[114] Fei Xia, Amir R Zamir, Zhiyang He, Alexander Sax, Jitendra Malik, and Silvio Savarese. Gibson env: Real-world perception for embodied agents. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 9068-9079, 2018.
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
A Propositions
A.1 Parallel Scan for Convolutional Recurrences
$$
c _ {i} * c _ {j} := \left(c _ {j, a} \circ c _ {i, a}, c _ {j, a} * c _ {i, b} + c _ {j, b}\right), \tag {14}
$$
where  $\circ$  denotes convolution of two kernels,  $*$  denotes convolution between a kernel and input, and  $+$  is elementwise addition.
Proof. Using that  $\circ$  is associative and the companion operator of  $*$ , i.e.  $(d \circ e) * f = d * (e * f)$  (see Blelloch [63], Section 1.4), we have:
$$
\begin{array}{l} \left(c _ {i} * c _ {j}\right) * c _ {k} = \left(c _ {j, a} \circ c _ {i, a}, c _ {j, a} * c _ {i, b} + c _ {j, b}\right) * \left(c _ {k, a}, c _ {k, b}\right) (15) \\ = \left(c _ {k, a} \circ \left(c _ {j, a} \circ c _ {i, a}\right), c _ {k, a} * \left(c _ {j, a} * c _ {i, b} + c _ {j, b}\right) + c _ {k, b}\right) (16) \\ = \left(\left(c _ {k, a} \circ c _ {j, a}\right) \circ c _ {i, a}, c _ {k, a} * \left(c _ {j, a} * c _ {i, b}\right) + c _ {k, a} * c _ {j, b} + c _ {k, b}\right) (17) \\ = \left(\left(c _ {k, a} \circ c _ {j, a}\right) \circ c _ {i, a}, \left(c _ {k, a} \circ c _ {j, a}\right) * c _ {i, b} + c _ {k, a} * c _ {j, b} + c _ {k, b}\right) (18) \\ = c _ {i} * \left(c _ {k, a} \circ c _ {j, a}, c _ {k, a} * c _ {j, b} + c _ {k, b}\right) (19) \\ = c _ {i} * \left(c _ {j} * c _ {k}\right) (20) \\ \end{array}
$$
A.2 Computational Cost of Parallel Scan for Convolutional Recurrences
Proof. Following Blelloch [63], given a single processor, the cost of computing the recurrence sequentially using the binary operator  $\circ$  defined in Proposition 1 is  $\mathcal{O}\big(L(T_{\circ} + T_{*} + T_{+})\big)$  where  $T_{\circ}$  refers to the cost of convolving two kernels,  $T_{*}$  is the cost of convolution between a kernel and input and  $T_{+}$  is the cost of elementwise addition. The cost of elementwise addition is  $T_{+} = \mathcal{O}(PHW)$ . For state kernels with resolution  $k_{A}$ ,  $T_{\circ} = \mathcal{O}(P^{3}k_{A}^{4})$  and  $T_{*} = \mathcal{O}(P^{2}k_{A}^{2}HW)$ . For pointwise convolutions this becomes  $T_{\circ} = \mathcal{O}(P^{3})$  and  $T_{*} = \mathcal{O}(P^{2}HW)$ . Thus, the cost of computing the recurrence sequentially using  $\circ$  is  $\mathcal{O}\big(L(P^3 + P^2 HW)\big)$ . Since there are work-efficient algorithms for parallel scans [105], the overall cost of the parallel scan is also  $\mathcal{O}\big(L(P^3 + P^2 HW)\big)$ .
Note that ConvS5's diagonalized parameterization discussed in Section 3.4 and Appendix B leads to  $T_{\circ} = \mathcal{O}(P)$  and  $T_{*} = \mathcal{O}(PHW)$ . Therefore the cost of applying the parallel scan with ConvS5 is  $\mathcal{O}(LPHW)$ .
A.3 Connection Between ConvSSMs and SSMs
$$
\boldsymbol {\mathcal {X}} ^ {\prime} (t) _ {i, j} = \mathbf {A} _ {\text {S S M}} \boldsymbol {\mathcal {X}} (t) _ {i, j} + \mathbf {B} _ {\text {S S M}} \boldsymbol {\mathcal {U}} _ {i m 2 c o l} (t) _ {i, j} \tag {21}
$$
where the state matrix,  $\mathbf{A}_{\mathrm{SSM}}\in \mathbb{R}^{P\times P}$ , and input matrix,  $\mathbf{B}_{\mathrm{SSM}}\in \mathbb{R}^{P\times (Uk_B^2)}$ , can be formed by reshaping the state kernel,  $\mathcal{A}$ , and input kernel,  $\mathcal{B}$ , respectively.
Proof. Let  $\mathbf{U}_{\mathrm{im2col}} \in \mathbb{R}^{Uk_b^2 \times HW}$  denote the result of performing the im2col operation on the input  $\mathcal{U}(t)$  for convolution with the kernel  $\mathcal{B}$ . Reshape this matrix into the tensor  $\mathcal{U}_{\mathrm{im2col}}(t) \in \mathbb{R}^{H \times W \times Uk_b^2}$ . Reshape  $\mathcal{U}_{\mathrm{im2col}}(t)$  once more into the tensor  $\mathcal{V}(t) \in \mathbb{R}^{H \times W \times U \times k_B \times k_B}$ .
Now, we can write the evolution for the individual channels of each pixel,  $\mathcal{X}'(t)_{i,j,k}$ , in (5) as