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)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: 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 towardsattaining human-level performance in many modalities such as text, vision, audioand 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 inmodeling short-range interactions, their performance on tasks requiring long rangereasoning 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: [GGR22]proposed the Structured State Space (S4) architecture delivering large gains overstate-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: Thecore proposition of S4 is the parameterization of state matrices via a diagonalplus 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 showthat one can match the performance of S4 even without the low rank correctionand 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, speechclassification on Speech Commands dataset, while being conceptually simpler andstraightforward 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 becomethe standard in natural language processing, exhibiting impressive amounts of linguistic and worldknowledge [BMR ` 20, BMH ` 21].

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 areasof 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 inputsof 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 modelis 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] (cf.[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 performanceon benchmarks created to formally evaluate and quantify a model’s ability to perform long-rangereasoning (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 modelthat 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 remarkableperformance 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: Forinstance, on Long Range Arena it advances the state-of-the-art by 19 accuracy points over the bestperforming 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 carryover 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, whichis 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 thismatrix, they parameterize the learned state spaces using a Diagonal Plus Low Rank (DLPR) 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.

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

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

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 beyondbasic 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 asingle page and is provided in §A.5 (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, andvisual-spatial reasoning over a wide range of modalities such as text, natural/synthetic images, andmathematical 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 acrossthe 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, DSSmaintains 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 againfound 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 modelinglong-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 modelsover their attention-based peers.

Original: Our code is available at https://anonymized .
Revised: Our code is available at https://anonymized .

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

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 1 on u is defined as a sequence-to-sequence map fromp 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 therecurrence 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 inputsand, instead, Equation 5 can be used to compute all elements of y in parallel, provided we havealready 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 simultaneouslyby 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 qqtime 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 viathe 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 computedefficiently 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 Ldistinct 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: Instead of directly using A , B , C , our idea is to use an alternateparameterization 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

Original: A ¨ k ∆ for all 0 ď k ă

Original: L (Equation 4).

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

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

Original: In this case Equation 4 simplifies to

Original: [Equation] 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: Our model is based 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 1 ˆ L to denote the kernel (Equation 4) of length L for statespace p A, B, C qover C and sample time ∆ .
Revised: We use the operator K ∆ ,L p A, B, C q P 93 R 1 ˆ L to denote the kernel (Equation 4) of length L for state space p A, B, C q and sample time ∆ .

Original: Proposition 1.
Revised: Proposition 1.

Original: Let K P C
Revised: Let K P R 1 ˆ

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 sampletime ∆ ą 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 ‰ 0and e

Original: Lλ i ∆ ‰ 1 .
Revised: 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: [Equation]

Original: [Equation]
Revised: [Equation]

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

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: [Equation]

Original: [Equation] A j j ! “ ř 8 j “ 0 p V Λ V ´ 1 q j j ! “ V ´ j “ 0 Λ j j !
Revised: softmax

Original: Input: parameters

Original: [Equation]
Revised: [Equation]

Original: R
Revised: q

Original: Output:

Original: [Equation]
Revised: [Equation]

Original: kernel

Original: [Equation] Proposition 1(b) [Equation] Ź ∆ is positive real
Revised: The proof of Proposition 1 is elementary and is provided in §A.1.

Original: [Equation] Ź outer product [Equation] Ź numerically stable (§A.2)
Revised: In the above equations, the lastequality follows by using Equation 4 to explicitly compute the expression for the kernel of thecorresponding diagonal state space.

Original: [Equation] Ź real part state space. 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 matrixthere 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: 2 More importantly, the expressions forthe kernels of the said diagonal state spaces no longer involve matrix powers but only a structuredmatrix-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 computethe 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 thematrix 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 modestvalue 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.

Revised: Input: parameters Λ re , Λ im P R N

Revised: [Equation]

Revised: w P C N , sample time ∆ log P R

Original: DSS EXP
Revised: Output: SSM kernel K P R L (Proposition 1(b)) [Equation] Ź [Table]

Revised: is positive real

Revised: [Equation]

Revised: Ź outer product

Revised: [Equation]

Revised: Ź numerically stable (§A.2)

Revised: [Equation]

Revised: Ź real part

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 partof 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 asexp 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 ofthe 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 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.

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 elementsof 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 ¨ Λ
Revised: Λ is computedas Λ re ` i ¨ Λ

Original: im , ∆ is computed as exp p ∆ log q P R ą 0 , and the kernel is then computed via Proposition1(b).
Revised: im , ∆ is computed as exp p ∆ log q P R ą 0 , and the kernel is then computed via Proposition1(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 , softmaxcan have singularities over C and slight care must be taken while computing it (e.g., softmax p 0 , iπ qis 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 replacethe 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 usingone 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 withparameters Λ 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 iscomputed 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 informationexchange 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 implementationof DSS SOFTMAX layer is provided in §A.5 (Figure 6).

Original: The implementation of DSS EXP layer is evensimpler  is omitted.
Revised: The implementation of DSS EXP layer is evensimpler 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 convolutionand 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 thekernels becomes important whereas for large batches more compute is spent on the convolution andthe 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: 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).

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. [Equation]
Revised: 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.$& % p 2 i ` 1 q 1 { 2 p 2 j ` 1 q 1 { 2 {i ă j ´ 1 {i “ j ´p 2 i ` 1 q 1 { 2 p 2 j ` 1 q 1 { 2 { 2 i ą j

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 ourexperiments, 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 DSSparameters 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 arenoted 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 insteadleveraged 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 duringinference, 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 statespace 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 SSMrecurrence 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 andhence 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 computedx 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 , wewould 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 beforgotten 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]

Original: [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 isnow a scalar.
Revised: where x k ´ 1 is 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, directlycomputing such terms can result in numerical instability and we must avoid exponentiating scalarswith 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 anintermediate 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]

Revised: [Table]

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∆ !  then r x k « u k and x k « u k { λ .
Revised: and hence if Re p λ q ! 0 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 previoustimesteps 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 ą 0then 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 tocapture 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, 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 ofmodels 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] 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 whereasother 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 vs80 . 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 Transformervariant ( 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 impressiveperformance 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 [Table] 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, 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.

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 invery 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-rangemodels.

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 shownin 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 hiddensize, 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 frameworkprovided 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: 4 Details aboutmodel 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 factorscontributing 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 aimedat 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 previousmodels, 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 modellocal 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 initializingeach element of Λ re and Λ im in DSS SOFTMAX by randomly sampling from N p 0 , 1 q .

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

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 implicationof 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 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.

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: [Table] 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.

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 withthe 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 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.

Original: Moreover, onsome tasks such as L IST O PS and I MAGE , using a truncated kernel still manages to outperform all
Revised: Moreover, onsome 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 capturinginteractions 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 andkernels 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 contrastbetween 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 thefirst 128 positions are significantly higher than the later positions indicating that these kernels aremostly 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 istrue, 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 decreasesignificantly even after limiting the kernel length to 128 , whereas P ATHFINDER and P ATH -X reportsignificant 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 sett 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, atwhich 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.4, 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: 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).

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 propertynot 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.4(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 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 , larger 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 andwe plan to address this issue in future work.  Discussion

Revised: Related work In a long line of work, several variants of the Transformer have been proposed toaddress its quadratic complexity (cf.

Revised: [GB21] and references therein).

Revised: Recently, Gu et al.

Revised: [GGR22]

Original: [Figure]
Revised: [Figure]

Revised: e r s × h i dden - s i z e

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

Original:  ( i 
Revised: cs i nh ( i m ag / [Figure] r cs i nh i m ag / 2 )

Original: [Figure]
Revised: [Figure]

Revised: introduced S4, a new type of model that leverages linear state spaces for contextualization insteadof attention.

Revised: Our work is inspired from S4 but uses a diagonal parameterization of state spaces.

Revised: As a result, our method is significantly simpler compared to S4 and we do not require (1) Padéapproximations to A “ e A ∆ (Euler, Bilinear, etc), (2) Woodbury Identity reductions to computematrix inverse, and (3) fourier analysis for computing the SSM kernel efficiently via truncatedgenerating functions.

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 classificationtasks.

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 pretrainDSS-based models  on  large amountsof raw data.
Revised: Another natural and important future direction is to pretrain models based on DSSover 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 DSSparameters Λ re , Λ im , ∆ log (§3.3) play an important role in the performance and convergence of themodel.

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 smallerinitialization 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 andremains for future work.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes] See end of Page 9.
Revised: [Yes] See §5.

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]
Revised: [No]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

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

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (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.3.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (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: (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 theywere chosen)?

Original: [Yes] See §A.3.
Revised: [Yes] See §A.3.(c)

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

Original: [No]
Revised: [No]

Original: (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: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

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

Original: If you are using existing assets (e
Revised: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...
Revised: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (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]
Revised: [N/A]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

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

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?
Revised: information or offensive content?

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

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

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

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

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

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?
Revised: spent on participant compensation?

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