Title: Convolutional State Space Models for Long-Range Spatiotemporal Modeling
Venue: NeurIPS 2023
Accepted: yes
Keywords: spatiotemporal modeling

ConvLSTM

RNN

state spaces

SSM

S4

S5

long-range dependencies

video prediction
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

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) 1 that combine the tensor modeling ideas of ConvLSTM with the long sequence modeling approaches of state space methods such as S4 and S5. First, we demonstrate how parallel scans can be applied to convolutional recurrences to achieve subquadratic parallelization and fast autoregressive generation. We then establish an equivalence between the dynamics of ConvSSMs and SSMs, which motivates parameterization and initialization strategies for modeling long-range dependencies. The result is ConvS5, an efficient ConvSSM variant for long-range spatiotemporal modeling. ConvS5 significantly outperforms Transformers and ConvLSTM on a long horizon Moving-MNIST experiment while training 3 × faster than ConvLSTM and generating samples 400 × 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.

1 Introduction

Developing methods that efficiently and effectively model long-range spatiotemporal dependencies is a challenging problem in machine learning. Whether predicting future video frames [1, 2], modeling traffic patterns [3, 4], or forecasting weather [5, 6], deep spatiotemporal modeling requires simultaneously capturing local spatial structure and long-range temporal dependencies. Although there has been progress in deep generative modeling of complex spatiotemporal data [7-12], most prior work has only considered short sequences of 20-50 timesteps due to the costs of processing long spatiotemporal sequences. Recent work has begun considering sequences of hundreds to thousands of timesteps [13-16]. As hardware and data collection of long spatiotemporal sequences continue to improve, new modeling approaches are required that scale efficiently with sequence length and effectively capture long-range dependencies.

k

, where

74

75

×

SSM

P

P

222

1:

L

205

i,j

217

i,j with a state update equation

represents some nonlinear function. For

A

96

74

pointwise state kernel

223

x

1

P

B

209

A

R

represents some nonlinear function. For spatiotemporal ConvS5 approach.

(

Uk

)

with a state update equation with a shared state matrix,

F

(

x

P

k

=

,

k

u

)

Block

, where

98

k

76

′

208

×

SSM

into a matrix

A

218

i,j pointwise state kernel

1

P

1

into a matrix operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

P

A

k

x

k

222

204

k

SSM

i,j

×

H

207

P

205

×

×

204

B

t

. Rearrange the input kernel i,j

SSM

U

R

210

1

,

u

k

)

, where

x

=

F

(

F

k

B

+

pointwise state kernel

R

2

221

95

k

B

(

R

)

)

with a shared state matrix,

A

, and input matrix,

, and input

97

, input kernel

)

. The differential equation of (7) is equivalent, after reshaping

SSM

216

X

B

()

2.2

75

206

(

A

F

U

example, one can represent a vanilla RNN (ignoring the bias term) as

k

P

P

B

S5 [36] that were developed for 1-D sequence modeling and whose insights we leverage for our

219

∈

R

(

t

)

=

A

represents some nonlinear function. For

P

(

F

()

205

2

×

SSM

W

A

′

B

R

X

223

t

)

(

x

P

222

P

()

, and input matrix, example, one can represent a vanilla RNN (ignoring the bias term) as

gnoring the bias term) as

)

1:

B

. Rearrange the input kernel

W

into a matrix

)

2

H

′

×

L

209

into a matrix with a state update equation

H

U

P

with a shared state matrix,

k

t

)

B

= tanh(

Ax

P

SSM

(

(

S5 [36] that were developed for 1-D sequence modeling and whose insights we leverage for our

R

212

97

∈

-

224

x

t

)

∈

t

)

,

′

)

x

(

t

208

212

and the hyperbolic tangent function

)

∈

R

P

(

with state matrix into

i,j

A

(

P

t

SSM

76

Bu

′

U

Uk

×

×

B

SSM

-

, and input matrix,

B

×

×

.

)

x

(

x

,

u

. See Figure 1.

74

H

P

Consider a ConvSSM state update as in (7) with pointwise state kernel cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

i,j

A

SSM

. See Figure 1.

SSM

P

Uk

)

′

k

P

×

A

U

B

W

SSM

∈

W

Uk

B

P

R

(

∈

′

R

P

P

220

′

k

Proposition 3.

U

150

221

R

×

×

×

-

1

75

B

LSTM [13] and GRU [14] define more sophisticated spatiotemporal ConvS5 approach.

216

98

207

×

h(

k

1

k

Ax

+

Bu

)

99

with state matrix

77

∈

R

P

2.2.1

HWP

A

SSM

A

223

X

∈

B

SSM

(

t

Rearrange the

SSM

208

P

P

1

. The differential equation of (7) is equivalent, after reshaping is applicable to any HxWxC tensor (which is the only case we consider). It is al

, where

t

(

t

)

∈

U

and

(

k

=

F

P

)

×

P

H

206

W

209

2

98

×

(

t

)

P

P

×

3.1

×

R

×

∈

, input matrix

SSM

×

applied elementwise.

×

2

1

P

U

(1)

′

×

k

1

k

2

×

×

, input kernel

R

. Rearrange the input kernel

B

209

P

(

Uk

)

i,j respectively, to the following differential equation

×

U

P

P

t

×

B

into a matrix

k

(

)

col

. See Figure 1.

SSM

represents some nonlinear function. For

B

i,j

F

()

+

)

, and input

P

B

B

B

example, one can represent a vanilla RNN (ignoring the bias term) as

R

206

′

224

B

into

B

R

P

2

∈

70

)

Linear State Space Models

(

U

(1)

A

, input matrix

×

P

96

R

77

76

R

77

k

i,j

209

B

x

-

(

t

)

X

215

P

∈

P

SSM

= tanh(

Ax

k

-

X

217

P

209

x

SSM

t

(

U

respectively, to the following differential equation

A

′

x

(

t

)

R

,

)

t

(

R

×

Uk

SSM

x

Background

∈

∈

(

t

i,j

(

t

)

im

2

into a matrix

(

k

= tanh(

im

R

Ax

k

′

and

t

)

(

2

HWP

H

R

×

×

X

t

)

+

Bu

)

)

Consider a ConvSSM state update as in (7) with pointwise state kernel

Proposition 3.

SSM

k

-

. The differential equation of (7) is equivalent, after reshaping

P

W

.

)

B

Uk

(

)

and the hyperbolic tangent function

∈

P

B

k

U

∈

k

B

U

′

′

B

×

X

(

t

)

i,j

+

P

i,j

Convolutional State Space Models

210

×

X

B

′

220

pointwise state kernel

×

. Rearrange the input kernel

∈

A

1

k

operations, and has a paper we can cite. I am open to suggestions, but was hav

2

222

k

(

t

)

and

X

∈

SSM

col

(

t

i,j

(

∈

t

)

B

×

=

A

SSM

76

Bu

P

(

k

)

∈

×

This section provides a brief background on the ideas of deep SSM methods such as S4 [33] and

∈

R

97

224

∈

+

B

210

208

P

P

applied elementwise.

213

formulations of

78

R

′

X

t

×

Other RNNs such as LSTM [13] and GRU [14] define more sophisticated

×

k

×

′

P

5

X

Rearrange the

)

, and input

-

×

and the hyperbolic tangent function

U

X

Uk

A

Block i,j

x

(

t

) +

)

P

H

W

HWP

B

A

∈

R

151

)

2

with state matrix

= tanh(

×

B

A

∈

B

U

k

k

X

P

P

1

1

P

and the hyperbolic tangent function

×

H

SSM

x

)

(

) =

(

)

)

,

x

W

i,j

x

′

P

(

Uk

B

)

)

(

t

′

(

(

t

207

R

t

respectively, to the following differential equation

1

into a matrix tanh()

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

×

×

P

P

A

SSM

Block

210

U

′

×

B

, input matrix

′

(1)

∈

A

t

x

R

)

Rearrange the

×

×

R

R

, input kernel

k

k

(

+

1

Uk

211

Consider a continuous-time tensor-valued input signal

t

)

(1)

Ax

R

i,j

t

)

k

1

+

Bu

)

H

′

col

k

Ax

(

U

×

207

-

+

u

(

t

)

i,j

B

SSM

R

×

B

211

A

B

)

k

B

2

B

221

Uk

×

223

k

U

B

(

′

i,j

HWP

W

P

(

(

t

′

∈

, width

∈

P

SSM

HWP

P

(

t

′

)

S5 [36] that were developed for 1-D sequence modeling and whose insights we leverage for our

Other RNNs such as LSTM [13] and GRU [14] define more sophisticated

213

)

t

(

F

225

SSM

P

210

U

×

t

)

×

Block

215

Uk

A

(

×

into

Bu

SSM

()

×

A

×

and the hyperbolic tangent function

100

formulations of

()

78

.

U

76

t

′

=

71

2

(

t

)

=

A

x

B

(

×

.

t

)

H

into a matrix

∈

222

SSM

77

×

R

99

X

2

210

R

P

with state matrix

P

99

F

applied elementwise.

Given a continuous input signal

Block

210

=

Block

A

where

×

R

R

with state matrix

, input matrix

217

∈

B

B

′

(

×

A

t

)

B

A

X

Ns with

B

U

SSM

i,j

(

cleaner way of describing what we need here. I feel like saying im2col and ci

×

=

152

t

)

R

. The differential equation of (7) is equivalent, after reshaping

W

U

X

pointwise state kernel

[JS: im2col

B

2

im is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

×

A

×

B

t

′

×

t

) +

x

Block

A

x

(

) =

(

t

B

∈

SSM

,

208

×

where

HWP

Block

U

×

×

B

+

5

Block

A

HWP

and

×

A

×

into a matrix

-

t

)

R

R

×

P

∈

∈

Other RNNs such as LSTM [13] and GRU [14] define more sophistica i,j

applied elementwise.

SSM

im

2

= tanh(

respectively, to the following differential equation

U

Linear State Space Models

X

are block-diagonal with

SSM

HWP

+

i,j

x

B

(

P

′

P

t

)

HWP

X

2.2.1

A

Linear State Space Models

P

×

P

This section provides the background necessary to introduce ConvSSMs and ConvS5 in Section 3.

.

im

u

′

i,j

(

into

x

(

t

)

,

x

(

t

)

R

×

208

R

×

X

HWP

R

′

98

×

2.2.1

Block

(

t

Block

225

A

B

R

×

U

)

76

, input matrix

B

and the hyperbolic tangent function

2

B

LSTM [13] and GRU [14] define more sophisticated

1, 12] (ConvRNNs) model spatiotemporal sequences

X

X

79

78

2

. Rearrange the input kernel

X

X

∈

U

′

and

(

t

)

(

t

)

(12)

B

×

(12)

B

with height

H

′

col

B

(12)

×

(

t

)

R

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

R

with state matrix

209

x

(

, a latent state

t

×

) +

R

R

×

×

HWP

and

∈

X

∈

2

P

W

′

H

212

R

ConvSSM

∈

+

B

×

(

2

B

and the hyperbolic tangent function and number of input features

X

)

t

(

col

2

SSM

HWP

224

212

211

′

U

Block

SSM

R

HWUk

U

U

P

2

) =

b

A

Block

A

R

211

∈

79

i,j

HWP

)

Uk

B

P

SSM

B

211

2

applied elementwise.

Convolutional recurrent neural networks

P

∈

∈

i,j

Uk

(

77

209

SSM

209

formulations of

∈

∈

153

im

t

)

i,j

∈

SSM

SSM

A

∈

225

∈

223

2

2

col

(

t

)

i,j

H

2

R

R

,

(1)

u

2

b

y

(

t

)

P

U

.

. We will define a continuous-time, linear convolutional state space are block-diagonal with

. The differential equation of (7) is equivalent, after reshaping

HWP

, input matrix formulations of

,

)

t

(

B

SSM

u

U

R

W

Proposition 3.

′

U

B

and an output signal into a matrix

W

′

,

(

M

U

[11, 12] (ConvRNNs) model spatiotemporal sequences

A

,

Consider a ConvSSM state update as in (7) with

R

and the hyperbolic tangent function

×

R

×

W

P

()

A

x

224

B

P

P

(

t

)

×

P

U

78

2

(

Uk

)

HWUk

F

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

210

∈

model (

x

(

t

)

,

P

into

Other RNNs such as LSTM [13] and GRU [14] define more sophisticated

∈

formulations of

∈

Other RNNs such as LSTM [13] and GRU [14] define more sophisticated

226

a continuous-time, linear SSM is defined by the differential equation:

B

226

101

F

()

.

HWP

Uk

′

B

A

B

Convolutional recurrent neural networks

X

×

SSM

on the diagonal blocks respectively, and

B

80

∈

213

HWUk

b

P

(

Uk

)

P

1

1

Uk

211

X

HWP

HWP

B

SSM

X

B

B

Given a continuous input signal

100

H

1

-D tensors and replacing the matrix-vector

77

on the diagonal blocks respectively, and

[11, 12] (ConvRNNs) model spatiotemporal sequences

1

()

columns formed by performing an im2col [74, 75] operation on the input

Given a continuous input signal

214

.

by replacing the 1-D states and inputs of RNNs with by replacing the 1-D states and inputs of RNNs with

)

100

=

Block

A

i,j

P

P

H

where

101

∈

) with state

212

×

R

W

U

applied elementwise.

∈

, derivative

×

SSM

R

×

and

×

is the result of concatenating the

×

×

Block

R

HWUk

79

B

H

and

A

P

P

u

u

R

u

(

t

)

SSM

R

.

∈

′

(

(

Block

(

t

×

B

Block

∈

x

x

U

′

∈

t

) =

A

HWUk

t

)

k

B

k

U

∈

)

R

HWP

2

b

and

H

W

P

is the result of concatenating the

(

b

(

t

)

∈

×

)

×

B

and an output signal

t

P

(

X

×

A

Other RNNs such as LSTM [13] and GRU [14] define more sophisticated are block-diagonal with

HWUk

t

)

Proposition 3.

H

∈

F

()

.

, a latent state

211

N

R

×

B

∈

x

(

)

×

im

2

t

, as the differential equation:

-D tensors and replacing the matrix-vector

W

227

×

×

and output

(12)

∈

′

(

t

x

(

t

) +

P

formulations of

R

R

2

212

SSM

78

olutional neural networks (CNNs). Given a length

216

79

∈

99

×

×

′

i,j

H

)

t

(

72

211

210

227

(

t

′

214

N

80

81

Ns with

1, 12] (ConvRNNs) model spatiotemporal sequences

80

B

u

(

t

)

,

respectively, to the following differential equation

SSM

and

1

)

)

t

Convolutional recurrent neural networks

×

R

R

columns formed by performing an im2col [74, 75] operation on the input im

R

col

2

′

(

t

)

(

210

-D tensors and replacing the matrix-vector

×

U

+

b

i,j

214

(

t

)

HWP

80

R

Consider a ConvSSM state update as in (7) with pointwise state kernel

, input kernel

1

∈

1

81

P

2

col

HWP

by replacing the 1-D states and inputs of RNNs with

L

P

(

′

t

)

[11, 12] (ConvRNNs) model spatiotemporal sequen

U

R

∈

×

and

×

R

B

P

1

, a latent state

x

(

is the result of concatenating the

P

Block

u

(

t

A

)

and an output signal

B

t

)

M

X

∈

Y

HWP

HWUk

2

y

(

t

)

R

∈

,

, input kernel

×

Block

∈

)

SSM

i,j

∈

and

×

on the diagonal blocks respectively, and

X

×

t

′

U

225

where a continuous-time, linear SSM is defined by the differential equation:

Linear State Space Models

Convolutional Recurrent Networks i,j

154

i,j

=

A

t

)

u

W

∈

×

×

(

t

)

multiplications with convolutions as in convolutional neural networks (CNNs). Given a length

P

78

F

multiplications with convolutions as in convolutional neural networks (CNNs). Given a length a continuous-time, linear SSM is defined by the differential equation:

Uk

226

R

2

79

N

N

B

+

Convolutional recurrent neural networks

×

)

t

B

X

Uk

SSM

(

2.2.1

×

W

2.1

columns formed by performing an im2col [74, 75] operation on the input

R

∈

) =

×

t

U

.

(

t

)

N

B

(

t

) =

A

R

′

k

k

B

x

on the diagonal blocks respectively, and

∈

,

(

t

)

M

∈

u

×

R

×

, and input

X

(

X

t

)

-D tensors and replacing the matrix-ve

x

(

t

) +

B

u

(

t

)

×

×

.

Rearrange the

y

P

P

U

∈

A

∈

U

(12)

∈

are block-diagonal with

b

HWUk

R

×

. Rearrange th

′

W

SSM

Block

2

) +

×

U

(

t

Bu

) +

U

212

x

′

(

Uk

H

′

P

W

U

×

Convolutional recurrent neural networks

101

216

2

X

t

) =

Ax

(

t

79

B

[11, 12] (ConvRNNs) model spatiotemporal sequences

(

81

-D tensors and replacing the matrix-vector

R

W

by replacing the 1-D states and inputs of RNNs with with height

, width

′

and

211

∈

A

82

81

212

k

k

×

×

×

U

′

×

Convolutional recurrent neural networks

H

′

U

82

83

olutional neural networks (CNNs). Given a length

, with height multiplications with convolutions as in convolutional neural networks (CNNs). Given a length

∈

X

217

f HxW] [JS: Interesting, can you point me to some examples?

82

×

212

215

updates its state,

83

H

213

L

W

U

84

f HxW] [JS: Interesting, can you point me to some examples?

1:

L

pes of Torch and JAX/Flax's conv2d operations and matches sequence of frames,

83

84

′

, width

′

and equation would be

features

84

R

214

point, but it is more typical to write WxH instead of HxW] [JS: Interesting, can you point me to some examples?

The

85

U

∈

pes of Torch and JAX/Flax's conv2d operations and matches

×

1

, with a state update equation convention matches the expected shapes of Torch and JAX/Flax's conv2d operations and matches

P

P

H

matrix, so I currently have a preference to keep it as is]

85

W

216

85

215

×

86

, with a state update equation the height and width of the rows

matrix, so I currently have a preference to keep it as is]

86

ConvRNN updates its state,

2

86

2

where

87

∗

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

and the tanh

X

X

216

×

k

=

217

217

X

k

218

219

is a spatial convolution operator with state kernel

220

features, input features, kernel height, kernel width] convention), input kernel

()

221

function is applied elementwise.

222

equations respectively. ConvRNNs have been used for biomedical, robotics, traffic modeling, weather

∈

and

A

Block

(

t

)

,

SSM

A

multiplications with convolutions as in convolutional neural networks (CNNs). Given a lengt

H

into a matrix

(7)

, width

A

is the result of concatenating the

[11, 12] (ConvRNNs) model spatiotemporal sequences

2

are block-diagonal with

Du

′

(

)

,

y

(

t

) =

Cx

(

t

R

P

P

1

1

212

U

)

∈

×

B

k

t

U

(

t

)

pointwise state kernel

(

t

t

)

k

H

′

×

∈

B

,

(3)

P

U

×

×

Du

(

t

)

U

R

, and input

)

B

,

b

(

t

)

′

SSM

and is the result of concatenating the

(

t

(3)

(8)

features, a ConvR

-

X

L

into a matrix

) +

-D tensors and replacing the matrix-vector

(

t

)

and

′

W

k

. The differential equation of (7) is equivalent, aft

[JS: i

×

U

(

∈

5

)

(

X

A

∗

X

)

.

R

A

B

∗

and

[11, 12] (ConvRNNs) model spatiotemporal sequences

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in

X

P

U

L

U

=

G

(

k

, with a state update equation features

U

∈

1

,

k

)

G

(

B

P

U

k

B

k

X

X

U

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up

L

, where

U

[SD: minor

B

∈

respectively, to the following differential eq a nonlinear function. For a concrete example analogous to (1), the vanilla ConvRNN state up

A

U

(

t

)

.

, an output matrix

×

u

×

k

×

D

k

×

D

U

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty pre

Block

R

×

R

x

(

t

) +

×

B

Block

u

(

1:

L

R

×

, the

R

is the

(12)

. Given a sampled sequence,

×

with the specially structured HiPPO matrices [47] and

SSM

are block-diagonal with and

A

equation would be https://github.com/anonymous

t

)

,

Consider a ConvSSM state update as in (7) with pointwise state kernel

×

2

∈

H

W

, is

.

H

L

′

W

′

U

H

equation would be would endow these

k

,

, and input

Block with the specially structured HiPPO matrices [47] and

1:

is the result of concatenating the

G

(2)

L

(

, a

R

×

×

×

,

W

′

′

b

HWP

HWUk

1

[WB: Define im2col before eq 12. im2cal is specific to images. Is there

U

∈

R

SSM

X

U

U

k

×

.

×

(

t

)

are bl

B

. Rearrange the input kernel

SSM

and

-

∈

U

k

1

k

+

u

(

t

would endow these

k

,

R

k

)

R

Rearran

HWUk

∈

×

)

B

2

(4)

= tanh(

= tanh(

b

,

[JS: im2col

k

X

Note the ConvSSM formulation is easier to work with than

)

,

where the discrete-time parameters are a function of the continuous-time parameters and a timescale

+

is applicable to any HxWxC tensor (which is the only case we consider).

-

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

into a

(

X

(9)

. The differential equation of (7) is equivalent, after reshaping

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

k

= DISCRETIZE

respectively, to the following differential equation

X

∈

X

B

∗

B

B

∗

U

)

(

A

,

B

,

∆)

U

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same is t

where

,

columns formed by performing an im2col [74, 75] operation on the in dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

with the specially structured HiPPO matrices [47] and

5

A

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

∗

X

operations, and has a paper we can cite. I am open to suggestions, but

(10)

B

would endow these denote that these kernels are in discrete-time.

Block

Note the ConvSSM formulation is easier to work with than

u

cleaner way of describing what we need here. I feel like saying im2col cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

ConvLSTM [11] and ConvGRU [12] are used by making similar changes to the 1-D LSTM and GRU

223

224

is a discretization method such as Euler, bilinear or zero-order hold [69].

t

)

and

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to des forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

2

b

(

t

)

,

Note that these definitions are analogous to those of ConvRNNs discussed in Section 2.1, since we

Consider a ConvSSM state update as in (7) with pointwise state kernel forecasting and numerous other applications [49-60] and many variants have been proposed [61-68].

235

236

is applicable to any HxWxC tensor (which is the only case we consider). It is also an officia

HWUk have replaced the SSM's vector-valued inputs, states, and outputs with tensors and replaced the

×

operations, and has a paper we can cite. I am open to suggestions, but was having a hard ti

×

HWUk into a matrix

P

[SD: I feel like the recurrence part of SSM is not very clearly explained and hence hard to follow/understand.

B

213

U

R

H

W

′

′

(

t

A

SSM

B

×

A

Block

∈

R

HWP

×

HWP

5

A

Block

∈

R

HWP

×

HWP

217

B

∈

R

P

×

U

×

k

B

×

k

B

X

(

t

)

∈

R

H

×

W

×

P

A

Block

∈

R

HWP

×

HWP

5

B

∈

B

B

215

X

(

t

)

∈

R

H

×

W

×

P

216

A

Block

∈

R

HWP

×

HWP

217

5

Proof.

See Appendix C.3.

228

Thus, parameterizing and initializing

A

SSM

with the specially structured HiPPO matrices [47] and

229

disretizing with a learnable timescale parameter to obtain

A

SSM

and

B

SSM

would endow these

230

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

231

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

232

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

233

advantageous dynamical properties.

Note the ConvSSM formulation is easier to work with than

234

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

235

convolution operations in modern deep learning frameworks [76-78].

236

3.4

Efficient ConvSSM for Long-Range Dependencies: ConvS5

237

In this section, we introduce

which combines ideas of parallelization of convolutional

238

recurrences (Section 3.2) and the SSM connection (Section 3.3).

ConvS5 is a ConvSSM that

239

leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition

Proof.

See Appendix C.3.

228

Thus, parameterizing and initializing

A

SSM

with the specially structured HiPPO matrices [47] and

229

disretizing with a learnable timescale parameter to obtain

A

SSM

and

B

SSM

would endow these

230

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

231

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

232

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

233

advantageous dynamical properties.

Note the ConvSSM formulation is easier to work with than

234

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

235

convolution operations in modern deep learning frameworks [76-78].

236

3.4

Efficient ConvSSM for Long-Range Dependencies: ConvS5

237

In this section, we introduce

which combines ideas of parallelization of convolutional

238

ConvS5 is a ConvSSM that

Proof.

See Appendix C.3.

228

Thus, parameterizing and initializing

A

SSM

with the specially structured HiPPO matrices [47] and

229

disretizing with a learnable timescale parameter to obtain

A

SSM

and

B

SSM

would endow these

230

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

231

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

232

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

233

advantageous dynamical properties.

Note the ConvSSM formulation is easier to work with than

234

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

235

convolution operations in modern deep learning frameworks [76-78].

236

3.4

Efficient ConvSSM for Long-Range Dependencies: ConvS5

237

In this section, we introduce

which combines ideas of parallelization of convolutional

238

recurrences (Section 3.2) and the SSM connection (Section 3.3).

ConvS5 is a ConvSSM that

239

leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition

SSM

∈

b

226

columns formed by performing an im2col [74, 75] operation on the input

U

(

t

)

.

227

Proof.

See Appendix C.3.

228

Thus, parameterizing and initializing

A

SSM

229

disretizing with a learnable timescale parameter to obtain

A

SSM

and

B

SSM

230

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

231

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

232

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

233

advantageous dynamical properties.

234

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

235

convolution operations in modern deep learning frameworks [76-78].

236

3.4

Efficient ConvSSM for Long-Range Dependencies: ConvS5

237

where

A

Block

∈

R

HWP

×

HWP

and

B

Block

∈

R

HWP

×

HWUk

2

b

225

B

SSM

on the diagonal blocks respectively, and

u

(

t

)

∈

R

HWUk

226

columns formed by performing an im2col [74, 75] operation on the input

227

Proof.

See Appendix C.3.

228

Thus, parameterizing and initializing

A

SSM

229

disretizing with a learnable timescale parameter to obtain

A

230

231

232

233

advantageous dynamical properties.

234

235

convolution operations in modern deep learning frameworks [76-78].

236

3.4

Efficient ConvSSM for Long-Range Dependencies: ConvS5

237

H

′

×

W

′

×

U

212

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

[JS: im2col

216

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

217

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

218

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

219

Proposition 3.

Consider a ConvSSM state update as in (7) with pointwise state kernel

A

∈

220

′

′

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

216

217

218

P

×

P

×

1

×

1

214

H

×

W

×

P

215

216

217

218

219

Proposition 3.

220

B

H

′

×

W

′

×

U

P

×

U

×

k

B

×

k

B

P

×

P

×

1

×

1

H

×

W

×

P

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

[JS: im2col

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

Proposition 3.

Consider a ConvSSM state update as in (7) with pointwise state kernel

A

∈

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

, and input

U

(

t

)

∈

R

H

′

×

W

′

×

U

.

Rearrange the

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

. Rearrange the input kernel

B

into a matrix

B

SSM

∈

R

P

×

(

Uk

2

B

)

. The differential equation of (7) is equivalent, after reshaping

X

(

t

)

and

X

′

(

t

)

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

[JS: im2col

216

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

217

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

218

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

219

Proposition 3.

Consider a ConvSSM state update as in (7) with pointwise state kernel

A

∈

220

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

, and input

U

(

t

)

∈

R

H

′

×

W

′

×

U

.

Rearrange the

221

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

. Rearrange the input kernel

B

into a matrix

222

B

SSM

∈

R

P

×

(

Uk

2

B

)

. The differential equation of (7) is equivalent, after reshaping

X

(

t

)

and

X

′

(

t

)

223

into

x

(

t

)

,

x

′

(

t

)

∈

R

HWP

respectively, to the following differential equation

224

x

′

(

t

) =

A

Block

x

(

t

) +

B

Block

u

(

t

)

,

(12)

2

B

211

H

′

×

W

′

×

U

212

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

[JS: im2col

216

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

217

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

218

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

219

Proposition 3.

Consider a ConvSSM state update as in (7) with pointwise state kernel

A

∈

220

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

, and input

U

(

t

)

∈

R

H

′

×

W

′

×

U

.

Rearrange the

221

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

. Rearrange the input kernel

B

into a matrix

222

B

SSM

∈

R

P

×

(

Uk

2

B

)

. The differential equation of (7) is equivalent, after reshaping

X

(

t

)

and

X

′

(

t

)

223

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

. Rearrange the input kernel

213

B

SSM

∈

R

P

×

(

Uk

2

B

)

. The differential equation of (7) is equivalent, after reshaping

214

into

x

(

t

)

,

x

′

(

t

)

∈

R

HWP

respectively, to the following differential equation

215

x

′

(

t

) =

A

Block

x

(

t

) +

B

Block

u

(

t

)

,

where

A

Block

∈

R

HWP

×

HWP

and

B

Block

∈

R

HWP

×

HWUk

2

b

216

B

SSM

on the diagonal blocks respectively, and

u

(

t

)

∈

R

HWUk

2

b

217

columns formed by performing an im2col [74, 75] operation on the input

U

(

t

)

.

218

Proof.

See Appendix C.3.

219

Thus, parameterizing and initializing

A

SSM

220

disretizing with a learnable timescale parameter to obtain

A

SSM

and

B

SSM

221

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

222

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

223

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

H

×

W

×

P

215

216

217

218

219

Proposition 3.

220

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

221

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

222

B

SSM

∈

R

P

×

(

Uk

2

B

)

223

into

x

(

t

)

,

x

′

(

t

)

∈

R

HWP

224

x

′

(

t

) =

A

Block

x

(

t

) +

where

A

Block

∈

R

HWP

×

HWP

and

B

Block

∈

R

HWP

225

B

SSM

on the diagonal blocks respectively, and

u

(

t

)

226

227

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

216

217

218

219

Proposition 3.

220

R

P

×

P

×

1

×

1

, input kernel

R

P

×

U

×

k

B

×

k

B

, and input

(

t

)

R

H

B

SSM

∈

R

P

×

(

Uk

2

B

)

214

into

x

(

t

)

,

x

′

(

t

)

∈

R

HWP

215

x

′

(

t

) =

where

A

Block

∈

R

HWP

×

HWP

and

B

216

B

SSM

on the diagonal blocks respectively, and

217

218

Proof.

See Appendix C.3.

219

Thus, parameterizing and initializing

A

220

N

-D tensors and replacing the matrix-vector

L

H

′

×

W

′

×

U

, with height

H

′

, width

W

′

and

U

features

[SD: minor

columns of a matrix, so I currently have a preference to keep it as is]

, a

∈

R

H

×

W

×

P

, with a state update equation

X

k

=

G

(

X

k

-

1

,

U

k

)

,

2

L

[SD: minor

, a

G

(

X

k

-

1

,

U

k

)

,

features

[SD: minor

, a

k

=

G

(

X

k

-

1

,

U

k

)

,

Given a continuous input signal

u

(

t

)

∈

R

U

, a latent state

x

(

t

)

∈

R

P

and an output signal

y

(

t

)

∈

R

M

,

100

a continuous-time, linear SSM is defined by the differential equation:

101

x

′

(

t

) =

Ax

(

t

) +

Bu

(

t

)

,

y

(

t

) =

Cx

(

t

) +

Du

(

t

)

,

(3)

and is parameterized by a state matrix

A

∈

R

P

×

P

, an input matrix

B

∈

R

P

×

U

, an output matrix

102

C

∈

R

M

×

P

and a feedthrough matrix

D

∈

R

M

×

U

. Given a sampled sequence,

u

1:

L

∈

R

L

×

U

, the

103

SSM can be discretized to define a discrete-time SSM

104

x

k

=

Ax

k

-

1

+

Bu

k

,

y

k

=

Cx

k

+

Du

k

,

(4)

where the discrete-time parameters are a function of the continuous-time parameters and a timescale

105

parameter,

∆

. We define

A

= DISCRETIZE

A

(

A

,

∆)

and

B

= DISCRETIZE

B

(

A

,

B

,

∆)

where

106

DISCRETIZE()

is a discretization method such as Euler, bilinear or zero-order hold [69].

107

[SD: I feel like the recurrence part of SSM is not very clearly explained and hence hard to follow/understand.

108

x

′

(

t

) =

Ax

(

t

) +

Bu

(

t

)

,

y

(

t

) =

Cx

(

t

) +

Du

(

t

)

,

(3)

and is parameterized by a state matrix

A

∈

R

P

×

P

, an input matrix

B

∈

R

P

×

U

, an output matrix

102

C

∈

R

M

×

P

and a feedthrough matrix

D

∈

R

M

×

U

. Given a sampled sequence,

u

1:

L

∈

R

L

×

U

, the

103

SSM can be discretized to define a discrete-time SSM

104

x

k

=

Ax

k

-

1

+

Bu

k

,

y

k

=

Cx

k

+

Du

k

,

(4)

where the discrete-time parameters are a function of the continuous-time parameters and a timescale

105

parameter,

∆

. We define

A

= DISCRETIZE

A

(

A

,

∆)

and

B

= DISCRETIZE

B

(

A

,

B

,

∆)

where

106

DISCRETIZE()

is a discretization method such as Euler, bilinear or zero-order hold [69].

107

[SD: I feel like the recurrence part of SSM is not very clearly explained and hence hard to follow/understand.

108

This is point is also important to understand the Parallel Scan and why that is important.]

[JS: Agreed this is

109

x

′

(

t

) =

Ax

(

t

) +

Bu

(

t

)

,

y

(

t

) =

Cx

(

t

) +

and is parameterized by a state matrix

A

∈

R

P

×

P

, an input matrix

B

102

C

∈

R

M

×

P

and a feedthrough matrix

D

∈

R

M

×

U

103

SSM can be discretized to define a discrete-time SSM

104

x

k

=

Ax

k

-

1

+

Bu

k

,

y

k

=

Cx

k

+

Du

105

parameter,

∆

. We define

A

= DISCRETIZE

A

(

A

,

∆)

and

B

106

DISCRETIZE()

107

108

This is point is also important to understand the Parallel Scan and why that is important.]

109

and is parameterized by a state matrix

A

∈

R

P

×

P

, an input matrix

B

∈

R

P

×

U

, an output matrix

102

C

∈

R

M

×

P

and a feedthrough matrix

D

∈

R

M

×

U

. Given a sampled sequence,

u

1:

L

∈

R

L

×

U

, the

103

SSM can be discretized to define a discrete-time SSM

104

x

k

=

Ax

k

-

1

+

Bu

k

,

y

k

=

Cx

k

+

Du

k

,

(4)

where the discrete-time parameters are a function of the continuous-time parameters and a timescale

105

parameter,

∆

. We define

A

= DISCRETIZE

A

(

A

,

∆)

and

B

= DISCRETIZE

B

(

A

,

B

,

∆)

where

106

DISCRETIZE()

is a discretization method such as Euler, bilinear or zero-order hold [69].

107

[SD: I feel like the recurrence part of SSM is not very clearly explained and hence hard to follow/understand.

108

This is point is also important to understand the Parallel Scan and why that is important.]

[JS: Agreed this is

109

important. Can you point toward the part that is most unclear or maybe what seems missing?]

110

X

k

= tanh(

A

∗

X

k

-

1

+

B

∗

U

k

)

,

(2)

X

k

=

G

(

A

∗

X

k

-

1

+

B

∗

U

k

)

,

(3)

A

∈

R

P

×

P

×

k

A

×

k

A

(using an [output

B

∈

R

P

×

U

×

k

B

×

k

B

More commonly, sophisticated updates such as

equation would be

85

X

k

= tanh(

A

∗

h

(

x

k

)

86

h

(

X

k

)

87

u

′

k

88

equation would be

85

X

k

= tanh(

A

∗

X

k

-

1

+

B

∗

U

k

)

,

(2)

h

(

x

k

)

86

h

(

X

k

)

87

u

′

k

88

U

′

89

equation would be

85

X

k

= tanh(

A

∗

X

k

-

1

+

B

∗

U

k

)

,

h

(

x

k

)

86

h

(

X

k

)

87

u

′

k

88

U

′

89

sequence of frames,

U

1:

L

∈

R

L

×

H

′

×

W

′

×

U

, with height

H

′

, width

W

′

and

U

features, a ConvRNN

X

k

∈

R

H

×

W

×

P

, with a state update equation

X

k

=

G

(

X

k

-

1

,

U

k

)

, where

G

()

is

a nonlinear function. For a concrete example analogous to (1), the vanilla ConvRNN state update

Implementation is available at:

https://github.com/anonymous

.

2

by replacing the 1-D states and inputs of RNNs with

N

80

multiplications with convolutions as in convolutional neural networks (CNNs). Given a length

81

sequence of frames,

U

1:

L

∈

R

L

×

H

′

×

W

′

×

U

, with height

H

′

, width

W

′

and

U

82

point, but it is more typical to write WxH instead of HxW] [JS: Interesting, can you point me to some examples?

83

The

H

×

W

convention matches the expected shapes of Torch and JAX/Flax's conv2d operations and matches

84

the height and width of the rows

×

columns of a matrix, so I currently have a preference to keep it as is]

85

ConvRNN updates its state,

X

k

∈

R

H

×

W

×

P

, with a state update equation

X

k

=

86

2

by replacing the 1-D states and inputs of RNNs with

N

-D tensors and replacing the matrix-vector

80

multiplications with convolutions as in convolutional neural networks (CNNs). Given a length

L

81

sequence of frames,

U

1:

L

∈

R

L

×

H

′

×

W

′

×

U

, with height

H

′

, width

W

′

and

U

features

[SD: minor

82

point, but it is more typical to write WxH instead of HxW] [JS: Interesting, can you point me to some examples?

83

The

H

×

W

convention matches the expected shapes of Torch and JAX/Flax's conv2d operations and matches

84

the height and width of the rows

×

columns of a matrix, so I currently have a preference to keep it as is]

, a

85

ConvRNN updates its state,

X

k

∈

R

H

×

W

×

P

, with a state update equation

X

k

=

G

(

X

k

-

1

,

U

k

)

,

86

2

sequence of frames,

U

1:

L

∈

R

L

×

H

′

×

W

′

×

U

, with height

H

′

, width

W

′

and

U

features

[SD: minor

point, but it is more typical to write WxH instead of HxW] [JS: Interesting, can you point me to some examples?

The

H

×

W

convention matches the expected shapes of Torch and JAX/Flax's conv2d operations and matches

the height and width of the rows

×

columns of a matrix, so I currently have a preference to keep it as is]

, a

ConvRNN updates its state,

X

k

∈

R

H

×

W

×

P

, with a state update equation

X

k

=

G

(

X

k

-

1

,

U

k

)

,

2

85

h

(

x

k

)

86

h

(

X

k

)

87

u

′

k

88

U

′

k

89

85

h

(

x

k

)

86

h

(

X

k

)

87

u

′

k

88

U

′

k

89

Y

(

t

) =

C

∗

X

(

t

) +

D

∗

U

(

t

)

where

∗

denotes the convolution operator,

A

∈

R

P

×

P

×

k

A

×

k

A

is the state kernel,

155

is

the

input kernel,

C

∈

R

U

×

P

×

k

C

×

k

C

is the output kernel, and

D

∈

156

feedthrough kernel. We assume proper padding is used to ensure the same spatial resolution,

157

maintained in the states and outputs. Similarly, given a sequence of

L

inputs,

158

we can define a discrete-time convolutional state space model as

159

X

k

=

A

∗

X

k

-

1

+

B

∗

U

k

Y

k

=

C

∗

X

k

+

D

∗

U

k

where

A

∈

R

P

×

P

×

k

A

×

k

A

and

B

∈

R

P

×

U

×

k

B

×

k

B

160

161

162

ConvRNN

SSM

ConvSSM

(this paper)

Nonlinear function

Linear function

sequence of frames,

U

1:

L

∈

R

L

×

H

′

×

W

′

×

U

, with height

82

updates its state,

X

k

∈

R

H

×

W

×

P

83

84

1

Implementation is available at:

2

Given a 1-D sequence of inputs

u

1:

L

∈

R

L

×

U

, an RNN updates its state

x

k

∈

R

P

at each time step

73

k

with a state update equation

x

k

=

F

(

x

k

-

1

,

u

k

)

, where

F

()

represents some nonlinear function.

74

For example, one can represent a vanilla RNN (ignoring the bias term) as

75

x

k

= tanh(

Ax

k

-

1

+

Bu

k

)

(1)

with state matrix

A

∈

R

P

×

P

, input matrix

B

∈

R

P

×

U

and the hyperbolic tangent function

76

applied elementwise.

Other RNNs such as LSTM [13] and GRU [14] define more sophisticated

77

formulations of

F

()

.

78

Convolutional recurrent neural networks

[11, 12] (ConvRNNs) model spatiotemporal sequences

79

by replacing the 1-D states and inputs of RNNs with

N

-D tensors and replacing the matrix-vector

80

multiplications with convolutions as in convolutional neural networks (CNNs). Given a length

L

81

U

, with height

, width

and

features, a ConvRNN

equation would be

85

X

k

= tanh(

A

∗

X

k

-

1

h

(

x

k

)

86

h

(

X

k

)

87

u

′

k

88

b

99

97

Consider a ConvSSM state update as in (7)

R

are block-diagonal with

1

∈

1

(

t

)

H

′

229

are block-diagonal with

×

U

2

b

HWUk which combines ideas of parallelization of convolutional

A

A

H

A

SS

is the result of concatenati

[JS: Agreed this is

R

W

U

′

′

[JS: Agreed this is

. Rearrange the input kernel

×

×

columns formed by performing an im2col [74, 75] operation on the input

SSM

.

and

P

P

Note the ConvSSM formulation is easier to work with than

W

W

U

′

U

In this section, we introduce

238

k

A

recurrences (Section 3.2) and the SSM connection (Section 3.3).

A

∈

U

into a matrix

SSM

∈

R

L

×

L

H

with the specially struct

′

×

W

W

.

Rearrange the

B

(

t

)

U

′

1:

L

U

Rearrange the

90

ConvS5 is a ConvSSM that

∈

P

U

(S5) layer. An S5 layer replaces S4's many SISO SSMs with one convolution operations in modern deep learning frameworks [76-78].

Efficient ConvSSM for Long-Range Dependencies: ConvS5

convolution operations in modern deep learning frameworks [76-78].

121

Ax

122

×

×

×

B

k

k

H

U

′

′

which combines ideas of parallelization of convolu

∈

A

forecasting and numerous other applications [49-60] and many variants have been proposed [

239

240

S5

This section provides a brief background on the id

equations respectively. ConvRNNs have been used for biomedical, robotics, traffic

∈

- - - - - - →

101

102

(S5) layer. An S5 layer replaces S4's many SISO SSMs with one simplified S4

recurrences (Section 3.2) and the SSM connection (Section 3.3).

(S5) layer. An S5 layer replaces S4's many SISO SSMs with one

236

y

scans, resulting in the multi-input, multi-output SSM as in 3 and is a fully recurrent, stateful model for both parallel and

S5 [36] that were developed for 1-D sequence modeling and whose insights we leverage for our

Du

=

Cx

+

k

scans, resulting in the

,

simplified S4

k

+

sequential computations.

121

3.4

,

237

Bu

k

122

k

1

k

100

238

columns formed by performing an im2col [74, 75] operation on

2

B

P

(

Uk

)

B

∈

B

k

k

U

230

rearrange leverages parallel scans and deep SSM parameterization/initialization schemes. Given Propo

simplified S4

123

(5)

=

x

where the discrete-time parameters are a function of the continuous-time parameters and a timescale

∆

parameter,

. We define

DISCRETIZE()

237

238

229

A

-

102

2.2

123

237

P

R

103

4

multi-input, multi-output SSM as in 3 and is a fully recurrent, stateful model for both parallel and

∆)

B

103

= DISCRETIZE

scans, resulting in the

121

In this section, we introduce

B

S5

×

227

R

(15)

×

.

×

×

S5 [36] that were developed for 1-D sequence mo

98

A

recurrences (Section 3.2) and the SSM connection (Section 3.3).

∈

- - - - - - →

B

100

Efficient ConvSSM for Long-Range Dependencies: ConvS5

We then run the discretized ConvSSM system of Eqs (9, 10), where the recurrence can be computed

HWP

which combines ideas of parallelization of convolutional multi-input, multi-output SSM as in 3 and is a fully recurrent, stateful model for both parallel and

Thus, parameterizing and initializing

231

forecasting and numerous other applications [49-60] and many variants have been

100

99

232

∈

2.2

ConvS5 is a ConvSSM that

Block with the specially structured HiPPO matrices [47] and

×

101

2.2

HWP

leverages parallel scans and deep SSM parameterization/initializatio with the specially structured HiPPO matrices [47] and

101

and

A

SSM

HWUk

2

b

SSM

Linear State Space Models

This section provides a brief background on the ideas of deep SSM methods such as S4 [

B

×

SSM

which combines ideas of parallelization of convolutional

3, we can implicitly parameterize a pointwise state kernel,

A

109

This section provides a brief background on would endow these

122

236

HWP

SSM disretizing with a learnable timescale parameter to obtain A SSM and B SSM would endow these 230 SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range 231 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 235 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 B SSM on the diagonal blocks respectively, and u ( t ) ∈ R HWUk b is the result of 226 columns formed by performing an im2col [74, 75] operation on the input U ( t ) . 227 Proof. See Appendix C.3. 228 Thus, parameterizing and initializing A SSM with the specially structured HiPP 229 disretizing with a learnable timescale parameter to obtain A SSM and B SSM 230 SSMs with the same dynamics as in S4/S5 methods that are important for m 231 dependencies. Due to the equivalence of Proposition 3, we can then rearrange 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional r 233 advantageous dynamical properties. Note the ConvSSM formulation is easier 234 forming the large, block-diagonal SSM state update explicitly and allows leveragi 235 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of parallelizatio 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 i 239 leverages parallel scans and deep SSM parameterization/initialization schemes. 240 6 Proof. See Appendix C.3. 228 Thus, parameterizing and initializing A SSM with the specially 229 disretizing with a learnable timescale parameter to obtain A 230 SSMs with the same dynamics as in S4/S5 methods that are 231 dependencies. Due to the equivalence of Proposition 3, we ca 232 the discrete ConvSSM state and input kernels of 9 to give the 233 advantageous dynamical properties. Note the ConvSSM for 234 forming the large, block-diagonal SSM state update explicitly an 235 convolution operations in modern deep learning frameworks [7 236 3.4 Efficient ConvSSM for Long-Range Dependencies: Co 237 In this section, we introduce ConvS5 which combines ideas 238 recurrences (Section 3.2) and the SSM connection (Section 239 leverages parallel scans and deep SSM parameterization/initial 240 6 In this section, 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 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 In this section, we introduce ConvS5 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 using parallel scans. Note that in practice this setup actually allows us to parameterize ConvS5 using 237 a diagonal state [34-36] which reduces the cost of applying the parallel scan in Proposition 2 to 238 O ( LPHW ) . See Appendix A for a more detailed discussion of parameterization, initialization and 239 discretization. 240 6 disretizing with a learnable timescale parameter to obtain 230 SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range 231 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 235 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 ∈ B ∈ R P × U × k B × k B in (7) using SSM parameters as used by S5 [36 233 R P × ( Uk 2 B ) . We can discretize these S5 SSM parameters as discussed 234 A S5 = DISCRETIZE A ( A S5 , ∆ ) , B S5 = DISCRETI and then rearrange to give the ConvS5 state update kernels: 235 A S5 ∈ R P × P rearrange - - - - - - → A ConvS5 ∈ R P × P B S5 ∈ R P × ( Uk 2 B ) rearrange - - - - - - → B ConvS5 ∈ R P × U We then run the discretized ConvSSM system of Eqs (9, 10), where th 236 using parallel scans. Note that in practice this setup actually allows us 237 a diagonal state [34-36] which reduces the cost of applying the par 238 O ( LPHW ) . See Appendix A for a more detailed discussion of para 239 discretization. 240 6 123 2.2.3 SSM Parameterization and Initialization 124 Parameterization and initialization are important aspects of deep SSMs such as S4 and S5 that 125 allows them to capture long-range dependencies. These methods rely on continuous-time SSM 126 parameters, explicitly performing a discretization step with a learnable timescale parameter, and 127 a deterministic initialization of the state matrix A with specially structured matrices inspired by 128 the HiPPO framework [47]. Many ablations have been performed highlighting the importance of 129 this setup for deep SSMs to achieve high performance on challenging long-range tasks [33, 36, 70]. 130 Concurrent work [71] has studied these initializations in more detail and provides insight into the 131 favorable initial eigenvalue distributions and normalization effects of this setup. 132 3 2.2.3 SSM Parameterization and Initialization 124 Parameterization and initialization are important aspects of deep SSMs such as S4 and S5 that 125 allows them to capture long-range dependencies. These methods rely on continuous-time SSM 126 parameters, explicitly performing a discretization step with a learnable timescale parameter, and 127 a deterministic initialization of the state matrix A with specially structured matrices inspired by 128 the HiPPO framework [47]. Many ablations have been performed highlighting the importance of 129 this setup for deep SSMs to achieve high performance on challenging long-range tasks [33, 36, 70]. 130 Concurrent work [71] has studied these initializations in more detail and provides insight into the 131 favorable initial eigenvalue distributions and normalization effects of this setup. 132 3 sequential computations. 123 2.2.3 SSM Parameterization and Initialization 124 Parameterization and initialization are important aspects of deep SSMs such as S4 and S5 that 125 allows them to capture long-range dependencies. These methods rely on continuous-time SSM 126 parameters, explicitly performing a discretization step with a learnable timescale parameter, and 127 a deterministic initialization of the state matrix A with specially structured matrices inspired by 128 the HiPPO framework [47]. Many ablations have been performed highlighting the importance of 129 this setup for deep SSMs to achieve high performance on challenging long-range tasks [33, 36, 70]. 130 Concurrent work [71] has studied these initializations in more detail and provides insight into the 131 favorable initial eigenvalue distributions and normalization effects of this setup. 132 3 2.2.3 SSM Parameterization and Initialization 124 Parameterization and initialization are important aspects of deep SSMs such as S4 and S5 that 125 allows them to capture long-range dependencies. These methods rely on continuous-time SSM 126 parameters, explicitly performing a discretization step with a learnable timescale parameter, and 127 a deterministic initialization of the state matrix A with specially structured matrices inspired by 128 the HiPPO framework [47]. Many ablations have been performed highlighting the importance of 129 this setup for deep SSMs to achieve high performance on challenging long-range tasks [33, 36, 70]. 130 Concurrent work [71] has studied these initializations in more detail and provides insight into the 131 favorable initial eigenvalue distributions and normalization effects of this setup. 132 3 A B is a discretization method such as Euler, bilinear or zero-order hold [69]. 2.2.2 State Space Layers: S4 and S5 Gu et al. [33] proposed the structured state space sequence (S4) layer, which efficiently models long sequences with subquadratic complexity in the sequence length. An S4 layer consists of many continuous-time parameterized single-input, single-output (SISO) SSMs, an explicit discretization step with learnable timescale parameters, and position-wise nonlinear activation functions applied to the SSM outputs to produce the layer output sequence. Multiple layers can be stacked to form deep sequence models that can represent nonlinear systems. An S4 layer can use its linear SSMs to process a sequence in parallel by using 1-D convolutions. It can also run its linear SSMs as RNNs to allow for fast autoregressive generation. Smith et al. [36] showed that with several architecture changes, the approach could be simplified and made more flexible by replacing the convolutions with parallel scans, resulting in the simplified S4 (S5) layer. An S5 layer replaces S4's many SISO SSMs with one multi-input, multi-output SSM as in 4 and is a fully recurrent, stateful model for both parallel and sequential computations. 2.2.3 SSM Parameterization and Initialization Parameterization and initialization are important aspects of deep SSMs such as S4 and S5 that allows them to capture long-range dependencies. These methods rely on continuous-time SSM parameters, explicitly performing a discretization step with a learnable timescale parameter, and a deterministic initialization of the state matrix A with specially structured matrices inspired by the HiPPO framework [47]. Many ablations have been performed highlighting the importance of this setup for deep SSMs to achieve high performance on challenging long-range tasks [33, 36, 70]. Concurrent work [71] has studied these initializations in more detail and provides insight into the favorable initial eigenvalue distributions and normalization effects of this setup. 3 2.2 Deep State Space Models 100 This section provides a brief background on the ideas of deep SSM methods su 101 S5 [36] that were developed for 1-D sequence modeling and whose insights w 102 spatiotemporal ConvS5 approach. 103 2.2.1 Linear State Space Models 104 Given a continuous input signal u ( t ) ∈ R U , a latent state x ( t ) ∈ R P and an output 105 a continuous-time, linear SSM is defined by the differential equation: 106 x ′ ( t ) = Ax ( t ) + Bu ( t ) , y ( t ) = Cx ( t ) + Du ( t ) , and is parameterized by a state matrix A ∈ R P × P , an input matrix B ∈ R P × U , 107 C ∈ R M × P and a feedthrough matrix D ∈ R M × U . Given a sampled sequence, 108 SSM can be discretized to define a discrete-time SSM 109 x k = Ax k -1 + Bu k , y k = Cx k + Du k , where the discrete-time parameters are a function of the continuous-time paramete 110 parameter, ∆ . We define A = DISCRETIZE A ( A , ∆) and B = DISCRETIZE 111 DISCRETIZE() is a discretization method such as Euler, bilinear or zero-order 112 2.2.2 State Space Layers: S4 and S5 113 Gu et al. [33] proposed the structured state space sequence (S4) layer, which 114 long sequences with subquadratic complexity in the sequence length. An S4 laye 115 continuous-time parameterized single-input, single-output (SISO) SSMs, an exp 116 step with learnable timescale parameters, and position-wise nonlinear activation fu 117 the SSM outputs to produce the layer output sequence. Multiple layers can be sta 118 sequence models that can represent nonlinear systems. An S4 layer can use its linea 119 a sequence in parallel by using 1-D convolutions. It can also run its linear SSMs 120 for fast autoregressive generation. Smith et al. [36] showed that with several arc 121 the approach could be simplified and made more flexible by replacing the convolu 122 scans, resulting in the simplified S4 (S5) layer. An S5 layer replaces S4's many SI 123 multi-input, multi-output SSM as in 3 and is a fully recurrent, stateful model for 124 sequential computations. 125 3 2.2.1 Linear State Space Models 104 Given a continuous input signal u ( t ) ∈ R U , a latent state x ( t ) ∈ R P and an output signal y ( t ) ∈ R M , 105 a continuous-time, linear SSM is defined by the differential equation: 106 x ′ ( t ) = Ax ( t ) + Bu ( t ) , y ( t ) = Cx ( t ) + Du ( t ) , (3) and is parameterized by a state matrix A ∈ R P × P , an input matrix B ∈ R P × U , an output matrix 107 C ∈ R M × P and a feedthrough matrix D ∈ R M × U . Given a sampled sequence, u 1: L ∈ R L × U , the 108 SSM can be discretized to define a discrete-time SSM 109 x k = Ax k -1 + Bu k , y k = Cx k + Du k , (4) where the discrete-time parameters are a function of the continuous-time parameters and a timescale 110 parameter, ∆ . We define A = DISCRETIZE A ( A , ∆) and B = DISCRETIZE B ( A , B , ∆) where 111 DISCRETIZE() is a discretization method such as Euler, bilinear or zero-order hold [69]. 112 2.2.2 State Space Layers: S4 and S5 113 Gu et al. [33] proposed the structured state space sequence (S4) layer, which efficiently models 114 long sequences with subquadratic complexity in the sequence length. An S4 layer consists of many 115 continuous-time parameterized single-input, single-output (SISO) SSMs, an explicit discretization 116 step with learnable timescale parameters, and position-wise nonlinear activation functions applied to 117 the SSM outputs to produce the layer output sequence. Multiple layers can be stacked to form deep 118 sequence models that can represent nonlinear systems. An S4 layer can use its linear SSMs to process 119 a sequence in parallel by using 1-D convolutions. It can also run its linear SSMs as RNNs to allow 120 for fast autoregressive generation. Smith et al. [36] showed that with several architecture changes, 121 the approach could be simplified and made more flexible by replacing the convolutions with parallel 122 scans, resulting in the simplified S4 (S5) layer. An S5 layer replaces S4's many SISO SSMs with one 123 multi-input, multi-output SSM as in 3 and is a fully recurrent, stateful model for both parallel and 124 sequential computations. 125 3 This section provides a brief background on the ideas of deep SSM methods such as S4 [33] and 101 S5 [36] that were developed for 1-D sequence modeling and whose insights we leverage for our 102 spatiotemporal ConvS5 approach. 103 2.2.1 Linear State Space Models 104 Given a continuous input signal u ( t ) ∈ R U , a latent state x ( t ) ∈ R P and an output signal y ( t ) ∈ R M , 105 a continuous-time, linear SSM is defined by the differential equation: 106 x ′ ( t ) = Ax ( t ) + Bu ( t ) , y ( t ) = Cx ( t ) + Du ( t ) , (3) and is parameterized by a state matrix A ∈ R P × P , an input matrix B ∈ R P × U , an output matrix 107 C ∈ R M × P and a feedthrough matrix D ∈ R M × U . Given a sampled sequence, u 1: L ∈ R L × U , the 108 SSM can be discretized to define a discrete-time SSM 109 x k = Ax k -1 + Bu k , y k = Cx k + Du k , (4) where the discrete-time parameters are a function of the continuous-time parameters and a timescale 110 parameter, ∆ . We define A = DISCRETIZE A ( A , ∆) and B = DISCRETIZE B ( A , B , ∆) where 111 DISCRETIZE() is a discretization method such as Euler, bilinear or zero-order hold [69]. 112 2.2.2 State Space Layers: S4 and S5 113 Gu et al. [33] proposed the structured state space sequence (S4) layer, which efficiently models 114 long sequences with subquadratic complexity in the sequence length. An S4 layer consists of many 115 continuous-time parameterized single-input, single-output (SISO) SSMs, an explicit discretization 116 step with learnable timescale parameters, and position-wise nonlinear activation functions applied to 117 the SSM outputs to produce the layer output sequence. Multiple layers can be stacked to form deep 118 sequence models that can represent nonlinear systems. An S4 layer can use its linear SSMs to process 119 a sequence in parallel by using 1-D convolutions. It can also run its linear SSMs as RNNs to allow 120 for fast autoregressive generation. Smith et al. [36] showed that with several architecture changes, 121 the approach could be simplified and made more flexible by replacing the convolutions with parallel 122 scans, resulting in the simplified S4 (S5) layer. An S5 layer replaces S4's many SISO SSMs with one 123 multi-input, multi-output SSM as in 3 and is a fully recurrent, stateful model for both parallel and 124 sequential computations. 125 3 Given a continuous input signal u ( t ) ∈ R U , a latent s 105 a continuous-time, linear SSM is defined by the diff 106 x ′ ( t ) = Ax ( t ) + Bu ( t ) , and is parameterized by a state matrix A ∈ R P × P , 107 C ∈ R M × P and a feedthrough matrix D ∈ R M × U . 108 SSM can be discretized to define a discrete-time SS 109 x k = Ax k -1 + Bu k , where the discrete-time parameters are a function of 110 parameter, ∆ . We define A = DISCRETIZE A ( A , 111 DISCRETIZE() is a discretization method such as 112 2.2.2 State Space Layers: S4 and S5 113 Gu et al. [33] proposed the structured state space 114 long sequences with subquadratic complexity in the 115 continuous-time parameterized single-input, single116 step with learnable timescale parameters, and positio 117 the SSM outputs to produce the layer output sequen 118 sequence models that can represent nonlinear system 119 a sequence in parallel by using 1-D convolutions. It 120 for fast autoregressive generation. Smith et al. [36] 121 the approach could be simplified and made more fle 122 scans, resulting in the simplified S4 (S5) layer. An S 123 multi-input, multi-output SSM as in 3 and is a fully 124 sequential computations. 125 3 S5 [36] that were developed for 1-D sequen 102 spatiotemporal ConvS5 approach. 103 2.2.1 Linear State Space Models 104 Given a continuous input signal u ( t ) ∈ R U , a 105 a continuous-time, linear SSM is defined by t 106 x ′ ( t ) = Ax ( t ) + Bu ( and is parameterized by a state matrix A ∈ 107 C ∈ R M × P and a feedthrough matrix D ∈ 108 SSM can be discretized to define a discrete-ti x k = Ax k -1 + B where the discrete-time parameters are a func 110 parameter, ∆ . We define A = DISCRETIZ 111 DISCRETIZE() is a discretization method s 112 2.2.2 State Space Layers: S4 and S5 113 Gu et al. [33] proposed the structured state 114 long sequences with subquadratic complexit 115 continuous-time parameterized single-input, 116 step with learnable timescale parameters, and 117 the SSM outputs to produce the layer output 118 sequence models that can represent nonlinear 119 a sequence in parallel by using 1-D convolut 120 for fast autoregressive generation. Smith et 121 the approach could be simplified and made m 122 scans, resulting in the simplified S4 (S5) layer. 123 multi-input, multi-output SSM as in 3 and is 124 sequential computations. 125 S5 [36] that were developed for 1-D sequence modeling and whose insights we leverage 102 spatiotemporal ConvS5 approach. 103 2.2.1 Linear State Space Models 104 Given a continuous input signal u ( t ) ∈ R U , a latent state x ( t ) ∈ R P and an output signal y ( t ) 105 a continuous-time, linear SSM is defined by the differential equation: 106 x ′ ( t ) = Ax ( t ) + Bu ( t ) , y ( t ) = Cx ( t ) + Du ( t ) , and is parameterized by a state matrix A ∈ R P × P , an input matrix B ∈ R P × U , an output 107 C ∈ R M × P and a feedthrough matrix D ∈ R M × U . Given a sampled sequence, u 1: L ∈ R L 108 SSM can be discretized to define a discrete-time SSM 109 x k = Ax k -1 + Bu k , y k = Cx k + Du k , where the discrete-time parameters are a function of the continuous-time parameters and a ti 110 parameter, ∆ . We define A = DISCRETIZE A ( A , ∆) and B = DISCRETIZE B ( A , B , ∆ 111 DISCRETIZE() is a discretization method such as Euler, bilinear or zero-order hold [69]. 112 2.2.2 State Space Layers: S4 and S5 113 Gu et al. [33] proposed the structured state space sequence (S4) layer, which efficiently 114 long sequences with subquadratic complexity in the sequence length. An S4 layer consists 115 continuous-time parameterized single-input, single-output (SISO) SSMs, an explicit discre 116 step with learnable timescale parameters, and position-wise nonlinear activation functions ap 117 the SSM outputs to produce the layer output sequence. Multiple layers can be stacked to for 118 sequence models that can represent nonlinear systems. An S4 layer can use its linear SSMs to 119 a sequence in parallel by using 1-D convolutions. It can also run its linear SSMs as RNNs t 120 for fast autoregressive generation. Smith et al. [36] showed that with several architecture c 121 the approach could be simplified and made more flexible by replacing the convolutions with 122 scans, resulting in the simplified S4 (S5) layer. An S5 layer replaces S4's many SISO SSMs 123 multi-input, multi-output SSM as in 3 and is a fully recurrent, stateful model for both para 124 sequential computations. 125 3 Convolutional recurrent networks (ConvRNNs) such as ConvLSTM [17] and ConvGRU [18] are common approaches for spatiotemporal modeling. These methods encode the spatial information using tensor-structured states. The states are updated with recurrent neural network (RNN) equations that use convolutions instead of the matrix-vector multiplications in standard RNNs (e.g., LSTM/GRUs [21, 22]). This approach allows the RNN states to reflect the spatial structure of the data while simultaneously capturing temporal dynamics. ConvRNNs inherit both the benefits and the weaknesses of RNNs: they allow fast, stateful autoregressive generation and an unbounded context window, but they are slow to train due to their inherently sequential structure and can suffer from the vanishing/exploding gradient problem [23]. Transformer-based methods [9, 13, 14, 24-27] operate on an entire sequence in parallel, avoiding these training challenges. Transformers typically require sophisticated compression schemes [28-30] to reduce the spatiotemporal sequence into tokens. Moreover, Transformers use an attention mechanism that has a bounded context window and whose computational complexity scales quadratically in sequence length for training and inference [31, 32]. More efficient Transformer methods improve on the complexity of attention [33-39], but these methods can fail on sequences with long-range dependencies [40, 13]. Some approaches combine Transformers with specialized training frameworks to address the attention costs [13]. However, recent work in deep state space models (SSMs) [19, 41, 42, 20, 43], like S4 [19] and S5 [20], has sought to overcome attention's quadratic complexity while maintaining the parallelizability and performance of attention and the statefulness of RNNs. These SSM layers have proven to be effective in various domains such as speech [44], images [45] and video classification [45, 46]; reinforcement learning [47, 48]; forecasting [49] and language modeling [50-53]. Inspired by modeling ideas from ConvRNNs and SSMs, we introduce convolutional state space models (ConvSSMs), which have a tensor-structured state like ConvRNNs but a continuous-time parameterization and linear state updates like SSM layers. See Figure 1. However, there are challenges to make this approach scalable and effective for modeling long-range spatiotemporal data. In this paper, we address these challenges and provide a rigorous framework that ensures both computational efficiency and modeling performance for spatiotemporal sequence modeling. First, we discuss computational efficiency and parallelization of ConvSSMs across the sequence for scalable training and fast inference. We show how to parallelize linear convolutional recurrences using a binary associative operator and demonstrate how this can be exploited to use parallel scans for subquadratic parallelization across the spatiotemporal sequence. We discuss both theoretical and practical considerations (Section 3.2) required to make this feasible and efficient. Next, we address how to capture long-range spatiotemporal dependencies. We develop a connection between

2.2.1

6

B

104

2

A

∈

are block-diago

Efficient ConvSSM for Long-Range Dependencies: ConvS5

B

,

∆)

where

R

sequential computations.

Thus, parameterizing and initializing

A

Block

(

A

,

sequential computations.

∈

229

3.4

where

= DISCRETIZE

(

A

,

3.4

225

and

R

.

B

R

′

×

×

R

227

×

×

(

t

∈

k

B

213

L

R

×

×

×

∈

∈

×

k

U

P

×

×

×

X

∈

R

×

ta

the dynamics of SSMs and ConvSSMs (Section 3.3) and leverage this, in Section 3.4, to introduce a parameterization and initialization design that can capture long-range spatiotemporal dependencies.

As a result, we introduce ConvS5 , a new spatiotemporal layer that is an efficient ConvSSM variant. It is parallelizable and overcomes difficulties during training (e.g., vanishing/exploding gradient problems) that traditional ConvRNN approaches experience. ConvS5 does not require compressing frames into tokens and provides an unbounded context. It also provides fast (constant time and memory per step) autoregressive generation compared to Transformers. ConvS5 significantly outperforms Transformers and ConvLSTM on a challenging long horizon Moving-MNIST [54] experiment requiring methods to train on 600 frames and generate up to 1,200 frames. In addition, ConvS5 trains 3 × faster than ConvLSTM on this task and generates samples 400 × faster than the Transformer. Finally, we show that ConvS5 matches or exceeds the performance of various state-of-the-art methods on challenging DMLab, Minecraft, and Habitat long-range video prediction benchmarks [13].

2 Background

This section provides the background necessary for ConvSSMs and ConvS5, introduced in Section 3.

2.1 Convolutional Recurrent Networks

Given a sequence of inputs u 1: L ∈ R L × U , an RNN updates its state, x k ∈ R P , using the state update equation x k = F ( x k -1 , u k ) , where F () is a nonlinear function. For example, a vanilla RNN can be represented (ignoring the bias term) as

with state matrix A ∈ R P × P , input matrix B ∈ R P × U 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 a length L sequence of frames, U 1: L ∈ R L × H ′ × W ′ × U , with height H ′ , width W ′ and U features, a ConvRNN updates its state, X k ∈ R H × W × P , with a state update equation X k = G ( X k -1 , U k ) , where G () is a nonlinear function. Analogous to (1), we can express the state update equation for a vanilla ConvRNN as

where ∗ is a spatial convolution operator with state kernel A ∈ R P × P × k A × k A (using an [output features, input features, kernel height, kernel width] convention), input kernel B ∈ R P × U × k B × k B and tanh() is applied elementwise. More complex updates such as ConvLSTM [17] and ConvGRU [18] are commonly used by making similar changes to the LSTM and GRU equations, respectively.

2.2 Deep State Space Models

This section briefly introduces deep SSMs such as S4 [19] and S5 [20] designed for modeling long sequences. The ConvS5 approach we introduce in Section 3 extends these ideas to the spatiotemporal domain.

Linear State Space Models Given a continuous input signal u ( t ) ∈ R U , a latent state x ( t ) ∈ R P and an output signal y ( t ) ∈ R M , a continuous-time, linear SSM is defined using a differential equation:

and is parameterized by a state matrix A ∈ R P × P , an input matrix B ∈ R P × U , an output matrix C ∈ R M × P and a feedthrough matrix D ∈ R M × U . Given a sequence, u 1: L ∈ R L × U , the SSM can be discretized to define a discrete-time SSM

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 = DISCRETIZE B ( A , B , ∆) where DISCRETIZE() is a discretization method such as Euler, bilinear or zero-order hold [55].

S4 and S5 Gu et al. [19] introduced the structured state space sequence (S4) layer to efficiently model long sequences. An S4 layer uses many continuous-time linear SSMs, an explicit discretization step with learnable timescale parameters, and position-wise nonlinear activation functions applied to the SSM outputs. Smith et al. [20] showed that with several architecture changes, the approach could be simplified and made more flexible by just using one SSM as in (3) and utilizing parallel scans. SSM layers, such as S4 and S5, take advantage of the fact that linear dynamics can be parallelized with subquadratic complexity in the sequence length. They can also be run sequentially as stateful RNNs for fast autoregressive generation. While a single SSM layer such as S4 or S5 has only linear dynamics, the nonlinear activations applied to the SSM outputs allow representing nonlinear systems by stacking multiple SSM layers [56-58].

SSM Parameterization and Initialization Parameterization and initialization are crucial aspects that allow deep SSMs to capture long-range dependencies more effectively than prior attempts at linear RNNs [59-61]. The general setup includes continuous-time SSM parameters, explicit discretization with learnable timescale parameters, and state matrix initialization using structured matrices inspired by the HiPPO framework [62]. Prior research emphasizes the significance of these choices in achieving high performance on challenging long-range tasks [19, 20, 56, 57]. Recent work [57] has studied these parameterizations/initializations in more detail and provides insight into this setup's favorable initial eigenvalue distributions and normalization effects.

2.3 Parallel Scans

We briefly introduce parallel scans, as used by S5, since they are important for parallelizing the ConvS5 method we introduce in Section 3. See Blelloch [63] for a more detailed review. A scan operation, given a binary associative operator · (i.e. ( a · b ) · c = a · ( b · c ) ) and a sequence of L elements [ a 1 , a 2 , ..., a L ] , yields the sequence: [ a 1 , ( a 1 · a 2 ) , ..., ( a 1 · a 2 · ... · a L )] .

Parallel scans use the fact that associative operators can be computed in any order. A parallel scan can be defined for the linear recurrence of the state update in (4) by forming the initial scan tuples c k = ( c k,a , c k,b ) := ( A , Bu k ) and utilizing a binary associative operator that takes two tuples q i , q j (either the initial tuples c i , c j or intermediate tuples) and produces a new tuple of the same type, q i · q j := ( q j,a ⊙ q i,a , q j,a ⊗ q i,b + q j,b ) , where ⊙ is matrix-matrix multiplication and ⊗ is matrix-vector multiplication. Given sufficient processors, the parallel scan computes the linear recurrence of (4) in O (log L ) sequential steps (i.e., depth or span) [63].

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 ConvSSM variant, ConvS5.

3.1 Convolutional State Space Models

Consider a continuous tensor-valued input U ( t ) ∈ R H ′ × W ′ × U with height H ′ , width W ′ , and number of input features U . We will define a continuous-time, linear convolutional state space model ( ConvSSM ) with state X ( t ) ∈ R H × W × P , derivative X ′ ( t ) ∈ R H × W × P and output Y ( t ) ∈ R H × W × U , using a differential equation:

where ∗ denotes the convolution operator, A ∈ R P × P × k A × k A is the state kernel, B ∈ R P × U × k B × k B is the input kernel, C ∈ R U × P × k C × k C is the output kernel, and D ∈ R U × U × k D × k D is the feedthrough kernel. For simplicity, we pad the convolution to ensure the same spatial resolution, H × W , is maintained in the states and outputs. Similarly, given a sequence of L inputs, U 1: L ∈ R L × H ′ × W ′ × U , we define a discrete-time convolutional state space model as

where A ∈ R P × P × k A × k A and B ∈ R P × U × k B × k B denote that these kernels are in discrete-time.

207

208

209

210

211

212

213

214

215

198

202

199

200

201

203

205

206

202

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

3.3

use this fact to leverage the parameterization and initialization ideas from deep SSMs discussed in

197

209

198

199

207

use this fact to leverage the parameterization and initialization ideas from d

We show in this section that an equivalence can be established between the dynamics of a ConvSSM

196

3

kernel. Computationally, we now have a construction

Connection to State Space Models

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as flattening the input and state tensors into vectors [73].

or compared to ConvLSTM with a

205

196

Connection to State Space Models

Thus, any ConvSSM can be described as

Thus, any ConvSSM can be described as

210

198

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and flattening the input and state tensors into vectors [73].

201

3

Proof.

3.3

209

a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

206

a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

197

197

using pointwise state kernels and the dynamics of a large block-diagonal SSM system. Thus, we can

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependen

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as

Thus, any ConvSSM can be described as

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependencies.

200

201

2.

Proof.

202

199

208

can be efficiently parallelized with subquadratic complexity with respect to the sequence length

Since the convolutions in 7-8 and 9-10 are linear operations, they can be

×

3.3

state kernels as described in the previous section provides an alternative SSM equivalence, which

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as flattening the input and state tensors into vectors [73].

210

198

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

3.3

lends a special structure that can leverage the deep SSM parameterization and initialization ideas state,

198

200

Proposition

The of

the convolutional

state kernels as described in the previous section provides an alternative SSM equivalence, which

207

197

We show in this section that an equivalence can be established bet

203

204

204

206

205

207

206

recurrence

Thus, any ConvSSM can be described as

Connection to State Space Models use this fact to leverage the parameterization and initialization ideas from deep SSMs discussed in

Connection to State Space Models a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

211

202

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently des

We show in this section that an equivalence can be e computational

199

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

201

209

Connection to State Space Models cost

flattening the input and state tensors into vectors [73].

in

Equation matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

(

t

) =

′

2

state kernels as described in the previous section provides an alternative SSM equivalence, which

202

210

P

Proposition

(

203

UPk

197

198

matrix-vector multiplications using large, circulant matrices consisting o

206

(

t

)

3.3

(

t

) +

a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

X

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

200

The

∗

A

matrix-vector multiplications using large, circulant matrices consisting of the kernel ele

X

cost

(

t

) +

of the

(

convolutional

t

)

recurrence discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

(

W

U

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependencies.

HWL

.

)

203

204

Connection to State Space Models

198

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

203

R

using pointwise state kernels and the dynamics of a large block-di

208

207

207

′

flattening the input and state tensors into vectors [73].

∗

t

) =

U

computational

211

i,j

P

201

We show in this section that an equivalence can be established between the dynamics of a Conv

X

O

3.3

B

(

state kernels as described in the previous section provides an alternative SSM equivalence, which

H

R

′

199

can be equivalently described as evolving according to a differential equation

208

state, in

Equation

?

a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise can be equivalently described as evolving according to a differential equation

(

t

)

state,

X

B

U

flattening the input and state tensors into vectors [73].

A

×

We show in this section that an equivalence can be established between the d

SSM

. See Figure 1.

∗

Thus, any ConvSSM can be described as

199

200

using pointwise state kernels and the dynamics of a l

2

X

Thus, any Conv

∗

A

X

can be equivalently described as evolving according to a differential equation

2.

, and input matrix,

HWL

We show in this section that an equivalence can be established between the dynamics of a Conv flattening the input and state tensors into vectors [73].

with a shared state matrix,

X

A

′

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

211

i,j

t

)

207

204

P

using pointwise state kernels and the dynamics of a large block-diagonal SSM system. Thus, w

′

(

(

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

202

SSM

Thus, any ConvSSM can be des

, and input matrix, a large SSM with a circulant dynamics matrix. However, we show here

×

212

200

B

)

)

t

(

with a shared state matrix,

(

UPk

P

SSM

A

212

) =

.

P

∈

SSM

(

P

R

∈

t

. See Figure 1.

203

SSM

207

201

202

B

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

(

×

W

P

(

t

X

)

t

U

205

R

206

∈

199

We show in this section that an equivalence can be established between the dynamics of a ConvSSM

i,j

)

t

t

t

(

)

Connection to State Space Models use this fact to leverage the parameterization and initialization id

208

205

208

209

209

196

)

t

B

SSM

i,j

)

∈

state kernels as described in the previous section provides an alternative SSM equivalence, which

208

R

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

R

. See Figure 1.

203

209

using pointwise state kernels and the dynamics of a large block-diagonal SSM system. Thus, we with a shared state matrix,

O

X

a large SSM with a circulant dynamics matrix. However, we show here that the use of

X

209

i,j using pointwise state kernels and the dynamics of a large block-diagonal SS

′

U

H

′

W

can be equivalently described as evolving according to a differential equation state,

use this fact to leverage the parameterization and in

(

′

(

×

×

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the state kernels as described in the previous section provides an alternative

B

×

, and input matrix,

H

) +

204

R

204

205

200

206

P

with a shared state matrix,

∈

204

i,j

(

t

A

′

∗

X

3.3

R

×

′

∗

1

′

(

t

i,j

)

P

P

U

U

lends a special structure that can leverage the deep SSM parameterizati

196

using pointwise state kernels and the dynamics of a large block-diagonal SSM system. Thus, we can i,j

)

Proof.

SSM

∈

196

213

207

214

use this fact to leverage the parameterization and initialization ideas from deep SSMs discuss

X

)

i,j

(

t

P

X

′

X

i,j

A

H

X

A

205

use this fact to leverage the parameterization and initialization ideas from i,j

R

can be equivalently described as evolving according to a differential equation

, and input matrix,

A

H

SSM

W

P

P

with a shared state matrix,

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependen

X

×

(

t

)

201

×

B

SSM

SSM

∈

×

213

P

205

206

. See Figure 1.

, and input matrix,

SSM

B

A

. See Figure 1.

210

R

3.3

Connection to State Space Models

)

X

W

P

state kernels as described in the previous section provides an alternative SSM equivalen

208

209

t

′

)

SSM

use this fact to leverage the parameterization and initialization ideas from deep SSMs discuss

210

196

197

198

199

200

201

202

203

204

205

206

207

∈

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as

X

state,

(

t

)

U

P

×

×

×

202

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

Section 2.2.3 to allow ConvSSMs to effectively model long-range

∈

′

×

′

(

(

t

t

P

P

)

)

Connection to State Space Models use this fact to leverage the parameterization and initialization ideas from deep SSMs discussed in

(

P

1

X

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependencies.

(

t

(

t

)

212

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependencies.

lends a special structure that can leverage the deep SSM parameterization and initializat

t

Section 2.2.3 to allow ConvSSMs to effectively mo

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivale

Connection to State Space Models

) +

(

discussed in Section 2.2.3 for modeling long-range dependencies. We s

t

)

matrix-vector multiplications using large, circulant matrices consisting of the ker

∈

A

R

∈

197

U

×

210

Proof.

P

205

X

i,j

207

208

3.3

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as

B

U

W

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pi

∗

Thus, any ConvSSM can be described as

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as

′

U

can be equivalently described as evolving according flattening the input and state tensors into vectors [73]. Thus, any ConvSSM can

A

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

Thus, any ConvSSM can be described as can be equivalently described as evolving according to a differential

state kernels as described in the previous section provides an alternative SSM equivalence, which a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

×

W

B

SSM

, and input matrix,

)

t

(

flattening the input and state tensors into vectors [73].

a large SSM with a circulant dynamics matrix. However, we show here that the

P

B

SSM

, and input matrix, a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

We show in this section that an equivalence can be established between the dynamics of a Conv

B

SSM

. See Figure 1.

Thus, any ConvSSM can be described as

U

col im

2

i,j

(

t

)

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

×

×

1

∗

state kernels as described in the previous section provides an alternative SSM equivalence, which

We show in this section that an equivalence can be established between the dynamics of a Conv a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

SSM

B

. See Figure 1.

state kernels as described in the previous section provides an alternative SSM equivalence, which state kernels as described in the previous section provides an alternative SSM eq

U

lends a special structure that can leverage the deep SSM parameterization and ini using pointwise state kernels and the dynamics of a large block-diagonal SSM system. Thus, we

lends a special structure that can leverage the deep SSM parameterization and initialization ideas using pointwise state kernels and the dynamics of a large block-diagonal SSM system. Thus, w

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

B

can be equivalently described as evolving according to a differential equation can be equivalently described as evolving according to a differential equation

[JS: im2col

×

5

B

SSM

B

. See Figure 1.

P

k

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

use this fact to leverage the parameterization and initialization ideas from deep SSMs discusse discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

use this fact to leverage the parameterization and initialization ideas from deep SSMs discuss

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as discussed in Section 2.2.3 for modeling long-range dependencies. We show that

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

[JS: im2col

W

P

A

SSM

, and input matrix,

B

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

can be equivalently described as evolving according to a differential equation can be equivalently described as evolving according to a diff

A

×

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependencies.

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependencies.

Thus, any ConvSSM can be described as

Thus, any ConvSSM can be described as

SSM

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

SSM

with a shared state matrix,

, and input matrix, state kernels as described in the previous section provides an alternative SSM equivalence, which

. See Figure 1.

SSM

B

state kernels as described in the previous section provides an alternative SSM equivalence, which

5

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

[JS: im2col a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

HWP

HWP

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

×

Consider a ConvSSM state update as in (7) with pointwise state kernel is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a lends a special structure that can leverage the deep SSM parameterization and initialization ideas

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

. See Figure 1.

[JS: im2col is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

H

W

U

′

′

∈

A

can be equivalently described as evolving according to a differential equation

×

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

∈

∈

A

×

.

Rearrange the

B

R

Rearrange the i,j

SSM

im

2

+

B

col can be equivalently described as evolving according to a differential equation

B

. See Figure 1.

Consider a ConvSSM state update as in (7) with pointwise state kernel

, and input matrix,

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to descri cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

. Rearrange the input kernel

SSM

, and input into a matrix

(

t

)

(

t

)

′

and

(

t

)

U

′

′

U

. The differential equation of (7) is equivalent, after reshaping

X

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official o into a matrix

U

U

X

H

W

P

×

′

and

(

t

)

X

P

(

t

)

)

i,j

+

B

U

∈

A

i,j

Rearrange the

Consider a ConvSSM state update as in (7) with pointwise state kernel operations, and has a paper we can cite. I am open to suggestions, but was having a hard time

X

U

∈

im

×

×

×

R

.

2

col

(

t

)

i,j

, and input

SSM

U

. Rearrange the input kernel into a matrix

(12)

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

. The differential equation of (7) is equivalent, after reshaping

U

(

t

)

(

t

)

State evolution of

ConvSSM

with pointwise state kernel

State pixels evolve according to

SSM

state update

Section 2.2.3 to allow ConvSSMs to effectively model long-range dependencies.

210

X

′

(

t

) =

A

∗

X

(

t

) +

B

∗

U

(

t

)

211

U

(

t

)

∈

R

H

′

×

W

′

×

U

212

X

′

(

t

)

∈

R

H

×

W

×

P

213

A

∈

R

P

×

P

×

1

×

1

214

B

∈

R

P

×

U

×

k

B

×

k

B

215

X

(

t

)

∈

R

H

×

W

×

P

216

A

Block

∈

R

HWP

×

HWP

217

X

′

(

t

)

∈

R

H

×

W

×

P

213

A

∈

R

P

×

P

×

1

×

1

214

B

∈

R

P

×

U

×

k

B

×

k

B

215

X

(

t

)

∈

R

H

×

W

×

P

216

A

Block

∈

R

HWP

×

HWP

217

5

B

∈

R

P

×

U

×

k

B

×

k

B

215

X

(

t

)

∈

R

H

×

W

×

P

216

A

Block

∈

R

HWP

×

HWP

217

A

∈

R

P

×

P

×

1

×

1

214

B

∈

R

P

×

U

×

k

B

×

k

B

215

X

(

t

)

∈

R

H

×

W

×

P

216

A

Block

∈

R

HWP

×

HWP

217

X

′

(

t

) =

A

∗

X

(

t

) +

B

∗

U

(

t

)

U

(

t

)

∈

R

H

′

×

W

′

×

U

X

′

(

t

)

∈

R

H

×

W

×

P

A

∈

R

P

×

P

×

1

×

1

R

P

×

U

×

k

B

×

k

B

X

′

(

t

) =

A

∗

X

(

t

) +

B

∗

U

(

t

)

211

U

(

t

)

∈

R

H

′

×

W

′

×

U

212

X

′

(

t

)

∈

R

H

×

W

×

P

213

A

∈

R

P

×

P

×

1

×

1

214

B

∈

R

P

×

U

×

k

B

×

k

B

215

X

(

t

)

∈

R

H

×

W

×

P

216

210

X

′

(

t

) =

A

∗

X

(

t

) +

B

∗

U

(

t

)

211

U

(

t

)

∈

R

H

′

×

W

′

×

U

212

X

′

(

t

)

∈

R

H

×

W

×

P

213

A

∈

R

P

×

P

×

1

×

1

214

B

∈

R

P

×

U

×

k

B

×

k

B

215

X

(

t

)

∈

R

H

×

W

×

P

216

A

R

HWP

×

HWP

217

X

′

(

t

) =

A

∗

X

(

t

) +

211

U

(

t

)

∈

R

H

′

×

W

′

×

U

212

X

′

(

t

)

∈

R

H

×

W

×

P

213

A

∈

R

P

×

P

×

1

×

1

214

R

P

×

U

×

k

B

×

k

B

215

X

′

(

t

) =

A

∗

211

U

(

t

)

∈

R

H

′

×

212

X

′

(

t

)

∈

R

H

×

213

A

∈

R

P

×

P

×

1

214

B

∈

R

P

×

U

×

k

215

X

(

t

)

∈

R

H

×

216

A

Block

∈

R

217

3.3

Connection to State Space Models

206

207

208

209

210

X

′

(

t

) =

A

∗

X

(

t

) +

B

∗

U

(

t

)

211

U

(

t

)

∈

R

H

′

×

W

′

×

U

212

H

W

P

3.3

Connection to State Space Models

206

207

208

209

210

X

′

(

t

) =

A

∗

X

(

t

) +

B

∗

U

(

t

)

211

U

(

t

)

∈

R

H

′

×

W

′

×

U

212

′

(

t

)

H

×

W

×

P

213

B

SSM

P

×

(

Uk

2

B

)

209

X

′

(

t

)

i,j

=

A

SSM

X

(

t

)

i,j

+

B

SSM

U

im

2

col

(

t

)

i,j

210

Uk

2

B

211

H

′

×

W

′

×

U

212

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

[JS: im2col

216

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

217

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

218

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

219

Proposition 3.

Consider a ConvSSM state update as in (7) with pointwise state kernel

A

∈

220

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

, and input

U

(

t

)

∈

R

H

′

×

W

′

×

U

.

Rearrange the

221

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

. Rearrange the input kernel

B

into a matrix

222

B

SSM

∈

R

P

×

(

Uk

2

B

)

. The differential equation of (7) is equivalent, after reshaping

X

(

t

)

and

X

′

(

t

)

223

into

x

(

t

)

,

x

′

(

t

)

∈

R

HWP

respectively, to the following differential equation

224

x

′

(

t

) =

A

Block

x

(

t

) +

B

Block

u

(

t

)

,

(12)

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

X

′

(

t

)

i,j

=

A

SSM

X

(

t

)

i,j

+

B

SSM

U

im

2

col

(

t

)

i,j

210

Uk

2

B

211

H

′

×

W

′

×

U

212

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

216

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

217

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

218

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

219

Proposition 3.

Consider a ConvSSM state update as in (7) with pointwise state kernel

220

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

, and input

U

(

t

)

∈

R

H

′

×

W

′

×

U

.

221

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

. Rearrange the input kernel

B

222

B

SSM

∈

R

P

×

(

Uk

2

B

)

. The differential equation of (7) is equivalent, after reshaping

X

(

t

)

223

into

x

(

t

)

,

x

′

(

t

)

∈

R

HWP

respectively, to the following differential equation

224

B

SSM

P

×

(

Uk

2

B

)

209

X

′

(

t

)

i,j

=

A

SSM

X

(

t

)

i,j

+

B

SSM

U

im

2

col

(

t

)

i,j

210

Uk

2

B

211

H

′

×

W

′

×

U

212

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

216

217

218

219

Proposition 3.

220

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

, and input

U

(

t

)

221

pointwise state kernel

A

into a matrix

A

SSM

∈

R

P

×

P

222

B

SSM

∈

R

P

×

(

Uk

2

B

)

223

into

x

(

t

)

,

x

′

(

t

)

∈

R

HWP

respectively, to the following differential equation

224

x

′

(

t

) =

A

Block

x

(

t

) +

B

Block

u

(

t

)

,

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

X

′

(

t

)

i,j

=

A

SSM

X

(

t

)

i,j

+

B

SSM

U

im

2

col

(

t

)

i,j

210

Uk

2

B

211

H

′

×

W

′

×

U

212

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

216

217

218

219

Proposition 3.

220

R

P

×

P

×

1

×

1

, input kernel

B

∈

R

P

×

U

×

k

B

×

k

B

221

pointwise state kernel

A

into a matrix

A

SSM

∈

R

222

B

SSM

R

P

×

(

Uk

2

B

)

223

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

X

′

(

t

)

i,j

=

A

SSM

X

(

t

)

i,j

+

210

Uk

2

B

211

H

′

×

W

′

×

U

212

P

×

U

×

k

B

×

k

B

213

P

×

P

×

1

×

1

214

H

×

W

×

P

215

216

217

218

219

Proposition 3.

220

R

P

×

P

×

1

×

1

, input kernel

B

221

197

198

flattening the input and state tensors into vectors [73].

199

200

201

202

203

state,

X

(

t

)

i,j

∈

R

P

204

with a shared state matrix,

A

SSM

, and input matrix,

B

SSM

. See Figure 1.

205

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

(

Uk

2

B

)

209

202

203

state,

X

(

t

)

i,j

∈

R

P

204

with a shared state matrix,

A

SSM

205

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

′

(

t

)

i,j

=

A

SSM

(

t

)

i,j

+

B

SSM

im

2

col

210

203

state,

X

(

t

)

i,j

∈

R

P

204

with a shared state matrix,

205

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

X

′

(

t

)

i,j

=

A

SSM

X

(

t

)

210

Uk

2

B

211

3.3

Connection to State Space Models

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

flattening the input and state tensors into vectors [73].

Thus, any ConvSSM can be described as

a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

state kernels as described in the previous section provides an alternative SSM equivalence, which

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

state,

X

(

t

)

i,j

∈

R

P

can be equivalently described as evolving according to a differential equation

with a shared state matrix,

A

SSM

, and input matrix,

B

SSM

. See Figure 1.

X

′

(

t

)

i,j

P

X

(

t

)

i,j

P

3.3

Connection to State Space Models

196

Since the convolutions in 7-8 and 9-10 are linear operations, they can be equivalently described as

197

matrix-vector multiplications using large, circulant matrices consisting of the kernel elements and

198

flattening the input and state tensors into vectors [73].

Thus, any ConvSSM can be described as

199

a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise

200

state kernels as described in the previous section provides an alternative SSM equivalence, which

201

lends a special structure that can leverage the deep SSM parameterization and initialization ideas

202

discussed in Section 2.2.3 for modeling long-range dependencies. We show that each pixel of the

203

state,

X

(

t

)

i,j

∈

R

P

can be equivalently described as evolving according to a differential equation

204

with a shared state matrix,

A

SSM

, and input matrix,

B

SSM

. See Figure 1.

205

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

3.3

Connection to State Space Models

196

197

198

flattening the input and state tensors into vectors [73].

199

200

201

202

203

state,

X

(

t

)

i,j

∈

R

P

204

with a shared state matrix,

A

SSM

, and input matrix,

B

SSM

. See Figure 1.

205

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

3.3

Connection to State Space Models

196

197

198

flattening the input and state tensors into vectors [73].

199

200

201

202

203

state,

X

(

t

)

i,j

∈

R

P

204

with a shared state matrix,

A

SSM

205

P

×

Uk

2

B

206

207

3.3

196

197

198

199

200

201

202

203

state,

X

(

t

)

i,j

∈

R

P

204

with a shared state matrix,

205

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

X

′

(

t

)

i,j

=

A

SSM

X

(

t

210

Uk

2

B

211

198

flattening the input and state tensors into vectors [73].

199

200

201

202

203

state,

X

(

t

)

i,j

∈

R

P

204

with a shared state matrix,

A

SSM

, and input matrix,

205

X

′

(

t

)

i,j

P

206

X

(

t

)

i,j

P

207

A

SSM

P

×

P

208

B

SSM

P

×

(

Uk

2

B

)

209

′

(

t

)

i,j

=

A

SSM

(

t

)

i,j

+

B

SSM

im

2

col

(

t

)

i,j

210

197

198

199

200

201

202

203

state,

X

(

t

)

i,j

∈

R

204

205

U

im2col

(

t

)

i,j

206

207

208

209

210

208

213

225

A

216

217

W

′

U

B

B

P

k

k

H

′

B

X

into a matrix

5

Rearrange the

[JS: im2col

×

.

(

t

)

×

×

A

×

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

pointwise state kernel cleaner way of describing what we need here. I feel like saying im2col and citing it is actual

∈

222

HWP

X

5

208

im i,j

+

B

SSM

2

col

B

224

R

209

223

i,j

R

(

t

)

respectively, to the following differential equation

2

into

x

(

t

)

,

x

′

(

t

)

R

2

212

×

b

HWP

HWP

Block

Block

HWUk i,j

t

)

A

SSM

=

A

HWP

∈

∈

B

215

SSM

P

P

x

(

t

) +

B

∈

×

. Rearrange the input kernel

HWUk

u

(

t

)

∈

SSM

X

)

∈

P

P

(

t

′

and operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

SSM

,

HWP

H

P

W

P

×

(

′

1

Uk

′

P

. The differential equation of (7) is equivalent, after reshaping

)

B

U

Block

∈

(12)

A

2

b

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

P

U

k

k

234

(

t

)

Block i,j

SSM

(

HWP

P

210

(

t

)

∈

′

×

×

×

HWP

HWP

x

(

t

) =

A

′

HWP

X

R

′

X

∈

X

232

∈

B

X

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

P

234

235

231

216

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

233

advantageous dynamical properties.

216

229

[JS: im2col

R

×

×

×

1

221

, input kernel

230

232

advantageous dynamical properties.

Note the ConvSSM formulation is easier to work with than cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

219

.

)

columns formed by performing an im2col [74, 75] operation on the input

218

217

B

B

A

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into

∈

B

208

216

217

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

A

∈

into a matrix

(

B

t

)

and

U

Note the ConvSSM formulation is easier to work with than

222

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

advantageous dynamical properties.

disretizing with a learnable timescale parameter to obtain

R

217

218

219

222

A

∈

H

Block

R

pointwise state kernel

A

Block

R

×

A

into a matrix

×

[JS: im2col

SSM

′

(

t

)

W

U

×

and

×

R

SSM

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

234

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

216

, and input

A

SSM

. Rearrange the input kernel

(

B

t

)

SSM

would endow these

[JS: im2col

2

Rearrange the

.

HWUk

U

(

t

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

×

Consider a ConvSSM state update as in (7) with pointwise state kernel

235

W

236

into a matrix

P

P

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

R

×

×

is applicable to any HxWxC tensor (which is the only case we consider). It is also a

P

233

217

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

220

215

221

∈

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

231

on the diagonal blocks respectively, and

P

B

SSM

×

[JS: im2col

A

Note the ConvSSM formulation is easier to work with than operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

B

b

∈

B

u

(

t

)

R

(

Uk

2

B

)

Note the ConvSSM formulation is easier to work with than respectively, to the following differential equation

is the result of c

223

H

Proposition 3.

P

(

X

∈

U

×

219

Proof.

See Appendix C.3.

B

SSM

R

224

X

′

t

)

and

(

t

)

. The differential equation of (7) is equivalent, after reshaping

HWP

H

′

W

′

convolution operations in modern deep learning frameworks [76-78].

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same operations, and has a paper we can cite. I am open to suggestions, but was having

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into columns formed by performing an im2col [74, 75] operation on the input

218

219

, input kernel

P

U

k

convolution operations in modern deep learning frameworks [76-78].

k

B

∈

P

1

1

[WB: Define im2col before eq 12. im2cal is specific to images. Is there another way to describe it?]

217

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

×

224

2

236

217

238

226

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation in GEMM

cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.]

operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a

)

Note the ConvSSM formulation is easier to work with than

U

(

t

)

(

t

.

)

Rearrange the

R

.

234

U

, and input advantageous dynamical properties.

(

t

)

∈

R

into

232

×

×

∈

x

(

t

)

,

x

′

218

cleaner way of describing what we need here. I feel like saying im2col and citing

3.4

×

5

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

SSM

,

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

236

B

P

×

∈

P

R

5

×

×

222

In this section, we introduce which combines ideas of parallelization of convolutional

pointwise state kernel

5 5 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 235 convolution operations in modern deep learning frameworks [76-78]. 236 233 advantageous dynamical properties. 234 235 operations, and has a paper we can cite. I am open to suggestions, but was having a hard time coming up with a 218 cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.] 219 216 217 218 219 Proposition 3. Consider a ConvSSM state update as in (7) with pointwise state kernel A ∈ 220 R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H ′ × W ′ × U . Rearrange the 221 pointwise state kernel A into a matrix A SSM ∈ R P × P . Rearrange the input kernel B into a matrix 222 B SSM ∈ R P × ( Uk 2 B ) . The differential equation of (7) is equivalent, after reshaping X ( t ) and X ′ ( t ) 223 into x ( t ) , x ′ ( t ) ∈ R HWP respectively, to the following differential equation Proposition 3. Consider a ConvSSM state update as in (7) with pointwise state kernel A ∈ 220 R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H ′ × W ′ × U . Rearrange the 221 Thus, parameterizing and initializing A SSM with the specially structured HiPPO matrices [47] and 220 disretizing with a learnable timescale parameter to obtain A SSM and B SSM 221 ∈ x ′ ( t ) = A Block x ( t ) + B Block u ( t 218 cleaner way of describing what we need here. I feel like saying im2col and citing it is actually pretty precise.] 219 Proof. See Appendix C.3. 219 3.2 Parallelizing Convolutional Recurrences

225

229

would endow these convolution operations in modern deep learning frameworks [76-78].

which combines ideas of parallelization of convolutional

R

×

225

into a matrix

Efficient ConvSSM for Long-Range Dependencies: ConvS5

∈

B

Consider a ConvSSM state update as in (7) with pointwise state kernel into a matrix

237

A

. Rearrange the input kernel

[

×

×

is applicable to any HxWxC tensor (which is the only case we consider). It is also an official operation

2

Proposition 3.

SSM

R

220

HWP

R

235

237

pointwise state kernel

A

In this section, we introduce

HWP

227

into a matrix

HWP

2

230

HWUk

b

HWP

HWP

231

disretizing with a learnable timescale parameter to obtain

235

239

P

P

HWP

where

. Rearrange the input kernel

A

HWUk

b

leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition forming the large, block-diagonal SSM state update explicitly and allows leveraging

SSM

and

233

236

B

SSM

b

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized where A Block ∈ R HWP × HWP and B Block ∈ R HWP × HWUk 2 b are block-diagonal with A SSM and 225 B SSM on the diagonal blocks respectively, and u ( t ) ∈ R HWUk 2 is the result of concatenating the columns formed by performing an im2col [74, 75] operation on the input U ( t ) . 227 B SSM ∈ R P × ( Uk 2 B ) . The differential equation of (7) is equivalent, after reshaping 223 into x ( t ) , x ′ ( t ) ∈ R HWP respectively, to the following differential equation 224 x ′ ( t ) = A x ( t ) + B u ( t ) , x ′ ( t ) = A Block x ( t ) + B Block u ( t ) where A Block ∈ R HWP × HWP and B Block R HWP × HWUk b 225 B on the diagonal blocks respectively, and u ( t ) R HWUk 2 Thus, parameterizing and initializing A SSM with the specially structured HiPPO matrices [47] and disretizing with a learnable timescale parameter to obtain A SSM and B SSM would endow these SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range 231 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 235 convolution operations in modern deep learning frameworks [76-78]. 236 Thus, parameterizing and initializing A SSM with the specially structured HiPPO matrices [47] and 229 disretizing with a learnable timescale parameter to obtain A SSM and B SSM would endow these 230 SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range 231 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 3, we can implicitly parameterize a pointwise state kernel, A ∈ R P × P × 1 × 1 and input kernel 232 B ∈ R P × U × k B × k B in (7) using SSM parameters as used by S5 [36], A S5 ∈ R P × P and B S5 ∈ 233 R P × ( Uk 2 B ) . We can discretize these S5 SSM parameters as discussed in Section 2.2.1 to give 234 A S5 = DISCRETIZE A ( A S5 , ∆ ) , B S5 = DISCRETIZE B ( A S5 , B S5 , ∆ ) , (13) and then rearrange to give the ConvS5 state update kernels: 235 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 235 convolution operations in modern deep learning frameworks [76-78]. 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 where A Block ∈ R HWP × HWP and B Block ∈ R HWP × HWUk b are block-diagonal with A SSM and 225 B SSM on the diagonal blocks respectively, and u ( t ) ∈ R HWUk 2 b is the result of concatenating the 226 columns formed by performing an im2col [74, 75] operation on the input U ( t ) . 227 Proof. See Appendix C.3. convolution operations in modern deep learning frameworks [76-78]. 227 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 228 q i ⊛ q j := ( q j,a ◦ q i,a , q j,a ∗ q i,b + q j,b ) , (9) where ◦ denotes convolution of two kernels, ∗ denotes convolution and + is elementwise addition.

(12)

and

SSM

,

the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same

(

t

)

∈

R

HWP

HWP

B

R

×

SSM

Block would endow these

R

×

P

forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized

226

228

226

convolution operations in modern deep learning frameworks [76-78].

228

230

231

2

X

are block-

b

is the r

ConvS5 is a ConvSSM that which combines ideas of parallelization of convolutional

Block

In this section, we introduce

In this section, we introduce

B

- - - - - - →

A

P

2

B

into a matrix and

B

HWP

HWUk are block-diagonal with

2

Block

R

×

2

×

P

×

1

×

1

A

P

convolution operations in modern deep learning frameworks [76-78].

Efficient ConvSSM for Long-Range Dependencies: ConvS5

238

239

240

rearrange rearrange

∈

∈

leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition columns formed by performing an im2col [74, 75] operation on the input

O

(

B

recurrences (Section 3.2) and the SSM connection (Section 3.3).

Block

×

k

P

U

∈

B

k

which combines ideas of parallelization

B

.

(14)

(15)

- - - - - - →

S5

S5

recurrences (Section 3.2) and the SSM connection (Section 3.3).

B

P

(

×

∈

R

×

×

∈

R

SSM

Uk

)

k

recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 238 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 240 6 Proposition 3. Consider a ConvSSM state update as in (7) with pointwise state kernel A ∈ 220 R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H ′ × W ′ × U . Rearrange the 221 pointwise state kernel A into a matrix A SSM ∈ R P × P . Rearrange the input kernel B into a matrix 222 B SSM ∈ R P × ( Uk B ) . The differential equation of (7) is equivalent, after reshaping X ( t ) and X ′ ( t ) 223 into x ( t ) , x ′ ( t ) ∈ R HWP respectively, to the following differential equation 224 x ′ ( t ) = A Block x ( t ) + B Block u ( t ) , (12) operations, and has a paper we can cite. I am open to suggestions, but was having a hard time comin 218 cleaner way of describing what we need here. I feel like saying im2col and citing it is actually prett 219 Proposition 3. Consider a ConvSSM state update as in (7) with pointwise state ker 220 R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H ′ × W ′ × U . Rear 221 pointwise state kernel A into a matrix A ∈ R P × P . Rearrange the input kernel B into 222 Proposition 3. Consider a ConvSSM state update as in (7) with poi 220 R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H 221 pointwise state kernel A into a matrix A SSM ∈ R P × P . Rearrange the in 222 B SSM ∈ R P × ( Uk 2 B ) . The differential equation of (7) is equivalent, after r 223 into x ( t ) , x ′ ( t ) ∈ R HWP respectively, to the following differential equati 224 B SSM ∈ R P × ( Uk B ) . The differential equation of (7) is equivalent, after reshaping X ( t ) and X ′ ( t ) 223 into x ( t ) , x ′ ( t ) ∈ R HWP respectively, to the following differential equation 224 x ′ ( t ) = A Block x ( t ) + B Block u ( t ) , (12) where A Block ∈ R × and B Block ∈ R × 2 are block-diagonal with A SSM and B SSM on the diagonal blocks respectively, and u ( t ) ∈ R HWUk 2 b is the result of concatenating the 226 columns formed by performing an im2col [74, 75] operation on the input U ( t ) . 227 Proof. See Appendix C.3. 228 x ′ ( t ) = A Block x ( t ) + B Block u ( t ) , (12) where A Block ∈ R HWP × HWP and B Block ∈ R HWP × HWUk 2 b are block-diagonal with A SSM and 225 B SSM on the diagonal blocks respectively, and u ( t ) ∈ R HWUk 2 b is the result of concatenating the 226 columns formed by performing an im2col [74, 75] operation on the input U ( t ) . Proof. See Appendix C.3. 228 Thus, parameterizing and initializing A SSM with the specially structured HiPPO matrices [47] and 229 B SSM ∈ R P × ( Uk B ) . The differential equation of (7) is equivalent, after reshaping X ( t ) and X ′ ( t ) 223 into x ( t ) , x ′ ( t ) ∈ R HWP respectively, to the following differential equation 224 x ′ ( t ) = A Block x ( t ) + B Block u ( t ) , (12) where A Block ∈ R × and B Block ∈ R × 2 are block-diagonal with A SSM and 225 B SSM on the diagonal blocks respectively, and u ( t ) ∈ R HWUk 2 b is the result of concatenating the 226 columns formed by performing an im2col [74, 75] operation on the input U ( t ) . 227 Proof. See Appendix C.3. 228 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 223 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 224 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 225 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 226 convolution operations in modern deep learning frameworks [76-78]. 227 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 228 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 229 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 230 ∈ ∈ B SSM on the diagonal blocks respectively, and u ( t ) ∈ R HWUk 2 b is the result of concatenating the 226 columns formed by performing an im2col [74, 75] operation on the input U ( t ) . 227 Proof. See Appendix C.3. 228 Thus, parameterizing and initializing A SSM with the specially structured HiPPO matrices [47] and 229 disretizing with a learnable timescale parameter to obtain A SSM and B SSM would endow these 230 SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range 231 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H ′ × W ′ × U . Rearrange the 221 pointwise state kernel A into a matrix A SSM ∈ R P × P . Rearrange the input kernel B into a matrix 222 B SSM ∈ R P × ( Uk 2 B ) . The differential equation of (7) is equivalent, after reshaping X ( t ) and X ′ ( t ) 223 into x ( t ) , x ′ ( t ) ∈ R HWP respectively, to the following differential equation 224 x ′ ( t ) = A Block x ( t ) + B Block u ( t ) , (12) Thus, parameterizing and initializing A SSM with the specially structured HiPPO 220 disretizing with a learnable timescale parameter to obtain A SSM and B SSM wo 221 SSMs with the same dynamics as in S4/S5 methods that are important for mod 222 dependencies. Due to the equivalence of Proposition 3, we can then rearrange th 223 the discrete ConvSSM state and input kernels of 9 to give the convolutional recu 224 advantageous dynamical properties. Note the ConvSSM formulation is easier to 225 ConvS5 leverages parallel scans to efficiently parallelize the recurrence in (7). As discussed in Section 2.3, this requires a binary associative operator. Given that convolutions are associative, we show: Proposition 1. Consider a convolutional recurrence as in (7) and define initial parallel scan elements c k = ( c k,a , c k,b ) := ( A , B ∗ U k ) . The binary operator ⊛ , defined below, is associative.

230

231

235

236

238

239

237

See Appendix C.3.

240

238

239

3.4

236

240

SSM

SSM

A

229

230

LPHW

)

P

U

B

k

B

227

228

leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition

∈

∈

. See Appendix A for a more detailed discussion of parameterization, initialization and

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

ConvS5 is a ConvSSM that recurrences (Section 3.2) and the SSM connection (Section 3.3).

recurrences (Section 3.2) and the SSM connection (Section 3.3).

238

239

B

240

P

(

Uk

SSMs with the same dynamics as in S4/S5 methods that are important for modeling long-range

×

2

B

)

discretization.

240 6 Proposition 2. Given the effective inputs B ∗ U 1: L ∈ R L × H × W × P and a pointwise state kernel A ∈ R P × P × 1 × 1 , the computational cost of computing the convolutional recurrence in Equation 7 with a parallel scan is O ( L ( P 3 + P 2 HW ) ) .

HWUk

b

P

is the result of concate in (7) using SSM parameters as used by S5 [36],

A

S5

∈

R

6

226

dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 235 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 columns formed by performing an im2col [74, 75] operation on the input U ( t ) . 227 Proof. See Appendix C.3. 228 Thus, parameterizing and initializing A SSM with the specially structured HiPPO matrices 229 disretizing with a learnable timescale parameter to obtain A SSM and B SSM would end 230 SSMs with the same dynamics as in S4/S5 methods that are important for modeling lo 231 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these mat 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work 234 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly 235 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of parallelization of con 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Pr Thus, parameterizing and initializing A SSM with the specially structured 229 disretizing with a learnable timescale parameter to obtain A SSM and 230 SSMs with the same dynamics as in S4/S5 methods that are important 231 dependencies. Due to the equivalence of Proposition 3, we can then rear 232 the discrete ConvSSM state and input kernels of 9 to give the convoluti 233 advantageous dynamical properties. Note the ConvSSM formulation is 234 forming the large, block-diagonal SSM state update explicitly and allows le 235 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of paralle 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). Co 239 leverages parallel scans and deep SSM parameterization/initialization sc 240 6 6 6 6 dependencies. Due to the equivalence of Proposition 3, we can then rearrange these matrices into 232 the discrete ConvSSM state and input kernels of 9 to give the convolutional recurrence the same 233 advantageous dynamical properties. Note the ConvSSM formulation is easier to work with than 234 forming the large, block-diagonal SSM state update explicitly and allows leveraging highly optimized 235 convolution operations in modern deep learning frameworks [76-78]. 236 3.4 Efficient ConvSSM for Long-Range Dependencies: ConvS5 237 In this section, we introduce ConvS5 which combines ideas of parallelization of convolutional 238 recurrences (Section 3.2) and the SSM connection (Section 3.3). ConvS5 is a ConvSSM that 239 leverages parallel scans and deep SSM parameterization/initialization schemes. Given Proposition 240 6 R × . We can discretize these S5 SSM parameters as discussed in Section 2.2 234 A S5 = DISCRETIZE A ( A S5 , ∆ ) , B S5 = DISCRETIZE B ( A S5 , B and then rearrange to give the ConvS5 state update kernels: 235 A S5 ∈ R P × P rearrange - - - - - - → A ConvS5 ∈ R P × P × 1 × 1 B S5 ∈ R P × ( Uk 2 B ) rearrange - - - - - - → B ConvS5 ∈ R P × U × k B × k B . We then run the discretized ConvSSM system of Eqs (9, 10), where the recurrence 236 using parallel scans. Note that in practice this setup actually allows us to parameteri 237 a diagonal state [34-36] which reduces the cost of applying the parallel scan in 238 O ( LPHW ) . See Appendix A for a more detailed discussion of parameterization, i 239 discretization. 240 6 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 A × k A resolution state kernels together. To maintain equivalence with the sequential scan, the resulting kernel will have resolution 2 k a -1 × 2 k a -1 . This implies that the state kernel will grow during the parallel scan computations for general kernels with a resolution greater than 1 × 1 . This allows the receptive field to grow in the time direction, a useful feature for capturing spatiotemporal context. However, this kernel growth is computationally infeasible for long sequences. We address this challenge by taking further inspiration from deep SSMs. These methods opt for simple but computationally advantageous operations in the time direction (linear dynamics) and utilize more complex operations (nonlinear activations) in the depth direction of the model. These nonlinear activations allow a stack of SSM layers with linear dynamics to represent nonlinear systems. Analogously, we choose to use 1 × 1 state kernels and perform pointwise state convolutions for the convolutional recurrence of (7). When we stack multiple layers of these ConvSSMs, the receptive field grows in the depth direction of the network and allows the stack of layers to capture the spatiotemporal context [64]. Computationally, we now have a construction that can be parallelized with subquadratic complexity with respect to the sequence length.

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.

SSM

240

233

ConvS5 is a ConvSSM that

×

R

)

t

u

(

×

R

b

∈

A

A

A

222

A

238

∈

201

X

Proof.

231

3.3 Connection to State Space Models

Since the convolutions in (5-6) and (7-8) are linear operations, they can be described equivalently as matrix-vector multiplications by flattening the input and state tensors into vectors and using large, circulant matrices consisting of the kernel elements [65]. Thus, any ConvSSM can be described as a large SSM with a circulant dynamics matrix. However, we show here that the use of pointwise state kernels, as described in the previous section, provides an alternative SSM equivalence, which lends a special structure that can leverage the deep SSM parameterization/initialization ideas discussed in Section 2.2 for modeling long-range dependencies. We show that each pixel of the state, X ( t ) i,j ∈ R P , can be equivalently described as evolving according to a differential equation with a shared state matrix, A SSM , and input matrix, B SSM . See Figure 2.

Proposition 3. Consider a ConvSSM state update as in (5) with pointwise state kernel A ∈ R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H ′ × W ′ × U . Let U im2col ( t ) ∈ R H × W × Uk 2 B be the reshaped result of applying the Image to Column (im2col) [66, 67] operation on the input U ( t ) . Then the dynamics of each state pixel of ( 5 ), X ( t ) i,j ∈ R P , evolve according to the following differential equation

where the state matrix, A SSM ∈ R P × P , and input matrix, B SSM ∈ R P × ( Uk 2 B ) , can be formed by reshaping the state kernel, A , and input kernel, B , respectively.

Thus, to endow these SSMs with the same favorable long-range dynamical properties as S4/S5 methods, we initialize A SSM with a HiPPO [62] inspired matrix and discretize with a learnable timescale parameter to obtain A SSM and B SSM . Due to the equivalence of Proposition 3, we then reshape these matrices into the discrete ConvSSM state and input kernels of (7) to give the convolutional recurrence the same advantageous dynamical properties. We note that if the input, output and dynamics kernel widths are set to 1 × 1 , then the ConvSSM formulation is equivalent to "convolving" an SSM across each individual sequence of pixels in the spatiotemporal sequence (this also has connections to the temporal component of S4ND [45] when applied to videos). However, inspired by ConvRNNs, we observed improved performance when leveraging the more general convolutional structure the ConvSSM allows and increasing the input/output kernel sizes to allow local spatial features to be mixed in the dynamical system. See ablations discussed in Section 5.3.

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, A ∈ R P × P × 1 × 1 and input kernel B ∈ R P × U × k B × k B in (5) using SSM parameters as used by S5 [20], A S5 ∈ R P × P and B S5 ∈ R P × ( Uk 2 B ) . We discretize these S5 SSM parameters as discussed in Section 2.2 to give

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

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.

4 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 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 strided 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.

5 Experiments

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

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-

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 × 16 vector-quantized (VQ) codes from the pretrained VQGANs [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.

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.

6 Discussion

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.

References

Chelsea Finn, Ian Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. Advances in neural information processing systems , 29, 2016.

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.

He Huang, Zheni Zeng, Danya Yao, Xin Pei, and Yi Zhang. Spatial-temporal ConvLSTM for vehicle driving intention prediction. Tsinghua Science and Technology , 27(3):599-609, 2021.

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.

Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, 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.11214 , 2022.

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):e2021MS002502, 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.

Aidan Clark, Jeff Donahue, and Karen Simonyan. Adversarial video generation on complex datasets. arXiv preprint arXiv:1907.06571 , 2019.

Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. Videogpt: Video generation using VQ-VAE and Transformers. arXiv preprint arXiv:2104.10157 , 2021.

Guillaume Le Moing, Jean Ponce, and Cordelia Schmid. CCVS: context-aware controllable video synthesis. Advances in Neural Information Processing Systems , 34:14042-14055, 2021.

Yu Tian, Jian Ren, Menglei Chai, Kyle Olszewski, Xi Peng, Dimitris N. Metaxas, and Sergey Tulyakov. A good image generator is what you need for high-resolution video synthesis. In International Conference on Learning Representations , 2021.

Pauline Luc, Aidan Clark, Sander Dieleman, Diego de Las Casas, Yotam Doron, Albin Cassirer, and Karen Simonyan. Transformation-based adversarial video prediction on large-scale data. arXiv preprint arXiv:2003.04035 , 2020.

Wilson Yan, Danijar Hafner, Stephen James, and Pieter Abbeel. Temporally consistent video Transformer for long-term video prediction. arXiv preprint arXiv:2210.02396 , 2022.

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. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XVII , pages 102-118. Springer, 2022.

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.

Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. Advances in Neural Information Processing Systems , 34:29246-29257, 2021.

Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional LSTM network: A machine learning approach for precipitation nowcasting. Advances in neural information processing systems , 28, 2015.

Nicolas Ballas, Li Yao, Chris Pal, and Aaron Courville. Delving deeper into convolutional networks for learning video representations. arXiv preprint arXiv:1511.06432 , 2015.

Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations , 2021.

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.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation , 9 (8):1735-1780, 1997.

Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoderdecoder for statistical machine translation. arXiv preprint arXiv:1406.1078 , 2014.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference on Machine Learning , pages 1310-1318. PMLR, 2013.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

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.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. ViViT: A video vision Transformer. In Proceedings of the IEEE/CVF international conference on computer vision , pages 6836-6846, 2021.

Agrim Gupta, Stephen Tian, Yunzhi Zhang, Jiajun Wu, Roberto Martín-Martín, and Li Fei-Fei. Maskvit: Masked visual pre-training for video prediction. arXiv preprint arXiv:2206.11894 , 2022.

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems , 30, 2017.

Jacob Walker, Ali Razavi, and Aäron van den Oord. Predicting video with VQVAE. arXiv preprint arXiv:2103.01950 , 2021.

Patrick Esser, Robin Rombach, and Bjorn 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.

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.05102 , 2022.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

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.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive Transformers with linear attention. In International Conference on Machine Learning , pages 5156-5165. PMLR, 2020.

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient Transformer. In International Conference on Learning Representations , 2020.

Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document Transformer. arXiv preprint arXiv:2004.05150 , 2020.

Ankit Gupta and Jonathan Berant. Gmat: Global memory augmentation for Transformers, 2020.

Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020.

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.

Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long Range Arena: A benchmark for efficient Transformers. In International Conference on Learning Representations , 2021.

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.

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.

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.

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 on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 7616-7633. PMLR, 17-23 Jul 2022.

Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4ND: Modeling images and videos as multidimensional signals with state spaces. In Advances in Neural Information Processing Systems , 2022.

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.

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.

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.

Linqi Zhou, Michael Poli, Winnie Xu, Stefano Massaroli, and Stefano Ermon. Deep latent state space models for time-series generation. arXiv preprint arXiv:2212.12749 , 2022.

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.

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.

Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint arXiv:2212.10544 , 2022.

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 arXiv:2302.10866 , 2023.

Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning of video representations using LSTMs. In International conference on machine learning , pages 843-852. PMLR, 2015.

Arieh Iserles. A first course in the numerical analysis of differential equations . 44. Cambridge university press, 2009.

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.

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 arXiv:2303.06349 , 2023.

Antonio Orvieto, Soham De, Caglar Gulcehre, Razvan Pascanu, and Samuel L Smith. On the universality of linear recurrences followed by nonlinear projections. arXiv preprint arXiv:2307.11888 , 2023.

Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations , 2018.

James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. In International Conference on Learning Representations , 2017.

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.

Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HiPPO: Recurrent memory with optimal polynomial projections. Advances in Neural Information Processing Systems , 33:1474-1487, 2020.

Guy Blelloch. Prefix sums and their applications. Technical report, Tech. rept. CMU-CS-90190. School of Computer Science, Carnegie Mellon, 1990.

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.

Vincent Dumoulin and Francesco Visin. A guide to convolution arithmetic for deep learning. arXiv preprint arXiv:1603.07285 , 2016.

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.

Yangqing Jia. Learning semantic image representations at a large scale. 2014.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Zhuang Liu, Hanzi Mao, Chao-Yuan 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.

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.

Reza Azad, Maryam Asadi-Aghbolaghi, 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.

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:113704, 2020.

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

Mohamadreza Bakhtyari and Sayeh Mirzaei. ADHD detection using dynamic connectivity patterns of EEG data and ConvLSTM with attention framework. Biomedical Signal Processing and Control , 76:103708, 2022.

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.

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.

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.

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 a new model. Advances in neural information processing systems , 30, 2017.

Mohammad Babaeizadeh, Chelsea Finn, Dumitru Erhan, Roy H. Campbell, and Sergey Levine. Stochastic variational video prediction. In International Conference on Learning Representations , 2018.

Yunbo Wang, Haixu Wu, Jianjin Zhang, Zhifeng Gao, Jianmin Wang, S Yu Philip, and Mingsheng Long. PredRNN: A recurrent neural network for spatiotemporal predictive learning. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45(2):2208-2225, 2022.

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.

Yunbo Wang, Lu Jiang, Ming-Hsuan Yang, Li-Jia Li, Mingsheng Long, and Li Fei-Fei. Eidetic 3D LSTM: A model for video prediction and beyond. In International conference on learning representations , 2019.

Wei Yu, Yichao Lu, Steve Easterbrook, and Sanja Fidler. Efficient and information-preserving future frame prediction and beyond. In International Conference on Learning Representations , 2020.

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 , 33:13714-13726, 2020.

Jue Wang, Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-form video understanding. arXiv preprint arXiv:2303.14526 , 2023.

Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning , pages 1747-1756. PMLR, 2016.

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, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

Özgün Çiçek, Ahmed Abdulkadir, Soeren S Lienkamp, Thomas Brox, and Olaf Ronneberger. 3D U-net: learning dense volumetric segmentation from sparse annotation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19 , pages 424-432. Springer, 2016.

Tobias Höppe, Arash Mehrjou, Stefan Bauer, Didrik Nielsen, and Andrea Dittadi. Diffusion models for video prediction and infilling. Transactions on Machine Learning Research .

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.13195 , 2021.

Ruslan Rakhimov, Denis Volkhonskiy, Alexey Artemov, Denis Zorin, and Evgeny Burnaev. Latent video Transformer. arXiv preprint arXiv:2006.10704 , 2020.

Younggyo Seo, Kimin Lee, Fangchen Liu, Stephen James, and Pieter Abbeel. HARP: Autoregressive latent video prediction with high-fidelity image generator. In 2022 IEEE International Conference on Image Processing (ICIP) , pages 3943-3947. IEEE, 2022.

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 on Computer Vision and Pattern Recognition , pages 3626-3636, 2022.

Masaki Saito and Shunta Saito. TGANv2: Efficient training of large models for video generation with multiple subsampling layers. arXiv preprint arXiv:1811.09245 , 2(6), 2018.

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.

Tim Brooks, Janne Hellsten, Miika Aittala, Ting-Chun Wang, Timo Aila, Jaakko Lehtinen, Ming-Yu Liu, Alexei Efros, and Tero Karras. Generating long videos of dynamic scenes. Advances in Neural Information Processing Systems , 35:31769-31781, 2022.

Wenyi Hong, Ming Ding, Wendi Zheng, Xinghan Liu, and Jie Tang. Cogvideo: Large-scale pretraining for text-to-video generation via Transformers. arXiv preprint arXiv:2205.15868 , 2022.

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.

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:1612.03801 , 2016.

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 arXiv:1907.13440 , 2019.

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.

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.

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.

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.

Mark Harris, Shubhabrata Sengupta, and John D Owens. Parallel prefix sum (scan) with CUDA. GPU gems , 3(39):851-876, 2007.

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.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

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.

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.

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.

Yann LeCun. The MNIST database of handwritten digits. http://yann. lecun. com/exdb/mnist/ , 1998.

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.

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.

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

Proposition 1. Consider a convolutional recurrence as in (7) and define initial parallel scan elements c k = ( c k,a , c k,b ) := ( A , B ∗ U k ) . The binary operator ⊛ , defined below, is associative.

where ◦ denotes convolution of two kernels, ∗ denotes convolution between a kernel and input, and + is elementwise addition.

Proof. Using that ◦ is associative and the companion operator of ∗ , i.e. ( d ◦ e ) ∗ f = d ∗ ( e ∗ f ) (see Blelloch [63], Section 1.4), we have:

A.2 Computational Cost of Parallel Scan for Convolutional Recurrences

Proposition 2. Given the effective inputs B ∗ U 1: L ∈ R L × H × W × P and a pointwise state kernel A ∈ R P × P × 1 × 1 , the computational cost of computing the convolutional recurrence in Equation 7 with a parallel scan is O ( L ( P 3 + P 2 HW ) ) .

Proof. Following Blelloch [63], given a single processor, the cost of computing the recurrence sequentially using the binary operator ⊛ defined in Proposition 1 is O ( L ( T ◦ + T ∗ + T + ) ) where T ◦ 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 + = O ( PHW ) . For state kernels with resolution k A , T ◦ = O ( P 3 k 4 A ) and T ∗ = O ( P 2 k 2 A HW ) . For pointwise convolutions this becomes T ◦ = O ( P 3 ) and T ∗ = O ( P 2 HW ) . Thus, the cost of computing the recurrence sequentially using ⊛ is O ( L ( P 3 + P 2 HW ) ) . Since there are work-efficient algorithms for parallel scans [105], the overall cost of the parallel scan is also O ( L ( P 3 + P 2 HW ) ) .

Note that ConvS5's diagonalized parameterization discussed in Section 3.4 and Appendix B leads to T ◦ = O ( P ) and T ∗ = O ( PHW ) . Therefore the cost of applying the parallel scan with ConvS5 is O ( LPHW ) .

A.3 Connection Between ConvSSMs and SSMs

Proposition 3. Consider a ConvSSM state update as in (5) with pointwise state kernel A ∈ R P × P × 1 × 1 , input kernel B ∈ R P × U × k B × k B , and input U ( t ) ∈ R H ′ × W ′ × U . Let U im2col ( t ) ∈ R H × W × Uk 2 B be the reshaped result of applying the Image to Column (im2col) [66, 67] operation on the input U ( t ) . Then the dynamics of each state pixel of ( 5 ), X ( t ) i,j ∈ R P , evolve according to the following differential equation

where the state matrix, A SSM ∈ R P × P , and input matrix, B SSM ∈ R P × ( Uk 2 B ) , can be formed by reshaping the state kernel, A , and input kernel, B , respectively.

Proof. Let U im2col ∈ R Uk 2 b × HW denote the result of performing the im2col operation on the input U ( t ) for convolution with the kernel B . Reshape this matrix into the tensor U im2col ( t ) ∈ R H × W × Uk 2 b . Reshape U im2col ( t ) once more into the tensor V ( t ) ∈ R H × W × U × k B × k B .

Now, we can write the evolution for the individual channels of each pixel, X ′ ( t ) i,j,k , in (5) as

Let A SSM ∈ R P × P be a matrix with rows, A SSM ,i ∈ R P , corresponding to a flattened version of the output features of A , i.e. A i ∈ R P × 1 × 1 . Similarly, reshape B into a matrix B SSM ∈ R P × ( Uk 2 B ) where the rows, B SSM ,i ∈ R Uk 2 B correspond to a flattened version of the output features of B , i.e. B i ∈ R U × k B × k B .

Then we can rewrite (22) equivalently as

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 A S5 = VΛ S5 V -1 ∈ R P × P where Λ S5 ∈ C P × P is a complex-valued diagonal matrix and V ∈ C P × P corresponds to the eigenvectors. Defining ˜ x ( t ) = V -1 x ( t ) , ˜ B = V -1 B , and ˜ C = CV we can reparameterize the SSM of (3) as the diagonalized system:

S5 uses learnable timescale parameters ∆ ∈ R P and the following zero-order hold (ZOH) disretization:

S5 Initialization S5 initializes its state matrix by diagonalizing A S5 as defined here:

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 A S5 as in (29) and diagonalize as A S5 = VΛ S5 V -1 . To apply ConvS5, we compute Λ S5 and B S5 using (27-28), and then form the ConvS5 state and input kernels:

See Listing 1 for an example of the core implementation. Note, the state kernel A ConvS5 is "diagonalized" in the sense that all entries in the state kernel are zero except A ConvS5 ,i,i = Λ S5 ,i,i ∀ i ∈ [ 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.

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

t=0

50

99

100

250

400

550

700

850

1000

1150

1299

Ground Truth

Transformer

ConvLSTM

Condition                                                   Predicted Frames

(a) Example 1

t=0

50

99

100

250

400

550

700

850

1000

1150

1299

Ground Truth

Transformer

ConvLSTM

Condition                                                   Predicted Frames

(b) Example 2

t=0

50

99

100

250

400

550

700

850

1000

1150

1299

Ground Truth

Transformer

ConvLSTM

Condition                                                   Predicted Frames

(c) Example 3

39

3

9

9

3

9

3

9

3

31

93

9

3

39

9

3

X

3

3

3

9

39

93

39

9

9

9

3

39

za

9

3

9

3

b

3

3

43

A

3

?

1

E

39

3

9

9

3

9

3

9

3

3

93

3

3

33

a

3

31

B

3

7

7

7

LoT

To

7

7

7

7v

G

7

7L

L7

C.2 3D Environments

±

±

±

±

t=0

73

143

144

163

182

202

221

240

260

279

299

Ground Truth

Transformer

S5

TECO-Transformer

TECO-S5

TECO-ConvS5

Condition                                                   Predicted Frames

t=0

17

35

36

68

101

134

167

200

233

266

299

Ground Truth

Transformer

S5

TECO-Transformer

TECO-S5

TECO-ConvS5

Condition                                                   Predicted Frames

t=0

73

143

144

163

182

202

221

240

260

279

299

Ground Truth

TECO-Transformer

TECO-ConvS5

Condition                                                   Predicted Frames

(a) 156 frames generated conditioned on 144 (action-conditioned).

Condition                                                   Predicted Frames

t=0

17

35

36

68

101

134

167

200

233

266

299

Ground Truth

TECO-Transformer

TECO-ConvS5

(b) 264 frames generated conditioned on 36 (action-conditioned).

Condition                                                   Predicted Frames

t=0

73

143

144

163

182

202

221

240

260

279

299

Ground Truth

TECO-Transformer

TECO-ConvS5

(a) 156 frames generated conditioned on 144 (action-conditioned).

Condition                                                   Predicted Frames

t=0

17

35

36

68

101

134

167

200

233

266

299

Ground Truth

TECO-Transformer

TECO-ConvS5

(b) 264 frames generated conditioned on 36 ( no action-conditioning).

±

±

±

±

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 × 3 kernels as implemented in Yan et al. [13].

D.2 Evaluation Metrics

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 × 10 -4 , 5 × 10 -4 , 1 × 10 -3 ] and chose the best model. For Transformer, we swept over model size, considering hidden dimensions of [512, 2014] and learning rates [ 1 × 10 -4 , 5 × 10 -4 , 1 × 10 -3 ] and chose the best model. We also observed better performance for the Transformer by convolving frames down to an 8 × 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 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 × 10 -4 , 5 × 10 -4 , 1 × 10 -3 ].

D.5 Long-Range 3D Environment Benchmarks

We follow the procedures from Yan et al. [13] and train models on the same pre-trained vectorquantized (VQ) 16 × 16 codes used by the baselines evaluated in that work. Models were trained to optimize a cross-entropy reconstruction loss between the predictions and true VQ codes. The evaluation of DMLab and Habitat involves both an action-conditioned 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.

All sequence models we trained used the same number of layers as the Transformer used in the TECO-Transformer trained by Yan et al. [13]. In addition, the TECO-Transformer, TECO-S5 and TECO-ConvS5 models we trained used the exact encoder/decoder configuration and MaskGit 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 a 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 × 16 to 8 × 8 ) by training an additional encoder/decoder with a codebook loss, L VQ . In addition, a MaskGit [98] dynamics prior loss, L prior , 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 a reconstruction loss, L recon . Models are trained to minimize the following total loss:

In addition, TECO includes the use of DropLoss [13], which drops out a percentage of random timesteps that are not decoded and therefore do not require computing the expensive L recon and L prior terms.

DMLab As mentioned above, the actions were randomly dropped out of sequences half the time (due to the two evaluation scenarios, action-conditioned 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, S5, ConvS5, Teco-Transformer, TECO-S5, TECO-ConvS5) by using this strategy when 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, in 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.

We trained each model, Transformer, Performer, S5, ConvS5, Teco-Transformer, TECO-S5, TECOConvS5, with three different learning rates [ 1 × 10 -4 , 5 × 10 -4 , 1 × 10 -3 ] and selected the best run for each model. See Tables 15-20 for more experiment configuration details.

The TECO-ConvSSM ablation used the exact same settings as TECO-ConvS5, except the state kernel was initialized with a random Gaussian and a lower learning rate of 1 × 10 -5 was required for stable training.

Minecraft and Habitat For Minecraft and Habitat, we only trained TECO-ConvS5 due to the costs of training on these datasets. See dataset details in Appendix E and reported compute costs in Yan et al. [13]. For Minecraft, we evaluated two different learning rates [ 1 × 10 -4 , 5 × 10 -4 ] and chose the best. For Habitat, we only performed one run with no further tuning. See Table 21 for further experiment configuration details.

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 at 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 et al. [81] and Su et al. [84] and generate 10,000 sequences for training.

E.2 DMLab

We use the DMLab long-range benchmark designed by Yan et al. [13] using the DeepMind Lab (DMLab) [99] simulator. The simulator generates random 3D mazes with random floor and wall textures. The benchmark consists of 40K action-conditioned, 300 frame videos at a 64 × 64 resolution. The videos are of an agent randomly navigating 7 × 7 mazes by choosing random points in the maze and navigating to 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 × 128 resolution. The videos are in Minecraft's marsh biome and the agent iterates walking forward for a 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

Weuse the Habitat long-range benchmark designed by Yan et al. [13] using the Habitat simulator [101]. The simulator renders trajectories using scans of real 3D scenes. Yan et al. [13] compiled 1400 indoor scans from HM3D [112], Matterport3D [113] and Gibson [114] to generate 200K action-conditioned, 300 frame videos with a 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nYes, the authors discuss limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

comments: 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.

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.

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.

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

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

References:
[1] Fast Autoregressive Transformers with Linear Attention
[2] Rethinking Attention with Performers
[3] cosFormer: Rethinking Softmax in Attention

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n+ **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.\n\nWEAKNESSES:\n- **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\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

comments: 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.

STRENGTHS:
1) The proposed method makes sense and is a natural extension of prior work.
2) 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.

2) 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.

3) 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].

[1] Gao et al,  Simvp: Simpler yet better video prediction, CVPR 2022.

QUESTIONS:
Please see the weaknesses.

LIMITATIONS:
No, they have not. However, I do not see a particular negative social impact associated for this work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-11