Simple Local Attentions Remain Competitive for Long-Context Tasks (2022)

Abstract: Many NLP tasks require processing long contexts beyond the length limit of pretrained models. In order to scale these models to longer text sequences, many efficient long-range attention variants have been proposed. Despite the abundance of research along this direction, it is still difficult to gauge the relative effectiveness of these models in practical use cases, e.g., if we apply these models following the pretrain-and-finetune paradigm. In this work, we aim to conduct a thorough analysis of these emerging models with large-scale and controlled experiments. For each attention variant, we pretrain large-size models using the same long-doc corpus and then finetune these models for real-world long-context tasks. Our findings reveal pitfalls of an existing widely-used long-range benchmark and show none of the tested efficient attentions can beat a simple local window attention under standard pretraining paradigms. Further analysis on local attention variants suggests that even the commonly used attention-window overlap is not necessary to achieve good downstream results — using disjoint local attentions, we are able to build a simpler and more efficient long-doc QA model that matches the performance of Longformer with half of its pretraining compute.

## introduction the quadratic complexity of transformer architectures makes it prohibitive to apply large stateof-the-art pretrained models to full-length documents. to efficiently handle longer text while still maintaining the capacity of attention-based models, a long list of efficient attention variants have been proposed and many claim to effectively capture long-range dependencies. typical paradigms of these architecture innovations involve learnable sparse attention patterns @xcite @xcite , fixed local patterns @xcite @xcite and attention matrix approximation methods @xcite @xcite . while most of these studies have reported numbers on long sequence inputs, they tend to adopt quite different benchmarks. for instance, reformer @xcite ) is tested on the 64k-chunk en-wik8 dataset for unidirectional language modeling; performer @xcite reports masked language modeling (mlm) perplexity on the pg-19 book corpus and protein sequences; linformer @xcite reports mlp perplexity with various input length, while most of the documents in their pretrain corpus are short documents. 2 the divergence of evaluation protocols makes it hard to compare the relative performance of each attention variant and it is also unknown how they perform well in more practical use cases, which typically involve large-scale pretraining and downstream finetuning. other lines of work such as longformer @xcite and etc @xcite conduct experiments on real-world long-context tasks such as long document qa and summarization. these methods only test fixed local attention patterns, i.e., each token can only attend to a small set of nearby tokens. to reduce the pretraining cost, these models are all initialized from the roberta @xcite ) checkpoint 3 before further long-doc pretraining. while this paradigm is useful to achieve strong downstream performance, it is not ideal for a fair comparison of all available attention mechanisms, since some of the models use different parametrization that is incompatible with the vanilla transformer attention. a recently proposed benchmark @xcite , named long-range arena (lra), aims to address the lack of unified evaluation with a bundle of long-sequence tasks. however, the textrelated tasks in this benchmark are either automatically generated or artificially lengthened by enforcing byte-level inputs, making them rather synthetic. with a fixed byte-level vocabulary and prespecified model size, all models are trained from scratch with the same epoch limit on each dataset. while the evaluation protocol is consistent across architectures, this setup still deviates from the common paradigm of applying transformer models, i.e., standard tokenization like bpe or wordpiece, large-scale pretraining followed and task-specific finetuning @xcite . thus, an important question yet to be addressed is whether the results on these artificial datasets are indicative of real-world long-context tasks. in this work, our goal is to better understand the effectiveness of various attention mechanisms through a systematic study on practical long-context tasks. instead of only relying on language modeling or synthetic tasks, we test each model under the standard pretraining-andfinetuning paradigm. for a fair comparison, we implement these attentions under a unified framework and test them using the same transformer architecture foot_3 used by roberta-large. we pretrain all models using a large corpus that contains mostly long documents and then finetune them on tasks like long-document question answering, full document retrieval, and text classification. our experiments show the discrepancies between the commonly used lra benchmark and downstream results (after pretraining). additionally, our analysis on the best local attention models allows us to further simplify these models and results in a more efficient long-context encoder. more specifically, the key findings of this paper include: ## preliminaries of tested attention variants we study three classes of efficient attentions: fixed local patterns. these methods restrict each token to only attend a local window of tokens. the long-range interactions are achieved by the depth of the model. we consider two variants of these models, the token-wise local window attention (local window) proposed in @xcite where each token attends to the same number of tokens on each side, and a simplified and easy-to-implement blockwise version (blockwise lw) @xcite where each token attends to tokens in the same block and half of the tokens in the left/right blocks. a visualization comparing these two models is shown in figure 1 . learnable sparse attention patterns. instead of relying on the inductive bias of locality, methods like reformer @xcite and sinkhorn attention @xcite allow the model to adaptively select tokens to attend to. briefly, reformer uses a learnable hashing function to bucket the sequence and each token only attends to tokens in the same bucket; sinkhorn uses a learnable sorting function to learn a permutation of the segments and each token will attend to tokens in its segment, and the corresponding segment after permutation. kernel-based/low-rank methods. this class of methods use matrix approximation methods to approximate the full attention function. for sequence length l and the hidden dimension d, linformer @xcite simply uses a projection matrix (l × k) to reduce the length of key and value feature matrix, i.e., from l × d to k × d (k l). nyström @xcite attention adopts a classic matrix approximation method which reconstructs the full attention matrix using a sampled sub-matrix. performer @xcite eliminates the need of explicitly calculating the l × l attention matrix by using a random feature method that can approximate the softmax kernel with only dot-product operations. hybrid attention. in addition to these representative methods in each class, our study also includes the more recent long-short attention @xcite which has a similar compression component as in linformer and combines it with local attentions. unlike linformer's compression component which is simply implemented as a standalone projection matrix, long-short proposes an inputdependent compression layer, which can adaptively reduce the sequence length. a note on global tokens. for many practical nlp tasks, e.g., classification or entailment, the final layer of the model usually requires a single sequence-level representation as input. for local attention models, it is common practice @xcite to mark a single or a small number of tokens as global tokens and allow these tokens to attend to and be attended by all other tokens. without incurring much computational cost, these global tokens are important to get better sequence representations and achieve good downstream results. while the mechanism of global tokens has not been used in models with learnable attention patterns, it is straightforward to augment reformer and sinkhorn with global tokens using gather operations in standard neural network packages, as their attention scores are still calculated by dot product and softmax operations. thus, in our experiments, except for the kernel-based/low-rank methods, we augment all other models with global tokens to offset the potential performance gap resulting from this trick. ## experiment setup we restrict our studies to encoder-only models and leave the analysis of generative models to future work. we begin by implementing a collection of efficient attentions with a unified framework @xcite , which allows us to plug these models into our pretraining-andfinetuning pipeline in a consistent fashion. ## conclusion we present a systematic study of recent proposed efficient attention variants on real-world long-context nlp tasks. in contrast to existing work, we are the first to test these models with a set of unified and large-scale experiments. our results reveal the gap between a widely used benchmark and practical downstream tasks after conducting largescale pretraining. among all the studied attention methods, we find that the simplest local attentions outperform other complex attention paradigms on downstream tasks. we also show that existing localattention models can be further simplified by removing the attention-window overlap, resulting in a faster model that achieves similar or better results. importantly, our work calls for more careful and practical evaluation protocols while developing long-context nlp models.