Efficient Long-Range Transformers: You Need to Attend More, but Not Necessarily at Every Layer (2023)

Abstract: Pretrained transformer models have demonstrated remarkable performance across various natural language processing tasks. These models leverage the attention mechanism to capture long- and short-range dependencies in the sequence. However, the (full) attention mechanism incurs high computational cost – quadratic in the sequence length, which is not affordable in tasks with long sequences, e.g., inputs with 8k tokens. Although sparse attention can be used to improve computational efficiency, as suggested in existing work, it has limited modeling capacity and often fails to capture complicated dependencies in long sequences. To tackle this challenge, we propose MASFormer, an easy-to-implement transformer variant with mixed attention spans. Specifically, MASFormer is equipped with full attention to capture long-range dependencies, but only at a small number of layers. For the remaining layers, MASformer only employs sparse attention to capture short-range dependencies. Our experiments on natural language modeling and generation tasks show that a decoder-only MASFormer model of 1.3B parameters can achieve competitive performance to vanilla transformers with full attention while significantly reducing computational cost (up to 75%). Additionally, we investigate the effectiveness of continual training with long sequence data and how sequence length impacts downstream generation performance, which may be of independent interest.

## introduction pre-trained transformer models have manifested superior performance in various natural language processing tasks such as natural language modeling (nlm) @xcite , natural language generation (nlg) @xcite and natural language understanding (nlu) @xcite @xcite . these models leverage the attention mechanism @xcite to compute the dependency score for each pair of tokens in an input sequence. some practical tasks require these transformer models to handle long-sequence inputs like 8k tokens. for example, chatbot systems gather longterm contexts of user interactions to generate informative texts @xcite . summarization for news, government reports, and academic papers request models to take inputs of long sequences to generate comprehensive summaries @xcite , otherwise models often miss important information. note that typical transformer models apply full attention to capture token dependencies pair-wise. it leads to a quadratic time and space complexity w.r.t. input length. however, such a complexity is prohibitive for long sequences. in particular, it incurs massive memory consumption during the back propagation. for example, a transformer model with 250m parameters consumes over 80g gpu memory when sequence length is 8k @xcite . to address this scalability issue, various approaches have been proposed to reduce the complexity. one approach is sparse attention, which restricts each token to attend a subset of tokens based on predefined sparsity patterns @xcite @xcite . for instance, block sparse attention @xcite divides the input sequence into several blocks, and only intra-block attention is performed. besides, sliding-window attention @xcite @xcite allows each token to attend to its neighboring tokens within a sliding window. these methods, though reducing the complexity of full attention, cannot sufficiently capture long-range dependencies. other variants, such as kernel approximation @xcite and low-rank approximation @xcite methods, share the similar spirit and drawbacks. to com-pensate for the lack of long-range dependencies, longt5 @xcite introduces global tokens that are obtained by average pooling on every block of tokens @xcite . however, the block pooling operations can weaken the signal of crucial tokens and prevent the long-range dependencies from being detected. in addition to these methods, state space models (ssms) prespecify global dependency patterns to capture the long-range dependencies only @xcite @xcite @xcite . these models can be regarded as linear recurrent neural networks with specifically designed fixed weights. as tailored for global dependencies, ssms fail to effectively capture local dependencies. in order to combine both local and global dependencies, spade @xcite and mega @xcite augment ssm layers into transformer layers equipped with local attention. however, state space methods require sophisticated implementation, and often encounter computational instability during the back propagation, especially when scaling up to large model size @xcite . spade and mega hence inherit these drawbacks. note that the aforementioned methods apply same attention mechanism for every layer. we challenge this conventional wisdom and propose a transformer variant -masformer (mixed attention span transformer). masformer utilizes full attention only at a subset of layers whereas employs sparse attention at the remaining layers. our design is motivated by the phenomenon -that most contexts in nlp data display a great deal of locality of reference @xcite . that is, most of information about a token can be derived from its neighboring tokens. in contrast, long-range dependencies among tokens are sparse and infrequent. consider an academic paper as an example. within a paragraph, there exist numerous short-term dependencies. neighboring tokens are closely connected to convey meaningful semantics. across paragraphs, there can be a small number of long-range dependencies. for example, tokens associated to the primary theme of the paper exhibit rare and weak dependencies across a long span. since long-range dependencies occur much less frequently, a few layers of full attention are adequate to capture them. in stark contrast, short-term dependencies are more frequent, necessitating local attention in the majority of layers to fully extract these signals. to demonstrate the effectiveness of masformer, we conduct experiments on natural language modeling (arxiv and pubmed @xcite ) and natural language generation (arxiv, @xcite and scrolls, @xcite ) tasks. specifically, we compare the performance of masformer to other attention methods using a pretrained gpt-2 model @xcite of 1.3 billion parameters. our empirical results demonstrate that masformer consistently outperforms baseline methods across different attention cost (i.e. the total number of computed attention scores). in particular, masformer can achieve comparable performance to full attention while significantly reducing the computational cost. for example, with 27% of its attention cost, masformer achieves a close r2 score as full attention on qmsum dataset. we also make additional discoveries with mas-former, which are of independent interest. firstly, we investigate the effectiveness of continual training for long sequence modeling. many publicly available models are pre-trained with sequences shorter than 2048, and often fail to perform well on longer sequences (e.g. 8k/16k tokens). to bridge the gap, we explore the option of continual training to adapt these models to long sequences, thereby avoiding pre-training from the scratch. we discuss its effectiveness with masformer in section 4.3. secondly, we showcase that increasing sequence length can yield more performance gains on downstream tasks than nlm tasks evaluated by perplexity. we are aware of the recent findings by @xcite that increasing context length exhibits limited impact on nlm perplexity. nevertheless, when applying masformer to downstream tasks like long-context summarization, we find that model performance benefits significantly from extending context length. such a difference arises from the fact that predicting the next tokens in nlm primarily relies on locality of reference. capturing infrequent long-range tokens can improve perplexity but not significantly. therefore, we emphasize the necessity to evaluate model performance on downstream tasks that require longrange dependencies. furthermore, our empirical evidence suggests that increasing the length can improve the performance only if models possess sufficient capability to handle additional long-range information. local attention, as a counterexample, often fails to capture long-range signals and hence benefits much less from long sequences. ## our approach we present our method -masformer, a longrange transformer variant that mixes different attention spans across layers. ## experiments we evaluate the effectiveness and efficiency of masformer on natural language modeling (arxiv and pubmed, @xcite al. (2018) ), natural language generation (arxiv @xcite , qmsum and govreport @xcite ). we choose the gpt-3 xl model architecture @xcite as our base model, which consists of 1.3 billion parameters and 24 layers and is pre-trained on pile @xcite for 300 billion tokens. gpt is a general purpose model that can be applied to many tasks instead of tailoring them for specific tasks. as such, it makes easy to control experiments and showcase the difference among various methods. implementation details. our base model uses absolute positional embeddings with maximum length 1024. to accommodate longer inputs, we duplicate its positional embeddings to have the maximum length as 8192 such that the model can handle sequences containing up to 8192 tokens. then, we implement different attention methods by modifying the attention pattern of the base model. we implement all the models with pytorch @xcite . all the experiments are conducted on nvidia a100 gpus. continual training details. after changing the attention pattern, we conduct the continual training for masformer and baseline methods on pile corpus @xcite to adapt the revised models to new attention patterns and long-sequence inputs. we leverage the casual language modeling (clm) objective to train the model for 50,000 steps with a warmup of 2000 steps. we set the input length as 8192 and use a batch size of 128 such that the models are optimized with 1m tokens per step. we use the constant learning 0.0001 for all methods. baseline. we compare masformer with the following methods: in the following experiments, we compare mas-former and the baseline methods across different attention cost c. that is, for all block sparse attention, we set the block size as @xmath0. for all sliding-window attention, we choose the window size as @xmath1. for masformer, we apply a small block size @xmath2. then we observe how their performance evolves when enhancing the attention cost c or input length n. experiment overview. we briefly summarize the experimental contents as follows: ## discussion gpt-neo @xcite introduces an attention pattern that alternates full and window attention. however, this models is not tailored for long sequences. it sets the local window size as 256 and has the maximum input length as 2048, unable to handle long sequences. instead, this attention pattern is applied heuristically in an attempt to reduce computational cost. however, as discussed in section 4.3.3, this approach is neither effective nor efficient as masformer when handling long sequences. as shown in in contrast, masformer presents an effective solution for efficient long-sequence modeling. it provides guidance on adapting existing pre-trained transformers to long inputs. meanwhile, it provides insights for designing large long-range models, especially for deeper models. by equipping only a subset of bottom layer with full attention, we can substantially mitigate computational cost. additionaly, the computation of masformer can be further optimized by leveraging system-level acceleration techniques (e.g., flashattention and xformer) that support both block and full attention. ## conclusion we propose an efficient long-range transformer -masformer that utilizes full attention at a few of bottom layers and employs sparse attention at the remaining layers. our empirical results on natural language modeling and generation tasks demonstrate that masformer can address the trade-off between computational cost and performance gains in a more efficient and effective way.