Simplified State Space Layers for Sequence Modeling
Authors: Jimmy T.H. Smith
Andrew Warrington
Scott Linderman
Venue: ICLR 2023 notable top 5% (2023)
Keywords: sequence models
state space
S4
RNN
transformers
long range arena

Abstract: Models using structured state space sequence (S4) layers have achieved state-of-the-art performance on long-range sequence modeling tasks. An S4 layer combines linear state space models (SSMs), the HiPPO framework, and deep learning to achieve high performance. We build on the design of the S4 layer and introduce a new state space layer, the S5 layer.  Whereas an S4 layer uses many independent single-input, single-output SSMs, the S5 layer uses one multi-input, multi-output SSM.  We establish a connection between S5 and S4, and use this to develop the initialization and parameterization used by the S5 model.  The result is a state space layer that can leverage efficient and widely implemented parallel scans, allowing S5 to match the computational efficiency of S4, while also achieving state-of-the-art performance on several long-range sequence modeling tasks.  S5 averages $87.4\%$ on the long range arena benchmark, and $98.5\%$ on the most difficult Path-X task.

Published as a conference paper at ICLR 2023
SIMPLIFIED STATE SPACE LAYERS FOR SEQUENCE MODELING
Jimmy T.H. And
Andrew Warrington'>2/ 22,3, Scott W. Linderman
"Equal contribution.
"Istitute for Computational and Mathematical Engineering, Stanford University.
WW Tsai Neurosciences Institute, Stanford University.
DDeparmeent of Statistics, Stanford University. ssmmth miith14, 4, awarring, scott. linderman @aaammaannnstttt @stanford. edu.
ABSTRACT
Models using structured state space sequence (S4) layers have achieved state-of- the-art performance long-range sequence modeling tasks. An S4 layer combines linear state space models (SSMs), the HiPPO framework, and deep learning to achieve high performance. We build on the design of the S4 layer and introduce new state space layer, the 55 layer. Whereas an S4 layer uses many independent single-input, single-output SSMs, the S5 layer uses one mulli--nppt,, multi-output SSM. We establish connection between S5 and S4, and use this to develop the initialization and parameterization used by the S5 model. The result is a state space layer that can leverage efficient and widely implemented parallel scans, allowing S5 match the computational efficiency of S4, while also achieving stae-of-the--as performance on several long-range sequence modeling tasks. S5 averages 87. .7.4% on the long range arena benchmark, and 98.5% on the most difficult Path-X task.
INTRODUCTION
Efficiently modeling long sequences is challenging problem in machine learning. Information crucial to solving tasks may be encoded jointly between observations that are thousands of timesteps apart. Specialized variants of recurrent neural networks (RNNs) (Arjovsky et al., 2016; Erichson et al., 2021; Rusch & Mishra, 2021; Chang et al., 2019), convolutional neural networks (CNNs) (Bai et al., 2018; Oord et al., 2016; Romero et al., 2022b), and transformers (Vaswani et al., 2017) have been developed to try to address this problem. In particular, many efficient transformer methods have been introduced (Choromanski et al., 2021; Katharopoulos et al., 2020; Kitaev et al., 2020; Beltagy et al., 2020; Gupta & Berant, 2020; Wang et al., 2020) to address the standard transformer's 's quadratic complexity in the sequence length. However, these more efficient transformers still perform poorly on very long-range sequence tasks (Tay et al., 2021).
Gu al. (2021a) presented alternative approach using structured state space sequence (S4) layers. An $4 layer defines nonlinear sequence-to-sequeee single-input, single-output (SISO) linear state space models (SSMs) (Gu al., 2021b), coupled together with nonlinear mixing layers. Each SSM leverages the HiPPO framework (Gu et al., 2020a) by initializing with specially constructed state matrices. Since the SSMs are linear, each layer can be equivalently implemented as convolution, which can then be applied efficiently by parallelizing across the sequence length. Multiple S4 layers can be stacked create deep sequence model. Such models have achieved significant improvements over previous methods, including on the long range arena (LRA) (Tay et al., 2021) benchmarks specifically designed stress test long-range sequence models. Extensions have shown good performance classification long movie clips (Islam & Bertasius, 2022).
transformation via bank of many independent
audio generation (Goel et al., 2022) and
We introduce a new state space layer that builds the S4 layer, the S5 layer, illustrated in Figure 1. S5 streamlines the S4 layer in two main ways. First, S5 uses one multi-input, multi-output (MIMO) SSM in place of the bank of many independent SISO SSMs in S4. Second, S5 uses an efficient and widely implemented parallel scan. This removes the need for the convolutional and frequency- domain approach used by $4, which requires non-trivial computation of the convolution kernel.

---

Published as a conference paper at ICLR 2023
A
W.l.
B
Diagonalize
A, B,
C
Parallel scan
Non- linearity
A
S5 SSM
Discretize
X.L
YEu
Wi:11
Learnable parameter
Operation
Instantiated variable
Message between layers
The computational components of an S5 layer for offline application to sequence. The S5 layer uses parallel scan a diagonalized linear SSM compute the SSM outputs Yı:L RR*H nonlinear activation function is applied to the SSM outputs to produce the layer outputs. similar diagram for S4 is included in Appendix B.
Figure I:
The resulting state space layer has the same computational complexity as S4, but operates purely recurrently and in the time domain.
We then establish mathematical relationship between S4 and S5. This connection allows us to inherit the HiPPO initialization schemes that are key to the success of S4. Unfortunately, the specific HiPPO matrix that S4 uses for initialization cannot be diagonalized in numerically stable manner for use in S5. However, in line with recent work on the DSS (Gupta et al., 2022) and S4D (Gu et al., 2022) layers, we found that diagonal approximation to the HiPPO matrix achieves comparable performance. We extend a result from Gu et al. (2022) to the MIMO setting, which justifies the diagonal approximation for use in S5, We leverage the mathematical relationship between S4 and S5 to inform several other aspects of parameterization and initialization, and we perform thorough ablation studies explore these design choices.
The final S5 layer has many desirable properties. It straightforward to implement (see Appendix A), enjoys linear complexity in the sequence length, and efficiently handle time-varying SSMs and irregularly sampled observations (which is intractable with the convolution implementation S4). S5 achieves state-of-the-art performance on a variety long-range sequence modeling tasks, with an LRA average of 87.4 87.4%, and 98.5% accuracy on the most difficult Path-X task.
BACKGROUND
We provide the necessary background in this section prior to introducing the S5 layer in Section 3,
2.1 LINEAR STATE SPACE MODELS
Continuous-time linear SSMs are the core component of both the S4 layer and the S5 layer. Given an input signal u(t) R a latent state x(t) Rº and an output signal y(t) E R linear continuous-time SSM is defined by the differential equation:
dx(() Ax(t) + Bu(t),
y(t) = Cx(t) + Du(t),
(1)
dt
and is parameterized by state matrix A E PxP an input matrix B e PxU an output matrix C RMXP and feedthrough matrix D MM For constant step size, A, the SSM can be discretized using, e.g. Fuler, bilinear or zero-order hold (ZOH) methods to define the linear recurrence
X Ax Bui
Y& = CX2 Dug:
(2)
where the discrete- time parameters are each function, specified by the discretization method, of the continuous-time parameters. See Iserles (2009) for more information on discretization methods.
2.2 PARALLELIZING LINEAR STATE SPACE MODELS WITH SCANS
We use parallel scans to efficiently compute the states of a discretized linear SSM. Given binary associative operator e (i.e. (a .
c)) and a sequence of elements [ai, a2 .., ar., the
'The full S5 implementation is available https: //gi thub. com/1 inderman' lab/55.
2

---

Published as a conference paper at ICLR 2023
scan operation (sometimes referred to as all all-prefix-sum) returns the sequence
[a1, (01 02):
(a1 02
az)]
Computing a length linear recurrence of a discretized SSM, X₁ A + Bu as in (2). is a specific example of scan operation. As discussed in Section 1.4 of Blelloch (1990), parallelizing the linear recurrence of the latent transitions in the discretized SSM above can be computed in a parallel time of (I's log L). assuming L processors, where To represents the cost of matrix-matrix multiplication. For general matrix A E RRPxP, To is ((P)). This be prohibitively expensive in deep learning settings. However, if A is diagonal matrix, the parallel time becomes ((P log T.) with I. processors and only requires (PT).) space. Finally, we note that efficient parallel scans are implemented in a work-efficient manner, thus the total computational cost of the parallel scan with a diagonal matrix is (PP)) operations. See Appendix H for more information on parallel scans.
2.3 S4: STRUCTURED STATE SPACE SEQUENCE LAYERS
The S4 layer (Gu et al., 2021a) defines nonlinear
wewetnccett-squunnce transformation, mapping from an imput sequence u1.z. E R to an output sequence u'y E R An S4 layer contains a bank of H independent single-input, single-output (SISO) SSMs with N-dimensional states. Each S4 SSM is applied to one dimension of the input sequence. This results in an independent linear transformation from each input channel to each preactivation channel. A nonlinear activation function is then applied to the preactivations. Finally, a position-wise linear mixing layer is applied to combine the independent features and produce the output sequence u'y Figure 4a in the appendix illustrates the view of the S4 layer as a bank of independent SSMs. Figure 2a shows an alternative view of S4 as one large SSM with state size HN and block-diagonal state, input and output matrices.
Each S4 SSM leverages the HiPPO framework for online function approximation (Gu et al., 2020a) by initializing the state matrices with HiPPO matrix (most often the HiPPO-LegS matrix). This was demonstrated empirically to lead to strong performance (Gu et al., 2021 and can be shown as approximating long-range dependencies with respect to an infinitely long, exponentially-decaaye measure (Gu et al., 2023). While the HiPPO-LegS matrix is not stably diagonalizable (Gu et al., 2021a), it can be represented as a normal plus low-rank (NPLR) matrix. The normal component, referred to HiPPO-N and denoted A Normall can be diagonalized. Thus, the HiPPO-LegS can be conjugated into a diagonal plus low-rank (DPLR) form, which S4 then utilizes to derive an efficient form of the convolution kernel. This motivates S4's DPLR parameterization.
Efficiently applying the S4 layer requires two separate implementations depending on context: a recurrent mode and convolution mode, For online generation, the SSM is iterated recurrently, much like other RNNs. However, when the entire sequence is available and the observations are evenly spaced, more efficient convolution mode used. This takes advantage of the ability to represent the linear recurrence as a one mee-imensional convolution between the inputs and convolution kernel for each of the SSMs. Fast Fourier transforms (FFTs) can then be applied to efficiently parallelize this application. Figure 4a in the appendix illustrates the convolution approach of the S4 layer for offline processing. We note that while parallel scans could, in principle, allow a recurrent approach to be used in offline scenarios, applying the parallel scan to all H of the N -ddimensional SSMs would in general be much more expensive than the convolution approach S4 actually uses.
The trainable parameters of each S4 layer are the H independent copies of the learnable SSM parameters and the ( parameters of the mixing layer. For each of the h E {1,..., H} S4 SSMs, given a scalar input signal u (M) e R. an S4 SSM uses an input matrix ((h) E C a DPLR parameterized transition matrix A (A) E CNNN an output matrix C(N) E CIxN and feedthrough matrix (M) E RIxl to produce a signal (*) (t) e R. To apply the S4 SSMs to discrete sequences, each continuous-time SSM discretized using constant timescale parameter (A) E R4. The learnable parameters of each SSM are the timescale parameter A (h) R_,, the continuous-time parameters (()), (()), ((,), and the DPLR matrix, parameterized by vectors ( CVV and (()), ((()) CN representing the diagonal matrix and low-rank terms respectively. For notational compactness we denote the concatenation of the H S4 SSM states at discrete time index as xt
(1:H) [(x(²)) ;..,,,".'" (x(()) and the H SSM outputs as yı [y(]] :..... y (H)
3

---

Published as a conference paper at ICLR 2023
Legend:
Layer imput, B
Learnable parameter
Operation
Iustantiated variable
Message between layers
independent $4 SSMs:
Single S4 layer
Ax..
Bu.
Emission:
Nonlinearity
Mixing layer:
V1
Cr
y'2
fly')
(18))
Next
Layer output, u'LL
(a) Internal structure of a single S4 layer (Gu ct al., 2021a) when viewed as block diock-da system.
Layer input,
S5 SSM:
Single S5
Ax,. + Bu,
Emission: Y1 = Cx
Y2
Noulinearity
Next
Layer output, u'u
(b) Internal structure a single layer.
Schematic of the internal structure of a discretized S4 layer (Gu et al., 2021a) (top) and S5 layer (bottom). Note is omitted for simplicity. We view an S4 layer as a single block-diagone SSM with latent state of size H.N, followed by nonlinearity and mixing layer to mix the independent features. (b) In contrast, the S5 layer uses dense, MIMO linear SSM with latent size P & HN.
Figure 2:
THE S5 LAYER
In this section present the S5 layer. We describe its structure, parameterization and computation, particularly focusing on how each of these differ from $4.
3.1 S5 STRUCTURE: FROM SISO TO MIMO
The S5 layer replaces the bank of SISO SSMs (or large block-diagonal system) in with multi- input, multi-output (MIMO) SSM, as in (1), with latent state size P, and input and output dimension H. The discretized version this MIMO SSM can be applied to a vector-valued input sequence un.L R to produce a vector-valued sequence of SSM outputs (or preactivations) Y1.L R Lxx using latent states X E R.P. A nonlinear activation function is then applied to produce a sequence of layer outputs u R See Figure 2b for illustration. Unlike S4, do not require an additional position-wise linear layer, since these features are already mixed. We note here that compared to the H N latent size of the block-diagonal SSM in the S4 layer, S5's latent size P can be significantly smaller, allowing for the use of efficient parallel scans, as we discuss in Section 3.3.
3.2 S5 PARAMETERIZATION DIAGONALIZED DYNAMICS
The parameterization the S5 layer's 's MIMO SSM motivated by the desire to use efficient parallel scans. As discussed in Section 2.2. diagonal state matrix is required efficiently compute the linear recurrence using a parallel scan. Thus, we diagonalize the system, writing the continuous time state matrix as A VAV-1 where A PxP denotes the diagonal matrix containing the eigenvalues
4

---

Published as a conference paper at ICLR 2023
and V CPxx corresponds to the eigenvectors. Therefore, diagonalize the continuous- time latent dynamics from (1) as
dVt)x(() AV-'x(1) + V Bu(().
(4)
dt
Defining x(t) = x(t), B V B, C CV gives a reparameterized system, dx(() dt
Ax(() + Bu(t),
y(t) = (t(( + Du(().
(5)
This is linear SSM with a diagonal state matrix. This diagonalized system can be discretized with timescale parameter A E IR., using the ZOH method give another diagonalized system with parameters
A - B-A- C=C, C, D
(6)
In practice, we use a vector learnable timescale parameters A E Rº (see Section 4.3) and restrict the feedthrough matrix D to be diagonal. The S5 layer therefore has the learnable parameters: B E C C CHXP, diag(D) E R diag(A) E CP, and A E RR..
Initialization Prior work showed that the performance of deep state space models are sensitive the initialization of the state matrix (Gu et al., 202 bb)). We discussed in Section 2.2 that state matrices must be diagonal for efficient application of parallel scans. We also discussed in Section 2.3 that the HiPPO-LegS matrix cannot be diagonalized stably, but that the HiPPO-N matrix can be. In Section 4 we connect the dynamics of $5 to S4 to suggest why initializing with HiPPO-like matrices may also work well the MIMO setting. We support this empirically, finding that diagonalizing the HiPPO-N matrix leads to good performance, and perform ablations in Appendix E to compare to other initializations. We note that DSS (Gupta et al., 2022) and S4D (Gu al., 2022) layers also found strong performance in the SISO setting by using diagonalization of the HiPPO-N matrix.
Conjugate Symmetry The complex cigenvalues of a diagonalizable matrix with real entries always occur conjugate pairs. We enforce this conjugate symmetry by using half the number of eigenvalues and latent states. This ensures real outputs and reduces the runtime and memory usage of the parallel scan by factor of two. This idea is also discussed Gu al. (2022).
3.3 S5 COMPUTATION: FULLY RECURRENT
Compared the large HN effective latent size of the block-diagonal S4 layer, the smaller latent dimension of the S5 layer (P) allows the use of efficient parallel scans when the entire sequence is available. The S5 layer can therefore be efficiently used as recurrence in the time domain for both online generation and offline processing. Parallel scans and the continuous- time parameterization also allow for efficient handling of irregularly sampled time series and other time-varying SSMs, by simply supplying a different A& matrix each step. We leverage this feature and apply S5 to irregularly sampled data in Section 6.3. In contrast, the convolution of the S4 layer requires a time invariant system and regularly spaced observations.
3.4 MATCHING THE COMPUTATIONAL EFFICIENCY OF S4 AND S5
key design desiderata for S5 matching the computational complexity of S4 for both online generation and offline recurrence. The following proposition guarantees that their complexities are of the same order if S5's latent size O(1)).
Proposition 1. Given an S4 layer with H input/output features, S5 layer with H input/output features and a latent size P = O(H) has the same order of magnitude complexity as an S4 layer in terms of both runtime and memory usage.
Proof. See Appendix C.1.
We also support this proposition with empirical comparisons in Appendix C.2.
5

---

Published as a conference paper at ICLR 2023
4 RELATIONSHIP BETWEEN S4 AND S5
We now establish a relationship between the dynamics of S5 and S4. In Section 4.1 we show that, under certain conditions, the outputs of the S5 SSM can be interpreted states computed by particular S4 system. This interpretation motivates using HiPPO initializations for S5, which diseuss in more detail in Section 4.2. In Section 4.3 discuss how the conditions required to relate the dynamics further motivate initialization and parameterization choices.
projection of the latent
4.1 DIFFERENT OUTPUT PROJECTIONS OF EQUIVALENT DYNAMICS
We compare the dynamics of S4 and S5 under some simplifying assumptions:
Assumption 1. We consider only H-dimensional to H H-dimensional sequence maps.
Assumption 2. We assume the state matrix of each S4 SSM is identical, A A(h) A CNNN
Assumption 3. We assume the timescales of each S4 SSM are identical, A (k) AER
Assumption 4. We assume that the same state matrix A is used in S5 as in S4 (also of Assumption
2). Note this also specifies the S5 latent size P N. We also assume the S5 input matrix is the
horizontal concatenation of the column input vectors used by S4: B (1) ... | B
We will discuss relaxing these assumptions shortly, but under these conditions it is straightforward to derive relationship between the dynamics of S4 and S5:
Proposition? 2. Consider an S5 layer, with state matrix A, input matrix B and some output matrix C (cf. Assumption 1); and an S4 layer, where each of the H S4 SSMs has state matrix A (cf. Assumption 2, 4) and input vector Bbb) (cf. Assumption 4). If the S4 and S5 layers are discretized with the same timescales (ef. Assumption 3), then the S5 SSM produces outputs, yk. equivalent to a linear
combination of the latent states of the H S4 SSMs, Ys Cequiv X
(1:H) where Cequiv = [ C..
Proof. See Appendix D.2.
Importantly, the S5 SSM outputs are not equal to the outputs of the block-diagonal S4 SSM. Instead they are equivalent the outputs of the block-diagonal S4 SSM with modified output matrix Cequiv Under the assumptions, however, the underlying state dynamics are equivalent. Recalling that initializing the $4 dynamics with HiPPO was key to performance (Gu al., 2021a), the relationship established in Proposition 2 motivates using HiPPO initializations for S5, as we now discuss.
4.2 DIAGONALIZABLE INITIALIZATION
Ideally, given the interpretation above, we would initialize S5 with the exact HiPPO-LegS matrix. Unfortunately, as discussed in Section 2.3, this matrix is not stably diagonalizable, as is required for the efficient parallel scans used for S5, However, Gupta et al. (2022) and Gu et al. (2022) showed empirically that removing the low rank terms and initializing with the diagonalized HiPPO-N matrix still performed well. Gu et al. (2022) offered a theoretical justification for the use of this normal approximation for single-input systems: in the limit of infinite state dimension, the linear ODE with HiPPO-N state matrix produces the same dynamics as an ODE with the HiPPO-LegS matrix. Using linearity, it is straightforward to extend this result to the multi-input system that S5 uses:
Corollary (Extension of Theorem 3 in Gu et al. (2022)). Consider ALegs E NNN A Normal
Legs
NNN, BLegs N Pegg R as defined in Appendix B.1.l. Given vector-valued inputs u(t) E ER", the ordinary differential equation
dx'(t) (t)
A
x'(t) + 2 BLegsu(t)) converges to Legs
Normal
dt
dx(() A (lagsx(1) + Bragsu(() as N N XX. dt
We include a simple proof of this extension Appendix D.3. This extension motivates the use of HiPPO-N to initialize S5's MIMO SSM. Note that S4D (the diagonal extension of S4) uses the same HiPPO-N matrix. Thus, when under the assumptions in Proposition 2, an S5 SSM in fact produces outputs that are equivalent to linear combination of the latent states produced by S4D's SSMs. Our empirical results in Section suggest that S5 initialized with the HiPPO-N matrix performs just as well as S4 initialized with the HiPPO-LegS matrix.

---

Published as a conference paper at ICLR 2023
4.3 RELAXING THE ASSUMPTIONS
We now revisit the assumptions required for Proposition 2, since they only relate constrained version of S5 to a constrained version of S4. Regarding Assumption 2, Gu al. (2021a) report that S4 models with tied state matrices can still perform well, though allowing different state matrices often yields higher performance. Likewise, requiring a single scalar timescale across all of the $4 SSMs, per Assumption 3, is restrictive. S4 typically learns different timescale parameters for each SSM (Gu et al., 2023) to capture different timescales in the data. To relax these assumptions, note that Assumption constrains S5 to have dimension P .N, and N is typically much smaller than the dimensionality of the imputs, H. Proposition established that S5 can match S4's complexity with P = O(H). By allowing for larger latent state sizes, Assumptions 2 and 3 can be relaxed, as discussed in Appendix D.4. We also discuss how this relaxation motivates a block-diagonal initialization with HiPPO-N matrices on the diagonal. Finally, to further relax the tied timescale assumptions, practice, we find improved performance by learning P different timescales (one per state). See Appendix D.5 for further discussion of this empirical finding and the ablations in Appendix E.1.
note that in
RELATED WORK
S5 is most directly related to S4 and its other extensions, which we have discussed thoroughly. However, there is prior literature that uses similar ideas those developed here. For example, prior work studied approximating nonlinear RNNs with stacks of linear RNNs connected by nonlinear layers, while also using parallel scans (Martin & Cundy, 2018). Martin & Cundy (2018) showed that several efficient RNNs, such as QRNNs (Bradbury et al., 2017) and SRUs (Lei et al., 2018), fall into class of linear surrogate RNNs that can leverage parallel scans. Kaul (2020) also used parallel scans for an approach that approximates RNNs with stacks of discrete time single-input, multi-output (SIMO) SSMs. However, S4 and $5 are the only methods to significantly outperform other comparable state-of-the-art nonlinear RNNs, transformers and convolution approaches. Our ablation study in Appendix E.2 suggests that this performance gain over prior attempts at parallelized linear RNNs is likely due to the continuous-ting parameterization and the HiPPO initialization.
EXPERIMENTS
We now compare empirically the performance of the S5 layer to the S4 layer and other baseline methods. We use the S5 layer as drop-in replacement for the S4 layer. The architecture consists of a linear input encoder, stacks of S5 layers, and a linear output decoder (Gu et al., 202 For all experiments we choose the S5 dimensions to ensure similar computational complexities as S4, following the conditions discussed in Section 3.3, as well as comparable parameter counts. The results present show that the S5 layer matches the performance and efficiency of the S4 layer. We include in the appendix further ablations, baselines and runtime comparisons.
6.1 LONG RANGE ARENA
The long range arena (LRA) benchmark (Tay et al., 2021) is suite of six sequence modeling tasks, with sequence lengths from 1,024 to over 16,000. The suite was specifically developed benchmark the performance of architectures on long-range modeling tasks (see Appendix G for more details), Table presents S5's LRA performance in comparison to other methods. S5 achieves the highest average score among methods that have linear complexity in sequence length (most notably $4, S4D, and the concurrent works: Liquid-S4 (Hasani al., 2023) and Mega-chunk (Ma et al., 2023)). Most significantly, S5 achieves the highest score among all models (including Mega (Ma et al., 2023)) on the Path-X task, which has by far the longest sequence length of the tasks in the benchmark.
6.2 RAW SPEECH CLASSIFICATION
The Speech Commands dataset (Warden, 2018) contains high-fidelity sound recordings of different human readers reciting word from a vocabulary of 35 words. The task is classify which word was spoken. We show in Table 2 that S5 outperforms the baselines, outperforms previous S4 methods and performs similarly to the concurrent Liquid-S4 method (Hasani al., 2023). As S4 and $5 methods
7

---

Published as a conference paper at ICLR 2023
Table Test accuracy on the LRA benchmark tasks (Tay al., 2021). X indicates the model did not exceed random guessing. We include an expanded table, Table 7, with full citations and enror bars in the appendix. We follow the procedure reported in Gu al. (2021a: 2022) and report means across three seeds for S4, S4D (as reported by Gu et al. (2021a; 2022)) and S5. Bold scores indicate highest performance, underlined scores indicate second placed performance. We also include the results for the concurrent methods Liquid-S4 (Hasani et al., 2023) and Mega (Ma et al., 2023). Unlike S4 methods and S3, the best Mega model retains the transformer's O(L?) complexity.
Model
ListOps
Text
Retrieval
Image (1,024)
Pathfinder (1,024)
Path-X (16,384)
Avg.
(Input length)
(2,048)
(4,096)
(4,000)
Transformer
36.37
64.27
57.46
42.44
71.40
X
53.66
Luma-256
37.25
64.57
79.29
47.38
77.72
59.37
H-Trans.-ID
49.53
78.69
63.99
46.05
68.78
X
61.41
CCNN
43.60
84.08
X
88.90
91.51
68.02
Mega (O(L2))
63.14
90.43
91.25
90.44
96.01
97.98
88.21
Mega-chunk (O(L)
58.76
90.19
90.97
85.80
94.41
93.81
85.66
S4D-LeegS
60.47
86.18
89.46
88.19
93.06
91.95
84.89
S4-LegS
59.60
86.82
90.90
88.65
94.20
96.35
86.09
Liquid-S4
62.75
89.02
91.20
89.50
94.8
96.66
87.32
S5
62.15
89.31
91.40
88.00
95.33
98.58
87.46
Table 2: Test accuracy on 35-way Speech Commands classification task (Warden, 2018). We include an expanded table, Table with error bars in the appendix. Training examples are one-second 16kHz audio waveforms. Last column indicates O-shot testing at 8kHz (constructed by naive decimation). As in Gu al. (2022), the mean across three random seeds reported. Performance for the baselines InceptionNet through to S4D-Lin are reported from Gu et al. (2022).
Model
Parameters
16kHz (16,000)
8kHz (8,000)
(Input length)
InceptionNet (Nonaka & Seita, 2021)
481K
61.24
05.18
ResNet-1 (Nonaka & Seita, 2021)
216K
77.86
08.74
XResNet-50 (Nonaka & Seita, 2021)
904K
83.01
07.72
ConvNet (Nonaka Seita, 2021)
26.2M
95.51
07.26
S4-LegS (Gu et al., 2021a)
307K
96.08
91.32
S4D-LegS (Gu et al., 2022)
306K
95.83
91.08
Liquid-S4 (Hasani et al., 2023)
224K
96.78
90.00
S5
280K
96.52
94.53
parameterized in continuous-time, these models can be applied to datasets with different sampling rates without the need for re- re-training, simply by globally re-scaling the timescale parameter A by the ratio between the new and old sampling rates. The result of applying the best S5 model trained 16kHz data, to the speech data sampled (via decimation) at 8kHz, without any additional fine-tuning, also presented in Table 2. $5 also improves this metric the baseline methods.
6.3 VARIABLE OBSERVATION INTERVAL
The final application we study here highlights how $5 can naturally handle observations received at irregular intervals. S5 does so by supplying a different A, value to the discretization at each step. We use the pendulum regression example presented by Becker et al. (2019) and Schirmer et al. (2022), illustrated Figure 3. The imput sequence a sequence of L = 50 images, each 24 x 24 pixels in size, that has been corrupted with correlated noise process and sampled at irregular intervals from continuous trajectory of duration = 100. The targets are the sine and cosine of the angle of the pendulum, which follows nonlinear dynamical system. The velocity is unobserved. We match the
8

---

Published as a conference paper at ICLR 2023
sin (0)
cos( (
Observed time point
10
20
40
50
70
80
100
Time step
Figure 3: Illustration of the pendulum regression example. Shown in the top row are the images used as input at the time points indicated. Shown on the bottom are the values of sin(0,) and cos(O:). where 0r the angle of the pendulum at time that are used as the regression targets.
Table 3: Regression MSE x10- (mean + std) and relative application speed the pendulum regression task on held-out test set. Performance for the baselines, mTAND through to CRU, are reported from Schirmer et al. (2022). We include expanded table, Table 9, and further details in the appendix. Results for CRU (our run) and S5 are across twenty seeds.
Model
Relative speed
Regression MSE (x10-3)
iTTAND (Shukla & Marlin, 2021)
12.2x
65.64 (4.05)
RKN (Becker et al., 2019)
1.9x
8.43 (0.61)
RKN-A, (Becker al., 2019)
1.9x
5.09 (0.40)
ODE-RNN (Rubanova et al., 2019)
1.0x
7,26 (0.41)
CRU (Schirmer et al., 2022)
1.0x
4.63 (1.07)
CRU (our run)
1.0x
3.94 (0.21)
S5
86x
3.41 (0.27)
architecture, parameter count and training procedure of Schirmer et al. (2022). Table 3 summarizes the results of this experiment. S5 outperforms CRU on the regression task, recovering a lower mean error. Furthermore, S5 is markedly faster than CRU on the same hardware.
6.4 PIXEL-LEVEL 1-D IMAGE CLASSIFICATION
Table 10 in Appendix F.4 shows results S5 on other common benchmarks including sequential MNIST, permuted sequential MNIST and sequential CIFAR (color). We see that S5 broadly matches the performance of S4, and outperforms range of state-of-the-art RNN-based methods.
7 CONCLUSION
We introduce the S5 layer for long-range sequence modeling. The S5 layer modifies the internal structure of the S4 layer, and replaces the frequency-domain approach used by S4 with purely recurrent, time-domain approach leveraging parallel scans. S5 achieves high performance while retaining the computational efficiency of S4. S5 also provides further opportunities. For instance, unlike the convolutional S4 methods, the parallel sean unlocks the ability to efficiently and easily process time- varying SSMs whose parameters can vary with time. Section 6.3 illustrated an example of this for sequences sampled at variable sampling rates. The concurrently developed method, Liquid-S4 (Hasani et al., 2023), uses an input-dependent bilinear dynamical system and highlights further opportunities for time-varying SSMs. The more general MIMO SSM design will also enable connections to be made with classical probabilistic state space modeling as well as more recent work parallelizing filtering and smoothing operations (Särkkä & Marria-Fermándee 2020). More broadly, we hope the simplicity and generality of the S5 layer can expand the use state space layers in deep sequence modeling and lead to formulations and extensions.

---

Published as a conference paper at ICLR 2023
ACKNOWLED AND DISCLOSURE OF FUNDING
We thank Albert Gu for his thorough and insightful feedback. We also acknowledge The Annotated S4 Blog (Rush & Karamcheti, 2022) which inspired our JAX implementation. This work was supported by grants from the Simons Collaboration on the Global Brain (SCGB 697092), the NIH BRAIN Initiative (U19NS1 13201 and ROINS113119, and the Sloan Foundation. Some the computation for this work was made possible by Stanford Data Science Microsoft Education Azure cloud credits.
REFERENCES
Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International Conference Machine Learning, pp. 1 1120-1128 PMLR, 2016.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv: /803.0177, 2018.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Trellis networks for sequence modeling. In International Conference Learning Representations, 2019.
Philipp Becker, Harit Pandya, Gregor Gehhardt, Cheng Zhao, C James Taylor, and Gerhard Neumann. Recurrent kalman networks: Factorized inference in high-dimenssionl deep feature spaces. In International Conference on Machine Learning, pp. 544-552, PMLR, 2019.
Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-documed transformer. arXiv preprint arXiv:2004.051 050000, 2020.
Guy Blelloch. Prefix sums and their applications. Technical report, Tech, rept. CMU-CS-90-190. School of Computer Science, Carnegie Mellon, 1990.
James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. In International Conference on Learning Representations, 2017.
Bo Chang, Minmin Chen, Eldad Haber, and Ed Chi. AntisymmetricRNN: A dynamical system view on recurrent neural networks. In International Conference on Learning Representations, 2019.
Shiyu Chang, Yang Zhang, Wei Han, Mo Yu, Xiaoxiao Guo, Wei Tan, Xiaodong Cui, Michael Witbrock, Mark A Hasegawa-Johnson, and Thomas S Huang. Dilated recurrent neural networks. Advances in Neural Information Processing Systems, 30, 2017.
Ricky Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances neural information processing systems, 31, 2018.
Narsimha Reddy Chilkuri and Chris Eliasmith. Parallelizing Legendre memory unit training. In International Conference on Machine Learning, pp. 1898-1909 PMLR, 2021.
Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder- approaches. Syntax, Semantics and Structure in Statistical Translation, pp. 103, 2014.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
Yann Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated comvolutional networks. In International Conference Machine Learning, pp. 933-941 PMLR, 2017.
Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau. GRU-ODE-Baree Bayes: Continuous modeling of sporadically time series, Advances in neural information processing systems, 32, 2019.
10

---

Published as a conference paper at ICLR 2023
N. Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael Mahoney. Lipschitz recurrent neural networks. In International Conference on Learning Representations, 2021.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Re. It's raw! Audio generation with state-space models. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp- 7616-7633. PMLR, 17-23 Jul 2022.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HiPPO: Recurrent memory with optimal polynomial projections. Advances in Neural Information Processing Systems, 33: 1474-1487, 2020a.
Albert Gu, Caglar Gulcehre, Thomas Paine, Matt Hoffman, and Razvan Pascanu. Improving the gating mechanism of recurrent neural networks. In International Conference Machine Learning. pp. 3800-3809. PMLR, 2020b.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in Neural Information Processing Systems, 34, 2021b.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, 2022.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Re. How to train your HIPPO: State space models with generalized orthogonal basis projections. In International Conference on Learning Representations, 2023.
Ankit Gupta and Jonathan Berant. Gmat: Global memory augmentation for transformers, 2020.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In Advances in Neural Information Processing Systems, 2022.
Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. In International Conference on Learning Representations, 2023.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8): 1735-1780, 1997.
Arieh Iserles. A first course in the numerical analysis of differential equations. 44. Cambridge university press, 2009.
Md Mohaiminul Islam and Gedas Bertasius. Long movie clip classification with state-space video models. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXXV, pp. 87-104, 2022.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers RNNs: Fast autoregressive transformers with linear attention. International Conference on Machine Learning, pp. 5156-5165. PMLR, 2020.
Shiva Kaul. Linear dynamical systems core computational primitive. Advances in Neural Information Processing Systems, 33: 1608-16820 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations, 2020.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University of Toronto, 2009.
Richard Ladner and Michael Fischer. Parallel prefix computation. Journal of the ACM (JACM), 27 (4):831-838, 1980.
11

---

Published as a conference paper at ICLR 2023
Sivaramakrishnan Lakshmivarahan and Sudarshan Dhall. Parallel computing using the prefix problem. Oxford University Press, 1994.
James Lee-Thorp. Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. FNet: Mixing tokens with
Fourier transforms. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4296-4313, 2022.
Tao Lei, Yu Zhang, Sida Wang, Hui Dai, and Yoav Artzi. Simple recurrent units for highly paral- lelizable recurrence. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 4470-4481, 2018.
Mario Lezcano-Casado and David Pertunez-Rubio Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In International Conference on Machine Learning, pp. 3794-3803. PMLR, 2019.
Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently recurrent neural network (INDRNN): Building a longer and deeper RNN. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 5457-5466, 2018.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long- range spatial dependencies with horizontal gated recurrent units. Advances in Neural Information Processing Systems, 31, 2018.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettle- moyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 34, 2021.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In International Conference on Learning Representations, 2023.
Andrew Maas, Raymond Daly, Peter Pham, Dan Huang, Andrew Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150, 2011.
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference Learning Representations, 2018.
Nikita Nangia and Samuel Bowman. ListOps: A diagnostic dataset for latent tree learning. NAACL HLT 2018, pp. 92, 2018
Naoki Nonaka and Jun Seita. In-depth benchmarking of deep neural network architectures for eeg diagnosis. In Machine Learning for Healthcare Conference, pp. 414 PMLR, 2021.
Aaron van 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: arXiv:1609.00 010.6669999999. 2016.
Dragomir Radev, Pradeep Muthukrishnan, and Vahed Qazvinian. The ACL anthology network corpus. ACL-IICNLP 2009, pp. 54, 2009.
David Romero, Robert-Jan Bruintjes, Jakub Mikolaj Tomezak, Erik Bekkers, Mark Hoogendoorn, and Jan van Gemert. Flexconv: Continuous kernel convolutions with differentiable kernel sizes. In International Conference on Learning Representations, 2021.
David Romero, David Knigge, Albert Gu, Erik Bekkers, Efstratios Gavves, Jakub Tomezak, and Mark Hoogendoorn. Towards a general purpose CNN for long range dependencies ND. arXiv preprint arXiv:2206.033: 033988 2022a.
David Romero, Anna Kuzina, Erik Bekkers, Jakub Mikolaj Tomezak, and Mark Hoogendoom. CKConv: Continuous kernel convolution for sequential data. In International Conference on Learning Representations, 2022b.
12

---

Published as a conference paper at ICLR 2023
Yulia Rubanova, Ricky Chen, and David Duvenaud. Latent ordinary differential equations for irreullary- time series. Advances in Neural Information Processing Systems, 32, 2019.
T. Konstantin Rusch and Siddhartha Mishra. Unicornn: A recurrent model for learning very long time dependencies. In International Conference on Machine Learning, pp. 9168-9178. PMLR, 2021.
Sasha Rush and Sidd Karamcheti. The Annotated S4. In Blog Track at ICLR 2022, 2022. URL https: / srush. gi ic/ ao//anotaateedsee
Simo Särkkä and Angel Transactions Automatic Control, 66(1):299-38) -009, 2020.
Temporal parallclization of Bayesian smoothers. IEEE
Mona Schirmer, Mazin Eltayeb, Stefan Lessmann, and Maja Rudolph. Modeling irregular time series with continuous recurrntt units. In International Conference on Machine Learning, pp. 19388-19405. PMLR, 2022.
Satya Narayan Shukla and Benjamin Marlin. Multi-time attention networks for irregularly sampled time series. In International Conference on Learning Representations, 2021.
Masayuki Tanaka. Weighted sigmoid gate unit for activation function of deep neural network. Pattern Recognition Letters, 135:354-2 -599,,5555-00 2020.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long Range Arena: A benchmark for efficient transformers. In International Conference on Learning Representations, 2021.
Trieu Trinh, Andrew Dai, Thang Luong, and Quoe Le. Learning longer-term dependencies in RNNs with auxiliary losses. In International Conference on Machine Learning, pp. 4965-4974. PMLR, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
Aaron Voelker, Ivana Kajié, and Chris Eliasmith. Legendre Memory Units: Continuous time representation in recurrent neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0477 047668 2020.
Pete Warden. Speech Commands: A dataset for limited- vocabulary speech recognition. arXiv preprint 88300902,, 2018.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystrônformer: A Nyström- based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 14138. NIH Public Access, 2021.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33:1 13:17283--0 11783-17293 17297, 2020.
Zhenhai Zhu and Radu Soricut. H-transformer-lee Id: Fast one-dimensional hierarchical attention for sequences. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3801-3815, 2021.
13

---

Published as a conference paper at ICLR 2023
APPENDIX FOR: SIMPLIFIED STATE SPACE LAYERS FOR SEQUENCE MODELING
CONTENTS:
Appendix A: JAX Implementation of S5 Layer.
Appendix B: S5 Layer Details.
Appendix C: Computational Efficiency of S5.
Appendix D: Relationship Between S4 and S5.
Appendix E: Ablations.
Appendix F: Supplementary Results.
Appendix G: Experiment Configurations.
Appendix H: Background on Parallel Scans for Linear Recurrences.
14

---

Published as a conference paper at ICLR 2023
A JAX IMPLEMENTATION OF S5 LAYER
Import Jax
import numpy parallel_scan jax.lax.asscciatvvv..ccom
def discretize (Lamada, a_tilde, velta):
*** Diseretize diagonalized, continucus-time Linear
Args:
Lanbda (complex64): diagonal state matrix
(2,)
B_tilde (complex64): input matrix
(2, B1
Delta (float32): diserecization step aizes
(P,)
Returns:
12
discretized Lambda_ban bar (complex64), B_bar (zomplex64)
(P,), (2
Toentity ty on.ones (Taamda Shappe (0])
Lambda_ba:
(Lambda Delta)
ba (1 / Lanbda (Lambda_kez Idensity) I.... (.... Nonel B_tilde
Lambda_bar, B_bar
binary opprratooppooooooooo felement_i, element_1)s
sinary operator for parallel scon of linear recurrence. Asaumes a diagonal natrix A.
Args:
21
element_i: tuple containing ALI and Bu_i at position element_ji: tuple containing A.j and Ba_) at position
(2,), (P,)
(2,) (P6)
Returns:
new element A_out, Ro_out )
ALi, element_i
26
A_1. Ru_] element_]
return A_1 ALA, A_ Bu_i 3u_j
29
def (_ambda_baas bar, H_bar, t_tilde, b, input_sequences
Compute the IxH output of discretized SSM veven 1xH input.
31
Args:
Lanbda_bax (complex61): discretized diagonal
natrix
(2)
33
B_bar (complex60): discretized input matrix
(P, 8]
C_tilde fcomplex61): output matrix
[C, P1
(2]oat32): feedthrough matrix
imput_sequena (flost32): input sequance of features
(L, 8]
Returns:
38
ys (Iloat32); the 5GM outpurs (65 layez preactivations! Prepare elements zequired initielize paraliel scan
(I, BI]
40
cambda_elenenst re.eet [Lambda_ bar [None, ...I, input_eequsnce. shape (0], axis=0)
41
Bu elements vmap (lanbda u: B_rer u) (iaput_ elements (Lambda_eiemente Bu_elements)
P), (L,
Compute
sequence given input sequence parallel LL, P) parallei_ opeaty., elements)
45
47
+ Compute SSM output sequence
48
ys vnap (lambda x, u: IC_tilde x al .real) (xs, input_seqqenze) ys
apply_85_layee (pazams, sequence)
52
Computes Lxil output sequence of an 95 layes given Lxil input sequence. Args:
54
params: tuple of the continuous sime 8SM paramaters
35
imput_sequence: input sequence of festures
(l, Hi
Returns:
57
The 55 layer output sequence
(I, B]
Lambda, B_tilde, C_tilde, log_Della- params
59
Lambda B_aar discretize (Lambda, B_tilde, ap.exp(log_Deita)) preactavations apply_ssm (Lmmmbdd a_bar, c_tilda, D, inpur_seqqence! return jax.nn.ge (preact
61
62
daf batch__ patoh_apply__st ayer (params, [pptt_aqquanees) Computes BxIxil output sequence of an 55 layer given Bxtxll input sequence. Args:
64
66
params: tuple of the continuous sine aSx parameters input_sequences
input feature sequerces
(8, R)
68
Returns:
Batch layez output sequences
(B, L, I))
jax. (apply_S), layer, in_axes-
(params, Input input_sequeett seu___seqeeoess
Listing 1: JAX implementation to apply a single S5 layer batch of input sequences.
15

---

Published as a conference paper at ICLR 2023
B S5 LAYER DETAILS
B.1 INITIALIZATION DETAILS
B.1.1 INITIALIZATION OF THE STATE MATRIX
Here we provide additional details to supplement the discussion initialization in Section 3.2. Gu et al. (2023) explains the ability of $4 to capture long-range dependencies when using the HiPPO- LegS matrix via decomposing the input with respect to an infinitely long, exponentially decaying measure. The HiPPO-LegS matrix and corresponding SISO input vector are defined as
(2n + (2n+ 1)k (2n+1) + 1)2/2, > k
(A nk
n+
k
0,
n
(blegs),. = (2n + 1)#
(8)
Note that in Section 4.2, the input matrix BLegs E RN used in Corollary 1 is formed by concatenating H copies of bLeg8 E RRN.
Theorem of Gu et al. (2021a) then shows that the HiPPO matrices in Gu et al. (2020a), Amrro E RNXN be represented with normal plus low-rank (NPLR) form consisting of normal matrix,
A Normal HiPPO = VAV* E N and a low-rank term
(9) for unitary V E CNNN diagonal A E CNN, and low-rank factorization P,Q E Nxr, The right hand side of this equation shows HiPPO matrices can be conjugated into a diagonal plus low-rank (DPLR) form. The HiPPO-LegS matrix can therefore be written in terms of the normal HiPPO-N matrix and low-rank term P Legs RN (Goel et al., 2022)
Normal AHIPPO A HiPPO - PQ = V - (VPA(("")
ALegs
A Legs
P PLeg
(10)
Legs
where
n> k (n 1 + 4))
A
Normal LegSat
k
(11)
(n+ + 2) 1 + 4)1/ n <
PLege
(n +
)
(12)
Our default is to set the S5 layer state matrix A A Normal E RRP. and take the eigendecomposi-
LegS
tion of this matrix to recover the initial A. We often find it beneficial to also use V and V-1 V initialize B and Č, as described below.
As mentioned in Section 4.3, we also found that performance on many tasks benefited from initializing the S5 state matrix as block-diagonal, with each block on the diagonal equal A Legs Rx where It here is less than the state dimension P. e.g. R = & when 4 blocks are used on the diagonal. We then take the eigendecomposition of this matrix to initialize A, well as B and Ć. We note that even in this case, B and Ö are still initialized in dense form and there no constraint that requires A to remain block-diagonal during learning. In the hyperparameter table in Appendix G, the J hyperparameter indicates the number of these HiPPO-N blocks used on the diagonal for initialization, where J = indicates we used the default case of initializing with single HiPPO-N matrix. We discuss the motivation for this block-diagonal initialization further in Appendix D.4.
1.1.2 INITIALIZATION OF INPUT, OUTPUT AND FEED-THROUGH MATRICES
In general, explicitly initialize the input matrix B and output matrix Č using the eigenvectors from the diagonalization of the initial state matrix. Specifically, we sample B and C and then initialize the (complex) learnable parameters B as B V-1B and Č = CV.
We initialize D E H by independently sampling each element from standard normal distribution.
16

---

Published as a conference paper at ICLR 2023
u
Stodda S SSM
SSM
Single S4 SSM
P
FFT
DPLR
Frequency domain
convolution kernel generation
A
U
Woodbury ideatity
FFT
B
FFFT
C
S4 SSM
Caucly kernel
K
IFFT
A
V
Nonlinearity
Legend:
Learnable parameter Instantiated variable
Operation
Mixing layer
Message between layers
u
(a) S4 layer (Gu al., 2021a) offline processing.
A
B
Diagonalize
A, B,
C
Parallel scan
Nou- liucarity
A
S5 SSM
Discretize
XLL
x
Wi
Learnable parameter
Operation
Instantiated variable
Message between layers
(b) SS layer offline processing, Duplicated from main text.
Figure 4: The computational components of the S4 layer (Gu et al., 2021a) (top) and the S5 layer (bottom) for offline application to a sequence. (a) The S4 layer applies an independent SSM each dimension of the input sequence u1.L E R if This requires a Cauchy kernel computation to compute the convolution kernel coefficients in the frequency domain. Convolutions computed using FFTs produce the independent SSM outputs YI:L RL*H A nonlinear activation function that includes a mixing layer is applied to the SSM outputs to produce the layer outputs. (b) (Reproduced from Figure 1) The S5 layer uses parallel scan on a diagonalized linear SSM to compute the SSM outputs YI.L E RLxH A nonlinear activation function is applied to the SSM outputs to produce the layer outputs.
1.1 INITIALIZATION OF THE TIMESCALES
Prior work (Gupta al., 2022; Gu al., 2023) found the initialization of this timescale parameter to be important. This is studied in detail in Gu et al. (2023). We sample these parameters in line with S4 and sample each element of log A E from uniform distribution on the interval [log 6 log ō mmax), where the default range is 6min = 0.001 and 6max = 0.1. The only exception is the Path-X experiment, where we initialize from 8min 0.0001 and &max 0.1 account for the longer timescales as discussed in Gu et al. (2023).
B.2 COMPARISON OF S4 AND S5 COMPUTATIONAL ELEMENTS
In Figure we illustrate a comparison of the computational details the $4 and $5 layers for efficient, parallelized offline processing.
17

---

Published as a conference paper at ICLR 2023
C COMPUTATIONAL EFFICIENCY OF S5
C.I THEORETICAL COMPUTATIONAL EFFICIENCY
Proposition 1. Given an S4 layer with H input/output features. an S5 layer with H input/ouput features and a latent size P = (H)) has the same order of magnitude complexity as an S4 layer in terms of both runtime and memory usage.
Proof. We first consider the case where the entire sequence is available and compare the S4 layer's convolution mode to the S5 layer's 's use of a parallel scan. We then consider the online generation case where each method operates recurrently.
Parallelized offline processing We consider the application of both the S4 and S5 layer to a vector-valued sequence U1:1 RLXH Note that there multiple ways to compute the convolution kernel for S4/S4D and the exact computational complexity depends on the implementation (Gu al., 2021a; 2022). However, the overall complexity of computing the S4 SSM outputs given the inputs is lower bounded by the FFTs used to convert into the frequency domain to apply the convolutions. Therefore, we can lower bound the cost of applying single SISO S4 SSM scalar input sequence as (L log L) operations and (L) space. The S4 layer then consists of H different S4 SSMs, and therefore requires O log L) operations and (( space. Finally, mixing the H activations at each timestep requires L independent matrix-vector multiplications. Thus, the S4 layer saueence-to-squennc transformation requires total of O + H log L) operations when run in the efficient convolution mode. Given H processors, these operations can be parallelized to minimum of (log H + log L) parallel time.
We now consider the S5 layer. As discussed in Section 2.2, the parallel scan requires ((P)) operations and (PL) space to perform the linear recurrence that computes the discretized states X1.L L In addition, O(PHL) operations are required for the independent matrix-vector multiplications that compute Bu RLXP and the SSM outputs Čx1.1 R Therefore, the S5 layer requires O (PHL + PL) operations, Given PH L processors, these operations can be parallelized to a minimum of Olog P log L) parallel time. Thus, we see that when the S5 state dimension P (H),, the S5 layer requires O (H² L + HL) operations compared to the S4 layer's O(H + HT. log L.) operations. Crucially, when P P= O(H), ) S4 and $5 each have parallel complexity of (log H log L) (when H processors are available). In addition, when P = OHH), the space complexity for the parallel scan is ((HL) which matches the space complexity of S4's FFTs. Both methods then also perform identical broadcasted matrix vector multiplications, and hence have the same space complexity.
Online generation For online generation, both the S4 and S5 layers are run recurrently. The S4 layer requires + operations per step due to its O(HN) DPLR matrix-vector multiplication (Gu et al., 202 and the ((H2) matrix-vector multiplication of its mixing layer. The S5 layer requires O(PH + P) operations per step due to its ((P) matrix- vector multiplication with its diagonal matrix and its O(PH) matrix-vector multiplications to compute Bu, E R&P and ČX2 E RR Thus, see the two approaches have the same per step complexity of O(H2) when P= O(H) and the individual S4 SSM state sizes N are OHH).
Thus, S4 and S5 have the same order computational complexity and memory requirements in both cases.
C.2 EMPIRICAL RUNTIME COMPARISON
Table 4 provides an empirical evaluation of the runtime performance, in terms of speed and memory, between $4, S4D and S5 across range of sequence lengths from the LRA tasks. We compared the JAX implementation of S5 to a JAX implementation of S4 and S4D, based on the JAX imple- mentation from Rush & Karamcheti (2022). For a fair comparison, we modified these existing JAX implementations of S4 and S4D to allow them both to enforce conjugate symmetry and use bidirectionality. For each task, models use bidirectionality and conjugate symmetry as reported in Gu et al. (2022). All models, except for the italicised S5 row, use the same input/output features and number of layers reported in Gu al. (2022). The S4 and S4D layers also the same S4 SSM
18

---

Published as a conference paper at ICLR 2023
Table 4: Benchmarking the runtime performance of S4, S4D and S5 three LRA tasks of varied sequence lengths using parameterizations described in Section C.2. For speeds, > 1x indicates faster than the S4D baseline. For memory, < 1x indicates less memory was used than the S4D baseline. The fifth line of each metric shows the performance of the actual S5 model used for the LRA result in Table for each task using the architecture reported Table 11.
Dataset (Input Length)
Model
Architecture
ListOps (2,048)
Text (4,096)
Path-X X (16,384)
Train step speed (1)
S4
From Gu et al. (2021a)
0.6 x
0.4 x
0.4
S4D
From Gu et al. (2022)
1.0
1.0
S5
(P=H) H) Matched to Gu et al. (2022)
0.7 X
0.9 X
1.9
S5
(P=N) Matched to Gu et al. (2022)
0.9 X
1.5 x
2.9
85
As in Table 11
1.1 x
1.0 x
4.7
Evaluation step speed (t)
From Gu et al. (2021a)
0.6
0.5
0.4
S4D
From Gu et al. (2022)
1.0 X
1.0 X
1.0
(P-H) Matched to Gu et al. (2022)
0.5 X
0.9 X
1.7
S5
(P=N) Matched to Gu et al. (2022)
0.6 X
1.2 x
1.8
55
As in Table 11
0.7 x
1.0 x
4.1 X
Memory allocation (1)
S4
From Gu et al. (2021a)
1.1 X
1.2 x
1.2
S4D
From Gu et al. (2022)
1.0 X
1.0
1.0
S5
(P-H) Matched Gu al. (2022)
0.9
1.2 X
1.1
S5
(P=N) Matched to Gu et al. (2022)
0.7 X
0.7 X
0.7
85
As in Table 11
0.7
1.0
0.9
latent size as reported in Gu et al. (2022). All methods used the same batch size and all comparisons were made using 16GB NVIDIA V100 GPU. Note we observed the JAX S4D implementation to in general be faster than the JAX S4 implementation (possibly due to this specific S4 implementation s use of the naive Cauchy kernel computation (Gu al., 2021a)). For this reason, we consider S4D as the baseline.
We consider three configurations of S5 for comparison. The first two configurations, corresponding to lines and for each metric in Table 4, show how the runtime metrics vary as S5's latent size is adjusted, with all other architecture choices equal to those of S4. In line 3 of each metric, we denote the S5 "Architecture" S5 sets the latent size equal to the number of input/output features, P H. This line empirically supports the complexity argument presented in Appendix C.I. In line 4 of each metric, we denote the S5 "Architecture" as (P=N) Matched to Gu et al. (2022)" to indicate that this configuration of S5 sets the latent size P equal to the latent size N S4 uses for each of its SISO SSMs. This line also corresponds to the constrained version of $5 that performs similarly to S4/S4D as presented in the ablation study in Table The runtime results of both of these configurations supports the claim in Section 4.3 that the latent size of S5 can be increased while maintaining S4's computational efficiency.
"(P-H) Matched to Gu et al. (2022)" to indicate that this configuration of
Finally, we include a third configuration of S5, presented in the fifth line of each metric and italicized. This configuration S5 uses the best architectural dimensions from Table and was used for the corresponding LRA results in Table 1.
Importantly, the broad takeaway from this empirical study that the runtime and memory usage of S5 and S4/S4D are broadly similar, as suggested by the complexity analysis in the main text.
19

---

Published as a conference paper at ICLR 2023
D RELATIONSHIP BETWEEN S4 AND S5
We now describe in more detail the connection between the S4 and S5 architectures. This connection allowed us to develop more performant architectures and extend theoretical results from existing work.
We break this analysis down into three parts:
1. In Section D.2 we prove Proposition 2. We exploit the linearity of the systems to identify that the latent states computed by the S5 SSM are equivalent to a linear combination of latent states computed by the H SISO S4 SSMs, and that the outputs of the S5 SSM are further linear transformation of these states. We then highlight how S4 and S5 effectively define different output matrices the block-diagonal perspective shown in Figure
2.
In Section D.3 we provide a simple extension of the proof provided by Gu et al. (2022). The original proof shows that in the SISO case, in the limit of large N, the dynamics arising from a (non-diagonalizzble) HiPPO-LegS matrix, are faithfully approximated by the (diagonalizable) normal component of the HiPPO-LegS matrix. We extend this proof to apply to the MIMO setting. This motivates initialization with the HiPPO-N matrix, which in-turn allows us use parallel scans efficiently.
3. In Section D.4 we conclude by showing that, by judicious choice of initialization of the S5 state matrix, S5 can implement multiple independent S4 systems and relax the assumptions made. We also discuss the vector of timescale parameters, which we found to improve performance.
We note that many of these results follow straightforwardly from the linearity of the recurrence.
D. ASSUMPTIONS
For these following sections will use the following assumptions, until otherwise stated: Assumption 1. We consider only H-dimensional H H-dimensional sequence maps.
Assumption 2. We assume the state matrix of each S4 SSM is identical, A (h) A E CNN
Assumption 3. We assume the timescales of each S4 SSM identical, A (h) A
Assumption 4. We assume that the same state matrix A is used in S5 in S4 (also ef. Assumption
2). Note this also specifies the S5 latent size P = =N We also assume the S5 input matrix is the horizontal concatenation of the column input vectors used by S4, B a (3 |... ((")).
D.2 DIFFERENT OUTPUT PROJECTIONS OF EQUIVALENT DYNAMICS
We provide proof of Proposition 2.
Proposition 2. Consider an $5 layer, with state matrix A, input matrix B and some output matrix C (cf. Assumption 1); and an S4 layer, where each of the H S4 SSMs has state matrix (cf. Assumption 2, 4) and input vector (h)) (cf. Assumption 4). If the S4 and S5 layers are discretized with the same timescales (ef. Assumption 3), then the S5 SSM produces outputs, Yk. equivalent to a linear (1:H) where Cequiv = [ C...
combination of the latent states of the H S4 SSMs, ys Cequiv x x
Proof. For a single $4 SSM, the discretized latent states can expressed function of the input sequence UI- E R Lx
(A)
(N) (N) U
(13)
For an S5 layer, the latent states are expressible as
X
Bu:
(14)
where we index as B A
B (1)
B
(1)
and u;
W
u
(H)
20

---

Published as a conference paper at ICLR 2023
Here make the observation:
(h)
X
h=1 X
(15)
where this result follows directly from the linearity of (13) and (14). This shows that (under the assumptions outlined above) the states of the MIMO S5 SSM are equivalent to the summation of the states across the different SISO S4 SSMs.
We can then consider the effect of the output matrix C. For S5, the output matrix is a single dense matrix
CX
(16)
We can substitute the relationship in (15) into (16) to cast the outputs of the MIMO S5 SSM in terms of the state of the H SISO S4 SSMs:
Yk C
3 h=l
X
(17)
11 (h) Σ 2==1 h=1
(18)
Denoting the vertical concatenation of the H S4 SSM state vectors (1:H) we see that the outputs of the S5 SSM are expressible as:
(1)
(H)
X;
x
Yk Cequiv (1:17) where Cequiv soquiv=[Cr =|C x
CC.
(19)
and hence are equivalent linear combination of the HN states computed by the H S4 SSMs.
This shows the outputs of the constrained S5 SSM under consideration (ef. Assumption 4) can be interpreted as a linear combination of the latent states computed by H constrained S4 SSMs with the same state matrices and timescale parameters. Note however, it does not show that the outputs of the S5 SSM directly equal the outputs of the effective diickndi S4 SSM. Indeed, they are not equal, and we can repeat this analysis for the S4 layer to concretely identify the difference. For comparison we assume that the output vector for each S4 SSM given as a row the S5 output
[ca) I... [[())
matrix, i.e. C
We can express the output of each S4 SSM as
(h) V/
(1h) (N) X
(20)
where y/ E R. We can then define the effective output matrix that operates on the entire latent space (the dashed box labelled C in Figure 2a) in as
(h)
h))
(C SC xx)) (N)
(21)
By inspecting (19) and (21), we can concretely express the difference in the equivalent output matrix used by both layers
(()
0
C(1
C(1
S1
Cequiv
= [[C |....... ... ...|....... (22)
(HH)
(
(
In S4, the effective output matrix consists of independent vectors on the leading diagonal (as is pictured in Figure 2a). In contrast, the effective output matrix used by S5 instead ties dense output matrices across the H S4 SSMs. As such, S5 can be interpreted as simply defining different projection of the H independent SISO SSMs than is used by S4. Note that both projection matrices have the same number of parameters.
Although the projection is different, the fact that the latent dynamics can still be interpreted as a linear projection of the same underlying S4 latent dynamics suggests that initializing the state dynamics in S5 with the HiPPO-LegS matrix may lead to good performance, similarly to what was observed in S4. We discuss this in the next section. We note that it not obvious whether tying the dense output matrices is any more less expressive than S4's use of single untied output vector for each SSM, and it is unlikely that one approach is universally better than the other. We also stress that one would never implement S4 using the block diagonal matrix in (22), or, implement $5 using the repeated
21

---

Published as a conference paper at ICLR 2023
matrix in (22). These matrices simply constructs for understanding the equivalence between S4 and S5.
Finally, note an alternative view: the biock-diagonal S4 with output matrix Cequiv of Proposition 2 is equivalent to a version of S4 that uses a bank of single-input, multi-output (SIMO) SSMs with tied state matrices, timescales and multi-channel output matrices C = E R#*N and sums the individual SSM outputs. See appendix of Gu et al. (2021b) for discussion of how the linear state space layer (LSSL), a predecessor of S4, can have multiple output channels.
D.3 DIAGONALIZABLE INITIALIZATION
Proposition 2 suggests that initializing with the HiPPO-LeggS matrix may yield good performance in S5, just it does in S4 (because the constrained version of S5 under consideration is effectively different linear projection of the same latent dynamics). However, the HiPPO-Legg matrix is not stably diagonalizable. Corollary allows us initialize MIMO SSMs with the diagonalizable HiPPO- N matrix to approximate the HiPPO-LegS matrix and expect the performance to be comparable.
Corollary (Extension of Theorem 3 in Gu et al. (2022)). Consider ALogs E RRNN,, A Normal
RNXN, BLegs RNNH PLegs RN as defined in Appendix B.1.l. Given vector-valued inputs
u(t) E RRH, the ordinary differential equation dx'(1) A Normalx x'(1) 2 BLeeguu(t) converges to
dt
dx(t) (t)gsx(t) + (tegsu(t) as N + 20.
Proof. Theorem in Gu et al. (2022) shows the following relationship for scalar input signals as N
00:
dx (t)
dt
A
Normal
LegS
x (1)
1 (M)) 2
(t), (t),
(23)
where the only modification we have made is introducing the (h) superscript to allow us to explicitly
index dimensions later. We define B
Lagg
....
B (H) We wish extend this the case of Blgg
vector-valued input signals.
We first recall (15), which shows that the latent states of the MIMO S5 SSM are the summation of the latent states of the H SISO S4 SSMs (to which Theorem 3 from Gu et al. (2021a) applies). Although we derived (15) discrete time, applies equally in continuous time:
x(t) = > h=1
H
x (t).
(24)
We therefore define the derivative of the S5 state
dx(()) = h=1 dt
H (1)
(25)
dt
Substituting (23) into this then yields:
dx(() dt
Σ H
A Legs
Normal
x (1)
1
B Logs (h)
(26)
h=1
2
A
Leg's
Normal
1=
x (b)
2
Σ
h=1
Lagss (t),
(27)
A Legs
x(1) + B (l)
(28)
This equivalence motivates initializing S5 state matrices with the diagonalizable HiPPO-N matrix and suggests that we can expect similar performance gains.
22

---

Published as a conference paper at ICLR 2023
D.4 RELAXING THE ASSUMPTIONS
Here we discuss how relaxing the constraint on S5's latent size from Assumption 4 helps to relax the assumptions the tied S4 SSM state matrices (Assumption 2) and timescales (Assumption 3) as well as the tied output matrices that result from Proposition 2.
We start by considering the case when the S5 SSM state matrix is block-diagonal. Consider S5
SSM with latent size JN = ( and block-diagonal A E R. RJNXJN, dense B E J dense C I # and J different timescale parameters A R As a result of the block-diagonal state matrix, this system has a latent state X₁= E R that can be partitioned into J different states (6)
X R.N. We can then partition this system into J different subsystems and discretize each subsystem with one of the A ( to get the following discretized system:
B (1)
A
B
C [........ c( ].
(29)
B (J)
where A ( be also be viewed as J independent N N-dimensional S5 SSM subsystems and the output of the overall system is the sum the output of the subsystems
E NRN " E RNNH and C() E R follows that this partitioned system can
Yk CX.
(30)
CU
(31)
It follows from Proposition 2 that the dynamics of each of these J S5 SSM subsystems can be related to the dynamics of a different S4 system from Proposition 2. Each of these $4 systems has its own bank of tied S4 SSMs (cf. Assumptions 2, 3). Importantly, each of the J S4 systems can have its own state matrix, timescale parameter and output matrix shared across the outputs of a JN dimensional S5 SSM can be equivalent the linear combination of the latent states of J different S4 systems from Proposition 2. This fact motivates the option to initialize block diikodia S5 state matrix with several HiPPO-N matrices the blocks, rather than just initializing with one larger HiPPO-N matrix. In practice we found the block-diagonal initialization to improve performance many tasks, see Appendix E.
S4 SSMs. Thus,
D.5 TIMESCALE PARAMETERIZATION
Finally, we take closer look the parameterization of the timescale parameters A. As discussed Section 4.3, $4 can learn a different timescale parameter for each S4 SSM, potentially allowing it to capture different timescales of the data. Further, the initialization of the timescales can be important (Gu et al., 2023; Gupta et al., 2022), and limiting to sampling a single initial parameter may lead to poor initialization. The discussion in the previous section motivates potentially learning different timescale parameters, one for each of the J subsystems. However, in practice, we found better performance when using P different timescale parameters, one for each of the states. On the one hand, this can be viewed simply as learning different scaling for each of the eigenvalues in the diagonalized system (see Eq. (6)). On the other hand, this could be viewed number of timescale parameters sampled at initialization, helping to combat the possibility of poor initialization. Of course, the system could learn use just single timescale by setting all of the timescales to be the same. See further discussion in the ablation study in Appendix E.
increasing the
23

---

Published as a conference paper at ICLR 2023
E ABLATIONS
We perform several ablations to empirically explore different aspects of S5.
E.
S5 LATENT SIZE, BLOCK-DIAGONAL INITIALIZATION, AND TIMESCALE PARAMETERIZATION
The discussion Section 4 and Appendix raises several interesting questions: How does S5 perform when the latent size P is restricted to be equal to the latent size N used by each of S4's SSMs? How important is the timescale parameterization discussed in Appendix D.5? How important the block-diagonal initialization? Table displays the results of an ablation study performed on the LRA tasks to get better sense of this. We consider 3 versions of S5.
The first version uses the same general architecture (e.g. number of input/output features H, number of layers, etc) as reported for the S4/S4D variants in Gu et al. (2022), sets the S5 SSM latent size P to be equal to the latent size = 64 used by each of the S4 SSMs, and uses only a single sealar timescale parameter A E IR. This is essentially the version of S5 we consider in Proposition 2. We observe that this constrained version of $5 actually performs well on most tasks, though struggles to perform comparably to the S4 basclines on Image and ListOps.
The second version of S5 is exactly the same as the first except we parameterize the timescale parameter as vector A IRN, We observe uniform improvements over the scalar timescale parameterization and this reflects our general findings when training S5.
Finally, the complexity analysis and runtime comparison in Appendix C.2 suggests the latent size of S5 can be increased while maintaining similar complexity and practical runtimes as S4. We include the unconstrained version of $5 reported in main results that uses the settings reported in the hyperparameter Table 11. These models were allowed to be parameterized with fewer input/output features H (to ensure similar parameter counts to the S4 baselines) and generally used larger latent sizes P > N. Further, we swept over the use of block-diagonal initialization or not and the number of blocks to use (where J 1 indicates no block-diagonal initialization was used). All models benefited from the Block-diagone initialization for the LRA tasks (See Table 11 )
Table Ablations on the LRA benchmark tasks (Tay et al., 2021). S4 results were taken from Gu et al. (2022; 2021a). Note that the total parameter count for all models in this table are commensurate, and hence variations performance cannot be attributed to a model having drastically more parameters.
Model
ListOps
Text Retrieval
Image (1,024)
Path Pathfinder (1,024)
Path-X
(Input length)
(2,048)
(4,096)
(4,000)
(16,384)
S4D-LegS
60.47
86.18
89.46
88.19 19
93.06
91.95
S4-LegS
59.60
86.82
90.90
88.65
94.20
96.35
$5 (P=N, 1, R
57.20
87.60
90.53
82.01
94.14
96.59
85 (P=N, J = A E I
58.65
88.12
90.76
85.04
94.53
97.49
(P: free, J21, E Rº, see Table 11)
62.15
89.31
91.40
88.00
95.33
98.58
E.2 IMPORTANCE OF HIPPO-N AND CONTINUOUS- PARAMETERIZATION
We perform further ablation study to gain insight into the differences between S5 and prior attempts at parallelized linear RNNs (discussed in Section 5) focusing on what appears to be the distinguishing features: continuous-ting parameterizations and HiPPO initializations. We compare different initializations of the state matrix: random Gaussian, random antisymmetric, and HiPPO-N. The antisymmetric initialization is interesting because prior work considered these matrices in RNNs for long-range dependencies (Chang et al., 2019), and because the HiPPO-LegS matrix can be parameterized in way related to antisymmetric matrices (Gu et al., 2021a). Moreover, compare to a setup more akin to the previous parallelized linear RNN work, we also consider a direct discrete- time parameterization of S5 that does not perform repeated discretization during training learn the
24

---

Published as a conference paper at ICLR 2023
Table 6: S5 Initialization and Parameterization Ablation Study. indicates the model did not improve over random guessing.
Model
Parameterization
Initialization
ListOps
Text
Path-X X
(Input length)
(2,048)
(4,096)
(16,384)
S5 (ablation)
Discrete
Gaussian
41.50
81.09
X
S5 (ablation)
Discrete
Antisymmetric
49.10
86.42
x
S5 (ablation)
Discrete
HiPPO-N
58.15
61.93
X
S5 (ablation)
Continuous
Gaussian
58.50
69.03
X
S5 (ablation)
Continuous
Antisymmetric
59.35
82.83
X
S5
Continuous
HiPPO-N
62.15
89.31
98.58
timescale parameter A. We present the results of this ablation study in Table 6 (along with S5). We consider three of the LRA tasks that vary in length and difficulty.
The main takeaway is that the only approach that consistently performs well all tasks, including the ability to solve Path-X, is the S5 approach that uses the continuous-time parameterization and HiPPO initialization. We also note that observed the discrete time/HiPPO-N matrix configuration to be difficult to train due to stability issues, typically requiring a much lower learning rate.
E.3 S4D INITIALIZATION ABLATIONS
Finally, Gu et al. (2022) propose several alternative diagonal matrices to the diagonalized HiPPO-N matrix, including the Snv-Inv and S4D-Lin matrices. They perform ablation the LRA tasks where they simply replace the diagonalized HiPPO-N matrix with the S4D-Inv and S4D-Lin matrices while keeping all other factors the same. We include these results in Table 7. In Table 7, also include results for a similar ablation S5 by using these matrices to initialize S5 in place of the HiPPO-N matrix while keeping all other factors constant. Both matrices perform well most tasks with the exception of the S4D-Lin matrix on Path-X. Interestingly, one of these runs reached 96.79% however the other runs did not exceed random guessing on this task. Future exploration of these and other matrices are an interesting direction for future work.
25

---

Published as a conference paper at ICLR 2023
F SUPPLEMENTARY RESULTS
We include further experimental results to supplement the results presented in the main text.
F.1 EXTENDED LRA RESULTS
Table 7:
Test accuracy on the LRA benchmark tasks (Tay al., 2021). X indicates the model did exceed random guessing. Citations refer to the original model. The results for Transformer through Performer are from (Tay al., 2021). We follow the procedure reported in Gu et al. (2021a; 2022) and report means across three seeds for $4, S4D (as reported by Gu et al. (2021a; 2022)) and S5, with the standard deviation in parenthesis.
Model
ListOps
Text
Retrieval
Image
Pathfinded
Path-X
Avg.
(Input length)
(2,048)
(4,096)
(4,000)
(1,024)
(1,024)
(16,384)
Transformer (Vaswani et al., 2017)
36.37
64.27
57,46
42.44
71.40
X
53.66
Reformer (Kitaev et al., 2020)
37.27
56.10
53.40
38.07
68.50
50.56
BigBird (Zaheer et al., 2020)
36.05
64.02
59.29
40.83
74.87
X
54.17
Linear Trans. (Katharopoulos al., 2020)
16.13
65.90
53.09
42.34
75.30
50.46
Performer (Choromanski et al., 2021)
18.01
65.40
53.82
42.77
77.05
X
51.18
FNet (Lee-Thhrpp al., 2022)
35.33
65.11
59.61
38.67
77.80
X
54.42
Nyströmformer (Xiong et al., 2021)
37.15
65,52
79.56
41.58
70.94
X
57.46
Luna-256 (Ma et al., 2021)
37.25
64.57
79.29
47.38
77.72
X
59.37
H-Transformer- 1D (Zhu & Soricut, 2021)
49.53
78.69
63.99
46.05
68.78
X
61.41
CCNN (Romero al., 2022a)
43.60
84.08
88.90
91.51
68.02
Mega (O)L))) (Ma et al., 2023)
63.14
90.43
91.25
90.44
96.01
97.98
88.21
Mega-chunk (O(L) (Ma et al., 2023)
58.76
90.19
90.97
85.80
94.41
93.81
85.66
DSS (Gupta et al., 2022)
60.60
84.80
87.80
85.70
84.60
87.80
81.88
S4D-LegS (Gu et al. 2022)
60.47 (0.34)
86.18 (0.43)
89.46 (0.14)
88.19 (0.26)
93.06 (1.24)
91.95
84.89
S4D-Inv (Gu et al., 2022)
60.18 (0.35)
87.34 (0.20)
91.09 (0.01)
87.83 (0.37)
93.78 (0.25)
92.80
85,50
S4D-Lin (Gu al., 2022)
60.52 (0.51)
86.97 (0.23)
(0.99))
87.93 (0.34)
93.96 (0.60)
X
78.39
S4-FouT (Gu et al., 2023)
57.88 (1.90)
86.34 (0.31)
89.66 (0.88)
89.07 (0.19)
94.46 (0.24)
X
77.90
S4-LegS/FouT (Gu et al., 2023)
60.45 (0.75)
86.78 (0.26)
(0.28)
89.00 (0.26)
94.44 (0.08)
X
78.50
S4-LegS (Gu et al., 2021a)
59.60 (0.07)
86.82 (0.13)
90.90
88.65 (0.23)
94.20 (0.25)
96.35
86.09
Liquid-S4 (Hasani et al., 2023)
62.75 (0.20)
89.02 (0.04)
91.20 (0.01)
89.50 (0.40)
94.80 (0.2)
96.66 (0.001)
87.32
I5-Inv (See Appendix E.3)
60.07 (0.26)
87.77 (0.29)
91.26 (0.12)
86.41 (0.17)
93.42 (0.42)
97.54 (0.74)
86.08
S5-Lin (See Appendix E.3)
59.98 (0.53)
88.15 (0.24)
91.31 (0.24)
86.05 (0.96)
94.31 (0.36)
65.60 (27.00)
80.90
S5
62.15 (0.23)
89.31 (0.15)
91.40 (0.05)
88.00 (0.22)
95.33 (0.26)
98.58 (0.17)
87.46
26

---

Published as a conference paper at ICLR 2023
F.2 EXTENDED SPEECH RESULTS
Table 8: Speech Commands classification task (Warden, 2018). Test accuracy on 35-way keyword spotting. Training examples one-second audio waveforms sampled 16kHz, one-dimensional sequence of length 16000. Last column indicates zero-shot testing at 8kHz where examples are constructed by naive decimation. The mean across three random seeds is reported, with the standard deviation in parenthesis. Performance for the baselines InceptionNet through to S4D-Lin are reported from Gu al. (2022).
Model
Parameters
16000Hz
8000Hz
InceptionNet (Nonaka Seita, 2021)
481K
61.24 (0.69)
05.18 (0.07)
ResNet-18 18 (Nonaka & Seita, 2021)
216K
77.86 (0.24)
08.74 (0.57)
XResNet-50 (Nonaka Seita, 2021)
904K
83.01 (0.48)
07.72 (0.39)
ConvNet (Nonaka & Seita, 2021)
26.2M
95.51 (0.18)
07.26 (0.79)
S4-LegS (Gu al., 2021a)
307K
96.08 (0.15)
91.32 (0.17)
S4-FouT (Gu et al., 2023)
307K
95.27 (0.20)
91.59 (0.23)
S4e(LeS//ouT)) (Gu al., 2023)
307K
95.32 (0.10)
90.72 (0.68)
S4D-LegS (Gu et al., 2022)
306K
95.83 (0.14)
91.08 (0.16)
S4D-Inv (Gu et al., 2022)
306K
96.18 (0.27)
91.80 (0.24)
S4D-Lin (Gu al., 2022)
306K
96.25 (0.03)
91.58 (0.33)
Liquid-S4 (Hasani et al., 2023)
224K
96.78 (0.05)
90.00 (0.25)
S5
280K
96.52 (0.16)
94.53 (0.10)
27

---

Published as a conference paper at ICLR 2023
F.3 PENDULUM EXTENDED RESULTS
We also evaluate two ablations: S5-drop uses the same S5 architecture, but drops the dependence on the inter-sample interval, i.e. A, 1.0. We expect this network to perform poorly as it has no knowledge of how long has elapsed between observations. 35-append uses the same S5 architecture, but appends the integration timestep to the thirty-dimensional image encoding, prior to being input into the dense S5 input layer. Hypothetically, we expect this network perform well S5, However, to do so, requires the S5 network to learn to process time, which may be difficult, especially more complex domains. We include these ablations in the bottom partition of Table 9.
Note that the runtimes quoted for the baseline methods (runtimes marked with a *) are as reported by Schirmer al. (2022). These times are the total time for a training epoch, and hence include any time spent batching data. We re-ran the CRU using the original PyTorch code on the same hardware as we run our JAX experiments on (labelled CRU (our run)). For these experiments we used single NVIDIA GeForce RTX 2080 Ti. For these runs (CRU (our run), S5, S3-drop and S5-append) we exclude the time spent batching the data to more faithfully compare the runtimes for the models themselves. Also note that our S5 experiments will benefit from JAX compilation, but that this is not sufficient to explain the difference in runtime.
Table 9: Test MSE x x10-3 and runtimes on the pendulum regression task. Performance for the baselines, mTAND through to CRU, are reported from Schirmer et al. (2022), with mean and standard deviations across five random seeds (standard deviation in parenthesis). Accompanying citation indicates the original citation for the method. We re-ran the CRU (labelled CRU (our run)) and ran our S5 methods across twenty random seeds. We report mean and variances of the MSE error on the held-out test set, using a model selected using the validation set MSE. We refer the reader to Schirmer et al. (2022) for full description of the baselines.
Model
Runtime (sec/train epoch)
Regression MSE (x 10
mTAND (Shukla & Marlin, 2021)
3
65.64 (4.05)
RKN (Becker et al., 2019)
20
8.43 (0.61)
RKN-A, (Becker et al., 2019)
20
5.09 (0.40)
GRU (Cho et al., 2014)
12a
9.44 (1.00)
GRU-A, (Cho et al., 2014)
12a
5.44 (0.99)
Latent ODE (Chen et al., 2018)
52
15.70 (2.85)
ODE-RNN (Rubanova al., 2019)
37
7.26 (0.41)
GRU-ODE-B (De Brouwer et al., 2019)
60
9.78 (3.40)
1-CRU (Schirmer al., 2022)
29*
6.16 (0.88)
CRU (Schirmer et al., 2022)
36
4.63 (1.07)
CRU (our run)
22
3.94 (0.21)
S5
0.25
3.41 (0.27)
S5-drop (ablation)
0.25
6,68 (0.38)
S5 55-append (ablation)
0.25
4.13 (0.43)
28

---

Published as a conference paper at ICLR 2023
F.4 PIXEL-LEVEL 1-D IMAGE CLASSIFICATION RESULTS
Table 10 presents results and citations of the pixel-level I-D image classification.
Table 10: Test accuracy on Pixel-level 1-D Image classification. Citations refer to the original model; additional citation indicates work from which this baseline reported.
Model
SMNIST (784)
PSMNIST
SCIFAR
(Input length)
(784)
(1024)
Transformer (Trinh et al., 2018; Vaswani et al., 2017)
98.9
97.9
62.2
CCNN (Romero et al., 2022a)
99.72
98.84
93.08
FlexTCN (Romero et al., 2021)
99.62
98.63
80.82
CKConv (Romero al., 2022b)
99.32
98.54
63.74
TrellisNet (Bai et al., 2019)
99.20
98.13
73.42
TCN (Bai al., 2018)
99.0
97.2
LSTM (Gu et al., 2020b; Hochreiter Schmidhuber, 1997)
98.9
95.11
63.01
r-LSTM (Trinh al., 2018)
98.4
95.2
72.2
Dilated GRU (Chang et al., 2017)
99.0
94.6
Dilated RNN (Chang al., 2017)
98.0
96.1
IndRNN (Li et al., 2018)
99.0
96.0
expRNN [Lezcano-Casado Martunez-Rubio, 2019)
98.7
96.6
UR-LSTM (Gu et al., 2020b)
99.28
96.96
71.00
UR-GRU (Gu et al., 2020b)
99.27
96.51
74.4
LMU (Voelker et al., 2019)
97.15
HiPPO-RNN (Gu al., 2020a)
98.9
98.3
61.1
UNIcoRNN (Rusch & Mishra, 2021)
98.4
LMU-FFT (Chilkuri & Eliasmith, 2021)
98.49
LipschitzRNN (Erichson et al., 2021)
99.4
96.3
64.2
LSSL (Gu et al., 2021b)
99.53
98.76
84.65
$4 (Gu et al., 2022: 2021a)
99.63
98.70
91.80
S4D (Gu al., 2022)
89.92
Liquid-S4 (Hasani et al., 2023).
92.02
99.65
98.67
90.10
29

---

Published as a conference paper at ICLR 2023
G EXPERIMENT CONFIGURATIONS
In this section describe the experimental details. This includes the model architecture, general hyperparameters, specifics for each task, and information about the datasets.
G.1 DEEP SEQUENCE MODEL ARCHITECTURE
For the experiments, we use the S5 layer as drop-in replacement for the S4 layer used in the sequence model architecture of (Gu et al., 2021a). On high level, this architecture consists of linear encoder (to encode the input at each time step into H features), multiple S5 layers, a mean pooling layer, linear decoder, and a Softmax operation for the classification tasks. The mean pooling layer compresses the output of the last S5 layer, of shape [batch size, sequence length, number of features (H)], across the sequence length dimension, so that single H H-dimensional encoding is available for softmax classification.
The baseline S4 models from Gu al. (2022; 2023) apply GLU activation (Dauphin et al., 2017) function to the S4 SSM outputs. To take advantage of the fact that the S5 SSM outputs have already been mixed throughout the MIMO SSM we use what is essentially weighted sigmoid gated unit (Tanaka, 2020) (a GLU activation without an additional linear transform). Specifically, given an S5 SSM output Yı E R. and dense matrix W E R the layer output of the activation function we apply is 4n. GyL O o( GELU(y:)).
Hyperparameter options such as dropout rate, using either layer normalization or batch normalization, and using either pre-norm post-norm applied between the layers. Exceptions to the basic architecture described here are mentioned in the individual experiment sections below.
G.2 DEFAULT HYPERPARAMETERS
Table presents the main hyperparameters used for each experiment. For all experiments we ensure the number of layers and layer input/output features H are less than or equal to the number of layers and layer input/output features reported Gu al. (2022) as well as ensuring comparable parameter counts.
In general, the models for most tasks used batch normalization and pre-norm. Exceptions are noted in the individual experiment sections below.
G.2.1 OPTIMIZERS AND LEARNING RATES
We follow the general optimization approach used by S4/S4D in Gu et al. (2022). We the A optimizer (Loshchilov Hutter, 2019) with a global learning rate. However, in general, no weight decay and a smaller learning rate (the SSM learning rate) is applied to A, B, A. All experiments used cosine annealing. Exceptions to these points noted in the individual experiment sections below.
G.2.2 BIDIRECTIONALITY
We follow Gu et al. (2022) and use bidirectional models for the LRA and speech tasks. Unidirectional (causal) models were used for the pendulum, sequential and permuted MNIST for fair comparison with prior methods that used unidirectional models.
G.3 TASK SPECIFIC HYPERPARAMETERS
Here we specify any task-specific details, hyperparameter or architectural differences from the defaults outlined above.
G.3.1
LISTOPS
Weight decay and the global learning were applied to B.
30

---

Published as a conference paper at ICLR 2023
Table 11: Hyperparameters used for the reported results. Depth: number layers. H: number of input/ouiput features. P: Latent size. J: number of blocks used for the initialization of A (see Section B.I.I). Dropout: dropout rate. LR: global learning rate. SSM LR: the SSM learning rate. B: batch size. Epochs: max epochs set for the run. WD: weight decay.
Depth
H
J
Dropout
LR
SSM LR
B
Epochs
WD
ListOps
128
8
0.0
0.003
0.001
50
40
0.04
Text
6
256
192
12
0.1
0.004
0.001
50
35
0.07
Retrieval
6
128
256
16
0.0
0.002
0.001
32
20
0.05
Image
512
384
3
0.1
0.005
0.001
50
250
0.07
Pathfinder
192
8
0.05
0.005
0.0009
64
200
0.03
Path-X
128
256
16
0.0
0.002
0.0006
32
75
0.06
Speech
96
128
16
0.1
0.008
0.002
16
40
0.04
Pendulum
8
0.0
0.012
0.003
32
100
SMNIST
0.1
0.008
0.002
150
PSMNIST
128
128
2
0.15
0.004
0.001
50
150
0.01
sCIFAR
6 512
384
3
0.1
0.0045
0.001
50
250
0.07
G.3.2 TEXT
No exceptions to the defaults for this run.
G.3.3 RETRIEVAL
This document matching task requires slightly different architecture from the other experiments, as discussed Tay al. (2021). We use the same configuration S4 (Gu et al., 2021a). Each string passed through the input encoder, S5 layers, and mean pooling layers. Denoting Xı as the output for the first document and X₂ as the output for the second document, four features are created and concatenated together (Tay et al., 2021)
X = [X X2, X2, -X2.
(32)
This concatenated feature is then fed to linear decoder and softmax function as normal.
G.3.4 IMAGE
Weight decay and the global learning rate were applied to B.
G.3.5 PATHFINDER
No exceptions to the defaults for this run.
G.3.6 PATH-X
Weight decay was applied to B.
G.3.7 SPEECH COMMANDS
No weight decay and the SSM learning rate were applied to for this run.
G.3.8 PENDULUM REGRESSION
We use the same encoder-decoder architecture as Schirmer et al. (2022). The encoder has layers: convolution, ReLU, max pool, convolution, ReLU, max pool, dense, ReLU, dense. The first con- volution layer has twelve features. 5 x kernel, and a padding of (2, 2). The second convolution layer has twelve features, window size of 2 x and a stride of 2. The dense layer has thirty hidden units. The linear readout layer outputs H 30 features. This is chosen to match the encoding size in Schirmer al. (2022),
3 kernel, stride of 2, and a padding of (1, 1). Both
pools use
31

---

Published as a conference paper at ICLR 2023
and is used for all layers. Separate mean and unconstrained variance decoders used, defined as a one-layer MLP with a hidden size of thirty. An elu+1 activation function used to constrain the variance to be positive.
Layer normalization and post-norm were used for this task.
For the timings presented in Table 3 and 9 batch size of 50, instead of the batch size of 32 used during training, to match the batch sizes reported by the baselines.
G.3.9 SEQUENTIAL MNIST
No exceptions to the defaults for this run.
G.3.10
PERMUTED SEQUENTIAL MNIST
Weight decay and the global learning rate were applied to B. Post-norm was used for this task.
G.3.11
SEQUENTIAL CIFAR
We trained a model with the exact hyperparameter settings as used for the LRA-IMAGE (grayscale sequential CIFAR) task with further tuning.
G.4 DATASET DETAILS
We provide more context and details for each of the LRA (Tay et al., 2021) and Speech Com- mands (Warden, 2018) datasets we consider. Note that we follow the same data pre-processing steps Gu et al. (2021a), which we also include here for completeness.
ListOps: A lengthened version of the dataset presented by Nangia & Bowman (2018). Given a nested set of mathematical operations (such as min and max) and integer operands in the range zero to nine, expressed in prefix notation with brackets, compute the integer result of the mathematical expression, e.g. [MAX29[MI44]]] - Characters are encoded one-hot vectors, with 17 unique values possible (opening brackets and operators are grouped into a single token). The sequences are of unequal length, and hence the end of shorter sequences is padded with fixed indicator value, padded to maximum length 2, 000. A reserved end-of- -oddoo--eeuence token is appended. There are 10 different classes, representing the integer result of the expression. There are 96, 000 training sequences, 2, 000 validation sequences, and 2, 000 test sequences. No normalization is applied.
Text: Based off of the iMDB sentiment dataset presented by Maas et al. (2011). Given movie review, where characters whether the movie review positive or negative. Characters are encoded as one-hot vectors, with 129 unique values possible. Sequences are of unequal length, and are padded to a maximum length of 4, 096. There are two different classes, representing positive and negative sentiment. There are 25, 000 training examples and 25, 000 test examples. No validation set is provided. No normalization is applied.
encoded as sequence of integer tokens, classify
Retrieval: Based off of the ACL Anthology network corpus presented by Radev et al. (2009). Given two textual citations, where characters are encoded as a sequence integer tokens, classify whether the two citations are equivalent. The citations must be compressed separately, before heing passed into a final classifier layer. This is to evaluate how effectively the network can represent the text. The decoder head then uses the encoded representation to complete the task. Characters are encoded into a one-hot vector with 97 unique values. Two paired sequences may be of unequal length, with maximum sequence length of 4, 000. There are two different classes, representing whether the citations are equivalent or not. There are 147, 086 training pairs, 18, 090 validation pairs, and 17, 437 test pairs. No normalization is applied.
Uses the CIFAR- 10 dataset presented by Krizhevsky (2009). Given a 32 x 32 grayscale CIFAR- image as dieeemmnni raster scan, classify the image into one of ten classes. Sequences are of equal length (1, 024). There are ten different classes. There are 45, 000 training examples, 000 validation examples, and 10, 000 test examples. RGB
Image:
32

---

Published as a conference paper at ICLR 2023
pixel values are converted to grayscale intensities, which are then normalized to have mean and unit variance (across the entire dataset).
Based off of the Pathfinder challenge introduced by Linsley et al. (2018). A 32 x 32 grayscale image image shows a start and an end point as a small circle. There are number of dashed lines on the image. The task to classify whether there a dashed line (or path) joining the start and end point. There are two different classes, indicating whether there is a valid path or not. Sequences are all of the same length (1, 024). There are 160, 000 training examples, 20, 000 validation examples, and 20, 000 test examples. The data is normalized to be in the range [-1,1].
Path- X: An "extreme" version of the Pathfinder challenge. Instead, the images are 128 x 128 pixels, resulting in sequences that are a factor of sixteen times longer. Otherwise identical to the Pathfinder challenge.
Speech Commands: Based on the dataset released by Warden (2018). Readers recite one of 35 words. The task is then to classify which of the 35 words was spoken from 16kHz one-dimensional audio recording. There are 35 different classes, each representing one of the words in the vocabulary. Sequences are all of the same length (16, 000). There are 24, 482 training examples, 5, 246 validation examples, and 5, 247 test examples. Data is normalized to be zero mean and have standard deviation of 0.2.
Speech Commands x 0.5: Temporally sub-sampled version of Speech Commands, where the validation and test datasets only are sub-sampled by fac- tor of 1/0.5, and are therefore shortened to length 8, 000. No subsequent padding is applied. The training dataset is not subsampled.
Sequential ial MNIST: (sMNIST) 10-way digit classification from a 28 x grayscale image of handwritten digit, where the input image is flattened into a 784-length scalar sequence.
Permuted Sequential MNIST: (pSMNIST) 10-way digit classification from a 28 x 28 grayscale image of a handwritten digit, where the input image is flattened into a 784-length scalar sequence. This sequence is then permuted using a fixed order.
Sequential CIFAR: (SCIFAR): 10-way image classification using the CIFAR-10 dataset. Identical to image, except that full colour images are input as a 1, 024-lehhtt input sequence, where each input is an (R, triple.
Pendulum Regress ion: Reproduced from Becker et al. (2019) and Schirmer al. (2022). The input sequence is a 24 x 24 grayscale rendering of a pendulum, driven by a random torque process. The images pixels are corrupted by noise process that is correlated in time. The pendulum is simulated for 100 timesteps, and 50 frames are irregularly sampled without replacement from the simulation. The objective the angle of the pendulum. A train/ mnin/validation/test split of 2, 000/ 1, 000/L, 000 is used.
estimate the sine and cosine of
33

---

Published as a conference paper at ICLR 2023
H BACKGROUND ON PARALLEL SCANS FOR LINEAR RECURRENCES
For the interested reader, this section provides more background on using a parallel scan for linear recurrence, as well as a simple example to illustrate how can compute the recurrence in parallel. The parallelization of scan operations has been well studied (Ladner & Fischer, 1980; Lakshmivarahan & Dhall, 1994; Blelloch, 1990), and many standard scientific computing libraries contain efficient implementations. We note the linear recurrence we consider here is specific instance of the more general setting discussed Section 1.4 of Blelloch (1990).
Computing general parallel requires defining two objects:
The initial elements the scan will operate on.
A binary associative operator used to combine the elements.
To compute length T. linear recurrence, Ar 1 + Bug, we will define the T. initial elements, C1:L, such that each element o is the tuple
C2 (Ck,as Ck,b) (A, Bua).
(33)
These CILL will be precomputed prior the scan. Having created the list of elements for the scan to operate on, we define the binary operator for the scan to use on this linear recurrence as
4
(Uia @ qi Gj.a
qi.b + (j,b)
(34)
where EK denotes an input element to the operator that could the initial elements C& or some intermediate result, @ denotes matrix-matrix multiplication, @ denotes matrix-vector multiplication and + denotes elementwise addition. We show that this operator is associative at the end of this section.
Simple example using binary operator We can illustrate how can be used to compute linear recurrence in parallel with simple example. Consider the system Ta Ar Bug, and a length L = 4 sequence of inputs U1:4. Assuming O, the desired latent states from this recurrence are:
x1 Buy
(35)
I2 ABuj Bu²
(36)
F3 A A'Bu, + AB Bug
(37)
T4 A 3 Bu, + A Bug A'Bug + ABu3 Bug
(38)
We first note that can be used to compute this recurrence sequentially. We can initialize the scan elements C1:4 in (33), and then sequentially scan over these elements to compute the output
elements 8; 8i-1 G. Defining &u == (I, where is the identity matrix, we have for our example:
81
5g
C = (I, 0) A, Buy) = (A1, A0+Bu, = (A, Bu,)
(39)
82 81
C2 = (A Buy) (A, Bug)
AB + Buz)
(40)
83 52 C3 = (A², AB + Buz) (A, Bua) = (A), A A²BuI ABuz + Bug)
(41)
84 83 C4 = AA", ABB AB + Bug) (A, Bu4)
(42)
(A*, A Bu A A²Buz ABu3 Buy).
(43)
Note that the second entry of each of the output tuples, 8,,, contains the desired x; computed above. Computing the scan in this way requires four sequential steps since each 8; depends 8;-1-
Now consider how we can use this binary operator to compute the recurrence in parallel. We will label the output elements of the parallel scan F1:4 and define To O). We will first compute the even indexed elements 72 and r4, and then compute the odd indexed elements 7i and r3. We start by applying the binary operator adjacent pairs of our initial elements C1:4 compute r2 and the
34

---

Published as a conference paper at ICLR 2023
intermediate result 91. and we then repeat this process to compute r4 by applying r2 and 91:
A
Buy (A, Buz) = (A", ABu1 Bu2)
(44)
C3
Buy)
A, Bua) = (A", ABu3 Bus)
(45)
Q4
ABwi + Buz) (A
ABu3 + Bu4)
(46)
A Bu, + AuB + ABug + Bma)-
(47)
Now we will compute the odd indexed elements rı and r3, using the even indexed ro and r2, as Fh Tk-1 1 Ch:
= 33
(I, 0) (A, B (A, Bu1) A ABu, Bug) (A, Bus) (A', A'Bu, + AB Bug)-
(48)
(49)
Note that the second entry of each of the output tuples, rk,b, corresponds to the desired Z5. Inspecting the required dependencies for each application of o, we see that 72 and the intermediate result qu can be computed in parallel. Once r2 and qa computed, 71, and We have therefore reduced the number of sequential steps required from four in the sequential scan version two in the parallel scan version. This reduction in sequential steps becomes important when the sequence length is large since, given sufficient processors, the parallel time scales logarithmically with the sequence length.
all be computed in parallel.
Associativity of binary operator associative:
Finally, for completeness, we show that the binary operator is
9
qk
(gta ① Qia qie qiob + 93.6) q
(50)
(gk,a (fia O R.a): ak.a @ (gi,a @ q.b. + 93,b) + 9r.b)
(51)
((gk,a qj.a) O hi Ok.a (qi,a q.,)) + th.a @ Qj.b + (c,b)
(52)
((gk,a C Lí,a) O Dha (4k, @ (g.) @ RtD + a. @ 4z + (b.b)
(53)
q: (gk.a @ qj.ar ak.a @ qz (c.b)
(54)
(gp 9kc
(55)
35


Reviews:
[{"review_id": "aVD0STwgAK6", "reviewer": "Reviewer_Gzd2", "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": "", "summary_of_the_paper": "This paper presents an extension of structured state space sequence (S4) layers [Gu et al. 2021], which the authors refer to as S5 layers.\nWhile S4 layers use a bank of (deterministic) linear continuous state-space models with subsequent mixing layers, S5 uses a single multi-input multi-output SSM that does not require mixing layers. Furthermore, S4 parallelises the (offline) computation for the sequence of linear operations using a 1D convolution identity and an implementation in frequency domain (FFT). This has been done for computational reasons (parallelisation). However, this paper shows that S5 layers can match the computational and memory complexity (also shown empirically), by using the scan operation. \nThe authors show connections between S4 and S5, which they use for building on the initialization of S4. The authors also show via extensive ablation studies (in the supplementary material), that the initialisation as well as the continuous-time formulation with learned discretisation timescale parameters are the main drivers for these models.\n", "summary_of_the_review": "Overall I enjoyed reading the paper, it is high-quality work, in terms of method, experimentation, and the presentation.\nI did not find any major problems or concerns. \n", "strength_and_weaknesses": "This paper is very well written, precise, and easy to follow. I found the background on S4 layers incl. the diagram very useful and more accessible than the original paper, and I also think the authors made a good choice in explaining their extension S5 by contrasting them to S4. It is also great to see the easy-to-implement jax code in the appendix.\n\nThe experimental evaluation is extensive, with many and the most important baselines, evaluated on commonly used long-range-dependency benchmarks. It is especially nice to see ablation studies in Appendix E, evaluating different model sizes (making it more or less similar to S4) as well as initialisation and continuous-time (vs. direct discrete-time implementation), showing that these are the most important features. \n\nI did not find major problems or concerns. \nBut I would be interested if the relation between S4 and S5 (App. D) could be made with less restrictions. Especially assumption 2 seems a bit strict? \nI would also have liked to see how the model deals with uncertainty/unpredictability, which is often the case in forecasting (e.g. the stock market would be an extreme case). The sequential image classification tasks are a bit unrealistic/artificial. But I do realise this is not the goal of the paper, and previous work from this domain also does not evaluate in such settings. ", "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": "TMKUIXNgy1s", "reviewer": "Reviewer_2DkJ", "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": "", "summary_of_the_paper": "The paper addresses the problem of long-range sequence modeling. The authors propose the s5 layer, which is built atop and modified version of structured state-space-sequence, s4. Utilizing parallel scans, s5 improves over s4 in terms of performance and complexity. They show improved results on LRA benchmark without incurring additional complexity comparing with S4. \n\n", "summary_of_the_review": "This is a well-written, well-motivated paper. ", "strength_and_weaknesses": "The proposed layer, s5, replaces the frequency-domain approach used by S4 with a recurrent, time-domain approach bypassing non-trivial steps needed for S4. The method achieves high performance with comparable complexity of s4. \n\nThe source code link provided in the paper isn't working, The authors should publish anonymized code for the review purposes with the correct link. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made 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": "SW08q9gywby", "reviewer": "Reviewer_w99z", "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": "", "summary_of_the_paper": "The paper introduces the new state space layer S5 which combines multiple single-input single-output state space models to a single multiple-input multiple-output state space model. The proposed state space layer can be efficiently implemented by a parallel scan. S5 achieves state-of-the-art results on multiple benchmarks, including the difficult Path-X task.", "summary_of_the_review": "The paper is very well written and easy to follow. All necessary background information is concisely introduced. The authors provide extensive analyses w.r.t. ablations, runtime and memory footprints. Some ambiguities and questions still remain, whereas the reviewer hopes they will be clarified during the rebuttal period.\n*The reviewer will likely increase the score if the weaknesses and questions are clarified in sufficient detail.*\n\n### Questions\n\n - What might be the main reason or intuition why S5 has better results than S4 w.r.t. scores, besides computational efficiency/runtime?\n - Why wasn’t S4 used as a baseline in Table 3?\n - What is the hyperparameter search space of the baseline methods?\n\n---\n### Update\n\nThe response of the authors addressed our concerns. Therefore, the score has been raised (as promised).", "strength_and_weaknesses": " - [+] The paper is very well written with a concise and clearly structured background section.\n - [+] All experiments are repeated multiple times with different seeds, leading to more stable and thus scientifically valuable results.\n - [+] Detailed and comprehensive comparison with S4 and its components with helpful illustrations and descriptions.\n - [+] The reviewer highly acknowledges the inclusion of concurrent work in the comparison.\n - [-] Although all experiments are repeated multiple times, no error bars are reported. Especially in cases where multiple methods perform similarly, a corresponding deviation measure proves to be really \n       helpful for judging.\n - [-] It is not clear what the reported scores in Table 1 are. Assuming the scores to be accuracy scores, the average accuracy provides only limited insight under the assumption that the tasks are of\n       varying difficulty (i.e. different number of classes with different relative class sizes, etc.). The reviewer suggests reporting the average rank for each method instead.\n - [-] Regarding the results reported in Table 3: are the results marked with an asterisk (*) computed on 5 seeds, and all others (i.e. „CRU (our run)“ and „S5“) on 20 seeds? If yes, this might skew the\n       result. The reviewer suggests using the same test protocol for all methods in this case. If not, the reviewer suggests clarifying the use of different seeds.\n\n", "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": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%