Efficient Hybrid Long Sequence Modeling with State Space Augmented Transformers

Authors: simiao  zuo
xiaodong  liu
jian  jiao
denis  charles
eren  manavoglu
tuo  zhao
jianfeng  gao
⋄ ⋄ microsoft
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 22:18:18

Abstract: Transformer models have achieved superior performance in various natural language processing tasks. However, the quadratic computational cost of the attention mechanism limits its practicality for long sequences. There are existing attention variants that improve the computational efficiency, but they have limited ability to effectively compute global information. In parallel to Transformer models, state space models (SSMs) are tailored for long sequences, but they are not flexible enough to capture complicated local information. We propose SPADE, short for State Space Augmented Transformer. Specifically, we augment a SSM into the bottom layer of SPADE, and we employ efficient local attention methods for the other layers. The SSM augments global information, which complements the lack of long-range dependency issue in local attention methods. Experimental results on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method. To further demonstrate the scalability of SPADE, we pre-train large encoder-decoder models and present fine-tuning results on natural language understanding and natural language generation tasks.

Main paper: 
 introduction: 
 With the rise of large language models, the difficulties of modeling long sequences have gained increasing attention. For instance, ChatGPT is capable of handling context that comprises up to 8k tokens, while GPT-4 (OpenAI, 2023) scales this ability up to 32k tokens. Conventional Transformer-based models rely on the attention mechanism (Vaswani et al., 2017) , which computes a dependency score for every pair of tokens in the input sequence. Thus, full attention has a quadratic time and space complexity with respect to the length of the sequence. However, such complexity proves computationally prohibitive for tasks that require modeling long sequences such as text summarization (Nallapati et al., 2016) and question-answering (Kwiatkowski et al., 2019) . In fact, we find that training a Transformer model (250M parameters) takes up over 80G of GPU memory when modelling an input sequence of length 8k.In addition, Transformer models that rely on full attention run the risk of overfitting due to the lack of structural biases (Lin et al., 2022) . The attention mechanism does not impose any structural prior over the input. As a result, order information (such as sinusoidal encoding) is required to train a Transformer model successfully. Transformer models, equipped with full attention mechanism, prove overly flexible resulting in overfitting to the noise. This significantly impacts the models' practicality in long sequence modeling, where the dependency signal is often weak, and signal-to-noise ratio is low (i.e., in long sequences, a majority of input tokens are useless). Empirical evidence shows that Transformer models without structural biases have a classification accuracy rate of 57.5% on a two-way classification task, nearly 30% less than state-of-the-art approaches that are equipped with powerful structural biases (see Section 4.1 for details).Several methods have been proposed to address the problem of full attention's quadratic complexity and the need to introduce structural biases to Transformer models. One of the primary approaches is employing approximation methods, which approximates full attention using fast algorithms with linear complexity. Low-rank approximations (Wang et al., 2020b) or kernel methods (Peng et al., 2021) , for instance, could simplify and speed up calculation of the attention score matrix (i.e., softmax(QK ⊤ / √ d) in Eq. 1). Nevertheless, although these methods mitigate the computational complexity of full attention, they inherit the problems associated with the lack of structural bias.To incorporate structural biases into Transformer, partial attention methods have been proposed. The methods can be further categorized in to sparse attention (Beltagy et al., 2020) and clustering (Kitaev et al., 2020) . In the former approach, individual tokens attend to only a subset of tokens determined by pre-defined sparsity patterns. In the latter, tokens are divided into clusters and intra-cluster attention is performed. However, introducing these structural biases restricts the models' ability to capture global information. Local window attention, for instance, assumes that each token depends only on its direct neighbors, causing long-range and global information to be lost.State space models (SSMs) introduce structural biases tailored for computing global information (Gu et al., 2022b) , in contrast to partial attention. Specifically, SSMs design fixed global dependency patterns that facilitate effective and efficient computation, and can be seen as linear recurrent neural networks with specifically designed fixed weights. Moreover, efficient algorithms have been developed to train these models. However, it should be noted that SSMs can be restrictive, because they do not capture local information as effectively as attention-based models, which explicitly compute dependencies among input tokens.We propose SPADE, short for State sPace AugmenteD TransformEr. SPADE is a multi-layer Transformer model that can effectively and efficiently capture complicated dependencies. Specifically, we augment a state space model (SSM) into the bottom layer of the model to integrate inputs with global information. Because the SSM only provides coarse global information, at the subsequent top layers of SPADE, we employ local attention methods to capture more complicated and refined local information. With this approach, the SSM induces a strong structural bias that augments global information and complements the long-range dependency issue in local attention methods. SPADE is flexible to accommodate different building blocks. For example, we can use SSMs such as S4 (Gu et al., 2022b) and S5 (Smith et al., 2022) in the bottom layer; while for subsequent layers, local attention algorithms such as sliding window attention and chunk attention can be applied.We demonstrate the efficiency and effectiveness of SPADE on various tasks. First, we show that SPADE outperforms existing approaches on the Long Range Arena (Tay et al., 2021b) benchmark, which is designed to test models' ability in modeling long sequences. Second, we show that in autoregressive language modeling, SPADE is not only significantly faster than the vanilla Transformer (Vaswani et al., 2017) , but also yields better performance. Third, we demonstrate the scalability of SPADE by conducting language model pre-training and fine-tuning experiments. Specifically, we pre-train an encoder-decoder model similar to T5 (Raffel et al., 2020) . And we fine-tune the model on various tasks, including natural language understanding and natural language generation benchmarks. In all the settings, SPADE outperforms the baselines. Our code1 and pre-trained model checkpoints2 are publicly available. 
attention mechanism: 
 Suppose the input to the layer is X ∈ R L×d , where L is the sequence length and d is the embedding dimension, then the attention mechanism outputsEQUATION)Here W q , W k , W v ∈ R d×d are learnable weights. The attention mechanism can simultaneously compute the alignment between any pair of input tokens, such that it models long-range dependencies better than recurrent neural networks. Specifically, denote the attention score matrix A = softmax(QK ⊤ / √ d) ∈ R L×L . Then, A ij captures the alignment between the i-th and the j-th input tokens.Continuous time state space model. A continuous time latent state space model maps a 1-dimensional input signal u(t) to a d s -dimensional latent state x(t), after which x(t) is mapped to a 1-dimensional output signal y(t). Concretely,x ′ (t) = Ax(t) + Bu(t), y(t) = Cx(t).(2)Here, A ∈ R d s ×d s , B ∈ R d s and C ∈ R d s .Existing works leverage Eq. 2 to model long sequences. For example, Gu et al. (2020) (u 0 , u 1 , • • • , u L ), where L is the sequence length. To facilitate modeling discrete data, the model in Eq. 2 can be discretized (using the bilinear method) by a step size ∆, such thatEQUATIONwhereA = (I -∆/2 • A) -1 (I + ∆/2 • A), B = (I -∆/2 • A) -1 ∆B, C = C.We unroll the above recurrent representation, after which we havey k = CA k Bu 0 + • • • + CABu k-1 + CBu k .This can be written as a convolutional representationy = K * u, where K ∈ R L = CB, CAB, • • • , CA L-1 B . (4)Here, K is the convolutional kernel, " * " is the discrete convolution operator, u represents the input sequence (u 0 , u 1 , • • • , u L ), and y represents the corresponding output sequence (y 0 , y 1 , • • • , y L ).In Eq. 4, the output y can be computed efficiently given that the convolution kernel K is known (e.g., using Fast Fourier Transform). However, computing the kernel is non-trivial.Most of existing algorithms have O(L 2 ) time and space complexity.Structured State Space Sequence model (S4). Gu et al. (2022b) develop the S4 model to efficiently compute Eq. 4. Specifically, C in Eq. 2 is randomly initialized, and A and B are initialized asEQUATION)whereP i = (i + 1/2) 1/2 , A (d s ) ij = -    (i + 1/2) 1/2 (j + 1/2) 1/2 , i > j, 1/2, i = j, -(i + 1/2) 1/2 (j + 1/2) 1/2 , i < j.Subsequently, the convolution kernel K in Eq. 4 can be computed efficiently with near linear time and space complexity. Then, for an input u, the S4 output y = K * u can be computed efficiently.More state space models. Despite the superior performance of S4, it suffers from known issues such as difficulty in scaling. Other state space models are subsequently developed, such as MEGA (Ma et al., 2022) , S4-D (Gu et al., 2022a ), H3 (Fu et al., 2022) , S5 (Smith et al., 2022) , GSS (Mehta et al., 2022) and Hyena (Poli et al., 2023) . Through experiments, we demonstrate that the proposed model is flexible to accommodate variants of SSMs. We use language modeling experiments (see Section 4.2), and the sequence length is 3k. 
method: 
 We first conduct experiments to demonstrate that SSMs do not model local information well. Then, we present SPADE, which effectively combines global and local information by augmenting SSMs into the Transformer architecture. The motivation behind SPADE is that even though SSMs perform well on several long sequence classification tasks (Gu et al., 2022b) , they perform poorly on language modeling, which is a fundamental task in natural language processing. To demonstrate such an observation, we compare S4 with Transformer with full attention and Transformer with local (window and chunk) attention. In local attention, each token can only attend to its neighboring tokens (see Figure 1 for illustrations). We conduct experiments on token-level language modeling. In this setting, local information is more important than global information. This is because in practice, we rarely see words (tokens) that are thousands of positions apart exhibit strong dependencies (Sukhbaatar et al., 2019) .Experimental results are illustrated in Figure 2 . We see that both Transformer with full attention and Transformer with local attention (e.g., window and chunk) outperforms S4. Notice that replacing full attention with local attention does not significantly hurt model performance, indicating that local information is more important in this setting. We remark that SSMs such as S4 produces a fixed dependency pattern, e.g., the convolution kernel in Eq. 4. Moreover, unlike attention, SSMs do not explicitly compute dependencies among tokens. Therefore, SSMs are not refined enough to capture local information, such that they perform poorly on language modeling tasks.We propose SPADE, which is a multi-layer Transformer model that can capture complicated global and local information. The overall architecture of SPADE is shown in Figure 3 (left). The proposed model employs a hierarchical structure. Specifically, at the bottom layer of SPADE (termed the global layer), we capture global dependencies using a SSM. Because the SSM only provides coarse global information, the subsequent local layers facilitate the model to handle more refined and complicated local dependencies. In other words, the SSM induces a strong structural bias that augments global information to the inputs.To instantiate the local layer, we replace the full attention in the conventional Transformer layer with off-the-shelf efficient local attention methods. SPADE is flexible to accommodate different approaches, such as window attention and chunk attention (see Figure 1 for illustrations).In the global layer (Figure 3 , right), given the input X to the layer, we have the output Y asY = FFN (LN(X a )) + X a , where X a = W LN(X local ), LN(X global ) + X, X local = Local (LN(X)) , X global = SSM (LN(X)) .Here, LN(•) denotes layer normalization (Ba et al., 2016) , FFN(•) denotes a two-layer feed-forward neural network, and W is a trainable weight that combines local and global representations. Notice that we apply normalization to X local and X global to align their scales.SPADE is flexible to accommodate different building blocks. For example, in the experiments we use S4 (Gu et al., 2022b) , S5 (Smith et al., 2022) and Hyena (Poli et al., 2023) as the SSM in the bottom layer. In the top layers, local attention algorithms such as sliding window attention and chunk attention can be used.We remark that because of the sequential nature of SSMs (Eq. 3), the global layer can encode positional information of the inputs. Therefore, we do not need additional fixed-length positional embedding techniques (Devlin et al., 2019) . 
experiments: 
 In the experiments, we implement all the models using PyTorch (Paszke et al., 2019) and Fairseq (Ott et al., 2019) . Training details such as hyper-parameter settings are deferred to the appendix.Dataset and models. We evaluate SPADE on Long Range Arena (LRA, Tay et al. 2021b) , which is a benchmark tailored for evaluating models' ability in modeling long sequences. Dataset details are presented in Appendix B.Following the setting in Ma et al. 2022 , we use small models (less than 2M parameters) for all the tasks. We limit the computational budget such that all the models are trained with similar speed for the same amount of time. To aggregate local information, we consider two approaches: window attention and chunk attention. We use S4 Gu et al. as 11 in the appendix). We see that both variants of SPADE (window and chunk) significantly outperform all the baselines in terms of average accuracy. For example, the window attention variant outperforms the best-performing baseline (MEGA-chunk) by 0.5%, and the chunk attention variant has a 1.8% performance gain. Therefore, SPADE is more suitable to model long sequences than existing approaches.Transformer (Vaswani et al., 2017) 18.8 Transformer (relative) (Shaw et al., 2018) 18.7Transformer (window) 19.7 Transformer (chunk) (Hua et al., 2022) 20.9 MEGA (chunk) (Ma et al., 2022) 19.8 S4 (Gu et al., 2022b) 23.2 S5 (Smith et al., 2022) 23.0 Hyena (Poli et al., 2023) 23 Dataset and models. We further evaluate our model by conducting language modeling experiments on Wikitext-103.The dataset contains English-language Wikipedia articles, and the total number of tokens is 103M. In all the experiments, we set the input sequence length to 3k and train for 286k steps. During testing, we set the sequence length to 3k and the context window size to 400.We use window attention as the local information extractor in SPADE. And we equip our model with three variants of SSMs: S4 (Gu et al., 2022b) , S5 (Smith et al., 2022) and Hyena (Poli et al., 2023) . For SPADE and Transformer-based baselines, we follow the settings in Baevski & Auli (2019) , where we use a large-scale Transformer model with 16 layers and about 250M parameters. For SSM-only baselines (e.g., S4), we also scale the models such that they contain approximately the same number of parameters.Experimental results are presented in Table 2 . From the results, we see that the proposed model achieves significant performance improvement and outperforms all the baselines. We remark that SPADE with window attention is not only significantly faster than the Transformer with full attention, but also yields a better performance.Remark. We remark that we do not need to train the S4 in the bottom layer of SPADE (S4+window) to achieve the performance in Table 2 . That is, we initialize the parameters in S4 using Eq. 5, and the parameters are frozen during training. This is because even without training, the initialization of S4 yields intriguing theoretical properties, which facilitates S4's ability to capture global information. 
language model pre-training: 
 We implement model pre-training using Fairseq, and we implement model fine-tuning using MT-DNN (Liu et al., 2019a; 2020b) . Note that all our experiments only use single task fine-tuning. Details such as pre-training settings and hyper-parameter settings are deferred to the appendix.To demonstrate the scalability of the proposed method, we pre-train an encoder-decoder variant of SPADE. The model architecture is the same as T5 base (Raffel et al., 2020) , except that we use post-layernorm instead of pre-layernorm to improve model performance (Liu et al., 2020a; Xiong et al., 2020) . The embedding dimension is 768, the hidden dimension of the FFN is 3072, the number of attention heads is 12, and both the encoder and the decoder have 12 layers. We add a S4 module to the bottom layer of SPADE, and the parameters of the S4 are fixed after initialization (Eq. 5). We use the window attention as the local information extractor, where we set the window size to 128. The model contains about 290M parameters.We consider two pre-training settings. In the first setting, we follow the pre-training settings in BERT (Devlin et al., 2019) , and we term the trained model SPADE base . In the second setting, we follow the pre-training settings in RoBERTa (Liu et al., 2019b) , and we term the trained model SPADE base++ . To facilitate fair comparisons, we also pre-train a T5 model using the second setting (the model yields better performance than the first setting), and we term the pre-trained model T5 base (re-imp).We remark that because the S4 module is not trained after proper initialization, and we do not use fixed-length positional embedding, our pre-trained model can be fine-tuned with any sequence length. For example, we can set the sequence length to 16k during fine-tuning, which is longer than the sequence length used in pre-training.We fine-tune the pre-trained models on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2019) , which is a collection of natural language understanding tasks. Dataset details are presented in Appendix B. We do not consider the long sequence setting in these tasks. In the experiments, all models are fine-tuned under the sequence length of 512. Experimental results are presented in Table 3 . We see that both variants of SPADE significantly outperforms T5 base . For example, T5 base has a 83.9 average accuracy on the MNLI dataset (average of MNLI-m and MNLI-mm); while SPADE base has a 87.2 average accuracy (+3.3) and SPADE base++ has a 89.4 average accuracy (+5.5). Recall that the sequence length is set to 512, which is the standard setting instead of the long-sequence setting. Therefore, the results indicate that SPADE is universal in that it is suitable to model both long and short sequences. We conduct language modeling experiments on Wikitext-103 with window attention (window=128).We fine-tune the pre-trained models on several abstractive summarization datasets. Dataset details are presented in Appendix B. We use ROUGE-2 as the evaluation metric.We compare SPADE with LongT5 (Guo et al., 2022) , which is a state-of-the-art model tailored for long sequences. Experimental results are summarized in Table 4 . From the results, we see that our model significantly outperforms LongT5. Note that SPADE base++ have about 290M parameters, while LongT5 large contains about 770M parameters and LongT5 xl contains about 3B parameters. From the results, we see that in all the tasks, our base-sized models have on par or better performance compared with LongT5 large . On the MultiNews dataset, our model even outperforms LongT5 xl , which is over ten times larger than our model. 
effectiveness of global information extractors: 
 Recall from Eq. 3 that SSMs are essentially linear recurrent neural networks which can be computed much more efficiently than conventional RNNs. In Figure 4 , we explore model variants with different global information extractors. Specifically, we equip Transformer with window attention with either S4 or LSTM (Hochreiter & Schmidhuber, 1997) .From the results, we see that RNNs can indeed capture global information. Also, we see that SSMs are more effective than RNNs. This is because RNNs suffer from known numerical problems such as forgetting and gradient explosion/vanishing (Pascanu et al., 2013) . On the other hand, the theoretical properties of SSMs (Gu et al., 2020; 2021; 2022b) can greatly alleviate such issues. We compare the efficiency of SPADE with other models: Transformer with full attention, Transformer with window attention, MEGA-chunk, and S4. The results are illustrated in Figure 5 . We see that SPADE is efficient in terms of both training speed and GPU memory usage. For example, when the sequence length is 6k, Transformer uses about 60GB of GPU memory, whereas SPADE with window attention only uses 27GB. Moreover, notice that SPADE also trains significantly faster than the vanilla Transformer under all settings. Notice that S4 may be less efficient than the vanilla Transformer (e.g., when the sequence length is 3k). This is because in Gu et al. 2022b , each layer of the model contains multiple S4 modules and expensive non-linear components. Therefore, the per-layer computational cost can exceed full attention when the sequence is not extremely long.We remark that even though we add a S4 module to the bottom layer of SPADE, such an additional module does not induce much computational overhead. We see that both the training speed and the memory usage of SPADE with window is only marginally different from those of window-attention Transformer. We have similar observations for the chunk attention variant.Recall that in SPADE, the bottom layer is equipped with a SSM and serves as the global layer, while the rest are local layers (see Figure 3 ). In Figure 6 , we empirically justify this design choice.We investigate the possibility of incorporating more global layers: we set the bottom 1 (the default choice), 4, 8, and 16 (all) layers as global layers. From the results, we see that model performance decreases as we use more global layers. This is because the SSM in the bottom layer captures and filters out global information, such that subsequent SSMs only introduce noise to the intermediate representations.We also investigate whether the global layer can be the top instead of the bottom layer in SPADE. From Figure 6 , we see that model performance drops significantly. This is because as a global information extractor, the global layer encodes positional information, on which the local attention modules rely. Therefore, using the global layer as the top layer is akin to using Transformer models without positional encoding, which will yield unsatisfactory performance.We examine how performance changes when the sequence length and window size change.From Figure 7 (left), we see that when we increase the sequence length from 512 to 3k, performance of Transformer with full attention increases. However, when we further increase the sequence length to 4k, model performance drastically drops. This is because in long sequences, the signal-to-noise ratio is low, such that the full attention may easily fit to the noise. From Figure 7 (right), we see that performance of Transformer with window attention increases when we increase the window size. Moreover, model performance is better with shorter sequences for the same window size. Such findings indicate that performance of window attention depends on the proportion of information within its perception.From Table 5 , we see that for the same sequence length, performance of SPADE increases when we increase the window size. Also, we see that performance of SPADE marginally decreases when we increase the sequence length from 4k to 6k. Recall from Figure 7 (left) that performance of Transformer with full attention drastically deteriorates when we increase the length from 3k to 4k. Such a result indicates that the proposed model is more suitable to model long sequences. 
conclusion and discussion: 
 We propose SPADE, a state space augmented Transformer model that targets long sequence modeling. SPADE is a multi-layer Transformer model, where the bottom layer is a global layer and the rest are local layers. In the global layer, we use a SSM to augment coarse global information, which are subsequently refined by the following local layers. SPADE is flexible to accommodate different SSMs such as S4 and S4 as the global information extractor. We instantiate the local layers with off-the-shelf efficient attention methods, such as sliding-window attention. The proposed model has near linear time and space complexity, facilitating it to handle long sequences. We conduct extensive experiments on the Long Range Arena (LRA) benchmark and language modeling datasets to demonstrate the effectiveness and efficiency of SPADE. We also pre-train encoder-decoder models to demonstrate the scalability of SPADE, and we perform fine-tuning experiments on natural language understanding (GLUE) and natural language generation (summarization) tasks. In all the experiments, SPADE exhibits superior performance.  
 
 Appendix: In Eq. 1, we have Q, K, V ∈ R L×d , such that computing the attention Attn(X) introduces O(L 2 ) time and space costs. Such quadratic costs are prohibitive when the sequence length L is large. There are various attempts to reduce the quadratic time and space complexity of the vanilla attention.One approach is to employ sparse attention. That is, each token only attends to a subset of all the tokens according to pre-defined patterns, e.g., neighboring tokens within a fixed size window. Some examples include Sparse Transformer (Child et al., 2019) , BlockBERT (Qiu et al., 2020) , Longformer (Beltagy et al., 2020) , ETC (Ainslie et al., 2020) , BigBird (Zaheer et al., 2020) , HEPOS (Huang et al., 2021) , and Poolingformer (Zhang et al., 2021) .Another approach is to use low-rank projection. For example, in Linformer (Wang et al., 2020b) , the attention mechanism in Eq. 1 becomes Attn(X) = softmax(Q(EK) ⊤ / √ d)(FV).Here, the two additional parameters satisfy E, F ∈ R r×L , where r is the projection rank such that r ≪ L. Similar methods include Nystr ömformer (Xiong et al., 2021) , Synthesizer (Tay et al., 2021a) , Transformer-LS (Zhu et al., 2021a) , and Luna (Ma et al., 2021) . However, these approaches face difficulty when handling causal tasks, such as auto-regressive language modeling. Specifically, in Eq. 1, we mask out the upper triangular part in the attention score matrix A ∈ R L×L such that each token can only attend to its previous tokens. However, this is implausible in Linformer since we project the L × L matrix to a L × r matrix.Kernel-based approaches can be used to approximate the full attention Attn(X). In these approaches, the quadratic-time softmax attention is replaced by fast linear-time kernel approximations (e.g., Gaussian and arc-cosine kernel). Some examples include Linear Transformer (Katharopoulos et al., 2020 ), Performer (Choromanski et al., 2021) , Random Feature Attention (Peng et al., 2021), and FMMformer (Nguyen et al., 2021) . Both low-rank projection and kernel-based approaches approximate the full attention, and thus, they often suffer from non-negligible approximation error.We can also adopt clustering-based approaches, where we divide Q or K into several clusters, and only perform inter-cluster attention. Such methods include Reformer (Kitaev et al., 2020) , Clusterformer (Wang et al., 2020a) , Sinkhorn Transformer (Tay et al., 2020) , Fast Transformer (Vyas et al., 2020) , Routing Transformer (Roy et al., 2021) , and FLASH (Hua et al., 2022) .Long Range Arena. We evaluate the effectiveness of the proposed model on Long Range Arena (LRA, Tay et al. 2021b) , which is a benchmark tailored for evaluating models' ability in modeling long sequences. The benchmark contains six tasks: ListOps, which tests the capability of modeling hierarchically structured data (Nangia & Bowman, 2018) ; byte-level text classification on the IMDB movie review dataset (Text, Maas et al. 2011 ); byte-level document retrieval on the ACL anthology network (Retrieval, Radev et al. 2013 ); pixel-level image classification on CIFAR-10 (Image, Krizhevsky et al. 2009) ; Pathfinder, which tests the capability in modeling spatial dependency (Linsley et al., 2018) ; and a longer version of Pathfinder (Path-X, Tay et al. 2021b) .For natural language understanding, we fine-tune the pre-trained models on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2019) , which is a collection of natural language understanding tasks. The benchmark includes two singlesentence classification tasks: CoLA (Warstadt et al., 2019) is a linguistic acceptability task; and SST-2 (Socher et al., 2013) is a binary classification task that classifies movie reviews to positive or negative. The benchmark also contains three similarity and paraphrase tasks: STS-B (Cer et al., 2017) is a text similarity task; MRPC (Dolan & Brockett, 2005 ) is a paraphrase detection task; and QQP is a duplication detection task. Additionally, there are natural language inference tasks: MNLI (Williams et al., 2018) ; QNLI (Rajpurkar et al., 2016) ; RTE (Dagan et al., 2006; Bar-Haim et al., 2006; Giampiccolo et al., 2007; Bentivogli et al., 2009) . Statistics of the GLUE benchmark is summarized in Table 7 .Summarization. For natural language generation, we fine-tune the pre-trained models on several abstractive summarization datasets. The sources and statistics are summarized in Table 6 . ⋄ SPADE base : We follow the pre-training settings in BERT (Devlin et al., 2019) . Specifically, we train the model on Wikipedia (Devlin et al., 2019) and BookCorpus (Zhu et al., 2015) .⋄ SPADE base++ : We follow the pre-training settings in RoBERTa (Liu et al., 2019b) . Specifically, we train the model on Wikipedia (Devlin et al., 2019) , BookCorpus (Zhu et al., 2015) , STORIES (Trinh & Le, 2018) , CC-News (Liu et al., 2019b) , and OpenWebText (Gokaslan et al., 2019) .For language model pre-training and fine-tuning experiments, we use Adam (Kingma & Ba, 2015) as the optimizer. Hyper-parameters for pre-training are detailed in Table 8 ; and hyper-parameters for fine-tuning are detailed in Table 9 .We follow the model architecture settings in Ma et al. 2022 . In all the experiments, we use Adam (Kingma & Ba, 2015) as the optimizer. We use a linear decay learning rate schedule.The rest of the hyper-parameters are detailed in Table 10 . https://github.com/microsoft/EfficientLongSequenceModelinghttps://github.com/namisan/mt-dnn 


Decision: The paper proposes a model architecture that combines a global layer based on SSM, followed by local layers using local attention (and MLPs) to model long sequences. The resulting architecture is evaluated on Long Range Arena (LRA), language modeling on wikitext-103, and an encoder-decoder model for natural language understanding and generation. The model shows better performance than Transformer and S4 on LRA, and slightly outperforms the standard Transformer on wikitext-103 and T5 on GLUE.
+ The model architecture is simple and the presentation is clear so it is easy to understand.
+ Strong performance on LRA.
+ Thorough analysis and ablations to justify the model architecture design.
- Language modeling evaluation on wikitext-103 and GLUE are relatively small datasets, where regularization may play a large role, compared to large datasets currently being used to train LLMs.
Overall the paper is a valuable contribution to the design of hybrid models, incorporating local attention and SSMs.

Reviews: review: paper_summary: paper_summary: This paper proposes SPADE, a hybrid architecture that consists of a global SSM layer followed by local attention layers. A variety of experiments are conducted with the architecture on long range tasks as well as standard language modeling tasks with the goal of demonstrating better efficiency with similar or better performance as compared to either SSMs or Transformers alone.
strength_weakness: strength_weakness: 1. There is a issue with discussion of recent related work. Both Based [1] and Griffin [2] went on arxiv with similar ideas before the colm deadline. They are close enough that they may be counted as concurrent, but they should definitely be addressed. There is also no reference or comparison to more recent state space models from the last year like mamba [3], all the SSMs referenced are a few years old. 

2. I am also worried about the choice and strength of baselines. First, there is no comparison to a flash-attention based transformer baseline, which resolves a lot of issues with transformers in terms of efficiency, especially at the <10k sequence lengths considered in the paper. I am not totally convinced that the proposed architecture would give substantial efficiency gains at these lengths and model sizes. Second, especially for the language modeling tasks there is a lack of recent baselines. 

3. I also would have liked to see experiments with decoder only causal models. Can the authors elaborate on why they chose to focus on encoder-decoder models?

[1] https://arxiv.org/abs/2402.18668

[2] https://arxiv.org/abs/2402.19427

[3] https://arxiv.org/abs/2312.00752

[4] https://arxiv.org/abs/2205.14135
paper_summary: paper_summary: In this paper, authors proposes the SPADE, aiming to combine SSMs with local attention to achieve a good tradeoff between long context understanding (capturing the global information) and computational complexity.  Comprehensive experiments are conducted to show the effectiveness of their methods.
strength_weakness: strength_weakness: 1. Can we say more about the philosophy to design this architecture. for example, why not add another global layers after some local layers, .e.g, global1 -> local1->global2->local2?
2. can we list the number parameters and latency for models listed in Table 1?
paper_summary: paper_summary: This paper focuses on improving the computational efficiency of Transformers. Designing an equally expressive model with a much more efficient handling of long context can have a great practical impact. 

Specifically, the authors propose to employ a state-space model at the bottom layer of a Transformer to capture long-range and positional information coupled with efficient local attention for the rest of the layers. The method is well-motivated and clearly described; its main novelty is in designing a hybrid model that address the limitations of individual Transformers and SSMs.

The evaluation is comprehensive, covering a diverse set of tasks including long-range tasks, language modeling, and language model pretraining followed by finetuning on NLU and NLG tasks. Comparisons against strong baselines such as S4 and MEGA demonstrate the effectiveness of the approach. One caveat may be that the experimental comparison with different baseline methods is not with the exact same experimental setup and resources (i.e. head-to-head).

Overall, the paper is clearly written and easy to follow. It is also positioned well with respect to prior work and provides sufficient details and artifacts for reproducibility that can enable future work.
strength_weakness: strength_weakness: - A direct head-to-head comparison with other efficient architectures is missing. This means running the models instead of comparing to the scores from the corresponding papers with comparable configuration.
- The positioning with respect to prior work needs some effort, especially compared to methods that combine global and local information.
questions: questions: - Is it possible to convert an existing Transformer model into a SPADE architecture instead of training it from scratch? That would be useful for impacting already trained models.
paper_summary: paper_summary: This work aims to reduce the quadratic computational costs of Transformer models by replacing the first layer by a hybrid layer which combines a state space model with a local attention module. 
In all other layers, the full attention module is replaced by a local attention module. 
Experiments have been conducted on different tasks leveraging different architectures and model sizes.
Results show consistent improvements over different baselines.
strength_weakness: strength_weakness: - Missing implementation details for the encoder-decoder architecture
- No larger scale experiments (but claiming to “demonstrate the scalability of the proposed method”)
- Related work in the appendix
questions: questions: - While appreciating the different model architectures, it would have been interesting to see results with a larger decoder-only model and/or results with SPADE_{large,xl}++
- MEGA also works for the machine translation task; since you trained an encoder-decoder model with your approach, results on machine translation would be interesting, too; in particular, since Vardasbi et al. (EAMT 2023) reported mixed results when employing S4 on this task
- Please move the related work section into the main paper

Comments: title: Reply to Rebuttal
Reply to Rebuttal
Reply to Rebuttal
comment: Thanks for the clarification and pointing me again to Section 6.3, I will raise my score to 6!
Thank you for the suggestions. We would like to address the following:

* The computation cost of the cross-attention layer depends on both the input (source) sequence length and output (target) sequence length, i.e., the attention matrix is not a square matrix. Because the output sequence length is usually short compared with the input sequence length, the computation cost of the cross-attention is only a fraction of the computation cost of the input self-attention. For example, the input length can be up to 16k in the experiments, while the output sequence length is only 128. That is, for the input self-attention, the size of the attention matrix is (16k * 16k), while for the cross-attention, the size of the attention matrix is only (16k * 128).

* Here are some results of S4+Full-attention. We can see that performance of S4+Full-attention is only slightly better than S4+Local-attention, but the computation cost is much higher.

  |  | Listops | Text | Retrieval |
  |---|---|---|---|
  | S4+Full | 60.75 | 90.81 | 91.20 |
  | S4+chunk | 60.50 | 90.69 | 91.17 |


* The ablation experiments of the position and number of attention modules are in Section 6.3 (Figure 6). By design, we consider a fixed-size (local attention) Transformer model and add SSM layers to the model. In Section 6.3, we illustrate why SSM layers have to be added to the bottom of the model, and we also demonstrate that only one SSM layer suffices to extract useful global information.
Thanks for your reply. I overall think this is a good work and will keep my score unchanged.
Thank you for the replies and the additional effort. Overall, I am positive about the work but I still have a concern regarding comparison to other models.

Please explicitly mention the re-implementation point and effort involved about its correctness because it's not clear in the paper.  It'd be also useful to describe in detail the model differences and provide information that shows that the experimental setups are comparable in terms of hyper-parameter budget and model size. 

Note that "head-to-head" comparison refers to using the exact same architecture and varying only a few components e.g. the global attention layer (this can be softmax attention, RNN attention, or any other efficient attention from the literature).
Thanks for the response and for adding the related work. I will raise my score to a 6 to reflect this change. Since the other comparisons have not been completed yet, I will not raise further for now.
Thanks for your rebuttal! 

- Cross-Attention: Can you be more specific on "output length is usually small" and "does not introduce much computation latency"?
- "S4+full-attention is not significantly better than S4+local-attention": Please provide details of your applied significance test.
- An ablation of the positions and number of different attention modules is still needed to verify the design choice.