Diagonal State Spaces are as Effective as Structured State Spaces
Authors: Ankit Gupta
Albert Gu
Jonathan Berant
Venue: NeurIPS 2022 Accept (2022)
Keywords: state spaces
long range models
efficient
Transformer
speech recognition
language modeling
time series model

Abstract: Modeling long range dependencies in sequential data is a fundamental step towards attaining human-level performance in many modalities such as text, vision, audio and video. While attention-based models are a popular and effective choice in modeling short-range interactions, their performance on tasks requiring long range reasoning has been largely inadequate. In an exciting result, Gu et al. (ICLR 2022) proposed the $\textit{Structured State Space}$ (S4) architecture delivering large gains over state-of-the-art models on several long-range tasks across various modalities. The core proposition of S4 is the parameterization of state matrices via a diagonal plus low rank structure, allowing efficient computation. In this work, we show that one can match the performance of S4 even without the low rank correction and thus assuming the state matrices to be diagonal. Our $\textit{Diagonal State Space}$ (DSS) model matches the performance of S4 on Long Range Arena tasks, speech classification on Speech Commands dataset, while being conceptually simpler and straightforward to implement.

Diagonal State Spaces are as Effective as Structured State Spaces
Ankit Gupta*
Albert Gu
Jonathan Berant Tel Aviv University joberant&cs tau. tau.ac.il
IBM Research ikitgupta.iittkanpun igmaa com
Stanford University albertguustanford.
Abstract
Modeling long range dependencies in sequential data is fundamental step towards attaining human-level performance in many modalities such as text, vision, audio and video. While attention-based models are popular and effective choice in modeling short-range interactions, their performance on tasks requiring long range reasoning has been largely inadequate. In exciting result, Gu al. (2022) proposed the Structured State Space (S4) architecture delivering large gains over fate-offtthearr models several long-range tasks across various modalities. The core proposition of S4 is the parameterization of state matrices via diagonal plus low rank structure, allowing efficient computation. In this work, show that one can match the performance of S4 even without the low rank correction and thus assuming the state matrices be diagonal. Our Diagonal State Space (DSS) model matches the performance of S4 on Long Range Arena tasks, speech classification on Speech Commands dataset, while being conceptually simpler and straightforward to implement.
Introduction
The Transformer architecture IVSP+ 17) has been successful across many areas of machine learning. Transformers pre-trained on large amounts of unlabelled text via denoising objective have become the standard in natural language processing, exhibiting impressive amounts of linguistic and world 20, knowledge [BMR 20, BMH+22) This recipe has also led to remarkable developments in the areas
of vision [RPG + 21, RKH 21] and speech [DXX18. BZMA20J.
The contextualizing component the Transformer is the multi-head attention layer which, for inputs of length T., has expensive 22 (.22) complexity. This becomes prohibitive tasks where the model required to capture long-range interactions various parts of a long input. To alleviate this issue, several Transformer variants have been proposed with reduced compute and memory requirements [QML+ 20, KKL200 RSVG21, BPC20, GB20, K KVPF20, WLK 20, CLD+21, ZS21, GDG [D] (see [TDBM20] for survey). Despite this effort, all these models have reported inadequate performance on benchmarks created to formally evaluate and quantify a model's I's ability to perform long-range reasoning (such Long Range Arena [TDA* +21] and SCROLLS (22).
In a recent breakthrough result, Gu et al. [GGR22] proposed S4, a sequence-to teaeeettt-ssuunne model that uses linear state spaces for contextualization instead of attention. It has shown remarkable performance on tasks requiring long-range reasoning in domains such as text, images and audio. For instance, Long Range Arena advances the state-of-the-ard by 19 accuracy points over the best performing Transformer variant. Its remarkable abilities are not limited to text and images and carry over to tasks such time-series forecasting, speech recognition and audio generation [GGDR22).
Despite S4's achievements, its design complex and is centered around the HiPPO theory, which is mathematical framework for long-range modeling [VKE19, 2DE+ GJG+21]. [GGR22]
work done while author was part of IBM AT Residency program.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

showed that state space models with various alternative initializations perform poorly in comparison to initializing the state-space parameters with a particular HiPPO matrix. In order to leverage this matrix, they parameterize the learned state spaces using Diagonal Plus Low Rank (DPLR) structure and, as a result, need to employ several reduction steps and linear algebraic techniques to be able to compute the state space oulput efficiently, making S4 difficult to understand, implement and analyze.
In this work, we show that is possible to match S4's performance while using a much simpler, fully- diagonal parameterization of state spaces. While confirm that random diagonal state spaces less effective, we observe that there do in fact exist effective diagonal state matrices: simply removing the low-rank component of the DPLR HiPPO matrix still preserves its performance. Leveraging this idea, our proposed Diagonal State Space (DSS) model enforces state matrices to be diagonal, making it significantly simpler to formulate, implement and analyze, while being provably expressive as general state spaces. In contrast to $4, DSS does not assume any specialized background beyond basic linear algebra and can be implemented in just few lines of code. Our implementation fits in single page and is provided in $A.7 (Figure 6).
We evaluate the performance of DSS on Long Range Arena (LRA) which is suite of sequence- level classification tasks with diverse input lengths (1K-16K) requiring similarity, structural, and visual-spatial reasoning over wide range of modalities such as text, natural/synthetic images, and mathematical expressions. Despite its simplicity, DSS delivers an average accuracy of 81.88 .88 across the 6 tasks of LRA,, comparable to the state-of-the-an performance of S4 (80.21). In addition, DSS maintains a comfortable 20 point lead over the best Transformer variant (81.8) (81.88 vs 61.41).
In the audio domain, evaluate the performance of DSS on raw speech classification. On the Speech Commands dataset [WarlB 8), which consists raw audio samples of length 16K, we again found the performance of DSS to be comparable to that of S4 (98.2 vs 98.1).
To summarize, our results demonstrate that DSS simple and effective method for modeling long-range range interactions in modalities such as text, images and audio. We believe that the effectiveness, efficiency and transparency of DSS can significantly contribute the adoption of state space models over their attention- peers. Our code is available at https: //g giihhu.. com/ ag/988/1///// dss..
2 Background
We start by reviewing the basics of time-invariant linear state spaces.
State Spaces A continuous- time state space model (SSM) parameterized by the state matrix Ae NNN and vectors B E Vx1 C RI*N N is given by the differential equation:
dx (t) Ar(t) + Bu(t) y(t) Cx(()
di
which defines a cenetiionto-functioon map (l() - (t). For a given value of time / TR, u(() e IR denotes the value of the input signal u, xx()= E Nx1 denotes the state vector and (() TR denotes the value of the output signal y. We call state space diagonal if it has a diagonal state matrix.
Discretization For given sample time A R=0, the discretization of continuous state space (Equation 1) assuming zero-order hold on is defined as a sequence-to-sequence map from (uo...., #L-1) e IR_ (Vo, ..... 1L-1) = y IR-- via the recurrence,
Dg ArR-I Bud Vs Czz
(2)
A= B (A TA- C-C.
Assuming .-1 0 for simplicity, this recurrence can be explicitly unrolled as
yk = Σ CA uk-j
(3)
1-0
For convenience, the scalars U.A' B are gathered to define the SSM kernel K IR_
K = (CB, CAB, C.AA 1-1 B) = (CeAtt (
- T)A-1B B )osk<L
(4)
assumes the value of sample of is held constant for duration of one sample interval A [Gaj03].

---

where the last equality follows by substituting the values of A, B, C from Equation 2. Hence,
k yk K5 KK-uk=j uz-j
(5)
Given an input sequence u E R it is possible to compute the output y E R sequentially via the recurrence in Equation 2. Unfortunately, sequential computation is prohibitively slow on long inputs and, instead, Equation 5 can be used to compute all elements y in parallel, provided we have already computed K.
Computing y from u and K is easy. Given an input sequence u E IR and the SSM kernel K E IR', naively using Equation 5 for computing would require ( multiplications. Fortunately, this can be done much more efficiently by observing that for the univariate polynomials
L-1
L-1
K(2)
Σ i=0
Kiz' and u(z) =
Σ u 1=0
ys is the coefficient of ** in the polynomial K(2) u(2), i.e. all ys 's can be computed simultaneously by multiplying two degree L - polynomials. It is well-known that this can be done in O log(L) time via Fast Fourier Transform (FFT) [CLRS09]. We denote this fast computation of Equation 5 via the discrete convolution as
= K u
(6)
Hence, given the SSM kernel K E R_, the output of a discretized state space can be computed efficiently from the input. The challenging part is computing K itself as it involves computing L distinct matrix powers (Equation 4). Similar to S4, instead of directly using A, B, C, our idea is to
an alternate parameterization of state spaces for which it would be easier to compute K.
3 Method
Having stated the necessary background, we now turn to the main contribution of our work.
3.1 Diagonal State Spaces
For general state matrix A it is expensive to compute the matrix exponentials AA for all 0 < k < (Equation 4). However, if A was diagonalizable over and we were given its eigendecomposition A = VAV- where A is diagonal, then we could easily compute e AAAA = Ve A as it is trivial to exponentiate diagonal matrices." In this case Equation 4 simplifies to
CA'B. = CC
()VAV " = (KA) (E IV
Leveraging the above observation, we base our model on the following proposition which asserts that, under mild technical assumptions, diagonal state spaces are expressive as general state spaces. We use the operator K (A, B,C) E ÇI*L to denote the kernel (Equation 4) of length I, for state space (A, B, C.) over C and sample time A.
Proposition 1. Let K clxL be the kernel of length L of a given state space (A, B, C) and sample time A > O, where A E NNN is diagonalizable over with eigenvalues .i...., Ay and Vi, X # 0 and LAA 1. Let P NxI. be Pt,- A, and A be the diagonal matrix with A1...., AN- Then there exist w. ClxN such that
(a) K = KAL(A, (1)EEEE,, w) = W. W - -) elementwise-exp(( (P),
(b) K KAL(A, ((eLAi-- 1)) w) w- A 11, row-softmax(P? (P).
c D0 # = D (AV-Jy V (I 44) V-1 For diag(A),..... Ax) we have A diag(X/ (iag(X),.... Ny) and hence e V. "ddaageepppp (Xt),....,, exp(A))) -V-1
3

---

Input: parameters Are. Aim E Rx", w E C sample time Alog E R (Proposition 1(b))
Output: SSM kernel K R²
I: A - Are + Aim
A - exp (Alog)
A is positive real
2: PNxL (A A)Nx1 (0,1....L. 1lIEE
E- outer product
3: row-softmax, (P)
t- numerically stable ($A.2)
4: K Re( (w/A)IxN-
t- real part"
Algorithm 1: DSS KERNEL (SKETCH)
The proof of Proposition 1 is provided in $A.1. In the above equations, the last equality follows by using Equation 4 to explicitly compute the expression for the kernel of the corresponding diagonal state space. Hence, for any given state space with a well-behaved state matrix there exists diagonal state space computing the same kernel.* More importantly, the expressions for the kernels of the said diagonal state spaces no longer involve matrix powers but only a structured matrix-vector product.
Proposition ((a) suggests that we can parameterize state spaces via A, 7 E CN and simply compute the kernel as shown in Proposition ((a). Unfortunately, in practice, the real part of the elements of A can become positive during training making the training unstable on long inputs. This is because the matrix elementwise-exp( P) contains terms large as exp( A, 1)) which even for modest value of can be very large. To address this, we propose two methods to model diagonal state spaces.
DSS In this variant, we use Proposition ((a) to model our state space but restrict the real part of elements of A to be negative. DSS has parameters Are Aim RN, w CN and Alog RR. A is computed -cemmentwise-cpp exp(Alog) E R= and the kernel is then computed via Proposition ((a).
+ i- Ain where --1 [GGDR22]. A is computed as
DSSExp provides remarkably simple computation of state space kernels but restricts the space of the learned (the real part must be negative).
not clear if such a restriction could detrimental for some tasks, and we now present an alternate method that provides the simplicity of Proposition ((a) while being provably as expressive as general state spaces.
5SS Instead of restricting the elements of A, another option for bounding the elements of elementwise-exp( (P) is to normalize each row by the sum of its elements, which leads us to Proposition 1(b). SSS has parameters Age, Aim E R w E C and Aoog E R. A is computed as Are + 1- Amm- A is computed as (to( RRs,, and the kernel is then computed via Proposition 1(b). A sketch of this computation is presented in Algorithm 1. We note that, unlike over R, softmax can have singularities over C and slight care must be taken while computing it (e.g.. softmax(0, im) is not defined). We instead use corrected version softmax, and detail in $A.2.
3.2 Diagonal State Space (DSS) Layer
We are now ready to describe the DSS layer. We retain the skeletal structure of $4 and simply replace the parameterization and computation of the SSM kernel by one of the methods described in $3.1.
Each DSS layer receives length-L sequence u of H -ddimensional vectors as input, i.c., u R L and produces an output y E IR xL The parameters of the layer are Are. Aiu E R.V. Alog E R and W E C For each coordinate h 1..... H, a state space kernel K, R is computed using of methods described in $3. 1. For example, in case of DSSSOFFMAX Algorithm 1 is used with parameters Are, Aim E R.N. W. e CN and (Aing)A E R. The output y/ E R& for coordinate is computed from 4y E _ and K₁ using Equation 6.
Following S4, we also add residual connection from y followed by GELU non-linearity [HG16]. This is followed by a position- wise linear projection Wout E IR #x'l to enable information exchange among the H coordinates.
It is possible for norms of the parameters of the resulting diagonal state spaces to be much larger than that of the original state space. For example, this occurs the HiPPO matrix [GGR22]. As the inputs in Equation 5 are over IR and we require the outputs to also be over R we explicitly cast the complex- complex-valued valued kernel produced by Proposition I K without affecting the soundness (details in $A.5).

---

The DSS layer be implemented in just few lines of code and PyTorch implementation of soSmmax layer is provided in $A.7 (Figure 6). The implementation DSS layer is even simpler and omitted.
Complexity For batch size B, sequence length and hidden size H. the DSS layer requires L) time and space to compute the kernels, OBBH L log(L)) time for the discrete convolution and L(BH²L) time for the output projection. For small batch size B, the time taken compute the kernels becomes important whereas for large batches more compute is spent on the convolution and the linear projection. The kernel part of DSS layer has 2N H 2HN real-valued parameters.
3.3 Initialization of DSS layer
The performance of state spaces models is known to be highly sensitive to initialization [GGR22]. In line with the past work, we found that carefully initializing the parameters of the DSS layer is crucial to obtain state-of-the-art performance ($4).
The real and imaginary parts of each element of W are initialized from No0, 1). Each element of Doog is initialized as e" where r ~ ((log(.001), (og( A CM is initialized using eigenvalues of the normal part of normal plus low-rank form of HiPPÓ matrix [GGR22]. Concretely, Are. Ain are initialized such that the resulting the vector those N eigenvalues of the following 2N x 2N matrix which have a positive imaginary part.
22 + 12 + 1)1/2/22 i<j
- 1/2
i=
(2i 12) 1)1/2/2
i>
Henceforth, would refer to the above initialization of A Skew-Hippo initialization." In all our experiments, we used the above initialization with N = 64. The initial learning rate of all DSS parameters was 10-3 and weight decay was not applied to them. Exceptions these settings noted in SA.3.
3.4 States of DSS via the Recurrent View
In $3, we argued that it can be slow to compute the output of state spaces via Equation 2 and instead leveraged Equation 5 for fast computation. But in situations such as autoregressive decoding during inference, it is more efficient to explicitly compute the states of a state space model similar to a linear RNN (Equation 2). We now show how to compute the states of the DSS models described in $3.1. Henceforth, we assume that we have already computed A and A.
DSS As stated in Proposition ((a), DSS computes KA,t (A, (1)1siEN, w). For this state space and sample time A, we use Equation to obtain its discretization
diag(e A14
AwA)
B B= (E 18iEN
where diag creates diagonal matrix of the scalars. We can now compute the states using the SSM recurrence Fa Ar + Bug (Equation 2). As A is diagonal, the N coordinates do not interact and hence can be computed independently. Let us assume x and say we have already computed T3-1- Then, for the i'th coordinate independently compute
Fik + 2, (e 1)uk
Fi,k
Note that in DSS Re(A,) and hence | | < 1. Intuitively, if [A//A v 0, we would have Fits -2,-11 and be able to copy history over many timesteps. On the other hand if Re(A,) then 24,4 - ~>, - and hence the information from the previous timesteps would be forgotten similar to a "forget" gate in LSMSs.
Following our work. [GGGR22] formalized the connection between Skew- Skew-Hippo initialization and the original HiPPO matrices and provided theoretical justification for its empirical properties.
5

---

DSSSOOFMAX As stated in Proposition 1(b), DSS computes KA,,(A, ((elris- )EiiE,, w). For this state space and sample time A, obtain the discretization
e: -1 11
= diag(e CANd) B
(g(edhaL- -1) 1ETEN
For the i'th coordinate we can independently compute
" (e
Ti,k A
(e
Let us drop the coordinate index for clarity obtain
ux(e>d 1) Fk c M(eAS1- 1)
where 2k-1 is now a scalar. As the expression involves the term eAAL, where L can be large, directly computing such terms can result in numerical instability and we must avoid exponentiating scalars with positive real part. We make two cases depending on the sign of Re(A) and compute Iy via an intermediate state 3g. follows. Let (Re(A)) > 0J {0, 1). Then,
Adplk (L-1))
(1(11pp)) 1
h = e A1(1-p) p) - + e kAdd uk
Th = Tx
edd(1-2p)I
The above equation be parsed as follows. If Re(A) & 0 then have
(ets - 11
AA_ FK-1 + Uk
TK
Ts A(e
and hence if Re(A) < then Te & uy and xa / X. this case, information from the previous timesteps would be ignored and the information used would be local. On the other hand if Re(A) > 0 then we would have
AA(k- (I-1)
AA
x Jx- e- KAA
Ik Z4
XAL
and hence if Re(A) >> Hence, the model would be able copy information even from extremely distant positions, allowing it to capture long-range dependencies.
then To uo and Hn s ãx-1 & uo. IkcL-1 & 0 and Ø1-1 R uo/A
Experiments
We evaluate the performance of DSS on sequence-level classification tasks over text, images and audio. Overall, we find its performance is comparable to S4.
Long Range Arena (LRA) LRA ITDA [I] is a standard benchmark for assessing the ability of models process long sequences. LRA contains 6 tasks with diverse input lengths 1K-16K, encom- passing modalities such as text and images. Several Transformer variants have been benchmarked on LRA but all have underperformed due to factors such their high compute-memon requirements, implicit locality bias and inability to capture long-range dependencies.
Table 1 compares DSS against S4, the Transformer variants reported in [TDA+21] as well as follow- up work. State space models (S4, DSS) shown in Table are left-to-right unidirectional whereas other models could be bidirectional.
Despite its simplicity, DSS delivers state-of-the-art performance on LRA. Its performance compa- rable to that of S4, with a modest improvement in test accuracy averaged across the 6 tasks (81.88 vs 80.21).7 In addition, DSS maintains a comfortable 20 point lead over the best performing Transformer variant (88.88 vs 61.41).
The large gap between S4 and DSS on TEXT is due to the use of a larger learning rate for Atog in DSS. For our S4 runs, decided to the official hyperparameters provided by [GGR22].

---

MODEL
LiSTOPS
TEXT
RETRIEVAL
IMAGE
PATHFINDER
PATH-X
AVG
(input length)
2000
2048
4000
1024
1024
16384
Transformer
36.37
64.27
57.46
42.44
71.40
X
53.66
Reformer
37.27
56.10
53.40
38.07
68.50
X
BigBird
36.05
64.02
59.29
40.83
74.87
X
54.17
Linear Trans.
16.13
65.90
53.09
42.34
75.30
X 50,46
Performer
18.01
65.40
53.82
42.77
77.05
X 51.18
FNet
35.33
65.11
59.61
38.67
77.80
X 54.42
Nyströmformer
37.15
65.52
79.56
41.58
70.94
X 57.46
Luna-256
37.25
64.57
79.29
47.38
77.72
X
59.37
H- 1--Tranfforme-ID
49.53
78.69
63.99
46.05
68.78
X
61.41
S4 (as [GGR22])
58.35
76.02
87.09
87.26
86.05
88.10
80.48
S4 (our run)
57.6
75.4
87.6
86.5
86.2
88.0
80.21
DSS (ours)
60.6
84.8
87.8
85.7
84.6
87.8
81.88
DSS (ours)
59.7
84.6
87.6
84.9
84.7
85.6
81.18
RXS (ours)
59.3
82.4
86.0
81.2
81.3
65.03
Table I: (Long Range Arena) Accuracy on the full suite of LRA tasks. (Top) Transformer variants reported in LRA, (Middle) other long-range models reported in the literature, (Bottum) state space models.
MODEL
MFCC
RAW
MODEL
MFCC
RAW
Transformer
90.75
X
WaveGAN-D
96.25
Performer
80.85
30.77
S4 (as in [GGR22])
93.96
98.32
ODE-RNN
65.9
X
S4 (our run)
98.1
NRDE
16.49
ExpRNN 82.13
11.6
DSSSOFTMAX (ours)
97.7
LipschitZRNN
88.38
X
DSS (ours)
98.2
CKConv 95.3
71.66
DSS EA--SSADD-S (ours)
97.7
Table (Speech Commands (SC)) Transformer, CTM. RNN. CNN, SSM models. (MFCC) Standard pre-processed MFCC features (ennth-1611) (Raw) Unprocessed signals (length- 16000). denotes not applicable computationally infeasible single GPU, All results except ours are reproduced from [GGR22).
Interestingly, despite being less expressive than DSS DSS also reports an impressive performance which makes it specially appealing given its simplicity during training and inference. We investigate an even simpler version of DSS denoted as DSS SC-LE which same DSS except that omit the term A (e 1) in the expression of the kernel shown in Proposition ((a) and instead compute it as i elementwise-expccC performance of DSSEXP-NO-SCALE generally inferior to that of DSS with the model failing on the challenging PATII-X X task which requires the model to capture extremely long-range interactions.
As shown in Table 1, the
Raw Speech Classification Audio is typically digitized using a high sampling rate resulting in very long sequences. This provides interesting domain for investigating the abilities of long-range models. We evaluate the performance of DSS on the Speech Commands (SC) dataset [Warl 8), consisting raw audio samples of length 16000. modeled as 10-way classification task. As shown in Table 2, the performance of all DSS variants is comparable to that of S4 (98.2 vs 98.1).
In all experiments and ablations, $4 and DSS identical model hyperparameters such hidden size, number of layers, etc. Our experimental setup was built on top of the training framework provided by the $4 authors and for $4 runs followed their official instructions." Details about model initialization, and hyperparameters are provided in $A.3.
4.1 Analyzing the Performance DSS
While the experimental results presented above encouraging, and clearly demonstrate the effec- tiveness DSS at modeling long-range dependencies, it is not clear what exactly are the main factors

---

MODEL
LISTOPS
TEXT
IMAGE
PATHFINDER PATH-X
SC
Non-State-Space best
49.5
78.7
47.4
77.8
X
96.5
60,6
84.8
85.7
84.6
87.8
97.7
DSS DOSSOO (random init) D0STONN (kernel length 128)
$7.6
80.54
72.13
77.5
X
96.8
51.6
75.41
83.8
65.1
X
96.7
DSSSOFTMAX (argmax 7. K -9--eeecettlle)
89
124
66
87
1269
81
Table 3: (Top) Ablation analysis of SOS "Non-State Space best" is the best performance among models in top and middle sections of Table 1. (Bottom) "argmax [K&] [ 1995-percentile" is the 95th percentile of fargmaxugk<<l [&] K is kernel in DSS } and is explained in $4.2.
contributing to its performance. To investigate this further, we performed an ablation analysis aimed at answering the following questions:
How significant is the Skew- Hiipo initialization ($3.3) to the model performance? Would initializ- ing A randomly work just as well?
the main source of superior performance of state space models (S4, DSS), compared previous models, their ability to model rongrrange dependencies? Would restricting DSS to only model local interactions hurt its performance on the above tasks?
Random Initialization To answer the first question, repeated the experiments after initializing each element of Are and Aim in DSS by randomly sampling from (O, 1).
Truncated Kernels To answer the second question, instead of constructing kernel of length equal to the length of the input, we restricted the length the kernel constructed DSS (Algorithm 1) to 128, significantly shorter than the length of the input. To understand the implication of this restriction recall Equation which states that yk
- K-- Uk-
For a given context size c 128, restricting K sc = 0 would imply
min(k.c-1)
yk Σ
K3 uk-
and hence the output yk at position would only depend on Ug. Uk-o+1. This would restrict each DSSSoFrMAx layer to only model local interactions and the model would require several layers have broader receptive field and capture long-range interactions.
As shown in Table 3, randomly initializing the A parameters of DSS leads to significant performance degradation the majority of tasks, with the model failing to perform on PATH-X. This is inline with the findings of [GGR22] who also reported the initialization of S4 to be critical to its performance. Interestingly, despite this performance reduction, DSS manages to outperform all non state-space- based models on every task.
Truncating the length of the kernel also leads to significant reduction in performance across most tasks (Table 3), suggesting that the superior performance state-space models on these tasks indeed be partly attributed to their ability to capture long-range dependencies. Moreover, some tasks such as LISTOPS and IMAGE, using a truncated kernel still manages to outperform all Transformer variants, which is surprising Transformer layers are known to be effective at capturing interactions at such short ranges.
4.2 Analysis of Learned DSS Parameters
To further explore the inner workings of DSS, we visually inspected the trained parameters and kernels of DSS
The kernels of all layers of the trained DSS are shown in Figure and reveal a stark contrast between the tasks. On the tasks IMAGE and SC, for almost all kernels, the absolute values of the first 128 positions are significantly higher than the later positions indicating that these kernels are mostly local. On the other hand for PATH-X, for a significant proportion of kernels the opposite true, indicating that these kernels are modeling long-range interactions.

---

listops
image
500
1000
400
1000
2000
02
00
200
400
800
800
1000
200
400
600
800
1000
200
400
600 pathx
800
1000
pathfinder
SC
num-layers loyers hidden
800
600
600
1000
1000
1800
200
400
800
820
1000
200
400
800
800
1000
200
400
800
800
1000
Figure Kernels trained SOSSOOAMA Each row pixels corresponds one the (numher-of- x H) kernels. For a row (kernel) K & R the element Kg E is shown as | | to enable visualization. To enable comparison across tasks, only the first 1024 positions are shown.
This partially explains why the performance of DSS SOFTMAX IMAGE and SC does decrease significantly even after limiting the kernel length to 128, whereas PATHFINDER and PATH-X report significant degradation (Table 3). The last row of Table shows the 95th percentile of the fargmaxngt<< JKl K is a kernel in DSS
i.e., the set of positions over all kernels, at which for given kernel its absolute value is maximized. As expected, most kernels of IMAGE and SC are local whereas for PATH-X they are mostly long range.
The values of the real and imaginary parts of A shown in Figure (and $A.6, Figure 5). We observe that, although all elements of A₁o are initialized as -0.5, their values can change significantly during training and become positive (see plots for LISTOPS, SC). It can also be observed that with more negative Re( () generally tends to have a larger Imn A), which interestingly a property satisfied by Skew-Hippo initialization.
listops
image
-1.D
-0.5
0.0
-1.5
-1.0
-0.5
0.0
-2
real
pathfinder
pathx
layer
0
3
5
-3
-2
roal
roal
Figure 3: Values of trained DSSSOPTMAX for tasks described in $4. For À x + it., we show y on
loglscale for better visualization by plotting À as (a, arcsinh(y/2)) (a, log(1/2 /w/2)² 1)). Black dots correspond to Skew-Hippo initialization (83.3).
The values of the trained Alog corresponding to the real and imaginary parts of A are shown in Figure 4 and, this case as well, change significantly during training. The values of A on short-range tasks such as IMAGE and SC generally tend to be larger compared to long-range tasks such as PATHFINdeR and PATH-X, inline with the intuition that, when Re(A) < 0, smaller values of A correspond to modeling long-range dependence ($3.4).

---

listops
retrieval
06018900
000010000 0H8.CER1011000
00 00000
layor
-10.0
-7.5
pathfinder
pathx
imag
-12
layor
-B
-4
Figure 4: Trained Atog in DSS for tasks described in $4. As noted in $A.3 we used separate Atog.re A im E R parameters to respectively scale the real and imaginary parts of A.
We note that the plot for LISTOPS reveals an outlier with value of 22 which after exponentiation in Algorithm would result in an extreme large A. This can potentially lead to training instabilities and we plan address this issue in future work.
Limitations and future work In this work, we evaluated DSS on sequence-lever classification tasks. In future work, plan to include token-level generation tasks such language modeling, forecasting, etc. Another important future direction is to pretrain DSS-based models on large amounts of raw data. Lastly, found that the initialization and learning rates of DSS parameters Are- Aim- Aing ($3.3) play an important role in the performance and convergence of the model. Informally, for tasks such as PATH-X X that require very long-range interactions, a smaller initialization of Doog was beneficial. An in-tepth analysis of this phenomenon could be helpful and remains for future work.
Acknowledgments and Disclosure of Funding
We thank Ramon Fernandez Astudillo for carefully reviewing the preliminary draft and suggesting several helpful edits. We thank Omer Levy, Achille Fokoue and Luis Lastras for their support. Our experiments were conducted on IBM's Cognitive Computing Cluster, with additional resources from Tel Aviv University. This research was supported by (1) IBM AI Residency program and (2) De- fense Advanced Research Projects Agency (DARPA) through Cooperative Agreement D20AC00004 awarded by the U.S. Department of the Interior (DOI), Interior Business Center, and (3) the Euro- pean Research Council (ERC) under the European Union Horizons 2020 research and innovation programme (grant ERC DELPHI 802800).
References
[BMH* + 22]
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean- Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego De Las Casas, Aurelia Guy. Jacob Menick. Roman Ring. Tom Hennigan, Saffron Huang. Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack Rac, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato. editors, Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings Machine Learning Research, pages 2206-2240. PMLR, 17-23 Jul 2022.
[BMR 20] Tom B. Brown, Benjamin Mann. Nick Ryder. Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranay Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voo Voss. Gretchen Krueger. Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott
10

---

Gray, Benjamin Chess, Jack Clark, Christopher Bemer, Sam McCandlish, Alec Radford, Tyya Sutskever, and Dario Amodei. Language models are few-shot learners. Hugo Larochllle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Ad- vances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[BPC20] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- transformer. ArXiv preprint, abs/2004.05 1500 2020.
[BZMA20] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A frame- work for self supervieeei learning of speech representations. Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Sysiems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020.
[CFG+21] Benjamin Charlier, Jean Feydy, Joan Alexis Glaunès, Camçois-Davite Collin, and Ghislain Durif. Kemel operations on the gpu, with autodiff, without memory overflows. Journal of Machine Leurning Reseurch, 22(74):1-6, 2021.
[CLD*21] Krzysztof Marein Choromanski, Valerii Likhosherstov. David Dohan, Xingyou Song, Andreca Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with per- formers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[CLRS09] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. Introduction to Algorithms. The MIT Press, 3rd edition, 2009.
[DXX18] Linhao Dong, Shuang Xu, and Bo Xu. Specch-transfommer: no-recurrence pequence-to-sequeee model for speech recognition. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2018, Calgary, AB, Canada, April 15-20, 2018, pages 5884-5888. IEEE, 2018.
[Gaj03] Zoran Gajic. Linear dynamic systems and signals. Prentice Hall/Pearson Education Upper Saddle River, 2003.
[GB20] Ankit Gupta and Jonathan Berant. GMAT: Global memory augmentation for transformers. ArXiv preprint, abs/2400.03377, 2020.
[GDE*20) Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Hugo Larochelle. Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020.
[1] Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant, Memory-efficee transformers via top-k attention. In Proceedings of the Second Workshop on Simple and Efficient Natural Language Processing, pages 39 52, Virtual, November 2021. Association for Computa- tional Linguistics.
[GGDR22] Karan Goel, Albert Gu, Chris Donahue, and Christopher Re. It's raw! Audio generation with state-space models. In Kamalika Chaudhuri, Stefanic Jegelka, L. Song, Csaha Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 7616-7633. PMLR, 17-23 Jul 2022.
[GGGR22] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. ArXiv, abs/2206. 11893, 2022.
[GGR22] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference Learning Representations (ICLR), 2022.
[GJG*21] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous time models with linear state-space layers. Advances in neural information processing systems, 34, 2021.
11

---

[HG16] Dan Hendryeks and Kevin Gimpel. Gaussian error linear units (gelus). ArXiv preprint, Abs/1606.08415, 2016.
[KKL20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[KVPF20] Angelos Katharopoulos, Apeory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns:: Fast autoregressive transformers with linear attention. Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5156-5165, PMLR, 2020.
(QML+20) Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. Findings of the Association for Computational Online, 2020. Association for Computational Lin-
Linguistics: EMNLP 2020, pages guistics.
[RKH*21] Alec Radford. Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh. Sandhini Agarwal, Girish Sastry. Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Iya Sutskever. Learning transferable visual models from natural language supervision. Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 8748 8763. PMLR, 2021.
[RPG*21] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 1-288831. PMLR, 2021.
[RSVG21] Aurko Roy; Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
[SSI*22] Uri Shaham, Elad Segal, Maor Ivgi. Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. Scrolls: Standardized comparison over long language sequences. In Proceedings the 2022 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2022.
[TDA 21] 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 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[TDBM20] Yi Tay, M. Dehghani, Dara Bahri, and Donald Metzler, Efficient transformers: A survey. ArXiv, abs/2009.06732, 2020.
[VKE19] Aaron Voelker, Ivana Kajic, and Chris Eliasmith. Legendre memory units: Continuous-time representation in recurrent neural networks. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence Alché-Buue, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 15544-15553, 2019.
[VSP* 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need, In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors. Advances in Neural Information Processing Systems 30: Annual Conference on Newral Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008, 2017.
[War18] Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. ArXiv. abs/1804.002000 2018
[WLK +2]
Sinong Wang. Belinda Z. Li. Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. ArXiv, abs/2006.00 047668 2020.
12

---

[7S21] Zhenhai Zhu and Radu Soricut. H-transformer-ID Fast one-dimensional hierarchical attention for sequences. 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), pages 3801-3815, Online, 2021. Association for Computational Linguistics.
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 end of $4.2.
(c) Did you discuss any potential negative societal impacts of your work? [No]
(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 set of assumptions all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See $A.1.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed reproduce the main ex- perimental results (either in the supplemental material or as a URL)? [Yes] See $1, SA.3.
(b) Did you specify all the training details (e.g., data splits, byperparameters, how they were chosen)? [Yes] See $A.3.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [No]
(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] See $A.33.
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? [No]
(c) Did you include any new assets either in the supplemental material as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
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 to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "dyBxYbIq243", "reviewer": "Reviewer_NmPw", "summary": "The paper builds upon the recently proposed S4 architecture, which was shown to be effective at modeling long-range dependencies. The authors propose simplifications to the S4 model by modifying the kernel such that, instead of the originally proposed diagonal-plus-low-rank kernel, they can remove the low rank correction. The theoretical contribution includes proving a proposition which says that under some mild assumptions, this approximation is equivalent to the original kernel formulation. Empirically, the authors show that it is competitive with the original S4 model on the Long Range Arena (LRA) benchmark.", "questions": "1. In the last paragraph of Section 2, the authors write that “our idea” is to use an alternate parameterization of state spaces. Isn’t this the original idea in S4?\n2. In Section 4.2 (paragraph 4), the authors note that during training, the parameters $\\text{Re}(\\lambda)$ and $\\text{Im}(\\lambda)$ move away from the skew-Hippo initialization. Does this mean that this is perhaps not the best initialization for these parameters?\n", "limitations": "The authors have discussed the limitations of their work which is sufficient.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The S4 model itself is quite new and interesting, and as such, further studies to analyze and simplify it are useful for the community at large.\n2. The 2 methods (DSS-exp and DSS-softmax) are theoretically motivated and demonstrate good empirical performance, where they achieve results comparable to the original S4 on most tasks.\n3. The paper is clear and well-written, and the background described in Section 2 is sufficient for new readers (not familiar with S4) to understand the paper. \n\nWeaknesses:\n1. My primary concern is that “simplicity” is subjective. The authors motivate DSS by saying that it simplifies computation of the SSM kernel. It would be a stronger result if they can show that it also provides quantitative gains, perhaps in terms of training time/flops.\n2. Proposition 1 is based on the assumption that the state matrix $A$ is diagonalizable, and the authors say that this is “a mild technical assumption.” Can they justify this assumption, for instance, by showing that it holds for some toy example where explicit computation of the state matrices is feasible?\n3. In Table 1, the large gap between S4 and DSS on the “TEXT” task is attributed to learning rate tuning. Could the authors also try to tune the baseline learning rate for S4, since these numbers can significantly change the “AVG” result?\n4. In Section 4.1, the authors try to analyze the effect of restricting DSS to local interactions. As is clear from Table 3, this would be very task-dependent for any model (not just S4), and it is not clear what is the insight gained from this experiment.\n5. In Section 4.2, the discussion in the last 3 paragraphs refer to figures in the Appendix. I think the main paper should be self-contained. Perhaps the authors could move the figures to the main paper by reducing some of the “Related work” content in Section 5.\n\n**Update (Aug 8):** After discussions and new time benchmarking results from the authors, I am increasing my rating below. \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": "LluHel8dF1b", "reviewer": "Reviewer_riDD", "summary": "The paper extends prior work on efficient state space models by simplifying the state transition matrix to be diagonal. The authors make the point that diagonal state space models (DSSMs) are almost as expressive as normal state spaces when considering complex diagonal matrices. This is due to the fact that almost all matrices can be Eigen-decomposed (diagonalized) in the complex plane. The paper introduces 2 parameterizations of DSSMs that are based on this insight and show strong empirical performance on the long-range arena as well as raw audio classification; tasks which require processing very long sequences. \n\nOverall, I think the contributions of this paper a very important for the evolution of the DSSM framework wrt Deep Learning. I have trouble with the presentation of the work though which should be addressed by the authors. It is harder to follow than it should be and the community would benefit from improvements on that end. I am happy to change my score after that accordingly.", "questions": "* Background discretization: There is no citation on where this comes from.\n* Background: Is the paragraph about computing y and u necessary? It breaks the reading flow and adds mental overhead. Maybe push to the appendix.\n* Proposition 1: A is not in C, is it?\n* Proposition 1: I fail to see how the resulting K is ensured to be real?\n* Initialization of delta_log: Why  is delta_log initialised as e^r with r in U(log(.001), log(.1))? An explanation in the paper would help.\n* Initialization of A: Why the eigen-values of only the normal part of the Hippo matrix. What's the reasoning behind this decision?\n* Figure 5: I don’t understand what are the delta_log values and how they refer to the real and imag parts of the lambda. Please explain and also extend the caption.\n* Figure 5: It is stated that the values of delta_log are larger on short-range tasks? How is this “inline with” larger values modeling long-range dependencies?\n* It seems like these models need to be tweaked quite a bit and for a practitioner it would be super helpful to see how brittle these models still are. That is, I would love to see a table/graph of results with non-optimal hparams.\n\n", "limitations": "Limitations have been addressed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n* Mostly well written\n* very practical insight and simplification\n* strong empirical results\n* insightful analysis of parameters after training\n\n**Weaknesses**\n\nI only found one weakness but it has to be addressed (IMO):\n*Clarity*: I find the main contribution in proposition 1 very hard to follow. This is due to the fact that it is posited without any derivation or intuition, leaving the \"elementary\" proof for the appendix. I think this is bad practice and puts the burden on the reader to gather all necessary background to understand how the proposition comes together. Given that this is so central to the paper I do not understand the choice of the authors to put the whole derivation of it in the appendix.\n\nStating the following at least on a high-level before the proposition would make it much easier to understand proposition 1:\n1) Most matrices diagonalize over the complex plane (via eigen-decomposition) —> Assume A is one of those —> almost no loss of expressibility.\n2) Matrix exponentials are trivial to compute for diagonalised matrices because (VDV^-1)^n = V D^n V^-1. Analogously this is true for power series e^A.\n3) V and V^-1 can be subsumed into B and C, show the resulting equation, show that B, C, V, V^-1 can be squashed into a single vector (w). \n\nThis would go a long way to make the main paper much more understandable and helping the reader understand where Proposition 1 comes from.\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": "Gwo-7YUrZWz", "reviewer": "Reviewer_zKuy", "summary": "The paper builds on structured state space (S4) models and shows that one can achieve similar or slightly better performance with just diagonal state matrices (avoiding the low-rank structure). This is interesting because then the model becomes conceptually simpler and easy to implement. In addition, the paper provides ablations covering initialization and kernel truncation, and an analysis of the learned parameters.", "questions": "Questions/suggestions:\n- I think the abstract could be largely improved. Some things to consider are: mentioning complexities, mentioning ablation experiments, mentioning analysis, etc.\n- Line 17: \"denoising objective\" --> This is a bit misleading, even more now that denoising diffusion models are in hype. Please find a better terminology.\n- Line 39: \"DLPR\" --> DPLR.\n- End of Sec 2: I think it would be better to present here the two views (convolutional and recurrent) instead of just one and then later mentioning the recurrent one.\n- Sec. 3: In general, I found it hard to follow with the notation used (and it is not wrong, it is just that it is not presented appropriately). $N$, $H$, and $L$ could be defined beforehand or in Sec 2 (and also give an idea of their relation: like $H<N<L$ if that was the case). I also think that there are some missing \"squeeze\" and \"unsqueeze\" operations in the algorithm in the Appendix.\n- Algorithm 1: Should there not be an outer product between $\\Delta$ and $\\Lambda$? Which type of division represents $/$ in line 4? Mention $H=1$?\n- Line 137: Isn't $W_{\\text{out}}$ the same as $C$ in Eq. 2? (I'm mostly asking for a clarification in the paper)\n- Line 152: Any intuition on the initialization of $r$? (same)\n- Lines 153-156: So in the end the HiPPO initialization is crucial (according to presented results). Therefore, the model could be considered to still rely on HiPPO theory? More explanation on that would be welcome (in the paper).\n- I find there are many differences between the training setups of S4 and DSS (and this can make the reader suspicious). Perhaps it would be instructive to report a dry-run of DSS without those differences (that is, using the S4 setup unaltered) in addition to the reported results.\n- Line 192: \"images, audio\" --> images, and audio.\n- Sec 4.1: I think that the section would be much easier to parse if the two ablations were presented one after the other. So first Random Initialization: Question, how, result, and then Truncated Kernels: Question, how, result.\n- Line 256: \"absolute values\" --> Better write \"magnitudes\"?\n- Line 280: Remove \"Related work\".\n\nStylistic issues:\n- Do not put references in the abstract.\n- I am not sure that the way to cite references is correct (not sure about NeurIPS policy here; typically one writes \"Surname et al.\" or just numbers).\n- Authors seem to be using \"cf.\" to mean \"see\". Note that \"cf.\" does not mean \"see\": https://en.wikipedia.org/wiki/Cf. (that is, cf. should be equivalent to \"compare\", but the way the authors write it I believe it should be \"see\").\n- A similar thing happens with \"provably\", where I think they mean \"probably\". It is not the same and the former implies that it can be proven mathematically: https://wikidiff.com/provably/probably\n", "limitations": "I am happy with the section on limitations.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I like the paper and recommend acceptance. I find it is moderately original, of very good quality, and significant. Clarity can be improved before the final submission.\n\nStrengths:\n* The idea is simple and attractive. Simplifications of existing processes is always welcome, especially if they match the accuracy of the previous, more complicated approach.\n* I think complexity is reduced. This is perhaps not emphasized enough in the abstract and introduction.\n* Good set of evaluations.\n* Insightful ablations and analysis.\n\nWeaknesses:\n* Presentation. At several points I found that the authors could have made a larger effort in polishing notation or structuring the paper (see some comments below).\n* Speed/complexity is only partially assessed. I'd suggest that the authors compare directly with complexities of the original S4 in the small subsection of section 3.2. Alternatively (or in addition), some speed differences can be measured on the same machine.\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": {}}]

Author rebuttal:

Decision: Accept