SlowGuess/NAACL_2022_OCR huggingface dataset:
Wenhan Xiong, Barlas Oğuz, Anchit Gupta, Xilun Chen, Diana Liskovich, Omer Levy, Wen-tau Yih, Yashar Mehdad
{xwhan,barlaso,anchit,xilun,dianaml,omerlevy,scottyih,mehdad}@fb.com
Many NLP tasks require processing long contexts beyond the length limit of existing pretrained models. To scale these models to longer text sequences, many efficient long-range attention variants have been recently proposed. Despite the abundance of research in this direction, it is 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 a widely-used long-range benchmark and show that the other efficient attentions fail to outperform the simple local-window attention after standard pretraining. Further analysis of 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 (Beltagy et al., 2020) with half of its pretraining compute. $^{1}$
The quadratic complexity of Transformer architectures makes it prohibitive to apply large state-of-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 (Kitaev et al., 2020; Tay et al., 2020; Roy et al., 2021), fixed 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 finetuning.
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 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 (Liu et al., 2019) checkpoint<sup>3</sup> 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
2Short documents are concatenated to form long sequences.
3By extending the position embeddings and reusing all other parameters.
The code to replicate our experiments can be found at https://github.com/pytorch/fairseq/tree/main/examples/xformers
Human Language Technologies, pages 1975 - 1986
with the vanilla transformer attention.
A recently proposed benchmark (Tay et al., 2021), named long-range arena (LRA), aims to address the lack of unified evaluation with a bundle of long-sequence tasks. However, the text-related 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 (Devlin et al., 2019). 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-and-finetuning paradigm. For a fair comparison, we implement these attentions under a unified framework and test them using the same Transformer architecture<sup>4</sup> 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:
- With proper tuning, we find that all the tested models can achieve similar level of performance on the LRA benchmark while their performance diverges significantly on large-scale pretraining and downstream tasks;
- In our experiments, the other attention paradigms barely outperform the class of sim
a) Local window
b) Blockwise LW
Figure 1: Attention pattern visualization of two types of local attentions: Left: Local window attention as in Longformer, with window size 2; Right: Blockwise local window attention with block size 2. The rows represent the tokens in the sequence and the columns represent the tokens being attended to.
ple local attentions on downstream tasks when using similar pretraining 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 efficient.
2 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 Beltagy et al. (2020) where each token attends to the same number of tokens on each side, and a simplified 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. 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.
4We only modify the attention calculation within the multihead attention blocks
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 \times k)$  to reduce the length of key and value feature matrix, i.e., from  $L \times d$  to  $k \times d$  ( $k \ll 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 \times 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 input-dependent 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 (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.
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 (Lefaudeauux et al., 2021), which allows us to plug these models into our pretraining-and-finetuning pipeline in a consistent fashion.
3.1 LRA Experiments
Following recent work on efficient long-range attentions, we take the LRA benchmark as our first 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 classification problems: ListOps requires the model to predict the correct output of an expression (10-way classification), 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<sup>5</sup>. 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 sufficient to achieve similar or better results than previously reported results. The final classification layer is added on top of the representations of [CLS] tokens which are prepended to each sequence.
3.2 Pretraining and Downstream Tasks
For practical NLP application, large-scale self-supervised 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 finetuning.
5The limit of training updates is arbitrarily set in LRA and various work have reported hugely improved results on the text matching task, simply by running more training steps.
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 find 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 finished 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 significant 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. See additional model-specific 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 classification, and document retrieval. For the first two tasks, we use TriviaQA and Hyperpartisan classification 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. (2020) 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 classification accuracy for Hyperpartisan. For the retrieval task, for the ease of experiments, we reported the mean reciprocal rank on the dev set<sup>7</sup>, which has been shown to correlate well with final 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-specific Architectures for Finetuning. We use standard architectures for the finetuning tasks: for extractive QA, a single-layer MLP span predictor is added on top of the output token representations; the classification task uses a binary MLP classifier 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 classification 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.
4.1 Models Perform Similarly in LRA
We report our reimplemented LRA results in Table 1. While previous work (Tay et al., 2021) has
For each question, the ground-truth document will be ranked with all documents (both positive and negative) corresponding to the dev-set questions.
$^{6}$ Linformer and Performer cannot reach reasonable perplexity without pre-layer normalization.
Model	ListOps	Text	Matching	Avg Acc	GFlops
Learnable attention pattern
Sinkhorn	37.6	63.8	80.4	60.6	0.289
LSH	37.9	62.5	80.5	60.3	0.273
Low-rank/kernel-based approximation
Linformer	37.7	61.9	78.4	59.3	0.271
Nystrom	37.9	66.1	81.0	61.7	0.256
Performer	37.1	66.1	79.8	61.0	0.205
Hybrid attention
Long-Short	37.7	65.7	81.6	61.7	0.199
Fixed attention pattern
Local Window	37.4	65.7	81.6	61.6	0.153
Blockwise LW	37.4	65.6	81.3	61.4	0.146
Table 1: LRA (the text-input subsets) results with our reimplementations. We did not observe significant performance gaps between different attention variants and simple local attentions remain strong compared to the best Nyström attention.
shown a clear performance gap between different models, we find that with proper tuning, the results of several models could be significantly improved, (e.g., Sinkhorn, Linformer, Reformer, Performer) and there is no significant 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.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.2 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, fixed local attentions remain to be strong baselines. However, in contrast to LRA, we observe local attentions are significantly 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 fixed 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 fixed 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 fixed local attention patterns, improvements on these text LRA tasks often do not transfer to the standard scaled pretraining-finetuning 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 Nystrom is significantly 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 findings 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
Models	MLM Pretraining	TriviaQA	Downstream Tasks
PPL ↓	k word/sec ↑	Doc Retrieval	Hyperpartisan
Learnable attention pattern
Sinkhorn	4.03	11.8	63.3/68.5	80.9	95.0
LSH	3.63	10.0	62.9/67.5	83.6	92.2
Low-rank/kernel-based approximation
Linformer	4.14	24.6	59.8/65.2	80.3	88.7
Nystrom	3.79	9.5	51.5/57.3	83.1	89.5
Performer	5.58	17.2	24.5/31.9	66.8	94.9
Hybrid attention
Long-Short	3.36	8.4	66.5/71.4	84.5	91.5
Fixed local attention pattern
Sliding Window	3.47	9.2	65.6/70.7	83.2	95.3
Blockwise LW	3.39	13.5	68.1/72.9	85.0	95.0
Table 2: MLM pretraining and downstream task results.
Model	MLM Train Perplexity
Linformer	4.31
Performer	6.36
Blockwise LW	4.04
Table 3: Training perplexity of our best fixed local attention and other faster attention variants. Each model uses similar GPU memory and training time.
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.3 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 findings, we aim to investigate the following questions: How effective are the long-range mechanism in local attention models? and Whether the studied long-context 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 field increases with model depth.