Efficiently Modeling Long Sequences with Structured State Spaces
Authors: Albert Gu
Karan Goel
Christopher Re
Venue: ICLR 2022 Oral (2022)
Keywords: sequence models
state space
RNN
CNN
Long Range Arena

Abstract: A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies.  Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of $10000$ or more steps.  A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) \( x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t) \), and showed that for appropriate choices of the state matrix \( A \), this system could handle long-range dependencies mathematically and empirically.  However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution.  We propose the Structured State Space sequence model (S4) based on a new parameterization for the SSM, and show that it can be computed much more efficiently than prior approaches while preserving their theoretical strengths.  Our technique involves conditioning \( A \) with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of a Cauchy kernel.  S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91\% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation $60\times$ faster (iii) SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X task of length 16k that all prior work fails on, while being as efficient as all competitors.

Published as a conference paper at ICLR 2022
EFFICIENTLY MODELING LONG SEQUENCES WITH STRUCTURED STATE SPACES
Albert Gu & Karan Goel & Christopher Ré
Department of Computer Science, Stanford University falbertqu, krtg]aattafooddcce chrismre@cs stanford.
ABSTRACT
A central goal of sequence modeling is designing single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies. Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of 10000 or more steps. A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) x' (t) Ax(t) + Bu(t), yt) = Cx(t) + Du(t), and showed that for appropriate choices of the state matrix A, this system could handle long- range dependencies mathematically and empirically. However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution. We propose the Structured State Space (S4) sequence model based on new parameterization for the SSM, and show that it can be computed much more efficiently than prior approaches while preserving their theoretical strengths. Our technique involves conditioning A with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of Cauchy kernel. S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91 % accuracy sequential CIFAR-10 with no data augmentation auxiliary losses, on par with larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation 60x faster (iii) SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X X task of length 16k that all prior work fails on, while being efficient all competitors
INTRODUCTION
A central problem in sequence modeling is efficiently handling data that contains long-range depen- dencies (LRDs). world time-series data often requires reasoning over tens of thousands of time steps, while few sequence models address even thousands of time steps. For instance, results from the long-range arena (LRA) benchmark (Tay et al., 202T) highlight that sequence models today perform poorly on LRD tasks, including one (Path-X X) where no model performs better than random guessing.
Since LRDs are perhaps the foremost challenge for sequence models, all standard model families such as continuous time models (CTMs), RNNs, CNNs, and Transformers include many specialized variants designed to address them. Modern examples include orthogonal and Lipschitz RNNs
(Arjovsky al. 2016 Erichson al. 202T) to combat vanishing gradients, dilated convolutions to increase context size (Bal et al., 2018 Oord et al.) 2016), and an increasingly vast family of efficient sequence length (Katharopoulos Choromanski et al. 2020). Despite being designed for LRDs, these solutions still perform poorly on audio classification (Gu et al.) 2021).
Transformers that reduce the quadratic dependence
al. 2020
challenging benchmarks such LRA (Tay al. 2021
An alternative approach to LRDs was recently introduced based on the state space model (SSM) (Fig. D. SSMs are a foundational scientific model used in fields such as control theory, computational neuroscience, and many more, but have not been applicable deep learning for concrete theoretical reasons. In particular, Gu et al. (2021) showed that deep SSMs actually struggle even on simple tasks, but can perform exceptionally well when equipped with special state matrices A recently derived
Code publicly available htttps///// http::////tttht..ccco///ysssssss sttp://seetttttttsssssssseeeeeee

---

Published as a conference paper at ICLR 2022
X
0
0
1
Ax + Bu
3
3
Ax + Bu
y = K +
= =x Du
y = Cx + Du
Continuous State Spaae
Long-Range Dependencies
Fast Discrete Representations
Figure 1: (Left) State Space Models (SSM) parameterized by matrices A, B, C, D map an input signal u(t) to output yt) through latent state xtt). (Center) Recent theory on continuous-time memorization derives special A matrices that allow SSMs to capture LRDs mathematically and empirically. (Right) SSMs can be computed either as a recurrence (left) or convolution (right). However. materializing these conceptual views requires utilizing different representations of its parameters (red, blue, green) which are very expensive to compute, S4 introduces a novel parameterization that efficiently swaps between these representations, allowing it to handle a wide range of tasks, be efficient at both training and inference, and excel at long sequences.
cap
solve a problem of continuous-tine memorization (Voelker et al. 2019 Gu et al. 2020a). Their Linear State Space Layer (LSSL) conceptually unifies the strengths of CTM, RNN and CNN models, and provides proof of concept that deep SSMs can address LRDs in principle.
Unfortunately, the LSSL. infeasible memory requirements induced by the state representation. For state dimension N and sequence length T., computing the latent state requires O L) operations and O space compared to a L + N) lower bound for both. Thus for reasonably sized models (e.g. N 256 in Gu et al. () the LSSL uses orders of magnitude more memory than comparably-sized RNNs or CNNs. Although theoretically efficient algorithms for the LSSL were proposed, we show that these are numerically unstable. In particular, the special A matrix is highly non-normal in the linear algebraic sense, which prevents the application of conventional algorithmic techniques. Consequently, although the LSSL showed that SSMs have strong performance, they are currently computationally impractical as general sequence modeling solution.
in practice because of prohibitive computation and
In this work, we introduce the Structured State Space (S4) sequence model based on the SSM that solves the critical computational bottleneck in previous work. Technically. S4 reparameterizes the structured state matrices A appearing in Voelker al. (2019); Gu al. (2020a) by decomposing them as the sum of low-rank and skew-symmetric term. Additionally, instead of expanding the standard SSM coefficient space, we compute truncated generating function frequency space, which can be simplified into multipole-like evaluation. Combining these two ideas, we show that the low-rank term can be corrected by the Woodbury identity while the skew-symmetric term can be diagonalized stably, ultimately reducing to a well-studied and thcoretically stable Cauchy kernel (Pan 2001 2017). This results in Ö(N + L) computation and O + L) memory usage, which is essentially hight for sequence models. Compared the LSSL, S4 up 30x faster with 400x less memory usage, while exceeding the LSSL'S performance empirically.
Empirically, S4 significantly advances the state-of-tee of-the-art for LRD. On the LRA benchmark for efficient sequence models, $4 as fast all baselines while outperforming them by 20+ points on average. S4 is the first model to solve the difficult LRA Path-X task (length-16384) achieving 88 % accuracy compared to 50 % random guessing for all prior work. On speech classification with 1ength-16000 sequences, S4 halves the test error (1.7%) of specialized Speech CNNs - by contrast, all RNN and Transformer baselines fail to learn (2 70% error).
Towards a general-purpose sequence model. Beyond LRD, a broad goal of machine learning is to develop a single model that can be used across wide range of problems. Models today are typically specialized to solve problems from a particular domain (e.g. images, audio, text, semess,res), and enable a narrow range of capabilities (e.g. efficient training, fast generation, handling irregularly sampled data). This specialization is typically expressed via domain-specific preprocessing. inductive biases, and architectures. Sequence models provide general framework for solving many these problems with reduced specialization - e.g. Vision Transformers for image classification with less
2

---

Published as a conference paper at ICLR 2022
2D information (Dosovitskiy al. 2020). However, most models such as Transformers generally still require substantial specialization per task to achieve high performance.
Deep SSMs in particular have conceptual strengths that suggest they may be promising as a general sequence modeling solution. These strengths include principled approach to handling LRDs, as well as the ability to move between continuous time, convolutional, and recurrent model representations, each with distinct capabilities (Fig. [T. Our technical contributions enable SSMs to be applied successfully to a varied set of benchmarks with minimal modification:
Large-scale generative modeling. On CIFAR-10 density estimation, S4 is competitive with the best autoregressive models (2.85 bits per dim). On WikiText-103 language modeling, S4 substantially closes the gap Transformers (within 0.8 perplexity), setting SoTA for attention-free models.
Fast autoregressive generation. Like RNNs, S4 can its latent state to perform 60x faster pixel/tkken generation than standard autoregressive models on CIFAR-10 and WikiText- 103.
Sampling resolution change. Like specialized CTMs, S4 can adapt to changes in time-series sampling frequency without retraining, e.g. at 0.5x frequency speech classification.
Learning with weaker inductive biases. With no architectural changes, S4 surpasses Speech CNNs speech classification, outperforms the specialized Informer model on time-series forecasting problems, and matches a 2-D ResNet on sequential CIFAR with over 90% accuracy.
BACKGROUND: STATE SPACES
Sections 2.1 to 2.4 describe the four properties of SSMs in Fig. [ the classic continuous time repre-
sentation, addressing LRDs with the HiPPO framework, the discrete-time recurrent representation, and the parallelizable convolution representation. In particular, Section [2.4] introduces the SSM convolution kernel K, which is the focus of our theoretical contributions in Section
2.1 STATE SPACE MODELS: CONTINUOUS- TIME LATENT STATE MODEL
The state space model is defined by the simple equation I It maps 1-D input signal u(() to an N-D latent state xr() before projecting to 1-D output signal y(1).
x' (1) = Ax(() Bu(()
(I)
y(t) Cx(()) + Du(t)
SSMs are broadly used in many scientific disciplines and related to latent state models such as Hidden Markov Models (HMM). Our goal is to simply use the SSM as black-box representation in deep sequence model, where A, B, C, D are parameters learned by gradient descent. For the remainder of this paper, we will omit the parameter D for exposition (or equivalently, assume D 0) because the term Du can be viewed as a skip connection and is easy to compute.
2.2 ADDRESSING LONG-RANGE DEPENDENCIES WITH HIPPO
Prior work found that the basic SSM I actually performs very poorly practice. Intuitively, one explanation is that linear first-order ODEs solve to an exponential function, and thus may suffer from gradients scaling exponentially in the sequence length (i.e., the vanishing/exploding gradients problem (Pascanu et al.] 2013)). ). To address this problem, the LSSL leveraged the HiPPO theory of continuous time memorization (Gu et al. 2020a). HiPPO specifies class of certain matrices A E R N X that when incorporated into [1. allows the state æ(t) to memorize the history of the input u(t). The most important matrix in this class is defined by equation (2). which will call the
HOPPO matrix. For example, the LSSL found that simply modifying an SSM from a random matrix A to equation improved its performance on the sequential MNIST benchmark from 60% to 98%.
(2n 1)2 + 1)1/2 if k
HiPPO Matrix)
Ank
n n+1
if k.
0
in < k
3

---

Published as a conference paper at ICLR 2022
2.3 DISCRETE-TIME SSM: THE RECURRENT REPRESENTATION
To be applied on a discrete input sequence (uo, ui....) instead of continuous function u(t), must be discretized by step size A that represents the resolution of the input. Conceptually, the inputs be viewed as sampling an implicit underlying continuous signal u(t), where uy u(kA).
To discretize the continuous- time SSM, we follow prior work using the bilinear method Tustin 1947), which converts the state matrix A into an approximation A. The discrete SSM is
3k AAt-1 + Bu A = (I (I- A) ' A/2 A)
(3)
yk Cak
B (I -)-
C-
Equation 3 is now a sequence-to-seequens map 14 - yı instead of function-to-fuuntion. Moreover the state equation now a recurrence in Zk. allowing the discrete SSM to computed like an RNN. Concretely, I& R can be viewed as a hidden state with transition matrix A.
Notationally, throughout this paper we use A, B,.... .... to denote discretized SSM matrices defined by B Note that these matrices are a function both as well as a step size A; we suppress this dependence for notational convenience when it is clear.
2.4 TRAINING SSMs: THE CONVOLUTIONAL REPRESENTATION
The recurrent SSM 3 is not practical for training on modern hardware due to its sequentiality. Instead, there well-known connection between linear fime-invariant (LTT) SSMs such as (I) and continuous convolutions. Correspondingly, (3) can actually be written as discrete convolution.
For simplicity let the initial state be I- O. Then unrolling explicitly yields
Zy = Bug
I1 = ABu + Buz
I2 = A Bug u + AB BBz
9/ = CBu
/1 CABu, + CBu
1/2 = CA + CABu, CBu2
This can be vectorized into convolution 44 with explicit formula for the convolution kernel
yk = CCA Buo + CA
Buu + ..... + CABu: + CBu
(4)
y K +
= - CA'
= CAB.CA CA CA'"BB
B).
(5)
(E[L]
In other words, equation 4 is single (non-circular) convolution and can be computed very efficiently with FFTs, provided that K is known. However, computing K in ( is non-trivial and is the focus We call K the SSM convolution kernel or filter.
of our technical contributions in Section 3
METHOD: STRUCTURED STATE SPACES (S4)
Our technical results focus on developing the S4 parameterization and showing how to efficiently compute all views of the SSM (Section[2:: the continuous representation (A, B, C*) ( the recurrent representation (A, B, C') and the convolutional representation K
Section 3.1 motivates our approach, which is based on the linear algebraic concepts of conjugation and diagonalization, and discusses why the naive application of this approach does not work. Section 3.2 gives an overview of the key technical components of our approach and formally defines the S4 parameterization. Section 3.3 sketches the main results, showing that S4 is asymptotically efficient (up to log factors) for sequence models. Proofs are in Appendices and
3.1 MOTIVATION: DIAGONALIZATION
The fundamental bottleneck computing the discrete-time SSM (3) is that it involves repeated matrix multiplication by A. For example, computing (5) naively as in the LSSL involves L successive multiplications by A. requiring 2 operations and O(NL) space.
To overcome this bottleneck, use structural result that allows simplify SSMs.
4

---

Published as a conference paper at ICLR 2022
Algorithm 1 S4 CONVOLUTION KERNEL (SKETCI))
Input: S4 parameters A, P. B, C CN and step size A Output: SSM convolution kernel K KL (A, C) for A A PQ (equation 5
Ö (I (I-X) A²) (
Truncate SSM generating function (SSMGF) length L
koo(w) Noi [a)
- [=]] (3 (A 1+w
- A) P
2: klo(w) k(((w)
Black-box Cauchy kernel
3: ((w) + 2
T+w (koo(w) kol (1)[1
Woodbury Identity
4: K [K(w)
= [[w]:w=eet
> Evaluate SSMGF at all roots unity w E Q2 > Inverse Fourier Transform
5: K < iFFT(K)
Lemma 3.1. Conjugation is an equivalence relation on SSMs (V
Proof. Write out the two SSMs with state denoted by x and i respectively:
x' = Ar + Bu
3' = V- + Bu
y = Cx
y = CVE
After multiplying the right side SSM by V, the two SSMs become identical with x = Vž. Therefore these compute the exact same operator + y, but with change of basis by V in the state
Lemma 3.1
motivates putting A into canonical form by conjugations] which is ideally more structured and allows faster computation. For example, if A were diagonal, the resulting computations become much more tractable. In particular, the desired K (equation (4) would be Vandermonde product which theoretically only needs O(( + L) log" (N + L)) arithmetic operations Paan 2001).
Unfortunately, the naive application of diagonalization does not work due to numerical issues. First, Vandermonde multiplication is itself famously ill-conditioned problem (Pan 2016). Furthermore, we derive the explicit diagonalization for the HiPPO matrix (2) and show it has entries exponentially large in the state size N, rendering the diagonalization numerically infeasible (e.g. CV in Lemma 3.] would not be computable). We note that Gu et al. (2021) proposed a different (unimplemented) algorithm to compute K faster than the naive algorithm. In Appendix numerically unstable for related reasons.
we prove that it is also
Lemma 3.2. The HiPPO matrix A in equation 2 is diagonalized by the matrix Vi particular, Vai,
(1-)) In
24d, Therefore V has entries of magnitude up to 24N/3
3.2 THE $4 PARAMETERIZATION NORMAL. PL.S LOW-RANK
The previous discussion implies that we should only conjugate by well-conditioned matrices V. The ideal scenario is when the matrix A is diagonalizable by perfectly conditioned (i.e., unitary) matrix. By the Spectral Theorem of linear algebra, this is exactly the class of normal matrices. However, this class of matrices is restrictive; in particular, it does not contain the HiPPO matrix (
We make the observation that although the HiPPO matrix is not normal, it can be decomposed as the sum of a normal and low-rank matrix. However, this still not useful by itself: unlike diagonal matrix, powering up this sum (in (5) is still slow and not easily optimized. We overcome this bottleneck by simultaneously applying three new techniques.
Instead of computing K directly, we compute its spectrum by evaluating its truncated generating function Σ/= K at the roots of unity Ç. K can then be found by applying inverse FFT.
This generating function is closely related to the matrix resolvent, and now involves matrix inverse instead of power. The low-rank term can now be corrected by applying the Woodbury identity which reduces (A + PQ*)- in terms of A truly reducing to the diagonal case.
Finally, we show that the diagonal matrix case is equivalent to the computation of a Cauchy kernel wy w well-studied problem with stable near-linear algorithms (Pan 2015 2017).
Note that although ultimately require A, conjugation commutes with discretization so refer A.
5

---

Published as a conference paper at ICLR 2022
Our techniques apply to any matrix that can be decomposed Normal Plus Low-Rank (NPLR).
Theorem 1. All HiPPO matrices from (Gu et al. 2020aa have NPLR representation
A VAV PQ- = (A- (V)))))
(6)
for unitary V E CN*N diagonal A, and low-rank factorization P.Q E R Nxx, These matrices
HiPPO- LegS, LegT, Lag? all satisfy = 1 or T = 2. In particular, equation is NPLR with r 1.
3.3 S4 ALGORITHMS AND COMPUTATIONAL COMPLEXITY
By equation (6). note that NPLR matrices can be conjugated into diagonal plus low-rank (DPLR) form (now over C instead of RR. Theorems 2 and B describe the complexities of SSMs where A is in DPLR form. $4 is optimal or near-optimal for both recurrent and convolutional representations.
Theorem (S4 Recurrence). Given any step size A, computing one step of the recurrence can be done in O(N) operations where N is the state size.
Theorem 2 follows from the fact that the inverse of a DPLR matrix is also DPLR (e.g. also by the Woodbury identity). This implies that the discretized matrix A the product of two DPLR matrices and thus has O(N) matrix multiplication. Appendix computes in closed DPLR form. Theorem 3 (S4 Convolution). Given any step size A, computing the SSM convolution filter K can be reduced Cauchy multiplies, requiring only O(N + L) operations and O(N + L) space.
Appendix C Definition formally formlly defines Cauchy matrices, which are related to rational interpolation problems. Computing with Cauchy matrices is an extremely well-studied problem in numerical analysis, with both fast arithmetic and numerical algorithms based on the famous Fast Multipole Method (FMM) (Pan 2001 2015 2017). The computational complexities of these algorithms under Proposition
various settings are described in Appendix C
We reiterate that Theorem is our core technical contribution, and its algorithm is the very motivation of the NPLR S4 parameterization. This algorithm is formally sketched in Algorithm
3.4 ARCHITECTURE DETAILS OF THE DEEP S4 LAYER
Concretely, S4 layer is parameterized follows. First initialize SSM with A set to the HiPPO matrix (2) By Lemma 3.1 and Theorem 1 this SSM is unitarily equivalent to some (A PQ" B, C)) for some diagonal A and vectors P. P,Q,B,C e C These comprise S4's 5N trainable parameters.
The overall deep neural network (DNN) architecture of S4 is similar to prior work. As defined above, S4 defines a map from IR' L R", i.e. I-D sequence map. Typically, DNNs operate on feature maps of size H instead of 1. S4 handles multiple features by simply defining H independent copies of itself, and then mixing the H/ features with a position-wise linear layer for total of ((11)) parameters per layer. Nonlinear activation functions are also inserted between these layers. Overall, S4 defines soquencc-tt-sequeeee exactly the same as related sequence models such as Transformers, RNNs, and CNNs.
map of shape (batch size, sequence length, hidden dimension),
4 EXPERIMENTS
Section 4.1 benchmarks $4 against the LSSL and efficient Transformer models. Section 4.2 validates S4 on LRDs: the LRA benchmark and raw speech classification. Section 4.3 investigates whether S4 can be used as general sequence model perform effectively and efficiently in wide variety of settings including image classification, image and text generation, and time series forecasting.
4.1 S4 EFFICIENCY BENCHMARKS
We benchmark that S4 can be trained quickly and efficiently, both compared to the LSSL, as well as
efficient Transformer variants designed for long-range sequence modeling. As outlined in Section 3
S4 theoretically much more efficient than the LSSL, and Table confirms that the S4 is orders of
magnitude more speed- and memory-efficient for practical layer sizes. In fact, S4's speed and memory
is competitive with the most efficient Transformer variants benchmarked by Tay al. (002
Linear Transformer (Katharopoulos et al. 2020) and Performer (Choromanski et al. 2020 in darameter-matched setting (Table 2 following the protocol of Tay al. (2021))

---

Published as a conference paper at ICLR 2022
Table 1:
Deep SSMs: The S4 parameterization with Algo- Table 2: Benchmarks vs. efficient Transformers
rithm
asymptotically more efficient than the I.SS..
LENGTH 1024
LENGTH 4096
Training STEP (MS)
MEMORY ALLOC. (MB)
Speed
Mem.
Speed
Mem.
Dim.
128
256
128
256
512
Transformer
1x
1x
1x
1x
LSSI.
9.32
20.6
140.7
222.1
1685
13140
Performer
1.23x
0.43x
3.79
0.086 x
S4
4.77
3.07
4.75
5.3
12.6
33.5
Linear Trans.
1.58x
0.37x
5,35x
0.067 x
Ratio
1.9x
6.7x
29.6x
42.0x
133x
392
S4
1.58x
0.43x
5.19x
0.091 x
Visualizations a trained S4 model on LRA Path-X, SSM convolution kernels K 16384 reshaped into a 128 128 image. (I.eft) Example from Path-X task, which involves deducing the markers are connected by path (Top) Filters from the first layer (Bottom) Filters from the last layer.
Figure 2:
Table 3: (Long Range Arena) Accuracy on full suite of LRA tasks, (Top) Original Transformer variants in (Bottom) Other models reported the literature.
LRA. Full results Appendix D.2
MODEL
LISTOPS
TEXT
RETRIEVAL
IMAGE
PATHFINDER
PATH-X X
AVG
Transformer
36.37
64.27
57.46
42.44
71.40
X
53.66
Reformer
37.27
56.10
53.40
38.07
68.50
X
50.56
BigBird
36.05
64.02
59.29
40.83
74.87
X
54.17
Linear Trans.
16.13
65.90
53.09
42.34
75.30
X
$0.46
Performer
18.01
65.40
53.82
42.77
77.05
X
51.18
FNet
35.33
65.11
59.61
38.67
77.80
X
54.42
Nyströmformer
37.15
65.52
79.56
41.58
70.94
X
57.46
Luna-256
37.25
64.57
79.29
47.38
77.72
X
59.37
58.35
76.02
87.09
87.26
86.05
88.10
80.48
4.2 LEARNING LONG RANGE DEPENDENCIES
As described in Sections 2.2 and 3.1 S4 uses principled approach to address LRDs based on the HiPPO theory of continuous-time memorization. Our goal in this section to validate that $4 achieves high performance on difficult tasks that require long-range reasoning. We focus here on two problems: (i) the Long-Range Arena, well-known benchmark designed test efficient sequence models on LRDs, and (ii) a speech classification problem as a real-world test of LRDs.
Long Range Arena (LRA). LRA (Tay al. 2021) contains 6 tasks with lengths 1K-16K steps, encompassing modalities and objectives that require similarity, structural, and visuospatial reasoning. Table compares S4 against the 11 Transformer variants from Tay et al. (202T) as well as follow-up work. S4 substantially advances the SoTA, outperforming all baselines on all tasks and averaging 80.48% compared to less than 60% for every baseline. Notably, S4 solves the Path-X X task, an extremely challenging task that involves reasoning about LRDs over sequences length 128 128 16384. All previous models have failed (i.e. random guessing) due to memory or computation bottlenecks, or simply being unable learn such long dependencies.
We analyze S4's performance on Path-X by visualizing its learned representations, in particular 1-D convolution kernels K which are the focus of technical results in Section 3 Fig. shows that S4

---

Published as a conference paper at ICLR 2022
learns variety of filters that display spatially consistent structure and demonstrate awareness of the 2-D nature of the data, In particular, the lower layers learn simple kernels that extract features from just few rows of local context while ignoring the rest of the image. On the other hand, higher layers aggregate information globally across full columns of the image at varying spatial frequencies. Filters in these higher layers span the entire context (16384 pixels), confirming S4's ability to learn LRDs.
Raw Speech Classification. Speech is typical real-world time series domain, involving signals sampled from an underlying physical process at high frequency. We perform speech classification using the Speech Commands dataset (Warden) 2018). While most sequence models for speech rely on extensive preprocessing (e.g. to MFCC features), we classify raw speech (length- 16000) following Romero et al. (2021). S4 achieves 98.3% accuracy, higher than all baselines that use the 100x shorter MFCC features, and validates that powerful LRD model is able to extract more information from the raw data and outperform hand-crafted pre- proccssscnnn Additionally, we include baseline CNN specifically designed for speech, the discriminator from the WaveGAN model (Donahue et al. 2019), which performs worse than S4 while having 90 more parameters and incorporating many more architectural heuristics (Appendix D.2
4.3 $4 AS A GENERAL SEQUENCE MODEL.
key goal of sequence modeling research is to develop a single model that can be applied in many domains (c.g. images, audio, text, time-series) with broad range of capabilities (c.g. efficient training, fast generation, handling irregularly sampled data). As a fundamental scientific model, SSMs a promising candidate that come with range of capabilities, and S4's strong results on LRD benchmarks spanning images, text, and speech are evidence of S4's potential as a general sequence model. In this section, we focus understanding this question in more depth by highlighting key strengths of S4 settings that usually require specialized models. The tasks we focus on (generative modeling, image classification, time-series forecasting) are considered as LRD tasks in the literature, and serve as additional validation that S4 handles LRDs efficiently.
Large-scale generative modeling. We investigate two well-studied image and text benchmarks to validate the scalability, flexibility, and efficiency of S4. These tasks require much larger models than our previous tasks- up to 250M parameters.
First, CIFAR density estimation
popular benchmark for autoregressive models, where images are flattened into a sequence of 3072 RGB subpixels that are predicted one by one. Table shows that with no 2D inductive bias, S4 competitive with the best models designed for this task.
Second, WikiText- 103 is an established benchmark for language modeling, an important task for large-scale sequence models where tokens are predicted sequentially based on past context. Although RNNs were the model of choice for many years, Transformers are now the dominant model in such applications that contain data that is inherently discrete. We show that alternative models
Table (Speech classification) Transformer, CTM, RNN, CNN, and SSM models. (MFCC) Standard pre-processed MFCC features (length-161). (Raw) Unprocessed signals (1ength-16000). (0.5x) x) Fre- quency change at test time. X denotes not applicable computationaly infeasible single GPU.
Table (Pixcl-level 1-D image classification) Transformer, RNN, CNN, and SSM models. Ex- tended results + citations in Appendix
SMNIST
PMNIST
SCIFAR
Transformer
98.9
97.9
62.2
MFCC
RAW
0.5x
LSTM
98.9
95.11
63.01
Transformer 90.75
X
X 30.68
T-L.TMM
98.4
95.2
72.2
Performer
80.85
30.77
UR-LSTM
99.28
96.96
71.00
UR-GRU
99.27
96.51
74.4
ODE-RNN
65.9
X
X
HiPPO-RNN
98.9
98.3
61.1
NRDE
89.8
16.49
15.12
LMU-FFT
98.49
ExpRNN
82.13
11.6
10.8
LipschitzRNN
99.4
96.3
64.2
LipschitZRNN
88.38
X
X
TCN
99.0
97.2
CKConv
95.3
71.66
65.96
TrellisNet
99.20
98.13
73.42
WaveGAN-D
96.25
X
CKCony
99.32
98.54
63.74
LSSL
93.58
X
X
LSSL
99.53
98.76
84.65
S4
93.96
98.32
96.30
S4
99.63
98.70
91.13
8

---

Published as a conference paper at ICLR 2022
Table 6: (CIFAR-10 density estimation) As generic Table 7: (Wikifext-103 language modeling) S4 ap- sequence model, S4 is competitive with previous autore-proaches the performance of Transformers with much gressive models (in bits per dim.) while incorporating no faster generation. (Top) Transformer baseline which 2D inductive bias, and has fast generation through its recur-our implementation is based on, with attention re- placed by S4. (Bottom) Attention-free models (RNNs and CNNs).
rence mode.
Model
bpd
2D bias
Images/ / sec
Transformer
3,47
None
0.32 (1x)
Model
Params
Test ppl.
Tokens sec
Linear Transf
3,40
None
17.85 (56x)
PixelCNN
3.14
2D conv.
Transformer
247M
20.51
0.8K (1x)
Row PixeIRNN
3,00
2D BiLSTM
GLU CNN
229M
37.2
PixeICN++
2.92
2D conv.
19.19 (59.97x)
AWD-QRNN
151M
33.0
Image Transf.
2.90
2D local attn.
0.54 (1.7x)
LSTM Hebb.
29.2
PixelSNAIL
2.85
2D conv, attn.
0.13 (0.4x)
TrellisNet
180M
29.19
Sparse Transf.
2.80
2D sparse attn.
Dynamic Conv.
255M
25.0
S4 (base)
2.92
None
20.84 (65.1x)
TaLK Conv.
240M
23.3
S4 (large)
2.85
None
3.36 (10.5x)
S4
249M
21.28
48K (60x)
to Transformers can still be competitive in these settings. By simply taking strong Transformer baseline (Baevski Auli, 2018) and replacing the self-attention layers, S4 substantially closes the gap to Transtormers (within 0.8 ppl), setting SoTA for attention-free models by over 2 ppl.
Fast autoregressive inference. A prominent limitation of autoregressive models is inference speed (e.g. generation), since they require pass over the full context for every new sample. Several methods have been specifically crafted to overcome this limitation such as the Linear Transformer, a hybrid Transformer/RNNN that switches to a stateful, recurrent view at inference time for speed.
As a stateful model, SSMs automatically have this ability (Fig. []. By switching to its recurrent representation (Section 2.3). S4 requires constant memory and computation per time step in contrast to standard autoregressive models which scale in the context length. On both CIFAR- --10 and WikiText-- 103, we report the throughput of various models at generation time, with S4 around 60x faster than a vanilla Transformer on both tasks (details in Appendix D.3.3).
Sampling resolution change. As continuous- time model. S4 automatically adapts to data sampled at different rates, a challenging setting for time series with a dedicated line of work (Rubanoval al. 2019 De Brouwer al. 2019 Romero al. (021). Without re-training. S4 achieves 96.3% accuracy at 0.5x the frequency on Speech Commands (Table 4. simply by changing its internal step size A (Section 2.3
Learning with weaker inductive bias. Beyond our results on speech (Section 4.2 we further validate that S4 can be applied with minimal modifications on two domains that typically require spe- cialized domain-speeific preprocessing and architectures. First, we compare $4 to the Informer (Zhoul iet al., (02II), a new Transformer architecture that uses a complex encoder-decoder designed for time-series forecasting problems. A simple application of S4 that treats forecasting sequence-to-sequenee transformation (Fig. 3) outperforms the Informer and other baselines on 40/50 settings across 5 forecasting tasks. Notably, S4 is better the longest setting in each task, e.g. reducing MSE by 37% when forecasting 30 days of weather data (Appendix D
masked
Finally, we evaluate S4 on pixel-level sequential image classification tasks (Table popular bench- marks which were originally LRD tests for RNNs (Arjovsky al., 2016). Beyond LRDs, these benchmarks point to a recent effort of the ML community to solve vision problems with reduced domain knowledge. in the spirit of models such Vision Transformers (Dosovitskiy et al. 2020) and MLP-Mixer (Tolstikhin et al. 2021) Sequential CIFAR is a particularly challenging dataset where outside of SSMs, all sequence models have gap of over 25% to simple 2-D CNN. By contrast, S4 is competitive with a larger ResNet18 (7.9M vs. 11.0M parameters), both with (93.16% vs. 95.62% or without (91.12% % vs. 89.46%) data augmentation. Moreover, it is much more robust to other architectural choices (c.g. 90.46% % vs. 79.52% when swapping BatchNorm for LayerNorm).
CONCLUSION
We introduce S4, a sequence model that uses a new parameterization for the state space model's continuous-time, recurrent, and convolutional views to efficiently model LRDs in a principled manner. Results across established benchmarks evaluating a diverse range of data modalities and model capabilities suggest that $4 has the potential to be effective general sequence modeling solution.

---

Published as a conference paper at ICLR 2022
REFERENCES
Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks, In The International Conference on Machine Learning (ICML), pp. 1120-1128, 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. arXiv preprint arXiv: 1809. 10853, 2018.
Shaojie Bai, Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv:1803.0127) 2018.
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Trellis networks for sequence modeling. In The International Conference Learning Representations (ICLR), 2019.
Shiyu Chang, Yang Zhang, Wei Han, Mo Yu, Xiaoxiao Guo. Wei Tan, Xiaodong Cui, Michael Witbrock, Mark Hasegawa-Johnson, and Thomas Huang. Dilated recurrent neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1904v.:904..0: 10509, 2019.
Narsimha Chilkuri and Chris Eliasmith. Parallelizing legendre memory unit training. The Interna- tional Conference on Machine Learning (ICML), 2021.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. The International Conference on Learning Representations (ICLR), 2020.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pp. 933-941. PMLR, 2017.
Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau. Gru-ode-bayes: Continuous modeling of sporadically-obsevye time series. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Chris Donahue, Julian McAuley, and Miller Puckette. Adversarial audio synthesis. In ICLR, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint 1XXi:20000.11299 2020.
N Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W Mahoney. Lipschitz recurrent neural networks. In International Conference on Learning Represen- tations, 2021.
Gene H Golub and Charles F Van Loan. Matrix computations, volume 3. JHU press, 2013.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tice Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020a. URL hash/ (f. b33 2aee122833c55coc
i neur wttpoorpps.cccppppeeeeeeem pttp://p///////////// Abstract. htmll
Albert Gu, Caglar Gulcehre, Tom Le Paine, Matt Hoffman, and Razvan Pascanu. gating mechanism of recurrent neural networks. In The International Conference on Machine Learning (ICML), 2020b.
Improving the
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-tim models with the structured learnable linear state space layer. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()) 1735-1780, 1997.
10

---

Published as a conference paper at ICLR 2022
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are Fast autoregressive transformers with linear attention. International Conference on Machine Learning, pp. 5156-5165. PMLR, 2020.
Mario Sizcano-Casadee and David Martínez-Rubic Cheap orthogonal constraints in neural networks: simple parametrization of the orthogonal and unitary group. In The International Conference on Machine Learning (ICML), 2019.
Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently recurrent neural network (IndRNN): Building longer and deeper RNN. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5457-5466, 2018.
Vasileios Lioutas and Yuhong Guo. Time-aware large kernel convolutions. In International Confer- ence on Machine Learning, pp. 6172-6183. PMLR, 2020.
Stephen Merity, Nitish Shirish Keskar, James Bradbury, and Richard Socher. Scalable language modeling: Wikitext-103 on single gpu in 12 hours. SysML, 2018.
Aaron den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprint arXiv:1609.03499 2016.
Victor Pan. Structured matrices and polynomials: unified superfast algorithms. Springer Science & Business Media, 2001.
Victor Pan. Fast approximate computations with cauchy matrices and polynomials. Mathematics of Computation, 8663009)2799-28222 2017.
Victor Y Pan. Transformations of matrix structures work again. Linear Algebra and Its Applications, 465:107-138, 2015.
Victor Y Pan. How bad are vandermonde matrices? SIAM Journal on Matrix Analysis and Applications, 37(2):676 694466966666666 2016.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural
networks. In International conference on machine learning, pp. 1310-1318, 2013.
Jack Rae, Chris Dyer, Peter Dayan, and Timothy Lillicrap. Fast parametric learning with activation memorization. The International Conference on Machine Learning (ICML), 2018.
Prajit Ramachandran, Tom Le Paine, Pooya Khorrami, Mohammad Babaeizadeh, Shiyu Chang, Yang Zhang, Mark Hasegawa-Johnson, Roy H Campbell, and Thomas S Huang. Fast generation for convolutional autoregressive models. arXiv preprint arXiv:1704.0660 2017.
David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomezak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. arXiv preprint arXiv:2102.02611 2021.
Yulia Rubanova, Tian Qi Chen, and David K Duvenaud. Latent ordinary differential equations for trregularly-sample time series. In Advances in Neural Information Processing Systems, pp. 5321-5331, 2019.
T Konstantin Rusch and Siddhartha Mishra. Unicormn: A recurrent model for learning very long time dependencies. The International Conference on Machine Learning (ICML), 2021.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications. arXiv preprint 1rXiv:1701.05511 2017.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena benchmark for efficient transformers. In International Conference on Learning Representations, 2021. [ //openreview. net/ pwppenrwweeiineeeeooouddaee/oe/ /ttp://rrmmiii//eeeeeree--
Ily Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lueas Beyer, Xiaohua Zhai Thomas Un- terthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al. Mip-mixer: An all-mip architecture for vision. arXiv preprint arXiv:2105.01601, 2021.

---

Published as a conference paper at ICLR 2022
Trieu H. Trinh, Andrew M Dai, Minh-Thang Luong, and Quoe Le. Learning longer-term depen- dencies in RNNs with auxiliary losses. In The International Conference on Machine Learning (ICML), 2018.
Arnold Tustin. A method of analysing the behaviour of linear systems in terms of time series. Journal of the Institution of Electrical Engineers- Parrt IIA: Automatic Regulators and Servo Mechanisms, 94(1):130-142, 1947.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz. Kaiser, and Iilia Polosukhin. Attention all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Aaron Voelker, Ivana Kajić, and Chris Eliasmith. Legendre memory units: Continuous timeeo represen- tation in recurrent neural networks. In Advances in Neural Information Processing Systems, pp. 15544-15553, 2019.
Aaron Russell Voelker. Dynamical systems in spiking neuromorphic hardware. PhD thesis, University of Waterloo, 2019.
Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. ArXiv, abs/1804.032009 2018.
Max A Woodbury. Inverting modified matrices. Memorandum report, 42:106, 1950.
Felix Wu, Angela Fan, Alexei Baevski, Yann N Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions, In The International Conference Learning Representa- tions (ICLR), 2019.
Haoyi Zhou, Shanghang Zhang. Jieqi Peng. Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting, In The Thirty- Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Virtual Conference, volume 35, pp.
1111111115. AAAI Press, 2021.
A DISCUSSION
Related Work. Our work is most closely related to a line of work originally motivated by a particular biologically-inspree SSM, which led to mathematical models for addressing LRDs. Voelker (2019); Voelker et al.] (2019) derived a non-trainable SSM motivated from approximating a neuromorphic spiking model, and Chilkuri & Eliasmith (202T) showed that it could be sped up at train time with a convolutional view. (Gu et al. (2020a) extended this special case to a general continuous-time function approximation framework with several more special cases of A matrices designed for long-range dependencies. However, instead of using a true SSM, all of these works fixed a choice of A and built RNNs around it. Most recently, Gu et al. (202T) used the full (T) explicitly as deep SSM model, exploring new conceptual views of SSMs, as well as allowing A to be trained. As mentioned in Section D their method used a naive instantiation of SSMs that suffered from an additional factor of N in memory and N2 in computation.
Beyond this work, our technical contributions (Section 3) on the S4 parameterization and algorithms are applicable to broader family of SSMs including these investigated in prior works, and our techniques for working with these models may be of independent interest.
Implementation. The computational core of S4's training algorithm is the Cauchy kernel discussed in Sections 3.2 and 3.3 and Appendix C.3 As described in Appendix C.3 Proposition []] there are many algorithms for it with differing computational complexities and sophistication. Our current implementation of S4 actually uses the naive O(NL) algorithm which is easily parallelized on GPUs and has more easily accessible libraries allowing it to implemented; we leverage the pykeops kernel operations. However, this library is a much more general library
library for memory-efficien that may not be optimized for the Cauchy kernels used here, and we believe that a dedicated CUDA implementation can be more efficient. Additionally, as discussed in this work, there are asymptotically faster and numerically stable algorithms for the Cauchy kernel (Proposition 5). However, these algorithms are currently not implemented for GPUs due to a lack of previous applications that require them. We believe that more efficient implementations of these self-contained computational kernels
12

---

Published as a conference paper at ICLR 2022
possible, and that S4 (and SSMs large) may have significant room for further improvements in efficiency.
Limitations and Future Directions. In this work, we show that S4 can address a wide variety of data effectively. However, it may not necessarily be the most suitable model for all types of data. For example, still found a gap compared to Transformers for language modeling. An interesting future direction exploring combinations of S4 with other sequence models complement their strengths. We are excited about other directions, including continuing to explore the benefits of S4 on audio data (e.g. pre-training generation settings), and generalizing HiPPO and S4 to higher-dimensional data for image and video applications.
B NUMERICAL INSTABILITY OF LSSL
This section proves the claims made in Section 3.1 about prior work. We first derive the explicit diagonalization of the HiPPO matrix, confirming instability because of exponentially large entries, We then discuss the proposed theoretically fast algorithm from (Gu et al.) 202T) (Theorem 2) and show that it also involves exponentially large terms and thus cannot be implemented.
B.I HIPPO DIAGONALIZATION
Proof of Lemma 3.2
The HiPPO matrix
2 is equal, up to sign and conjugation by a diagonal matrix,
1
-1
1
-3
3
-5
4
A
1
-3
5
7
-1
-5
7
-9
1
-3
5
-7
-11
-1
-5
7
-9
11
-13
(-1)m-*(2)) (2k 1)
n
Ank k+1
n n
Our goal is to show that this A is diagonalized by the matrix
1
3
V i
6
1
10
15
7
15
35
28
or in other words that columns of this matrix are eigenvectors of A. Concretely, we will show that the j-th column of this matrix vb with elements
( -12)) (2))
<
i2j
an eigenvector with eigenvalue j ++ In other words we must show that for all indices [M].
(Av Σ Amoo (j (1+1)e
If l < j, then for all inside the sum, either k < i or i < j. In the first case A 0 and in the second case o.55) 0, both sides of equation ( equal 0.
13

---

Published as a conference paper at ICLR 2022
It remains to show the case k 2 j, which proceeds by induction k. Expanding equation the formula for A yields
using
k-1 ((1))-1(21+ (2i (1(((tt)) (it))
k (kt
(Av) ( Σ Am
+ (k: 1k+1)
1)
2j
2j
In the base case k j. the sum disappears and we are left with (Av -i+11 = (i+1)a as desired.
Otherwise, the sum for (Av) edge terms. The result follows from applying the inductive hypothesis and algebraic simplification:
(1) is the same the sum for (Av) k- but with sign reversed and few
(1)
k + J
k 1 1
+ j
(Av) -(Av) - (2h 1)
+ h
+ (k ( 1)
2j
2j
2j
h: 1-1+ + 1
h: 1 + j
)
+ (k: + 1) (h+J))
(k+tt))
-(++1)
(1: 1)
2j
2j
2j
k 1 1
(kttt)))
(1 +
+ (Ie +
2j
2j
(k + j)!
+
(1 + k)
+ (K: + 1)
(k - 1 1 ()((2))
(k: 1- - j){(22)) + (1: +1 ((+11)0) (k+))
(k: j)!
1)
2j
(1+1)
- (k - 1) (k -()(()))
(I: (k+11)
2j
k k
[k+
(j (k: + 1)
2j )
+ (l: +1)
2j
( 1)0
B.2 FAST BUT UNSTABLE LSSL ALGORITIMM
Instead of diagonalization, Gu et al. 2021 Theorem 2) proposed a sophisticated fast algorithm to compute
K (A,AB, (CB, CAB. "AB,
CA' B).
This algorithm runs in O(N log² N + L log L) operations and O(N + L) space. However, we now show that this algorithm is also numerically unstable.
There are several reasons for the instability of this algorithm, but most directly we can pinpoint particular intermediate quantity that they use.
Definition 1. The fast LSL algorithm computes coefficients of p(z), the characteristic polynomial of A, as an intermediate computation. Additionally, computes the coefficients of its inverse, p(x)- (mod )²).
We now claim that this quantity is numerically unfeasible. We narrow down to the case when A = I is the identity matrix. Note that this case is actually in some sense the most typical case: when discretizing the continuous matrix A is brought closer to the identity. For example, with the Euler discretization A I AA, we have A -> as the step size O.
SSM to discrete-time by a step-size A, the discretized transition
Lemma B.I. When A I, the fast LSSL algorithm requires computing lerms exponentially large in N.
Proof. The characteristic polynomial of I is
p(x) det I rI = (1 - N
14

---

Published as a conference paper at ICLR 2022
2
These coefficients have size up V=N/2'
The inverse of p(æ:) has larger coefficients. can be calculated in closed form by the generalized binomial formula:
20
N +
(1 - σ)
k
k=U
Taking this (mod )²), the largest coefficient is
(N +
L L- (N N
N 22
(L (1-1)(L- -2)
12) ... ...(L...... + +1) (N - 1)!
1
When L = N N-1 this is
[2(N -1)
22N
NN-1
T
already larger than the coefficients of (1 - xxx,, and only increases as L grows.
C $4 ALGORITHM DETAILS
This section proves the results of Section 3.3 providing complete details of our efficient algorithms for S4.
Appendices to prove Theorems respectively.
C.I NPLR REPRESENTATIONS OF HиPPO MATRICES
We first prove Theorem I showing that all HiPPO matrices for continuous- tim memory fall under the S4 normal plus low-rank (NPLR) representation.
Proof of Theorem We consider each of the three cases HiPPO-LagT, HiPPO-LegT, and HiPPO-
LegS separately. Note that the primary HiPPO matrix defined in this work (equation 2 is the HiPPO-LegT matrix.
HiPPO-LagT. The HiPPO-LagT matrix is simply
n
Ank
2
n
n2
>
k
i
1
2
A
1
1
Adding the matrix of all 2> which rank 1, yields
7
2
This matrix is now skew-symmetric. Skew-symmetric matrices are a particular case of normal matrices with pure-imaginary eigenvalues.
Gu et al. 2020a) also consider a case of HiPPO corresponding to the generalized Laguerre polynomi- als that generalizes the above HiPPO-LagT case. In this case, the matrix A (up conjugation by diagonal matrix) ends up being close to the above matrix, but with a different element on the diagonal. After adding the rank-1 1 correction, it becomes the above skew-symmetric matrix plus multiple of
15

---

Published as a conference paper at ICLR 2022
the identity. Thus after diagonalization by the same matrix in the LagT case, is still reduced to diagonal plus low-rank (DPLR) form, where the diagonal now pure imaginary plus real constant.
HiPPO-LegS. We restate the formula from equation for convenience.
(2n + 1)2/ 1)2/2
if n > k
Ank
n+
if n k.
if n
Adding (2n + 12) + 1)1/2 to the whole matrix gives
(2n + ()k 11
if > k if k
(2n + 12 + 1)1/2
if n < k
Note that this matrix is not skew-symmetric, but is I S where S is skew-symmetric matrix. This is diagonalizable by the same unitary matrix that diagonalizes S. HiPPO-LegT.
Up the diagonal scaling, the LegT matrix
1
1
-1
1
1
By adding -1 to this matrix and then the matrix
2
2
2
2
the matrix becomes
-2
2
2
2
which is skew-symmetric. In fact, this matrix is the inverse of the Chebyshev Jacobi.
An alternative way to see this is as follows. The Leg'T matrix is the inverse of the matrix
-1
This can obviously be converted to a semmmtytettiic matrix by adding a rank 2 term. The inverses of these matrices are also rank-2 differences from each other by the Woodbury identity.
A final form is
0
1
07
0
0
17
-1
0
1
0
0
+
0
1
-1
0
1
0
0
This has the advantage that the rank-2 correction is symmetric (like the others), but the normal skew-symmetric matrix is now 2-quasiseparable instead of 1-quasiseparable.
16

---

Published as a conference paper at ICLR 2022
C.2 COMPUTING THE S4 RECURRENT VIEW
We prove Theorem 2 showing the efficiency of the S4 parameterization for computing one step the recurrent representation (Section 2.3
Recall that without loss of generality, we can assume that the state matrix A = A - PQ" is diagonal plus low-rank (DPLR), potentially over C. Our goal in this section to explicitly write out closed form for the discretized matrix A.
Recall from equation that
= (T- (I+ 4/2 A)
B (I- A) A)-ABB
We first simplify both terms the definition of A independently. Forward discretization. The first term is essentially the Euler discretization motivated in Section 2.3
1+ A
=
A PQ")
2
A
2
+ (A
PQQ)
2
A
Ao
2
where Ao defined as the term in the final brackets.
Backward discretization. The second term is known as the Backward Euler's method. Although this inverse term is normally difficult to deal with, in the DPLR case we can simplify it using Woodbury's s Identity (Proposition
PO'))
I
A
=
(A
2
2
2
A + PQ"
A
A
2
D
- DP
DP) QDD
A
2
A1
A
where D = -A)
A) and A₁ defined as the in the final brackets. Note that (1 Q" DP) is actually a scalar in the case when the low-rank term has rank 1.
S4 Recurrence. Finally, the full bilinear discretization can be rewritten in terms of these matrices as
A1
2
B A₁ AB B. A
The discrete-time SSM
becomes
A + Buk
A₁ A1 + 2A, Bug
Vk = C.a.
Note that Ao. A, are accessed only through matrix-vects multiplications. Since they are both DPLR, they have O(N) matrix-vector multiplication, showing Theorem 2
C.3 COMPUTING THE CONVOLUTIONAL VIEW
The most involved part of using SSMs efficiently is computing K. This algorithm sketched in Section 3.2 and is the main motivation for the $4 parameterization. In this section, we define the necessary intermediate quantities and prove the main technical result.
17

---

Published as a conference paper at ICLR 2022
The algorithm for Theorem 3 falls in roughly three stages, leading Algorithm I
Assuming A has been conjugated into diagonal plus low-rank form, we successively simplify the problem of computing K by applying the techniques outlined in Section 3.2
Remark C.I. We note that for the remainder of this section, transpose C be column vector of shape C.N or C Nx1 instead of matrix or row vector CI'NN as in (T). In other words the SSM
(1) = Ax(t) + Bu(t) u(t) = a(() + Du(t)-
(8)
This convention is made so that C has the same shape as B, P, Q and simplifies the implementation of S4.
Reduction Diagonalization By Lemma 3.1 we can switch the representation by conjugating with any unitary matrix. For the remainder of this section, we can assume that A is (complex) diagonal plus low-rank (DPLR).
Note that unlike diagonal matrices, a DPLR matrix does not lend itself to efficient computation of K. The reason is that K computes terms C' A B which involve powers of the matrix A. These trivially computable when A is diagonal, but is no longer possible for even simple modifications to diagonal matrices such as DPLR.
Reduction 1: SSM Generating Function To address the problem of computing powers of A. we introduce another technique. Instead of computing the SSM convolution filter K directly, we introduce a generating function on its coefficients and compute evaluations of it.
Definition 2 (SSM Generating Function). We define the following quantities:
The SSM convolution function is K(A,B, (C CCB, C,C AB,...) and the (truncated)
SSM filter of length L
Kr, -A, AB ......,AAA C'A [-1
B) ER
(9)
The SSM generating function at node z is
R(2; CC: (I-Aa=--1
Az)
(10)
and the truncated SSM generating function node is
L-1 k, =C= E:
A2)
(11)
1=0
The truncated SSM generating function
functioon at nodes 12 C is
A,, = (Kn(WRi A, (,B,C) kc(M]
(12)
Intuitively, the generating function essentially converts the SSM convolution filter from the time domain to frequency domain. Importantly, it preserves the same information, and the desired SSM convolution filter can be recovered from evaluations its generating function.
Lemma C.2. The SSM function Kr, (A, B, C)) can be computed from the SSM generating function (9; A, B, C) at the roots of unity Ω (exp( - -arp- 4 k [T]] stably in OL L)) operations.
Proof. For convenience define
K K,K,
K = ((,,,,,, k()) = K((;,,),;
18

---

Published as a conference paper at ICLR 2022
Note that
L-1 exp (-2mi 3k:
Kj = Σ Kx
I.
k=0
Note that this is exactly the same as the Discrete Fourier Transform (DFT):
=
Therefore K can be recovered from K with single inverse DFT, which requires ( log L) operations with the Fast Fourier Transform (FFT) algorithm.
Reduction 2: Woodbury Correction The primary motivation of Definition is that it turns powers of Á into a single inverse of A (equation (T). While DPLR matrices cannot be powered efficiently due to the low-rank term, they can be inverted efficiently by the well-known Woodbury identity.
Proposition 4 (Binomial Inverse Theorem or Woodbury matrix identity Woodbury (1950); Golub Van Loan ). ). Over commutative ring R. let A E R NNNN and U,V E R xp, Suppose A and A UV are invertible. Then I, V A U E RPXP invertible and
(A (A+ A A U + V -V--
With this identity, can convert the SSM generating function DPLR matrix A into one on just its diagonal component.
Lemma C.3. Let = A PQ be a diagonal plus low-rank representation. Then for any root of unity E Q. the truncated generating, function satisfies
= 1 "CBBBBB-CCHRPPLLEHBEE Q
Č= C-I(I - -A")
2 1 1-22-
R(z:A)
Al+z
Proof. Directly expanding Definition! yields
K2 A,B, C' C' AB. C² A" /. B
= = (1- =C' 7") (I- B
- Az) -BB B
where :- Č- C (I (I-A")
We can now explicitly expand the discretized SSM matrices A and B back in terms of the original SSM parameters A and B. Lemma C.4 provides an explicit formula, which allows further simplifying
2
121-2
C (T - Az B =
A B
1+2
A1+2
2
8x
(21
(21-z
A
B
I+2
(1
2
(
Q" R())
1+
The last line applies the Woodbury Identity (Proposition 4h where R(z) = ( - A)
The previous proof used the following self-contained result to back out the original SSM matrices from the discretization.
19

---

Published as a conference paper at ICLR 2022
Lemma C.4. Let A, B be the SSM matrices A, B discretized by the bilinear discretization with step size A. Then
C (I-
B=
2A C*
2 2
AA
B
1+ 2
1+
Proof.
Recall that the bilinear discretization that we use (equation
3
is
A
A
I
A
A
2
2
A
B =
A
AB
2
The result is proved algebraic manipulations.
A
A A
C (I - Az) -1 B C
A
-
I+
1
/
2
B
2
2
2
2
(1
|
A 2
A
4 A
C
A
A
1
B
2
A
C
I(1
2)
A(1
+
AB
2
AA
C
B
1-2
2 1+
2A C'
[ 1
2
I
AA
B
1+2
1
Note that in the S4 parameterization, instead of constantly computing Č = C c( (1- we can simply reparameterize our parameters to learn Ö directly instead of C, saving minor computation cost and simplifying the algorithm.
Reduction 3: Cauchy Kernel We have reduced the original problem of computing K to the problem of computing the SSM generating function Kr (S2; A, B, C) the case that A a diagonal matrix. We show that this is exactly the same as a Cauchy kernel, which is a well-studied problem with fast and stable numerical algorithms.
Definition 3. A Cauchy matrix or kernel on nodes Ω = (wi) E AM and A = (A²) E C.N is
1
M MXN- (Q2. (MM))EE(M)JEN
My
X2 Wi A
The computation time of Cauchy matrix-vector product of size M x N is denoted by C(M,N).
Computing with Cauchy matrices is an extremely well-studied problem in numerical analysis, with both fast arithmetic algorithms and fast numerical algorithms based on the famous Fast Multipole
Method (FMM) (Pan 2001 2015 2017).
Proposition 5 (Cauchy). A Cauchy kernel requires O(M N) space, and operation count
O (MN)
naively in exact arithmetic numerically to precision E.
C(M,N)
O ((M+ Nog N)) O ((M + N) log + No
Corollary C.5. Evaluating Q' R(S;; A)P (defined in Lemma C.3 for any set of nodes Q E CL diagonal matrix A, and vectors P. can be computed in C(L, N) operations and O + N) space, where C(I, N) Ö(T + N) is the cost of Cauchy matrix-vector multiplication.
20

---

Published as a conference paper at ICLR 2022
Proof. For any fixed w E Q we want to compute
Σ, Computing this over wi is therefore
exactly a Cauchy matrix-vector multiplication.
This completes the proof of Theorem 3
In Algorithm| D note that the work is dominated by Step
2 which has a constant number of calls to a black-box Cauchy kernel, with complexity given by Proposition 5
D EXPERIMENT DETAILS AND FULL RESULTS
This section contains full experimental procedures and extended results and citations for our experi- mental evaluation in Section 4 Appendix D.I corresponds to benchmarking results in Section 4.1
Appendix D.2 corresponds to LRD experiments (LRA and Speech Commands) in Section 4.2
and Appendix D.3 corresponds to the general sequence modeling experiments (generation, image classification, forecasting) in Section 4.3
D.1 BENCHMARKING
Benchmarking results from Table and Table were tested on a single A100 GPU.
Benchmarks against LSSL For given dimension H, single LSSL $4 layer was constructed with H hidden features. For LSSL, the state size N was set to H as done in (Gu et al. (02T). For S4, the state size N was set to partheter-match the LSSL, which was a state size of NV' due to differences I in the parameterization. Table benchmarks single forward+backward pass a single layer.
Benchmarks against Efficient Transformers Following TTay et al. 2021 the Transformer mod- els had layers, hidden dimension 256 with 4 heads, query/key/value projection dimension 128, and batch size 32, for a total of roughly 600k parameters. The S4 model was parameter tied while keeping the depth and hidden dimension constant (leading state size of N 256).
We note that the relative orderings of these methods can vary depending on the exact hyperparameter settings.
D.2 LONG-RANGE DEPENDENCIES
This section includes information for reproducing our experiments the Long- Mange Arena and Speech Commands long-range dependency tasks.
Long Range Arena Table 8 ct al. 202I).
contains extended results table with all 11 methods considered in Tay
For the S4 model, hyperparameters for all datasets are reported Table 9 For all datasets, we used the AdamW optimizer with a constant learning rate schedule with decay validation plateau. However, the learning rate on HiPPO parameters (in particular A, P, Q, B, C, A) were reduced to maximum starting LR of 0.001, which improves stability since the HiPPO equation is crucial to performance.
The S4 state size was always fixed to N = 64.
As $4 is a seuuence-to-seqeeee model with output shape (batch, length, dimension) and LRA tasks classification, mean pooling along the length dimension was applied after the last layer.
We note that most of these results were trained for far longer than what was necessary achieve SotA results (e.g., the Image task reaches SotA in epoch). Results often keep improving with longer training times.
Hardware. All models were run on single GPU. Some tasks used an A 1100 GPU (notably, the Path-X experiments), which has a larger max memory of 40Gb. To reproduce these on smaller GPUs, the batch size can be reduced gradients
be accumulated for batches.
Path-X. We remark that an earlier version of this paper reported higher score for Path-X. This earlier version used a different variant of the dataset, due to nisunderstanding of the properties of the dataset. More specifically, we found that S4 scored 93. 68% accuracy version that involved
21

---

Published as a conference paper at ICLR 2022
Table 8: Full results for the Long Range Arena (LRA) benchmark for long-ran rangrrange dependencies in sequence
models. (Top): Original Transformer variants in LRA. (Bottom): Other models reported in the literature.
Model
LiSTOPs
TEXT
RETRIEVAL
IMAGE
PATHFINDER
PATH-X X
AVG
Random
10.00
50.00
50.00
10.00
50.00
50.00
36.67
Transformer
36.37
64.27
57.46
42.44
71.40
53.66
Local Attention
15.82
52.98
53.39
41.46
66.63
46.71
Sparse Trans.
17.07
63.58
59.59
44.24
71.71
51.03
Longformer
35.63
62.85
56.89
42.22
69.71
X
52.88
Linformer
35.70
53.94
S2.27
38.56
76.34
X
51.14
Reformer
37.27
56.10
53,40
38.07
68.50
50.56
Sinkhorn Trans.
33.67
61.20
53.83
41.23
67.45
51.23
Synthesizer
36,99
61.68
54.67
41.61
69.45
$2.40
BigBird
36.05
64.02
59.29
40.83
74.87
54.17
Linear Trans.
16.13
65.90
53.09
42.34
75.30
X
50,46
Performer
18.01
65.40
53.82
42.77
77.05
X
51.18
FNet
35.33
65.11
59.61
38.67
77.80
54,42
Nyströmformer
37.15
65.52
79.56
41,58
70.94
X
57.46
Luna-256
37.25
64.57
79.29
47.38
77.72
X
59.37
S4
58.35
76.02
87.09
87.26
86.05
88.10
80.48
Table
The values of the best hyperparameters found for classification datasets; LRA (Top) and images/speech (Bottom). LR learning rate and WD is weight decay. BN and LN refer to Batch Normalization and Layer Normalization.
Depth Features I! Norm Pre-norm Dropout
LR
Batch Size
Epochs
WD
Patience
ListOps
128
BN
False
100
0.01
Text
64
BN
True
0.001
50
20
Retrieval
256
BN
True
0.002
64
20
20
Image
512
LN
False
0.004
50
200
0.01
20
Pathfinder
256
BN
True
0.1
0.004
100
200
10
Path-X
256
BN
0.0005
32
20
CIFAR-10
1024
LN
False
0.25
0.01
50
200
0.01
20
Speech Commands (MFCC)
256
LN
False
0.2
0.01
100
50
Speech Commands (Raw)
128
taking the 256 x 256 resolution version of the Pathfinder dataset and averaging every 2 x 2 square; erroneously thought that this version of the dataset equivalent the original Path-X.
After discussions with the LRA authors, we discovered that this not equivalent to the 128 128 resolution Pathfinder dataset (the correct Path-X), which is in fact much harder. In fact, Path-X X is so difficult that 2D CNN without global receptive field (e.g. ResNet-18 achieve above chance. This fact led to the original misunderstanding, as we could not solve this image classification task with ResNet and thought the data might have errors.
ResNet-34) also cannot
Speech Commands We provide details of sweeps run for baseline methods run by us-numbers for all others method are taken from Gu et al. (2021). The best hyperparameters used for S4 are included Table
Transformer Vaswani al. 2017 For MFCC, we swept the number of model layers (2, 4]. dropout {0, 0.13 } and learning rates f0.001, 0.0005). We used 8 attention heads, model dimension 128, prenorm, positional encodings, and trained for 150 epochs with batch size of 100. For Raw, the Transformer model's memory usage made training impossible.
Performer [Choromanski er al. 2020) For MFCC, we swept the number of model layers {2,4). dropout (0, 0.1} and learning rates (0.001, 0.0005). We used attention heads, model dimension 128, prenorm, positional encodings, and trained for 150 epochs with a batch size of 100. For Raw, we used model dimension of 128, attention heads, prenorm, and batch size of 16. We reduced the number of model layers to 4, so the model would on the single GPU. We trained for 100 epochs with learning rate of 0.001 and no dropout.
22

---

Published as a conference paper at ICLR 2022
ExpRNN (Lezcano- Casado & Martinez-Rubia 2019 For MFCC, swept hidden sizes {256, 512) and learning rates {0.001, 0.002, J.0005]. Training was run for 200 epochs, with a single layer model using a batch size of 100. For Raw, we swept hidden sizes {32, 64) and learning rates {0.001, 0.0005) (however, ExpRNN failed to learn).
Lipschitz RNN MErichson et al. 2021) For MFCC, we swept hidden sizes (256, 512} and learning rates {0.001, 0.002, 0.0005]. Training was run for 150 epochs, with single layer model using batch size of 100. For Raw, we found that LipschitZRNN was too slow to train on single GPU (requiring a full day for 1 epoch of training alone).
WaveGAN Discriminator (Donahue et al. 2019) The WaveGAN-D in Table actually improved version of the discriminator network from the recent WaveGAN model for speech (Donahue et al.) 2019). This CNN actually did not work well out-of-the-box, and we added several features to help perform better. The final model is highly specialized compared our model, and includes:
Downsampling or pooling between layers, induced by strided convolutions, that decrease the sequence length between layers.
A global fully-connected output layer; thus the model only works for input sequence length and does not work on MFCC features or the frequency-shift setting in Table F
Batch Normalization essential, whereas $4 works equally well with either Batch Normal-
ization or Layer Normalization.
Almost 90x as many parameters as the S4 model (26.3M vs. 0.3M).
D.3 GENERAL SEQUENCE MODELING
This subsection corresponds to the experiments in Section 4.3
Because of the number of experiments
in this section, use subsubsection dividers for different tasks to make it easier to follow: CIFAR-10
density estimation Appendix [D..I]] WikiText- 103 language modeling Appendix D.3.2 autoregres-
sive generation Appendix [D3.3] sequential image classification Appendix D.3.4 and time-series
forecasting Appendix D.3.5
D.3.1 CIFAR DENSITY ESTIMATION
This task used different backbone than the rest of our experiments. We used blocks of alternating S4 layers and position-wise feed-forward layers (in the style of Transformer blocks). Each feed- forward intermediate dimension was set to 2x the hidden size of the incoming S4 layer. Similar to Salimans al. 2017), we used UNet-style backbone consisting of B identical blocks followed by a downsampling layer. The downsampling rates were 3, 4, 4 (the chosen because the sequence consists of RGB pixels). The base model had B with starting hidden dimension 128, while the large model had = 16 with starting hidden dimension 192.
We experimented with both the mixture of logistics from (Salimans et al. 2017) as well as a simpler 256-way categorical loss. We found they were pretty close and ended up using the simpler softmax loss along with using input embeddings.
We used the LAMB optimizer with learning rate 0.005. The base model had no dropout, while the large model had dropout 0.1 before the linear layers inside the S4 and FF blocks.
D.3.2 WIKITEXT- 103 LANGUAGE MODELING
The RNN baselines included in Table 7 are the AWD-QRNN (Merity al., 2018), an efficient linear gated RNN, and the LSTM + Cache + Hebbian + MbPA (Rac et al. 2018), the best performing pure RNN in the literature. The CNN baselines are the CNN with GLU activations (Dauphin et al. 2017). the TrellisNet (Bai et al. 2019). Dynamic Convolutions (Wu et al.) 2019). and TalK
Convolutions (Lioutas & Guo 2020)
The Transformer baseline (Baevski Auli 2018). which uses Adaptive Inputs with a tied Adaptive Softmax. This model is standard high-performing Transformer baseline on this benchmark, used for example by Lioutas & Guo (2020) and many more.
the same Transformer backbone in (Baevski & Auli 2018). The model consists of 16 blocks of S4 layers alternated with position-wiss feedforward layers, with a feature dimension of 1024. Because our S4 layer has around 1/4 the number of parameters self-attention layer
Our S4 model
23

---

Published as a conference paper at ICLR 2022
with the same dimension, made two modifications to match the parameter count better: (i) used a GLU activation after the S4 linear layer (Section (3.4 (ii) we used two S4 layers per block. Blocks use Layer Normalization in the pre-norm position. The embedding and softmax layers were the Adaptive Embedding from (Baevski Auli 2018) with standard cutoffs 20000, 40000, 200000.
Evaluation was performed similarly to the basic setting in (Baevski & Auli 2018). Table 5, which reported in (Baevski & Auli 2018) that include more context at training and evaluation time and improves the score. Because such evaluation protocols are orthogonal to the basic model, we do not consider them and report the base score from (Baevski & Auli 2018) Table
involves sliding non-overlapping windows of width 1024 tokens. Other settings
Instead of SGD+Momentum with multiple cosine learning rate annealing cycles, our S4 model was trained with the simpler AdamW optimizer with single cosine learning rate cycle with a maximum of 800000 steps. The initial learning rate was set to 0.0005. We used 8 A100 GPUs with a batch size of 8 per gpu and context size 1024. We used no gradient clipping and a weight decay of 0.1. Unlike (Baevski & Auli, 2018) which specified different dropout rates for different parameters, constant dropout rate of 0.25 throughout the network, including before every linear layer and on the residual branches.
used a
D.3.3 AUTOREGRESSIVE GENERATION SPEED
Protocol. To account for different model sizes and memory requirements for each method, benchmark generation speed by throughput, measured in images per second (Table or tokens per second (Table 77. Each model generates images single A100 GPU, maximizing batch size to fit in memory. (For CIFAR-10 generation we limited memory to 16Gb, to be more comparable to the Transformer and Linear Transformer results reported from (Katharopoulos et al. .020
The Transformer and Linear Transformer baselines reported in Table 6 are the results reported directly from Katharopoulos et al. 22200. Note that the Transformer number is the one in their Appendix, which implements the optimized cached implementation of self-attention.
Baselines.
For all other baseline models, we used open source implementations of the models to benchmark generation speed. For the PixeICNN++, we used the fast cached version by Ramachandran al.] (2017), which sped up generation by orders of magnitude from the naive implementation. This code was only available in TensorFlow, which may have slight differences compared to the rest of the baselines which were implemented in PyTorch.
We were unable to run the Sparse Transformer (Child et al. 2019) model due to issues with their custom CUDA implementation of the sparse attention kernel, which we were unable to resolve.
The Transformer baseline from Table was run using a modified GPT-2 backbone from the Hug- gingFace repository, configured to recreate the architecture reported in (Baevski & Auli 2018). These numbers actually slightly favorable the baseline, embedding or softmax layers, whereas the number reported for $4 is the full model.
did not include the timing of the
D.3.4 PIXEL-LEVEL SEQUENTIAL IMAGE CLASSIFICATIN
Our models were trained with the Adam optimizer for up to 200 epochs. Hyperparameters for the CIFAR-10 model is reported in Table
For our comparisons against ResNet-18, the main differences between the base models that S4 uses LayerNorm by default while ResNet uses BatchNorm. The last ablation in Section [4.3 swaps the normalization type, using BatchNorm for S4 and LayerNorm for ResNet, to ablate this architectural difference. The experiments with augmentation take the base model and train with mild data augmentation: horizontal flips and random crops (with symmetric padding).
D.3.5
TIME SERIES FORECASTING COMPARED TO INFORMER
We include a simple figure (Fig. contrasting the architecture of S4 against that of the Informer (Zhou et al. 2021).
In Fig. the goal is forecast contiguous range of future predictions (Green, length F) given range of past context (Blue, length C') ). We simply concatenate the entire context with a sequence of masks set to the length of the forecast window. This input is single sequence of length C + F that
24

---

Published as a conference paper at ICLR 2022
Table 10: (Pixel-level image classification.) Citations refer the original model; additional citation indicates work from which this baseline is reported.
Model
SMNIST
PMNIST
SCIFAR
Transformer VVawwani et al
2017
Trinh et al. 2018
98.9
97.9
62.2
CKConv [Romero et al. 2021
99.32
98.54
63.74
TrellisNet (Bai cl 2019)
99.20
98.13
73.42
TCN (Bai et al. 2018
99.0
97.2
LSTM (Hochreiter Schmidhuber
1997.
Gu
2020b
95.11
63.01
L-LSTM Trinh of al. 2018
98.4
95.2
72.2
Dilated GRU CChang el al.
2017
99.0
94.6
Dilated RNN (Chang al.
2017
98.0
96.1
IndRNN (Li ct al. (2018)
99.0
96.0
expRNN Lezcano-Casado & Martinez-Rubi
2019
98.7
96.6
UR-LSTM
99.28
96.96
71.00
UR-GRU (Gu et al. 2020h
99.27
96.51
74.4
LMU (Voelker el al.
2019)
97.15
HiPPO-RNN (Gu ct al.
2020a
98.9
98.3
61.1
UNIcoRNN (Rusch & Mishra
2021
98.4
LMUFFT (Chikuu Elasmmh
98.49
LipschitZRNN (Erchssn
2021)
99.4
96.3
64.2
S4
99.63
98.70
91.13
Context
Forecast
Book
Outputs
Conserntel
Feauan Map
Encoder
Decoder
Forecast
Mart werl
Muliskeat
Attetiin
S4
Modeet Malistee
Deperikass" pymaadd
foll aal
Context
Inputs:
Nuct
Inpuls:
Nost as Detie
Figure 3: Comparison of $4 and specialized time-series models for forccasting tasks. (Top Left) The forecasting task involves predicting future values of a time-series given past context. (Bottom Left) We perform simple forecasting using sequence model such as S4 as black box. (Right) Informer uses an encoder-decoder architecture designed specifically for forecasting problems involving customized attention module (figure taken from Zhou al. (202
through the simple deep S4 model used throughout this work, which maps output of length + F. We then use just the last F features as the forecasted predictions.
Tables 11 and 12 contain full results on all 50 settings considered by Zhou et al. (2021 S4 sets the best results 40 out of 50 of these settings.
25

---

Published as a conference paper at ICLR 2022
Methods
S4
Informer
Informer'
Lonnraas
Reformer
LSTMa
DeeAR
ARIMA
Prophet
Metric
MSG
MAG|
MSE
MAB
MSE
MAE
MSE
MAE
MSE
MAE
MSE
MAE
MSE
MAE
MSE
MAE
24
0.061
0.191
0.098
0.247
0.092 0.246
0.103
0.259
0.223
0.389
0.114 0.272
0.107
0.280
0.108
0.284
0.115
0.275
ETTh
48
0.079
0.220
0.158
0.319
0.161
0.322
0.167 0.328
0.284
0.445
0.193 0:358
0.162
0.327
0.175
0.424
0.168
0.330
168
0,104
0.258
0.183
0.346
0.187
0.355
0.207
0.375
1.522
1.191
0.236
0.392
0.239
0.422
0.396
0.504
1.224
0.763
336
0.080
0.229
0.222
0.215
0.369
0.230
0.398
1,860
1.124
0.590
0.698
0.445
0.552
0.468
0.593
1.549
1.820
720
0.116
0.271
0.269
0.435
0.257
0.421
0.273
0.463
2.112
1.436
0.683
0.768
0.658
0.707
0.689
0.766
2.735
3.253
24
0.095
0.234
0.093
0.240
0.099
0.241
0.102
0.255
0.263
0.437
0.155
0.307
0.098
0.263
3.554
0.445
0.199
0.381
ET
48
0.191
0.346
0.155
0.314
0.159
0.317
0.169
0.348
0.458
0.545
0.190
0. 3.348
0.163
0.341
3.190
0.474
0.504
0.462
168
0.167
0,333
0,132
0.389
0.235
0.390
0.246
0.422
1,029
0.879
0.385
0.514
0.255
0.414
2,800
0.595
2.145
1.068
0.189
0.361
0,263
0,417
0.25%
0.423
0.267
0.437
1,668
1.228
0.558
0.606
0.6014
0.
2.753
0.738
2.09
2.543
0.187
0,358
0.277
0,451
0.285
0.442
0.303
0.493
2,030
1.721
0.640
0 0.681
0.429
0.580
2.878
4,664 3.355
24
0.024
0.117
0.039
0.137
0
0.160
0.065
0.202
0.095
0.228
0.121
0.233
0.091
0.243
0.090
0,206
0.120
0.290
Weather Ellaa
0,051
0.174
0,099
0303
0.066
0.194
0.078
0.220
0.249
0.390
0.305
0:411
0.219
0.362
0.179
0.306
0.133
0.305
96
0.086
0.224
P,190
0,373
0.187
0.384
0.199
0.386
0.920
0.767
0.287
0.420
0.364
0.496
0.272
0.399
0.194
0.396
0.327
10P00
0,554
0,409
0.548
0.411
0.572
1.108
1.245
0.524
0.5K4
0.795 0,948
0.462
0.55%
0.452
0.574
672
0.292
0.466
0,512
0,644
0.519
0.665
0.598
0.
1.793
1.528
1.064
0.873
2.437
1
0.639
0.697
2.747
1.174
24
0.125
0.254
0.117
0,251
0.119
0.256
0.136
0.279
0.231
0.401
0.131
0.254
0.128
0.274
0.219
0.355
0,302
0.433
48
0
0,305
0.178
0.318
0,185
0.316
0.206
0.356
0.328
0.423
0.190
0.334
0.203
0.353
0:273
0,409
0.445
0.536
168
0.19%
0.333
0,266
0.398
0,269
0.404
0,309
0.439
0.654
0.634
0.341
0.448
0.291
0.451
0.503
0.599
2.441
1.142
0.300
0.417
0.310
0.422
0.359
0.484
1.792
1,093
0.554
0.5K5
0.644
0.728
0.730
1.987
2.468
0.245
0.375
0.159
0.466
0.471
0.388
0.499
2.087
1.534
0.866
0 0.809
0.499
0.596
1.062
0.943
3.859
1.144
48
0.222
0.350
0.239
0.359
0.238
0.368
0.280
0.429
0.971
0.884
0.493
0.539
0.204
0.357
0.879
D. 0.764
0.524
0.595
168
0.331
0.421
0.447
0,503
0,442
0,514
0,454
0.529
1,671
1 1.587
0.723
0.655
0,315
0.436
1.032
0.833
2.725
1.273
336
0.328
0.422
0.489
0.528
0.501
0.552
0.514
0.563
3.528
2. 196
1.212
0.89K
0.414
0.519
1,136 0.876
2. 2.246
3.077
0,428
0.494
0.571
0.543
0.578
0.558
0.609
4.891
4.047
1,511
0.966
0.563
0.595
1.251
0.933
4.243
1.415
960
0.432
0.497
0.582
0.608
0.594
0.638
0.624
0.645
7,019
5.105
1,545
1.006
0.657
0.683
1,370
0.982
6.901
4.264
Count
22
0
0
0
0
2
0
0
Table 11: Univariate long sequence time-series forecasting results on four datasets (five cases).
Methods
S4
Informer
Informer'
LogTrans
Reformer
LSTMa
LSThct
Metric
MSE
MAF|
MSE
MAF|
MSF
MAE
MSR
MAF
MSF
MAR
MSE
MAF
MSE
MAE
24
0.525
0.542
0.577
0.549
0.620
0.577
0.686
0.604
0.991
0.754
0,650
0.624
1.293
0.901
48
0.641
0.615
0.685
0.625
0.692
0.671
0.766
0.757
1.313
0.906
0.702
0.675
1.456
0.960
168
0,980
0.779
0.931
0.752
0.947
0.797
1.002
0.846
1.824
1.138
1.212
0.867
1.997
1.214
336
1.407
0.910
1.128
0.873
1.094
0.813
1.362
0.952
2.117
1.280
1.424
0.994
2.655
1.369
720
1.162
0.842
1.215
0.896
1.241
0.917
1.397
1.291
2.415
1.520
1.960
1.322
2.143
1.380
24
0.871
0.736
0.720
0.665
0.753
0.727
0.828
0.750
1.531
1.613
1.143
0.813
2.742
1.457
ET FTTha
48
1.240
0.867
1.457
1.001
1.461
1.077
1.806
1.034
1.871
1.735
1.671
1.221
3.567
1 1.687
168
2.580
1.255
3.489
1.515
3.485
1.612
4.070
1.681
4.660
1.846
4.117
1.674
3.242
2.513
336
1,980
1.128
2.723
1.340
2.626
1.285
3.875
1.763
4.028
1.688
3,434
1.549
2.544
2.591
720
2.650
1.340
3.467
1.473
3.548
1.495
3.913
1.552
5,381
2.015
3,963
1.788
4.625
3.709
24
0.426
0.487
0.323
0.369
0.306
0.371
0.419
0.412
0.724
0.607
0.621
0.629
1.968
1.170
48
0.580
0.565
0.494
0.503
0.465
0.470
0.507
0.583
1.098
0.777
1.392
0.939
1.999
1.215
96
0.699
0.649
0.678
0.614
0.681
0.612
0.768
0.792
1.433
0.945
1.339
0,913
2.762
1.542
288
0.824
0.674
1.056
0.786
1.162
0.879
1.462
1.320
1.820
1.094
1.740
1.124
1.257
2.076
672
0.846
0.709
1.192
0.926
1.231
1.103
1,669
1.461
2.187
1.232
2.736
1.555
1.917
2.941
24
0.334
0.385
0.335
0.381
0.349
0.397
0.435
0.477
0.655
0.583
0.546
0.570
0.615
0.545
Weuther
48
0.406
0 0.444
0.395
0.459
0.386
0.433
0.426
0.495
0.729
0.666
0.829
0.677
0.660
0.589
168
0.525
0.527
0.608
0.567
0.613
0.582
0.727
0.671
1.318
0.855
1.038
0.835
0.748
0.647
336
0.531
0.539
0.702
0.620
0.707
0.634
0.754
0.670
1.930
1.167
1.657
1.059
0.782
0.683
720
0.578
0.578
0.831
0.731
0.834
0.741
0.885
0.773
2.726
1.575
1.536
1.109
0.851
0.757
48
0.255
0.352
0.344
0.393
0.334
0.399
0.355
0.418
1.404
0,999
0.486
0.572
0.369
0.445
168
0.283
0.373
0.368
0.424
0.353
0.420
0.368
0.432
1.515
1.069
0.574
0.602
0.394
0.476
BCL
336
0.292
0.382
0.381
0.431
0.381
0.439
0.373
0.439
1.601
1.104
0.886
0.795
0.419
0.477
720
0,289
0.377
0.406
0.443
0.391
0.438
0.409
0.454
2.009
1.170
1.676
1.095
0.556
0.565
960
0.299
0.387
0.460
0.548
0.492
0.550
0.477
0.589
2.141
1.387
1.591
1.128
0.605
0.599
Count
18
6
0
0
0
0
Table 12: Multivariate long sequence time-series forecasting results on four datasets (five cases).
D.4 VISUALIZATIONS
We visualize the convolutional filter K learned by S4 for the Pathfinder and CIFAR-10 tasks in Appendix D.4
26

---

Published as a conference paper at ICLR 2022
Figure 4: (Convolutional filters on Pathfinder) A random selection of filters learned by S4 in first layer (top 2 rows) and last layer (bottom 2 rows) of the best model.
27


Reviews:
[{"review_id": "XBv0jDVPyU", "reviewer": "Reviewer_ncZn", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper seems well written both regards to clarity and citations. Contentwise, the theoretical and experimental parts are interesting and relevant. The novel contribution, which is efficient computation of the discretized convolution kernel, is highly technical with details given in the appendix, but the authors did a good job at summarizing the key ideas. On the LRA benchmark, which was originally introduced to benchmark scalable transformer variants, the model sets a new high score on all problems, outperforming transformers on their home turf. The model also outperforms competitors on a raw audio dataset, including a SoTA CNN variant specialized to audio. Overall, the performance is compelling on all considered tasks, in respect of both accuracy and efficiency.\n\nI have just a few questions/remarks:\n\n1) The authors mention in Section 2.2 that linear SSMs can perform poorly due to vanishing/exploding gradients. The given HiPPO matrix is given by the sum of a normal and a low-rank (i.e. NPLR) matrix, and in particular it is not unitary. I wonder then how it gets around the aforementioned problem? I expect the answer can be found by looking up the cited papers, but if there is a short answer to this, it might be good to include it in the discussion.\n2) Related to the previous question, it turns out the authors only use the HiPPO matrix as an initialization (although as a very sensible one), but then the algorithm is free to learn any NPLR matrix. Is it true in general for NPLR matrices that they result in well-behaved SSMs in the previous sense, or does that only hold for a neighborhood around the HiPPO matrix? Additionally, I would be interested in what happens (e.g. how performance changes) when a) the HiPPO matrix  $\\mathbf{A}$ is fixed throughout the training, and only the parameters $\\mathbf{B}, \\mathbf{C} \\in \\mathbb{R}^N$ are trainable, and b) if the model is initialized from a random NPLR matrix, but not the HiPPO?\n3) Is the model able to handle multidimensional input signals $\\mathbf{u}(t) \\in \\mathbb{R}^d$?\n4) In Section 3.4, it is mentioned that a multidimensional feature map of size H is created by defining H independent copies of S3. I am wondering if it would be more efficient if these copies shared the input-to-state ($\\mathbf{B}$) and state-to-state ($\\mathbf{A}$) mappings, and only differed in the state-to-feature mapping ($\\mathbf{C}$)? Do the authors expect this would negatively affect the results? Somehow it seems wasteful to me for each feature-coordinate $y^i(t)$ to have its own separate $N$-dimensional state representation, instead of sharing a common one (perhaps with a larger state size $N^\\prime >> N$).\n5) The S3 itself seems to be a linear model, which made me wonder where the nonlinearities are introduced into the deep model? Perhaps, is there an activation placed on the feature map after linearly mixing H independent copies of S3?\n6) In Section 4.3 paragraph _Irregularly sampled data_, it seems what the authors really mean is a resolution change. As far as I know, irregular sampling means the data is sampled on a highly non-uniform time grid, which means that in eq.(3) the step size $\\Delta_t$ would become time dependent. It looks like this might make the computation of the convolution kernel in eq.(5) a bit problematic (the discretized $\\bar{\\mathbf{A}}_{\\Delta_t}$ matrices might not commute with eachother).", "summary_of_the_paper": "The paper extends previous work on linear state space models (SSM), where the state transition matrix is fixed to be a highly structured _HiPPO_ matrix, which has provably beneficial properties for memorizing long-term information from continuous-time signals. The main contribution of the paper is regarding the computational aspect of the model. Namely, a novel approach is proposed for computing the convolutional kernel associated to the discretized SSM unrolled over time. This is done by first showing that the state transition matrix can be decomposed as the sum of a normal (i.e. diagonalizable with orthonormal eigenvectors) and a low-rank (rank 1 or 2) matrix, and this representation is then combined with techniques from numerical linear algebra to reduce the problem to computing the Cauchy kernel (i.e. a well-studied problem). The resulting structured state space model (S3) is then placed into a deep neural network setting, and extensive experiments are carried out on various tasks, such as: 1) Long Range Arena, a benchmark collection for scalable transformers, 2) raw speech classification (length-16k audio signals), 3) generative modelling on CIFAR-10 and WikiText-103, 4) sequential image classification on sMNIST, pMNIST and sCIFAR. Overall, the model seems to perform very well on each task either performing close to SoTA or setting a new high score.", "summary_of_the_review": "This is a solid paper with a novel technical contribution that utilizes nontrivial insights for efficient matrix computations, and with a strong experiments section. The experiments demonstrate not only that the model can be an efficient alternative to transformers on tasks requiring long-range reasoning, but that it also shows promise as a generic sequence model that can be applied across a broad range of tasks. Overall, there are not many drawbacks of the paper for me, other than some unanswered questions in my mind.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Njj9Mravgr", "reviewer": "Reviewer_Bxrv", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper presents a strong and clear theory for the proposed reparameterization. Its advantages over the prior work linear state-space layers (LSSL) are explained very well theoretically and presented empirically. \nThe convolutional kernel defined by LSSL is a convolutional interpretation to unrolling SSMs over time, granting parallelization to SSMs during training. The proposed model S3 resolves the bottleneck in this formulation by reducing it to a Cauchy kernel, resulting in a significant improvement in the space and computation complexity. \nThe model is also compared against the state-of-the-art models in a diverse set of benchmarks. Especially the performance in the Long Range Arena (LRA) benchmark is superior. \n\nI have the following questions for the authors.\n- The base LSSL model is not included in most of the benchmarks, particularly in the LRA. Is it due to the scalability problem of LSSL? \n- Is it expected to get improved results over the LSSL (see Tab. 4 and 5)? Does the proposed reparameterization guarantee a “more” optimal state matrix A? \n- The diagonal matrix $\\Lambda$ and the vectors $p$ and $q$ are trainable parameters that construct the state matrix A which is initialized to be a HiPPO matrix (Section 3.4). Do these trainable parameters preserve the structure of the matrix A to be HiPPO during and after training? Is it necessary? \n- In autoregressive generation tasks (i.e., the model is fed with its own prediction), the sequence models (RNNs, TCNs, etc.) are likely to suffer from error accumulation problem as the prediction horizon increases. Could the authors comment on S3’s performance in a similar task? Can S3 alleviate this problem? \n\n\n-- Post-rebuttal edit --\n\nI thank the authors for their rebuttal. I read other reviews and the author responses. It is clear that this is both empirically and theoretically a strong paper. The improvements over the baselines are substantial. Looking forward to seeing the follow-up work.", "summary_of_the_paper": "This paper presents a novel parameterization for the established state-space models (SSM), which tackles the scalability problem in linear state-space layers (LSSL) in modeling long-range dependencies for very long sequences. The proposed technique reparameterizes the structured state matrix in LSSL that allows the state to memorize the past, a key component following the continuous-time memorization theory. A complete theoretical description of this reparameterization and empirical evaluations on a diverse set of benchmarks are presented. The proposed model, S3, achieves astonishing results both in terms of performance and computational efficiency. ", "summary_of_the_review": "This is theoretically and empirically a solid paper. The evaluations show that its superior performance is not limited to modeling long-term dependencies only, but it can be a strong alternative to the established sequence models. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "-jIuhxDHIKC", "reviewer": "Reviewer_7JCQ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* The motivation for the approach is clear both from a high-level as well as mathematical viewpoint.  The authors do a good job walking through the theoretical justification and solution of the approach and proposed algorithm.\n\n* The discussion and analysis is performed on a broad domain of tasks.  As the authors discuss, many approaches today are narrowly focused around a single task/domain which limits our broader understanding of LRD modeling.  By using such a varied set of experiments, the usefulness of this approach is well highlighted.\n\n* The introduction is very well written.  The general reader may be less familiar with some of these past works and the authors do an excellent job highlighting the motivation and current status of efforts in this line of study.  Similarly the background is clear and detailed without being unnecessarily complicated or verbose.  \n\n* The performance is compelling from both an efficiency and accuracy standpoint across a number of tasks.\n\n------------------------\n\nWeaknesses:\n* Limitations and next steps aren’t explicitly discussed.  It would be helpful for the authors to include even a few sentences on this.\n\n* A known issue with approaches like transformers is the need for huge datasets (even by deep learning standards). How does the performance of this approach vary as a function of dataset size?  Does this method work in a low-data regime?\n\n* The results section is more difficult to follow, especially compared to the rest of the text.  Many of the results are included in the appendix (which is fine especially given their extensiveness).  However, some of the table references are strange.  Table 4 is given, but not directly referenced in the main text (should go with Raw Speech Classification on p.8).  Table 5 is not referenced until after Table 6.  The section \"Irregularly Sampled Data\" doesn't explicitly reference where those results can be found (i.e. Table 4).  This makes the results section a little bit unclear as the reader has to the text-to-table mapping.   Table 4 caption is a bit unclear (i.e. requires going to the text to interpret what things mean).\n\n* Figure 2 (feature visualization): the authors state that the visualization of the low-level features shows that context over a small area is being learned and the rest of the image is ignored.  However, two of the activations seem to span the entire kernel.  Similarly, one of the higher-level filters is over only a small area.  This seems to be more of a general trend (low learns mostly, but not completely localized, high learns mostly, but not completely global), but not a hard and fast rule.  Is there a hypothesis as to why only some, but not all filters in these layers follow these trends?  With the current visualization, a number of the filters appear identical (just a few solid rows at the top)- are these degenerate or is this a limitation of the visualization?  There is less “hierarchical” structural differences in the convolutional filters (Figure 5 appendix).  Could the authors explain this?\n\n* Minor: Figure 1, I missed the green K several times (spent a while looking for the green in the figure).  Perhaps altering the perceptual qualities of the font could help draw attention to these variables more.   \n", "summary_of_the_paper": "The authors propose a sequence modeling approach called the structured state space model (S3) which parameterized the SSM in a more computationally efficient manner.  This is done through decomposing the structured state matrices A into a low-rank and skew-symmetric term and expanding the SSM in frequency space and using a multipole-like evaluation.  This approach maintains much of the gains and efficiencies of past SSM approaches while being more computationally stable and efficient as demonstrated on several tasks across broad domains (speech, images, text).  The paper first steps through the theoretical motivations and justifications of this approach and then performs a number of experiments to demonstrate the competitive or superior performance on a wide range of LRD tasks.  ", "summary_of_the_review": "My recommendation is to accept this work.  This paper proposes a novel parameterization for solving SSM which provides computational efficiency and accuracy gains.  The approach is explained and justified theoretically and the performance is evaluated on a number of tasks across different domains.  The text is clear and well written.  The authors explicitly state the goal of finding general methods that work on a broad range of tasks and this approach offers a good step toward solving that problem under realistic computational constraints.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agreed this was a very strong submission: it was clearly written, was theoretically and experimentally interesting, and had excellent motivation. A clear accept. Authors: you've already indicated that you've updated the submission to respond to reviewer changes, if you could double check their comments for any recommendation you may have missed on accident that would be great! The paper will make a great contribution to the conference!!
Decision: Accept (Oral)