Simple Local Attentions Remain Competitive for Long-Context Tasks
Meta AI
Abstract
Many NLP tasks require processing long contexts beyond the length limit of existing pretrained models.
To scale these models to longer text sequences, many efﬁcient longrange attention variants have been recently proposed.
Despite the abundance of research in this direction, it is difﬁcult to gauge the relative effectiveness of these models in practical use cases, e.g. , if we apply these models following the pretrain-and-ﬁnetune paradigm.
In this work, we aim to conduct a thorough analysis of these emerging models with largescale and controlled experiments.
For each attention variant, we pretrain large-size models using the same long-doc corpus and then ﬁnetune these models for real-world long-context tasks.
Our ﬁndings reveal pitfalls of a widelyused long-range benchmark and show that the other efﬁcient attentions fail to outperform the simple local-window attention after standard pretraining.
Further analysis of localattention 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 efﬁcient long-doc
QA model that matches the performance of
Longformer (Beltagy et al., 2020) 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 efﬁciently handle longer text while still maintaining the capacity of attention-based models, a long list of efﬁcient 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 (Kitaev et al., 2020; Tay et al., 2020; Roy et al., 2021), ﬁxed local patterns (Beltagy et al., 2020; Ainslie et al., 2020; Zaheer et al., 2020) and
attention matrix approximation methods (Wang et al., 2020; Choromanski et al., 2021; Xiong et al., 2021).
While most of these studies have reported numbers on long sequence inputs, they tend to adopt quite different benchmarks.
For instance, Reformer (Kitaev et al., 2020) is tested on the 64k-chunk enwik8 dataset for unidirectional language modeling; Performer (Choromanski et al., 2021)
reports masked language modeling (MLM) perplexity on the PG-19 book corpus and protein sequences; Linformer (Wang et al., 2020) reports MLP perplexity with various input length, while most of the documents in their pretrain corpus are short documents.
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 ﬁnetuning.
Other lines of work such as Longformer (Beltagy et al., 2020) and ETC (Ainslie et al., 2020) conduct experiments on real-world long-context tasks such as long document QA and summarization.
These methods only test ﬁxed 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 (Liu et al., 2019) 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 (Tay et al., 2021), named long-range arena (LRA), aims to address the lack of uniﬁed evaluation with a bundle of long-sequence tasks.
However, the textrelated tasks in this benchmark are either automatically generated or artiﬁcially lengthened by enforcing byte-level inputs, making them rather synthetic.
With a ﬁxed byte-level vocabulary and prespeciﬁed 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-speciﬁc ﬁnetuning (Devlin et al., 2019).
Thus, an important question yet to be addressed is whether the results on these artiﬁcial 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-andﬁnetuning paradigm.
For a fair comparison, we implement these attentions under a uniﬁed framework and test them using the same Transformer architecture 4 used by RoBERTa-large.
We pretrain all models using a large corpus that contains mostly long documents and then ﬁnetune them on tasks like long-document question answering, full document retrieval, and text classiﬁcation.
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 efﬁcient long-context encoder.
More speciﬁcally, the key ﬁndings of this paper include:
• With proper tuning, we ﬁnd that all the tested models can achieve similar level of performance on the LRA benchmark while their performance diverges signiﬁcantly on large-scale pretraining and downstream tasks; • In our experiments, the other attention paradigms barely outperform the class of sim-
[Figure]
ple local attentions on downstream tasks when using similar
compute
• As a result of our further analysis of the best performing attention variants, we are able to build a long-doc QA model that is on-par with Longformer while being 2x more efﬁcient.
Preliminaries of Tested Attention Variants
We study three classes of efﬁcient 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 Beltagy et al.
(2020) where each token attends to the same number of tokens on each side, and a simpliﬁed and easy-to-implement blockwise version ( Blockwise
LW )
(Zaheer et al., 2020) 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 (Kitaev et al., 2020) and Sinkhorn
Attention (Tay et al., 2020) allow the model to adaptively select tokens to attend to.
Brieﬂy, 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 (Wang et al., 2020) 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 (cid:28) L ).
Nyström (Xiong et al., 2021) attention adopts a classic matrix approximation method which reconstructs the full attention matrix using a sampled sub-matrix.
Performer (Choromanski et al., 2021) 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 (Zhu et al., 2021) 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. , classiﬁcation or entailment, the ﬁnal layer of the model usually requires a single sequence-level representation as input.
For local attention models, it is common practice (Beltagy et al., 2020; Zaheer et al., 2020) 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 efﬁcient attentions with a uniﬁed framework (Lefaudeux et al., 2021), which allows us to plug these models into our pretraining-andﬁnetuning pipeline in a consistent fashion.
3. LRA Experiments
Following recent work on efﬁcient long-range attentions, we take the LRA benchmark as our ﬁrst set of experiments.
As our focus here is on NLP tasks, we consider a subset of LRA tasks with text inputs, i.e. , the ListOps, IMDB sentiment analysis, and text matching tasks.
All tasks are formulated as classiﬁcation problems: ListOps requires the model to predict the correct output of an expression (10way classiﬁcation), sentiment analysis is to predict the positive/negative labels of IMDB reviews and text matching aims to predict citation links between papers.
We follow the hyperparameter settings of recent work (Xiong et al., 2021; Zhu et al., 2021).
Two-layer Transformer encoders are used across all tasks and enough training updates are allowed to ensure convergence 5 .
Note that this is different from the setup proposed in the original LRA benchmark, where different tasks adopt different model sizes.
It is observed from recent work that two-layer models with smaller dimensions are sufﬁcient to achieve similar or better results than previously reported results.
The ﬁnal classiﬁcation layer is added on top of the representations of [CLS] tokens which are prepended to each sequence.
3. Pretraining and Downstream Tasks
For practical NLP application, large-scale selfsupervised training has become an indispensable ingredient to fully unlock the power of Transformer models.
In terms of the experiment scale and testing settings, there is a clear gap between LRA’s setup and how we apply state-of-the-art Transformer models in practice.
For the second set of experiments, we aim to test these models at scale and investigate whether the results on the LRA benchmark are accurate indicators for real-world long-context tasks after standard large-scale pretraining and ﬁnetuning.
Pretraining Resource.
Following Beltagy et al.
(2020), we compile a corpus that contains mostly long documents, including Stories (Trinh and Le, 2018), RealNews (Zellers et al., 2019), Books corpus (Zhu et al., 2015) and English Wikipedia.
To make the experiments manageable and relevant for standard GPU hardware, we restrict each model’s memory usage close to the 16GB threshold when taking 4,096 tokens in each training batch.
We control the batch size and training update across all models: we use a batch size of 256 sequences ( 2 20 tokens) and pretrain each model using the standard masked language modeling objective for 100k updates.
We ﬁnd that all models’ training curves almost stabilize after this amount of training steps.
We use 32 A100 GPUs for pretraining and all model runs are ﬁnished within around 2 days.
Pretraining Architecture
In contrast to Longformer (Beltagy et al., 2020) and Bigbird (Zaheer et al., 2020) where the models are initialized from
RoBERTa before pretraining on long documents, we pretrain these models from scratch, as our goal here is to ensure fair comparison and not all architectures can reuse weights from a standard transformer model.
In particular, Nyström and Performer do not use the standard dot-product and softmax to compute attention probabilities, making their parameters not compatible with common models like RoBERTa or BERT.
Furthermore, other models like Linformer or LongShort introduce additional parameters inside the attention module.
In our initial experiments, we observe initializing from the RoBERTa put these models at a signiﬁcant disadvantage compared to other models (e.g., local window attention) that are more compatible with vanilla transformers.
Apart from the expanded position embedding matrix and the attention blocks, the architecture hyperparameters are consistent with RoBERTa-large.
For both LRA and the large-scale experiments, we adopt the pre layer-normalization trick (Xiong et al., 2020) for feedforward and attention blocks.
This usually results in better performance in LRA and turns out to be essential for several models in the pretraining experiments.
6 See additional model-speciﬁc architecture settings and models’ average memory usage in the Appendix.
Downstream Datasets and Metrics.
We consider practical tasks that naturally involve long documents.
We test extractive QA over long documents, long document classiﬁcation, and document retrieval.
For the ﬁrst two tasks, we use TriviaQA and Hyperpartisan classiﬁcation respectively, both of which have been used in existing long Transformer work (Beltagy et al., 2020).
For full document retrieval, we construct the dataset based on recent open-domain QA work (Lee et al., 2019) that uses passage-level retrievers.
We take an existing passage corpus from Karpukhin et al.
and reconstruct the document-level corpus.
We consider a document to be positive if it includes the answer passage.
We reported token-level answer exact match and F1 for extractive QA and the classiﬁcation accuracy for Hyperpartisan.
For the retrieval task, for the ease of experiments, we reported the mean reciprocal rank on the dev set 7 , which has been shown to correlate well with ﬁnal retrieval metric like answer recall (Oguz et al., 2021).
We conduct grid search for all tasks and report the best dev results.
Given the small size of the Hyperpartisan dataset, we reported averaged results from 4 random seeds.
Task-speciﬁc
Architectures for
Finetuning.
We use standard architectures for the ﬁnetuning tasks: for extractive QA, a single-layer MLP span predictor is added on top of the output token representations; the classiﬁcation task uses a binary MLP classiﬁer that takes the [CLS] vector as input.
For retrieval, we share the query and document encoder using our pretrained models and use dot-product of the [CLS] vectors as the similarity score.
For models that are compatible with global tokens, we use all the question tokens as global tokens in the QA task and use a single global token at the start of the sequences for both classiﬁcation and retrieval.
Except for the
Hyperpartisan dataset, the document lengths of the other two datasets usually exceed 4,096 tokens after tokenization.
In these cases, we drop the tokens outside the models’ position range.
We put further implementation details and each task’s length statistics in the Appendix.
Results and Analysis 4. Models Perform Similarly in LRA
We report our reimplemented LRA results in Table 1.
While previous work (Tay et al., 2021) has
[Table]
shown a clear performance gap between different models
[Table]
we ﬁnd that with proper tuning, the results of several models could be signiﬁcantly improved, ( e.g. , Sinkhorn, Linformer, Reformer, Performer) and there is no signiﬁcant performance gap between any of the models when using a similar level of compute (measure by FLOPS).
It is worth noting that these improved results are not obtained by increasing the complexity of models ( e [Table] g. , by using larger bucket size in Sinkhorn), as our implementation either uses similar or smaller size models compared to existing work.
Also note that while the single global token we added to Sinkhorn and LSH might be essential for some performance gains, it only brings trivial computation overhead.
4. Pretraining and Downstream Tasks
We now evaluate these models on practical benchmarks that involve real-world long documents.
As shown in Table 2, after we scale up the experiments and control the memory consumption of each model, we see more clear differences between these models than what we observe in LRA.
Clearly, ﬁxed local attentions remain to be strong baselines.
However, in contrast to LRA, we observe local attentions are signiﬁcantly better than the other attention variants, for both pretraining perplexity and downstream task results.
The only exception in terms of the pretraining perplexity is the hybrid Long-Short attention, which already integrates a local attention component: it achieves better perplexity than ﬁxed local attentions, but the downstream results are at most on par with much simpler models like Blockwise LW.
It is worth noting that while we only control the training updates and memory usage in Table 2, the conclusion still holds if we control the training time of each model:
We compare the training perplexity of Blockwise
LW attention and other faster models with ﬁxed training time in Table 3.
Even though our LRA experiments also study tasks with text inputs, we see clear discrepancies between the two sets of experiments.
Apart from models with ﬁxed local attention patterns, improvements on these text LRA tasks often do not transfer to the standard scaled pretraining-ﬁnetuning experiments.
For instance, while Performer can outperform most of the non-local attention methods on LRA, it performs poorly on both large-scale
MLM and downstream long-context tasks.
Similarly, while Nyström is signiﬁcantly better than LSH in LRA on average, we observe the opposite trend in Table 2.
Among the three tasks, only ListOps is loosely aligned with the MLM perplexity.
However, the gaps between each model on this task are still too narrow to be indicative.
Given that large-scale pretraining has become the gold-standard paradigm to build state-of-the-art NLP models.
Our ﬁndings here call for a more careful and reliable evaluation of lots of existing and emerging long-range attentions.
On the other hand, our results also reveal that the local context [Table] MLM Pretraining Downstream
[Table] [Table] might still be highly essential even in long context tasks. In the following section, we conduct further analysis on local attention models and attempt to identify the key ingredients of building strong NLP models for downstream long-context tasks.
4. Analysis on Local Attentions
As we have seen in §4.2, models that compute exact attention for local contexts around each token achieve better results.
Moreover, the Blockwise
LW variant performs the best even it does not guarantee a balanced left and right context window for each token.
Given these intriguing ﬁndings, we aim to investigate the following questions: How effective are the long-range mechanism in local attention models?
and Whether the studied longcontext tasks still mostly rely on locality bias?
Ablation Study.
In the Blockwise LW model, there are two mechanisms that enable long-range connections: the global tokens and the attention window overlap, i.e. , each token will additionally attend to half the tokens in the neighboring blocks, and the receptive ﬁeld increases with model depth.
While both are adopted as common practice in existing work (Zaheer et al., 2020; Beltagy et al., 2020), we study the isolated effect of each component in both pretraining and ﬁnetuning experiments.
For the non-overlap variant, we increase the block size by a factor of 2 such that the amount of tokens each token attends to remains the same.
We show the results in Table 4.
Surprisingly, we see different stories in terms of MLM pretraining and downstream tasks.
While both mechanisms are useful for achieving lower MLM perplexity, only the global-token mechanism seems important for downstream tasks.
Note that in the document retrieval tasks, removing both mechanisms results in slightly better performance.
Now the model is only able to use the ﬁrst block of the whole document for retrieval.
While this seems to suggest that this task is highly local and involves strong positional bias 8 , the gap might be too trivial to be conclusive.
Additionally, we only use a single global token for this task, it is likely that assigning more global tokens, e.g. , at passage boundaries, could bring additional improvements.
Investigating the particular task further is beyond scope of this work.
In terms of the effect of attention-window overlap, it is expected that this scheme is crucial for lower perplexity: it not only enables more distant dependencies but also reduces the number of "boundary tokens" which can only attend to one side of the context.
However, it is counter-intuitive that the [Table] overlapping attention links between neighboring blocks, which adds more long-range information, result in worse downstream performance.
Also, note that this observation is consistent for all the tasks we studied.
There are two possible implications of this ﬁnding: 1) the tested tasks still highly depend on locality bias, i.e. , most of the important information can be captured solely from the local bias, or 2) the overlapping scheme is not effective at capturing the long-range dependency in downstream tasks.
To conﬁrm either hypothesis, we conduct another set of experiments with models that have access to different sizes of context.
On Locality Bias.
We take the non-overlapping variant and experiment with various block sizes to see whether longer context is important to studied tasks.
We show the results in Table 5 and the pretraining curves in Figure 2.
While the long-range connections brought by the attention overlap is not helpful for downstream results, we see that increasing the local block sizes does consistently improve both pretraining and downstream performance although the improvement becomes modest beyond block size 256.
It is also interesting that the models with smaller block sizes converge faster at the early stage of pretraining.
This suggests a staged pretraining process might be more efﬁcient than directly training from long sequences, which aligns with Press et al.
(2021)’s ﬁnding on unidirectional LMs.
Overall, this set of experiments suggests that increasing model’s capabilities to capture a longer context is generally helpful for both pretraining and downstream tasks.
However, using overlapping attention windows is not an effective way to make use of more context.
Thus, we hypothesize the MLM perplexity improvements of overlapping local attentions might mainly come from the reduction of the “boundary" tokens instead of the ability to capture long-range dependencies.
For downstream tasks, the issue of “boundary" tokens is not that essential and the introduction of the overlapping attention windows might disrupt the effective modeling of local context, as the attention module needs to extract both local and distant information from the
[Table]
[Figure]
same set of tokens.
Initializing from
Existing
Short
Models.
While we train all models from scratch for the sake of fair comparison, existing state-of-the-art long context models like Longformer (Beltagy et al., 2020) or BigBird (Zaheer et al., 2020) usually initialize their longer models from an extensively pretrained short model like RoBERTa (Liu et al., 2019).
With simple techniques like positional embedding copying, a strong long-context encoder can be initialized without the need of pretraining from scratch.
To test our ﬁndings from the above analysis in this setting, we follow the same scheme Blocksize [Figure] [Table]
Longformer (64k) 6.6k 73.1/77.
Blockwise LW [Table] but use the non-overlapping block attention as discussed in §4.3. We compare this model with Longformer (based on Sliding Window attention) as it uses the same long-doc corpus and pretrain-and-ﬁnetune pipeline ( e.g., packages and downstream data processing) as our experiments.
Same as our setting in §4.2, here we control the batch size and number of training updates: we use a batch size of 64 and train the model for 64k steps.
Note that as we drop the attention window overlaps, the model is 2x more efﬁcient than Longformer:
Given the same window/block size B and sequence length L , the complexity of the non-overlapping block attention is L × B compared to Longformer’s 2 L × B .
We show the
TriviaQA results in Table 6, where the speed is measured by words per second during pretraining.
With only half of the pretraining compute, our model with disjoint attention blocks achieves slightly better performance than Longformer.
This conﬁrms that our ﬁndings of the attention overlap from the above section are still valid when the models are not trained from scratch.
Related Work
Long-Range Context in Language Models.
Various studies have investigated the effective usage of distant context in unidirectional language models.
Khandelwal et al.
(2018) look into the context usage of LSTM LMs and ﬁnd that these models are only capable to make full use of the nearbytokens and the longer range context is only roughly captured, i.e. , excluding detailed information such as word orders.
Similarly, O’Connor and Andreas (2021) studies the mid- and long-range context usage in transformer LMs, by manipulating the ordering and lexical information in the text.
Their experiments show that while long-range context is usually helpful, most of the usable information is carried by local ordering statistics and non-function words instead of detailed content like sentence orders.
These observations provide a possible explanation of our ablation experiments in §4.3 that adding overlaps to attention windows does not yield better downstream results, despite allowing the capture of more long-range interaction.
Press et al.
(2021) observe diminishing returns as they increase the context length when using sliding windows at infer- ence time.
They propose a staged training paradigm that train LMs from smaller context to longer ones.
This paradigm can more efﬁciently use the training compute and achieves lower perplexity compared to directly training with long sequences.
Given that models with smaller attention windows converge faster at early training steps (Figure 2), the staged training might also beneﬁt MLM pretraining but further investigation is required to validate whether it can also bring downstream improvements.
Other Long-Range Architectures.
Instead of modifying the attention calculation, other work proposes to augment transformers with parametric long-term memories.
Transformer-XL (Dai et al., 2019) maintains frozen activations of previous tokens in memory and uses them as additional inputs.
To handle the shift of positional information of these activations, it also requires a relative position encoding mechanism which brings additional computation cost.
The Compressive Transformer (Rae et al., 2020) takes a similar scheme but proposes to use compression modules to account for even further memories.
Both methods cannot be directly applied to long-context understanding tasks.
Under the scheme of kernel-based methods, Katharopoulos et al.
; Peng et al.
; Schlag et al.
(2021) also attempt to linearize the softmax with kernel methods.
The core ideas of these methods are similar to Performer and they only differ in the choice of kernel functions.
Outside of the transformer families, a recent work (Lei, 2021) proposes to augment recurrent LMs with minimal attention blocks.
It is more efﬁcient while achieving stronger
LM perplexity compared to Transformer-XL.
However, it is still unknown whether this model scales as well as transformer architectures.
Conclusion
We present a systematic study of recent proposed efﬁcient attention variants on real-world long-context
NLP tasks.
In contrast to existing work, we are the ﬁrst to test these models with a set of uniﬁed 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 ﬁnd that the simplest local attentions outperform other complex attention paradigms on downstream tasks.
We also show that existing localattention models can be further simpliﬁed 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.