Multi-Head State Space Model for Sequence Modeling
Authors: Yassir Fathullah
Chunyang Wu
Yuan Shangguan
Junteng Jia
Wenhan Xiong
Jay Mahadeokar
Chunxi Liu
Yangyang Shi
Ozlem Kalinli
Mike Seltzer
Mark Gales
Venue: Submitted to ICLR 2023 (2023)
Keywords: multi-head
state space
transformer
stateformer
sequence model
rnn

Abstract: Recently, state space models (SSMs) have shown promising results on sequence modeling tasks. 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. 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. Furthermore, we propose a novel variant of the Transformer, referred to as the Stateformer, which combines MSSMs with attention. Experiments on large-scale automatic speech recognition (ASR) and language modeling tasks show the MSSM outperforming a range of attention-based baselines. The Stateformer further improves performance, achieving the state-of-the-art performance on the LibriSpeech ASR task.

Under review as a conference paper at ICLR 2023
MULTI-HEAD STATE SPACE MODEL FOR SEQUENCE MODELING
Anonymous authors Paper under double-blind review
ABSTRACT
Recently, state space models (SSMs) have shown promising results on sequence modeling tasks. 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. In this paper, we propose multi-head state space model (MSSM), in which parallel heads are in- troduced to learn different temporal dynamics on sequence data. Furthermore, we propose novel variant of the Transformer, referred as the Stateformer, which combines MSSMs with attention. Experiments on large-scale automatic speech recognition (ASR) and language modeling tasks show the MSSM outperforming a range of attention-based baselines. The Stateformer further improves perfor- mance, achieving the state-of-the-arr performance the LibriSpeech ASR task.
INTRODUCTION
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 Liu al. b021b) and biological applications such as protein sequence modeling (Jumper et al, 2021 Rives et al. 2021). 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). This is unlike previous state-of-the recurrent and convolutional neural networks which have linearly and logarithmically scaling path lengths respectively. 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). 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 much of the original transformer performance as possible (Katharopoulos et al., 2020: Kitaev et al. 2020 Beltagy et al., 2020 Zaheer al. 2020). 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 al. 2020 Shen al. 2021 Zhang al. 2019).
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 his- torically has been widely used in many time-series and control problems (Brogan 1991 Hyndman al. 2008 Durbin Koopman 2012). 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 Rangaauram et al. 2018). Furthermore, a simplified linear time-invariant (LTT) SSM, which closely related to fully linear recurrent neural network (RNN) layers, has 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).
One potential limitation of existing works in this direction is that SSMs are usually initialized and used in a homogeneous way. For instance, in Gu al. (2022a). the S4 approach equips SSMs with careful options on parameter initialization for long-range modeling ability. This design would force the model to capture similar temporal dynamics on different features. As long and short term

---

Under review as a conference paper at ICLR 2023
dependency would both be useful in sequence modeling, in this paper. develop multi ieati-headdee state space model, which consists of parallel heads to learn different temporal dynamics on sequence data. We investigate the use of multi-head state space models in sequence tasks both replacement and complement to attention. Since SSMs have been shown to be promising on long sequences, we hypothesize such model is able handle both short and long-term dependencies and can operate as an effective alternative to attention. Following are our main technical contributions around the state space model:
1. Stacked and multi-head generalization. We extend the SSM approach by allowing multi- head parallel processing of projected lower-dimensional signals and stack such a layer (within a residual block) for better performance.
2. Head gating. We propose an inter-head gating approach allowing different SSMs within the multi-head layer to communicate.
3. Combination with attention. We also augment the transformer architecture by simply in- cluding bidirectional SSM residual block prior the attention for computationally unre- stricted applications and state-of-the the-art performance; referred to as the Stateformer.
With these contributions, we advance the state of attention-fres models on large-scale speech recog- nition and language modeling tasks, and show that our novel multi-head state space model outper- forms strong attention- based baselines. We also show that combining the multi-head state space with attention can achieve state-of-the-art performance in large-scale speech recognition.
2
BACKGROUND
This section will cover theory and related work regarding the state space model and its use in deep learning. We first start by briefly covering the state space model and block diagonal structured restriction. We further mention a convolution point of view for parallel training (Section 2.1. We finish with brief discussion of the state space layer as used in neural nets (Section 2.2
2.1 STATE SPACE MODEL: THE LINEAR RNN
A specific realization of the state space model (SSM) is linear time-invariant (LTT) model (Brogan 1991 that transforms some input signal ul)) E D. to some output y)() E RCD, through some hidden process a+) E RDD according to
(t() Aa(t) + Bu(t),
y()) = Car(1) + Du(()
where Du(t) can be neglected since this is simple term compute. To ensure such model is compatible with discrete signals such as sequences of speech frames or sub-word unit embeddings, simple discretization, e.g. the bilinear method (Tustin 1947) can be utilized:
2t Axx-1 + Bug, A [1 AA/2) A/2)),
(2)
Cac::
B B= (I- AB, -
which simply represents a fully linear recurrent neural network (RNN) cell. Note that the discretiza- tion matrix E Dr is diagonal with positive entries, but can subsumed by A, B due to scale invariance. Furthermore, while such model is highly flexible, providing it with structure can increase its efficiency. Therefore, one can restrict the matrices A, B, C' to have a block diagonal structure as follows:
A
0
0
C
0
A
0
0
B
0
C2
0
4
B
..
C
(3)
An.
BD./
Cn.
where the structure of this block recurrent model is based on the number of input dimensions D The original multi-varial state space model can now be seen as D; smaller independent sub-SSMs, where each one of these smaller models + eC mappings running in parallel.
A block diagonal treatment of structured state spaces is also covered in Smith al. 2022
2

---

Under review as a conference paper at ICLR 2023
Denoting the sub-hidden dimension by Dah Dr/D; obtain A, E R D.r B, E PAxX and C, E RCx The skip-connection can also be split into blocks of Dr E RCx Similarly, the discretization matrix can now be seen diagonal matrix with isotropic blocks A., = A of the same size as the transition blocks A., meaning that each sub-state space model has its own discretization time-step that can be tuned its own input channel.
Furthermore, due to the model being fully linear, one can utilize the widely known duality between convolutions and recurrences, resulting in efficient expression:
CA'B). k= (CB, CA'B y y=
where L is the input sequence length. 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 recur- rent mode (Equation 2), efficiently and incrementally generating predictions. To further improve the efficiency, a specific diagonal plus low-rank restriction of the sub-transition blocks can be used Gu et al. (022a),
A. An + Pu Pa. qn
Dab Vn (..
(5)
where A, is diagonal matrix. This allows the kernel K to be computed efficiently utilizing the Woodbury matrix identity with several fast implementations (Pan 2001 2014).
2.2 PRACTICAL LAYER DESIGN
The structured state space model described in Equation 2 is general linear R.D. - RCD, mapping but processes each channel independently due to the structuring in Equation 3 To amend both these
problems, non-linear activation is introduced followed by point-wise linear layer (see Figure la for a pre-norm residual unit) to mix the information across input channels:
f(u) = Linear(Act(SSMMuu)
(6)
Note the output dimension expansion C' is set to 1 when using element-wise activations such as GELU (Hendrycks & Gimpel 2016) or set to when using gating mechanisms such as GLU Dauphin al. 2017).
Drop. & Add
Act Linear
Drop. & Add
SSM Layer
Act & Linear
Act Linear
SSM Layer
SSM Layer
LayerNorm
LayerNorm
(a) Basic block used Goel et (2022).
(b) Stacked proposal.
Figure 1: The standard block used in previous work and an example of double stacked block. Each block is followed by feed-forward net similar to transformers.
To further stabilize the state space, Goel et al. (2022) used the more stable Hurwitz formulation An A, P.Pp', ensuring the eigenvalues of the transition matrix remain negative when A,, is negative definite. Practically, the parameters are stored in log-form: time-step A exp(*a.) and diagonal A,, = exp(*A,)) Regarding the initialization of this system, Gu et al. 2020) developed several options for how the transition and input matrices should be set equip the model with long-range modeling ability which can be motivated from a memorization signals point of view. However, all of the parameters were made trainable, discarding the initialization.
3

---

Under review as a conference paper at ICLR 2023
3 MULTI-HEAD STATE SPACE MODEL
This section introduces the novel multi-head state space model as a natural extension to the linear time-invariant structured state space model described above. First, we discuss configuration for the state space layer and its intialization (Section (.1). Second, we propose novel ideas for how the architecture be improved, including multi- head generalization (Section 3.2) Third, new
gating mechanism is introduced which was found notably beneficial (Section 3.3 Furthermore,
bidirectional architecture is discussed (Section 3.4) and how it can be used with attention based models (Section 3.5). Finally, the compute complexity of state space models scales only linearly with sequence length allowing it to be used prior sub-sampling layers without much added cost (Section 3.6..
3.1 TIED PARAMETERS & INITIALIZATION
The original design in Equation 3 describes 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. 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. In this work we opt for a slightly more parameter efficient blockwise isotropic implementation, similar to to et al. (2022b), where we restrict the copies to have shared input and transition matrices across all blocks:
An An P.P' = Ao = Ao PoPo
Vn E (1,..,D),
(7)
B Bo
but keep independent output Cn and skip Dn matrices. We also let the learning rate remain the same as the remainder of the network. By still having different time-step discretizations for each dimension of the imput sequence, temporal features of different scale can be extracted while sharing the recurrence parameters. Furthermore, we choose to simply random initialize all parameters of the SSM by sampling from the unit Gaussian distribution (O, 1), apart from the diagonal matrix 1). This is to ensure that the initial eigen-
which is shifted Ao values of the transition matrix are not too small and can capture enough of the input signal''s 's history. Compared with Gu et al. (2020). the random initialization is simple but shows good performance in practice. The effectiveness of random initialization was also observed in Mehta et al. (2022). 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.
exp( AAo): AAa ~ N
3.2 STACKED MULTI-HEAD GENERALIZATION
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). Therefore, the design is expanded by stacking y an example in Figure 1b This would balance the parameters in terms of the ratio between FFN and SSM layers and give the whole model more depthwise/temporal flexibility.
fi))((()))) each with its own set of parameters,
Moreover, we found it useful to project the input into several different lower-dimensional representa- tions which are passed head attention, see Figure 2a Not only does having shared recurrence parameters across a structured SSM stabilize training, but gives each structured model in multi-head framework the flexibil- ity to learn both meaningful time-steps and different types of temporal dynamics. This stacked multi-head state space model (MSSM) requires an additional initial projection layer compared its single-head implementation, leading to a small increase in parameters. In practice, the projection would take a sequence of D,-dimensional features and project down to E {2, 8,... differ- ent, D; = D&/ H dimmensiion signals, each of which is processed by isotropic blockwise diagonal SSMs. Since, the effective dimension D; = D; remains the same, the cost of a multi-headed approach boils down to an extra linear projection layer.
independent structured state space models, in design similar multi-
3.3 INTER-HEAD GATED LINEAR UNIT
By default, prior work has used the GELU activation function. some experiments (Gu et al, 2022al Goel et al.] 2022 Mehta al. 2022 the GLU activation (with C 2) was also found
4

---

Under review as a conference paper at ICLR 2023
Add
Project
SSM Layers
Cat & Project
Project
SSM Layers
Sigmoid
Project
SSM Layer
SSM Layer
LayerNorm
Project
(a) Stacked and multi-head block.
(b) Inter-head gating a multi-head block.
Figure 2: Multi-hcad generalization and novel inter-head gating idca. Instead of gating the two- dimensional output of a single SSM, this approach is based elementwise gating of head using the corresponding output of different head.
beneficial which was achieved by gating the channels of an individual sub-state space model:
Ug (Mma), IR? Vn E (1,...,
However, multi-head state space architecture with H headed outputs {ym E R many more gating possibilities. We propose an inter-head gating (IHG) approach in which the output is computed by mixing the heads according to (where a refers to the sigmoid):
Tha== offers
(l) yne (yylhtH/22) Vc = (1,2)
(1,...,,/2)
allowing different heads to communicate and gate each other, generally leading to improved results. Clearly, flattening out the output Equation Ileads to the original D dddimenniional output, and note that the number of heads in this approach has to be even, see Figure 2b While our approach is wildly different, the concept of letting heads communicate has also been proposed in Shazeer et al. (2200) in the context of attention where linear layers mix information across heads, and was found to improve performance. However, unlike the fast-mentioned, our approach allows for communication between different heads without any extra parameters.
3.4 BIDIRECTIONAL DESIGN
So far the proposed architecture has a unidirectional design however, for better performance, bidi- rectional context is needed. Therefore, simply concatenate the output of independent forward and backward stacked multi-head modules and project it back down to the original input dimension (Cornegruta et 2016 Goel al. 2022). This new module is then wrapped in pre-norm block similar to the architectures above:
y + Cat (MMMM(M)((). Rev(MSSM(Rev(22))
(10)
+- Linear(Act(y))
Since there is linear layer the bidirectional block, the final linear layer of the MSSM modules removed. This attention free design has many advantages over multi-head attention such as memory requirements scaling linearly with sequence length, making computationally feasible for longer sequences.
3.5 STATEFORMER: STATE SPACE AUGMENTED TRANSFORMER
A pure multi-head state space model architecture is attractive due to its efficiency and ability to capture both short and long range dependencies. However, since the state space model is equivalent to a linear RNN, it is expressively more limited, and still plagued by linearly scaling path length. Therefore, for scenarios where is less restricted by compute limitations, 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
propose model

---

Under review as a conference paper at ICLR 2023
Layer Norm
Bidirectional MSSM
Dropout & Add
Layer Norm
Multi-Head Attention
Dropout & Add
Layer Norm
Feed Forward Net
Dropout & Add
Figure 3: Stateformer: Bidirectional state space augmented transformer encoder example. It simply has an additional block prior to the attention unit with bidirectional MSSM as seen Equation
This type of design is inspired by the Conformer (Gulati et al. 2020 where dedicated convolu- tional block inserted to handle local relations. However, we hypothesize that the advantage of the Stateformer, is the MSSM with its variable length kernel and the multi-head attention both han- dle short and long-term dependencies and can complement each other. Furthermore, the attention enables maximum path lengths of I, resulting in a possibly more flexible model.
3.6 MULTI-SCALE ARCHITECTURE
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. For example, both Listen Attend and Spell (Chan 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. In this work we investigate the multi-scale architecture, using MSSM layers prior to the striding, for sub- sub-samp sampling. There are two advantages of this multi-scale design: (1) MSSMs scale linearly and can potentially handle longer sequences; (2) processing the sequence prior sub-sampling can introduce better contextual information.
4 EXPERIMENTS
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. Baseline and proposed models are implemented using an extension of the Fairseq framework (Ott et al. 2019).
4.1
AUTOMATIC SPEECH RECOGNITION
We first evaluate the proposed models on automatic speech recognition using the widely known LibriSpeech corpus (Panayotov et al. 2015), consisting of about 960 hours of speech data sampled at 16kHz. All speech recognition models were built in the transducer framework (Graves) 2012), which has three components: encoder, predictor and joiner. We train various transducers by keeping the predictor and joiner fixed, and compare attention and S4 (Gu et al. 2022a) baselines, and the proposed models the encoder.
In these experiments, we used sliding window of 25ms with a 10ms frame shift to extract 80- dimension log Mel-filter bank features which were immediately linearly projected dimension of 128. Depending on the model, we used one of the following sub-samplers (SS) in order to reduce the frame rate 40ms and increase the dimensionality to 512:
Conv: 2D-convolution network with effective stride of and 512 output channels (Wang et al. 2020b] Gulati et al. 2020).
TR: Time-reduction layer which stacks each 4 consecutive frames (stride = 4).
MS: Multi-scale front introduces MSSM layers prior time-reductions (Section 3.6 MSSM blocks are passed on to TR(stride = 2) followed by MSSM blocks and a final TR(stride 2).
The attention-based baselines used the convolutional sub-sampler and consist of 20-36 Transformer or Conformer (Gulati et al., 2020) blocks with 312-dimensional features, 8 attention heads, and fordafo nets with dimension expansion of 4. Different to Gulati et al. (2020), baseline did not use macaron style block as it was not found useful, and had the convolutional module prior to the attenion with kerne size of 31. Also, S4 model prepared, for it is
Conformer

---

Under review as a conference paper at ICLR 2023
single-head SSM baseline with careful parameter initialization; the scaled Legendre (LegS) ini- tialization is used. The S4 baseline and our proposed models simply replace the attention block of the transformer layer, and the configuration of this layer, such number of stacks, heads, of inter-head gating and sub-sampler design, investigated; saA = 32 is fixed for all experiments. The Stateformer uses the same parameters the transformer combined with the best MSSM con- figuration. The predictor will remain fixed for all experiments and contain three 512-dimensional LSTM layers with layer-norm and dropout. The outputs of both encoder and predictor are pro- jected to 1024 dimensions and fed into an additive joiner with single linear layer of [ 4097 sentence-piece (Kudo& & Richardson 2018) output units.
All models used the Adam optimizer (Kingma & Ba 2015) with a learning rate linearly warming up to the peak value 8000 iterations, fixed until 60th epoch and thereafter, exponentially decayed by a factor of 0.96 each epoch. A dropout value of 0.10 is used in all encoders and 0.30 all predictors and the batch size is set based on occupying maximal GPU memory. All models were trained up to 200 epochs using 32 NVIDIA A100 GPUs. In evaluation, decoding results generated using the beam search algorithm on the transducers without any external language model. Best checkpoints for test-sets are selected the respective dev-sets.
WER% (L)
Layers
Params
Stack
develean
devother
testelean
testother
56.8M
2.57
7.13
2.79
6.86
16
66.3M
2
2.36
6.88
2.52
6.59
20
67.6M
2.42
6.92
2.67
6.56
Table 1: Impact of stack depth of state space model in one layer on Librispeech.
First, the impact of SSM stack depth is examined. As reported in Table [ the double stacked 16-layer configuration (illustrated in Figure Tb) achieves similar performance to the 20-layer non- stacking model. This demonstrates that the double stacked design more effective and will be the default topology used in subsequent experiments.
WER% (4) testelean
Model
Params
IHG
develean
devother
testother
77.5M
Conv TR
2.45
5.82
2.62
6.15
Transformer 20L
76.7M
2.96
7.09
3.05
7.18
Conformer 20I.
92.7M
Conv
2.17
5.54
2.43
5.45
78.8M
Conv
2.36
6.60
2.67
6.47
S4-L.cgS 20I.
78.1M
TR
2.45
6.88
2.71
6.72
66.3M
X
TR
2.36
6,88
2.52
6.59
74.7M
X
TR
2.33
6.92
2.58
6.49
MSSM 16T.
74.7M
TR
2.13
6.81
2.47
6.44
74.7M
2.19
6.38
2.42
6.25
74.7M
TR
2.17
6.49
2.43
6.20
79.2M
MS
2.14
6.12
2.39
5.99
Stateformer 16L.
96.1M
MS
2.06
5.01
2.27
5.07
Table 2: Performance of base configuration Librispeech. #H is the number of heads; IHG is inter-head gating (Section 3.3 SS stands for the sub-sampling scheme used the bottom layer.
Second, in Table 2 we compare the impact of multi-head, inter-head gating and sub- sub-sampler designs on the 16-layer MSSM models. Although directly increasing the number the number of heads from I to 2 shows limited improvement, by turning inter-head gating one can observe consistent gains on multi-head models. Furthermore, additional improvement can be achieved when increasing the number of heads from 2 to 8, keeping IHG activated. Continuing, we replace the time-reduction layer with a multi-scale sub-sampler which notably improves performance. Note the second row is transformer model with simple time-reduction (instead of convolutional) sub-sampler, which can be seen as fair baseline for MSSMs without the multi-scale option. In this case, the proposed models significantly outperform the respective transformer baselines of similar size. Next couple

---

Under review as a conference paper at ICLR 2023
attention with the best MSSM configuration to form the Stateformer (described in Section (.5)) The final row Table 2 indicates some degree complementarity, that further gains can achieved when combining both of these modeling techniques.
WER % (4)
Model
Params
devclean
devother
testclean
testother
Bi-LSTM (Zhang et al.
2020a
130M
3.2
7.8
Transformer
Zhang et al.
2020a
139M
2.4
5.6
Liu et al.
202
160M
2.2
4.7
ContextNet Han ct
2020
112.7M
2.0
4.6
2.1
4.6
Branchformer (Peng el al. 2022
116.2M
2.2
5.5
2.4
5.5
Conformer
Gulati et al.
2020
118.8M
1.9
4.4
2.1
4.3
Zhang et al.
2020b
2.0
4.7
2.2
4.8
Peng et al.
2022
116.2M
2.2
5.6
2.5
5.5
Our Baselines
Transformer 36L
129.
2.16
5,28
2.32
5.34
Conformer 24L
133.7M
1.95
4.84
2.21
5.04
Our Proposed Models
MSSM 32L
140.3M
1.80
4.96
2.01
4.61
Statcformer 25L
139.8M
1.76
4.37
1.91
4.36
Table 3: Performance of large configurations on Librispeech, compared with best results literature.
We also scale the model size up to 140M parameters verify the performance of large config- uration. Baselines and the proposed models this large configuration are trained using auxiliary classifiers similar to Szegedy al. 2015), in which intermediate encoder outputs trained pre- dict frame labels every 4 layers. Table 3 compares the proposed models with large-size baselines and state-of -ft-----thhtrrrre numbers previous works Pl Both the 32-layer MSSM and 25-layer Stateformer outperform our large-size baselines and a range of attention- based results in literature.
4.2 MASKED LANGUAGE MODELING
To explore our model's effectiveness on other modalities, we consider the widely used masked lan- guage modeling task long sequences (Devlin et al. 2018). We follow Beltagy et al. (2020) in preparing a scale long-document corpus, including Stories (Trinh & Le 2018). Real- News (Zellers al., 2019), Books (Zhu et al. (015) and English Wikipedia. This corpus, consisting of 6.6 billion tokens, is loaded as in |Bllagy et al. (020)) terms of 4,096-token long sequences.
It is prohibitive to apply the full attention in transformers to the 4096-token sequences due to the quadratic memory complexities. Therefore, we pick a few feasible transformer variants as baselines, including sliding-window attention (Beltagy al.] 2020), Performer (Choromanski et al. 22220) and Linformer (Wang et al. 2020a). The sliding window has a window size of 256; Performer uses random feature dimension of 256 and ReLU as kernel function; Linformer has compression ratio of 8. The MSSM model with Dah. = 32 has 8 heads; remaining hyperparameters are shared with attention baselines.
To validate that our model works well at scale, we test the attention and MSSM layers using 24- layer model similar in size RoBERTa-large (Liu al. 2019). To make a fair comparison, training stops after the same number of data batches. Following Xiong et al. (202I), baselines are trained using the standard masked language modeling objective for 100K steps with per GPU batch size of 8 on 32 A100 G 3 The MSSM is trained for 50K steps with a per GPU batch size of on 64 GPUs, for rapid development. Perplexities (PPLs) are reported on the validation set.
Table 4 summarizes the perplexity of different models. The proposed MSSM achieves performance on-par with attention baselines, but shows notably improved perplexity with further training (20K
"The Conformer performance varies significantly in prior work. We list few recent works for comparison.
"The training curves of these models nearly stable after 100k steps.
8

---

Under review as a conference paper at ICLR 2023
Model
Params
PPL (L)
Performer
359.1M
5.58
Linformer
459.7M
4.14
Sliding window
434.7M
3,47
MSSM
3.47
416. IM
+more updates
3.22
Table 4: Performance on masked language modeling. The final row shows the performance of MSSM with additional 20k training updates.
additional updates). Although the MSSM is expected to learn information from the full sequence, the performance gap to the sliding window is smaller than expected. To better understand the behavior of the proposed model, analyze the learned state space kernels in Section 4.3
4.3 MULTI-HEAM INTRODUCES MULTI-MODAL AND MULTI-SCALE MODELING POWER
Now that we have empirically demonstrate the effectiveness of multi-head state space models, we want answer the following research question: what does the mulli heads learn from data? Since the state space model can be presented as convolution (Equation 4), the shape of the kernel filters can be studied to answer this question.
singlebeeed
NUM muttinead (8)
muthi head continuc..
1.00
0.75
Nommatzed
20
0.40
1.5
0.25
Time Stops
Time Stops
Time Stops
Time Stops
Figure 4: Visualization of MSSM kernels learned on LibriSpeech and masked language modeling tasks: (a): Kernel learned by single-head SSM ASR. (b): MSSM kernels learned ASR. (c): Kernels learned MLM, showing multi-modal capability. (d): Kernels learned on MIM,, showing multi-scale capability. (c) and (d) together are heads from the same MSSM module. To simplify the discussion, we skip plotting heads with similar patterns: are plotted on ASR while 7 on MLM.
We first look into the learned multi-head kernels from ASR model (Figure 4B), and compare them with the learned single-head kernel (Figure 4A). We find the the single head kernel learned simple monotonic decaying pattern, which computes the weighted average of history in a sliding window. while the first head is still
On the other hand, the multi-head kernels much more interesting monotonic decaying, the kernels from the second and third heads involve context with different context concentration. In other words, multi-head state space models are able introduce multi- modal modeling power.
The kernels from the MLM tasks bring additional insights. 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. Also, this localized behavior can be viewed similarly to that of the sliding-winde window attention. Additionally, kemels from the next four heads (Figure 4D) are monotonic with different decay rates, which captures history with multi-scale granularity.
CONCLUSION
We propose the multi-head state space model (MSSM) for sequence modeling in this paper. Par- allel heads of SSMs, together with inter-head gating, are introduced, enabling multiple temporal dynamics to be effectively learned on sequence data. addition, we combine the MSSM and at- tention models to form new type of transformer architecture, referred to Stateformer. On the LibriSpeech speech recognition task, the proposed MSSM outperforms transformer baseline by a large margin. Besides, the Stateformer further improves the performance, achieving state-of-the- art performance. On the masked language modeling task, the MSSM achieves on-par performance compared with range of attention-based baselines trained large corpus.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. In International Conference on Learning Representations (ICLR), 2015.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. 051::2004400150. 2020.
William L Brogan. Modern control theory; 3rd ed. Prentice-Hall, Englewood Cliffs, NJ, 1991. URL cttp///ddcceeeeceeeeeeeeeeeeeeee cern. wwtp:ccccccc//c.222222332
https::/eds.cemn.cott
William Chan, Navdeep Jaitly, Quoc Le, and Oriol Vinyals. Listen, attend and spell: A neural network for large vocabulary conversational speech recognition. In 2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2016.
Peng-Jen Chen, Ann Lec, Changhan Wang, Naman Goyal, Angela Fan, Mary Williamson, and Jiatao Gu. Facebook Al's WMT20 news translation task submission. In Proceedings of the Fifth Conference on Machine Translation, 2020.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Savelie Cornegrua,, Robert Bakewell, Samuel Withey, and Giovanni Montana. Modelling radiolog- ical language with bidirectional long short-term memory networks. 2016.
Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks, International Conference on Machine Learning, 2017.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: aet:. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbom, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x words: Transformers for image recognition at scale. In International Conference Learning Representations, 2021.
James Durbin and Siem Jan Koopman. Time Series Analysis by State Space Methods. Oxford University Press, 05 2012. ISBN 9780199641 doi: 10. 0011003/appoofos00777799999999 0001.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state- space models. International Conference on Machine Learning (ICML), 2022.
Alex Graves. Sequence transduction with recurrent neural networks. In International Conference of Machine Learning (ICML) Workshop on Representation Learning, 2012.
Albert Gu, Tri Dao. Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33, 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao. Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-time time models with linear state-space layers. Ad- vances in neural information processing systems, 34, 2021.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré, On the parameterization and initialization of diagonal state space models. arXiv preprint arXiv:2206. /193, 2022b.
Anmol Gulati, James Qin, Chung-Cheng Chiu. Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, and Ruoming Pang. Conformer: Convolution-augmennte transformer for speech recognition. Interspeech, 2020.
10

---

Under review as a conference paper at ICLR 2023
Wei Han, Zhengdong Zhang, Yu Zhang, Jiahui Yu, Chung Cheng Chiu, James Qin, Anmol Gulati, Ruoming Pang, and Yonghui Wu. Contextnet: Improving convolutional neural networks for automatic speech recognition with global context. arXiv preprint arXiv:2005.03191, 2020.
Dan Hendrycks and Kevin Gimpel. Gaussian error
linear
units (gelus). arXiv preprint
160.:01606008415, 2016.
Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, and Jurgen Schmidhuber. Gradient flow in recur- rent nets: the difficulty of learning long-term dependencies. 2001.
Rob Hyndman, Anne B Koehler, Keith Ord, and Ralph Snyder. Forecasting with Exponential
Smoothing: The State Space Approach. Springer Berlin, Heidelberg, 2008. ISBN 978-3-5400
71916-8. doi: http://do..rrr11101100/////7////
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Juunyasuvunakool, and Russ Bates et al. Highly accurate protein structure prediction with alphafold. In Nature, 2021.
R. E. Kalman. A New Approach to Linear Filtering and Prediction Problems. Journal of Basic Engineering, 82(1):35 03 1960. ISSN 0021-9223. doi: 10.1115/1.366252.. URL https: //doi www.10..//000.....o.o 1115/1 3662552
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th Interna- tional Conference on Machine Learning, 2020.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference for Learning Representations (ICLR), 2015.
Nikita Kitaev, L ukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Learning Representations (ICLR), 2020.
Taku Kudo and John Richardson. Sentencepiece: simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv: 1808.088000022200 2018.
Chunxi Liu, Frank Zhang, Due Le, Suyoun Kim, Yatharth Saraf, and Geoffrey Zweig. Improving rnn transducer based asr with auxiliary tasks. In 2021 IEEE Spoken Language Technology Workshop (SLT), pp. 172-179. IEEE, 2021a.
Yinhan Liu, Myle Out, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907:1197.11999 2019.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV), 2021b.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language model- ing via gated state spaces. 2022. URL https: http//axvvv...///////eee//a.at.a org/ abs abs/aaassssss 6 13947
Nathan Ng, Kyra Yee, Alexei Baevski, Myle Ott, Michael Auli, and Sergey Edunov. Facebook FAIR's WMT19 news translation task submission. Proceedings of the Fourth Conference on Machine Translation (Volume 2: Shared Task Papers, Day 1), 2019.
Myle Ott, Sergey Edunov, Alexci Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
Victor Pan. Structured matrices and polynomials: BirkhäaseerSpringer, Boston, 2001. ISBN 08 1817642404.
Unified superfast algorithms.
Victor Pan. Fast approximate computations with cauchy matrices, polynomials and rational func- tions. In Computer Science Theory and Applications, Cham, 2014. Springer International Pub- lishing. ISBN 978-3-399-06666--8

---

Under review as a conference paper at ICLR 2023
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: An ast corpus based on public domain audio books. In International Conference on Acoustics, Speech and Signal Processing (ICASSP),001)
Yifan Peng, Siddharth Dalmia, Jan Lane, and Shinji Watanabe. Branchformer: Parallel mlp-attention architectures capture local and global context for speech recognition and understanding. In International Conference on Machine Learning, pp. 17627-17643. PMLR, 2022.
Syama Sundar Rangapuram, Matthias W Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski. Deep state space models for time series forecasting. In S. Bengio, H. Wal- lach, H. Larochelle, K. Grauman, N. Cesa- Resa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning 250 million protein sequences. 2021.
Matthias Seeger. David Salinas, and Valentin Flunkert. Bayesian intermittent demand forecasting for large inventories. In Advances in Neural Information Processing Systems, 2016.
Matthias Seeger, Syama Rangapuram, Yuyang Wang. David Salinas, Jan Gasthaus, Tim Januschowski, and Valentin Flunkert. Approximate bayesian inference in linear state space mod- els for intermittent demand forecasting at scale. 09 2017.
Noam Shazeer, Zhenzhong Lan, Youlong Cheng, Nan Ding, and Le Hou. Talking-heads attention. In arXivv2003.024336 2020.
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In WACV, 2021.
Jimmy T.H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.049233 2022.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1-9, 2015.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 2022.
Chau Tran, Shruti Bhosale, James Cross, Philipp Koehn, Sergey Edunov, and Angela Fan. Face- book AI's WMT21 news translation task submission. In Proceedings of the Sixth Conference on Machine Translation, 2021.
Trieu H Trinh and Quoe Le. A simple method for commonsense reasoning. arXiv preprint arXiv:1806.0268 02847, 2018
Arnold Tustin. A method of analysing the behaviour of linear systems terms of time series. In Journal of the Institution of Electrical Engineers-Par IIA: Automatic Regulators and Servo Mechanisms, 1947.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances Neural Infor- mation Processing Systems, 2017.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention 0406.0688...
with linear complexity. arXiv preprint arXiv:2006.0u 04768, 2020a.
Yongqiang Wang. Abdelrahman Mohamed, Due Le, Chunxi Liu, Alex Xiao, Jay Mahadeokar, Hongzhao Huang, Andros Tjandra, Xiaohui Zhang, Frank Zhang, et al. Transformer-basee acous- tic modeling for hybrid speech recognition. In ICASSP 2020-2022 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp.
IEFF, 2020b.
12

---

Under review as a conference paper at ICLR 2023
Wenhan Xiong, Barlas Oğuz, Anchit Gupta, Xilun Chen, Diana Liskovich, Omer Levy, Wen-tau Yih, and Yashar Mehdad. Simple local attentions remain competitive for long-context tasks. arXiv preprint arXiv:2112.0722 2021.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences, Advances in Neural Information Processing Systems, 2020.
Rowan Zellers, Ari Holtzman, Hannah Rashkin, Yonatan Bisk, Ali Farhadi, Franziska Roesner, and Yejin Choi. Defending against neural fake news. Advances in neural information processing systems, 32, 2019.
Qian Zhang, Han Lu, Hasim Sak, Anshuman Tripathi, Erik McDermott, Stephen Koo, and Shankar Kumar. Transformer transducer: A streamable speech recognition model with transformer coders and rnn-t loss. 2020a.
Songyang Zhang, Shipeng Yan, and Xuming He. LatentGNN: Learning efficient non-local rela- tions for visual recognition. In Proceedings of the 36th International Conference on Machine Learning(/CML), 2019.
Yu Zhang, James Qin, Daniel Park, Wei Han, Chung-Cheng Chiu, Ruoming Pang, Quoc V Le, and Yonghui Wu. Pushing the limits of semi-supervised learning for automatic speech recognition. arXiv preprint arXiv:2010.10500 2020b.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pp. 19-27, 2015.
13


Reviews:
[{"review_id": "c2ZIbPQqk-", "reviewer": "Reviewer_qqzq", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes some architectural changes to the S4 model and applies it to speech recognition and language modeling tasks.\n", "summary_of_the_review": "This paper proposes several modifications to the S4 model and architecture to achieve state-of-the-art results in speech recognition. Weaknesses are a lack of clarity in the details of the model, and a major lack of discussion and comparison to related work.", "strength_and_weaknesses": "Amendment after rebuttal:\n\nThe original review pointed out that many of the contributions in this work are equivalent to flags available in the public S4 model. The authors have since submitted supplementary code showing that while based on the public model, several of these newer modifications are implemented somewhat differently.\n\nThis points to one of the main weaknesses of the submission, that the proposed changes are each not particularly novel. The other main weakness is a lack of discussion and comparison to very closely related prior work.\n\nThe strengths of the paper lie in the empirical results, which show strong performance in a number of speech recognition benchmarks as well as language modeling.\n\nA much more detailed review and discussion is contained in the comments, with suggestions for improvement in the experiments and presentation.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TtScEWwxLb4", "reviewer": "Reviewer_F93h", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Tasks: Automatic speech recognition (ASR, measure WER), language modeling (measure PPL)\n\nTwo new models are proposed, which are based on state space models (SSMs).\n\nThe newly proposed models are: MSSM and Stateformer. The MSSM purely uses the SSM, with residual connections and some layer norm. The Stateformer is like a Transformer, i.e. it uses multi-head attention, and it has a new block type which embeds the MSSM.\n\nOn ASR, on Librispeech, without using an external language model, new state-of-the-art results are achieved.\n\nFor language modeling, the new model performs just the same as the sliding window attention model.\n", "summary_of_the_review": "It's an interesting extension of the SSM, applied to speech recognition and language modeling, and it achieves state-of-the-art (SOTA) results on Librispeech without external LM.\n\nThere are some weaknesses though. Some of them can be improved now by better clarification, but not all of them, where more experiments should be done.\n\nStill, due to the good results for speech recognition, I think it's interesting enough.", "strength_and_weaknesses": "\nStrengths:\n\n- Good state-of-the-art results with the Stateformer on speech recognition.\n\nWeaknesses:\n\n- The code is not released. No information is given on the code.\n- No real improvement for language modeling.\n- Speech recognition lacks experiments also with an external language model. It's good to also see the numbers without, but using an external language model is the more realistic setting for speech recognition.\n- Subsampling seems to be non-standard, and suboptimal. See below.\n- How does it compare when replacing the SSM by a RNN or LSTM? It is said that SSMs are equivalent to linear RNNs. So what exactly is the advantage of a SSM? This does not become clear. Such experiments should be done to really see the difference.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OgJL8yDlC8a", "reviewer": "Reviewer_PQ8K", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper combines continuous time state-space models with transformer themes.  The method is applied to large language and speech modelling tasks, and is compared to some baselines.  The method appears to perform well.  Some analysis of the temporal characteristics of the learned model are explored.  As far as I can tell, no code or supplementary materials were submitted. ", "summary_of_the_review": "Overall I do not think this paper is currently at the requisite level for publication.  Obvious links to existing work are missing.  I am not confident the core contribution itself is actually novel.  Standard evaluation benchmarks are missing.  The clarity of the written communication is poor.  \n\nI am willing to entertain the idea that my understanding of S4, or my understanding of particularly Equ. (8, 9) are incorrect.  If this is the case (and it is indeed how the authors claim it to be) then this would improve the paper slightly in my eyes, but probably not enough for me to rate the paper as acceptance worthy.  \n\nI also note that this paper is one of as many as half a dozen papers submitted to ICLR advancing S4.  \n", "strength_and_weaknesses": "# Strengths:\nThe paper appears to perform favourably on the tasks presented.  \n\n# Weaknesses:\nThe authors highlight three core contributions of the paper: multi-head parallel attention, head gating, and combination with attention through a bi-directional SSM.  I do not agree with any of these claims.  \n\n(1.1)  I am not convinced that stacking the layers is a better architecture, as the “ablations” you provide are a bit confusing (unclear how performances change with overall depth).  I also don’t understand how you allow multi-head parallel processing where S4 does not? \n\n(1.2)  I do not believe you have correctly assessed how S4 implements attention, and so I do not believe you have “contributed” this.  (see comment below)\n\n(1.3)  S4 uses a bidirectional SSM, which is not mentioned in your submission.  I also don’t fully understand why a transformer would _require_ bidirectionality in the SSM layer, as the transformer itself is all-to-all.  Ultimately, I find the overall architecture a little confusing.  I’d like Figure 3 to be explained a little more fully, maybe with some accompanying math close to or within the figure, to concretely define the spaces that each element operates on, its computational complexity etc.  \n\n\n## (2)  Connections with S4:  \nI find the presentation and connections to previous work, specifically in the arc of S4 papers, to be chronically lacking.  \n\n(2.1)  _I implore the authors to comment on this prior to me making my final judgement, as I may have misunderstood their claims._  Most significantly, I believe the description of S4 and the use of GLU activation is incorrect.  In Equ. (8), the authors claim that S4 uses GLU activations applied to each two-dimensional output channel per input channel.  To my understanding, this is incorrect, as S4 applies GLU to the entire $y_k^{1:H}$ vector.  The interhead gating proposed in (9) is actually a special case of the GLU activation used by S4, where the number of output dimensions is strictly tied to twice the number of input channels with size-two block-diagonal weights.  As this is the core contribution of the paper, this is a critical point. \n\n(2.2)  The bidirectional design (Section 3.4) is also proposed by S4.  This is not mentioned.  To the reader, it appears that this bidirectional design is novel.  The parameterization used in Equ (7) is also a special case of the parameterization in Goel+ [2021], but this is not highlighted. \n\n(2.3)  The block-diagonal interpretation was also explored by Smith+ [2022] (although I accept that this is probably acceptable as concurrent work).  \n\n(2.4)  Was S4 evaluated using the original codebase released by Gu+ [2021], or did you use your own S4 implementation? \n\n(2.5)  Broader discussion of LSSL, S4, HTTYH, S4D, DSS, Sashimi etc is chronically lacking.\n\n## (3) Core claim on contributions:  \n(This also relates to the exposition of S4 above) I do not understand the author's claim that S4 cannot capture different timescales.  The whole premise of S4 is that each SISO SSM is imbued with its own timescale, and hence can integrate information over different timescales.  Recent work, particularly How To Train Your HiPPO [Gu+, June 2022, arXiv], has explored this further, studying the kernel that S4 is implicitly using.  This paper, although recent, should still have been cited, since it tackles very similar core ideas as here.  It is also not clear to me how multi-head attention actually facilitates the claimed improvement in the ability to capture timescales.  Maybe this is true, but with the very limited intuition provided in Section 4.3, I don’t agree that better performance on benchmark tasks can be directly attributed to the new mechanism.  I am not convinced that the authors didn’t simply tune a larger model, with more parameters, flexibility, expressive non-linearities etc, until they got better performance – which then they post-hoc attribute to the capturing of timescales.  \n\n(3.1)  You state: “... long and short term dependency would both be useful in sequence modelling…”.  I do not believe you have backed the claim that yours offers this whereas S4 does not  –  this was pretty much the whole point of S4.  \n\n(3.2)  As an example, I would look at the eigen spectra of each learned latent SSM, for both SSM and the Stateformer, since this describes the temporal dynamics of the filter.\n\n## (4)  Empirical validation:  \nI am not overly familiar with the particular benchmarks that the authors present.  I would like to know why they didn’t use the LRA benchmark [Tay+, 2021] or the Speech Commands dataset [Warden, 2018], as these seem like much more standard and widespread benchmarks.  These also don’t require numerous A100’s, and so also help reproducibility.  \n\n(4.1)  How were hyperparameters selected for the different methods? \n\n(4.2)  Table 1 is very unclear to me.  Does the 16L-2S model actually have 32 layers?  Or are 16 layers arranged into eight blocks of two.  \n\n(4.3)  It is unclear what each variant in the table is exploring.  It appears to me that a soup of different layers and architectures have been tested and the best one is being reported as a substantial scientific contribution.  \n\n(4.4)  The evaluation in Section 4.3 is not explained very well.  I do not understand how the curves in Figure 4 were generated, and so I cannot evaluate the validity of the claims.  I am not even entirely sure what the curves are showing. \n\n\n\n# Minor Weakness / Typographical Errors:\n(a) Figures and tables should be floated to the top or bottom of pages.\n\n(b) Footnotes should be avoided if possible.  \n\n(c) States are missing a time index, e.g. in Equ. (8) it should be $y_{k, n1}$, where $k$ is the time index in the sequence.  \n\n(d) I think that the claim in Section 3.2 that the shared recurrence parameters stabilise training is unsupported.  I also don’t think the second half of that claim (that multi-head frameworks give it the flexibility to learn multiple timescales) is fully supported either.  This may be me misunderstanding the analysis presented in Section 4.3.  \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7uOkZq5Gpzq", "reviewer": "Reviewer_aXdZ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed a multi-head state space model (MSSM) and the Stateformer which is a combination of bidirectional MSSM and self-attention. The performance of MSSM and Stateformer are evaluated on Librispeech for ASR. Further, the effectiveness of MSSM is explored on the MLM task.", "summary_of_the_review": "See above.\nThe major concern is about the experimental evaluations.", "strength_and_weaknesses": "**Strengths**\n\nThe paper is generally well-written.\n\nThis work is a quick application of the recently proposed Linear State-Space Layer (LSSL). The ideas of multi-head and a combination of bidirectional MSSM and self-attention seem to be novel.\n\nAs compared to existing SSM methods, this paper introduces a simpler parameter initialization manner and the multi-head structure. In the librispeech experiment, the Stateformer obtains the SOTA performance.\n\n**Weaknesses**\n\n\"Stateformer 25L\" performs close to \"Gulati et al. 2020\" (better only over testclean). So the real benefit taken by Stateformer could be questioned. Can the authors argue against this concern?\n\nTraining with the auxiliary classifiers is not common in ASR. As said in the paper, \"Our Baselines\" and \"Our proposed models\" in this large configuration are trained using auxiliary classifiers. It brings a further concern that the superior performance of \"Stateformer 25L\" may be partly due to the use of auxiliary classifiers. It is better to have an ablation study.\n\nHow the perplexities are calculated for masked language models?\n\nCompared to the improved results in ASR, the improvement on masked language modeling seems weaker. MSSM needs 20K additional updates to perform better than Sliding window. The authors should compare different models for the same number of training iterations.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject