On the Parameterization and Initialization of Diagonal State Space Models
Authors: Albert Gu
Karan Goel
Ankit Gupta
Christopher Ré
Venue: NeurIPS 2022 Accept (2022)
Keywords: Deep learning
sequence model
state space model
S4
HiPPO
diagonal state space

Abstract:   State space models (SSM) have recently been shown to be very effective as a deep learning layer as a promising alternative to sequence models such as RNNs, CNNs, or Transformers.
  The first version to show this potential was the S4 model, which is particularly effective on tasks involving long-range dependencies by using a prescribed state matrix called the HiPPO matrix.
  While this has an interpretable mathematical mechanism for modeling long dependencies,
  it also requires a custom representation and algorithm that makes the model difficult to understand and implement.
  On the other hand, a recent variant of S4 called DSS showed that restricting the state matrix to be fully diagonal can still preserve the performance of the original model when using a specific initialization based on approximating S4's matrix.
  This work seeks to systematically understand how to parameterize and initialize diagonal state space models.
  While it follows from classical results that almost all SSMs have an equivalent diagonal form, we show that the initialization is critical for performance.
  First, we explain why DSS works mathematically, as the diagonal approximation to S4 surprisingly recovers the same dynamics in the limit of infinite state dimension.
  We then systematically describe various design choices in parameterizing and computing diagonal SSMs, and perform a controlled empirical study ablating the effects of these choices.
  Our final model S4D is a simple diagonal version of S4 whose kernel computation requires just 3 lines of code and performs comparably to S4 in almost all settings, with state-of-the-art results in image, audio, and medical time-series domains, and 85\% average on the Long Range Arena benchmark.


On the Parameterization and Initialization of Diagonal State Space Models
Albert Gu', Ankit Gupta', Karan Goel', Christopher Ré
Department of Computer Science, Stanford University
$ IBM Research
(albertgu, kngg) edu, chrismreècs. wtt.arreeeststtorrrreoeaao stanford.eddedoo ankitgupta. makittuuta.iitkaapaaa iitkanpur@gmail.comm
Abstract
State space models (SSM) have recently been shown be very effective deep learning layer as promising alternative to sequence models such as RNNs, CNNs, Transformers. The first version show this potential was the S4 model, which is particularly effective on tasks involving long-range dependencies by using a prescribed state matrix called the HiPPO matrix. While this has interpretable mathematical mechanism for modeling long dependencies, introduces a custom representation and algorithm that can be difficult to implement. On the other hand, a recent variant of S4 called DSS showed that restricting the state matrix to be fully diagonal can still preserve the performance of the original model when using a specific initialization based on approximating S4's matrix. This work seeks to systematically understand how to parameterize and initialize such diagonal state space models. While follows from classical results that almost all SSMs have an equivalent diagonal form. we show that the initialization is critical for performance. We explain why DSS works mathematically, by showing that the diagonal restriction of $4's matrix surprisingly recovers the same kernel the limit of infinite state dimension. We also systematically describe various design choices in parameterizing and computing diagonal SSMs, and perform a controlled empirical study ablating the effects of these choices. Our final model S4D is a simple diagonal version of S4 whose kernel computation requires just lines code and performs comparably to S4 in almost all settings, with state-of-the-an art results for image, audio, and medical time-series domains, and averaging 85% on the Long Range Arena benchmark.
Introduction
A core class of models in modern deep learning are sequence models, which are parameterized mappings operating on arbitrary sequences of inputs. Recent approaches based on state space models (SSMss) have outperformed traditional deep sequence models such as recurrent neural networks (RNNs), convolutional neural networks (CNNs), and Transformers, in both computational efficiency and modeling ability. In particular, the S4 model displayed strong results on a range of sequence modeling tasks, especially on long sequences [9] Its ability to model from being defined with a particular state matrix called the "HiPPO matrix" [6], which allows S4 to be viewed convolutional model that decomposes input
dependencies arises
orthogonal system of smooth
basis functions 10
However, beyond its theoretical interpretation, actually computing S4 as a deep learning model requires a sophisticated algorithm with many linear algebraic techniques that are difficult to understand and implement. These techniques were necessitated by parameterizing its state matrix as a diagonal plus low-rank (DPLR) matrix, which is necessary to capture HiPPO matrices. A natural question is
36th Conference on Neural Information Processing Systems (NeurIPS 2022).

---

y(1)
C
A
x(()
-1-
0.0
0.5
1.0
1.5
2.0
def s4d_kermel (A, B, C, dt, L):
B
vandermonde (arange (L) Mone I d'd'')
return sun (vanderoonde 8" (exp 11/A, -1)
S4D Recurrent View
S4D Convolution Kernel
Figure 1: S4D is a diagonal SSM which inherits the strengths of S4 while being much simpler. (Left) The diagonal structure allows it to be viewed as a collection of I-dimensional SSMs. (Right) As a convolutional model, S4D has simple interpretable convolution kemeel which can be implemented in two lines of code. Colors denote independent 1-D SSMs; purple denotes trainable parameters.
whether simplifications of this parameterization and algorithm are possible. In particular, removing the low-rank term would result in diagonal state space model (diagonal SSM) that is dramatically simpler to implement and understand.
Although it is known that almost all SSMs have an equivalent diagonal and therefore (complex) diagonal SSMs are fully expressive algebraically- finding good initialization is critical. Gu et al. [ showed that it is difficult find performant diagonal SSM, and that many alternative parameterizations of the state matrix including by random diagonal matrices
may not represent all SSMs numerically, and
much less effective empirically, which motivated the necessity of the more complicated HiPPO matrix. However, recently Gupta [[]] made the empirical observation that a variant of S4 using particular diagonal matrix is nearly as effective as the original S4 method. This matrix is based on the original HiPPO matrix and is defined by simply chopping off the low-rank
term in the DPLR representation.
The discovery of performant diagonal state matrices opens up new possibilities for simplifying deep state space models, and consolidating models such as S4 and DSS to understand and improve them. First, the strongest version of DSS computes the SSM with complex-valued softmax that complicates the algorithm, and is actually less efficient than S4. Additionally, DSS and S4 differ in several auxiliary aspects parameterizing SSMs that can conflate performance effects, making more difficult to isolate the core effects of diagonal versus DPLR state matrices. Most importantly, DSS relies on initializing the state matrix to particular approximation S4's HiPPO matrix. While S4's matrix has mathematical interpretation for addressing long-range dependencies, the efficacy of the diagonal approximation to it remains theoretically unexplained.
In this work, we seek systematically understand how to train diagonal SSMs. We introduce the S4D method, a diagonal SSM which combines the best of S4's computation and parameterization and DSS's initialization, resulting in method that is extremely simple, theoretically princpled, and empirically effective.
First, we describe S4D, a simple method outlined by S4 for computing diagonal instead of DPLR matrices, which based on Vandermonde matrix multiplication and is even simpler and more efficient than the DSS. Outside of the core state matrix, we categorize different representations of the other components of SSMs, introducing flexible design choices that capture both S4 and DSS and allow different SSM parameterizations to be systematically compared (Section 3}.
We provide a new mathematical analysis of DSS's initialization, showing that the diagonal ap- proximation of the original HiPPO matrix surprisingly produces the same dynamics as S4 when the state size goes to infinity. We propose even simpler variants of diagonal SSMs using different initializations of the state matrix (Section 4).
We perform controlled study of these various design choices across many domains, tasks, and sequence lengths, and additionally compare diagonal (S4D) versus DPLR (S4) variants. Our best S4D methods are competitive with S4 on almost all settings, with near state-of-the -ot-----ttttttt-- results

---

image, audio, and medical time series benchmarks, and achieving 85% on the Long Range Arena benchmark (Section
2 Background
Continuous State Spaces Models
S4 investigated state space models TT that are parameterized maps on signals u(t) + (t). These SSMs are linear time-invariant systems that can be represented either as a linear ODE (equation (T) or convolution (equation (
x' xt) Ar(t) + Bu(t)
K(t) Ce'AB B v(1) = (K * u)(
(2)
v(t) Cx(()
Here the parameters are the state matrix A E CNNN and other matrices B E CNX11 C & CIXN, In the case of diagonal SSMs, A is diagonal and we will overload notation so that A,, Bn: Cn denotes the entries of the parameters.
An intuitive way to view the convolution kemne (2) is to interpret it as a linear combination (controlled by C) of basis kernels K,, (1) (controlled by A, B)
N-1
(tt) = Σ (n))()) n=0
(t) = e, B
(3)
We denote this basis (t) Kt,b(t) A if necessary to disambiguate; note that it is vector of N functions. In the case of diagonal SSMs, each function (t) is just e B,,
S4: Structured State Spaces As a deep learning model. SSMs have many elegant properties with concrete empirical and computational benefits [8]. For example. the convolutional form (2) can be converted into temporal recurrence that is substantially faster for autoregressive applications (5).
However, making SSMs effective required overcoming two key challenges: choosing appropriate values for the matrices, and computing the kernel ② efficiently.
First, Gu et al. showed that naive instantiations of the SSM do not perform well, and instead relied on a particular (real-valued) matrix A called the HiPPO-LegS matrix were derived so that the basis kernels (t) have closed-form formulas [,, ), where (t) are normalized Legendre polynomials. Consequently, the SSM has mathematical interpretation of
( These matrices
decomposing the input signal u(t) onto a set of infinitelyloon basis functions that are orthogonal respect to an axponentiallyddccaii measure, giving it long-range modeling abilities [10].
Second, S4 introduced particular parameterization that decomposed this A matrix into the sum of a normal and rank-1 1 matrix (5). which can be unitarily conjugated into (complex) diagonal plus rank-1 matrix. Leveraging this structured form, they then introduced a sophisticated algorithm for efficiently computing the convolution kernel (2) for state matrices that are diagonal plus low-rank (DPLR).
n n>6 (m + 1 4)122
(2n + 12k + 1)} n+1
> k n = k n <<< k
A (N)
Ank
n
(n 1) + 2/2/22
<<>
0
(4)
A = (N)) -PP.. A2) == cig (A (N)
B, = 11 (n+ (HiPPO-LegS matrix used S4)
(Normal / diagonal plus low-rank form)
(5)
DSS: Diagonal State Spaces S4 was originally motivated by searching for a diagonal state matrix, which would be even more structured and result in very simple computation of the SSM. However, the HiPPO-LegS matrix cannot be stably transformed into diagonal form 9 Lemma 3.2], and they were unable to find any diagonal matrices that performed well, resulting in the DPLR formulation.
Gupta [1 made the surprising empirical observation that simply removing the low-rank portion of the DPLR form of the HiPPO-LegS matrix results in a diagonal matrix that performs comparably
HiiPOO also specifies formulas for B, but the state matrix A is more important. There are many other HiPPO instantiations besides LegS, but HiPPO-LegS is the main one that S4 uses and the term "HiPPO matrix" without the suffix refers to this one.
3

---

to the original S4 method. More precisely, their initialization is the diagonal matrix A A or the diagonalization of A (N) in (5). They termed A (N) the skew-HiPPO matrix, which we will also call the normal-HiPPO matrix. To be more specific and disambiguate these variants, may also call A the HiPPO-LegS-N or HiPPO-N matrix and ( the HiPPO-LegS-D D or HiPPO-D matrix.
In addition to this initialization, they proposed method for computing diagonal SSM kernel. Beyond these two core differences, several other aspects of their parameterization differ from S4's.
In Sections parameterize and compute the diagonal state space, and explain the theoretical interpretion of this particular diagonal A matrix.
B and 4 we systematically study the components of DSS: we categorize different ways to
3 Parameterizing Diagonal State Spaces
We describe various choices for the computation and parameterization of diagonal state spaces, Our categorization of these choices leads to simple variants of the core method. Both DSS and our proposed S4D can be descrihed using a combination of these factors (Section 3.4..
3.1 Discretization
The continuous-time SSM can be represented continuous convolution y(1) (K u)(()
So Ce Bu(t s) ds.
In discrete time, view input sequence uo, ,i.... uniformly-space samples from an underly- ing function ut) and must approximate this integral. Standard methods for doing so that preserve the convolutional structure of the model exist. The first step is to discretize the parameters. Two simple choices that have been used in prior work include
(Bilinear)
A = [I -1-A - A/2A)
(ZOH) Á = exp(AA)
B = (I -I- A/2A) AB
B (AA) - "eexp(A A) - I) AB.
With these methods, the discrete-time SSM output is just
y 2 K
where K (CB.CAB..
CA B).
(6)
These integration rules have both been used in prior works (e.g. LMU and DSS use ZOH 226
while S4 and its predecessors use bilinear 6 8
In Section 5 we show that there is little empirical difference between them. However, we note that there curious phenomenon where the bilinear transform actually perfectly smooths out the kernel used in DSS to match the S4 kernel (Section 4 Fig. 2db. We additionally note that numerical integration a rich and well-studied topic and more stable methods of approximating the convolutional integral may exist. For example, it is well-known that simple rules like the Trapezoid rule [18] can dramatically reduce numerical integration error when the function has bounded second derivative.
3.2 Convolution Kernel
The main computational difficulty of the original S4 model is computing the convolution kernel K. This is extremely slow for general state matrices A, and S4 introduced a complicated algorithm for DPLR state matrices. When A is diagonal, the computation nearly trivial. By (6)
N-1
K.
Σ C. = R= RR=Ba B 1==0
C) (A
where V1. (A) A
(7)
where is Hadamard product, is matrix multiplication, and V is known Vandermonde matrix.
Time and Space Complexity The naive way compute is by materializing the Vandermonde matrix (A( and performing matrix multiplication, which requires O(NL) time and space.
However, Vandermonde matrices are well-studied and theoretically the multiplication can be com- puted in Ö(N + L) operations and O(N L) space. In fact, Vandermonde matrices closely related to Cauchy matrices, which are the computational core of S4's DPLR algorithm, and have identical complexity [17).

---

Proposition 1. The time and space complexity of computing the kernel of diagonal SSMs is equal to that of computing DPLR SSMs.
We note that on modern parallelizable hardware such as GPUs, a simple fast algorithm is to compute 7 with naive summation (using (NNL) operations), but without materializing the Vandermonde matrix (using O(N + L) space). Just as with S4, this may require implementing custom kernel in some modern deep learning frameworks such as Py Torch to achieve the space savings.
3.3 Parameterization
Parameterization of A. Note that the kernel (t) = Ce blows up to oc as - oo if A has any eigenvalues with positive real part. Goel et al. [5] found that this is serious constraint that affects the stability of the model, especially when using the SSM autoregressively. They propose to force the real part of A to be negative, also known as the left-half plane condition in classical controls, by parameterizing the real part inside an exponential function A
exp(ARe) + Asm-
We note that instead of exp, any activation function can be used long as its range is bounded on one side, such as ReLU, softplus, etc. The original DSS does not constrain the real part of A, which is sufficient for simple tasks involving fixed-length sequences, but could become unstable in other settings.
Parameterization of B, C. Another choice in the parameterization is how to represent B and C. Note that the computation of the final discrete convolution kernel K depends only on the elementwise product B o C (equation (7). Therefore DSS chose to parameterize this product directly, which they call W, instead of B and C' individually.
However, observe that this is equivalent to keeping independent B and C, and simply freezing B = 1 while training C. Therefore, just as S4 has separate parameters A, B, and C and uses a fixed initialization for A and B, S4D also proposes separate A, B, and C and uses fixed initializations for A (discussed in Section 44) and B (set to 1). Then the difference between S4D and DSS is simply that DSS does not train B. In our ablations, we show that training B gives a minor but consistent improvement performance.
3.4 S4D: the Diagonal Version of S4
A key component of our exposition is disentangling the various choices possible in representing and computing state space models. With this categorization, different choices can be mixed and matched to define variants of the core method. Table [ compares S4, DSS, and S4D, which have a core structure and kernel computation, but have various choices of other aspects of the parameterization.
Table 1: (Parameterization choices for Structured SSMs.) Aside from the core structure of A and the computation of its convolution kernel, SSMs have several design choices which are consolidated in S4D.
Method Structure
Kernel Computation
Discretization Constraint R(A) Trainable B
Initialization of A
S4 DSS
DPLR diagonal diagonal
Cauchy
Bilincar ZOH
exp id (none)
Yes
HiPPO
softmax
No
HiPPO-D
S4D
Vandermonde
either
exp/ / ReLU
optional
various
4
Initialization of Diagonal State Matrices
The critical question remains: which diagonal state matrices A are actually effective? We comment on the limitations of diagonal SSMs, and then provide three instantiations of S4D that perform well empirically.
Expressivity and Limitations of Diagonal SSMs. We first present simplified view on the ex- pressivity of diagonal SSMs mentioned by [.]. First, it is well-known that almost all matrices diagonalize over the complex plane. Therefore it is critical to use complex-valued matrices in order to use diagonal SSMs.
Proposition 2. The set D C NNN of diagonalizable matrices is dense in NNN and has full measure (i.e. its complement has measure 0).
5

---

It is also well known that the state space (A BB C') is exactly equivalent to (i.e. expresses the same map u + y) the state space (V-1 AV. V- B, CV), known in the SSM literature as a state space transformation. Therefore Proposition 2 says that (almost) all SSMs are equivalent to a diagonal SSM.
However, emphasize that Proposition 2 is about expressivity which does not guarantee strong performance of a trained model after optimization. For example, Gu et al. [9] and Gupta [T]] show that parameterizing A dense real matrix diagonal complex matrix, which are both fully expressive classes, performs poorly if randomly initialized.
Second, Proposition 2 does not take into account numerical representations of data, which was the original reason S4 required a low-rank correction term instead of pure diagonalization [9] Lemma 3.2]. In Section [.2] we also show that two different initializations with the same spectrum (i.e., are equivalent to the same diagonal A) can have very different performance.
S4D-LegS. The HiPPO-LegS matrix has DPLR representation A PP and Gupta [1]] showed that simply approximating it with A works quite well (5). Our first result is providing a clean math- ematical interpretation of this method. TheoremBshove a surprising fact that does not hold in general for DPLR matrices (Appendix A.I). and arises out of the special structure of this particular matrix.
Theorem Let = ( PP and B be the HiPPO-LegS matrices, and K (t, be its basis. As the state size N + oc, the SSM basis K A ,B/2(t) A(0,B/22(t limits to K (t,((t) (Fig.
Note that (N) is then unitarily equivalent to A (D), which preserves the stability and timescale 10 of the system.
We define S4D-LegS to be the S4D method for this choice of diagonal A A (D), Theorem explains the empirical results in [11 whereby this system performed quite close to S4, but was usually slightly worse. This is because DSS is a variant of S4D-LegS, which by Theorem 3 is noisy approximation to S4-LegS. Fig. Diilusstras this result, and also shows a curious phenomenon involving different discretization rules that is open for future work.
S4D-Inv. To further simplify S4D-LegS, we analyze the structure of A diag(A) in more detail. The real part is easy to understand, which follows from the analysis 19: R(A) - 1. Let the imaginary part be sorted, i.c. ((A,, is the n-th largest (positive) imaginary component. We empirically deduced the following conjecture for the asymptotics of the imaginary part.
Conjecture 4. As N + co, (A)) + N2 + where = 0.5236 is constant. For fixed N, the other eigenvalues satisfy inverse sculing in ((A),. O(n
Fig. 3 empirically supports this conjecture. Based on Conjecture 4 we propose the initialization S4D-Inv to use the following inverse-law diagonal matrix which closely approximates S4D-LegS.
2n +1 -1)
N
(S4D-Inv) A,
(8)
(S4D-Lin) A,
2 + inn (9)
S4D-Lin. While SnD-Inv can be seen as an approximation to the original S4-LegS, we propose an even simpler scaling law for the imaginary parts that can be seen an approximation of S4-FouT (I10), where the imaginary parts are simply the Fourier series frequencies (i.e. matches the diagonal part of the DPLR form of S4-Fou')). Fig. (Right) illustrates the S4D-Lin basis A B, which are simply damped Fourier basis functions.
5 Experiments
Our experimental study shows that S4D has strong performance in wide variety of domains and tasks, including the well-studied Long Range Arena (LRA) benchmark where the best S4D variant competitive with S4 on all tasks and significantly outperforms all non-SSM baselines.
We begin with controlled ablations of the various representations of diagonal state space models. Sec- tions and 5.2 ablate the proposed methods for parameterizing, computing, and initializing diagonal SSMs from Sections and Section 5.3 show full results of larger models on standard benchmarks,
Methodology and Datasets. In order to study the effects of different S4 and S4D variants in a controlled setting, propose the following protocol. We focus three datasets covering varied

---

(a) S4-LegS (original S4 (b) S4D-LegS (N 256, (c) S4D-LegS (N 1024, (d) S4D-LegS (N = 1024, kernel)
ZOHI)
ZOHI)
Bilinear)
Figure 2: (Visualization of Theorem (a) The particular (A, B) matrix chosen in S4 results in smooth basis functions A B with a closed form formula in terms of Legendre polynomials. By the HiPPO theory, convolving against these functions has a mathematical interprctation as orthogonalizing against an exponentially-dccnn measure. (b, c) By special properties of this state matrix, removing the low-rank term of its NPLR representation produces the same basis functions N oc, explaining the empirical effectiveness of DSS. (c) Curiously, the bilinear transform instead of ZOH smooths out the kernel to exactly match S4- S4-LegS N grows.
1200
Logs S4b4t
Figure 3: (S4D eigenvalues.) All S4D methods have eigenval- ues A.i. S4D-LegS theoretically approximates dynamics of the original (non-diagonal) S4 (Bluc), and has eigenvalues following inverse law A.
S4D-Lin
Inuginay part
1000
$4D Quad
800
54D-inz2
(Orange). The precise law important: other scaling laws with the same range, including an inverse law with different constant (Purple) and quadratic law (Red), perform empirically worse (Section (.2). A very different linear law based on Fourier frequencies also performs well (Green).
600
400
elgerwalue, descending Imaginary part
range of data modalities (image pixels, biosignal time series, audio waveforms), sequence lengths (IK, 4K, 16K), and tasks (classification and regression with bidirectional and causal models).
Sequential CIFAR (sCIFAR), CIFAR-I0 10 images are flattened into sequence of length 1024, and a bidirectional sequence model is used to perform 10-way classification.
BIDMC Vital Signs. EKG and PPG signals length 4000 are used to predict respiratory rate (RR), heart rate (HR), and blood oxygen saturation (SpO2). We focus on SpO2 in this study.
Speech Commands (SC) I-second raw audio waveform comprising 16000 samples is used for 35-way spoken word classification. We use an autoregressive (AR) model to vary the setting; this causal setting more closely imitates autoregressive speech generation, where SSMs have shown recent promise (5]
We fix a simple architecture and training protocol that works generically. The architecture has 4 layers and hidden dimension H 128, resulting in 100K parameters. All results are averaged over multiple seeds (full protocol and results including std. reported in Appendix B
5.1 Parameterization, Computation, Discretization
Given the same diagonal SSM matrices A, B, there are many variants of how to parameterize the matrices and compute the SSM kernel described in Section 3 described in Table [ Results are in Table 2 and show that:
We ablate the different choices
(i Computing the model with softmax instead of Vandermonde product does not make much difference
(ii Training B is consistently slightly better
(iii) Different discretizations (Section 3.1 do not make noticeable difference
(iv)
Unrestricting the real part of A (Section 3.3) may be slightly better
These ablations show that for a fixed initialization (A, B), different aspects of parameterizing SSMs make little difference overall. This justifies the parameterization and algorithm S4D uses (Section which preserves the choices of the original S4 model and is simpler than DSS. For the remaining
We note that a line of prior work including S4 14] 9 all used smaller cl-ssass subset of SC, our results the full dataset are not directly comparable.

---

Discrefization part
SCIFAR
5C (AR)
(5902)
Mothod
SC (AR)
(5p02)
0Ap
65.20
Saftmax
89.90
0.1299
Vondermonde
84.11 78
04102
85
90.06
0.1170
20H
85.07
80.03
0.1305
Vandermardde
85,
00.34
0.1274
RULu
84.98
90.03
0.1232
65.15
90.13
0.1289
Table 2: Ablations of different parameterizations of diagonal SSMs using S4D-Inv. (Left) trainability and computation: (Right) discretization and parameterization.
of the experiments in Section 5.2 and Section 5.3
we fix the S4D parameterization and algorithm
described in Section Note that this computes exactly the same kernel as the original S4 algorithm when the low-rank portion is set to 0, allowing controlled comparisons of the critical state matrix A
for the remainder of this section.
5.2 S4D Initialization Ablations
The original S4 model proposed a specific formula for the A matrix, and the first diagonal version [II] used a specific matrix based on it. Our new proposed variants S4D-Inv and S4D-Lin also define precise formulas for the initialization of the A matrix ( This raises the question of whether the initialization of the still needs to be so precise, despite the large simplifications from the original version. We perform several natural ablations these initializations, showing that even simple variations of the precise formula can degrade performance.
Imaginary part scaling factor. The scaling rules for the imaginary parts of S4D-Inv and S4D-Lin are simple polynomial laws, but how is the constant factor chosen and how important is it? These constants are based on approximations HiPPO methods (e.g. Conjecture 4). Note that the range of imaginary components for S4D-Inv and S4D-Lin are quite different (Fig. 3) the largest imaginary
part is 2 for Inv and N for S4D-Lin.
We consider scaling all imaginary parts by a constant factor of 0.01 or 100.0 to investigate whether the constant matters. Note that this preserves the overall shape the basis functions (Fig. [ dashed lines) and simply changes the frequencies, and it is not obvious that this should degrade performance. However, both changes substantially reduce the performance of S4D in all settings.
Randomly initialized imaginary part. Next, we consider choosing the imaginary parts ran- domly. For S4D-Inv, we keep the real parts equal to 2 and set each imaginary component to A, + 2u+1 -1 1) for N [1, 1]. Note that when u is equally spaced in [0, 1] instead of uniformly random, this exactly recovers S4D-Inv (8) so this is a sensible random approximation to it.
Similarly, we consider a variant of S4D-Lin that replaces the n in 9 with N (1(00, 1].
Table 3al (Random Imag) shows that this small change causes minor degradation in performance. We additionally note that the randomly initialized imaginary ablation can be interpreted follows. Fig. 3 shows the asymptotics of the imaginary parts of SSM matrices, where the imaginary parts of the eigenvalues correspond to y-values corresponding to uniformly spaced nodes on the x-axis. This ablation then replaces the uniform spacing on the x-axis with uniformly random x values.
Randomly initialized real part.
We considering initializing the real part of each eigenvalue as -2[[0 1] instead of fixing them to -1. Table Left, Random Real) shows that this also causes minor but consistent degradation in performance on the ablation datasets. Finally, we also consider randomizing both real and imaginary parts, which degrades performance even further.
Ablation: Other S4D matrices.
Other simple variants of initializations show that it is not just the range of the eigenvalues but the actual distribution that important (Fig. 3). Both S4D-Inv2 and S4D-Quad have real part p and imaginary part satisfying the same maximum value as Conjecture 4 The S4D-Inv2 initialization uses the same formula as S4D-Inv, but replaces 2n 1 in the denominator with 1. The S4D-Quaa initialization uses a polynomial law with power 2 instead of - (S4D-Inv) or 1 (S4D-Lin).
We include two additional methods here that are not based on the proposed S4D-Inv or S4D-Lin methods. First, S4D-Rand uses randomly initialized diagonal A, and validates that it performs
8

---

Table (Initialization and Trainability ablations)
(AR)
(AR)
Frozen (A, B)
Acc (first)
Acc (best)
Acc (firet)
Acc (best)
RMSE (best)
S4D-Lin
85.12
90.66
0.128
S4-LegS
53.63
86.19
33.87
65.33
0.1049
Scale 0.01
-7.27
-1.92
+0.040
Scale 100
-.991
-4.04
+0.077
S4-LegT
54.76
86.30
8.77
57.35
0.1106
Random imag
0.42
3308
0.001
54-FouT
56.28
86.05
9.27
69.57
0.1072
Random Real
0.73
-0.87
+0.011
S4-LegS-FouT
54.38
86.53
34.06
83.37
0.0887
-128
54D-LegS
50.87
84.81
22.76
77.18
0.0960
S4D-Inv
84.79
00.27
0:114
S4D-Lin
51.75
84.96
19.09
75.58
0.0985
Scale 0.01
5.08
-0.08
+0.026
Traineble
Scale 100
-7.77
-2.31
+0.034
Random imag
-0.99
-0.52
+0.010
S4-Leg5
54.23
86.29
62.19
90.68
0. 0.1033
Random Real
0.12
22188
+0.032
S4-LegT
56.16
86.12
55.86
90.42
0.1146
Random Both
-1.55
-0.55
+00024
55.89
85.93
60.56
90.83
0.1136
S4D-Inv2
-2.62
39.84
+0.005
54-Leg5-FouT
55.00
86.18
61.76
91.01
0.0970
S4D-Quad
-1.83
-0.62
+0.024
S4D-LegS
S4D-Random
6.32
-1.95
+0.034
54D-Inv
53.42
84.59
45.73
89.69
0.1132
94D-Real
5.45
-10.17
+0.066
S4D-Lin
(a) Ablations of the initialization of the diag onal A matrix in S4D. Very simple changes (b) Results for all S4 S4D methods on the ablation datasets, that largely preserve the structure of the diag- when the A and B matrices are either frozen (Top) or trained onal eigenvalues all degrade performance.
(Bottom). Diagonal state matrices are highly competitive with full DPLR versions, achieving strong results on all datasets.
poorly, in line with earlier findings I1I. Second, S4D-Real uses a particular real initialization with An -(n + 1). This is the exact same spectrum as the original S4(-LegS) method, which validates that it is not just the diagonalization that matters, highlighting the limitations of Proposition 2
5.3 Full Comparisons of S4D and S4 Methods
Trainable A, B matrices.
Table 3b shows the performance of all S4D and S4 variants 110 on the ablations datasets. We observe several Interesting phenomena:
(i) Freezing the matrices performs comparably to training them on sCIFAR and BIDMC, but is substantially worse on SC. We hypothesize that this results from A being poorly initialized for SC, so that at initialization models do not have context over the entire sequence, and training A and B helps adjust for this. As further evidence, the finite window methods S4-LegT and S4-Fou'T (defined in 10 have the most limited context and suffer the most when A is frozen.
(ii) The full DPLR versions are often slightly better than the diagonal version throughout the entire training curve. We report the validation accuracy after 1 epoch of training on SCIFAR and SC illustrate this phenomenon. Note that this is not consequence of having more parameters (Appendix B
Large models ablation datasets.
Finally, relax the strict requirements model size and regularization for the ablation datasets, and show the performance of S4 and S4D variants on the test sets with a larger model (architecture and training details in Appendix B when the model size and regularization is simply increased (Table 4.. We note that results for each dataset are better than the original $4 model, which was already oftt-off-theart on these datasets 8
Long Range Arena.
We use the same hyperparameter setting for the state-of-the-art S4 model in 10 on the Long Range Arena benchmark for testing long dependencies in sequence models. S4D variants highly competitive on all datasets except Path-X, and outperform the S4 variants on several of them. On Path-X using this hyperparameter setting with bidirectional models, only S4D-Inv, approximation to the original S4-LegS model, achieves above random chance, and has an average of 85% on the full LRA suite, more than 30 points better than the original Transformer (24..
simpler

---

Table 4: (Ablation datasets: Full results with larger models.) For Speech Commands, show both an autoregressive model as in ablations, and an unconstrained bidirectional model.
MODEL
SCIFAR
SC
BIDMC
TEST
B1.
HR
RR
SPO2
S4-LegS
91.80 (0.43)
93.60 (0.13)
96.08 (0.15)
0.332 (0.013) 0.247 (0.062)
0.090 (0.006)
S4-Fou'T
91.22 (0.25)
91.78 (0.10)
95.27 (0.20)
0.339 (0.020)
0.301 (0.030)
0.068 (0.003)
S4D-LegS
89.92 (1.69)
93.57 (0.09)
95.83 (0.14)
0.367 (0.001)
0.248 (0.036)
0.102 (0.001)
S4D-Inv
90.69 (0.06)
93.40 (0.67)
18.18 (0.27)
0.373 (0.024)
0.254 (0.022)
0.110 (0.001)
SiD-Lin
90.42 (0.03)
93.37 (0.05)
96.25 (0.03)
0.379 (0.006) 0.226 (0.008)
0.114 (0.003)
Table (Long Range Arena) Accuracy on full suite L.A tasks. Hyperparameters in Appendix
MODEL
ListOps
TEXT
RETRIEVAL
IMAGE
PATHFINDER
PATH-X
AVG
LegS
59.60 (0.07)
86.82 (0.13)
90.90 (0.15)
88.65 (0.23)
94.20 (0.25)
96.35
86.09
S4-FouTT
57.88 (1.90)
86.34 (0.31)
89.66 (0.88)
89.07 (0.19)
94,46 (0.24)
X
77.90
S4D-LegS
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
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
60.52 (0.51)
86.97 (0.23)
90.96 (0.09)
87.93 (0.34)
93.96 (0.60)
X
78.39
$4 (original) 58.35 Transformer
76.02
87.09
87.26
86.05
88.10 10
80.48
36.37
64.27
$7.46
42.44
71.40
X
53.66
6
Conclusion
State space models based S4 are promising family of models for modeling many types of sequen- tial data, with particular strengths for continuous signals and long-range interactions. These models a large departure from conventional sequence models such RNNs, CNNs, and Transformers, with many new ideas and moving parts. This work provides a more in- exposition for all aspects of working with S4-style models, from their core structures and kernel computation algorithms, to miscellaneous choices their parameterizations, to new theory and methods for their initialization. We systematically analyzed and ablated each of these components, and provide recommendations for building state space model that is as simple as possible, while as theoretically principled and empirically effective as S4. We believe that S4D can be a strong generic sequence model for variety of domains, that opens new directions for state space models theoretically, and is much more practical to understand and implement for practitioners.
Acknowledgments
We gratefully acknowledge the support of DARPA under Nos. FA86501827865 (SDH) and FA86501 827882 (ASED); NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under No. N00014177 12266 (Unifying Weak Supervision); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Swiss Re, Brown Institute for Media Innovation. Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, Fannie and John Hertz Foun- dation, National Science Foundation Graduate Research Fellowship Program, Texas Instruments, and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC, VMWare, and Infosys. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, endorsements, either expressed implied, of DARPA, NIH, ONR, or the U.S. Government.
10

---

References
[1] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization, arXiv preprint arXiv:1607.06450, 2016.
[2] Shaojie Bai, Zico Kolter, and Vladlen Koltun. Trellis networks for sequence modeling. The International Conference on Learning Representations (ICLR), 2019.
[3] Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933-941. PMLR, 2017.
[4] Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W Mahoney. Lipschitz recurrent neural networks. In International Conference on Learning Representations, 2021.
[5] Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. In The International Conference Machine Leurning (ICML), 2022.
[6] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[7] Albert Gu, Caglar Gulcehre, Tom Le Paine, Matt Hoffmann, and Razvan Pascanu. Improving the gating mechanism of recurrent neural networks. In The International Conference on Machine Learning (ICML), 2020.
[8] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with the structured learnable linear state space layer. Advances in Neural Information Processing Systems (NeurIPS), 2021.
[9] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022.
[10] Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How train your hippo: State space models with generalized basis projections. arXiv preprint arXiv:2206. 12037, 2022.
[11] Ankit Gupta. Diagonal state spaces are effective as structured state spaces. Advances in Neural Information Processing Systems (NeurIPS), 2022.
[12] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()) 1785-1780, 1997.
[13] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pages 448-456. PMLR, 2015.
[14] Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time series. arXiv preprint arXiv:2005. 002.02000..0926,, 2020.
[15] James Morrill, Cristopher Salvi, Patrick Kidger, James Foster, and Terry Lyons. Neural rough differential equations for long time series. The International Conference on Machine Learning (ICML), 2021.
[16] Naoki Nonaka and Jun Seita. In-depth benchmarking of deep neural network architectures for eeg diagnosis. In Machine Learning for Healthcare Conference, pages 414-439. PMLR, 2021.
[17] Victor Pan. Structured matrices and polynomials: unified superfast algorithmss Springer Science Business Media, 2001.
[18] Authony Ralston and Philip Rabinowitz. A first course in numerical analysis. Courier Corpora- tion, 2001.

---

[19] David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomezak, and Mark Hoogendoom. Ckconv: Continuous kernel convolution for sequential data. arXiv preprint arXiv:2102.02611 2021.
[20] David W Romero, Robert-Jan Bruintjes, Jakub M Tomezak, Erik J Bekkers, Mark Hoogendoomn, and Jan C van Gemert. Flexconv: Continuous kernel convolutions with differentiable kernel sizes. In The International Conference on Learning Representations (ICLR), 2022.
[21] T Konstantin Rusch and Siddhartha Mishra. Unicornn: A recurrent model for learning very long time dependencies. The International Conference on Machine Learning (ICML), 2021.
[22] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05200) 2020.
[23] Chang Wei Tan, Christoph Bergmeir, Francois Petitjean, and Geoffrey I Webb. Time series extrinsic regression. Data Mining and Knowledge Discovery, pages 1-29, 2021. doi: hitps:
//dbi.org/10.1100///1666 00000066610001-05-----11
[24] 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 Reprexentations, 2021. URL
ggrC2k
[25] Trieu Trinh, Andrew M Dai, Minh- Minh-Thang Luong, and Quoc V Le. Learning longer-tern term dependencies in RNNs with auxiliary losses. In The International Conference on Machine Learning (ICML), 2018.
[26] Aaron Voelker, Ivana Kajić, and Chris Eliasmith. Legendre memory units: Continuous-timg representation in recurrent neural networks. In Advances in Neural Information Processing Systems, pages 15544-15553, 2019.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section c.g. not matching the baseline S4 Path-X.
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix A
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (cither in the supplemental material or as a URL)? [Yes] The code is a simple modification from the original S4 [9] repository and is publicly available.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See the Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] The experiment code and configs publically available with resource and timing information reported.
4. If you using existing assets (e.g., code, data, models) curating/releasing assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
12

---

(c) Did you include any new assets either in the supplemental material or as URL? [No] Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(d)
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] We used standard benchmarks and synthetic
data.
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "fEoH7WWkxQ1", "reviewer": "Reviewer_MC9C", "summary": "The authors examine a state-space modeling approach that has gained popular in recent literature — the main idea of that line of research concerned with the expressive capability, and ability to capture long-term memory dependence through the use of linear SSMs defined in continuous time.  The authors note that practical use of HiPPo, the structured state-space architecture (S4) and diagonal state-space architectures (DSS) may present complications due to various numerical instabilities associated with bad initializations or parameterizations. To be more specific, they systematically analyze the diagonal approximation to S4 — they prescribe different initializations that perform better numerically for which they also back up theoretically.  Through an ablation study, they are also able to demonstrate the practical utility of their analysis and suggestions.\n", "questions": "If my statement about the dimensionality of the variables specified was wrong, hopefully you could clear that up.  In, Eq. (5) A^D as defined, is always just the matrix of 1/2 on the diagonal, is that correct?\n\nI don’t think the DSS acronym was specified in text (is it diagonal state-space ?).  Is bar{C} as defined in Eq. (6) just equal to the original C row vector?\n\nIn Eq. (7) should the ordering of B and C be swapped? Stemming from my confusion earlier, I thought that B was a column vector and C a row vector.  Maybe subscripts here would help i.e. y_l = (u * K)_l or something more economical that clears up confusion.  A reference for the equality using the Vandermonde matrix would probably help as well in my opinion.  \n\nOn page 5, section 4 below prop. 1 you call the randomly initialized matrices fully expressive — in what sense?\n", "limitations": "In my opinion, the contributions of the paper are somewhat incremental;  that’s not to say the authors contributions are not novel — but a great deal of writing space is used to describe the previous research on HiPPo.  I think more of the authors contributions could have been thrown into the spotlight, rather than to rehash the cited literature.  Furthermore, since a large part of the paper is about how to select good, in some sense, dynamics matrices — how come there is no discussion on parameter training? I would think for a practitioner, it may be easier to train the parameters of the matrices in question, rather than doing a large amount of model selection.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I think the paper addresses some general issues that would be good for someone looking to employ the previously introduced HiPPo matrix into their modeling schemes.  For example, the authors address numerical concerns that are often brushed over in many papers — even though not addressing them can be catastrophic i.e. parameterization of the dynamics such that the matrix is stable.  The memory saving tip in 3.3 for storing B and C is useful as well.  In section 3.4 I like the initialization prescribed, and while its a nice concept I think the presentation could just be slightly tidied up — a comparison to balancing would be useful too, see [E. E. Osborne, “On pre-conditioning of matrices”].  I thought some of the ablation studies were also very useful; since the point of the paper is to prescribe practical guidelines to follow its nice the authors swept through different variables such as the discretization method, normalization, and parameterizations.  \n\nOne of the first difficulties I encountered was with notation.  If inline with the papers cited regarding HiPPo, then it would have been useful to specify that for example, C is a row vector, B is a column vector, and y(t) and u(t) are one-dimensional observations and control inputs respectively.   The lack of conclusion is also somewhat concerning, I think a succinct summary of the experimental and theoretical results would have been a good way to end the paper.  Another big weakness I think lays in the total contributions of the authors; a lot of rehashing of previous results is done, but the authors in my opinion do not devote enough space in the main text to outline their current contributions.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "UX7u1k_c-B", "reviewer": "Reviewer_E3uH", "summary": "The authors set out to understand why some diagonal variants of the S4 state-space model work well, then introduce S4D as a specific approach building off of DSS. Their main mathematical result shows that a purely diagonal update reproduces the dynamics of S4’s diagonal plus low rank update in the limit of large latent size.\n\nBased on their theoretical insights, they improve on existing diagonal state-space models (DSS) by carefully choosing initialization, discretization, and related parameters in the S4D family of models. The authors also report detailed ablations to confirm the theoretical conclusions, and demonstrate S4-equivalent quality on a variety of benchmarks.\n", "questions": "Do you have a guess or narrative for what goes wrong when a parameterization is expressive but doesn’t optimize well? (“[W]e emphasize that Proposition 1 is about expressivity which does not guarantee strong performance of a trained model after optimization. For example, Gu et al. [7] and Gupta [8] show that randomly initialized dense real matrices or diagonal complex matrices, which are both fully expressive, perform much worse than S4.”)\n\nThe ability to fully diagonalize the SSM update creates a strongly privileged basis; have the authors explored interpretability of latent states in this basis?\n", "limitations": "If there are significant limitations of diagonal SSMs, including S4D, relative to the original S4, it would certainly be worth mentioning them as negative results. But so far it appears not to be the case.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is clear despite packing information very densely and including both theoretical and empirical results.\n\nThe ablations are strong and careful, and presented concisely and clearly (I haven’t seen this particular graphical presentation of ablations before, and I think it works well.)\n\nOverall, I think it was unclear before this paper why diagonal SSMs worked but related methods like linear RNNs didn’t; this paper begins to answer that question.\n\nThe approach presented in this paper is to compute the kernel using a Vandermonde matrix, but the authors don't make the connection between this approach and the parallel scan implementation of the linear recurrence. The parallel scan approach has some advantages over explicitly computing the kernel, including support for unbounded sequence length (important for online tasks like reinforcement learning) and variable/dynamic sampling rate.\n\nThe abstract mentions “three lines of code”; these should be included in the main paper.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Sgta6uGnnc2", "reviewer": "Reviewer_s9cB", "summary": "Detailed analyses of some variants of the so-called S4 model are presented. The focus is on the use of a diagonal state-transition matrix instead of a diagonal plus low-rank matrix, time discretization scheme, spectral radius constraint of $A$, parametrization of $B$ and $C$, normalization of convolution kernels corresponding to SSM, and initialization scheme. The authors provide empirical ablation studies in these regards. They also provide the rationale for why the diagonal approximation of a so-called HiPPO-LegS matrix works well.", "questions": "**1.**\nIn Figures 1 and 2, please specify what is drawn there not only in the caption but also by using axis labels and legends. The current way of presentation is not so kind.\n\n**2.**\nHow the proposed models, which are sequence-to-sequence models, were applied to the sequence-to-label problems is not clearly stated. Maybe it is just the same as the previous works, but it would be nice to briefly mention the model architecture somewhere in the main text of the paper.\n\n**3.**\nThe proof of Theorem 2. I found no obvious errors in the step-by-step reasoning. However, the intermediate result of $u(t) = \\mathbf{B}^T x(t)$ is quite counterintuitive. I understand this happens only for $N \\to \\infty$ with the HiPPO-LegS $\\mathbf{B}$ matrix, but I wonder if there is any more intuitive interpretation of why such a relation is the case here.", "limitations": "I don't have particular comments in this regard.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I did not find serious shortcomings. The analyses would be of interest to a limited community.\n\n### Strengths\n\n- The experiments are intensive and well organized.\n- While I could not check all the correctness of the proofs, the discussion around Theorem 2 is interesting.\n\n### Weaknesses\n\n- The target audience is limited.\n- The evaluation is done only in terms of the final performance of classification or regression tasks. Some qualitative and direct analyses of the models' outputs, if possible (and perhaps on much simpler input data), would be interesting.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EYK0ZszoA39", "reviewer": "Reviewer_zJVF", "summary": "Paper simplifies diagonal state space models by i) analysing Eigen-values of Hippo matrices and proposing appropriate initialisation schemes, ii) simplifying the parameterisation of the real parts (to be negative) of the Eigen-values (for stability) and iii) introducing a simple normalisation scheme that ensures variance preservation of the resulting SSM. Experiments are conducted to verify all these simplifications and the authors show that they do not lead to inferior performance. A very important take away message from the paper is that diagonalization alone doesn’t make a performant SSM but it needs proper parameterisation/initialization. Overall the paper is a great extension to the growing literature on SSMs for Deep Learning. ", "questions": "* “Second, Proposition 1 does not take into account numerical representations of data, which was  the original reason S4 required a low-rank correction term instead of a pure diagonalization.“   —> This is quite vague, pleas elaborate or strike.\n* How do you go from conjecture 3 to S4D-Inv?\n* Table 1: What S4D variant is used?\n* Curious to see what would happen if only A were left untrained.\n* S4D-Trap introduced but never ablated?\n* More information on how robust S4D is to choice of optimization hparams is would be interesting.\n* Simplicity is really nice but it is not clear to me from the paper why S4D should be used over S4 for any practical reasons? Do you see any speed ups? Results seem to still lag behind S4.", "limitations": "The authors haven't addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n* good and mostly clear presentation\n* solid mathematical foundation for simplifications\n* convincing empirical results\n\n## Weaknesses\n* No dicusssion/limitations\n* Mostly small scale, toyish tasks evaluated\n* The results don't show any advantages of S4D over S4", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept