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

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Modeling long range dependencies in sequential data is a fundamental step towardsattaining human-level performance in many modalities such as text, vision, audioand 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 inmodeling short-range interactions, their performance on tasks requiring long rangereasoning 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 overstate-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: Thecore proposition of S4 is the parameterization of state matrices via a diagonalplus 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 efficient computation.

Original: In this work, we showthat one can match the performance of S4 even without the low rank correctionand 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, speechclassification on Speech Commands dataset, while being conceptually simpler andstraightforward to implement.
Revised: 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.

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 becomethe standard in natural language processing, exhibiting impressive amounts of linguistic and worldknowledge [BMR ` 20, BMH ` 21].
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 areasof 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 inputsof 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 modelis 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 performanceon benchmarks created to formally evaluate and quantify a model’s ability to perform long-rangereasoning (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 modelthat 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 remarkableperformance 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: Forinstance, on Long Range Arena it advances the state-of-the-art by 19 accuracy points over the bestperforming 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 carryover 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, whichis 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 thismatrix, they parameterize the learned state spaces using a Diagonal Plus Low Rank (DPLR) structureand, as a result, need to employ several reduction steps and linear algebraic techniques to be able tocompute 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 efficiently, making S4 difficult 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 areless effective, we observe that there do in fact exist effective diagonal state matrices: simply removingthe low-rank component of the DPLR HiPPO matrix still preserves its performance.
Revised: 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.

Original: Leveraging thisidea, our proposed Diagonal State Space (DSS) model enforces state matrices to be diagonal, makingit significantly simpler to formulate, implement and analyze, while being provably as expressive asgeneral state spaces.
Revised: 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.

Original: In contrast to S4, DSS does not assume any specialized background beyondbasic 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 asingle page and is provided in §A.7 (Figure 6).
Revised: Our implementation fits 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, andvisual-spatial reasoning over a wide range of modalities such as text, natural/synthetic images, andmathematical expressions.
Revised: 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.

Original: Despite its simplicity, DSS delivers an average accuracy of 81 . 88 acrossthe 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, DSSmaintains 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 classification.

Original: On the
Revised: On the

Original: Speech Commands dataset [War18], which consists of raw audio samples of length 16 K , we againfound 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 modelinglong-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 modelsover their attention-based peers.
Revised: 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.

Original: Our code is available at https://anonymized .
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:

Original: [Equation]
Revised: [Equation]

Original: (1) which defines a function-to-function map u p t q ÞÑ y p t q .
Revised: which defines 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 Rdenotes 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 denotesthe 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 1 on u is defined as a sequence-to-sequence map fromp 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 defined 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 define 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 therecurrence 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 inputsand, instead, Equation 5 can be used to compute all elements of y in parallel, provided we havealready 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 efficiently by observing that for the univariate polynomials

Original: [Equation]
Revised: [Equation]

Original: u p z

Original: [Equation]

Original: z
Revised: p

Original: [Equation]
Revised: [Equation]

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

Original: q ¨ u p z q , i.e. all y k ’s can be computed simultaneouslyby 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 qqtime 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 viathe 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 computedefficiently from the input.
Revised: Hence, given the SSM kernel K P R L , the output of a discretized state space can be computed efficiently from the input.

Original: The challenging part is computing K itself as it involves computing Ldistinct 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 touse 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
Revised: A “ V Λ V ´ 1 , where Λ is diagonal

Original: , then we could easily compute e A ¨ k ∆ “ V e Λ ¨ k ∆ V ´ 1 as it istrivial to exponentiate diagonal matrices.
Revised: , 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 simplifies to

Original: [Equation]
Revised: [Equation]

Original: Leveraging the above observation, we base our model on the following proposition which assertsthat, 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 statespace p A, B, C q over C and sample time ∆ .
Revised: 1 ˆ L to denote the kernel (Equation 4) of length L for state space 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 sampletime ∆ ą 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 @
Revised: , λ N and @

Original: i , λ i ‰ 0and e
Revised: i , λ i ‰and

Original:  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: [Equation]
Revised: (a) K “ K ∆ ,L p Λ , p 1 q 1 ď i ď N , r w q “ r w ¨ Λ ´ 1 p e Λ∆

Original: [Equation]
Revised: ´ I q ¨ elementwise - exp p P q , (b) K “ K ∆ ,L p Λ , pp e Lλ i ∆ ´ 1 q ´ 1 q 1 ď i ď N , w q “ w ¨ Λ ´ 1 ¨ row - softmax p P q .

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 byusing Equation 4 to explicitly compute the expression for the kernel of the corresponding diagonal
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

Original: [Equation]
Revised: [Equation]

Original: A j j ! “ ř 8 j “ 0 p V Λ V ´ 1 q j j !
Revised: Input: parameters Λ re , Λ im P R N , w P C N , sample time ∆ log P R Output: SSM kernel K P R L (Proposition 1(b))

Original: “ V ´ j “ 0 Λ j j ! Input: parameters
Revised: 1: Λ Ð Λ [Equation] ` i

Original: [Equation]
Revised: [Equation]

Original: R
Revised: Λ

Original: Output:
Revised: [Equation] ∆ Ð

Original: [Equation]
Revised: [Equation]

Original: kernel
Revised: ∆ log q [Table] 2: P N ˆ

Original: [Equation]
Revised: [Equation]

Original: Proposition 1(b)
Revised: Ð p ∆

Original: [Equation]
Revised: [Equation]

Original: Ź ∆ is positive real
Revised: Λ q N [Equation] ˚ r 0 , 1 , . . . L ´

Original: [Equation]
Revised: [Equation]

Original: Ź outer product
Revised: s

Original: [Equation]
Revised: [Equation]

Revised: ˆ L

Original: Ź numerically stable (§A.2) 
Revised: [Table] outer product 3: S Ð row - softmax ϵ p P q [Table] (§A.2) 4: K Ð Re p p w { Λ q 1 ˆ N ¨ S q

Original: [Equation]
Revised: [Table]

Original: Ź real part state space.
Revised: state space.

Original: Hence, for any given state space with a well-behaved state matrix there exists a diagonalstate 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: 3 More importantly, the expressions for the kernels of the saiddiagonal 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 computethe 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 thematrix elementwise - exp p P q contains terms as large as exp p λ i ∆
Revised: This is because the matrix elementwise - exp p P q contains terms as large as exp p λ i ∆

Original: p L ´ 1 qq which even for a modestvalue of L can be very large.
Revised: 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 partof 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 asexp 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 ą 0 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 ofthe 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 detrimentalfor some tasks, and we now present an alternate method that provides the simplicity of Proposition1(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 elementsof 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 computedas Λ re ` i ¨ Λ im , ∆ is computed as exp p ∆ log q P R ą 0 , and the kernel is then computed via Proposition1(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 , softmaxcan have singularities over C and slight care must be taken while computing it (e.g., softmax p 0 , iπ qis 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 defined).

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 replacethe 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 usingone 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 withparameters Λ 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 iscomputed 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 informationexchange 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 implementationof 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 evensimpler 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 convolutionand 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 thekernels becomes important whereas for large batches more compute is spent on the convolution andthe 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: Inline with the past work, we found that carefully initializing the parameters of the DSS layer is crucialto 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 ofthe 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 areinitialized such that the resulting Λ is the vector of those N eigenvalues of the following 2 N ˆ 2 Nmatrix 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 ourexperiments, 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 DSSparameters 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 arenoted 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 insteadleveraged 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 duringinference, 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 efficient 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 
Revised: DSS EXP

Original: 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 statespace 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 SSMrecurrence 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 andhence 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 computedx 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 , wewould 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 beforgotten 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.

Original: DSS 
Revised: DSS SOFTMAX

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

Original: Lλ i ∆ ´ [Equation] 1 ď i ď N , w q .
Revised: Lλ i ∆ ´ 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]

Revised: , e λ N ∆ q , B “ ˆ e λ i ∆ ´λ i

Revised: p

Revised: e λ i ∆ L ´ 1 q ˙ 1 ď i ď N .

Original: [Table]
Revised: [Table]

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, directlycomputing such terms can result in numerical instability and we must avoid exponentiating scalarswith 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 anintermediate 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 previoustimesteps 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 ą 0then we would have
Revised: On the other hand if Re p λ q ąthen we would have

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: 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 { λ .

Original: Hence, the model would be able to copy information even from extremely distant positions, allowingit to capture long-range dependencies.
Revised: 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 andaudio.
Revised: We evaluate the performance of DSS on sequence-level classification tasks over text, images and audio.

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

Original: Long Range Arena (LRA) LRA [TDA ` 21] is a standard benchmark for assessing the ability ofmodels 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]
Revised: [Table]

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 whereasother 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 . 88 vs80 . 21 ).
Revised: Its performance is comparable to that of S4, with a modest improvement in test accuracy averaged across the 6 tasks ( 81 . 88 vs 80 . 21 ).

Original: In addition, DSS maintains a comfortable 20 point lead over the best performing Transformervariant ( 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: Interestingly, despite being less expressive than DSS SOFTMAX , DSS EXP also reports an impressiveperformance 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 sameas DSS EXP except that we omit the term Λ ´ 1 p e Λ∆ ´ I q in the expression of the kernel shown in
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 Λ∆ ´ 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 r w ¨ elementwise - exp p P q .

Original: As shown in Table 1, theperformance of DSS EXP - NO - SCALE is generally inferior to that of DSS EXP with the model failing on thechallenging 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: Raw Speech Classification Audio is typically digitized using a high sampling rate resulting invery long sequences.
Revised: Raw Speech Classification 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-rangemodels.
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 classification task.

Original: As shownin 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 hiddensize, 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 frameworkprovided 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 official instructions.

Original: 7 Details aboutmodel 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 factorscontributing to its performance.
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 contributing to its performance.

Original: To investigate this further, we performed an ablation analysis aimedat answering the following questions:
Revised: To investigate this further, 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 significant is the Skew-Hippo initialization (§3.3) to the model performance?

Original: Would initializing Λ randomly work just as well?• Is the main source of superior performance of state space models (S4, DSS), compared to previousmodels, their ability to model long-range dependencies?
Revised: Would initializing Λ randomly work just as well? • 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 modellocal 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 initializingeach element of Λ re and Λ im in DSS SOFTMAX by randomly sampling from N p 0 , 1 q .
Revised: 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 .

Original: [Table]
Revised: [Table]

Original: Truncated Kernels To answer the second question, instead of constructing a kernel of lengthequal 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 , significantly shorter than the length of the input.

Original: To understand the implicationof this restriction recall Equation 5 which states that y k “ ř kj “ 0 K j ¨ u k ´ j .
Revised: To understand the implication of this restriction recall Equation 5 which states that y k “ ř kj “ 0 K j ¨ u k ´ j .

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 restricteach DSS SOFTMAX layer to only model local interactions and the model would require several layersto 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 field and capture long-range interactions.

Original: As shown in Table 3, randomly initializing the Λ parameters of DSS leads to significant performancedegradation 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 significant performance degradation on the majority of tasks, with the model failing to perform on P ATH -X.

Original: This is inline withthe findings of [GGR22] who also reported the initialization of S4 to be critical to its performance.
Revised: This is inline with the findings 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 mosttasks (Table 3), suggesting that the superior performance of state-space models on these taskscan indeed be partly attributed to their ability to capture long-range dependencies.
Revised: 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.

Original: Moreover, onsome 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 capturinginteractions 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 andkernels 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 contrastbetween 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 thefirst 128 positions are significantly higher than the later positions indicating that these kernels aremostly local.
Revised: 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.

Original: On the other hand for P ATH -X, for a significant proportion of kernels the opposite istrue, indicating that these kernels are modeling long-range interactions.
Revised: 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.

Original: This partially explains why the performance of DSS SOFTMAX on I MAGE and SC does not decreasesignificantly even after limiting the kernel length to 128 , whereas P ATHFINDER and P ATH -X reportsignificant degradation (Table 3).
Revised: 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).

Original: The last row of Table 3 shows the 95th percentile of the sett argmax 

Revised: The last row of Table 3 shows the 95th percentile of the set t argmax

Original: ď k ă L | K k | : K is a kernel in DSS SOFTMAX u , i.e., the set of positions over all kernels, atwhich for a given kernel its absolute value is maximized.
Revised: ď 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 4).
Revised: The values of the real and imaginary parts of Λ are shown in Figure 3 (and §A.6, Figure 4).

Original: Weobserve that, although all elements of Λ re are initialized as ´ 0 . 5 , their values can change significantlyduring 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 significantly 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 propertynot satisfied by Skew-Hippo initialization.
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 satisfied by Skew-Hippo initialization.

Original: The values of the trained ∆ log corresponding to the real and imaginary parts of Λ are shown in §A.6(Figure 5) 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 §A. (Figure 5) and, in this case as well, change significantly during training.

Original: The values of ∆ log onshort-range tasks such as I MAGE and SC generally tend to be larger compared to long-range taskssuch as 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).
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 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: 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 andwe 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: [Figure]
Revised: [Figure]

Original: [Figure]

Original: ( i 
Revised: i dden

Original: [Figure] 
Revised: [Figure] [Figure]

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

Original: In this work, we evaluated DSS on sequence-level classificationtasks.
Revised: In this work, we evaluated DSS on sequence-level classification 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 amountsof 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, fortasks such as P ATH -X that require very long-range interactions, a smaller initialization of ∆ log wasbeneficial.
Revised: Informally, for tasks such as P ATH -X that require very long-range interactions, a smaller initialization of ∆ log was beneficial.

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: Checklist
Revised: Acknowledgments and Disclosure of Funding

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

Original: For all authors...
Revised: Tel Aviv University. 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

Original:  (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 reflect 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 Page 9.
Revised: [Yes] See end of Page 9.

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 tothem?
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.3.
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., typeof 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]
Revised: [No]

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

Original: [N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: [N/A] (d) Did you discuss whether 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 identifiable information or offensive content?

Original: [N/A] 
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? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

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

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: We restate Proposition 1 for convenience. Proposition. Let K P C

Original: spent on participant compensation?
Revised: ˆ L be the kernel of length L of a given state space p A, B, C q and sample time ∆ ą 0 , where A P C N ˆ N is diagonalizable over C with eigenvalues λ 1 , . . . , λ N and @ i , λ i ‰and

Revised: e Lλ i ∆ ‰ 1 .

Revised: Let P P C N ˆ L

Revised: be P

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

Revised: , λ N .

Revised: Then there exist r w

Revised: [Equation]

Revised: w P C 1 ˆ N such that

Revised: [Equation]

Revised: Proof.

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

Revised: , λ N P C .

Revised: From

Revised: Equation 4 we have

Revised: [Equation]

Revised: where

Revised: [Equation]

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

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

Revised: Then,

Revised: [Equation]

Revised: where the last equality follows from

Revised: p z

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

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

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

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

Revised: [Equation]

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

Revised: [Equation]

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

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

Revised: Let w P C N be defined as

Revised: [Equation]

Revised: Then from Equation 9,

Revised: [Equation]

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

Revised: [Equation]

Revised: i,k

Revised: [Equation]

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

Revised: [Equation]

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

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

Revised: [Equation]

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

Revised: A.2 Numerically Stable softmax

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

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

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

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

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

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

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

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

Revised: , x L ´ 1 ´

Revised: c q .

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

Revised: E

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

Revised: • max :

Revised: Given p x 0 , . . .

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

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

Revised: ϵ

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

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

Revised: ϵ q ´ 1 . • softmax ϵ :

Revised: Given p x 0 , . . .

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

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

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

Revised: ϵ p x qq k “ e r x k ¨ reciprocal

Revised: ϵ

Revised: [Equation]

Revised: softmax ϵ is always bounded and differentiable.

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

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

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

Revised: Claim 1 (SSM Softmax) .

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

Revised: Let ω “ exp p

Revised: ´

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

Revised: ´ 1 .

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

Original: {+, c ¨ p L ´ 1 qq “ inverseFFT ˆ 1 ´ e n ´ pe ` p p ´ ne q ω k ˙ 
+}
Revised: {+, c ¨ p L ´ 1 qq “ inverseFFT ˆ 1 ´ e n ´ pe ` p p ´ ne q ω k ˙ 
+}

Original: {+ď k ă L “ inverseFFT ˆ r `r ` ω k ˙ 
+}
Revised: {+ď k ă L “ inverseFFT ˆ r `r ` ω k ˙ 
+}

Revised: ď k ă L .

Revised: Proof.

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

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

Revised: For the map

Revised: [Equation]

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

Revised: [Equation]

Revised: We have,

Revised: [Equation]

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

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

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

Revised: For the map

Revised: [Equation]

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

Revised: [Equation]

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

Revised: Moreover, we have

Revised: [Equation]

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

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

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

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

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

Revised: A.3 Experimental Setup

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

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

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

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

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

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

Revised: ∆ log , im q Λ im .

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

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

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

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

Revised: We made the following exceptions for

Revised: DSS trainings:

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

Revised: 02 instead of 10 ´ 3 .

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

Revised: 02 instead of 10 ´ 3 .

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

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

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

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

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

Revised: 0001 q , log p .

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

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

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

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

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

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

Revised: NVIDIA 3090 GPU.

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

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

Revised: Table 5.

Revised: For large batches the time taken to perform the FFT-based convolution and the feedforward part dominates and hence

Revised: the speedups of DSS over S4 are less pronounced.

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

Revised: 8 ˆ speedup over S4.

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

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

Revised: [Table]

Revised: [Table]

Revised: A.5 Additional Remarks

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

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

Revised: y k

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

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

Revised: j q ¨ u k ´ j .

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

Revised: [Figure]

Revised: l a y

Revised: [Figure]

Revised: [Figure]

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

Revised: [H],

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

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

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

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

Revised: * pos #

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

Revised: Lambda.real >

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

Revised: P_max=

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

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

Revised: S = P.exp()

Revised: [H N L] dt_Lambda_neg =

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

Revised: =

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

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

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

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

Revised: # [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

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

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

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

Revised: [Equation]

Revised: [Equation]

Revised: ] # yi =

Original: {+ui*K
+}
Revised: {+ui*K
+}

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