Original: M ULTI -H EAD S
Revised: M ULTI -H EAD S

Original: TATE
Revised: TATE

Original: S PACE M
Revised: S PACE M

Original: ODEL FOR S EQUENCE
Revised: ODEL FOR S EQUENCE

Original: M ODELING
Revised: M ODELING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Recently, state space models (SSMs) have shown promising results on sequence modeling tasks.
Revised: Recently, state space models (SSMs) have shown promising results on sequence modeling tasks.

Original: However, a potential challenge of existing works is that SSMs are usually introduced or initialized in a homogeneous way, encouraging the model to only capture similar temporal dynamics on different features.
Revised: However, a potential challenge of existing works is that SSMs are usually introduced or initialized in a homogeneous way, encouraging the model to only capture similar temporal dynamics on different features.

Original: In this paper, we propose a multi-head state space model (MSSM), in which parallel heads are introduced to learn different temporal dynamics on sequence data.
Revised: In this paper, we propose a multi-head state space model (MSSM), in which parallel heads are introduced to learn different temporal dynamics on sequence data.

Original: Furthermore, we propose a novel variant of the Transformer, referred to as the Stateformer, which combines MSSMs with attention.
Revised: Furthermore, we propose a novel variant of the Transformer, referred to as the Stateformer, which combines MSSMs with attention.

Original: Experiments on large-scale automatic speech recognition (ASR) and language modeling tasks show the MSSM outperforming a range of attention-based baselines.
Revised: Experiments on large-scale automatic speech recognition (ASR) and language modeling tasks show the MSSM outperforming a range of attention-based baselines.

Original: The Stateformer further improves performance, achieving the state-of-the-art performance on the LibriSpeech ASR task.
Revised: The Stateformer further improves performance, achieving the state-of-the-art performance on the LibriSpeech ASR task.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Transformers (Vaswani et al., 2017) and attention based models (Bahdanau et al., 2015) have for more than half a decade shown state-of-the-art performance on a wide range of tasks anywhere from speech recognition (Zhang et al., 2020a; Gulati et al., 2020) and neural machine translation (Ng et al., 2019; Chen et al., 2020; Tran et al., 2021) to computer vision (Dosovitskiy et al., 2021;
Revised: Transformers (Vaswani et al., 2017) and attention based models (Bahdanau et al., 2015) have for more than half a decade shown state-of-the-art performance on a wide range of tasks anywhere from speech recognition (Zhang et al., 2020a; Gulati et al., 2020) and neural machine translation (Ng et al., 2019; Chen et al., 2020; Tran et al., 2021) to computer vision (Dosovitskiy et al., 2021;

Original: Liu et al., 2021b) and biological applications such as protein sequence modeling (Jumper et al., 2021; Rives et al., 2021).
Revised: Liu et al., 2021b) and biological applications such as protein sequence modeling (Jumper et al., 2021; Rives et al., 2021).

Original: One of the strengths of these models comes from having a significantly smaller maximum path lengths across time, which is the shortest path linking the first encoder input and final decoder output (Hochreiter et al., 2001).
Revised: One of the strengths of these models comes from having a significantly smaller maximum path lengths across time, which is the shortest path linking the first encoder input and final decoder output (Hochreiter et al., 2001).

Original: This is unlike previous state-of-the-art recurrent and convolutional neural networks which have linearly and logarithmically scaling path lengths respectively.
Revised: This is unlike previous state-of-the-art recurrent and convolutional neural networks which have linearly and logarithmically scaling path lengths respectively.

Original: However, a well known drawback of transformers is the quadratic space and time complexity of the self-attention layer, restricting its applicability to fields requiring longer sequences or to devices with strict limitations on compute resources (Tay et al., 2022).
Revised: However, a well known drawback of transformers is the quadratic space and time complexity of the self-attention layer, restricting its applicability to fields requiring longer sequences or to devices with strict limitations on compute resources (Tay et al., 2022).

Original: To combat both of these issues, a wide range of restricted sparse attention mechanisms have been proposed, all aiming at both reducing the computational cost and scaling to longer sequences but retaining as much of the original transformer performance as possible (Katharopoulos et al., 2020; Kitaev et al., 2020; Beltagy et al., 2020; Zaheer et al., 2020).
Revised: To combat both of these issues, a wide range of restricted sparse attention mechanisms have been proposed, all aiming at both reducing the computational cost and scaling to longer sequences but retaining as much of the original transformer performance as possible (Katharopoulos et al., 2020; Kitaev et al., 2020; Beltagy et al., 2020; Zaheer et al., 2020).

Original: There has also been a notable amount effort been poured into designing alternative more efficient attention schemes instead of using the dot-product approach (Wang et al., 2020a; Choromanski et al., 2020;
Revised: There has also been a notable amount effort been poured into designing alternative more efficient attention schemes instead of using the dot-product approach (Wang et al., 2020a; Choromanski et al., 2020;

Original: Shen et al., 2021; Zhang et al., 2019).
Revised: Shen et al., 2021; Zhang et al., 2019).

Original: Meanwhile, the machine learning community is paying more attention to a well established signal processing and control theory technique, the state space model (SSM) (Kalman, 1960), which historically has been widely used in many time-series and control problems (Brogan, 1991; Hyndman et al., 2008; Durbin & Koopman, 2012).
Revised: Meanwhile, the machine learning community is paying more attention to a well established signal processing and control theory technique, the state space model (SSM) (Kalman, 1960), which historically has been widely used in many time-series and control problems (Brogan, 1991; Hyndman et al., 2008; Durbin & Koopman, 2012).

Original: More recently, the linear time-variant state space model has also been used within neural networks for improving time-series forecasting (Seeger et al., 2017; 2016;
Revised: More recently, the linear time-variant state space model has also been used within neural networks for improving time-series forecasting (Seeger et al., 2017; 2016;

Original: Rangapuram et al., 2018).
Revised: Rangapuram et al., 2018).

Original: Furthermore, a simplified linear time-invariant (LTI) SSM, which is closely related to fully linear recurrent neural network (RNN) layers, has a well-known convolution equivalent making it particularly attractive for parallelized training while inference mode can utilize its fast recurrent formulation (Brogan, 1991; Gu et al., 2021).
Revised: Furthermore, a simplified linear time-invariant (LTI) SSM, which is closely related to fully linear recurrent neural network (RNN) layers, has a well-known convolution equivalent making it particularly attractive for parallelized training while inference mode can utilize its fast recurrent formulation (Brogan, 1991; Gu et al., 2021).

Original: One potential limitation of existing works in this direction is that SSMs are usually initialized and used in a homogeneous way.
Revised: One potential limitation of existing works in this direction is that SSMs are usually initialized and used in a homogeneous way.

Original: For instance, in Gu et al.
Revised: For instance, in Gu et al.

Original: (2022), the S4 approach equips SSMs with careful options on parameter initialization for long-range modeling ability.
Revised: (2022), the S4 approach equips SSMs with careful options on parameter initialization for long-range modeling ability.

Original: This design would force the model to capture similar temporal dynamics on different features.
Revised: This design would force the model to capture similar temporal dynamics on different features.

Original: As long and short term dependency would both be useful in sequence modeling, in this paper, we develop a multi-head state space model, which consists of parallel heads to learn different temporal dynamics on sequence data.
Revised: As long and short term dependency would both be useful in sequence modeling, in this paper, we develop a multi-head state space model, which consists of parallel heads to learn different temporal dynamics on sequence data.

Original: We investigate the use of multi-head state space models in sequence tasks as both a replacement and complement to attention.
Revised: We investigate the use of multi-head state space models in sequence tasks as both a replacement and complement to attention.

Original: Since SSMs have been shown to be promising on long sequences, we hypothesize such a model is able to handle both short and long-term dependencies and can operate as an effective alternative to attention.
Revised: Since SSMs have been shown to be promising on long sequences, we hypothesize such a model is able to handle both short and long-term dependencies and can operate as an effective alternative to attention.

Original: Following are our main technical contributions around the state space model:
Revised: Following are our main technical contributions around the state space model:

Original: Stacked and multi-head generalization .
Revised: Stacked and multi-head generalization .

Original: We extend the SSM approach by allowing multihead parallel processing of projected lower-dimensional signals and stack such a layer (within a residual block) for better performance.
Revised: We extend the SSM approach by allowing multihead parallel processing of projected lower-dimensional signals and stack such a layer (within a residual block) for better performance.

Original: Head gating .
Revised: Head gating .

Original: We propose an inter-head gating approach allowing different SSMs within the multi-head layer to communicate.
Revised: We propose an inter-head gating approach allowing different SSMs within the multi-head layer to communicate.

Original: Combination with attention .
Revised: Combination with attention .

Original: We also augment the transformer architecture by simply including a bidirectional SSM residual block prior to the attention for computationally unrestricted applications and state-of-the-art performance; referred to as the Stateformer .
Revised: We also augment the transformer architecture by simply including a bidirectional SSM residual block prior to the attention for computationally unrestricted applications and state-of-the-art performance; referred to as the Stateformer .

Original: With these contributions, we advance the state of attention-free models on large-scale speech recognition and language modeling tasks, and show that our novel multi-head state space model outperforms strong attention-based baselines.
Revised: With these contributions, we advance the state of attention-free models on large-scale speech recognition and language modeling tasks, and show that our novel multi-head state space model outperforms strong attention-based baselines.

Original: We also show that combining the multi-head state space with attention can achieve state-of-the-art performance in large-scale speech recognition.
Revised: We also show that combining the multi-head state space with attention can achieve state-of-the-art performance in large-scale speech recognition.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: This section will cover theory and related work regarding the state space model and its use in deep learning.
Revised: This section will cover theory and related work regarding the state space model and its use in deep learning.

Original: We first start by briefly covering the state space model and a block diagonal structured restriction.
Revised: We first start by briefly covering the state space model and a block diagonal structured restriction.

Original: We further mention a convolution point of view for parallel training (Section 2.1).
Revised: We further mention a convolution point of view for parallel training (Section 2.1).

Original: We finish with a brief discussion of the state space layer as used in neural nets (Section 2.2).
Revised: We finish with a brief discussion of the state space layer as used in neural nets (Section 2.2).

Original: 2. S TATE S PACE M ODEL : T HE L INEAR RNN
Revised: 2. S TATE S PACE M ODEL : T HE L INEAR RNN

Original: A specific realization of the state space model (SSM) is a linear time-invariant (LTI) model (Brogan, 1991) that transforms
Revised: A specific realization of the state space model (SSM) is a linear time-invariant (LTI) model (Brogan, 1991) that transforms

Original: some input signal u ( t ) ∈
Revised: some input signal u ( t ) ∈

Original: R D i to some output y ( t ) ∈ R CD i through some hidden process x ( t ) ∈
Revised: R D i to some output y ( t ) ∈ R CD i through some hidden process x ( t ) ∈

Original: R D h according to
Revised: R D h according to

Original: [Equation]
Revised: [Equation]

Original: where Du ( t ) can be neglected since this is a simple term to compute.
Revised: where Du ( t ) can be neglected since this is a simple term to compute.

Original: To ensure such a model is compatible with discrete signals such as sequences of speech frames or sub-word unit embeddings, a simple discretization, e.g. the bilinear method (Tustin, 1947) can be utilized:
Revised: To ensure such a model is compatible with discrete signals such as sequences of speech frames or sub-word unit embeddings, a simple discretization, e.g. the bilinear method (Tustin, 1947) can be utilized:

Original: [Equation]
Revised: [Equation]

Original: which simply represents a fully linear recurrent neural network (RNN) cell
Revised: which simply represents a fully linear recurrent neural network (RNN) cell

Original: Note that the discretization matrix ∆ ∈ R D
Revised: Note that the discretization matrix ∆ ∈ R D

Original: h × D h is diagonal with positive entries, but can be subsumed by A , B due to scale invariance.
Revised: h × D h is diagonal with positive entries, but can be subsumed by A , B due to scale invariance.

Original: Furthermore, while such a model is highly flexible, providing it with structure can increase its efficiency.
Revised: Furthermore, while such a model is highly flexible, providing it with structure can increase its efficiency.

Original: Therefore, one can restrict the matrices A , B , C to have a block diagonal structure as follows:
Revised: Therefore, one can restrict the matrices A , B , C to have a block diagonal structure as follows:

Original: [Equation]
Revised: [Equation]

Original: where the structure of this block recurrent model is based on the number of input dimensions D i .
Revised: where the structure of this block recurrent model is based on the number of input dimensions D i .

Original: The original multi-variate state space model can now be seen as D i smaller independent sub-SSMs, where each one of these smaller models are R → R C mappings running in parallel.
Revised: The original multi-variate state space model can now be seen as D i smaller independent sub-SSMs, where each one of these smaller models are R → R C mappings running in parallel.

Original: Denoting the sub-hidden dimension by D sh
Revised: Denoting the sub-hidden dimension by D sh

Original: = D h
Revised: = D h

Original: /D
Revised: /D

Original: i we obtain A n ∈ R D sh
Revised: i we obtain A n ∈ R D sh

Original: × D sh , B n ∈ R D sh ×and C n ∈ R C
Revised: × D sh , B n ∈ R D sh ×and C n ∈ R C

Original: × D sh .
Revised: × D sh .

Original: The skip-connection can also be split into blocks of D n ∈ R C × 1 .
Revised: The skip-connection can also be split into blocks of D n ∈ R C × 1 .

Original: Similarly, the discretization matrix can now be seen as a diagonal matrix with a isotropic blocks ∆ n = ∆ n I of the same size as the transition blocks A n , meaning that each sub-state space model has its own discretization time-step that can be tuned to its own input channel.
Revised: Similarly, the discretization matrix can now be seen as a diagonal matrix with a isotropic blocks ∆ n = ∆ n I of the same size as the transition blocks A n , meaning that each sub-state space model has its own discretization time-step that can be tuned to its own input channel.

Original: Furthermore, due to the model being fully linear, one can utilize the widely known duality between convolutions and recurrences, resulting in an efficient expression:
Revised: Furthermore, due to the model being fully linear, one can utilize the widely known duality between convolutions and recurrences, resulting in an efficient expression:

Original: [Equation]
Revised: [Equation]

Original: where L is the input sequence length.
Revised: where L is the input sequence length.

Original: Therefore, state space models and linear RNNs can be trained in convolution mode (Equation 4), fully utilizing the power of GPUs, and be deployed using recurrent mode (Equation 2), efficiently and incrementally generating predictions.
Revised: Therefore, state space models and linear RNNs can be trained in convolution mode (Equation 4), fully utilizing the power of GPUs, and be deployed using recurrent mode (Equation 2), efficiently and incrementally generating predictions.

Original: To further improve the efficiency, a specific diagonal plus low-rank restriction of the sub-transition blocks can be used Gu et al.
Revised: To further improve the efficiency, a specific diagonal plus low-rank restriction of the sub-transition blocks can be used Gu et al.

Original: [Equation]
Revised: [Equation]

Original: where Λ n is a diagonal matrix.
Revised: where Λ n is a diagonal matrix.

Original: This allows the kernel K to be computed efficiently utilizing the
Revised: This allows the kernel K to be computed efficiently utilizing the

Original: Woodbury matrix identity with several fast implementations (Pan, 2001; 2014).
Revised: Woodbury matrix identity with several fast implementations (Pan, 2001; 2014).

Original: 2. P RACTICAL L AYER D ESIGN
Revised: 2. P RACTICAL L AYER D ESIGN

Original: The structured state space model described in Equation 2 is a general linear R D i
Revised: The structured state space model described in Equation 2 is a general linear R D i

Original: → R CD i mapping but processes each channel independently due to the structuring in Equation 3.
Revised: → R CD i mapping but processes each channel independently due to the structuring in Equation 3.

Original: To amend both these problems, a non-linear activation is introduced followed by a point-wise linear layer (see Figure 1a for a pre-norm residual unit) to mix the information across input channels:
Revised: To amend both these problems, a non-linear activation is introduced followed by a point-wise linear layer (see Figure 1a for a pre-norm residual unit) to mix the information across input channels:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Note the output dimension expansion C is set to 1 when using element-wise activations such as
Revised: Note the output dimension expansion C is set to 1 when using element-wise activations such as

Original: GELU (Hendrycks & Gimpel, 2016) or set to 2 when using a gating mechanisms such as GLU (Dauphin et al., 2017).
Revised: GELU (Hendrycks & Gimpel, 2016) or set to 2 when using a gating mechanisms such as GLU (Dauphin et al., 2017).

Original: [Figure]
Revised: [Figure]

Original: To further stabilize the state space, Goel et al.
Revised: To further stabilize the state space, Goel et al.

Original: (2022) used the more stable Hurwitz formulation
Revised: (2022) used the more stable Hurwitz formulation

Original: A n = Λ n − p n p Tn , ensuring the eigenvalues of the transition matrix remain negative when Λ n is negative definite.
Revised: A n = Λ n − p n p Tn , ensuring the eigenvalues of the transition matrix remain negative when Λ n is negative definite.

Original: Practically, the parameters are stored in log-form: time-step ∆ n = exp( z ∆ n ) and diagonal Λ n = − exp( z Λ n ) .
Revised: Practically, the parameters are stored in log-form: time-step ∆ n = exp( z ∆ n ) and diagonal Λ n = − exp( z Λ n ) .

Original: Regarding the initialization of this system, Gu et al.
Revised: Regarding the initialization of this system, Gu et al.

Original: (2020) developed several options for how the transition and input matrices should be set to equip the model with long-range modeling ability which can be motivated from a memorization of signals point of view.
Revised: (2020) developed several options for how the transition and input matrices should be set to equip the model with long-range modeling ability which can be motivated from a memorization of signals point of view.

Original: However, all of the parameters were made trainable, discarding the initialization.
Revised: However, all of the parameters were made trainable, discarding the initialization.

Original:  M ULTI -H EAD S TATE S PACE M ODEL
Revised:  M ULTI -H EAD S TATE S PACE M ODEL

Original: This section introduces the novel multi-head state space model as a natural extension to the linear time-invariant structured state space model as described above.
Revised: This section introduces the novel multi-head state space model as a natural extension to the linear time-invariant structured state space model as described above.

Original: First, we discuss our configuration for the state space layer and its intialization (Section 3.1).
Revised: First, we discuss our configuration for the state space layer and its intialization (Section 3.1).

Original: Second, we propose novel ideas for how the architecture can be improved, including a multi-head generalization (Section 3.2).
Revised: Second, we propose novel ideas for how the architecture can be improved, including a multi-head generalization (Section 3.2).

Original: Third, a new gating mechanism is introduced which was found notably beneficial (Section 3.3).
Revised: Third, a new gating mechanism is introduced which was found notably beneficial (Section 3.3).

Original: Furthermore, a bidirectional architecture is discussed (Section 3.4) and how it can be used with attention based models (Section 3.5).
Revised: Furthermore, a bidirectional architecture is discussed (Section 3.4) and how it can be used with attention based models (Section 3.5).

Original: Finally, the compute complexity of state space models scales only linearly with sequence length allowing it to be used prior to sub-sampling layers without much added cost (Section 3.6).
Revised: Finally, the compute complexity of state space models scales only linearly with sequence length allowing it to be used prior to sub-sampling layers without much added cost (Section 3.6).

Original: 3. T IED P ARAMETERS & I NITIALIZATION
Revised: 3. T IED P ARAMETERS & I NITIALIZATION

Original: The original design in Equation 3 describes a block diagonal simplification to a multi-variate SSM for purposes of efficiency, which can also be seen as several small independent SSMs running in parallel.
Revised: The original design in Equation 3 describes a block diagonal simplification to a multi-variate SSM for purposes of efficiency, which can also be seen as several small independent SSMs running in parallel.

Original: However, this SSM required a smaller separate learning rate for its parameters compared to the rest of the network to achieve the best possible performance.
Revised: However, this SSM required a smaller separate learning rate for its parameters compared to the rest of the network to achieve the best possible performance.

Original: In this work we opt for a slightly more parameter efficient blockwise isotropic implementation, where we restrict the copies to have shared input and transition matrices across all blocks:
Revised: In this work we opt for a slightly more parameter efficient blockwise isotropic implementation, where we restrict the copies to have shared input and transition matrices across all blocks:

Original: [Equation]
Revised: [Equation]

Original: but keep independent output C n and skip D n matrices, and let the learning rate remain the same as the remainder of the network.
Revised: but keep independent output C n and skip D n matrices, and let the learning rate remain the same as the remainder of the network.

Original: By still having different time-step discretizations for each dimension of the input sequence, temporal features of different scale can be extracted while sharing the recurrence parameters.
Revised: By still having different time-step discretizations for each dimension of the input sequence, temporal features of different scale can be extracted while sharing the recurrence parameters.

Original: Furthermore, we choose to simply random initialize all parameters of the SSM by sampling from a the unit Gaussian distribution N (0 , 1) , apart from the diagonal matrix which is shifted Λ 0 = − exp
Revised: Furthermore, we choose to simply random initialize all parameters of the SSM by sampling from a the unit Gaussian distribution N (0 , 1) , apart from the diagonal matrix which is shifted Λ 0 = − exp

Original: ( z Λ 0 ) , z Λ 0 ∼ N (ln( scale ) , 1 ) .
Revised: ( z Λ 0 ) , z Λ 0 ∼ N (ln( scale ) , 1 ) .

Original: This is to ensure that the initial eigenvalues of the transition matrix are not too small and can capture enough of the input signal’s history.
Revised: This is to ensure that the initial eigenvalues of the transition matrix are not too small and can capture enough of the input signal’s history.

Original: Compared with Gu et al.
Revised: Compared with Gu et al.

Original: (2020), the random initialization is simple but shows good performance in practice.
Revised: (2020), the random initialization is simple but shows good performance in practice.

Original: Also, it provides the potential for the multi-head state space model (see section below) in breaking the symmetry between heads instead of using the same deterministic initialization.
Revised: Also, it provides the potential for the multi-head state space model (see section below) in breaking the symmetry between heads instead of using the same deterministic initialization.

Original: 3. S TACKED & M ULTI -H EAD G ENERALIZATION
Revised: 3. S TACKED & M ULTI -H EAD G ENERALIZATION

Original: The basic architecture described above is highly parameter efficient, dominated by the linear layer, and is significantly smaller than multi-head attention and feed-forward nets (FFN).
Revised: The basic architecture described above is highly parameter efficient, dominated by the linear layer, and is significantly smaller than multi-head attention and feed-forward nets (FFN).

Original: Therefore, the design is expanded by stacking y =
Revised: Therefore, the design is expanded by stacking y =

Original: f s ( f s − 1 ( ... f 1 ( u )))
Revised: f s ( f s − 1 ( ... f 1 ( u )))

Original: each with its own set of parameters, see an example in Figure 1b.
Revised: each with its own set of parameters, see an example in Figure 1b.

Original: This would balance the parameters in terms of the ratio between FFN and
Revised: This would balance the parameters in terms of the ratio between FFN and

Original: SSM layers and give the whole model more depthwise/temporal flexibility.
Revised: SSM layers and give the whole model more depthwise/temporal flexibility.

Original: Moreover, we found it useful to project the input into several different lower-dimensional representations which are passed on to independent structured state space models, in a design similar to multihead attention, see Figure 2a.
Revised: Moreover, we found it useful to project the input into several different lower-dimensional representations which are passed on to independent structured state space models, in a design similar to multihead attention, see Figure 2a.

Original: Not only does having shared recurrence parameters across a structured
Revised: Not only does having shared recurrence parameters across a structured

Original: SSM stabilize training, but it gives each structured model in a multi-head framework the flexibility to learn both meaningful time-steps and different types of temporal dynamics.
Revised: SSM stabilize training, but it gives each structured model in a multi-head framework the flexibility to learn both meaningful time-steps and different types of temporal dynamics.

Original: This stacked multi-head state space model (MSSM) requires an additional initial projection layer compared to its single-head implementation, leading to a small increase in parameters.
Revised: This stacked multi-head state space model (MSSM) requires an additional initial projection layer compared to its single-head implementation, leading to a small increase in parameters.

Original: In practice, the projection would take a sequence of D i -dimensional features and project it down to H ∈ { 2 , 4 , 8 , . . .
Revised: In practice, the projection would take a sequence of D i -dimensional features and project it down to H ∈ { 2 , 4 , 8 , . . .

Original: } different, ¯ D i =
Revised: } different, ¯ D i =

Original: D i /H
Revised: D i /H

Original: -dimensional signals, each of which is processed by isotropic blockwise diagonal
Revised: -dimensional signals, each of which is processed by isotropic blockwise diagonal

Original: SSMs.
Revised: SSMs.

Original: Since, the effective dimension H · ¯ D i
Revised: Since, the effective dimension H · ¯ D i

Original: = D i remains the same, the cost of a multi-headed approach boils down to an extra linear projection layer.
Revised: = D i remains the same, the cost of a multi-headed approach boils down to an extra linear projection layer.

Original: 3. I NTER -H EAD G ATED L INEAR U NIT
Revised: 3. I NTER -H EAD G ATED L INEAR U NIT

Original: By default, prior work has used the GELU activation function.
Revised: By default, prior work has used the GELU activation function.

Original: In some experiments (Gu et al., 2022;
Revised: In some experiments (Gu et al., 2022;

Original: Goel et al., 2022; Mehta et al., 2022)
Revised: Goel et al., 2022; Mehta et al., 2022)

Original: the GLU activation (with C = 2 ) was also found beneficial
Revised: the GLU activation (with C = 2 ) was also found beneficial

Original: [Figure]
Revised: [Figure]

Original: which was achieved by gating the channels of an individual sub-state space model:
Revised: which was achieved by gating the channels of an individual sub-state space model:

Original: [Equation]
Revised: [Equation]

Original: However, a multi-head state space architecture with H headed outputs { y ( h ) n
Revised: However, a multi-head state space architecture with H headed outputs { y ( h ) n

Original: [Equation]
Revised: [Equation]

Original: offers many more gating possibilities.
Revised: offers many more gating possibilities.

Original: We propose an inter-head gating (IHG) approach in which the output is computed by mixing the heads according to (where σ refers to the sigmoid):
Revised: We propose an inter-head gating (IHG) approach in which the output is computed by mixing the heads according to (where σ refers to the sigmoid):

Original: [Equation]
Revised: [Equation]

Original: allowing different heads to communicate and gate each other, generally leading to improved results.
Revised: allowing different heads to communicate and gate each other, generally leading to improved results.

Original: Clearly, flattening out the output in Equation 9 leads to the original D i -dimensional output, and note that the number of heads in this approach has to be even, see Figure 2b.
Revised: Clearly, flattening out the output in Equation 9 leads to the original D i -dimensional output, and note that the number of heads in this approach has to be even, see Figure 2b.

Original: While our approach is wildly different, the concept of letting heads communicate has also been proposed in Shazeer et al.
Revised: While our approach is wildly different, the concept of letting heads communicate has also been proposed in Shazeer et al.

Original: (2020) in the context of attention where linear layers mix information across heads, and was found to improve performance.
Revised: (2020) in the context of attention where linear layers mix information across heads, and was found to improve performance.

Original: However, unlike the last-mentioned, our approach allows for communication between different heads without any extra parameters.
Revised: However, unlike the last-mentioned, our approach allows for communication between different heads without any extra parameters.

Original: 3. B IDIRECTIONAL D ESIGN
Revised: 3. B IDIRECTIONAL D ESIGN

Original: So far the proposed architecture has a unidirectional design however, for better performance, bidirectional context is needed.
Revised: So far the proposed architecture has a unidirectional design however, for better performance, bidirectional context is needed.

Original: Therefore, we simply concatenate the output of independent forward and backward stacked multi-head modules and project it back down to the original input dimension.
Revised: Therefore, we simply concatenate the output of independent forward and backward stacked multi-head modules and project it back down to the original input dimension.

Original: This new module is then wrapped in a pre-norm block similar to the architectures above:
Revised: This new module is then wrapped in a pre-norm block similar to the architectures above:

Original: [Equation]
Revised: [Equation]

Original: Since there is a linear layer in the bidirectional block, the final linear layer of the MSSM modules are removed.
Revised: Since there is a linear layer in the bidirectional block, the final linear layer of the MSSM modules are removed.

Original: This attention free design has many advantages over multi-head attention such as memory requirements scaling linearly with sequence length, making it computationally feasible for longer sequences.
Revised: This attention free design has many advantages over multi-head attention such as memory requirements scaling linearly with sequence length, making it computationally feasible for longer sequences.

Original: 3. S TATEFORMER : S TATE S PACE A UGMENTED T RANSFORMER
Revised: 3. S TATEFORMER : S TATE S PACE A UGMENTED T RANSFORMER

Original: A pure multi-head state space model architecture is attractive due to its efficiency and ability to capture both short and long range dependencies.
Revised: A pure multi-head state space model architecture is attractive due to its efficiency and ability to capture both short and long range dependencies.

Original: However, since the state space model is equivalent to a linear RNN, it is expressively more limited, and is still plagued by a linearly scaling path length.
Revised: However, since the state space model is equivalent to a linear RNN, it is expressively more limited, and is still plagued by a linearly scaling path length.

Original: Therefore, for scenarios where one is less restricted by compute limitations, we propose a model combining the MSSM with attention by simply inserting a pre-norm (bidirectional) block prior to the self-attention unit in the transformer architecture, referred to as the Stateformer , see Figure 3.
Revised: Therefore, for scenarios where one is less restricted by compute limitations, we propose a model combining the MSSM with attention by simply inserting a pre-norm (bidirectional) block prior to the self-attention unit in the transformer architecture, referred to as the Stateformer , see Figure 3.

Original: [Figure]
Revised: [Figure]

Original: This type of design is inspired by the Conformer (Gulati et al., 2020) where a dedicated convolutional block is inserted to handle local relations.
Revised: This type of design is inspired by the Conformer (Gulati et al., 2020) where a dedicated convolutional block is inserted to handle local relations.

Original: However, we hypothesize that the advantage of the
Revised: However, we hypothesize that the advantage of the

Original: Stateformer, is the MSSM with its variable length kernel and the multi-head attention both can handle short and long-term dependencies and can complement each other.
Revised: Stateformer, is the MSSM with its variable length kernel and the multi-head attention both can handle short and long-term dependencies and can complement each other.

Original: Furthermore, the attention enables maximum path lengths of 1, resulting in a possibly more flexible model.
Revised: Furthermore, the attention enables maximum path lengths of 1, resulting in a possibly more flexible model.

Original: 3. M ULTI -S CALE A RCHITECTURE
Revised: 3. M ULTI -S CALE A RCHITECTURE

Original: In many long sequence problems, it is standard practice to sub-sample or stride the input sequence before being fed to the main model for reducing computation.
Revised: In many long sequence problems, it is standard practice to sub-sample or stride the input sequence before being fed to the main model for reducing computation.

Original: For example, both Listen Attend and
Revised: For example, both Listen Attend and

Original: Spell (Chan et al., 2016) and transducer (Graves, 2012) based models for speech recognition utilize this type of sub-sampling to achieve lower frame rates and shorter sequences.
Revised: Spell (Chan et al., 2016) and transducer (Graves, 2012) based models for speech recognition utilize this type of sub-sampling to achieve lower frame rates and shorter sequences.

Original: In this work we investigate the multi-scale architecture, using MSSM layers prior to the striding, for sub-sampling.
Revised: In this work we investigate the multi-scale architecture, using MSSM layers prior to the striding, for sub-sampling.

Original: There are two advantages of this multi-scale design: (1) MSSMs scale linearly and can potentially handle longer sequences; (2) processing the sequence prior to sub-sampling can introduce better contextual information.
Revised: There are two advantages of this multi-scale design: (1) MSSMs scale linearly and can potentially handle longer sequences; (2) processing the sequence prior to sub-sampling can introduce better contextual information.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We investigate our proposed models on large-scale automatic speech recognition (ASR) and masked language modeling (MLM) tasks, and compare with a range of transformer-based schemes on each task.
Revised: We investigate our proposed models on large-scale automatic speech recognition (ASR) and masked language modeling (MLM) tasks, and compare with a range of transformer-based schemes on each task.

Original: Baseline and proposed models are implemented using an extension of the Fairseq framework (Ott et al., 2019).
Revised: Baseline and proposed models are implemented using an extension of the Fairseq framework (Ott et al., 2019).

Original: 4. A UTOMATIC S PEECH R ECOGNITION
Revised: 4. A UTOMATIC S PEECH R ECOGNITION

Original: We first evaluate the proposed models on automatic speech recognition using the widely known
Revised: We first evaluate the proposed models on automatic speech recognition using the widely known

Original: LibriSpeech corpus (Panayotov et al., 2015), consisting of about 960 hours of speech data sampled at 16kHz.
Revised: LibriSpeech corpus (Panayotov et al., 2015), consisting of about 960 hours of speech data sampled at 16kHz.

Original: All speech recognition models were built in the transducer framework (Graves, 2012), which has three components: encoder, predictor and joiner.
Revised: All speech recognition models were built in the transducer framework (Graves, 2012), which has three components: encoder, predictor and joiner.

Original: We train various transducers by keeping the predictor and joiner fixed, and compare attention-based and S4 (Gu et al., 2022) baselines, and the proposed models on the encoder.
Revised: We train various transducers by keeping the predictor and joiner fixed, and compare attention-based and S4 (Gu et al., 2022) baselines, and the proposed models on the encoder.

Original: In these experiments, we used a sliding window of 25ms with a 10ms frame shift to extract 80dimension log Mel-filter bank features which were immediately linearly projected to a dimension of 128.
Revised: In these experiments, we used a sliding window of 25ms with a 10ms frame shift to extract 80dimension log Mel-filter bank features which were immediately linearly projected to a dimension of 128.

Original: Depending on the model, we used one of the following sub-samplers (SS) in order to reduce the frame rate to 40ms and increase the dimensionality to 512:
Revised: Depending on the model, we used one of the following sub-samplers (SS) in order to reduce the frame rate to 40ms and increase the dimensionality to 512:

Original: • Conv : 1D-convolution network with an effective stride of 4 and 512 output channels (Wang et al., 2020b; Gulati et al., 2020).
Revised: • Conv : 1D-convolution network with an effective stride of 4 and 512 output channels (Wang et al., 2020b; Gulati et al., 2020).

Original: • TR : Time-reduction layer which stacks each 4 consecutive frames (stride = 4).
Revised: • TR : Time-reduction layer which stacks each 4 consecutive frames (stride = 4).

Original: • MS : Multi-scale front introduces MSSM layers prior to time-reductions (Section 3.6).MSSM blocks are passed on to TR(stride = 2) followed by 4 MSSM blocks and a final TR(stride = 2).
Revised: • MS : Multi-scale front introduces MSSM layers prior to time-reductions (Section 3.6).MSSM blocks are passed on to TR(stride = 2) followed by 4 MSSM blocks and a final TR(stride = 2).

Original: The attention-based baselines used the convolutional sub-sampler and consist of 20-36 Transformer or Conformer (Gulati et al., 2020)
Revised: The attention-based baselines used the convolutional sub-sampler and consist of 20-36 Transformer or Conformer (Gulati et al., 2020)

Original: blocks with 512-dimensional features, 8 attention heads, and feed-forward nets with a dimension expansion of 4.
Revised: blocks with 512-dimensional features, 8 attention heads, and feed-forward nets with a dimension expansion of 4.

Original: Different to Gulati et al.
Revised: Different to Gulati et al.

Original: (2020), our Conformer baseline did not use a macaron style block as it was not found useful, and had the convolutional module prior to the attenion with a kernel size of 31.
Revised: (2020), our Conformer baseline did not use a macaron style block as it was not found useful, and had the convolutional module prior to the attenion with a kernel size of 31.

Original: Also, an S4 model was prepared, for it is a single-head SSM baseline with careful parameter initialization; the scaled Legendre (LegS) initialization is used.
Revised: Also, an S4 model was prepared, for it is a single-head SSM baseline with careful parameter initialization; the scaled Legendre (LegS) initialization is used.

Original: The S4 baseline and our proposed models simply replace the attention block of the transformer layer, and the configuration of this layer, such as number of stacks, heads, use of inter-head gating and sub-sampler design, is investigated; D sh = 32 is fixed for all experiments.
Revised: The S4 baseline and our proposed models simply replace the attention block of the transformer layer, and the configuration of this layer, such as number of stacks, heads, use of inter-head gating and sub-sampler design, is investigated; D sh = 32 is fixed for all experiments.

Original: The Stateformer uses the same parameters as the transformer combined with the best MSSM configuration.
Revised: The Stateformer uses the same parameters as the transformer combined with the best MSSM configuration.

Original: The predictor will remain fixed for all experiments and contain three 512-dimensional
Revised: The predictor will remain fixed for all experiments and contain three 512-dimensional

Original: LSTM layers with layer-norm and dropout.
Revised: LSTM layers with layer-norm and dropout.

Original: The outputs of both encoder and predictor are projected to 1024 dimensions and fed into an additive joiner with a single linear layer of |Y| =sentence-piece (Kudo & Richardson, 2018) output units.
Revised: The outputs of both encoder and predictor are projected to 1024 dimensions and fed into an additive joiner with a single linear layer of |Y| =sentence-piece (Kudo & Richardson, 2018) output units.

Original: All models used the Adam optimizer (Kingma & Ba, 2015) with a learning rate linearly warming up to the peak value in 8000 iterations, fixed until the 60th epoch and thereafter, exponentially decayed by a factor of 0.96 each epoch.
Revised: All models used the Adam optimizer (Kingma & Ba, 2015) with a learning rate linearly warming up to the peak value in 8000 iterations, fixed until the 60th epoch and thereafter, exponentially decayed by a factor of 0.96 each epoch.

Original: A dropout value of 0.10 is used in all encoders and 0.
Revised: A dropout value of 0.10 is used in all encoders and 0.

Original: in all predictors and the batch size is set based on occupying maximal GPU memory.
Revised: in all predictors and the batch size is set based on occupying maximal GPU memory.

Original: All models were trained up to 200 epochs using 32 NVIDIA A100 GPUs.
Revised: All models were trained up to 200 epochs using 32 NVIDIA A100 GPUs.

Original: In evaluation, decoding results are generated using the beam search algorithm on the transducers without any external language model.
Revised: In evaluation, decoding results are generated using the beam search algorithm on the transducers without any external language model.

Original: Best checkpoints for test-sets are selected on the respective dev-sets.
Revised: Best checkpoints for test-sets are selected on the respective dev-sets.

Original: [Table]
Revised: [Table]

Original: First, the impact of SSM stack depth is examined.
Revised: First, the impact of SSM stack depth is examined.

Original: As reported in Table 1, the double stacked 16-layer configuration (illustrated in Figure 1b) achieves similar performance to the 20-layer nonstacking model.
Revised: As reported in Table 1, the double stacked 16-layer configuration (illustrated in Figure 1b) achieves similar performance to the 20-layer nonstacking model.

Original: This demonstrates that the double stacked design is more effective and will be the default topology used in subsequent experiments
Revised: This demonstrates that the double stacked design is more effective and will be the default topology used in subsequent experiments

Original: [Table]
Revised: [Table]

Original: Second, in Table
Revised: Second [Table] in Table

Original: [Table]
Revised: [Table]

Original: we compare the impact of multi-head, inter-head gating and sub-sampler designs on the 16-layer MSSM models.
Revised: we compare the impact of multi-head, inter-head gating and sub-sampler designs on the 16-layer MSSM models.

Original: Although directly increasing the number the number of heads from 1 to 2 shows limited improvement, by turning on inter-head gating one can observe consistent gains on multi-head models.
Revised: Although directly increasing the number the number of heads from 1 to 2 shows limited improvement, by turning on inter-head gating one can observe consistent gains on multi-head models.

Original: Furthermore, additional improvement can be achieved when increasing the number of heads from 2 to 4, keeping IHG activated.
Revised: Furthermore, additional improvement can be achieved when increasing the number of heads from 2 to 8, keeping IHG activated.

Original: Continuing, we replace the time-reduction layer with a multi-scale sub-sampler which notably improves performance.
Revised: Continuing, we replace the time-reduction layer with a multi-scale sub-sampler which notably improves performance.

Original: Note the second row is a transformer model with a simple time-reduction (instead of a convolutional) sub-sampler, which can be seen as a fair baseline for MSSMs without the multi-scale option.
Revised: Note the second row is a transformer model with a simple time-reduction (instead of a convolutional) sub-sampler, which can be seen as a fair baseline for MSSMs without the multi-scale option.

Original: In this case, the proposed models significantly outperform the respective transformer baselines of similar size.
Revised: In this case, the proposed models significantly outperform the respective transformer baselines of similar size.

Original: Next we couple attention with the best MSSM configuration to form the Stateformer (described in Section 3.5).
Revised: Next we couple attention with the best MSSM configuration to form the Stateformer (described in Section 3.5).

Original: The final row of Table 2 indicates some degree of complementarity, that further gains can be achieved when combining both of these modeling techniques.
Revised: The final row of Table 2 indicates some degree of complementarity, that further gains can be achieved when combining both of these modeling techniques.

Original: [Table]
Revised: [Table]

Original: We also scale the model size up to 140M parameters to verify the performance of a large configuration.
Revised: We also scale the model size up to 140M parameters to verify the performance of a large configuration.

Original: Baselines and the proposed models in this large configuration are trained using auxiliary classifiers similar to Szegedy et al.
Revised: Baselines and the proposed models in this large configuration are trained using auxiliary classifiers similar to Szegedy et al.

Original: (2015), in which intermediate encoder outputs are trained to predict frame labels every 4 layers.
Revised: (2015), in which intermediate encoder outputs are trained to predict frame labels every 4 layers.

Original: Table 3 compares the proposed models with large-size baselines and state-of-the-art numbers in previous works 1 .
Revised: Table 3 compares the proposed models with large-size baselines and state-of-the-art numbers in previous works 1 .

Original: Both the 32-layer MSSM and 25-layer Stateformer outperform our large-size baselines and a range of attention-based results in literature.
Revised: Both the 32-layer MSSM and 25-layer Stateformer outperform our large-size baselines and a range of attention-based results in literature.

Original: 4. M ASKED L ANGUAGE M ODELING
Revised: 4. M ASKED L ANGUAGE M ODELING

Original: To explore our model’s effectiveness on other modalities, we consider the widely used masked language modeling task on long sequences (Devlin et al., 2018).
Revised: To explore our model’s effectiveness on other modalities, we consider the widely used masked language modeling task on long sequences (Devlin et al., 2018).

Original: We follow Beltagy et al.
Revised: We follow Beltagy et al.

Original: (2020) in preparing a large-scale long-document corpus, including Stories (Trinh & Le, 2018), RealNews (Zellers et al., 2019), Books (Zhu et al., 2015) and English Wikipedia.
Revised: (2020) in preparing a large-scale long-document corpus, including Stories (Trinh & Le, 2018), RealNews (Zellers et al., 2019), Books (Zhu et al., 2015) and English Wikipedia.

Original: This corpus, consisting of 6.6 billion tokens, is loaded as in Beltagy et al.
Revised: This corpus, consisting of 6.6 billion tokens, is loaded as in Beltagy et al.

Original: (2020) in terms of 4,096-token long sequences.
Revised: (2020) in terms of 4,096-token long sequences.

Original: It is prohibitive to apply the full attention in transformers to the 4096-token sequences due to the quadratic memory complexities.
Revised: It is prohibitive to apply the full attention in transformers to the 4096-token sequences due to the quadratic memory complexities.

Original: Therefore, we pick a few feasible transformer variants as baselines, including sliding-window attention (Beltagy et al., 2020), Performer (Choromanski et al., 2020) and
Revised: Therefore, we pick a few feasible transformer variants as baselines, including sliding-window attention (Beltagy et al., 2020), Performer (Choromanski et al., 2020) and

Original: Linformer (Wang et al., 2020a).
Revised: Linformer (Wang et al., 2020a).

Original: The sliding window has a window size of 256; Performer uses a random feature dimension of 256 and ReLU as kernel function; Linformer has a compression ratio of 8.
Revised: The sliding window has a window size of 256; Performer uses a random feature dimension of 256 and ReLU as kernel function; Linformer has a compression ratio of 8.

Original: The MSSM model with D sh = 32 has 8 heads; remaining hyperparameters are shared with attention baselines.
Revised: The MSSM model with D sh = 32 has 8 heads; remaining hyperparameters are shared with attention baselines.

Original: To validate that our model works well at scale, we test the attention and MSSM layers using a 24layer model similar in size to RoBERTa-large (Liu et al., 2019).
Revised: To validate that our model works well at scale, we test the attention and MSSM layers using a 24layer model similar in size to RoBERTa-large (Liu et al., 2019).

Original: To make a fair comparison, training stops after the same number of data batches.
Revised: To make a fair comparison, training stops after the same number of data batches.

Original: Following Xiong et al.
Revised: Following Xiong et al.

Original: (2021), baselines are trained using the standard masked language modeling objective for 100K steps with a per GPU batch size of 8 on 32 A100 GPUs 2 .
Revised: (2021), baselines are trained using the standard masked language modeling objective for 100K steps with a per GPU batch size of 8 on 32 A100 GPUs 2 .

Original: The MSSM is trained for 50K steps with a per GPU batch size of 8 on
Revised: The MSSM is trained for 50K steps with a per GPU batch size of 8 on

Original: GPUs, for rapid development.
Revised: GPUs, for rapid development.

Original: Perplexities (PPLs) are reported on the validation set.
Revised: Perplexities (PPLs) are reported on the validation set.

Original: Table 4 summarizes the perplexity of different models.
Revised: Table 4 summarizes the perplexity of different models.

Original: The proposed MSSM achieves performance on-par with attention baselines, but shows notably improved perplexity with further training (20K additional updates).
Revised: The proposed MSSM achieves performance on-par with attention baselines, but shows notably improved perplexity with further training (20K additional updates).

Original: Although the MSSM is expected to learn information from the full sequence, the performance gap to the sliding window is smaller than expected.
Revised: Although the MSSM is expected to learn information from the full sequence, the performance gap to the sliding window is smaller than expected.

Original: To better understand the behavior of the proposed model, we analyze the learned state space kernels in Section 4.3.
Revised: To better understand the behavior of the proposed model, we analyze the learned state space kernels in Section 4.3.

Original: [Table]
Revised: [Table]

Original: 4. M ULTI - HEAD I NTRODUCES M ULTI - MODAL AND M ULTI - SCALE M ODELING P OWER
Revised: 4. M ULTI - HEAD I NTRODUCES M ULTI - MODAL AND M ULTI - SCALE M ODELING P OWER

Original: Now that we have empirically demonstrate the effectiveness of multi-head state space models
Revised: Now that we have empirically demonstrate the effectiveness of multi-head state space models

Original: [Figure]
Revised: [Figure]

Original: we want to answer the following research question: what does the multi heads learn from data?
Revised: we want to answer the following research question: what does the multi heads learn from data?

Original: Since the state space model can be presented as a convolution (Equation 4)
Revised: Since the state space model can be presented as a convolution (Equation 4)

Original: [Figure]
Revised: [Figure]

Original: the shape of the kernel filters can be studied to answer this question.
Revised: the shape of the kernel filters can be studied to answer this question.

Original: [Figure]
Revised: [Figure]

Original: We first look into the learned multi-head kernels from ASR model (Figure 4B)
Revised: We first look into the learned multi-head kernels from ASR model (Figure 4B)

Original: [Figure]
Revised: [Figure]

Original: and compare them with the learned single-head kernel (Figure 4A).
Revised: and compare them with the learned single-head kernel (Figure 4A).

Original: We find the the single head kernel learned a simple monotonic decaying pattern, which computes the weighted average of history in a sliding window.
Revised: We find the the single head kernel learned a simple monotonic decaying pattern, which computes the weighted average of history in a sliding window.

Original: On the other hand, the multi-head kernels are much more interesting — while the first head is still monotonic decaying, the kernels from the second and third heads involve context with different context concentration.
Revised: On the other hand, the multi-head kernels are much more interesting — while the first head is still monotonic decaying, the kernels from the second and third heads involve context with different context concentration.

Original: In other words, multi-head state space models are able to introduce multimodal modeling power.
Revised: In other words, multi-head state space models are able to introduce multimodal modeling power.

Original: The kernels from the MLM tasks bring additional insights.
Revised: The kernels from the MLM tasks bring additional insights.

Original: To begin with, kernels from the first three heads shows the same pattern as the ASR kernels (Figure 4C), reaffirming the multi-modal modeling power on the MLM tasks.
Revised: To begin with, kernels from the first three heads shows the same pattern as the ASR kernels (Figure 4C), reaffirming the multi-modal modeling power on the MLM tasks.

Original: Also, this localized behavior can be viewed similarly to that of the sliding-window attention.
Revised: Also, this localized behavior can be viewed similarly to that of the sliding-window attention.

Original: Additionally, kernels from the next four heads (Figure 4D) are monotonic with different decay rates, which captures history with multi-scale granularity.
Revised: Additionally, kernels from the next four heads (Figure 4D) are monotonic with different decay rates, which captures history with multi-scale granularity.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We propose the multi-head state space model (MSSM) for sequence modeling in this paper.
Revised: We propose the multi-head state space model (MSSM) for sequence modeling in this paper.

Original: Parallel heads of SSMs, together with inter-head gating, are introduced, enabling multiple temporal dynamics to be effectively learned on sequence data.
Revised: Parallel heads of SSMs, together with inter-head gating, are introduced, enabling multiple temporal dynamics to be effectively learned on sequence data.

Original: In addition, we combine the MSSM and attention models to form a new type of transformer architecture, referred to as Stateformer.
Revised: In addition, we combine the MSSM and attention models to form a new type of transformer architecture, referred to as Stateformer.

Original: On the
Revised: On the

Original: LibriSpeech speech recognition task, the proposed MSSM outperforms transformer baseline by a large margin.
Revised: LibriSpeech speech recognition task, the proposed MSSM outperforms transformer baseline by a large margin.

Original: Besides, the Stateformer further improves the performance, achieving state-of-theart performance.
Revised: Besides, the Stateformer further improves the performance, achieving state-of-theart performance.

Original: On the masked language modeling task, the MSSM achieves on-par performance compared with a range of attention-based baselines trained on a large corpus.
Revised: On the masked language modeling task, the MSSM achieves on-par performance compared with a range of attention-based baselines trained on a large corpus.