Original: Diagonal State Spaces are as Effective as Structured State Spaces
Revised: Diagonal State Spaces are as Effective as Structured State Spaces

Original: Abstract
Revised: Abstract

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In an exciting result, Gu et al.
Revised: In an exciting result, Gu et al.

Original: (2022) proposed the Structured State Space (S4) architecture delivering large gains over state-of-the-art models on several long-range tasks across various modalities.
Revised: (2022) proposed the Structured State Space (S4) architecture delivering large gains over state-of-the-art models on several long-range tasks across various modalities.

Original: The core proposition of S4 is the parameterization of state matrices via a diagonal plus low rank structure, allowing efficient computation.
Revised: The core proposition of S4 is the parameterization of state matrices via a diagonal plus low rank structure, allowing efﬁcient computation.

Original: 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.
Revised: 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.

Original: 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.
Revised: Our Diagonal State Space (DSS) model matches the performance of S4 on Long Range Arena tasks, speech classiﬁcation on Speech Commands dataset, while being conceptually simpler and straightforward to implement.

Original:  Introduction
Revised:  Introduction

Original: The Transformer architecture [VSP ` 17] has been successful across many areas of machine learning.
Revised: The Transformer architecture [VSP ` 17] has been successful across many areas of machine learning.

Original: Transformers pre-trained on large amounts of unlabelled text via a denoising objective have become the standard in natural language processing, exhibiting impressive amounts of linguistic and world knowledge [BMR ` 20, BMH ` 22].
Revised: Transformers pre-trained on large amounts of unlabelled text via a denoising objective have become the standard in natural language processing, exhibiting impressive amounts of linguistic and world knowledge [BMR ` 20, BMH ` 22].

Original: This recipe has also led to remarkable developments in the areas of vision [RPG ` 21, RKH ` 21] and speech [DXX18, BZMA20].
Revised: This recipe has also led to remarkable developments in the areas of vision [RPG ` 21, RKH ` 21] and speech [DXX18, BZMA20].

Original: The contextualizing component of the Transformer is the multi-head attention layer which, for inputs of length L  has an expensive Ω p L 2 q complexity.
Revised: The contextualizing component of the Transformer is the multi-head attention layer which, for inputs of length L , has an expensive ⌦ p L 2 q complexity.

Original: This becomes prohibitive on tasks where the model is required to capture long-range interactions of various parts of a long input.
Revised: This becomes prohibitive on tasks where the model is required to capture long-range interactions of various parts of a long input.

Original: To alleviate this issue, several Transformer variants have been proposed with reduced compute and memory requirements [QML ` 20, KKL20, RSVG21, BPC20, GB20, KVPF20, WLK ` 20, CLD ` 21, ZS21, GDG ` 21] (see [TDBM20] for a survey).
Revised: To alleviate this issue, several Transformer variants have been proposed with reduced compute and memory requirements [QML ` 20, KKL20, RSVG21, BPC20, GB20, KVPF20, WLK ` 20, CLD ` 21, ZS21, GDG ` 21] (see [TDBM20] for a survey).

Original: Despite this effort, all these models have reported inadequate performance on benchmarks created to formally evaluate and quantify a model’s ability to perform long-range reasoning (such as Long Range Arena
Revised: Despite this effort, all these models have reported inadequate performance on benchmarks created to formally evaluate and quantify a model’s ability to perform long-range reasoning (such as Long Range Arena

Original: [TDA ` 21] and SCROLLS [SSI ` 22]).
Revised: [TDA ` 21] and SCROLLS [SSI ` 22]).

Original: In a recent breakthrough result, Gu et al.
Revised: In a recent breakthrough result, Gu et al.

Original: [GGR22] proposed S4, a sequence-to-sequence model that uses linear state spaces for contextualization instead of attention.
Revised: [GGR22] proposed S4, a sequence-to-sequence model that uses linear state spaces for contextualization instead of attention.

Original: It has shown remarkable performance on tasks requiring long-range reasoning in domains such as text, images and audio.
Revised: It has shown remarkable performance on tasks requiring long-range reasoning in domains such as text, images and audio.

Original: For instance, on Long Range Arena it advances the state-of-the-art by 19 accuracy points over the best performing Transformer variant.
Revised: For instance, on Long Range Arena it advances the state-of-the-art by 19 accuracy points over the best performing Transformer variant.

Original: Its remarkable abilities are not limited to text and images and carry over to tasks such as time-series forecasting, speech recognition and audio generation [GGDR22].
Revised: Its remarkable abilities are not limited to text and images and carry over to tasks such as time-series forecasting, speech recognition and audio generation [GGDR22].

Original: Despite S4’s achievements, its design is complex and is centered around the HiPPO theory, which is a mathematical framework for long-range modeling [VKE19, GDE ` 20, GJG ` 21].
Revised: Despite S4’s achievements, its design is complex and is centered around the HiPPO theory, which is a mathematical framework for long-range modeling [VKE19, GDE ` 20, GJG ` 21].

Original: [GGR22] showed that state space models with various alternative initializations perform poorly in comparison to initializing the state-space parameters with a particular HiPPO matrix.
Revised: [GGR22] showed that state space models with various alternative initializations perform poorly in comparison to initializing the state-space parameters with a particular HiPPO matrix.

Original: In order to leverage this matrix, they parameterize the learned state spaces using a 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 output efficiently, making S4 difficult to understand, implement and analyze.
Revised: In order to leverage this matrix, they parameterize the learned state spaces using a 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 output efﬁciently, making S4 difﬁcult to understand, implement and analyze.

Original: In this work, we show that it is possible to match S4’s performance while using a much simpler, fullydiagonal parameterization of state spaces.
Revised: In this work, we show that it is possible to match S4’s performance while using a much simpler, fullydiagonal parameterization of state spaces.

Original: While we confirm that random diagonal state spaces are 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.
Revised: While we conﬁrm that random diagonal state spaces are 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.

Original: 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 as expressive as general state spaces.
Revised: Leveraging this idea, our proposed Diagonal State Space (DSS) model enforces state matrices to be diagonal, making it signiﬁcantly simpler to formulate, implement and analyze, while being provably as expressive as general state spaces.

Original: In contrast to S4, DSS does not assume any specialized background beyond basic linear algebra and can be implemented in just a few lines of code.
Revised: In contrast to S4, DSS does not assume any specialized background beyond basic linear algebra and can be implemented in just a few lines of code.

Original: Our implementation fits in a single page and is provided in §A.7 (Figure 6).
Revised: Our implementation ﬁts in a single page and is provided in §A.7 (Figure 6).

Original: We evaluate the performance of DSS on Long Range Arena (LRA) which is a suite of sequencelevel classification tasks with diverse input lengths ( 1 K - 16 K ) requiring similarity, structural, and visual-spatial reasoning over a wide range of modalities such as text, natural/synthetic images, and mathematical expressions.
Revised: We evaluate the performance of DSS on Long Range Arena (LRA) which is a suite of sequencelevel classiﬁcation tasks with diverse input lengths ( 1 K - 16 K ) requiring similarity, structural, and visual-spatial reasoning over a wide range of modalities such as text, natural/synthetic images, and mathematical expressions.

Original: Despite its simplicity, DSS delivers an average accuracy of 81 . 88 across the 6 tasks of LRA, comparable to the state-of-the-art performance of S4 ( 80 . 21 ).
Revised: Despite its simplicity, DSS delivers an average accuracy of 81 . 88 across the 6 tasks of LRA, comparable to the state-of-the-art performance of S4 ( 80 . 21 ).

Original: In addition, DSS maintains a comfortable 20 point lead over the best Transformer variant ( 81 . 88 vs 61 . 41 ).
Revised: In addition, DSS maintains a comfortable 20 point lead over the best Transformer variant ( 81 . 88 vs 61 . 41 ).

Original: In the audio domain, we evaluate the performance of DSS on raw speech classification.
Revised: In the audio domain, we evaluate the performance of DSS on raw speech classiﬁcation.

Original: On the
Revised: On the

Original: Speech Commands dataset [War18], which consists of raw audio samples of length 16 K , we again found the performance of DSS to be comparable to that of S4 ( 98 .
Revised: Speech Commands dataset [War18], which consists of raw audio samples of length 16 K , we again found the performance of DSS to be comparable to that of S4 ( 98 .

Original: 2 vs 98 . 1 ).
Revised: 2 vs 98 . 1 ).

Original: To summarize, our results demonstrate that DSS is a simple and effective method for modeling long-range interactions in modalities such as text, images and audio.
Revised: To summarize, our results demonstrate that DSS is a simple and effective method for modeling long-range interactions in modalities such as text, images and audio.

Original: We believe that the effectiveness, efficiency and transparency of DSS can significantly contribute to the adoption of state space models over their attention-based peers.
Revised: We believe that the effectiveness, efﬁciency and transparency of DSS can signiﬁcantly contribute to the adoption of state space models over their attention-based peers.

Original: Our code is available at https://github.com/ag1988/dss .
Revised: Our code is available at https://github.com/ag1988/dss .

Original:  Background
Revised:  Background

Original: We start by reviewing the basics of time-invariant linear state spaces.
Revised: We start by reviewing the basics of time-invariant linear state spaces.

Original: State Spaces
Revised: State Spaces

Original: A continuous-time state space model (SSM) parameterized by the state matrix
Revised: A continuous-time state space model (SSM) parameterized by the state matrix

Original: A P R N ˆ N and vectors B P R N ˆ 1 , C P R 1 ˆ N is given by the differential equation:
Revised: A P R N ˆ N and vectors B P R N ˆ 1 , C P R 1 ˆ N is given by the differential equation:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: which defines a function-to-function map u p t q ÞÑ y p t q .
Revised: which deﬁnes a function-to-function map u p t q ﬁ Ñ y p t q .

Original: For a given value of time t P R , u p t q P R denotes the value of the input signal u , x p t q P R N ˆ 1 denotes the state vector and y p t q P R denotes the value of the output signal y .
Revised: For a given value of time t P R , u p t q P R denotes the value of the input signal u , x p t q P R N ˆ 1 denotes the state vector and y p t q P R denotes the value of the output signal y .

Original: We call a state space diagonal if it has a diagonal state matrix.
Revised: We call a state space diagonal if it has a diagonal state matrix.

Original: Discretization For a given sample time ∆ P R ą 0 , the discretization of a continuous state space (Equation 1) assuming zero-order hold 2 on u is defined as a sequence-to-sequence map from p u 0 , . .
Revised: Discretization For a given sample time  P R ° 0 , the discretization of a continuous state space (Equation 1) assuming zero-order hold 2 on u is deﬁned as a sequence-to-sequence map from p u 0 , . .

Original: , u L ´ 1 q “ u P R L to p y 0 , . . .
Revised: , u L ´ 1 q “ u P R L to p y 0 , . . .

Original: , y L ´ 1 q “ y P R L via the recurrence,
Revised: , y L ´ 1 q “ y P R L via the recurrence,

Original: [Equation]
Revised: [Equation]

Original: Assuming x ´ 1 “ 0 for simplicity, this recurrence can be explicitly unrolled as
Revised: Assuming x ´ 1 “ 0 for simplicity, this recurrence can be explicitly unrolled as

Original: [Equation]
Revised: [Equation]

Original: For convenience, the scalars CA k B are gathered to define the SSM kernel K P R L as
Revised: For convenience, the scalars CA k B are gathered to deﬁne the SSM kernel K P R L as

Original: [Equation]
Revised: [Equation]

Original: where the last equality follows by substituting the values of A , B , C from Equation 2.
Revised: where the last equality follows by substituting the values of A , B , C from Equation 2.

Original: Hence,
Revised: Hence,

Original: [Equation]
Revised: [Equation]

Original: Given an input sequence u P R L , it is possible to compute the output y P R L sequentially via the recurrence in Equation 2.
Revised: Given an input sequence u P R L , it is possible to compute the output y P R L sequentially via the recurrence in Equation 2.

Original: Unfortunately, sequential computation is prohibitively slow on long inputs and, instead, Equation 5 can be used to compute all elements of y in parallel, provided we have already computed K .
Revised: Unfortunately, sequential computation is prohibitively slow on long inputs and, instead, Equation 5 can be used to compute all elements of y in parallel, provided we have already computed K .

Original: Computing y from u and K is easy.
Revised: Computing y from u and K is easy.

Original: Given an input sequence u P R L and the SSM kernel K P R L , naively using Equation 5 for computing y would require O p L 2 q multiplications.
Revised: Given an input sequence u P R L and the SSM kernel K P R L , naively using Equation 5 for computing y would require O p L 2 q multiplications.

Original: Fortunately, this can be done much more efficiently by observing that for the univariate polynomials
Revised: Fortunately, this can be done much more efﬁciently by observing that for the univariate polynomials

Original: [Equation]
Revised: [Equation]

Original: y k is the coefficient of z k in the polynomial K p z
Revised: y k is the coefﬁcient of z k in the polynomial K p z

Original: q ¨ u p z q , i.e. all y k ’s can be computed simultaneously by multiplying two degree L ´ 1 polynomials.
Revised: q ¨ u p z q , i.e. all y k ’s can be computed simultaneously by multiplying two degree L ´ 1 polynomials.

Original: It is well-known that this can be done in O p L log p L qq time via Fast Fourier Transform (FFT) [CLRS09].
Revised: It is well-known that this can be done in O p L log p L qq time via Fast Fourier Transform (FFT) [CLRS09].

Original: We denote this fast computation of Equation 5 via the discrete convolution as
Revised: We denote this fast computation of Equation 5 via the discrete convolution as

Original: [Equation]
Revised: [Equation]

Original: Hence, given the SSM kernel K P R L , the output of a discretized state space can be computed efficiently from the input.
Revised: Hence, given the SSM kernel K P R L , the output of a discretized state space can be computed efﬁciently from the input.

Original: The challenging part is computing K itself as it involves computing L distinct matrix powers (Equation 4).
Revised: The challenging part is computing K itself as it involves computing L distinct matrix powers (Equation 4).

Original: Similar to S4, instead of directly using A , B , C , our idea is to use an alternate parameterization of state spaces for which it would be easier to compute K .
Revised: Similar to S4, instead of directly using A , B , C , our idea is to use an alternate parameterization of state spaces for which it would be easier to compute K .

Original:  Method
Revised:  Method

Original: Having stated the necessary background, we now turn to the main contribution of our work.
Revised: Having stated the necessary background, we now turn to the main contribution of our work.

Original: 3.1 Diagonal State Spaces
Revised: 3.1 Diagonal State Spaces

Original: For a general state matrix A it is expensive to compute the matrix exponentials e
Revised: For a general state matrix A it is expensive to compute the matrix exponentials e

Original: A ¨ k ∆ for all 0 ď k ă
Revised: A ¨ k  for all 0 § k †

Original: L (Equation 4).
Revised: L (Equation 4).

Original: However, if A was diagonalizable over C and we were given its eigendecomposition
Revised: However, if A was diagonalizable over C and we were given its eigendecomposition

Original: A “ V Λ V ´ 1 , where Λ is diagonal , then we could easily compute e A ¨ k ∆ “ V e Λ ¨ k ∆ V ´ 1 as it is trivial to exponentiate diagonal matrices.
Revised: A “ V ⇤ V ´ 1 , where ⇤ is diagonal, then we could easily compute e A ¨ k  “ V e ⇤ ¨ k  V ´ 1 as it is trivial to exponentiate diagonal matrices.

Original: In this case Equation 4 simplifies to
Revised: In this case Equation 4 simpliﬁes to

Original: [Equation]
Revised: [Equation]

Original: Leveraging the above observation, we base our model on the following proposition which asserts that, under mild technical assumptions, diagonal state spaces are as expressive as general state spaces.
Revised: Leveraging the above observation, we base our model on the following proposition which asserts that, under mild technical assumptions, diagonal state spaces are as expressive as general state spaces.

Original: We use the operator K ∆ ,L p A, B, C q P C
Revised: We use the operator K  ,L p A, B, C q P C

Original: 1 ˆ L to denote the kernel (Equation 4) of length L for state space
Revised: 1 ˆ L to denote the kernel (Equation 4) of length L for state space

Original: p A, B, C q over C and sample time ∆ .
Revised: p A, B, C q over C and sample time  .

Original: Proposition 1.
Revised: Proposition 1.

Original: Let K P C
Revised: Let K P C

Original: ˆ L be the kernel of length L of a given state space
Revised: ˆ L be the kernel of length L of a given state space

Original: p A, B, C q and sample time ∆ ą 0 , where A P C N ˆ N is diagonalizable over C with eigenvalues λ 1 , . . .
Revised: p A, B, C q and sample time  ° 0 , where A P C N ˆ N is diagonalizable over C with eigenvalues  1 , . . .

Original: , λ N and @ i , λ i ‰and
Revised: ,  N and @ i ,  i ‰and

Original: e Lλ i ∆ ‰ 1 .
Revised: e L  i  ‰ 1 .

Original: Let P P C N ˆ L
Revised: Let P P C N ˆ L

Original: be P
Revised: be P

Original: i,k “ λ i k ∆ and Λ be the diagonal matrix with λ 1 , . ..
Revised: i,k “  i k  and ⇤ be the diagonal matrix with  1 , . .

Original: , λ N .
Revised: ,  N .

Original: Then there exist r w, w P C 1 ˆ N such that
Revised: Then there exist r w, w P C 1 ˆ N such that

Original: (a) K “ K ∆ ,L p Λ , p 1 q 1 ď i ď N , r w q “ r w ¨ Λ ´ 1 p e Λ∆ ´ I q ¨ elementwise - exp p P q , (b) K “ K ∆ ,L p Λ , pp e 
Revised: (a) K “ K  ,L p ⇤ , p 1 q 1 § i § N , r w q “ r w ¨ ⇤ ´ 1 p e ⇤ ´ I q ¨ elementwise - exp p P q , (b) K “ K  ,L p ⇤ , pp e L  i

Original: Lλ i ∆ ´ 1 q ´ 1 q 1 ď i ď N , w q “ w ¨ Λ ´ 1 ¨ row - softmax p P q .
Revised:  ´ 1 q ´ 1 q 1 § i § N , w q “ w ¨ ⇤ ´ 1 ¨ row - softmax p P q .

Original: 3 e A [Equation] ř 8 j “ 0 A j j ! [Equation]
Revised: 3 e A “ ∞ 8 j “ 0 A j j ! “ ∞ 8 j “ 0 p V ⇤ V ´ 1 q j j !

Original: ř 8 j “0 p V Λ V ´ 1 q j j! “ V ´ řj “ 0 Λ j j !
Revised: “ V ´ ∞ 8 j “ 0 ⇤ j j !

Original: ¯ V ´ 1 .
Revised: ¯ V ´ 1 .

Original: For Λ “ diag p λ 1 ,. . .
Revised: For ⇤ “ diag p  1 ,

Original: , λ N q we

Original: have
Revised: we

Original: Λ j “ diag p λ j 1 , . Input: parameters Λ re , Λ im P R N , w P C N , sample time ∆ log P R
Revised: Input: parameters ⇤ re , ⇤ im P R N , w P C N , sample time  log P R

Original: Output: SSM kernel K P R L (Proposition 1(b)) 
Revised: Output: SSM kernel K P R L (Proposition 1(b)) 1:

Revised: ⇤ – ⇤

Original: [Equation]
Revised: [Equation]

Revised: ` i ¨ ⇤

Original:  q 
Revised: [Equation] ,  – [Equation] log q ô  is positive real 2: P N ˆ L – p

Original: [Table]
Revised: [Equation]

Original: 2: P
Revised: q N

Original: [Equation]
Revised: [Equation]

Revised: ˚ r 0 , 1 ,

Original: [Table]
Revised: [Equation]

Original: outer product 3: S Ð row - softmax ϵ p P q Ź
Revised: L ´

Original: [Table]
Revised: [Equation]

Original: stable (§A.2) 4: K Ð Re p p w
Revised: s

Original: [Equation]
Revised: [Equation]

Original: Λ q 1 ˆ N ¨ S  q 
Revised: ˆ L ô outer product 3: S – row - softmax ✏ p P q ô

Original: [Table]
Revised: [Table]

Revised: ) 4: K – Re p p w { ⇤ q 1 ˆ N ¨ S q ô real part

Original: The proof of Proposition 1 is provided in §A.1.
Revised: The proof of Proposition 1 is provided in §A.1.

Original: 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.
Revised: 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.

Original: Hence, for any given state space with a well-behaved state matrix there exists a diagonal state space computing the same kernel.
Revised: Hence, for any given state space with a well-behaved state matrix there exists a diagonal state space computing the same kernel.

Original: 4 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.
Revised: 4 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.

Original: Proposition 1(a) suggests that we can parameterize state spaces via Λ , r w P C N and simply compute the kernel as shown in Proposition 1(a).
Revised: Proposition 1(a) suggests that we can parameterize state spaces via ⇤ , r w P C N and simply compute the kernel as shown in Proposition 1(a).

Original: Unfortunately, in practice, the real part of the elements of Λ can become positive during training making the training unstable on long inputs.
Revised: Unfortunately, in practice, the real part of the elements of ⇤ can become positive during training making the training unstable on long inputs.

Original: This is because the matrix elementwise - exp p P q contains terms as large as exp p λ i ∆ p L ´ 1 qq which even for a modest value of L can be very large.
Revised: This is because the matrix elementwise - exp p P q contains terms as large as exp p  i  p L ´ 1 qq which even for a modest value of L can be very large.

Original: To address this, we propose two methods to model diagonal state spaces.
Revised: To address this, we propose two methods to model diagonal state spaces.

Original: DSS EXP
Revised: DSS EXP

Original: In this variant, we use Proposition 1(a) to model our state space but restrict the real part of elements of Λ to be negative.
Revised: In this variant, we use Proposition 1(a) to model our state space but restrict the real part of elements of ⇤ to be negative.

Original: DSS EXP has parameters Λ re , Λ im P R N , r w P C N and ∆ log P R .
Revised: DSS EXP has parameters ⇤ re , ⇤ im P R N , r w P C N and  log P R .

Original: Λ is computed as ´ elementwise - exp p Λ re q ` i ¨ Λ im where i “ ?
Revised: ⇤ is computed as ´ elementwise - exp p ⇤ re q ` i ¨ ⇤ im where i “ ?

Original: ´ 1 [GGDR22].
Revised: ´ 1 [GGDR22].

Original: ∆ is computed as exp p ∆ log q P R ą 0{+° 
+} and the kernel is then computed via Proposition 1(a).
Revised:  is computed as exp p  log q P R {+° 
+} and the kernel is then computed via Proposition 1(a).

Original: DSS EXP provides a remarkably simple computation of state space kernels but restricts the space of the learned Λ (the real part must be negative).
Revised: DSS EXP provides a remarkably simple computation of state space kernels but restricts the space of the learned ⇤ (the real part must be negative).

Original: It is not clear if such a restriction could be detrimental for some tasks, and we now present an alternate method that provides the simplicity of Proposition 1(a) while being provably as expressive as general state spaces.
Revised: It is not clear if such a restriction could be detrimental for some tasks, and we now present an alternate method that provides the simplicity of Proposition 1(a) while being provably as expressive as general state spaces.

Original: DSS SOFTMAX
Revised: DSS SOFTMAX

Original: Instead of restricting the elements of Λ , another option for bounding the elements of elementwise - exp p P q is to normalize each row by the sum of its elements, which leads us to
Revised: Instead of restricting the elements of ⇤ , another option for bounding the elements of elementwise - exp p P q is to normalize each row by the sum of its elements, which leads us to

Original: Proposition 1(b).
Revised: Proposition 1(b).

Original: DSS SOFTMAX has parameters Λ re , Λ im P R N , w P C N and ∆ log P R .
Revised: DSS SOFTMAX has parameters ⇤ re , ⇤ im P R N , w P C N and  log P R .

Original: Λ is computed as Λ re ` i ¨ Λ im , ∆ is computed as exp p ∆ log q P R ą 0 , and the kernel is then computed via Proposition 1(b).
Revised: ⇤ is computed as ⇤ re ` i ¨ ⇤ im ,  is computed as exp p  log q P R ° 0 , and the kernel is then computed via Proposition 1(b).

Original: A sketch of this computation is presented in Algorithm 1.
Revised: A sketch of this computation is presented in Algorithm 1.

Original: We note that, unlike over R , softmax can have singularities over C and slight care must be taken while computing it (e.g., softmax p 0 , iπ q is not defined).
Revised: We note that, unlike over R , softmax can have singularities over C and slight care must be taken while computing it (e.g., softmax p 0 , i ⇡ q is not deﬁned).

Original: We instead use a corrected version softmax ϵ and detail it in §A.2.
Revised: We instead use a corrected version softmax ✏ and detail it in §A.2.

Original: 3.2 Diagonal State Space (DSS) Layer
Revised: 3.2 Diagonal State Space (DSS) Layer

Original: We are now ready to describe the DSS layer.
Revised: We are now ready to describe the DSS layer.

Original: We retain the skeletal structure of S4 and simply replace the parameterization and computation of the SSM kernel by one of the methods described in §3.1.
Revised: We retain the skeletal structure of S4 and simply replace the parameterization and computation of the SSM kernel by one of the methods described in §3.1.

Original: Each DSS layer receives a length- L sequence u of H -dimensional vectors as input, i.e., u P R H ˆ L , and produces an output y P R H ˆ L .
Revised: Each DSS layer receives a length- L sequence u of H -dimensional vectors as input, i.e., u P R H ˆ L , and produces an output y P R H ˆ L .

Original: The parameters of the layer are Λ re , Λ im P R N , ∆ log P R H and
Revised: The parameters of the layer are ⇤ re , ⇤ im P R N ,  log P R H and

Original: W P C H ˆ N .
Revised: W P C H ˆ N .

Original: For each coordinate h “ 1 , . . .
Revised: For each coordinate h “ 1 , . . .

Original: , H , a state space kernel K h P R L is computed using one of methods described in §3.1.
Revised: , H , a state space kernel K h P R L is computed using one of methods described in §3.1.

Original: For example, in case of DSS SOFTMAX , Algorithm 1 is used with parameters Λ re , Λ im P R N , W h P C N and p ∆ log q h P R .
Revised: For example, in case of DSS SOFTMAX , Algorithm 1 is used with parameters ⇤ re , ⇤ im P R N , W h P C N and p  log q h P R .

Original: The output y h P R L for coordinate h is computed from u h P R L and K h using Equation 6.
Revised: The output y h P R L for coordinate h is computed from u h P R L and K h using Equation 6.

Original: Following S4, we also add a residual connection from u to y followed by a GELU non-linearity [HG16].
Revised: Following S4, we also add a residual connection from u to y followed by a GELU non-linearity [HG16].

Original: This is followed by a position-wise linear projection W out P R H ˆ H to enable information exchange among the H coordinates.
Revised: This is followed by a position-wise linear projection W out P R H ˆ H to enable information exchange among the H coordinates.

Original: The DSS layer can be implemented in just a few lines of code and our PyTorch implementation of DSS SOFTMAX layer is provided in §A.7 (Figure 6).
Revised: The DSS layer can be implemented in just a few lines of code and our PyTorch implementation of DSS SOFTMAX layer is provided in §A.7 (Figure 6).

Original: The implementation of DSS EXP layer is even simpler and is omitted.
Revised: The implementation of DSS EXP layer is even simpler and is omitted.

Original: Complexity For batch size B , sequence length L and hidden size H , the DSS layer requires
Revised: Complexity For batch size B , sequence length L and hidden size H , the DSS layer requires

Original: O p NHL q time and space to compute the kernels, O p BHL log p L qq time for the discrete convolution and O p BH 2 L q time for the output projection.
Revised: O p NHL q time and space to compute the kernels, O p BHL log p L qq time for the discrete convolution and O p BH 2 L q time for the output projection.

Original: For small batch size B , the time taken to compute the kernels becomes important whereas for large batches more compute is spent on the convolution and the linear projection.
Revised: For small batch size B , the time taken to compute the kernels becomes important whereas for large batches more compute is spent on the convolution and the linear projection.

Original: The kernel part of DSS layer has 2 N ` H ` 2 HN real-valued parameters.
Revised: The kernel part of DSS layer has 2 N ` H ` 2 HN real-valued parameters.

Original: 3.3 Initialization of DSS layer
Revised: 3.3 Initialization of DSS layer

Original: The performance of state spaces models is known to be highly sensitive to initialization [GGR22].
Revised: The performance of state spaces models is known to be highly sensitive to initialization [GGR22].

Original: 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).
Revised: 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).

Original: The real and imaginary parts of each element of W are initialized from N p 0 , 1 q .
Revised: The real and imaginary parts of each element of W are initialized from N p 0 , 1 q .

Original: Each element of ∆ log is initialized as e r where r „ U p log p .
Revised: Each element of  log is initialized as e r where r „ U p log p .

Original: 001 q , log p .
Revised: 001 q , log p .

Original: 1 qq .
Revised: 1 qq .

Original: Λ P C N is initialized using eigenvalues of the normal part of normal plus low-rank form of HiPPO matrix [GGR22].
Revised: ⇤ P C N is initialized using eigenvalues of the normal part of normal plus low-rank form of HiPPO matrix [GGR22].

Original: Concretely, Λ re , Λ im are initialized such that the resulting Λ is the vector of those N eigenvalues of the following 2 N ˆ 2 N matrix which have a positive imaginary part.
Revised: Concretely, ⇤ re , ⇤ im are initialized such that the resulting ⇤ is the vector of those N eigenvalues of the following 2 N ˆ 2 N matrix which have a positive imaginary part.

Original: [Equation]
Revised: [Equation]

Original: Henceforth, we would refer to the above initialization of Λ as Skew-Hippo initialization.
Revised: Henceforth, we would refer to the above initialization of ⇤ as Skew-Hippo initialization.

Original: In all our experiments, we used the above initialization with N “ 64 .
Revised: In all our experiments, we used the above initialization with N “ 64 .

Original: The initial learning rate of all DSS parameters was 10 ´ 3 and weight decay was not applied to them.
Revised: The initial learning rate of all DSS parameters was 10 ´ 3 and weight decay was not applied to them.

Original: Exceptions to these settings are noted in §A.3.
Revised: Exceptions to these settings are noted in §A.3.

Original: 3.4 States of DSS via the Recurrent View
Revised: 3.4 States of DSS via the Recurrent View

Original: 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.
Revised: 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.

Original: 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
Revised: But in situations such as autoregressive decoding during inference, it is more efﬁcient to explicitly compute the states of a state space model similar to a linear

Original: RNN (Equation 2).
Revised: RNN (Equation 2).

Original: We now show how to compute the states of the DSS models described in §3.1.
Revised: We now show how to compute the states of the DSS models described in §3.1.

Original: Henceforth, we assume that we have already computed Λ and ∆ .
Revised: Henceforth, we assume that we have already computed ⇤ and  .

Original: DSS EXP
Revised: DSS EXP

Original: As stated in Proposition 1(a), DSS EXP computes K ∆ ,L p Λ , p 1 q 1 ď i ď N , r w q .
Revised: As stated in Proposition 1(a), DSS EXP computes K  ,L p ⇤ , p 1 q 1 § i § N , r w q .

Original: For this state space and sample time ∆ , we use Equation 2 to obtain its discretization
Revised: For this state space and sample time  , we use Equation 2 to obtain its discretization

Original: [Equation]
Revised: [Equation]

Original: where diag creates a diagonal matrix of the scalars.
Revised: where diag creates a diagonal matrix of the scalars.

Original: We can now compute the states using the SSM recurrence x k “ Ax k ´ 1 ` Bu k (Equation 2).
Revised: We can now compute the states using the SSM recurrence x k “ Ax k ´ 1 ` Bu k (Equation 2).

Original: As A is diagonal, the N coordinates do not interact and hence can be computed independently.
Revised: As A is diagonal, the N coordinates do not interact and hence can be computed independently.

Original: Let us assume x ´ 1 “ 0 and say we have already computed x k ´ 1 .
Revised: Let us assume x ´ 1 “ 0 and say we have already computed x k ´ 1 .

Original: Then, for the i ’th coordinate independently compute
Revised: Then, for the i ’th coordinate independently compute

Original: [Equation]
Revised: [Equation]

Original: Note that in DSS EXP , Re p λ i q ă 0 and hence | e λ i ∆ | “ | e Re p λ i q ∆ | ă 1 .
Revised: Note that in DSS EXP , Re p  i q † 0 and hence | e  i  | “ | e Re p  i q  | † 1 .

Original: Intuitively, if | λ i | ∆ « 0 , we would have x i,k « x i,k ´ 1 and be able to copy history over many timesteps.
Revised: Intuitively, if |  i |  « 0 , we would have x i,k « x i,k ´ 1 and be able to copy history over many timesteps.

Original: On the other hand if
Revised: On the other hand if

Original: Re p λ i q ∆ ! 

Revised: Re p  i q  !

Original: then x i,k « ´ λ ´ 

Revised: then x i,k « ´  ´

Original: i u k
Revised: i u k

Original: and hence the information from the previous timesteps would be forgotten similar to a “forget” gate in LSTMs.
Revised: and hence the information from the previous timesteps would be forgotten similar to a “forget” gate in LSTMs.

Revised: formalized the connection between Skew-Hippo initialization and the original HiPPO matrices and provided a theoretical justiﬁcation for its empirical properties.

Original: DSS SOFTMAX
Revised: DSS SOFTMAX

Original: As stated in Proposition 1(b), DSS SOFTMAX computes K ∆ ,L p Λ , pp e 
Revised: As stated in Proposition 1(b), DSS SOFTMAX computes K  , L p ⇤ , pp e L  i

Original: Lλ i ∆ ´ 1 q ´ 1 q 1 ď i ď N , w q .
Revised:  ´ 1 q ´ 1 q 1 § i § N , w q .

Original: For this state space and sample time ∆ , we obtain the discretization
Revised: For this state space and sample time  , we obtain the discretization

Original: [Equation]
Revised: [Equation]

Original: For the i ’th coordinate we can independently compute
Revised: For the i ’th coordinate we can independently compute

Original: [Equation]
Revised: [Equation]

Original: Let us drop the coordinate index i for clarity to obtain
Revised: Let us drop the coordinate index i for clarity to obtain

Original: [Equation]
Revised: [Equation]

Original: where x k ´ 1 is now a scalar.
Revised: where x k ´ 1 is now a scalar.

Original: As the expression involves the term e λ ∆ L , where L can be large, directly computing such terms can result in numerical instability and we must avoid exponentiating scalars with a positive real part.
Revised: As the expression involves the term e   L , where L can be large, directly computing such terms can result in numerical instability and we must avoid exponentiating scalars with a positive real part.

Original: We make two cases depending on the sign of Re p λ q and compute x k via an intermediate state r x k as follows.
Revised: We make two cases depending on the sign of Re p  q and compute x k via an intermediate state r x k as follows.

Original: Let p “ I r Re p λ q ą 

Revised: Let p “ I r Re p  q °

Original: s P t 0 , 1 u .
Revised: s P t 0 , 1 u .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: The above equation can be parsed as follows.
Revised: The above equation can be parsed as follows.

Original: If Re p λ q ď 0 then we have
Revised: If Re p  q § 0 then we have

Original: [Equation]
Revised: [Equation]

Original: and hence if Re p λ q ∆ ! {+
+}
Revised: and hence if Re p  q  ! {+
+}

Original: then r x k « u k and x k « u k { λ .
Revised: then r x k « u k and x k « u k {  .

Original: In this case, information from the previous timesteps would be ignored and the information used would be local.
Revised: In this case, information from the previous timesteps would be ignored and the information used would be local.

Original: On the other hand if Re p λ q ąthen we would haver
Revised: On the other hand if Re p  q ° 0 then we would have

Original: [Equation]
Revised: [Equation]

Original: λ
Revised: and hence if Re p  q  " 0 then r x 0 « u 0 and r

Original: [Equation]
Revised: [Equation]

Original: ´e
Revised: u

Original: [Equation]
Revised: [Equation]

Original: ´and hence if Re p λ q ∆ " 0 then r x 0 « u 0 and r x k « r x k ´ 1 « u 0 , x k ă L ´ 1 « 0 and x L ´ 1 « u 0 { λ .
Revised: , x [Equation]

Original:  Hence, the model would be able to copy information even from extremely distant positions, allowing it to capture long-range dependencies.
Revised: and x L ´ [Equation] « u 0 {  . Hence, the model would be able to copy information even from extremely distant positions, allowing it to capture long-range dependencies.

Original:  Experiments
Revised:  Experiments

Original: We evaluate the performance of DSS on sequence-level classification tasks over text, images and audio.
Revised: We evaluate the performance of DSS on sequence-level classiﬁcation tasks over text, images and audio.

Original: Overall, we find its performance is comparable to S4.
Revised: Overall, we ﬁnd its performance is comparable to S4.

Original: Long Range Arena (LRA) LRA [TDA ` 21] is a standard benchmark for assessing the ability of models to process long sequences.
Revised: Long Range Arena (LRA) LRA [TDA ` 21] is a standard benchmark for assessing the ability of models to process long sequences.

Original: LRA contains 6 tasks with diverse input lengths 1 K - 16 K , encompassing modalities such as text and images.
Revised: LRA contains 6 tasks with diverse input lengths 1 K - 16 K , encompassing modalities such as text and images.

Original: Several Transformer variants have been benchmarked on
Revised: Several Transformer variants have been benchmarked on

Original: LRA but all have underperformed due to factors such as their high compute-memory requirements, implicit locality bias and inability to capture long-range dependencies.
Revised: LRA but all have underperformed due to factors such as their high compute-memory requirements, implicit locality bias and inability to capture long-range dependencies.

Original: Table 1 compares DSS against S4, the Transformer variants reported in [TDA ` 21], as well as followup work.
Revised: Table 1 compares DSS against S4, the Transformer variants reported in [TDA ` 21], as well as followup work.

Original: State space models (S4, DSS) shown in Table 1 are left-to-right unidirectional whereas other models could be bidirectional.
Revised: State space models (S4, DSS) shown in Table 1 are left-to-right unidirectional whereas other models could be bidirectional.

Original: Despite its simplicity, DSS delivers state-of-the-art performance on LRA.
Revised: Despite its simplicity, DSS delivers state-of-the-art performance on LRA.

Original: Its performance is comparable to that of S4, with a modest improvement in test accuracy averaged across the 6 tasks ( 81 .
Revised: Its performance is comparable to that of S4, with a modest improvement in test accuracy averaged across the 6 tasks ( 81 .

Original: 88 vs 80 . 21 ).
Revised: 88 vs 80 . 21 ).

Original: In addition, DSS maintains a comfortable 20 point lead over the best performing Transformer variant ( 81 . 88 vs 61 . 41 ).
Revised: In addition, DSS maintains a comfortable 20 point lead over the best performing Transformer variant ( 81 . 88 vs 61 . 41 ).

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Interestingly, despite being less expressive than DSS SOFTMAX , DSS EXP also reports an impressive performance which makes it specially appealing given its simplicity during training and inference.
Revised: Interestingly, despite being less expressive than DSS SOFTMAX , DSS EXP also reports an impressive performance which makes it specially appealing given its simplicity during training and inference.

Original: We investigate an even simpler version of DSS EXP , denoted as DSS EXP - NO - SCALE , which is same as DSS EXP except that we omit the term Λ ´ 1 p e Λ∆
Revised: We investigate an even simpler version of DSS EXP , denoted as DSS EXP - NO - SCALE , which is same as DSS EXP except that we omit the term ⇤ ´ 1 p e ⇤

Original: ´ I q in the expression of the kernel shown in
Revised: ´ I q in the expression of the kernel shown in

Original: Proposition 1(a) and instead compute it as r w ¨ elementwise - exp p P q .
Revised: Proposition 1(a) and instead compute it as [Equation] w ¨ elementwise - exp p P q .

Original: As shown in Table 1, the performance of DSS EXP - NO - SCALE is generally inferior to that of DSS EXP with the model failing on the challenging P ATH -X task which requires the model to capture extremely long-range interactions
Revised: As shown in Table 1, the performance of DSS EXP - NO - SCALE is generally inferior to that of DSS EXP with the model failing on the challenging P ATH -X task which requires the model to capture extremely long-range interactions

Original: [Table]
Revised: [Table]

Original: Raw Speech Classification Audio is typically digitized using a high sampling rate resulting in very long sequences.
Revised: Raw Speech Classiﬁcation Audio is typically digitized using a high sampling rate resulting in very long sequences.

Original: This provides an interesting domain for investigating the abilities of long-range models.
Revised: This provides an interesting domain for investigating the abilities of long-range models.

Original: We evaluate the performance of DSS on the Speech Commands (SC) dataset [War18], consisting of raw audio samples of length 16000 , modeled as a 10 -way classification task.
Revised: We evaluate the performance of DSS on the Speech Commands (SC) dataset [War18], consisting of raw audio samples of length 16000 , modeled as a 10 -way classiﬁcation task.

Original: As shown in Table 2, the performance of all DSS variants is comparable to that of S4 ( 98 .
Revised: As shown in Table 2, the performance of all DSS variants is comparable to that of S4 ( 98 .

Original: 2 vs 98 . 1 ).
Revised: 2 vs 98 . 1 ).

Original: In all experiments and ablations, S4 and DSS use identical model hyperparameters such as hidden size, number of layers, etc.
Revised: In all experiments and ablations, S4 and DSS use identical model hyperparameters such as hidden size, number of layers, etc.

Original: Our experimental setup was built on top of the training framework provided by the S4 authors and for our S4 runs we followed their official instructions.
Revised: Our experimental setup was built on top of the training framework provided by the S4 authors and for our S4 runs we followed their ofﬁcial instructions.

Original: 8 Details about model initialization, and hyperparameters are provided in §A.3.
Revised: 8 Details about model initialization, and hyperparameters are provided in §A.3.

Original: 4.1 Analyzing the Performance of DSS
Revised: 4.1 Analyzing the Performance of DSS

Original: While the experimental results presented above are encouraging, and clearly demonstrate the effectiveness of DSS at modeling long-range dependencies, it is not clear what exactly are the main factors 
Revised: While the experimental results presented above are encouraging, and clearly demonstrate the effectiveness of DSS at modeling long-range dependencies, it is not clear what exactly are the main factors [Table] [Table]

Original: [Table]
Revised: [Table]

Original: contributing to its performance.
Revised: contributing to its performance.

Original: To investigate this further
Revised: To investigate this further

Original: [Table]
Revised: [Table]

Original: we performed an ablation analysis aimed at answering the following questions:
Revised: we performed an ablation analysis aimed at answering the following questions:

Original: • How significant is the Skew-Hippo initialization (§3.3) to the model performance?
Revised: • How signiﬁcant is the Skew-Hippo initialization (§3.3) to the model performance?

Original: Would initializing Λ randomly work just as well?
Revised: Would initializing ⇤ randomly work just as well?

Original: • Is the main source of superior performance of state space models (S4, DSS), compared to previous models, their ability to model long-range dependencies?
Revised: • Is the main source of superior performance of state space models (S4, DSS), compared to previous models, their ability to model long-range dependencies?

Original: Would restricting DSS to only model local interactions hurt its performance on the above tasks?
Revised: Would restricting DSS to only model local interactions hurt its performance on the above tasks?

Original: Random Initialization To answer the first question, we repeated the experiments after initializing each element of Λ re and Λ im in DSS SOFTMAX by randomly sampling from N p 0 , 1 q .
Revised: Random Initialization To answer the ﬁrst question, we repeated the experiments after initializing each element of ⇤ re and ⇤ im in DSS SOFTMAX by randomly sampling from N p 0 , 1 q .

Original: Truncated Kernels To answer the second question, instead of constructing a kernel of length equal to the length L of the input, we restricted the length of the kernel constructed in DSS SOFTMAX (Algorithm 1) to 128 , significantly shorter than the length of the input.
Revised: Truncated Kernels To answer the second question, instead of constructing a kernel of length equal to the length L of the input, we restricted the length of the kernel constructed in DSS SOFTMAX (Algorithm 1) to 128 , signiﬁcantly shorter than the length of the input.

Original: To understand the implication of this restriction recall Equation 5 which states that y k 
Revised: To understand the implication of this restriction recall Equation 5 which states that y k “

Original: [Equation]
Revised: [Equation]

Original: For a given context size c “ 128 , restricting K ě c “ 0 would imply
Revised: For a given context size c “ 128 , restricting K • c “ 0 would imply

Original: [Equation]
Revised: [Equation]

Original: and hence the output y k at position k would only depend on u k , . .
Revised: and hence the output y k at position k would only depend on u k , . .

Original: , u k ´ c ` 1 .
Revised: , u k ´ c ` 1 .

Original: This would restrict each DSS SOFTMAX layer to only model local interactions and the model would require several layers to have a broader receptive field and capture long-range interactions.
Revised: This would restrict each DSS SOFTMAX layer to only model local interactions and the model would require several layers to have a broader receptive ﬁeld and capture long-range interactions.

Original: As shown in Table 3, randomly initializing the Λ parameters of DSS leads to significant performance degradation on the majority of tasks, with the model failing to perform on P ATH -X.
Revised: As shown in Table 3, randomly initializing the ⇤ parameters of DSS leads to signiﬁcant performance degradation on the majority of tasks, with the model failing to perform on P ATH -X.

Original: This is inline with the findings of [GGR22] who also reported the initialization of S4 to be critical to its performance.
Revised: This is inline with the ﬁndings of [GGR22] who also reported the initialization of S4 to be critical to its performance.

Original: Interestingly, despite this performance reduction, DSS manages to outperform all non state-spacebased models on every task.
Revised: Interestingly, despite this performance reduction, DSS manages to outperform all non state-spacebased models on every task.

Original: Truncating the length of the kernel also leads to a significant reduction in performance across most tasks (Table 3), suggesting that the superior performance of state-space models on these tasks can indeed be partly attributed to their ability to capture long-range dependencies.
Revised: Truncating the length of the kernel also leads to a signiﬁcant reduction in performance across most tasks (Table 3), suggesting that the superior performance of state-space models on these tasks can indeed be partly attributed to their ability to capture long-range dependencies.

Original: Moreover, on some tasks such as L IST O PS and I MAGE , using a truncated kernel still manages to outperform all
Revised: Moreover, on some tasks such as L IST O PS and I MAGE , using a truncated kernel still manages to outperform all

Original: Transformer variants, which is surprising as Transformer layers are known to be effective at capturing interactions at such short ranges.
Revised: Transformer variants, which is surprising as Transformer layers are known to be effective at capturing interactions at such short ranges.

Original: 4.2 Analysis of Learned DSS Parameters
Revised: 4.2 Analysis of Learned DSS Parameters

Original: To further explore the inner workings of DSS, we visually inspected the trained parameters and kernels of DSS SOFTMAX .
Revised: To further explore the inner workings of DSS, we visually inspected the trained parameters and kernels of DSS SOFTMAX .

Original: The kernels of all layers of the trained DSS SOFTMAX are shown in Figure 2 and reveal a stark contrast between the tasks.
Revised: The kernels of all layers of the trained DSS SOFTMAX are shown in Figure 2 and reveal a stark contrast between the tasks.

Original: On the tasks I MAGE 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.
Revised: On the tasks I MAGE and SC, for almost all kernels, the absolute values of the ﬁrst 128 positions are signiﬁcantly higher than the later positions indicating that these kernels are mostly local.

Original: On the other hand for P ATH -X, for a significant proportion of kernels the opposite is true, indicating that these kernels are modeling long-range interactions.[Figure] y e r s × h i dden - s i
Revised: On the other hand for P ATH -X, for a signiﬁcant proportion of kernels the opposite is true, indicating that these kernels are modeling long-range interactions.

Original: z e [Figure] This partially explains why the performance of DSS SOFTMAX on I MAGE and SC does not decrease significantly even after limiting the kernel length to 128 , whereas P ATHFINDER and P ATH -X report significant degradation (Table 3).
Revised: This partially explains why the performance of DSS SOFTMAX on I MAGE and SC does not decrease signiﬁcantly even after limiting the kernel length to 128 , whereas P ATHFINDER and P ATH -X report signiﬁcant degradation (Table 3).

Original: The last row of Table 3 shows the 95th percentile of the set t argmax [-
 ď-] k ă L | K k | : K is a kernel in DSS SOFTMAX u , i.e., the set of positions over all kernels, at which for a given kernel its absolute value is maximized.
Revised: The last row of Table 3 shows the 95th percentile of the set t argmax [-
 ď-]0 § k † L | K k | : K is a kernel in DSS SOFTMAX u , i.e., the set of positions over all kernels, at which for a given kernel its absolute value is maximized.

Original: As expected, most kernels of I MAGE and
Revised: As expected, most kernels of I MAGE and

Original: SC are local whereas for P ATH -X they are mostly long range.
Revised: SC are local whereas for P ATH -X they are mostly long range.

Original: The values of the real and imaginary parts of Λ are shown in Figure 3 (and §A.6, Figure 5).
Revised: The values of the real and imaginary parts of ⇤ are shown in Figure 3 (and §A.6, Figure 5).

Original: We observe that, although all elements of Λ re are initialized as ´ 0 . 5 , their values can change significantly during training and can become positive (see plots for L IST O PS , SC).
Revised: We observe that, although all elements of ⇤ re are initialized as ´ 0 . 5 , their values can change signiﬁcantly during training and can become positive (see plots for L IST O PS , SC).

Original: It can also be observed that a λ with a more negative Re p λ q generally tends to have a larger Im p λ q , which interestingly is a property not satisfied by Skew-Hippo initialization.[Figure] a r cs i nh ( i m ag / 2 )
Revised: It can also be observed that a  with a more negative Re p  q generally tends to have a larger Im p  q , which interestingly is a property not satisﬁed by Skew-Hippo initialization.

Original: [Figure]
Revised: a

Original: The values of the trained ∆ log corresponding to the real and imaginary parts of Λ are shown in Figure 4 and, in this case as well, change significantly during training.
Revised: The values of the trained  log corresponding to the real and imaginary parts of ⇤ are shown in Figure 4 and, in this case as well, change signiﬁcantly during training.

Original: The values of ∆ log on short-range tasks such as I MAGE and SC generally tend to be larger compared to long-range tasks such as
Revised: The values of  log on short-range tasks such as I MAGE and SC generally tend to be larger compared to long-range tasks such as

Original: P ATHFINDER and P ATH -X, inline with the intuition that, when Re p λ q ă 0 , smaller values of ∆ correspond to modeling long-range dependence (§3.4).2.5 0.0     l a y e r text    listops      retrieval    0 real image [Figure] l a y e r pathfinder
Revised: P ATHFINDER and P ATH -X, inline with the intuition that, when Re p  q † 0 , smaller values of  correspond to modeling long-range dependence (§3.4).

Original: [Figure]

Original: sc [Figure] [Figure] We note that the plot for L IST O PS reveals an outlier with a value of 22 which after exponentiation in
Revised: We note that the plot for L IST O PS reveals an outlier with a value of 22 which after exponentiation in

Original: Algorithm 1 would result in an extreme large ∆ .
Revised: Algorithm 1 would result in an extreme large  .

Original: This can potentially lead to training instabilities and we plan to address this issue in future work.
Revised: This can potentially lead to training instabilities and we plan to address this issue in future work.

Original: Limitations and future work
Revised: Limitations and future work

Original: In this work, we evaluated DSS on sequence-level classification tasks.
Revised: In this work, we evaluated DSS on sequence-level classiﬁcation tasks.

Original: In future work, we plan to include token-level generation tasks such as language modeling, forecasting, etc.
Revised: In future work, we plan to include token-level generation tasks such as language modeling, forecasting, etc.

Original: Another important future direction is to pretrain DSS-based models on large amounts of raw data.
Revised: Another important future direction is to pretrain DSS-based models on large amounts of raw data.

Original: Lastly, we found that the initialization and learning rates of DSS parameters Λ re , Λ im , ∆ log (§3.3) play an important role in the performance and convergence of the model.
Revised: Lastly, we found that the initialization and learning rates of DSS parameters ⇤ re , ⇤ im ,  log (§3.3) play an important role in the performance and convergence of the model.

Original: Informally, for tasks such as P ATH -X that require very long-range interactions, a smaller initialization of ∆ log was beneficial.
Revised: Informally, for tasks such as P ATH -X that require very long-range interactions, a smaller initialization of  log was beneﬁcial.

Original: An in-depth analysis of this phenomenon could be helpful and remains for future work.
Revised: An in-depth analysis of this phenomenon could be helpful and remains for future work.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: We thank Ramon Fernandez Astudillo for carefully reviewing the preliminary draft and suggesting several helpful edits.
Revised: We thank Ramon Fernandez Astudillo for carefully reviewing the preliminary draft and suggesting several helpful edits.

Original: We thank Omer Levy, Achille Fokoue and Luis Lastras for their support.
Revised: We thank Omer Levy, Achille Fokoue and Luis Lastras for their support.

Original: Our experiments were conducted on IBM’s Cognitive Computing Cluster, with additional resources from
Revised: Our experiments were conducted on IBM’s Cognitive Computing Cluster, with additional resources from

Original: Tel Aviv University.
Revised: Tel Aviv University.

Original: This research was supported by (1) IBM AI Residency program and (2) Defense Advanced Research Projects Agency (DARPA) through Cooperative Agreement D20AC awarded by the U S. Department of the Interior (DOI), Interior Business Center and (3) the European Research Council (ERC) under the European Union Horizons 2020 research and innovation programme (grant ERC DELPHI 802800)
Revised: This research was supported by (1) IBM AI Residency program and (2) Defense Advanced Research Projects Agency (DARPA) through Cooperative Agreement D20AC awarded by the U.S. Department of the Interior (DOI), Interior Business Center, and (3) the European Research Council (ERC) under the European Union Horizons 2020 research and innovation programme (grant ERC DELPHI 802800).

Original: Checklist
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] See end of §4.2.
Revised: [Yes] See end of §4.2.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: [No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 2.
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] (b) Did you include complete proofs of all theoretical results?
Revised: [Yes] (b) Did you include complete proofs of all theoretical results?

Original: [Yes] See §A.1.
Revised: [Yes] See §A.1.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] See §1, §A.3.
Revised: [Yes] See §1, §A.3.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] See §A.3.
Revised: [Yes] See §A.3.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [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)?
Revised: [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)?

Original: [Yes] See §A.3.
Revised: [Yes] See §A.3.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [No] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [No] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A] (d) Did you discuss whether
Revised: [N/A] (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: and how consent was obtained from people whose data you’re using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

Original: [N/A] 5.
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?[N/A] A Supplemental Material A.1 Diagonal State Spaces
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: We restate Proposition 1 for convenience.

Original: Proposition.

Original: Let K P C

Original: ˆ L be the kernel of length L of a given state space

Original: p A, B, C q and sample time ∆ ą 0 , where A P C N ˆ N is diagonalizable over C with eigenvalues λ 1 , . . .

Original: , λ N and @

Original: i , λ i ‰and

Original: e Lλ i ∆ ‰ 1 .

Original: Let P P C N ˆ L

Original: be P

Original: i,k “ λ i k ∆ and Λ be the diagonal matrix with λ 1 , . . .

Original: , λ N .

Original: Then there exist r w

Original: [Equation]

Original: w P C 1 ˆ N such that

Original: [Equation]

Original: Proof.

Original: Let A be diagonalizable over C as A “ V Λ V ´ 1 with eigenvalues λ 1 , . . .

Original: , λ N P C .

Original: From

Original: Equation 4 we have

Original: [Equation]

Original: where

Original: [Equation]

Original: For CV P C 1 ˆ N and V ´ 1 B P C N ˆ 1 let p CV

Original: q J ˚ p V ´ 1 B q “ r w P C N be the element-wise product of CV and V ´ 1 B .

Original: Then,

Original: [Equation]

Original: where the last equality follows from

Original: p z

Original: L ´ 1 q “ p z ´ 1 qp z 0 ` . . .

Original: ` z L ´ 1 q and using z L ‰ 1 .

Original: Let P P C N ˆ L be the matrix P i,k “ λ i ¨ k ∆ and let E “ elementwise - exp p P q .

Original: It is easy to verify that Equation 7 can be re-written as a vector-matrix product as

Original: [Equation]

Original: Similarly, for the state space p Λ , p 1 q 1 ď i ď N , r w qq and sample time ∆ its kernel r K can be obtained from Equation 4 as

Original: [Equation]

Original: which is also the expression for K k (Equation 7).

Original: This proves part (a) and we now consider part (b).

Original: Let w P C N be defined as

Original: [Equation]

Original: Then from Equation 9,

Original: [Equation]

Original: Let S “ row - softmax p P q denote the matrix obtained after applying softmax on the rows of P , i.e.

Original: [Equation]

Original: i,

Original: [Equation]

Original: It is easy to verify that Equation 10 can be expressed as a vector-matrix product

Original: [Equation]

Original: Similarly, for the state space p Λ , pp e

Original: Lλ i ∆ ´ 1 q ´ 1 q 1 ď i ď N , w qq and sample time ∆ its kernel p K can be obtained from Equation 4 as

Original: [Equation]

Original: which is also the expression for K k (Equation 7).

Original: A.2 Numerically Stable softmax

Original: As noted in §3.1, softmax can have singularities over C .

Original: To address this issue, we use a simple correction to make it well-defined over the entire domain:

Original: • softmax : Given p x 0 , . . .

Original: , x L ´ 1 q “ x P C L , let softmax p x q P C L be defined as p softmax

Original: p x qq k “ e x k p e x 0 ` . . .

Original: ` e x L ´ 1 q ´ 1 .

Original: Note that for any c P C , softmax p x 0 , . . .

Original: , x L ´ 1 q “ softmax p x 0 ´ c, . . .

Original: , x L ´ 1 ´

Original: c q .

Original: Unlike over R , softmax can have singularities over C as sum of exponentials can vanish.

Original: E

Original: g. e 0 ` e iπ “ 0 and hence softmax p 0 , iπ q is not defined.

Original: • max :

Original: Given p x 0 , . . .

Original: , x L ´ 1 q “ x P C L , let max p x q be the x i with the maximum real part, i.e. x argmax i Re p x i q .

Original: • reciprocal ϵ : Given x P C and ϵ P R ą 0 , let reciprocal

Original: ϵ

Original: p x q “ x x ¨ x ` ϵ where x is the complex conjugate of x .

Original: The denominator is always in R ě ϵ and | reciprocal ϵ | ď p 2 ?

Original: ϵ q ´ 1 . • softmax ϵ :

Original: Given p x 0 , . . .

Original: , x L ´ 1 q “ x P C L let m “ max p x q and r x i “ x i ´ m .

Original: Note that | e r x i | ď 1 .

Original: Given ϵ P R ą 0 , let softmax ϵ p x q P C L be p softmax

Original: ϵ p x qq k “ e r x k ¨ reciprocal ϵ ˜ L

Original: [Equation]

Original: 1 ÿ r “e r x r ¸ .

Original: softmax ϵ is always bounded and differentiable.

Original: In our implementation, we use softmax ϵ with ϵ “ 10 ´ 7 .

Original: SSM Softmax In our current implementation of softmax p x q (Figure 6), we exploit the specific structure of x that arises in Algorithm 1.

Original: We now describe an alternate method based on FFT which also uses this specific structure and might lead to a faster implementation in the future.

Original: Claim 1 (SSM Softmax) .

Original: Given c P C , let p “ I r Re p c q ą 0 s , n “ 1 ´ p , e “ exp p c ¨ p n ´ p qq and r “ p n ´ pe q{p p ´ ne q .

Original: Let ω “ exp p

Original: ´

Original: 2 πi { L q where i “ ?

Original: ´ 1 .

Original: Then, softmax p c ¨ 0 , . .

Original: , c ¨ p L ´ 1 qq “ inverseFFT

Original: [Equation]

Original: p p

Original: [Equation]

Original: ne

Original: [Equation]

Original: ˙

Original: [Equation]

Original: ď

Original: [Equation]

Original: “ inverseFFT

Original: [Equation]

Original: r ` ω k ˙ 0 ď k ă L .

Original: Proof.

Original: There are 2 cases depending on sign of Re p c q .

Original: Case 1 ( p “ 0 , n “ 1 ): In this case we have e “ exp p c q .

Original: For the map

Original: [Equation]

Original: we get the coefficients of F p z q as

Original: [Equation]

Original: We have,

Original: [Equation]

Original: where last equality follows from ω L “ 1 .

Original: Case 2 ( p “ 1 , n “ 0 ):

Original: In this case we have e “ exp p´ c q .

Original: For the map

Original: [Equation]

Original: we get the coefficients of F p z q as

Original: [Equation]

Original: as e ´ 1 “ exp p c q .

Original: Moreover, we have

Original: [Equation]

Original: where last equality follows from ω L “ 1 .

Original: Finally, the second equality of the main Claim follows from 1 ´ e “ n ´ pe ` p ´ ne .

Original: The computation of softmax in Claim 1 is numerically stable and we always exponentiate scalars with a negative real part.

Original: The computed function has singularities at c P t´ 2 πik

Original: { L, 0 ď k ă L u .

Original: A.3 Experimental Setup

Original: We now describe the training details for DSS and S4 on LRA and Speech Commands (§4).

Original: Sequence Classification Head: Both LRA and Speech Commands are sequence classification tasks.

Original: The final layer of the DSS stack outputs a sequence which is aggregated into a single vector via mean pooling along the length dimension.

Original: Exceptions to this were T EXT and P ATHFINDER tasks where the rightmost token was used as the aggregate.

Original: We used a separate ∆ log , re , ∆ log , im P R parameters to respectively scale the real and imaginary parts of Λ .

Original: I.e. for a given Λ P C N , we computed ∆ ˚ Λ as exp p ∆ log , re q Λ re ` i ¨ exp p

Original: ∆ log , im q Λ im .

Original: For all datasets, we used AdamW optimizer with a constant learning rate schedule with decay on validation plateau.

Original: However, for the DSS parameters (§3.2) initial learning rate was 10 ´ 3 and weight decay was not used, with a few exceptions noted below.

Original: We used hyperparameters such as model sizes, number of update steps, etc as recommended by the

Original: S4 authors on their official repository and are listed in Table 4.

Original: We made the following exceptions for

Original: DSS trainings:

Original: • L IST O PS : learning rate of ∆ log was 0 .

Original: 02 instead of 10 ´ 3 .

Original: • T EXT : learning rate of ∆ log 0 .

Original: 02 instead of 10 ´ 3 .

Original: • I MAGE : we used seed 0 and trained for 200 epochs instead of 100 .

Original: [Table]
Revised: [N/A]

Original: • P ATHFINDER : we used Patience “ 13 .

Original: • P ATH -X: we used batch size 16 and trained for 35 epochs.

Original: ∆ log was initialized as e r where r „ U p log p .

Original: 0001 q , log p .

Original: 01 qq and its learning rate was 10 ´ 4 .

Original: This was beneficial in early convergence of the model.

Original: For our experiments, the test accuracy that we report in §4 was measured at the checkpoint with the highest validation accuracy.

Original: All our experiments were conducted on a single A100 GPU (40GiB).

Original: A.4 Benchmarking Running Times of DSS and S

Original: We compared the running time of DSS with that of the official S4 implementation on a single

Original: NVIDIA 3090 GPU.

Original: Implementations of both S4 and DSS versions below utilize the PyKeOps library for memory efficiency [CFG ` 21].

Original: • Running times on the P ATH -X task with input length L “ 16384 are summarized in Table

Original: For large batches the time taken

Original: to perform the FFT-based convolution and the feedforward part dominates and hence the speedups of DSS over S4 are less pronounced.

Original: For smaller B , the time taken to compute the kernel becomes significant and DSS provides upto 1 .

Original: 8 ˆ speedup over S4.

Original: • We also isolated and benchmarked the time taken to solely compute the S4 kernel vs the DSS kernel (including the time to perform a backward pass on the sum of all kernel entries).

Original: As summarized in Table 6, for large L , DSS kernel is more than 2 ˆ faster than S4 kernel.

Original: [Table]

Original: [Table]

Original: A.5 Additional Remarks

Original: Casting the kernel to R In this work, we require the inputs u and outputs y to always be over

Original: R , i.e., we assume Equation 5 to be with an explicit casting operation

Original: y k

Original: “ Re p ř kj “ 0 K j ¨ u k ´ j q .

Original: For u over R , this further implies y k “ ř kj “ 0 Re p K

Original: j q ¨ u k ´ j .

Original: Hence, we can explicitly cast the complex-valued kernel produced by Proposition 1 to R during training (Algorithm 1, final step) and cast the outputs over C to R during inference in §3.

Original: [Table]

Original: [Figure]

Original: a y

Original: [Figure]

Original: def reciprocal(x, epsilon=1e-7): x_conj = x.conj() # conjugate return x_conj / (x*x_conj + epsilon) def dss_kernel(L): # L: kernel length # Lambda: [N 2], log_dt:

Original: [H],

Original: W: [H N 2] (floats) Lambda, log_dt, W = get_layer_parameters() # complex parameter stored as 2 floats denoting real, # imaginary parts as ADAM moments are non-linear # convert reals to complex Lambda, W = map(torch.view_as_complex, (Lambda, W))

Original: # [N], [H N] dt_Lambda = log_dt.exp().unsqueeze(-1)

Original: Lambda # [H L] pos = torch.arange(L, device=W.device)

Original: # [L] P = dt_Lambda.unsqueeze(-1)

Original: * pos #

Original: [H N L] # fast softmax using structure of P Lambda_gt_0 =

Original: Lambda.real >

Original: # [N] if Lambda_gt_0.any(): withtorch.no_grad():

Original: P_max=

Original: dt_Lambda * (Lambda_gt_0 * (L-1))

Original: # [H N] P = P - P_max.unsqueeze(-1)

Original: S = P.exp()

Original: [H N L] dt_Lambda_neg =

Original: dt_Lambda * (1 - 2*Lambda_gt_0) # [H N] # 1 / S.sum(-1)

Original: =

Original: = num / den num = dt_Lambda_neg.exp() -

Original: # [H N] den = (dt_Lambda_neg * L).exp() -

Original: # [H N] W = W * num * reciprocal(den * Lambda)

Original: # [H N] # mixture of softmaxes return torch.einsum('hn,hnl->hl', W, S).real

Original: # [H L] def state_space(u): # u: batch of input sequences # B: batch size, H: hidden size, L: sequence length B, H, L = u.shape

Original: # compute state space kernel for each of H coordinates K = dss_kernel(L) # [H L]

Original: # multiply two degree L-1 polynomials # (u0 + u1*z ... uL-1*z^L-1)(K0 + K1*z ... KL-1*z^L-1)

Original: # zero-pad them to degree 2L-1 to avoid wrap-around K_f = torch.fft.rfft(K, n=2*L)

Original: [Equation]

Original: [Equation]

Original: ] # yi =

Original: [-ui*K
-]
Revised: [-ui*K
-]

Original: + ... u0*Ki # residual connection, non-linearity, output projection not shown return y