What are the method proposed for extending transformer models to efficiently process long context?
Long-context efficient text generation transform-
ers (Tay et al., 2021, 2022) extend earlier trans-
former models (Vaswani et al., 2017) for processing
long sequences, often using a sparse self-attention
architecture. Examples include the Longformer
Encoder-Decoder (LED) (Beltagy et al., 2020),
and LongT5 (Guo et al., 2022). These models
demonstrated that single-text approaches be can
adapted to multi-document tasks by concatenating multiple documents into a single sequence and
processing them using their sparse attention patterns. They sparsify the full self-attention matrix
of transformers by using a combination of a lo-
calized sliding window (called local attention), as
well as a global attention pattern on a few spe-
cific input locations. LED is build upon the BART
model (Lewis et al., 2020) by using additional po-
sitional embeddings and global attention weights,
and introduces the global attention mode that op-
erates over pre-selected tokens. LongT5 extends
the T5 model (Raffel et al., 2020) by using a simi-
lar technique introduced in the ETC and BIGBIRD
models (Ainslie et al., 2020; Zaheer et al., 2020),
relieving the requirement to manually select global
tokens by automatically globalizing the aggregated
representations of groups of tokens.
Longformer: The Long-Document Transformer
Transformer-based models are unable to process long sequences due to their self-attention operation, which scales quadratically with the sequence length. To address this limitation, we introduce the Longformer with an attention mechanism that scales linearly with sequence length, making it easy to process documents of thousands of tokens or longer. Longformer's attention mechanism is a drop-in replacement for the standard self-attention and combines a local windowed attention with a task motivated global attention. Following prior work on long-sequence transformers, we evaluate Longformer on character-level language modeling and achieve state-of-the-art results on text8 and enwik8. In contrast to most prior work, we also pretrain Longformer and finetune it on a variety of downstream tasks. Our pretrained Longformer consistently outperforms RoBERTa on long document tasks and sets new state-of-the-art results on WikiHop and TriviaQA.
[{"first":"Iz","last":"Beltagy","middle":[]},{"first":"Matthew","last":"Peters","middle":["E."]},{"first":"Arman","last":"Cohan","middle":[]}]
Long Range Arena: A Benchmark for Efficient Transformers
Transformers do not scale very well to long sequence lengths largely because of quadratic self-attention complexity. In the recent months, a wide spectrum of efficient, fast Transformers have been proposed to tackle this problem, more often than not claiming superior or comparable model quality to vanilla Transformer models. To this date, there is no well-established consensus on how to evaluate this class of models. Moreover, inconsistent benchmarking on a wide spectrum of tasks and datasets makes it difficult to assess relative model quality amongst many models. This paper proposes a systematic and unified benchmark, LRA, specifically focused on evaluating model quality under long-context scenarios. Our benchmark is a suite of tasks consisting of sequences ranging from $1K$ to $16K$ tokens, encompassing a wide range of data types and modalities such as text, natural, synthetic images, and mathematical expressions requiring similarity, structural, and visual-spatial reasoning. We systematically evaluate ten well-established long-range Transformer models (Reformers, Linformers, Linear Transformers, Sinkhorn Transformers, Performers, Synthesizers, Sparse Transformers, and Longformers) on our newly proposed benchmark suite. LRA paves the way towards better understanding this class of efficient Transformer models, facilitates more research in this direction, and presents new challenging tasks to tackle. Our benchmark code will be released at https://github.com/google-research/long-range-arena.
[{"first":"Yi","last":"Tay","middle":[]},{"first":"Mostafa","last":"Dehghani","middle":[]},{"first":"Samira","last":"Abnar","middle":[]},{"first":"Yikang","last":"Shen","middle":[]},{"first":"Dara","last":"Bahri","middle":[]},{"first":"Philip","last":"Pham","middle":[]},{"first":"Jinfeng","last":"Rao","middle":[]},{"first":"Liu","last":"Yang","middle":[]},{"first":"Sebastian","last":"Ruder","middle":[]},{"first":"Donald","last":"Metzler","middle":[]}]
ETC: Encoding Long and Structured Inputs in Transformers
Transformer models have advanced the state of the art in many Natural Language Processing (NLP) tasks. In this paper, we present a new Transformer architecture, Extended Transformer Construction (ETC), that addresses two key challenges of standard Transformer architectures, namely scaling input length and encoding structured inputs. To scale attention to longer inputs, we introduce a novel global-local attention mechanism between global tokens and regular input tokens. We also show that combining global-local attention with relative position encodings and a Contrastive Predictive Coding (CPC) pre-training objective allows ETC to encode structured inputs. We achieve state-of-the-art results on four natural language datasets requir-ing long and/or structured inputs.
[{"first":"Joshua","last":"Ainslie","middle":[]},{"first":"Santiago","last":"Ontanon","middle":[]},{"first":"Chris","last":"Alberti","middle":[]},{"first":"Vaclav","last":"Cvicek","middle":[]},{"first":"Zachary","last":"Fisher","middle":[]},{"first":"Philip","last":"Pham","middle":[]},{"first":"Anirudh","last":"Ravula","middle":[]},{"first":"Sumit","last":"Sanghai","middle":[]},{"first":"Qifan","last":"Wang","middle":[]},{"first":"Li","last":"Yang","middle":[]}]
Big Bird: Transformers for Longer Sequences
Transformers-based models, such as BERT, have been one of the most successful deep learning models for NLP. Unfortunately, one of their core limitations is the quadratic dependency (mainly in terms of memory) on the sequence length due to their full attention mechanism. To remedy this, we propose, BigBird, a sparse attention mechanism that reduces this quadratic dependency to linear. We show that BigBird is a universal approximator of sequence functions and is Turing complete, thereby preserving these properties of the quadratic, full attention model. Along the way, our theoretical analysis reveals some of the benefits of having $O(1)$ global tokens (such as CLS), that attend to the entire sequence as part of the sparse attention mechanism. The proposed sparse attention can handle sequences of length up to 8x of what was previously possible using similar hardware. As a consequence of the capability to handle longer context, BigBird drastically improves performance on various NLP tasks such as question answering and summarization. We also propose novel applications to genomics data.
[{"middle":[],"last":"Zaheer","first":"Manzil"},{"middle":[],"last":"Guruganesh","first":"Guru"},{"middle":[],"last":"Dubey","first":"Avinava"},{"middle":[],"last":"Ainslie","first":"Joshua"},{"middle":[],"last":"Alberti","first":"Chris"},{"middle":[],"last":"Ontanon","first":"Santiago"},{"middle":[],"last":"Pham","first":"Philip"},{"middle":[],"last":"Ravula","first":"Anirudh"},{"middle":[],"last":"Wang","first":"Qifan"},{"middle":[],"last":"Yang","first":"Li"},{"middle":[],"last":"Ahmed","first":"Amr"}]
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Recent work has shown that either (1) increasing the input length or (2) increasing model size can improve the performance of Transformer-based neural models. In this paper, we present a new model, called LongT5, with which we explore the effects of scaling both the input length and model size at the same time. Specifically, we integrated attention ideas from long-input transformers (ETC), and adopted pre-training strategies from summarization pre-training (PEGASUS) into the scalable T5 architecture. The result is a new attention mechanism we call {\em Transient Global} (TGlobal), which mimics ETC's local/global attention mechanism, but without requiring additional side-inputs. We are able to achieve state-of-the-art results on several summarization tasks and outperform the original T5 models on question answering tasks.
[{"first":"Mandy","last":"Guo","middle":[]},{"first":"Joshua","last":"Ainslie","middle":[]},{"first":"David","last":"Uthus","middle":[]},{"first":"Santiago","last":"Ontanon","middle":[]},{"first":"Jianmo","last":"Ni","middle":[]},{"first":"Yun-Hsuan","last":"Sung","middle":[]},{"first":"Yinfei","last":"Yang","middle":[]}]
Longformer: The Long-Document Transformer
Transformers (Vaswani et al., 2017) have achieved state-of-the-art results in a wide range of natural language tasks including generative language modeling  and discriminative language understanding (Devlin et al., 2019). This success is partly due to the self-attention component which enables the network to capture contextual information from the entire sequence. While powerful, the memory and computational requirements of self-attention grow quadratically with sequence length, making it infeasible (or very expensive) to process long sequences on current hardware.
Longformer: The Long-Document Transformer
To address this limitation, we present Longformer, a modified Transformer architecture with Figure 1: Longformer's memory usage scales linearly with the sequence length, unlike the full self-attention mechanism that runs out of memory for long sequences on current GPUs. Longformer's GPU-kernel is nearly as fast as the highly optimized full self-attention operation, and nearly 6X faster than naive Pytorch. a self-attention operation that scales linearly with the sequence length, making it versatile for processing long documents (Fig. 1). This is an advantage for natural language tasks such as long document classification, question answering (QA), and coreference resolution, where existing approaches partition or shorten the long context into smaller sequences that fall within the typical 512 token limit of BERT-style pretrained models. Such partitioning could potentially result in loss of important cross-partition information, and to mitigate this problem, existing methods often rely on complex architectures to address such interactions. On the other hand, our proposed Longformer is able to build contextual representations of the entire context using multiple layers of attention, reducing the need for task-specific architectures.
Longformer: The Long-Document Transformer
Recent work has addressed the computational inefficiency of Transformers on long sequences (see Tab. 1). However, they primarily focus on autoregressive language modeling, while the application of long document transformers to document-level NLP tasks in the transfer learning setting (Dai and Le, 2015;Peters et al., 2018;Howard and Ruder, 2018;Devlin et al., 2019) has remained largely unexplored. We address this gap and show that Longformer's attention mechanism can act as a drop-in replacement for the self-attention mechanism in pretrained Transformers, and leads to gains across a suite of document NLP tasks. Longformer's attention mechanism is a combination of a windowed local-context self-attention and an end task motivated global attention that encodes inductive bias about the task. Through ablations and controlled trials we show both attention types are essential -the local attention is primarily used to build contextual representations, while the global attention allows Longformer to build full sequence representations for prediction.
Longformer: The Long-Document Transformer
We first evaluate Longformer on autoregressive character-level language modeling using a combination of windowed and a new dilated attention pattern, allowing the model to process sequences of up to 32K characters on modern GPUs. We achieve state-of-the-art results on text8 and enwik8 benchmark datasets, demonstrating the effectiveness of Longformer in long document modeling.
Longformer: The Long-Document Transformer
Then, to evaluate Longformer's ability to replace the full self-attention operation of existing pretrained models, we pretrain it with the masked language modeling (MLM) objective, continuing from the RoBERTa (Liu et al., 2019) released checkpoint. After pretraining, we apply it to downstream language tasks through finetuning and demonstrate that Longformer consistently outperforms RoBERTa on a wide range of document-level natural language tasks including text classification, QA, and coreference resolution, achieving state-ofthe-art results on two of these datasets.
Longformer: The Long-Document Transformer
The original Transformer model has a self-attention component with O(n 2 ) time and memory complexity where n is the input sequence length and thus, is not efficient to scale to long inputs. To address this challenge, we sparsify the full self-attention matrix according to an "attention pattern" specifying pairs of input locations attending to one another. Unlike the full self-attention, our proposed attention pattern scales linearly with the input sequence, making it efficient for longer sequences. In the following, we discuss the design and implementation of this attention pattern. Sliding Window Given the importance of local context (Kovaleva et al., 2019), our attention pattern employs a fixed-size window attention surrounding each token. Using multiple stacked layers of such windowed attention results in a large receptive field, where top layers have access to all input locations and have the capacity to build representations that incorporate information across the entire input. More formally, in this attention pattern, given a fixed window size w, each token attends to 1 2 w tokens on each side (Fig. 2b). The computation complexity of this pattern is O(n×w), which scales linearly with input sequence length n. To make this attention pattern efficient, w should be small compared with n. However, as mentioned above, a model with typical multiple stacked transformer will have a large receptive field. This is analogues to CNNs where stacking layers of small kernels leads to high level features that are built from a large portion of the input (receptive field) . In our case, with a transformer of layers, the receptive field size is × w (assuming w is fixed for all layers). Depending on the application, it might be helpful to use different values of w for each layer to balance between efficiency and model representation capacity ( §4.1).
Longformer: The Long-Document Transformer
To further increase the receptive field without increasing computation, the sliding window can be "dilated". This is analogues to dilated CNNs (van den Oord et al., 2016) where the window has gaps of size dilation d (Fig. 2c). Assuming a fixed d and w for all layers, the receptive field is × d × w, which can reach tens of thousands of tokens even for small values of d.
Longformer: The Long-Document Transformer
In multi-headed attention, each attention head computes a different attention score. We found settings with different dilation configurations per head improves performance by allowing some heads without dilation to focus on local context, while others with dilation focus on longer context.
Longformer: The Long-Document Transformer
In state-of-the-art BERT-style models for natural language tasks, the optimal input representation differs from language modeling and varies by task. For masked language modeling (MLM), the model uses local context to predict the masked word, while for classification, the model aggregates the representation of the whole sequence into a special token ([CLS] in case of BERT). For QA, the question and document are concatenated, allowing the model to compare the question with the document through self-attention.
Longformer: The Long-Document Transformer
In our case, the windowed and dilated attention are not flexible enough to learn task-specific representations. Accordingly, we add "global attention" on few pre-selected input locations. Importantly, we make this attention operation symmetric: that is, a token with a global attention attends to all tokens across the sequence, and all tokens in the sequence attend to it. Fig. 2d shows an example of a sliding window attention with global attention at a few tokens at custom locations. For example for classification, global attention is used for the [CLS] token while in QA global attention is provided on all question tokens. Note that, since the number of such tokens is small relative and independent of n (the total number of tokens in the input sequence) the complexity of the combined local and global attention is still O(n).
Longformer: The Long-Document Transformer
While specifying global attention is task specific, it is much simpler than existing task specific approaches that chunk/shorten the input into smaller sequences and often use complex architecture to combine information across these chunks. Further, it increases the representational power of the model as it allows building contextual representations across the entire sequence.
Longformer: The Long-Document Transformer
Linear Projections for Global Attention Recall that given the linear projections Q, K, V , the Transformer model (Vaswani et al., 2017) computes attention scores as follows:
Longformer: The Long-Document Transformer
We use two sets of projections, Q s , K s , V s to compute attention scores of sliding window attention, and Q g , K g , V g to compute attention scores for the global attention. The additional projections provide flexibility to model the different types of attention, which we show is critical for best performance on downstream tasks. Q g , K g , V g are all initialized with values that match Q s , K s , V s .
Longformer: The Long-Document Transformer
In regular transformers, attention scores are computed as in Eqn. 1. The expensive operation is the matrix multiplication QK T because both Q and K have n (sequence length) projections. Our dilated sliding window attention pattern is not straightforward to implement using modern deep learning libraries. Implementing it requires a form of banded matrix multiplication (matrix multiplication where the output is all zero except certain diagonals) that is not supported in existing deep learning libraries like PyTorch/Tensorflow. In addition, naive implementations with loops is unusably slow. To address this challenge, we provide a highly efficient custom CUDA kernel that implements these attention operations and allows parallelizing the operation on GPU threads.
Longformer: The Long-Document Transformer
Tensor Virtual Machine (TVM) We build our custom CUDA kernel using TVM (Chen et al., 2018), a deep learning compiler stack that compiles high level description of a function into optimized device-specific code. Using TVM, we describe our form of banded matrix multiplication in high-level python constructs, then TVM generates the corresponding CUDA code and compiles it for GPUs. Fig. 1 compares the runtime and memory of the full self-attention, our TVM implementation of the dilated sliding window, and a naive implementation of this attention. It is clear the full attention implementation is fast but it takes significant amount of memory because it needs to store all n 2 values. The naive implementation with loops is not memory consuming because it only stores the non-zero values, however it is significantly slow and impractical to use. Finally, our implementation is both fast and memory efficient because it only computes and stores the non-zero values. 3
Longformer: The Long-Document Transformer
Autoregressive or left-to-right language modeling is loosely defined as estimating the probability distribution of an existing token/character given its previous tokens/characters in an input sequence. This task is considered one of the fundamental tasks in natural language and recent prior work on modeling long sequences using transformers has relied on this task as their primary evaluation Rae et al., 2020;Sukhbaatar et al., 2019). Similarly, we develop and evaluate our model on autoregressive language modeling.
Longformer: The Long-Document Transformer
For autoregressive language modeling we use our dilated sliding window attention. Following Sukhbaatar et al. (2019) we use differing window sizes across the layers. In particular, we use small window sizes for the lower layers and increase window sizes as we move to higher layers. This allows the top layers to learn higher-level representation of the entire sequence while having the   lower layers capture local information. In addition, it provides balance between efficiency (smaller window sizes are less computationally expensive due to fewer nonzero values) and performance (larger window sizes have richer representation power and often result in performance improvements). We do not use dilated sliding windows for lower layers to maximize their capacity to learn and utilize the immediate local context. For the higher layers, we use a small amount of increasing dilation only on 2 heads. This gives the model the ability to directly attend to distant tokens without sacrificing local context.
Longformer: The Long-Document Transformer
Task and Datasets We focus on character-level autoregressive language modeling because the sequences are naturally longer than those of wordlevel language modeling, making them suitable for evaluating our model. We used text8 and enwik8 (Mahoney, 2009) for evaluation, two standard and widely used character-level language modeling datasets.
Longformer: The Long-Document Transformer
Training Ideally, we would like to train our model on the largest window size and sequence length we can fit in a modern GPU memory. However, we found that the model needs a large number of gradient updates to learn the local context first; before learning to utilize longer context. To accom-modate this, we adopt a staged training procedure where we increase the attention window size and sequence length across multiple training phrases. In particular, in the first phase we start with a short sequence length and window size, then on each subsequent phase, we double the window size and the sequence length, and halve the learning rate. This makes training fast, while keeping the slow part (longest sequences and window sizes) to the end. We train the model over 5 total phases with starting sequence length of 2,048 and ending sequence length of 23,040 on the last phase (see Appendix A for detailed configurations of each phase).
Longformer: The Long-Document Transformer
Evaluation At evaluation time, we are able to run our model on sequence of length 32,256. Following prior work Sukhbaatar et al., 2019), during evaluation we split the dataset into overlapping sequences of size 32,256 with a step of size 512, and report the performance on the last 512 tokens on the sequence.
Longformer: The Long-Document Transformer
Hyperparameters Our model only specifies how the self-attention component works, and it is agnostic to the other design choices for the transformer model. We use relative position embeddings with sinusoidal weights as in . We use two different model sizes, a small (12 layers, 512 hidden size) model as in , and a large (30 layers, 512 hidden size) model as in . We employed mixed precision training (floating points 16 and 32) using apex 4 to reduce memory consumption and speed-up training. However, we kept the attention computation in fp32 to avoid numerical instability issues. 5 We used gradient checkpointing (Chen et al., 2016) to reduce memory usage, and ran our experiments on 48GB RTX8000 GPUs. Refer to Appendix A for a more detailed list of hyperparameters.
Longformer: The Long-Document Transformer
Tab. 2 and 3 summarize evaluation results on text8 and enwik8 datasets. We achieve a new state-of-the-art on both text8 and enwik8 using the small models with BPC of 1.10 and 1.00 on text8 and enwik8 respectively, demonstrating the effectiveness of our model.
Longformer: The Long-Document Transformer
For large models, given how expensive these experiments are, and following recent work (Kitaev et al., 2020;Rae et al., 2020), we are only   , and matches or slightly underperforms recent models that have more than twice the number of parameters. It is worth noting that Adaptive Span (Sukhbaatar et al., 2019) and Compressive Transformer (Rae et al., 2020) are not good fit for the pretrainingfinetuning paradigm as discussed in §2.
Longformer: The Long-Document Transformer
To show the importance of the design choices of our attention patterns, we tried different variants and report their controlled experiment results. To make the ablation study more manageable, we train each configuration for 150K steps 6 with phase 1 configuration on a small model on text8, then report the BPC performance on the dev set. The top of Tab. 4 demonstrates the impact of different ways of configuring the window sizes per layer. We observe that increasing the window size from the bottom to the top layer leads to the best performance, arranging them in the reverse way leads to worse performance, and using a fixed window size (the average of window sizes of the other configuration) leads to a performance that it is in between. The bottom of Tab. 4 shows the impact of adding dilation. Adding some dilation to two heads leads to some improvement compared with no dilation at all.
Longformer: The Long-Document Transformer
Current state-of-the-art systems for many NLP tasks finetune a pretrained model with task supervision (e.g. BERT). One of our main motivations is to develop such a model suitable for long document tasks. To do so, we pretrained Longformer  We pretrain Longformer with masked language modeling (MLM), where the goal is to recover randomly masked tokens in a sequence. Since MLM pretraining is expensive, we continue pretraining from the RoBERTa (Liu et al., 2019) released checkpoint, while only making the minimal changes necessary to support Longformer's attention mechanism. Note that our attention pattern can be plugged into any pretrained transformer model without the need to change the model architecture.
Longformer: The Long-Document Transformer
Attention Pattern We use the sliding window attention with window size of 512 on all layers. This matches RoBERTa's sequence length, and therefore uses the same amount of computation as RoBERTa. 8
Longformer: The Long-Document Transformer
Position Embeddings RoBERTa uses learned absolute position embeddings with the maximum position being 512. To support longer documents, we add extra position embeddings to support up to position 4,096. To leverage RoBERTa's pretrained weights, instead of randomly initializing the new position embeddings, we initialize them by copying the 512 position embeddings from RoBERTa multiple times as analysis of BERT's attention heads shows a strong learned bias to attending to local context, including the previous or next token (Clark et al., 2019). Using the copy initialization preserves this local structure everywhere except at the partition boundaries. Despite its simplicity, we found this to be a very effective (see Tab. 6), allowing Longformer pretraining to rapidly converge with a small number of gradient updates. 7 Sequences up to 16K are possible on current GPUs. 8 We tried adding the additional dilation pattern on a few heads as in §4.1 but found it to hurt performance, likely because it is not compatible with the pretrained RoBERTa weights. We suspect that retraining such model from scratch might be needed to get improved performance Pretraining Data In order to allow the model to learn long dependencies in pretraining, we compiled a corpus of long documents. Some of these data sources were also included in the original RoBERTa pretraining including the Books corpus (Zhu et al., 2015) plus English Wikipedia. We additionally included one third of a subset of the Realnews dataset (Zellers et al., 2019) with documents longer than 1,200 tokens as well as one third of the Stories (Trinh and Le, 2018) corpus. Our goal was to include a mix of long and short documents to both allow the model to learn longer dependencies while not to forget information from the original RoBERTa pretraining. The statistics of the pretraining data is shown in Tab. 5.
Longformer: The Long-Document Transformer
We train 9 two sizes of Longformer, a base model and a large model. Both models are trained for 65K gradient updates with sequences length 4,096, batch size 64 (2 18 tokens), maximum learning rate of 3e-5, linear warmup of 500 steps, followed by a power 3 polynomial decay. The rest of the hyperparameters are the same as RoBERTa.
Longformer: The Long-Document Transformer
Tab. 6 shows the BPC on the development set of our training corpus. The first row shows a 1.846 BPC using RoBERTa-base, which is comparable to the 1.880 BPC reported on the RoBERTa paper on their corpus. This indicates our training corpus is from a distribution close to that used to train RoBERTa. The following two rows show the performance of Longformer before pretraining with randomly initialized position embeddings and with copied position embeddings. The significant difference indicates the importance of the copy initialization, and the relative small difference between the RoBERTa BPC and the initialized BPC indicates that our sliding window attention is working well with the RoBERTa weights. The following two rows show the impact of continuing pretraining. Traininig for 2K steps improves BPC from 1.957 to 1.753, which further decreases to 1.705 after 65K steps, demonstrating the model is learning to better utilize the sliding window attention and longer context. Similar patterns are observed with RoBERTa-large and Longformer-large.
Longformer: The Long-Document Transformer
We apply Longformer to multiple long document tasks, including QA, coreference resolution and classification. Tab. 7 shows the evaluation datasets 9 using fairseq (Ott et al., 2019)   have contexts significantly longer than 512 wordpieces. Our primary goal is to evaluate whether our attention mechanism can act as a replacement for the standard self-attention mechanism in BERT style models, and to perform controlled trials against a strong baseline. We are also interested in evaluating whether we can replace complicated task specific models necessitated by BERT's limited context with simpler models that just concatenate all available context into a single sequence. Our baseline is a RoBERTa based model that breaks the context into the longest possible segment, passes each individually through RoBERTa, and concatenates the activations for further processing. For QA tasks, we also concatenate the question to each segment so that RoBERTa can condition it's contextual representations of the context on the question. The Longformer variant replaces the RoBERTa self-attention mechanism with our windowed attention used during pretraining, plus a task motivated global attention. The global attention uses additional linear projections ( §3.1).
Longformer: The Long-Document Transformer
Question answering We consider three datasets with long contexts: WikiHop (Welbl et al., 2018), Wikipedia setting of TriviaQA (Joshi et al., 2017), and HotpotQA, distractor dev setting (Yang et al., 2018). 10 For WikiHop and TriviaQA we follow the simple QA model of BERT (Devlin et al., 2019), and concatenate question and documents into one long sequence, run it through Longformer, then have a  dataset-specific prediction layer. WikiHop uses a classification layer for the candidate while Trivi-aQA uses the loss function of Clark and Gardner (2017) to predict answer span. We include global attention to question tokens and answer candidates for WikiHop and to question tokens for TriviaQA. For HotpotQA, we experiment with two different models, a single stage multitask model that jointly predicts evidence and answer spans, and a twostage model that first extracts evidence paragraphs and passes them to a second stage for answer extraction. The single stage model combines a span extraction loss with a question type (yes/no/span) classification head over the first token and an evidence extraction loss predicted from special tokens at the end of sentences and paragraphs. We also include global attention to sentence and paragraph tokens. Note that both approaches are significantly simpler than recent SOTA models that include pipeline approaches of multiple stages and customized architectures (e.g., (Tu et al., 2019;Tu et al., 2020;Groeneveld et al., 2020)). See Appendix B for further details about the models and hyperparameters.
Longformer: The Long-Document Transformer
Coreference Resolution We use OntoNotes (Pradhan et al., 2012), and the model from Joshi et al. (2019), a modification of the system from Lee et al. (2018) to replace ELMo with BERT. The Longformer system is a straightforward adaption of the baseline model by replacing RoBERTa with Longformer and extending the sequence length. We didn't use global attention for this task.
Longformer: The Long-Document Transformer
We evaluate on IMDB (Maas et al., 2011) and Hyperpartisan news detection (Kiesel et al., 2019) datasets. 11 IMDB is a standard sentiment classification datasets consisting of movie reviews. While most documents in this dataset are short, about 13.6% of them are 11 For Hyperpartisan we split the training data into train/dev/test sets using standard 90/10/10 splits and performed each experiment five times with different seeds to control variability associated with the small dataset. larger than 512 wordpieces (Tab. 7). Documents in Hyperpartisan are relatively long, and it is small with only 645 documents making it a good test for Longformer's ability to adapt to limited data. We use global attention on the [CLS] token.
Longformer: The Long-Document Transformer
Main Result Tab. 8 summarizes the results of all our finetuning experiments. We observe that Longformer consistently outperforms the RoBERTa baseline. Its performance gain is especially obvious for tasks that require long context such as WikiHop and Hyperpartisan. For TriviaQA, the improvement is more modest as the local context is often sufficient to answer the question. In the case of HotpotQA, the supporting fact auxiliary supervision allows models to easily find relevant contexts and then focus on local context, leading to smaller gains. This is contrasted with WikiHop that only includes distant supervision of intermediate reasoning chains, where our approach excels by reasoning over the entire context. On the IMDB and OntoNotes datasets the performance gains are smaller. For IMDB, the majority of the dataset consists of short documents and thus it is expected to see smaller improvements. For OntoNotes, we found that the distance between any two mentions is typically quite small so that a baseline that processes smaller chunks separately is able to stitch together mentions into coreference chains without considering cross chunk interactions.
Longformer: The Long-Document Transformer
Longformer-large for QA We also evaluate the performance of Longformer-large on long context QA tasks. Tab. 9 shows that our Longformer-large achieves new state-of-the-art results on WikiHop and TriviaQA by large margins (3.6 and 4 points respectively). Tab. 10 summarizes results of Hot-potQA, and, as expected, using Longformer-large improves the result compared to Longformer-base. The two-stage model improves the results even further, likely because of the increased capacity that allows each stage to specialize on one task.   (Fang et al., 2019) 81.0 87.9 73.0 C2F Reader (Shao et al., 2020) --73.9 This model matches performance of TAP2 (Glaß et al., 2019), which is the the best performing single model that doesn't use a form of graph neural networks (GNN; Kipf and Welling, 2017). However, our model is simpler than TAP2 where they use a three-stage approach consisting of extracting paragraphs, evidence sentences and finally answer spans with an additional specialized span selection pretraining. All the better performing models for HotpotQA (Shao et al., 2020;Fang et al., 2019) use multi-stage approaches plus GNNs or graph network of entities, which seem to encode an important inductive bias for the task. 12 Furthermore, the 10 paragraphs come from 10 different documents and the concatenated context is therefore not coherent. In this case we suspect that Longformer has less chance to use information from pretraining due to difference between the pretraining and finetuning input structure. Additional pretraining tasks could help further improve results.
Longformer: The Long-Document Transformer
Tab. 11 presents an ablation study for WikiHop. All results use Longformer-base, trained for 5 epochs with identical hyperparameters except where noted. Longformer benefits from longer sequences, global attention, separate projection matrices for global attention, MLM pretraining, and longer training.
Longformer: The Long-Document Transformer
In addition, when configured as in RoBERTa-base (seqlen: 512, and n 2 attention) Longformer performs slightly worse then RoBERTa-base, confirm- 12 We can encode this inductive bias using global attention over entities, but we leave this for future work. 13 We report numbers from dev set scores of published papers. Missing values are not publicly available. Leaderboard test results will be included in the future.
Longformer: The Long-Document Transformer
Accuracy / ∆  ing that performance gains are not due to additional pretraining.
Longformer: The Long-Document Transformer
We present Longformer, a transformer-based model that is scalable for processing long documents and that makes it easy to perform a wide range of document-level NLP tasks without chunking/shortening the long input and without complex architecture to combine information across these chunks. Longformer employs an attention pattern that combines local and global information while also scaling linearly with the sequence length. Longformer achieves state-of-the-art results on the character-level language modeling tasks of text8 and enwik8. When pretrained, Longformer consistently outperforms RoBERTa on long document tasks and sets new state-of-the-art results on Wiki-Hop and TriviaQA. For future work, we would like to explore other attention patterns that are more efficient by dynamically adapting to the input. We also would like to apply our model to other relevant long document tasks such as summarization.
Longformer: The Long-Document Transformer
We evaluate on text8 and enwik8, both contain 100M characters from Wikipedia split into 90M, 5M, 5M for train, dev, test. Our implementation is based on the Transformer-XL  code 14 with the memory mechanism disabled. Our hyperparameters and stage configurations are listed in Tab. 12. Our CUDA kernel supports the autoregressive mode where each token attends to a window of previous tokens only. Our implementation also includes a version of the relative position embedding that is compatible with our dilated sliding window attention.
Longformer: The Long-Document Transformer
We ran the small model experiments on 4 RTX8000 GPUs for 16 days. For the large model, we ran experiments on 8 RTX8000 GPUs for 13 days. Most of our hyperparameter search is similar to the ablation in Tab. 4 where we run the configuration for 150K steps on text8. We experimented with absolute position embeddings and learned position embeddings, dropout values of [0.1, 0.2] (small model) and [0.1, 0.4] (large model), prelayernorm and post-layernorm (Xiong et al., 2020), learning rate (LR) of phase1 of values [2.5e-5, 5e-4, 1e-4] constant and cosine LR schedules, and different configurations for dilation (on all heads, on 2 heads, no dilation). Number of gradient updates/phase reported in Tab. 12 is determined by running each phase until the validation BPC stops getting better.
Longformer: The Long-Document Transformer
All the QA and classification models are implemented using PyTorch-Lightning 15 .
Longformer: The Long-Document Transformer
WikiHop Instances in WikiHop consist of: a question, answer candidates (ranging from two candidates to 79 candidates), supporting contexts (ranging from three paragraphs to 63 paragraphs), and the correct answer. The dataset does not provide any intermediate annotation for the multihop reasoning chains, requiring models to instead infer them from the indirect answer supervision.
Longformer: The Long-Document Transformer
To prepare the data for input to Longformer and RoBERTa, we first tokenize the question, answer candidates, and support contexts using RoBERTa's wordpiece tokenizer.
Longformer: The Long-Document Transformer
Then After preparing the input data, we compute activations from the top layer of each model as follows. We take the question and answer candidates and concatenate them to as much context as possible up to the model sequence length (512 for RoBERTa, 4,096 for Longformer), run the sequence through the model, collect the output activations, and repeat until all of the context is exhausted (for all models except Longformer-large, where we just include the first 4,096 length sequence due to memory requirements). Then all activations for all chunks are concatenated into one long sequence. In the case of Longformer, we use global attention to the entire question and answer candidate sequence.
Longformer: The Long-Document Transformer
For prediction, we attach a linear layer to each [ent] that outputs a single logit, average over all logits for each candidate across the chunks, apply a softmax and use the cross entropy loss with the correct answer candidate.
Longformer: The Long-Document Transformer
Training used the Adam optimizer with linear warmup over 200 gradient updates to a maximum LR, and linear decay over the remainder of training. We used gradient accumulation to effective batch size of 32 instances, checking the development accuracy every 250 gradient updates and reported the maximum development accuracy. Other hyperparameters (dropout, weight decay) were identical to RoBERTa pretraining.
Longformer: The Long-Document Transformer
In general, we ran minimal hyperparameter trials, but for fair comparison between Longformer and RoBERTa ran an identical hyperparameter search with Longformer-base and RoBERTa-base. This consisted of a grid search of LR in [2e-5, 3e-5, 5e-5] and number epochs in [5,10,15]. The best Longformer-base configuration used lr=3e-5, 15 epochs. We ran two hyperparameter trials for Longformer-large, lr=3e-5 and number epochs in [5,15] (the 5 epoch model had higher dev accuracy of 77.6, and was the single model submitted to the public leaderboard for test set evaluation). All mod-  For prediction, we add one layer that predicts the beginning and end of the answer span. Because of the distant supervision nature of the training data (no gold answer spans), we use the loss function of Clark and Gardner (2017) which works like an OR that the model only needs to get one answer span right, not all of them.
Longformer: The Long-Document Transformer
Hyperparameters of the best configuration are listed in Tab. 13. All other hyperparameters are similar to RoBERTa's. For hyperparameter search, we only tuned LR for the RoBERTa baseline and tried rates [3e-5, 5e-5, 1e-4], then used the best, which is 3e-5, for all subsequent experiments with no further tuning. We trained the Longformer-large with the best configuration once and submitted its output to the leaderboard. We ran our experiments on 32GB V100 GPUs. Small model takes 1 day to train on 4 GPUs, while large model takes 1 day on 8 GPUs.
Longformer: The Long-Document Transformer
HotpotQA dataset involves answering questions from a set of 10 paragraphs from 10 different Wikipedia articles where 2 paragraphs are relevant to the question and the rest are distractors. It includes 2 tasks of answer span extraction and evidence sentence identification. Our model for HotpotQA combines both answer span extraction and evidence extraction in one joint model. We also experimented with a two-stage model with similar setup that uses longformer first for evidence extraction and then extract answers using a BERT-based QA model. The second stage of our two-stage model is based on Groeneveld et al. (2020). Similar to Wikihop and TriviaQA, to prepare the data for input to Longformer and RoBERTa, we concatenate question and then all the 10 paragraphs in one long context. We particularly use the following input format with special tokens: " [ are special tokens representing sentences and paragraphs. The special tokens were added to the RoBERTa vocabulary and randomly initialized before task finetuning. For Longformer, we use global attention to input tokens as well as sentence and paragraph tokens. For answer span extraction we use BERT's QA model (Devlin et al., 2019) with addition of a question type (yes/no/span) classification head over the first special token ([CLS]). For evidence extraction we apply 2 layer feedforward networks on top of the representations corresponding to sentence and paragraph tokens to get the corresponding evidence prediction scores and use binary cross entropy loss to train the model. We combine span, question classification, sentence, and paragraphs losses and train the model in a multitask way using linear combination of losses. Our experiments are done on RTX8000 GPUs and training each epoch takes approximately half a day on 4 GPUs. We trained the model using Adam optimizer with linear warmup (1000 steps) and linear decay. We used minimal hyperpareter tuning using LRs of 3e-5 and 5e-5 and epochs of 3 to 7 and found the model with LR of 3e-5 and 6 epochs to work best. We conduct the same hyperparameter search for the RoBERTa baseline as well. The rest of hyperparameters are reported in Tab 13.  Coreference model details The coreference model is a straightforward adaptation of the coarseto-fine BERT based model from Joshi et al. (2019). After preprocessing each document with the RoBERTa wordpiece tokenizer, it splits each document into non-overlapping segments up to the maximum sequence length, then concatenates the activations for the coarse-to-fine clustering stage that forms coreference clusters. The maximum sequence length was 384 for RoBERTa-base, chosen after three trials from [256,384,512] using the default hyperparameters in the original implementation. 16 For Longformer-base the sequence length was 4,096. Similar to the original implementation, 16 https://github.com/mandarjoshi90/coref different learning rates were used for the pretrained RoBERTa parameters and the randomly initialized task parameters. Using a larger learning rate in the task parameters allows the optimizer to adjust them farther from their randomly initialized values without destroying the information in the pretrained RoBERTa parameters. Hyperparameter searches were minimal and consisted of grid searches of  and task LR in [1e-4, 2e-4, 3e-4] for both RoBERTa and Longformer for a fair comparison. The best configuration for Longformer-base was RoBERTa lr=1e-5, task lr=1e-4. All other hyperparameters were the same as in the original implementation. Training takes about 10 hours on a single GPU.
Longformer: The Long-Document Transformer
Our implementation is a superhack that involves PyTorch and Tensorflow sharing a single process and GPU. To avoid re-implementing the complicated coarse-to-fine logic from Tensorflow in PyTorch (that involves a highly optimized custom GPU kernel originally released by Lee et al. (2018)), we devised a system where the lower transformer portion of the model passes activations and gradients back and forth between PyTorch and Tensorflow. The input tensors are first run through the transformer in PyTorch, the activations are collected from the top layer, transferred from GPU to CPU then from CPU to Tensorflow and back to GPU to run the coarse-to-fine clustering and compute the loss. Then gradients are back propogated in Tensorflow to the top of the transformer and the process reversed to transfer them to PyTorch for back propogation through the remainder of the model. Separate optimizers are maintained with identical LR schedules for parameter updates. The overhead in this approach is minimal compared to the overall cost of running the model.
Longformer: The Long-Document Transformer
Text classification For classification, following BERT, we used a simple binary cross entropy loss on top of a first [CLS] token with addition of global attention to [CLS]. We used Adam optimizer with batch sizes of 32 and linear warmup and decay with warmup steps equal to 0.1 of the total training steps. For both IMDB and Hyperpartisan news we did grid search of LRs [3e-5, 5e-5] and epochs [10,15,20] and found the model with [3e-5] and epochs 15 to work best. Experiments were done on a single RTX8000 GPU.
Long Range Arena: A Benchmark for Efficient Transformers
Transformers (Vaswani et al., 2017) are ubiquitously state-of-the-art across many modalities, from language (Devlin et al., 2018;Raffel et al., 2019;Child et al., 2019) to images (Tan & Bansal, 2019;Lu et al., 2019) to protein sequences (Rives et al., 2019). A common weakness of Transformers is their quadratic memory complexity within the self-attention mechanism that restricts their potential application to domains requiring longer sequence lengths. To date, a dizzying number of efficient Transformer models ('xformers') have been proposed to tackle this problem (Liu et al., 2018;Kitaev et al., 2020;Tay et al., 2020b;Katharopoulos et al., 2020). Many of these models demonstrate comparable performance to the vanilla Transformer model while successfully reducing the memory complexity of the self-attention mechanism. An overview of this research area can be found in (Tay et al., 2020c).
Long Range Arena: A Benchmark for Efficient Transformers
Transformers (Vaswani et al., 2017) are ubiquitously state-of-the-art across many modalities, from language (Devlin et al., 2018;Raffel et al., 2019;Child et al., 2019) to images (Tan & Bansal, 2019;Lu et al., 2019) to protein sequences (Rives et al., 2019). A common weakness of Transformers is their quadratic memory complexity within the self-attention mechanism that restricts their potential application to domains requiring longer sequence lengths. To date, a dizzying number of efficient Transformer models ('xformers') have been proposed to tackle this problem (Liu et al., 2018;Kitaev et al., 2020;Tay et al., 2020b;Katharopoulos et al., 2020). Many of these models demonstrate comparable performance to the vanilla Transformer model while successfully reducing the memory complexity of the self-attention mechanism. An overview of this research area can be found in (Tay et al., 2020c).
Long Range Arena: A Benchmark for Efficient Transformers
Comparing the evaluation and experimental setup of many of these papers, we can make the following observations. Firstly, there is no unifying consensus on what makes an acceptable test bed for benchmarking efficient Transformers. There is also a large diversity in the types of tasks adoptedevery single model is evaluated on a different set of tasks and datasets, which makes comparison of different models as well as an assessment of their relative strengths and weaknesses difficult. Secondly, the benchmarks used for evaluation are often arbitrarily chosen, without much consideration to whether the task is suitable for evaluating long-range modeling. Thirdly, many papers tend to conflate the effectiveness of the inductive bias with the benefits of pretraining Zaheer et al., 2020;, which tends to obfuscate the true value of the architecture. Pretraining itself is a computationally expensive endeavour and de-coupling inductive bias research from pretraining would make xformer research more accessible.
Long Range Arena: A Benchmark for Efficient Transformers
In this paper, we propose a new benchmark, Long-Range Arena (LRA), for the purpose of benchmarking sequence models under the long-context scenario. We design a benchmark suite comprised of both synthetic probing tasks and real-world tasks and provide relative comparisons for ten recently proposed efficient Transformer models including Sparse Transformers (Child et al., 2019), Reformer (Kitaev et al., 2020), Linformer , Longformer (Beltagy et al., 2020), Sinkhorn Transformers (Tay et al., 2020b), Performers (Choromanski et al., 2020b), Synthesizers (Tay et al., 2020a), Linear Transformers (Katharopoulos et al., 2020), and BigBird (Zaheer et al., 2020). This is the most comprehensive and extensives side-by-side evaluation of this class of models.
Long Range Arena: A Benchmark for Efficient Transformers
In this paper, we propose a new benchmark, Long-Range Arena (LRA), for the purpose of benchmarking sequence models under the long-context scenario. We design a benchmark suite comprised of both synthetic probing tasks and real-world tasks and provide relative comparisons for ten recently proposed efficient Transformer models including Sparse Transformers (Child et al., 2019), Reformer (Kitaev et al., 2020), Linformer , Longformer (Beltagy et al., 2020), Sinkhorn Transformers (Tay et al., 2020b), Performers (Choromanski et al., 2020b), Synthesizers (Tay et al., 2020a), Linear Transformers (Katharopoulos et al., 2020), and BigBird (Zaheer et al., 2020). This is the most comprehensive and extensives side-by-side evaluation of this class of models.
Long Range Arena: A Benchmark for Efficient Transformers
While the focus of this benchmark is the ability of these architectures to reason in long-context scenarios, we are also fundamentally interested in understanding the capabilities and properties of these xformer architectures when exposed to different types of data and conditions. Hence, our benchmark is purposefully designed to be capability probing, i.e, we select datasets and tasks with certain innate structure. For example, can these architectures model long sequences that are intrinsically hierarchical or that contain some form of spatial structure? In general, we are especially interested in the relative performance of these xformer models across diverse circumstances. We hope that understanding these better will inspire research on more efficient architectures in the future. While the focus of this paper is on efficient Transformer models, our benchmark is also model agnostic and can also serve as a benchmark for long-range sequence modeling.
Long Range Arena: A Benchmark for Efficient Transformers
Aside from comparing the quality of these models, we also conduct extensive efficiency and memory usage analysis of these models. We believe such a side-by-side performance benchmark will be valuable to the community, providing deeper insight on the practical efficiency of these methods. Overall, we propose a unified framework for enabling easy side-by-side comparisons of efficient Transformer models and broadly speaking, long-range sequence models in general. Our framework, which we open source, is written in JAX/FLAX 1 .
Long Range Arena: A Benchmark for Efficient Transformers
This section introduces the Long-Range Arena (LRA) benchmark (pronounced el-ra). We implement our benchmark (which includes the task, evaluators, and models) in Python 3 and Jax/Flax and open-source our code 2 -making it easy to extend and to build on top of our work.
Long Range Arena: A Benchmark for Efficient Transformers
For creating the Long-Range Arena benchmark, we established a set of desiderata:
Long Range Arena: A Benchmark for Efficient Transformers
1. Generality: All efficient Transformers models should be applicable to our tasks. For instance,
Long Range Arena: A Benchmark for Efficient Transformers
given that not all xformer models are able to perform autoregressive decoding , we include tasks that only require encoding.
Long Range Arena: A Benchmark for Efficient Transformers
2. Simplicity: The tasks should have a simple setup. All factors that make comparisons difficult should be removed. This encourages simple models instead of cumbersome pipelined approaches. For instance, we avoid including any particular data augmentation and consider pretraining to be out of scope of this benchmark.
Long Range Arena: A Benchmark for Efficient Transformers
The tasks should be difficult enough for current models to ensure there is room for improvement to encourage future research in this direction.
Long Range Arena: A Benchmark for Efficient Transformers
The input sequence lengths should be reasonably long since assessing how different models capture long-range dependencies is a core focus of LRA.
Long Range Arena: A Benchmark for Efficient Transformers
Probing diverse aspects: The set of tasks should assess different capabilities of models like their ability to model relations and hierarchical/spatial structures, generalization capability, etc.
Long Range Arena: A Benchmark for Efficient Transformers
6. Non-resource intensive and accessible: The benchmarks should be deliberately designed to be lightweight so as to be accessible to researchers without industry-grade computing resources.
Long Range Arena: A Benchmark for Efficient Transformers
This section describes the tasks in the LRA benchmark. Note that these tasks are specifically designed for the purpose of assessing different aspects of efficient Transformer models. Further details about each task can be found in the appendix.
Long Range Arena: A Benchmark for Efficient Transformers
In this task, we are interested in the capability of modeling hierarchically structured data in a longcontext scenario. This benchmark task is a longer variation of the standard ListOps task proposed in (Nangia & Bowman, 2018), which was designed to investigate the parsing ability of neural models.
Long Range Arena: A Benchmark for Efficient Transformers
In this task, we are interested in the capability of modeling hierarchically structured data in a longcontext scenario. This benchmark task is a longer variation of the standard ListOps task proposed in (Nangia & Bowman, 2018), which was designed to investigate the parsing ability of neural models.
Long Range Arena: A Benchmark for Efficient Transformers
The dataset is comprised of sequences with a hierarchical structure and operators MAX, MEAN, MEDIAN and SUM MOD that are enclosed by delimiters (brackets). An example (much shorter) sequence is as follows: In our task we use a version of ListOps of sequence lengths of up to 2K to test the ability to reason hierarchically while handling long contexts. Naturally, in the above example the model needs to access all tokens and model the logical structure of the inputs in order to make a prediction. The task is a ten-way classification task and is considerably challenging.
Long Range Arena: A Benchmark for Efficient Transformers
This task using real-world data represents a common use case of efficient Transformers, which are often needed to process long documents. Text classification in particular is associated with many real-world applications such as spam, fraud, and bot detection and commercial document classification, among others (Howard & Ruder, 2018).
Long Range Arena: A Benchmark for Efficient Transformers
This task using real-world data represents a common use case of efficient Transformers, which are often needed to process long documents. Text classification in particular is associated with many real-world applications such as spam, fraud, and bot detection and commercial document classification, among others (Howard & Ruder, 2018).
Long Range Arena: A Benchmark for Efficient Transformers
This task also benchmarks the ability of the models to deal with compositionality as it is required to compose characters into words into higher-level phrases. Compared to ListOps, boundaries are less well defined and need to be learned from the data, which is a challenging problem in its own right (Kawakami et al., 2019).
Long Range Arena: A Benchmark for Efficient Transformers
This task also benchmarks the ability of the models to deal with compositionality as it is required to compose characters into words into higher-level phrases. Compared to ListOps, boundaries are less well defined and need to be learned from the data, which is a challenging problem in its own right (Kawakami et al., 2019).
Long Range Arena: A Benchmark for Efficient Transformers
We consider the byte/character-level setup of this task in order to simulate a longer input sequence, which also makes the task considerably more challenging. 3 Note that this setup differs significantly from character-level language modeling (char LM). In char LM, it would suffice to read nearby context to determine the next character, e.g., a model is very likely to predict 'e' after having seen the prefix 'appl'. For byte-level text classification, the model needs to reason with compositional, unsegmented data in order to solve a meaningful real-world task. We use the IMDb reviews (Maas et al., 2011) dataset, which is a commonly used dataset to benchmark document classification. We use a fixed max length of 4K for this task, which is truncated or padded when necessary. This is a binary classification task with accuracy as the metric.
Long Range Arena: A Benchmark for Efficient Transformers
We further evaluate a model's ability to encode and store compressed representations that are useful for matching and retrieval. Learning the similarity score between two vectors is a common problem in machine learning and is useful for a wide array of applications (Guo et al., 2016).
Long Range Arena: A Benchmark for Efficient Transformers
We further evaluate a model's ability to encode and store compressed representations that are useful for matching and retrieval. Learning the similarity score between two vectors is a common problem in machine learning and is useful for a wide array of applications (Guo et al., 2016).
Long Range Arena: A Benchmark for Efficient Transformers
Hence, this task is mainly about modeling a similarity score between two documents in a 'two tower setup' in which compressed representations are concatenated and passed into a linear classifier. Note that we deliberately prevent models from using cross attention. This task thus serves as a test of how well models are able to compress long sequences into representations suitable for similarity-based matching.
Long Range Arena: A Benchmark for Efficient Transformers
We use the ACL Anthology Network (AAN; Radev et al., 2013) dataset, which identifies if two papers have a citation link, a common setup used in long-form document matching (Jiang et al., 2019;.
Long Range Arena: A Benchmark for Efficient Transformers
We use the ACL Anthology Network (AAN; Radev et al., 2013) dataset, which identifies if two papers have a citation link, a common setup used in long-form document matching (Jiang et al., 2019;.
Long Range Arena: A Benchmark for Efficient Transformers
Similar to the text classification setup, we use a byte/character level setup, which challenges the model to compose and aggregate information over longer contexts. We use a sequence length of 4K for each document, which makes the total text length 8K for this task. This is a binary classification task with accuracy as the metric.
Long Range Arena: A Benchmark for Efficient Transformers
This task is an image classification task, where the inputs are sequences of pixels. In other words, an N × N image is flattened to a sequence of length N 2 pixels.
Long Range Arena: A Benchmark for Efficient Transformers
Similar to how the previous tasks require capturing the hierarchical structure in the data, this task requires the model to learn the 2D spatial relations between input pixels, while presented as a 1D sequence of symbols.
Long Range Arena: A Benchmark for Efficient Transformers
We focus on assessing Transformer models that are designed to process a sequence of discrete symbols, so we do not allow extra modules such as a CNN stem that embeds pixel-level inputs.
Long Range Arena: A Benchmark for Efficient Transformers
To simplify the setup, we map the input images to a single gray-scale channel where each pixel is represented with an 8-bit pixel intensity (vocabulary size of 256). In LRA, we use the CIFAR-10 dataset (Krizhevsky, 2009) for the image classification task.
Long Range Arena: A Benchmark for Efficient Transformers
(a) A positive example.
Long Range Arena: A Benchmark for Efficient Transformers
(b) A negative example. Figure 1: Samples of the Pathfinder task.
Long Range Arena: A Benchmark for Efficient Transformers
The Pathfinder challenge (Linsley et al., 2018;Kim* et al., 2020) was first introduced for learning long-range spatial dependencies. It is a synthetic visual task motivated by cognitive psychology (Houtkamp & Roelfsema, 2010).
Long Range Arena: A Benchmark for Efficient Transformers
The Pathfinder challenge (Linsley et al., 2018;Kim* et al., 2020) was first introduced for learning long-range spatial dependencies. It is a synthetic visual task motivated by cognitive psychology (Houtkamp & Roelfsema, 2010).
Long Range Arena: A Benchmark for Efficient Transformers
The task requires a model to make a binary decision whether two points represented as circles are connected by a path consisting of dashes. We show a positive example of two connected points and a negative example of two unconnected points in Figure 1. The dataset also contains distractor paths, which makes this setup challenging. We model this task by treating images as sequences of pixels. In this task, images are of dimensions (32 × 32), which make up a sequence length of 1024.
Long Range Arena: A Benchmark for Efficient Transformers
Finally, we consider an extreme version of Pathfinder (Pathfinder-X) where examples consist of 16K pixels (i.e., images of 128 × 128).
Long Range Arena: A Benchmark for Efficient Transformers
The key goal here is to observe if a model would fail to solve the 16K extreme version even if it can successfully learn the standard version of 1024 tokens. This is an interesting litmus test to see if the same algorithmic challenges bear a different extent of difficulty when sequence lengths are much longer. We include this in our benchmark as Path-X.
Long Range Arena: A Benchmark for Efficient Transformers
One of the main goals of the LRA benchmark is assessing the ability of different efficient Transformer models to capture long-range dependencies. The tasks and setups are designed with this goal in mind. In order to have a quantitative estimate of the spatial extent needed to be considered by an attention mechanism to encode the inputs, we define required attention span.
Long Range Arena: A Benchmark for Efficient Transformers
Given a trained attention-based model and a sequence of tokens as inputs, the required attention span of an attention module is computed as the mean distance between the query token and the attended tokens, scaled by attention weights. Here, we compute the mean required attention span over all attention modules in our best vanilla Transformer model for each task, averaged over 1K random samples from the validation set.  Given the purpose of LRA, we found required attention span serves as a good proxy for how difficult a task is for Transformer-based models. 4 3 EXPERIMENTAL RESULTS
Long Range Arena: A Benchmark for Efficient Transformers
This section describes the models we evaluate on our LRA benchmark. We base our evaluation on ten recently proposed efficient Transformer models. Aside from the standard vanilla Transformer (Vaswani et al., 2017) and a simple local attention baseline, we compare Sparse Transformers (Child et al., 2019), Longformers (Beltagy et al., 2020), Linformers , Reformers (Kitaev et al., 2020), Sinkhorn Transformers (Tay et al., 2020b), Synthesizers (Tay et al., 2020a), BigBird (Zaheer et al., 2020), Linear Transformers (Katharopoulos et al., 2020), and Performers (Choromanski et al., 2020a).
Long Range Arena: A Benchmark for Efficient Transformers
This section describes the models we evaluate on our LRA benchmark. We base our evaluation on ten recently proposed efficient Transformer models. Aside from the standard vanilla Transformer (Vaswani et al., 2017) and a simple local attention baseline, we compare Sparse Transformers (Child et al., 2019), Longformers (Beltagy et al., 2020), Linformers , Reformers (Kitaev et al., 2020), Sinkhorn Transformers (Tay et al., 2020b), Synthesizers (Tay et al., 2020a), BigBird (Zaheer et al., 2020), Linear Transformers (Katharopoulos et al., 2020), and Performers (Choromanski et al., 2020a).
Long Range Arena: A Benchmark for Efficient Transformers
We believe these ten models to represent a diverse cross-section of recent efficient Transformer models.
Long Range Arena: A Benchmark for Efficient Transformers
We note that it is non-trivial and almost impossible to conduct a perfectly fair evaluation of all models. The large search space motivates us to follow a set of fixed hyperparameters (number of Best model is in boldface and second best is underlined. All models do not learn anything on Path-X task, contrary to the Pathfinder task and this is denoted by FAIL. This shows that increasing the sequence length can cause seriously difficulties for model training. We leave Path-X on this benchmark for future challengers but do not include it on the Average score as it has no impact on relative performance.
Long Range Arena: A Benchmark for Efficient Transformers
layers, heads, embedding dimensions, etc.) for all models. The best performance and relative order of the models may change if we aggressively tune hyperparameters for all models. Hence, the results provided in this paper are not meant to be a final authoritative document on which xformer is the best. Instead, we provide a starting point for future research and strive to be as fair as possible.
Long Range Arena: A Benchmark for Efficient Transformers
In order to do so, we plan to release the code with all the hyperparameters and implementation details. Additionally, we intend for our paper to be a living document and encourage researchers (authors and the broader community) to contribute and continue updating this paper (with rules and limitations described in the appendix). We also implemented all models to the best of our abilities. We often consulted with the original developers of the included models.
Long Range Arena: A Benchmark for Efficient Transformers
Based on our results, we observe that (1) all proposed tasks in LRA are considerably challenging and (2) there are meaningful differences in model performance across different xformer models.
Long Range Arena: A Benchmark for Efficient Transformers
Results on ListOps The ListOps task (10-way classification) has proven to be reasonably difficult with the best models obtaining only 37%. The considerable gap to random chance shows that models are indeed learning the task. We notice that the inductive bias of the xformer model plays a substantial role on this task in which approximately half the xformer models are able to get > 30% performance while the remainder of the models only get slightly above random chance. This may imply that certain efficiency-inspired inductive biases may be better at handling hierarchical data than others. For instance, the results from our experiments seem to suggest that kernel-based models (e.g., Performer, Linear Transformers) are possibly not as effective on hierarchically structured data.
Long Range Arena: A Benchmark for Efficient Transformers
Results on Text Classification Byte-level classification is shown to be difficult and challenging especially when no pretraining or contextual embeddings are used. The best model only obtains 65.90 accuracy. The Linear Transformer performs well on this task, along with the Performer model. Contrary to the ListOps task, it seems like fast kernel-based models do well on this task.
Long Range Arena: A Benchmark for Efficient Transformers
The scores of different models on this task are also rather low (average of 55%), indicating the difficulty of the task. The vanilla Transformer model only achieves 57.46% accuracy with some xformer variants scoring very close to random chance. The best performing model is the Sparse Transformer and the second best is BigBird. We find that models that follow fixed sparse patterns to do well on this task. Models that are based on low-rank factorization and kernels perform relatively worse.  Results on Image Classification On the image classification task, most models perform quite similarly (low variance amongst model performance). The best model on this task is the Sparse Transformer, followed by the Performer. Linformer and Reformers do not do well on this task. On a related note, we also observed most of models struggle generalizing to the test even though they manage to overfit the training set. While we extensively tried different regularization techniques on every single model, there is a rather large gap between their performance on train and test set (More details in Appendix).
Long Range Arena: A Benchmark for Efficient Transformers
Results on Pathfinder / Path-X Results show that all models achieve reasonable performance on the Pathfinder task. The average performance is 72 and the best model Performer obtains 77.05% accuracy. The Local Attention model performs the worse out of all models.
Long Range Arena: A Benchmark for Efficient Transformers
All models failed to solve the Path-X task, achieving at best 50%. We find this intriguing because this is essentially an identical task to the standard Pathfinder, albeit with much longer sequence lengths. Hence, we observe that the extreme length of the task can significantly obstruct a model from leaning anything meaningful. We leave Path-X in our benchmark suite, hoping to spur future progress in modeling sequences at extreme lengths.
Long Range Arena: A Benchmark for Efficient Transformers
In this section, we report efficiency metrics of our runs. For simplicity, we use the byte-level text classification benchmark and report run times and memory consumption of the sequence lengths {1K, 2K, 3K, 4K}. We use a batch size of 32 for all runs and conduct experiments on 4x4 TPU V3 Chips. We emphasize that this is again largely conditioned on hardware and implementation details (more details can be found in the appendix). Table 2 reports our efficiency benchmarks on the xformer models. We note that low-rank and kernel-based models are generally the fastest. The overall fastest model is the Performer model (Choromanski et al., 2020a), which is 5.7× faster than Transformers on the 4k sequence length. Linformer  and Linear Transformers (Katharopoulos et al., 2020) come in a close second and are almost as fast as Performers (at 5.5 to 5.6× faster). Based on our implementation, the slowest model is the Reformer model (Kitaev et al., 2020) that is about 80% the speed of vanilla Transformer at 4K sequence lengths and half the speed at 1K sequence length.
Long Range Arena: A Benchmark for Efficient Transformers
In this section, we report efficiency metrics of our runs. For simplicity, we use the byte-level text classification benchmark and report run times and memory consumption of the sequence lengths {1K, 2K, 3K, 4K}. We use a batch size of 32 for all runs and conduct experiments on 4x4 TPU V3 Chips. We emphasize that this is again largely conditioned on hardware and implementation details (more details can be found in the appendix). Table 2 reports our efficiency benchmarks on the xformer models. We note that low-rank and kernel-based models are generally the fastest. The overall fastest model is the Performer model (Choromanski et al., 2020a), which is 5.7× faster than Transformers on the 4k sequence length. Linformer  and Linear Transformers (Katharopoulos et al., 2020) come in a close second and are almost as fast as Performers (at 5.5 to 5.6× faster). Based on our implementation, the slowest model is the Reformer model (Kitaev et al., 2020) that is about 80% the speed of vanilla Transformer at 4K sequence lengths and half the speed at 1K sequence length.
Long Range Arena: A Benchmark for Efficient Transformers
The model with the smallest memory footprint in our benchmarks is the Linformer model, coming in at 0.99GB per TPU device as compared to 9.48GB per TPU device for the vanilla Transformers at N = 4K. That is about a 10x reduction in memory footprint. Similar to speed, Performers and Linear Transformers are also relatively compact and are almost as compact as Linformers. Other models (Local Attention, Reformers, BigBird, Synthesizers) are still less memory hungry compared to vanilla Transformers but are relatively less efficient (memory consumption wise) compared to Linformers, Performers, and Linear Transformers. We also notice that the memory consumption of models such as Linformer and Performer scales very well, with the memory usgae at 3K and 4K being approximately equal.
Long Range Arena: A Benchmark for Efficient Transformers
Based on our analysis, the best qualitative performance in terms of LRA score, i.e. integrated across all five tasks, is the BigBird model. While BigBird does not do extremely well on any individual task compared to other models, it has consistently good performance across all tasks. Performers and Linear Transformers have strong performance on some tasks but their average is lowered by the ListOps task. Figure 3 shows the trade-off between qualitative performance, model speed, and memory footprint. While BigBird performs well, its speed is almost similar to the vanilla Transformer. On the other hand, a model like Local Attention is fast at the cost of lower quantitative performance. Among these models, the kernel-based variants, i.e., Performer, Linformer, and linear Transformer seem to be able to make a better trade-off in terms of speed and performance, while having reasonable memory usage. 
Long Range Arena: A Benchmark for Efficient Transformers
The pervasiveness of Transformer models, along with its well-known trait of being memory intensive, has spurred on a large number of innovations on this front. Early work in this area has typically considered a fixed pattern (local window) approach (Liu et al., 2018;Parmar et al., 2018). More advanced models have been proposed recently, including combined patterns (Child et al., 2019;Ho et al., 2019;Beltagy et al., 2020;Zaheer et al., 2020), learned patterns (Kitaev et al., 2020;Roy et al., 2020), and recent models based on kernels (Katharopoulos et al., 2020;Choromanski et al., 2020a) or low-rank approximations . For the sake of brevity, we refer interested readers to (Tay et al., 2020c) for a detailed survey of this line of research.
Long Range Arena: A Benchmark for Efficient Transformers
The pervasiveness of Transformer models, along with its well-known trait of being memory intensive, has spurred on a large number of innovations on this front. Early work in this area has typically considered a fixed pattern (local window) approach (Liu et al., 2018;Parmar et al., 2018). More advanced models have been proposed recently, including combined patterns (Child et al., 2019;Ho et al., 2019;Beltagy et al., 2020;Zaheer et al., 2020), learned patterns (Kitaev et al., 2020;Roy et al., 2020), and recent models based on kernels (Katharopoulos et al., 2020;Choromanski et al., 2020a) or low-rank approximations . For the sake of brevity, we refer interested readers to (Tay et al., 2020c) for a detailed survey of this line of research.
Long Range Arena: A Benchmark for Efficient Transformers
Generative Modeling / Language Modeling This generative modeling task requires predicting the next character, word, or pixel and is a staple in xformer evaluations (Roy et al., 2020;Kitaev et al., 2020). However, it has been debated how much long-range signal such tasks actually encode (Rae & Razavi, 2020).
Long Range Arena: A Benchmark for Efficient Transformers
Generative Modeling / Language Modeling This generative modeling task requires predicting the next character, word, or pixel and is a staple in xformer evaluations (Roy et al., 2020;Kitaev et al., 2020). However, it has been debated how much long-range signal such tasks actually encode (Rae & Razavi, 2020).
Long Range Arena: A Benchmark for Efficient Transformers
LSTM language models augmented with attention have been shown to rarely attend beyond seven preceding words of context (Daniluk et al., 2017) and samples from LSTM language models are known to quickly devolve into generic text. On the other hand, recent models such as the Transformer-XL (Dai et al., 2019) have been observed to be sensitive to a context of around 900 tokens and samples from large-scale models (Radford et al., 2019) maintain a consistent theme over much longer sequences. Even such recent models, however, can be improved by limiting the range of attention (Rae & Razavi, 2020). In sum, while standard language modelling datasets contain some long-range signal, which is required to perform long-range coreference resolution, reasoning with events, discourse understanding, etc. (Ruder et al., 2019) it seems to be overshadowed by the much stronger signal of short-term word co-occurrences and is thus difficult to evaluate. 5
Long Range Arena: A Benchmark for Efficient Transformers
LSTM language models augmented with attention have been shown to rarely attend beyond seven preceding words of context (Daniluk et al., 2017) and samples from LSTM language models are known to quickly devolve into generic text. On the other hand, recent models such as the Transformer-XL (Dai et al., 2019) have been observed to be sensitive to a context of around 900 tokens and samples from large-scale models (Radford et al., 2019) maintain a consistent theme over much longer sequences. Even such recent models, however, can be improved by limiting the range of attention (Rae & Razavi, 2020). In sum, while standard language modelling datasets contain some long-range signal, which is required to perform long-range coreference resolution, reasoning with events, discourse understanding, etc. (Ruder et al., 2019) it seems to be overshadowed by the much stronger signal of short-term word co-occurrences and is thus difficult to evaluate. 5
Long Range Arena: A Benchmark for Efficient Transformers
Question Answering Another commonly used evaluation task is question answering (QA; Zaheer et al., 2020). Open-domain QA in particular typically requires the model to answer questions based on long contexts such as entire Wikipedia documents (Joshi et al., 2017;Kwiatkowski et al., 2019) or even books (Kočiský et al., 2018). Other datasets are explicitly designed to require multiple 'hops' of reasoning (Welbl et al., 2018;Yang et al., 2018). Successful approaches are often highly engineered, computationally expensive systems that require pre-training and a separate retrieval model (Lee et al., 2019;Guu et al., 2020).
Long Range Arena: A Benchmark for Efficient Transformers
Question Answering Another commonly used evaluation task is question answering (QA; Zaheer et al., 2020). Open-domain QA in particular typically requires the model to answer questions based on long contexts such as entire Wikipedia documents (Joshi et al., 2017;Kwiatkowski et al., 2019) or even books (Kočiský et al., 2018). Other datasets are explicitly designed to require multiple 'hops' of reasoning (Welbl et al., 2018;Yang et al., 2018). Successful approaches are often highly engineered, computationally expensive systems that require pre-training and a separate retrieval model (Lee et al., 2019;Guu et al., 2020).
Long Range Arena: A Benchmark for Efficient Transformers
Natural Language Understanding / GLUE tasks Evaluation on natural language understanding (NLU) tasks is also common . Examples in most of these datasets such as MultiNLI (Williams et al., 2018) and SST (Socher et al., 2013) consist of single sentences and less than 100 tokens on average.
Long Range Arena: A Benchmark for Efficient Transformers
Natural Language Understanding / GLUE tasks Evaluation on natural language understanding (NLU) tasks is also common . Examples in most of these datasets such as MultiNLI (Williams et al., 2018) and SST (Socher et al., 2013) consist of single sentences and less than 100 tokens on average.
Long Range Arena: A Benchmark for Efficient Transformers
We proposed Long Range Arena (LRA), a new benchmark for evaluating progress on efficient Transformer research. Our new benchmark is challenging and probes at model capabilities in dealing with diverse data types and structures such as text, mathematics, and visual data. Our benchmark comprises of tasks ranging from 1K to 16K tokens. For the first time, we conduct an extensive side-by-side comparison of ten recently proposed efficient Transformer models. The experimental results show that these tasks are very challenging even for long-range Transformer models. The overall results show that there is no one-size-fits-all solution and trade-offs have to be made in terms of model quality and speed/memory. We plan to open source our code and benchmarks to facilitate future benchmarking, research and model development.
Long Range Arena: A Benchmark for Efficient Transformers
A.1 LRA TASKS This section describes the details and hyperparameters of each task. We also plan to release the configuration files along with the implementation of the models and benchmarks, that can be used to reproduce the results reported in the paper.
Long Range Arena: A Benchmark for Efficient Transformers
A.1.1 LISTOPS Following the generation steps in (Nangia & Bowman, 2018), we generate our own long version of this task. We use a sequence length of 2k for this task. All our xformer models have an embedding dimension of 512, 8 heads, 6 layers and a feed-forward dimensions of 2048. We train all models for 5K steps. The [CLS] token is used and mapped into a 10 class Softmax layer for classification.
Long Range Arena: A Benchmark for Efficient Transformers
We use the IMDb reviews dataset (Maas et al., 2011) and a sequence length of {1K, 2K, 3K, 4K} tokens for all models. We pick the best results across these four sequence lengths. We use a [cls] token for prediction. All the [cls] tokens from xformer encoders are passed into a two layered MLP with ReLU activations. The MLP emits a 2-class logits for binary classification. We optimize the softmax cross entropy loss function. All xformer models are parameterized by the same number of layers, heads and hidden dimensions, namely 8 heads, 512 hidden dimensions and d = 2048 for positional FFN layers. We use 6 layers for all xformers. The learning rate is 0.05 with weight decay of 0.1. We use Adam with warmup. All models are trained for 20K steps and a batch size of 32.
Long Range Arena: A Benchmark for Efficient Transformers
We use the ACL anthology network for a related article matching task. We use a sequence length of 4K per document (8K tokens in total for two sequences). The two encoders share parameters. Similar to document classification, we use the [cls] token from xformer encoders. Let X 1 be the [cls] token embedding from document 1 and X 2 be the [cls] token embedding from document 2, the final score is computed via:
Long Range Arena: A Benchmark for Efficient Transformers
where MLP(.) is a two layered MLP with relu activation functions. In lieu of the much longer sequence length, we use a batch size of 32, embedding dimension of 128, 4 heads, a FFN dimension of 512 and 4 layers. Model is trained with Adam for 5K steps with a learning rate of 0.5.
Long Range Arena: A Benchmark for Efficient Transformers
We use the gray-scaled (single channel) CIFAR10 as the image classification dataset, with 10 classes. The resolution of input images is 32 × 32 and after flattening the input images, we feed our xformer encoders with a sequence of 1024 pixels. Similar to our other classification tasks, there is a classifier head on top of the xformer encoder, consisting of a two-layer MLP with ReLU activation. Softmax cross-entropy has been used for optimizing the parameters of the models. We trained our models for 200 epochs and have done extensive sweeps over different hyper-parameters and found the following values leading to the best average performance across all xformers: 3 layers, 4 heads, 128 as the hidden dimensions of FFN blocks, 64 as the query/key/value hidden dimensions, and finally the learning rate of 0.01.
Long Range Arena: A Benchmark for Efficient Transformers
For the image classification benchmark, in Section 3, we mentioned that most of the models struggle generalizing to the test set. Table 3 presents the train and test accuracy for different models and for almost all these models, the gap between the two scores is considerably high.
Long Range Arena: A Benchmark for Efficient Transformers
While this task can be simple to solve for convectional models (e.g., accuracy of wide-resnet on gray-scale CIFAR10 with no data augmentation is 89.21) it is rather difficult for Transformer-based  Table 3: Test and train accuracy of different models on Image Classification task. models with this setup. Naturally, one can find ways to improve the performance with a different setup. For instance, in our setup, models are not informed about the oridinality of pixel intensities and consume them as independent symbols. We observed that learning embedding that reflects this property is rather hard for most of these models (Figure ). If we simply replace the embedding layer with a CNN stem, we see imitate boost in the performance (e.g. replacing the embedding layer of a vanilla Transformer with a convectional stem, with 3 × 3 kernel, we get accuracy of 75.32).
Long Range Arena: A Benchmark for Efficient Transformers
Another modification that can lead to better performance is to incorporate spatial representation that are translation invariant in Transformer models (e.g., adding 2D relative positional embedding to a vanilla transformer, we get accuracy of 61.72). However, adding these sorts of changes make the setup digress from the original point of this task in our benchmark.  On the left, we can see the pairwise similarity of learned embeddings for pixel intensities. Although there is a higher similarity for close pixel values, the patterns from these learned embeddings do not perfectly reflect the ordinality of the pixel intensities. On the right, we can see the pairwise similarity of positional embeddings for different input positions. We can see that the lower the distance between two pixels is, the more similar are their learned positional embeddings. However, the spatial closeness in y axis is more preserved in the learned embedding than the distances in the x axis.
Long Range Arena: A Benchmark for Efficient Transformers
A.3 PATHFINDER Pathfinder task probes the ability of models to detect long range spatial dependencies between input features. To solve the task, a model requires to identify the target contour and trace it from one end to the other. Although Pathfinder is visually a simple task, it has been show that the clutter and variations in path shape makes the task difficult for CNN models (Linsley et al., 2018;Kim* et al., 2020).
Long Range Arena: A Benchmark for Efficient Transformers
A.3 PATHFINDER Pathfinder task probes the ability of models to detect long range spatial dependencies between input features. To solve the task, a model requires to identify the target contour and trace it from one end to the other. Although Pathfinder is visually a simple task, it has been show that the clutter and variations in path shape makes the task difficult for CNN models (Linsley et al., 2018;Kim* et al., 2020).
Long Range Arena: A Benchmark for Efficient Transformers
The Pathfinder task is a binary classification task and the resolution of input images is 32 × 32. Similar to image classification task, we feed our xformer encoders with a sequence of 1024 pixels after flattening the input images. The classifier head on top of the xformer encoder is also a twolayer MLP with ReLU activation and we use Softmax cross-entropy loss for the optimization. We trained our models for 200 epochs. The hyper-parameters used for the xformer model are as follow: 4 layers, 8 heads, 128 as the hidden dimensions of FFN blocks, 128 as the query/key/value hidden dimensions, and the learning rate of 0.01.
Long Range Arena: A Benchmark for Efficient Transformers
Given that transformers have many units with global receptive field, they have better potential for solving the task, compared to models with local receptive fields. Figure 5 shows the attention distributions for a set of examples given on token (CLS token) as the query. We can see that the attention module collects information from different positions in input to be able to trace the target path. Each map presents the attention distribution, given the CLS token at the final layer as the query, averaged across all heads in a vanilla Transformer model. Note that for visualization, we use attention-rollout (Abnar & Zuidema, 2020) for more precise input attribution.
Long Range Arena: A Benchmark for Efficient Transformers
Given that transformers have many units with global receptive field, they have better potential for solving the task, compared to models with local receptive fields. Figure 5 shows the attention distributions for a set of examples given on token (CLS token) as the query. We can see that the attention module collects information from different positions in input to be able to trace the target path. Each map presents the attention distribution, given the CLS token at the final layer as the query, averaged across all heads in a vanilla Transformer model. Note that for visualization, we use attention-rollout (Abnar & Zuidema, 2020) for more precise input attribution.
Long Range Arena: A Benchmark for Efficient Transformers
We have also included a Pathfinder-X in LRA, which is similar to Pathfinder, but inputs are in higher resolutions, i.e. longer input sequences. On Pathfinder-X, we have tried two setups for training our models, first training models from scratch, second evaluating models that are trained on Pathfinder.
Long Range Arena: A Benchmark for Efficient Transformers
In both cases, we found out none of the models are able to deal with/generalize to 16K input length.
Long Range Arena: A Benchmark for Efficient Transformers
This section describes the details of our implementation. The code is primarily written in JAX and FLAX. In this section, we note specific details about certain implementations of models. We plan to release hyperparameters in a form of readme or script later.
Long Range Arena: A Benchmark for Efficient Transformers
While most of the fine-grained details is planned to be available in the released code, we provide a brief overview of some settings of the xformer models being evaluated. For local attention, we do not use overlapping blocks. For local attention within Sinkhorn Transformer blocks, we do not overlap windows either. For Linformers, the projections are shared between key and values but not across multiple layers. For Performer models, our implementation uses FAVOR+, the more recent version in the paper Choromanski et al. (2020b).
Long Range Arena: A Benchmark for Efficient Transformers
This section describes several special cases in our implementation details. The diverse suite of Transformers come with a plethora of hardware constraints and implementation details. To succeed, a Transformer model needs to also 'win' the hardware lottery (Hooker, 2020), i.e., having readily supported ops, kernels or accelerator support to take advantage of its technical design. This section discusses some of the trade-offs and edge cases that make comparison of several models challenging.
Long Range Arena: A Benchmark for Efficient Transformers
In the end, we argue that simplicity is a virtue and not requiring any special support is a positive thing for an efficient Transformer model.
Long Range Arena: A Benchmark for Efficient Transformers
On CUDA kernels CUDA kernels are cumbersome and are specific to GPU hardware, making it difficult to implement or use on TPU pods. Generally, these are considered to be undesirable and inconvenient in practical applications. Hence, Sparse Transformer and Longformer are implemented with equivalent implementations to emulate for performance. This is by applying an equivalent mask. For this reason, we do not benchmark Sparse Transformer and Longformer for speed.
Long Range Arena: A Benchmark for Efficient Transformers
Reformer's Implementation Having optimized ops to support many of Reformer's functionality is crucial. Hence, Reformer is implemented slightly differently from other Transformer models. Instead of computing tensors with batch size dimensions B and head dimensions H, (i.e., B × H × N × d), we compute the attention function for tensors of N × d dimensions. After which, we parallelize this function via VMAP over the batch and head dimensions.
Long Range Arena: A Benchmark for Efficient Transformers
We welcome re-evaluation of our models on any task. However, we consider some hyperparameters to be immutable to ensure fair comparison with all models. In the case of proposing new models, the LRA table in the paper can be copied as it is as long as (1) the model size remains unchanged, (2) no pretraining is conducted, (3) no alterations to the fundamental setups (e.g., changing char level to word level or adding spatial information to the image task). We will provide more details at https://github.com/google-research/long-range-arena.
ETC: Encoding Long and Structured Inputs in Transformers
Models based on Transformers (Vaswani et al., 2017), such as BERT (Devlin et al., 2018), or other variants (Yang et al., 2019;Lan et al., 2019;Raffel et al., 2019) have yielded state-of-the-art results in many NLP tasks such as language modeling (Child et al., 2019;Sukhbaatar et al., 2019;Rae et al., 2019;Kitaev et al., 2020), question answering (Lan et al., 2019;Beltagy et al., 2020), and summarization (Zhang et al., 2019). We present the Extended Transformer Construction (ETC) architecture 1 , targeting two limitations of the original models: (1) scaling input length, (2) encoding structured inputs.
ETC: Encoding Long and Structured Inputs in Transformers
The computational and memory complexity of attention in the original Transformer scales quadratically with the input length, typically limiting input length to around 512 tokens. While 512 might be enough for some tasks (e.g., co-reference resolution seems to benefit from even smaller input lengths (Joshi et al., 2019)), this is problematic in 1 Code at goo.gle/research-etc-model others. Consider question answering (QA) tasks that require reasoning across multiple documents (e.g., the HotpotQA dataset (Yang et al., 2018)) all of which must simultaneously fit in the model input.
ETC: Encoding Long and Structured Inputs in Transformers
Other examples are summarization, or QA on long documents. Many approaches have been proposed to address this, like hierarchical processing (Zhang et al., 2019), sparse attention (Child et al., 2019), and segment-level recurrence (Dai et al., 2019).
ETC: Encoding Long and Structured Inputs in Transformers
A second limitation is that few models focus on structured inputs, by which we refer to any underlying graph or hierarchical structure among the input tokens. Although ETC can encode more general graph structure, in this work we focus on representing hierarchical structure in NLP tasks, not usually modeled by Transformer models. For example, text is organized into sentences and paragraphs, and while these have a sequential order, different input documents might not hold any order between them (e.g., the HotpotQA dataset). Additionally, web text contains markup and is laid out using a DOM tree, giving additional structure. We show ETC can represent these and other types of structure, like linking different entity mentions.
ETC: Encoding Long and Structured Inputs in Transformers
To address these challenges, we present a novel attention mechanism called global-local attention, which divides the input into two sequences (which we call the global input and the long input). This mechanism introduces local sparsity to reduce the quadratic scaling of the attention mechanism. When this is coupled with relative position encodings (Shaw et al., 2018), it allows for handling structured inputs in a natural way. Additionally, unlike previous Transformer extensions, ETC can be initialized from existing pre-trained standard BERT models (which together with a GPU/TPUfriendly implementation, allows for efficient model training) 2 . Our results show that initializing from RoBERTa (Liu et al., 2019) significantly improves performance. Finally, we show that by adding a pre-training Contrastive Predictive Coding (CPC) task (Oord et al., 2018), performance improves even further for tasks where structure is important, as CPC plays the role of a masked language model (MLM) task, but at a sentence level of granularity.
ETC: Encoding Long and Structured Inputs in Transformers
We report experiments on four datasets: Natural Questions (NQ) (Kwiatkowski et al., 2019), Hot-potQA (Yang et al., 2018), WikiHop (Welbl et al., 2018), and OpenKP (part of MS MARCO) (Xiong et al., 2019), which have long and/or structured inputs. We set a new state of the art in all of them.
ETC: Encoding Long and Structured Inputs in Transformers
Moreover, although in this paper we strictly focus on ETC, in a related model called BigBird (Zaheer et al., 2020), we experimented with an alternative set of ideas to handle long inputs and its extensions to a decoder for text generation. The focus of BigBird is on the idea of adding random sparse attention patterns to global-local attention, and on showing under which conditions models like BigBird/ETC are universal approximators of sequence functions and are Turing complete. While the key ideas and techniques required to achieve the state-of-the-art results mentioned above for QA tasks are the focus of this paper, the reader is referred to the BigBird work for a joint evaluation of ETC (referred to as BigBird-ETC in that work) and the idea of random sparse attention patterns.
ETC: Encoding Long and Structured Inputs in Transformers
Many variants of the original Transformer model (Vaswani et al., 2017) have been proposed for scaling up training (RoBERTa, Liu et al., 2019), the internal representation (ALBERT, Lan et al., 2019), or both (T5, Raffel et al., 2019), outperforming BERT (Devlin et al., 2018) in tasks such as GLUE (Wang et al., 2018), SQuAD (Rajpurkar et al., 2016) or RACE (Lai et al., 2017). However, these models typically limit inputs to n = 512 tokens due to the O(n 2 ) cost of attention. We classify prior approaches to scale up attention into four categories: sparse attention, recurrence, hierarchical mechanisms, and compressed attention.
ETC: Encoding Long and Structured Inputs in Transformers
Sparse Attention involves limiting each token to attend to a subset of the other tokens. For example, the Sparse Transformer (Child et al., 2019) used predefined attention patterns for both text and image generation. They showed that attending only a new model developed concurrently to ETC, which also allows initialization from BERT/RoBERTa. to previous pixels in the same row or column was enough to generate high quality images, while keeping attention cost at O(n √ n). In the Adaptive Attention Span Transformer (Sukhbaatar et al., 2019) each attention head is associated with a decaying learnable masking function, which limits the number of tokens it can attend to. They show that lower layers learn to use short attention spans, and only in higher layers are attention spans longer. Sparse attention has also been used to increase the interpretability of attention heads by allowing attention to assign exactly zero weight to certain input tokens (Correia et al., 2019). The Reformer (Kitaev et al., 2020) model finds the nearest neighbors of the attention query (those input tokens that would result in the highest attention weights) using locality sensing hashing (Andoni et al., 2015) and only uses those for attention. This reduces attention cost to O(n log(n)). The Routing Transformer (Roy et al., 2020) learns dynamic sparse attention patterns using online k-means, reducing complexity to O(n 1.5 ). Finally, the most related approach to the work presented in this paper is Longformer (Beltagy et al., 2020), developed concurrently to ETC, and which features a very similar global-local attention mechanism as ETC's but does not directly encode graph or hierarchical structure (more detailed comparison in Section 3).
ETC: Encoding Long and Structured Inputs in Transformers
Recurrence incorporates elements of recurrent neural networks into Transformer models to lengthen their attention span. Transformer-XL (Dai et al., 2019) takes this approach, dividing the input sequence into segments and then processing these segments one at a time. At each layer, the model attends to the layer immediately below for both the current and previous input segments. The effect is that layer k is influenced by the current segment and the k − 1 previous segments, as shown in the top-right of Figure 1.
ETC: Encoding Long and Structured Inputs in Transformers
In Hierarchical Mechanisms the input sequence is split into blocks that are ingested independently to produce summary embeddings that represent the whole block. Then, separate layers ingest the concatenation of these embeddings. For example, HIBERT (Zhang et al., 2019) uses this idea at the sentence level for extractive summarization (illustrated in the bottom-left of Figure 1). Hierarchical attention in Transformers has also been applied to other NLP tasks such as neural machine translation (Maruf et al., 2019). Moreover, notice that this idea of processing the input hierarchically  is not specific to Transformer models, and it has been applied to recurrent neural network models both at the level of sentences (Yang et al., 2016;Miculicich et al., 2018) and blocks (Shen et al., 2018).
ETC: Encoding Long and Structured Inputs in Transformers
Compressed Attention takes the idea of hierarchical attention one step further by selectively compressing certain parts of the input. The BP-Transformer (Ye et al., 2019) model builds a binary partitioning tree over the input, and only lets the model attend to the leaves (the raw tokens) for nearby tokens, and higher nodes in the tree (summaries of groups of tokens) as tokens grow more distant (see Figure 1, middle top). Other ideas include memory compressed attention (Liu et al., 2018) where groups of k tokens are compressed via a convolution filter before they are attended to, and the Star Transformer (Guo et al., 2019), where each token can attend only to its immediate left/right neighbors and to a separate special auxiliary token that represents a summary of the whole input (see Figure 1, left). The Compressive Transformer (Rae et al., 2019) integrates this idea into Transformer-XL by compressing tokens in the input that are far away. The model benefits from detailed attention to nearby tokens, while using summarized information for more distant tokens (see Figure 1, lower right).
ETC: Encoding Long and Structured Inputs in Transformers
Our model follows the original Transformer architecture (Vaswani et al., 2017), with key modifi-cations to tackle long and structured inputs: relative position encoding, global-local attention, and a CPC pre-training task, explained below. In this paper, we consider only the encoder side of the Transformer, and leave the decoder for future work.
ETC: Encoding Long and Structured Inputs in Transformers
Inspired by the work of Shaw et al. (2018), ETC replaces absolute position encodings with relative position encodings, which provide information about the relative position of tokens in the input sequence with respect to one another. Given the input sequence x = (x 1 , ..., x n ), we can see it as a labeled fully connected and directed graph, where l ij is the label of the edge that connects x i to x j . Given a maximum clipping distance k, Shaw et al. define 2k + 1 relative position labels: l −k , ..., l k . The label of the edge between two input tokens depends only on their relative position j − i. For input pairs with j − i ≥ k, label l k is given, and with j − i ≤ −k, l −k is given. Each label then becomes a learnable vector a K l , which modifies the attention mechanism (equations in the next section) 3 .
ETC: Encoding Long and Structured Inputs in Transformers
Relative position encodings are independent of input length, so it is easy to adapt a model to greater input lengths than seen during pre-training. As other recent work (Shaw et al., 2019), ETC's attention mechanism uses relative position labels not just for relative positions in a sequence but also to express arbitrary pairwise token relations useful for Figure 2: Sparsity diagram showing which attention queries (rows) can attend to which attention keys (columns) a) for standard Transformer attention with input size n; b) for global-local attention with input sizes n g , n l , and radius r; c) how the l2l attention piece is reshaped into a much smaller attention matrix, limited by local radius. structured inputs, as explained below.
ETC: Encoding Long and Structured Inputs in Transformers
Global-local attention is a generalization of several of the models presented above. ETC receives two separate input sequences: the global input x g = (x g 1 , ..., x g ng ) and the long input x l = (x l 1 , ..., x l n l ). Typically, the long input contains the input a standard Transformer would receive, while the global input contains a much smaller number of auxiliary tokens (n g n l ). Attention is then split into four separate pieces: global-to-global (g2g), global-tolong (g2l), long-to-global (l2g), and long-to-long (l2l). Attention in the l2l piece (the most computationally expensive piece) is restricted to a fixed radius r n l . To compensate for this limited attention span, the tokens in the global input have unrestricted attention, and thus long input tokens can transfer information to each other through global input tokens. Accordingly, g2g, g2l, and l2g pieces of attention are unrestricted.
ETC: Encoding Long and Structured Inputs in Transformers
This concept is illustrated in Figure 2, where each cell (row i, column j) is shaded grey if token x i can attend to token x j . As we can see, in a regular Transformer, attention is unrestricted (full n × n attention). ETC, illustrated in Figure 2b, however, restricts the l2l piece to a local radius, significantly reducing computational and memory complexity for very long inputs. Conceptually, the l2l attention piece is reshaped into a n l × (2r + 1) matrix as illustrated in Figure 2c. 4 If r = 1 and n g = 1, we recover exactly the Star Transformer (Section 2). Similarly, placing all the tokens in the global input and setting n l = 0 yields standard Transformer attention. Attention in ETC is O(n g (n g + n l ) + n l (n g + 2r + 1)). If we assume n g = O(2r + 1), we see attention is linear in the size of the long input: O(n 2 g + n g n l ). To provide flexible attention and help with structured inputs, per-instance Boolean attention matrices M g2g , M g2l , M l2g , and M l2l exist, with zeroes for those pairs of tokens that should not attend to one another. Each g2g attention head works as follows. Given the global input x g = (x g 1 , ..., x g ng ), which is a sequence of token representations x g i ∈ R dx , the output of attention is z g = (z g 1 , ..., z g ng ), where each z g i ∈ R dz is calculated as follows:
ETC: Encoding Long and Structured Inputs in Transformers
where: M g2g is a binary attention mask, W Q , W K , and W V are learnable weight matrices, and a K ij are learnable vectors representing the relative position labels, and C is a large constant (C = 10000 in our experiments to follow the same convention as BERT). Attention for the other 3 pieces is analogous. We experiment with having separate W K and W V across all four attention pieces, or sharing them. And for W Q , we experiment with having one for g2g and g2l, and a separate one for l2g and l2l; or sharing them also. To recover BERT as a special case when r is large enough to remove sparsity, attention is actually only split into 2 pieces internally instead of 4, as g2g+g2l can be computed jointly (top half of Figure 2c), and l2g+l2l can also be computed jointly (bottom half of Figure 2c). A single softmax is used to jointly calculate α g2g ij and α g2l ij , and another for α l2g ij and α l2l ij . Thus, the output of global-local attention is a sequence of length n g and one of length n l . These sequences go through a layer normalization and feed forward layer in the same way as in the standard transformer.
ETC: Encoding Long and Structured Inputs in Transformers
Let us illustrate how ETC can be used to encode long inputs. A general way to handle long inputs in ETC is to place the entire sequence of input tokens (e.g., word pieces) in the long input, and then assuming some sort of division into segments (e.g., sentences), place one auxiliary token in the global input per segment in the long input. We then use one relative position label to link the global segment tokens with the word piece tokens that belong to them, and a different label for those that do not. Moreover, as we will show in the experiments below, we have seen that using the M g2l attention masks to perform hard masking in one direction (g2l) can bring performance gains in some datasets. This last asymmetric hard-masking is illustrated in Figure 3a, where we used different colors to indicate different relative position labels. In this way, although tokens in the long input can only attend to the local neighborhood defined by the radius k, they can indirectly attend to all other tokens in the input sequence via the global tokens.
ETC: Encoding Long and Structured Inputs in Transformers
A standard Transformer resembles a graph neural network (Scarselli et al., 2008) over a fully connected graph g; see Ye et al. (2019). Thanks to the combination of global-local attention and relative position labels, ETC exploits this relation to encode structured inputs. Given the input x = (x 1 , ..., x n ), we use the term structure to refer to the relations that exist between the tokens in x. When x is a plain ordered sequence, the only relation is the sequential order of tokens, which is the only structure captured by BERT (encoded by absolute position encodings, used to modify attention). We define structured inputs as those that have additional relations between the tokens beyond sequential order. In principle, we could think of inputs with arbitrary graph structure (such as chemical molecule graphs), but here we focus on structure in NLP tasks.
ETC: Encoding Long and Structured Inputs in Transformers
ETC is particularly well suited to capture hierarchical structure thanks to three mechanisms. First, as originally conceived, the vocabulary of relative position labels is used to represent token relative positions. However, seeing a Transformer as a graph neural network over a graph g (with one vertex per token in x, and edges representing their relations), we can expand this vocabulary to label some edges with labels for relations such as is-a, part-of, or others. Second, the division between long and global input induces a natural structure where the global input contains summary tokens of sets of tokens in x (a 2-level hierarchy). However, we can also have tokens summarizing sets of summary tokens (constructing a 3-level hierarchy, or beyond). Third, if some pairs of tokens should not have an edge between them, this can be captured with the M g2g , M g2l , M l2g , M l2l masks. An illustration of all these concepts is shown in Figure 3b, which uses masking and relative position labels to represent a context-sentence-token hierarchy that includes within-context order of sentences but no order between contexts. Another example would be social community graphs structure, where we could partition the graph into components, use M l2l to constrain attention to within components, and add per-component global tokens, linked to allow information to propagate from one component to another in a hierarchical way.
ETC: Encoding Long and Structured Inputs in Transformers
We use two pre-training tasks: (1) a masked language model (MLM) task with whole word masking (if one word piece token is masked, then all other tokens of the same word are masked); and (2) instead of using BERT's next sentence prediction (NSP), we adapt Contrastive Predictive Coding (CPC) (Oord et al., 2018) for ETC.
ETC: Encoding Long and Structured Inputs in Transformers
The goal of CPC is to predict subsequent inputs in latent space, i.e., to predict internal hidden representations of blocks of tokens. We adapted this idea in ETC by using global input sentence sum-mary tokens. Given an input sequence containing n sentences, we mask all the tokens corresponding to a subset of sentences (but leave the sentence summary tokens in the global input). We then train the model to minimize the difference between the hidden representation of the global sentence summary tokens for the masked sentences with respect to that of a global summary token that can see the unmasked sentence and nothing else. We use a Noise Contrastive Estimation (NCE) loss as in the work of Oord et al. (2018) (details in the appendices).
ETC: Encoding Long and Structured Inputs in Transformers
Having described ETC, we can now compare it with Longformer (Beltagy et al., 2020), which uses a similar attention mechanism, except Longformer has a single input sequence with some tokens marked as global (the only ones that use full attention). The key differences are that (1) ETC's combination of global-local attention with relative position encodings and flexible masking enables it to encode structured inputs in a similar way as graph neural networks do; (2) global tokens in Longformer are never pre-trained with anything like our CPC loss, and thus their use must be learned during fine-tuning.
ETC: Encoding Long and Structured Inputs in Transformers
ETC and BERT share enough similarities that BERT parameters are useful to perform a warm start. The parameters are compatible because the global-local attention mechanism includes BERT as a special case if the input is small enough or the local radius is large enough to eliminate sparsity. Moreover, when lifting weights from BERT into an ETC model with separate W Q , W K , and W V projection matrices, BERT's parameters are just copied over to the different matrices of ETC.
ETC: Encoding Long and Structured Inputs in Transformers
Although pre-training is still required to adapt the weights to use global tokens and relative position encodings, we show that initializing from RoBERTa results in significant performance improvements compared to pre-training from scratch. Specifically, we initialized from the RoBERTa checkpoints reported in the work of Rothe et al. (Rothe et al., 2020).
ETC: Encoding Long and Structured Inputs in Transformers
This section focuses on evaluating our two main contributions: (1) long inputs, and (2) structure in text inputs, as well as initialization from existing BERT models. We chose four datasets (Table 1) with long inputs or interesting input structure.  NQ (Kwiatkowski et al., 2019): in Google's Natural Questions (NQ) dataset the input consists of a question and a full Wikipedia article. The task is to identify both a short answer (a few words from the article) and a long answer (e.g., a whole paragraph), if they exist within the article (and otherwise, return null answers). Performance is measured based on the F1 score of the model predictions with respect to the human generated answers.
ETC: Encoding Long and Structured Inputs in Transformers
HotpotQA (Yang et al., 2018) is a question answering dataset where the goal is to combine evidence from multiple contexts. We use the distractor setting, where 10 paragraphs are provided: two of them contain useful information to answer the question, and the rest are distractors. The task is both to answer the question, and also to identify the supporting facts that are relevant to answer the questions (at a sentence granularity).
ETC: Encoding Long and Structured Inputs in Transformers
WikiHop (Welbl et al., 2018) is similar in structure to HotpotQA. The contexts correspond to portions of Wikipedia articles, and the goal is to answer about properties of an entity that cannot be found in the entity's article. Each instance contains a query, a collection of candidate answers, and a collection of contexts from which to obtain information to select among the candidate answers.
ETC: Encoding Long and Structured Inputs in Transformers
OpenKP (Xiong et al., 2019) is a keyphrase extraction dataset. Each document contains up to 3 short keyphrases to be identified. We selected this dataset as the input is not flat text sequences, but websites, including the hierarchical and spatial relations between the different DOM elements on the website, as well as other visual properties.
ETC: Encoding Long and Structured Inputs in Transformers
We use two basic configurations: base and large. Base uses 12 layers, 768 hidden size, 12 attention heads, local attention radius r = 84, and relative position maximum distance k = 12. Large uses 24 layers, 1024 hidden size, 16 heads, r = 169, and k = 24. We used 128, 230 and 460 global tokens for models with 512, 4096 and 8192 long  Pre-training: We place all word piece tokens in the long input and add one auxiliary token per sentence to the global input. We defaulted to BERT's 30k English uncased word piece vocabulary. Models were pre-trained using the original BERT datasets, except that documents with fewer than 7 sentences were filtered out. Unless stated otherwise, base models were pre-trained with the same total number of tokens as the original BERT, and for large models, twice as many. We used the LAMB optimizer (You et al., 2019) with learning rate set to √ 8 × 10 −3 . Fine-tuning: we put all input tokens in the long input (CLS, question, and context tokens for QA datasets), and use relative position labels to encode structure (see Section 3.4). Global input has a CLS token, tokens mirroring the question tokens in long, and one summary token per paragraph/sentence (or VDOM block in OpenKP). OpenKP had no CLS nor question tokens. For WikiHop, we also add one global token per candidate answer, and used a different relative position label to link these tokens to their string-matched mentions in the text (more details in the appendices).
ETC: Encoding Long and Structured Inputs in Transformers
NQ: We used NQ to study the different parts of ETC via ablations. Results are shown in Table 2. The first three rows show baseline models: BERTbase, BERT-large, and RikiNet (Liu et al., 2020) (one of the best models in the NQ leaderboard). 5 With gradient checkpointing, ETC can scale beyond this, but we limit our experiments to 8192 tokens for this paper.
ETC: Encoding Long and Structured Inputs in Transformers
BERT's performance is comparable to ETC using input length of 512. The smaller local radius of ETC (84) puts ETC at a disadvantage with respect to BERT, but other ETC improvements, such as dynamic whole word masking seem to compensate.
ETC: Encoding Long and Structured Inputs in Transformers
The rest of Table 2 shows performance under different ablations. Our default configuration (marked with a "-" in the configuration column) is ETC-base with long input length of 4096 tokens, using CPC, hard g2l masking, and separate W Q , W K , and W V matrices for long/global inputs. We tested the following ablations: shared (sharing all model parameters for attention across both the global and long inputs), no CPC (removing the CPC pre-training task), no hard g2l (not having a hard g2l mask), and fixed blocks (which configures the global input to just have one global token per 97 long input tokens, to keep the same proportion as without fixed-blocks, ignoring sentence boundaries, and not having any other tokens in the global input for pre-training or fine-tuning). Sharing W Q , W K , and W V and removing CPC significantly hurt the performance of ETC in NQ 7 . Using fixed blocks, surprisingly, seems to slightly help without CPC.
ETC: Encoding Long and Structured Inputs in Transformers
Increasing long input from 512 to 4096 significantly helped performance, and going to 8192 increased performance further to 0.740 / 0.542, highlighting the importance of longer inputs. Increasing the local radius, relative position vocabulary, or the amount of pre-training all helped performance (especially the latter, reaching 0.746 / 0.558). Moving to a large model also helped, especially when lifting from RoBERTa (both large models used the 7 Separate projection matrices were also found to be helpful in other models, like Longformer (Beltagy et al., 2020).     Table 3 shows our results in HotpotQA and WikiHop. We show two Longformer models as baselines (which is currently the state-of-the-art model in WikiHop), as well as ablations to study the effect of structure in the results. In particular, we consider a flat structure ablation where: (1) we do not break long input attention by context boundaries, (2) we limit relative position labels between global and long tokens to representing only sentence-level relationships (this removes any special attention in WikiHop between candidate answers and their mentions).
ETC: Encoding Long and Structured Inputs in Transformers
Our results show that both our base and large models outperform their corresponding Longformer models in both HotpotQA and WikiHop. Besides parameter counts, the main factors that can explain this difference in performance are the different pre-training strategies and the different handling of structure in ETC and Longformer. Removing the CPC pre-training task, and not using a hard g2l mask significantly hurt the performance of the model in HotpotQA, going from a performance of 0.751 / 0.869 for the baseline model to 0.722 / 0.857 using none of those features. Using a flat structure (but keeping CPC and hard g2l) did not seem to hurt in HotpotQA. WikiHop shows a slightly different picture, as it seems that hard g2l masking and especially flat structure hurt performance in this dataset. Our best model is the base configuration without hard g2l masking, which achieves an accuracy of 75.9. Interestingly, sharing W Q , W K , and W V seems to help performance in WikiHop. This is our smallest dataset, and maybe the added capacity of the model without sharing parameters leads it to overfit.
ETC: Encoding Long and Structured Inputs in Transformers
OpenKP: Table 4 shows our results on the OpenKP dataset, using RoBERTa-JointKPE (Sun et al., 2020) as the baseline, which is currently #1 in the leaderboard. This is an interesting structured dataset, and thus, we performed additional ablations to investigate the effect of removing such structural information. Our results show that even the most constrained ETC model already achieves very good performance (0.399), and scaling to 4096 length seems to give a slight boost. Using hard g2l also helps, and adding the visual features brings the largest benefit. Finally, we see that using a large model, and especially lifting weights from RoBERTa improve results significantly. As with WikiHop, sharing W Q , W K , and W V does not hurt performance. Our default model uses the first occurrence of a keyphrase, but we saw that using the maximum logit of all occurrences (max loss) improved results.
ETC: Encoding Long and Structured Inputs in Transformers
Finally, Table 5 shows official results on the leaderboards of each dataset. The model submitted to the leaderboards was the model with best dev set results (shown at the bottom of the respective results tables, lifting weights from RoBERTa). We set a new state of the art in WikiHop and OpenKP, NQ long answer, and HotpotQA Support F1. Remarkably, our submissions were all single model, outperforming the leaderboard ensemble models.
ETC: Encoding Long and Structured Inputs in Transformers
This paper introduced the Extended Transformer Construction (ETC), an architecture designed to (1) scale up the input length (linearly with input), and (2) encode structured inputs. ETC allows lifting weights from existing BERT models, improving results significantly. The key ideas are a new globallocal attention mechanism, coupled with relative position encodings and a CPC pre-training task.
ETC: Encoding Long and Structured Inputs in Transformers
We showed that significant gains can be obtained thanks to increased input sequence length. The ability to represent dataset structure in ETC further improves the model quality. We hypothesize that CPC helps the model train the usage of the higherlevel global input summary tokens, as CPC plays a role akin to MLM, but at the global input level. Notice that although our datasets contain a limited amount of structure (compared to graph datasets), our experiments show that ETC was able to exploit this existing structure.
ETC: Encoding Long and Structured Inputs in Transformers
As future work, we would like to investigate complementary attention mechanisms like those of Reformer (Kitaev et al., 2020) or Routing Transformer (Roy et al., 2020), push scalability with ideas like those from RevNet (Gomez et al., 2017), and study the performance of ETC in datasets with even richer structure. Appendix A: Implementation Details
ETC: Encoding Long and Structured Inputs in Transformers
This appendix provides further details on the TPU/GPU-friendly implementation of global-local attention. Our implementation of sliding window local attention is similar to the approach in the local attention 1d layer in Ten-sor2Tensor 8 , but with the addition of flexible masking, relative position encoding, and global tokens as side keys/values. We use a simple example to describe the internal blocking logic. Let's say the input corresponds to embeddings for the following word pieces, each represented by a letter: ABCDEFG.
ETC: Encoding Long and Structured Inputs in Transformers
As usual, we project these embeddings into queries, keys, and values, yielding the following (for each attention head):
ETC: Encoding Long and Structured Inputs in Transformers
Queries:
ETC: Encoding Long and Structured Inputs in Transformers
Let's say we want to perform sliding window local attention with local radius r = 2. Internally, we split the input into blocks of length r + 1 (3 in our example) and add padding blocks to the left and right, resulting in the following five blocks for the queries (and similarly for keys and values), with 0 representing padding: 000 A q B q C q D q E q F q G q 00 000 Conceptually we'd like to compare each query with the 2r + 1 (5 in our example) surrounding keys, as follows:
ETC: Encoding Long and Structured Inputs in Transformers
But materializing each window of keys would be memory-intensive. Instead, we allow each block of queries to attend to 3 blocks of keys (the same block, and the blocks immediately to the left and right), resulting in the following:
ETC: Encoding Long and Structured Inputs in Transformers
Queries Keys
ETC: Encoding Long and Structured Inputs in Transformers
Now each query can potentially see a few more tokens than it's strictly allowed to by the local ra-8 https://arxiv.org/abs/1803.07416 dius r. For example, A q takes a dot product with D k , E k and F k , but this is easy to simply mask out, yielding the same sliding window local attention result. In this way, the blocking mechanism saves memory at the expense of some extra compute.
ETC: Encoding Long and Structured Inputs in Transformers
The values are also divided into the same blocks as the keys (concatenating 3 at a time), and standard scaled dot product attention is applied independently for each row in the table below, where Keys have been truncated for brevity:
ETC: Encoding Long and Structured Inputs in Transformers
Queries Keys Values
ETC: Encoding Long and Structured Inputs in Transformers
Efficient Relative Attention Implementation
ETC: Encoding Long and Structured Inputs in Transformers
To efficiently implement relative position encoding (a.k.a. relative attention), we take an approach similar to the optimization in Music Transformer (Huang et al., 2018) but generalized to allow arbitrary pairwise labels rather than adhering to a relative position pattern. We briefly describe our implementation in the case of full attention (with a single sequence length n), but the approach naturally extends to the case of the four attention pieces used in ETC. The original relative attention work (Shaw et al., 2018) reported O(n 2 d z ) memory overhead (Section 3.3) by materializing a K ij for every query-key pair while sharing a K ij across all heads (or O(hn 2 d z ) if not sharing across heads), where d z is the dimension per head and h is the number of heads. We instead take a dot product between each query vector and all unique a K ij vectors in the relative attention vocabulary. Then we gather these scalar results for each query-key pair. This avoids the O(n 2 d z ) memory overhead and allows us to use different a K ij per attention head with no additional activation memory cost. Note that our relative attention vocabulary sizes are noticeably smaller than n, so our implementation reduces the number of dot products required for relative attention also.
ETC: Encoding Long and Structured Inputs in Transformers
We adapted the original formulation of CPC for ETC by modeling it as a dual encoder problem. We have two instances of the same ETC model g 1 and g 2 (using the same weights). g 1 is the main encoder we are training, and we divide its long input into segments (e.g., sentences) and have one global token in global input for each segment. We mask some segments in the long input, and encode those segments independently using g 2 (by having as input just the tokens of that segment in the long input, and a single global token in the global input). Then, we train g 1 and g 2 so that the encodings of the global tokens corresponding to the masked segments should be as similar as possible as the encoding of the global segment token obtained via g 2 . We use within-batch random negatives for this process as well, and use a Noise Contrastive Estimation (NCE) loss, in the same way as in the original CPC work (Oord et al., 2018).
ETC: Encoding Long and Structured Inputs in Transformers
Our default pre-training procedure used the same Wikipedia and Books training corpora as BERT, but we filtered to remove those documents with fewer than 7 sentences. Models were pre-trained for 33 epochs to match the amount of pre-training of the original BERT model, which used batches of 256 sequences of 512 tokens each, and pre-trained for 1,000,000 iterations. Instead, we used batches of 512 sequences of 4096 tokens each and pre-trained for 63,000 iterations. The ETC-large models were pre-trained for 66 epochs by using a batch size of 1024 instead. When lifting weights from RoBERTa, we found that decreasing the learning rate to 2 × 10 −3 improved model quality.
ETC: Encoding Long and Structured Inputs in Transformers
When pre-training models, we split any input documents that are longer than the long input length. For efficiency, we also concatenate as many shorter documents as will fit into the 512/4096/8192 window and mask attention to prevent them from influencing each other. This results in a roughly 3x speedup in pre-training time for 4096-token models, highlighting once more the advantage of flexible masking.
ETC: Encoding Long and Structured Inputs in Transformers
When pre-training with CPC, we randomly select 10% of sentences to be masked for the CPC task. Subsequently, 15% of the remaining tokens are masked for MLM.
ETC: Encoding Long and Structured Inputs in Transformers
In the models where we use both MLM and CPC, we used a 0.8 weight for MLM and a 0.2 weight for CPC to combine them into a single loss.
ETC: Encoding Long and Structured Inputs in Transformers
Data Download Link: https://ai.google.
ETC: Encoding Long and Structured Inputs in Transformers
Data Pre-processing: Following Alberti's BERT implementation (Alberti et al., 2019), long input in NQ contains a CLS token followed by the question word pieces, then a separator followed by the long document, a final separator, and padding. Global input contains a CLS token, one special "question" token per token in the question, and then one special "segment" token per paragraph (long answer candidate) in the long input. Moreover, since the ground truth indexes in this dataset are word indexes, in order to be able to align tokens with words, sentences are first tokenized by words, and then each word is given to the BERT/RoBERTa tokenizer.
ETC: Encoding Long and Structured Inputs in Transformers
Fine-Tuning: After pre-training, all models were fine-tuned with a hyperparameter sweep consisting of learning rates in {3 × 10 −5 , 5 × 10 −5 } and number of epochs in {3, 5} ({2, 3} for large models) with a batch size of 64 on the NQ training set using the Adam optimizer. The model is trained to predict four logits coming out of the long input tokens: long answer start, long answer end, short answer start, and short answer end. A final prediction (predicted from the long input CLS token embedding) is the answer type (null, yes, no, short, long). For NQ instances that are longer than long input size, a sliding window approach is used (with stride 128 for input lengths of 512, 2048 for input lengths of 4096, and 4096 for input lengths of 8192).
ETC: Encoding Long and Structured Inputs in Transformers
Model Selection: we performed a single hyperparameter sweep for each model (i.e., we tested a single random seed per parameter configuration). The best model was selected as the highest average F1 score on dev at the end of fine-tuning. Our best model (large, lifting from RoBERTa) was trained for 2 epochs with learning rate 3 × 10 −5 .
ETC: Encoding Long and Structured Inputs in Transformers
Inference: Final predictions are then aggregated similarly as in the work of Alberti et al. (2019), but with two improvements: First, instead of predicting start and end of short/long answer separately, we first select the best start position, and then select the best end location that occurs after the start position. For short answer, we also filter out all end positions further than 38 words from the start position. Second, when the logit for a yes/no answer is higher than the logits for short, long or null answer, we replace the short answer with a corresponding yes/no.
ETC: Encoding Long and Structured Inputs in Transformers
Data Download Link: https://hotpotqa. github.io/ Data Pre-processing: Only 90,431 out of the 90,447 instances were used for training, as we model the task as extractive QA and thus filtered out the 16 instances where the answer could not be found in the contexts. Long input in HotpotQA is organized as follows: CLS token followed by question tokens followed by all the context tokens. Each context is represented as the concatenation of its title, and then all the sentences. Global input has a CLS token, and then one token per question token (as in NQ), followed by global tokens representing the contexts. For every context, in global, we have one token representing the whole context, and then one per sentence. We did not use any windowing approach for longer instances, and just fit as many tokens as possible within the 4096 long input. Global input length was set to 256.
ETC: Encoding Long and Structured Inputs in Transformers
Fine-Tuning: After pre-training, all base models were fine-tuned with a hyperparameter sweep consisting of learning rates in {3 × 10 −5 , 5 × 10 −5 }, number of epochs in {3, 5, 7, 9}, batch size in {32, 64}, and supporting fact threshold in linspace(0, 1, 11) on the training set using the Adam optimizer. Large models were tested with learning rate in {1 × 10 −5 , 2 × 10 −5 , 3 × 10 −5 , 5 × 10 −5 , 7×10 −5 }, number of epochs in {2, 5, 9, 13}, and batch size in {32, 64, 128}.
ETC: Encoding Long and Structured Inputs in Transformers
Model Selection: we performed a single hyperparameter sweep to determine the best parameter configuration for each model, and then we tried 3 different random seeds for the best configuration. The best model was selected as the one with the best joint F1 score on dev. Our best model (large, lifting from RoBERTa) was trained for 5 epochs, with a learning rate of 3 × 10 −5 , batch size of 32 and supporting fact threshold of 0.4.
ETC: Encoding Long and Structured Inputs in Transformers
In order to make predictions, supporting facts are predicted using a single dense layer taking the global input embeddings as input with a threshold over the output logits. Output type is predicted with a single dense layer from the global CLS token. Answer spans where predicted also with dense layers, but using the long input embeddings as inputs, using the following criteria: begin/end positions must be in sentences or titles, begin/end must be in the same sentence/title, spans must belong to a supporting fact, begin must be before end, and spans cannot exceed a maximum answer length of 30 tokens. Within spans satisfying those criteria, a single span with top begin prob * end prob is selected.
ETC: Encoding Long and Structured Inputs in Transformers
Data Download Link: https://qangaroo.cs. ucl.ac.uk/ Data Pre-processing: Global and Long input was set similarly as in HotpotQA, except that global input was set to 430, and that instead of a CLS token in global, we have one token per candidate answer (WikiHop provides a list of candidate answers, and the model needs to select among them). We used a relative position label (the same used to link sentence summary tokens with its corresponding tokens) to link candidate answers to their mentions in the text, where mentions are determined only by string matching. Also, as in HotpotQA, no sliding window was used, and instances were just cropped to a length of 4096. A MaxHeap was used to ensure that in case a context is truncated, truncation happens from the contexts with the larger number of sentences.
ETC: Encoding Long and Structured Inputs in Transformers
Fine-Tuning: After pre-training, all base models were fine-tuned with a hyperparameter sweep consisting of learning rates in {1 × 10 −5 , 2 × 10 −5 , 3 × 10 −5 , 4 × 10 −5 , 5 × 10 −5 }, and number of epochs in {5, 10, 15} with a batch size of 64 on the training set using the Adam optimizer. For large models, we narrowed down the hyperparameter sweep to learning rates in {2 × 10 −5 , 3 × 10 −5 , 4 × 10 −5 , 5 × 10 −5 }, and number of epochs in {5, 10}. For this dataset we also experimented with the LAMB optimizer (in addition to Adam), which was used for our leaderboard submission.
ETC: Encoding Long and Structured Inputs in Transformers
Model Selection: we performed a single hyperparameter sweep to determine the best parameter configuration for each model (i.e., single random seed per parameter configuration). The best model was selected as the one with the highest accuracy on dev. Our best model (large, lifting from RoBERTa) was trained for 10 epochs, with a learning rate of 5 × 10 −5 . Finally, for the final leaderboard submission, we selected the checkpoint of the model that had the highest dev set accuracy.
ETC: Encoding Long and Structured Inputs in Transformers
Inference: For final prediction, we used a dense layer from the global input embeddings, after that, the candidate with the highest logit is selected as the final prediction.
ETC: Encoding Long and Structured Inputs in Transformers
Data Download Link: https://microsoft. github.io/msmarco/ Data Pre-processing: Long input in OpenKP contains all the word pieces of the input. One global token per VDOM node was added to the global input (notice this is like the pre-training setup, except instead of sentences we have VDOM nodes as the higher-level units). No sliding windowing was used, and we simply truncate instances to whichever of max tokens in long input or max VDOM tokens in global ends up being more constraining. Long input length was set to 4096 and global input length to 512 by default in this dataset. After url deduplication and skipping examples without keyphrases in the truncated document, there were 133,374 valid training examples. Regarding visual features, we embed font sizes based on 24 bucket ranges, and we also construct an embedding for the cross of "block", "heading", and "bolded" Boolean features in the input data. All other visual features were treated as dense features, with the floating point features clipped to reasonable ranges and re-scaled to the [−1, 1] interval. These dense features are then transformed to the same embedding space as the other embeddings, and all visual feature embeddings are added to both the relevant long and global input tokens.
ETC: Encoding Long and Structured Inputs in Transformers
Fine-Tuning: After pre-training, all models were fine-tuned with a hyperparameter sweep consisting of learning rates in {3 × 10 −5 , 5 × 10 −5 } and number of epochs in {2, 3} with a batch size of 64 on the OpenKP training set using the Adam optimizer. To generate predictions, we first sum the embeddings (from the long input) of all the word pieces for each word to form word embeddings. Then we run convolutions with kernel size 1, 2, 3, 4, and 5 to form the respective n-gram embeddings. Finally, a dense linear layer is used to form logits for all the n-grams and concatenate them together for one combined softmax. The loss is cross entropy where the ground truth probabilities are divided equally among the keyphrase labels (up to 3). By default we used the first occurrence of each keyphrase as the label. Our improved "max loss" takes the max of logits across all occurrences of the same keyphrase in the text, rather than just the first occurrence.
ETC: Encoding Long and Structured Inputs in Transformers
Model Selection: we performed a single hyperparameter sweep to determine the best parameter configuration for each model (i.e., single random seed per parameter configuration). The best model was selected as the one with the highest F1@3 on dev. Our best model (large, max loss, lifting from RoBERTa) was trained for 2 epochs, with a learning rate of 3 × 10 −5 .
ETC: Encoding Long and Structured Inputs in Transformers
Time Hardware ETC (share) 11h 13m 256 core TPU v3 ETC 11h 46m 256 core TPU v3 ETC-large 63h 41m 512 core TPU v3  Inference: During inference, we select the top 5 keyphrases ordered by logits, removing any duplicates. All keyphrases were treated as uncased for the purpose of deduplication.
ETC: Encoding Long and Structured Inputs in Transformers
When lifting weights from BERT or RoBERTa, the weights that can be lifted are (for every Transformer layer): feed forward layer, W Q , W K , W V (since BERT/RoBERTa only have one copy of such matrices, in models where we use different matrices for global and long inputs, we initialize both sets of matrices to the same BERT/RoBERTa weights), attention output projection, and layer normalization. Additionally, we can also lift the token embedding matrix. Absolute position embeddings and next sentence prediction weights from BERT are discarded. After that, weights for the layers necessary for the CPC loss, and those involved in relative position encodings are randomly initialized.
ETC: Encoding Long and Structured Inputs in Transformers
For lifting to be possible, the number of layers, hidden size, number of attention heads, and size of the feed forward intermediate layers of the BERT/RoBERTa model need to match with the ETC model.   Figure 5: Illustration of structure encoding with ETC (using the same example shown in Figure 3b). Top: each box represents an input token, and arrows represent attention. The different colors and dash patterns in arrows represent the different relative position labels. Bottom: illustration of where each token would appear in the input of ETC.
ETC: Encoding Long and Structured Inputs in Transformers
Memory: To gauge headroom for scaling input lengths beyond what we used in this paper, we ran some additional experiments on TPU v3 hardware with gradient checkpointing and removing the extra gradient moments required by optimizers like Adam and LAMB. Fixing global input length to 512 tokens, we were able to push base models to long input size of 22656, and large models to long input size of 8448 before running out of memory on a single TPU v3 core. We leave for future work experimentation with more memory-efficient optimizers like Adafactor (Shazeer and Stern, 2018) and model-parallelism techniques in ETC.
ETC: Encoding Long and Structured Inputs in Transformers
Compute: As stated above, the computational complexity of attention in ETC is O(n g (n g + n l ) +  n l (n g +2r +1)) and if we assume n g = O(2r +1), this results in a complexity of O(n 2 g +n g n l ), which is linear in the size of the long input. Table 6 shows pre-training times in our experiments, together with the hardware used in each experiment. Table 7 shows the fine-tuning times taken by the baseline ETC model on the different datasets. Notice that pre-training is the most computational intensive part, and thus, we used significantly more hardware. In order to gain further insights into the common use case of running the models using GPUs, Figure 4 shows a comparison of the wall-time used per step when using a single NVIDIA Tesla V100 GPU as the input length increases, for both BERT and ETC in their base configurations. As the plot shows ETC is initially slower, but it becomes faster than BERT for input lengths larger than about 1500. Moreover, the BERT plot ends earlier due to memory constraints. Finally, notice that the ETC wall time is not linear in this figure, as we also increased the size of the global input together with the long input.
ETC: Encoding Long and Structured Inputs in Transformers
Parameters: Finally, Table 8 shows the total number of parameters of the ETC model for the different configurations used in our experiments. The most important consideration is that the number of trainable parameters does not depend on the input length. As a matter of fact, it only depends on: the embedding dimensionality (d), the number of layers (l), and the number of relative position labels (which depends on k). The parameter count also depends on the fully connected feed forward intermediate size, but this is 4d by convention and for all ETC models in this paper. Our baseline model uses separate W Q , W K , W V , and output projection matrices for global and long inputs, resulting in about 50% more parameters than BERT. But the configuration with shared W Q , W K , W V , and output projection has a similar number of parameters as BERT. Parameter count for BERT base is reported from the original paper (Devlin et al., 
Big Bird: Transformers for Longer Sequences
Models based on Transformers [92], such as BERT [23,64], are wildly successful for a wide variety of Natural Language Processing (NLP) tasks and consequently are mainstay of modern NLP research. Their versatility and robustness are the primary drivers behind the wide-scale adoption of Transformers. The model is easily adapted for a diverse range of sequence based tasks -as a seq2seq model for translation [92], summarization [67], generation [16], etc. or as a standalone encoders for sentiment analysis [84], POS tagging [66], machine reading comprehension [94], etc. -and it is known to vastly outperform previous sequence models like LSTM [38]. The key innovation in Transformers is the introduction of a self-attention mechanism, which can be evaluated in parallel for each token of the input sequence, eliminating the sequential dependency in recurrent neural networks, like LSTM. This parallelism enables Transformers to leverage the full power of modern SIMD hardware accelerators like GPUs/TPUs, thereby facilitating training of NLP models on datasets of unprecedented size. This ability to train on large scale data has led to surfacing of models like BERT [23] and T5 [76], which pretrain transformers on large general purpose corpora and transfer the knowledge to down-stream task. The pretraining has led to significant improvement in low data downstream tasks [52] as well as tasks with sufficient data [102] and thus have been a major force behind the ubiquity of transformers in contemporary NLP.
Big Bird: Transformers for Longer Sequences
The self-attention mechanism overcomes constraints of RNNs (namely the sequential nature of RNN) by allowing each token in the input sequence to attend independently to every other token in the sequence. This design choice has several interesting repercussions. In particular, the full self-attention have computational and memory requirement that is quadratic in the sequence length. Using commonly available current hardware and model sizes, this requirement translates to roughly being able to handle input sequences of length 512 tokens. This reduces its direct applicability to tasks that require larger context, like question answering [61], document summarization [21], etc.
Big Bird: Transformers for Longer Sequences
However, while we know that self-attention and Transformers are useful, our theoretical understanding is rudimentary. What aspects of the self-attention model are necessary for its performance? What can we say about the expressivity of Transformers and similar models? Apriori, it was not even clear from the design if the proposed self-attention mechanism was as effective as RNNs. For example, the self-attention does not even obey sequence order as it is permutation equivariant. This concern has been partially resolved, as Yun et al. [105] showed that transformers are expressive enough to capture all continuous sequence to sequence functions with a compact domain. Meanwhile, Pérez et al. [73] showed that the full transformer is Turing Complete (i.e. can simulate a full Turing machine). Two natural questions arise: Can we achieve the empirical benefits of a fully quadratic self-attention scheme using fewer inner-products? Do these sparse attention mechanisms preserve the expressivity and flexibility of the original network?
Big Bird: Transformers for Longer Sequences
In this paper, we address both the above questions and produce a sparse attention mechanism that improves performance on a multitude of tasks that require long contexts. We systematically develop BigBird, an attention mechanism whose complexity is linear in the number of tokens (Sec. 2). We take inspiration from graph sparsification methods and understand where the proof for expressiveness of Transformers breaks down when full-attention is relaxed to form the proposed attention pattern. This understanding helped us develop BigBird, which is theoretically as expressive and also empirically useful. In particular, our BigBird consists of three main part:
Big Bird: Transformers for Longer Sequences
• A set of g global tokens that attend on all parts of the sequence. • For each query q i , a set of r random keys that each query will attend to.
Big Bird: Transformers for Longer Sequences
• A block of local neighbors w so that each node attends on their local structure.
Big Bird: Transformers for Longer Sequences
Our design leads to empirically high performing attention mechanism scaling to much longer sequence lengths (8x) on standard hardware (∼16GB memory) for large size models. Use of gradient checkpointing [15] could potentially allow for handling even longer sequences.
Big Bird: Transformers for Longer Sequences
To summarize, our main contributions are:
Big Bird: Transformers for Longer Sequences
1. BigBird satisfies all the known theoretical properties of full transformer (Sec. 3). In particular, we show that adding extra tokens allows one to express all continuous sequence to sequence functions with only O(n)-inner products. Furthermore, we show that under standard assumptions regarding precision, BigBird is Turing complete. 2. Empirically, we show that the extended context modelled by BigBird greatly benefits variety of NLP tasks. In particular, we achieve state of the art results for question answering and document summarization on a number of different datasets. Summary of these results are presented in Sec. 4.
Big Bird: Transformers for Longer Sequences
3. Lastly, we introduce a novel application of attention based models where long contexts are beneficial: extracting contextual representations of genomics sequences like DNA. With longer masked LM pretraining, BigBird improves performance on downstream tasks such as promoter-region and chromatin profile prediction (Sec. 5).
Big Bird: Transformers for Longer Sequences
In this section, we describe the BigBird model using the generalised attention mechanism that is used in each layer of transformer operating on an input sequence X = (x 1 , ..., x n ) ∈ R n×d . The generalized attention mechanism is described by a directed graph D whose vertex set is [n] = {1, . . . , n}. The set of arcs (directed edges) represent the set of inner products that the attention mechanism will consider. Let N (i) denote the out-neighbors set of node i in D, then the i th output vector of the generalized attention mechanism is defined as 
Big Bird: Transformers for Longer Sequences
where Q h , K h : R d → R m are query and key functions respectively, V h : R d → R d is a value function, σ is a scoring function (e.g. softmax or hardmax) and H denotes the number of heads. Also note X N (i) corresponds to the matrix formed by only stacking {x j : j ∈ N (i)} and not all the inputs.
Big Bird: Transformers for Longer Sequences
If D is the complete digraph, we recover the full quadratic attention mechanism of Vaswani et al. [92]. To simplify our exposition, we will operate on the adjacency matrix A of the graph D even though the underlying graph maybe sparse. To elaborate, A ∈ [0, 1] n×n with A(i, j) = 1 if query i attends to key j and is zero otherwise. For example, when A is the ones matrix (as in BERT), it leads to quadratic complexity, since all tokens attend on every other token. This view of self-attention as a fully connected graph allows us to exploit existing graph theory to help reduce its complexity. The problem of reducing the quadratic complexity of self-attention can now be seen as a graph sparsification problem. It is well-known that random graphs are expanders and can approximate complete graphs in a number of different contexts including in their spectral properties [81,39]. We believe sparse random graph for attention mechanism should have two desiderata: small average path length between nodes and a notion of locality, each of which we discuss below.
Big Bird: Transformers for Longer Sequences
Let us consider the simplest random graph construction, known as Erdős-Rényi model, where each edge is independently chosen with a fixed probability. In such a random graph with justΘ(n) edges, the shortest path between any two nodes is logarithmic in the number of nodes [18,44]. As a consequence, such a random graph approximates the complete graph spectrally and its second eigenvalue (of the adjacency matrix) is quite far from the first eigenvalue [9,10,6]. This property leads to a rapid mixing time for random walks in the grpah, which informally suggests that information can flow fast between any pair of nodes. Thus, we propose a sparse attention where each query attends over r random number of keys. Mathematically this means that A(i, ·) = 1 for r randomly chosen keys (see Fig. 1a).
Big Bird: Transformers for Longer Sequences
The second viewpoint which inspired the creation of BigBird is that in most contexts within NLP and computational biology have data which displays a great deal of locality of reference. In this phenomenon, a great deal of information about a token can be derived from its neighboring tokens. Most pertinently, Clark et al. [20] investigated self-attention models in NLP tasks and concluded that that neighboring inner-products are extremely important. The concept of locality, proximity of tokens in linguistic structure, also forms the basis of various linguistic theories such as transformational-generative grammar. In the terminology of graph theory, clustering coefficient is a measure of locality of connectivity, and is high when the graph contains many cliques or near-cliques (subgraphs that are almost fully interconnected). Simple Erdos-Renyi random graphs do not have a high clustering coefficient [85], but a class of random graphs, known as small world graphs, exhibit high clustering coefficient [95]. A particular model introduced by Watts and Strogatz [95] is of high relevance to us as it achieves a good balance between average shortest path and the notion of locality. The generative process of their model is as follows: Construct a regular ring lattice, a graph with n nodes each connected to w neighbors, w/2 on each side. In other words we begin with a sliding window on the nodes. Then a random subset (k%) of all connections is replaced with a random connection. The other (100k)% local connections are retained. However, deleting such random edges might be inefficient on modern hardware, so we retain it, which will not affect its properties. In summary, to capture these local structures in the context, in BigBird, we define a sliding window attention, so that during self attention of width w, query at location i attends from i − w/2 to i + w/2 keys. In our notation, A(i, i − w/2 : i + w/2) = 1 (see Fig. 1b).  As an initial sanity check, we performed basic experiments to test whether these intuitions are sufficient in getting performance close to BERT like models, while keeping attention linear in the number of tokens. We found that random blocks and local window were insufficient in capturing all the context necessary to compete with the performance of BERT as seen in Tab. 1.
Big Bird: Transformers for Longer Sequences
The final piece of BigBird is inspired from our theoretical analysis (Sec. 3), which is critical for empirical performance. More specifically, our theory utilizes the importance of "global tokens" (tokens that attend to all tokens in the sequence and to whom all tokens attend to (see Fig. 1c). These global tokens can be defined in two ways:
Big Bird: Transformers for Longer Sequences
• BigBird-itc: In internal transformer construction (itc), we make some existing tokens "global", which attend over the entire sequence. Concretely, we choose a subset G of indices (with g := |G|), such that A(i, :) = 1 and A(:, i) = 1 for all i ∈ G. • BigBird-etc: In extended transformer construction (etc), we include additional "global" tokens such as CLS. Concretely, we add g global tokens that attend to all existing tokens. In our notation, this corresponds to creating a new matrix B ∈ [0, 1] (N +g)×(N +g) by adding g rows to matrix A, such that B(i, :) = 1, and B(:, i) = 1 for all i ∈ {1, 2, . . . g}, and B(g
Big Bird: Transformers for Longer Sequences
This adds extra location to store context and as we will see in the experiments improves performance.
Big Bird: Transformers for Longer Sequences
The final attention mechanism for BigBird ( Fig. 1d) has all three of these properties: queries attend to r random keys, each query attends to w/2 tokens to the left of its location and w/2 to the right of its location and they contain g global tokens (The global tokens can be from existing tokens or extra added tokens). We provide implementation details in App. D.
Big Bird: Transformers for Longer Sequences
In this section, we will show that that sparse attention mechanisms are as powerful and expressive as full-attention mechanisms in two respects. First, we show that when sparse attention mechanisms are used in a standalone encoder (such as BERT), they are Universal Approximators of sequence to sequence functions in the style of Yun et al. [105]. Second, we will show that sparse encoder-decoder transformers are Turning Complete (assuming the same conditions defined in [73]).
Big Bird: Transformers for Longer Sequences
Complementing the above positive results, we also show that moving to a sparse-attention mechanism incurs a cost, i.e. there is no free lunch. In Sec. 3.4, we show lower bounds by exhibiting a natural task where any sufficiently sparse mechanism will require polynomially more layers.
Big Bird: Transformers for Longer Sequences
The complete Transformer encoder stack is nothing but the repeated application of a singlelayer encoder (with independent parameters). We denote the class of such Transformer encoders stack, defined using the generalized encoder (Sec. 2), by T H,m,q D which consists of H-heads with head size m and q is the hidden layer size of the output network, and the attention layer is defined by the directed graph D.
Big Bird: Transformers for Longer Sequences
The key difference between our proposed attention mechanism to that of Vaswani et al.
Big Bird: Transformers for Longer Sequences
[92], Yun et al. [105] is that we add a special token at the beginning of each sequence and assign it a special vector. We will refer to this as x 0 . Therefore our graph D will have vertex set {0} ∪ [n] = {0, 1, 2, . . . , n}. We will assume that this extra node and its respective vector will be dropped at the final output layer of transformer. To avoid cumbersome notation, we will still treat transformer as mapping sequences X ∈ R n×d to R n×d . We will also allow the transformer to append position embeddings E ∈ R d×n to matrix X in the input layer.
Big Bird: Transformers for Longer Sequences
Finally, we need to define the function class and distance measure for proving universal approximation property. Let F CD denote the set of continuous functions f : [0, 1] n×d → R n×d which are continuous with respect to the topology defined by p norm. Recall for any p ≥ 1, Our main theorem is that the sparse attention mechanism defined by any graph containing S is a universal approximator:
Big Bird: Transformers for Longer Sequences
Theorem 1. Given 1 < p < ∞ and > 0, for any f ∈ F CD , there exists a transformer with sparse-attention, g ∈ T H,m,q D such that d p (f, g) ≤ where D is any graph containing star graph S.
Big Bird: Transformers for Longer Sequences
To prove the theorem, we will follow the standard proof structure outlined in [105].
Big Bird: Transformers for Longer Sequences
Step 1: Approximate F CD by piece-wise constant functions. Since f is a continuous function with bounded domain [0, 1) n×d , we will approximate it with a suitable piece-wise constant function. This is accomplished by a suitable partition of the region [0, 1) into a grid of granularity δ to get a discrete set G δ . Therefore, we can assume that we are dealing with a functionf :
Big Bird: Transformers for Longer Sequences
Step 2: Approximate piece-wise constant functions by modified transformers. This is the key step of the proof where the self-attention mechanism is used to generate a contextual-mapping of the input. Informally, a contextual mapping is a unique code for the pair consisting of a matrix (X, x i ) and a column. Its uniqueness allows the Feed forward layers to use each code to map it to a unique output column.
Big Bird: Transformers for Longer Sequences
The main technical challenge is computing the contextual mapping using only sparse attention mechanism. This was done in [105] using a "selective" shift operator which shift up entries that are in a specific interval. Key to their proof was the fact that the shift, was exactly the range of the largest entry to the smallest entry.
Big Bird: Transformers for Longer Sequences
Creating a contextual mapping with a sparse attention mechanism is quite a challenge. In particular, because each query only attends to a few keys, it is not at all clear that sufficient information can be corralled to make a contextual embedding of the entire matrix. To get around this, we develop a sparse shift operator which shifts the entries of the matrices if they lie in a certain range. The exact amount of the shift is controlled by the directed sparse attention graphg D. The second key ingredient is the use of additional global token. By carefully applying the operator to a set of chosen ranges, we will show that each column will contain a unique mapping of the full mapping. Therefore, we can augment the loss of inner-products in the self attention mechanism by using multiple layers and an auxiliary global token.
Big Bird: Transformers for Longer Sequences
Step 3: Approximate modified transformers by original Transformers: The final step is to approximate the modified transformers by the original transformer which uses ReLU and softmax.
Big Bird: Transformers for Longer Sequences
We provide the full details in App. A.
Big Bird: Transformers for Longer Sequences
Transformers are a very general class. In the original paper of Vaswani et al. [92], they were used in both an encoder and a decoder. While the previous section outlined how powerful just the encoders were, another natural question is to ask what the additional power of both a decoder along with an encoder is? Pérez et al. [73] showed that the full transformer based on a quadratic attention mechanism is Turing Complete. This result makes one unrealistic assumption, which is that the model works on arbitrary precision model. Of course, this is necessary as otherwise, Transformers are bounded finite state machines and cannot be Turing Complete.
Big Bird: Transformers for Longer Sequences
It is natural to ask if the full attention mechanism is necessary. Or can a sparse attention mechanism also be used to simulate any Turing Machine? We show that this is indeed the case: we can use a sparse encoder and sparse decoder to simulate any Turing Machine.
Big Bird: Transformers for Longer Sequences
To use the sparse attention mechanism in the transformer architecture, we need to define a suitable modification where each token only reacts to previous tokens. Unlike the case for BERT, where the entire attention mechanism is applied once, in full transformers, the sparse attention mechanism at decoder side is used token by token. Secondly the work of Pérez et al. [73], uses each token as a representation of the tape history and uses the full attention to move and retrieve the correct tape symbol. Most of the construction of Pérez et al. [73] goes through for sparse attentions, except for their addressing scheme to point back in history (Lemma B.4 in [73]). We show how to simulate this using a sparse attention mechanism and defer the details to App. B.
Big Bird: Transformers for Longer Sequences
We demonstrate a natural task which can be solved by the full attention mechanism in O(1)layers. However, under standard complexity theoretic assumptions, this problem requires Ω(n)-layers for any sparse attention layers withÕ(n) edges (not just BigBird). (HereÕ hides poly-logarthmic factors.)
Big Bird: Transformers for Longer Sequences
Consider the simple problem of finding the corresponding furthest vector for each vector in the given sequence of length n. Formally, Task 1. Given n unit vectors {u 1 , . . . , u n }, find f (u 1 , . . . , u n ) → (u 1 * , . . . , u n * ) where for a fixed j ∈ [n], we define j * = arg max k u k − u j 2 2 . Finding vectors that are furthest apart boils down to minimize inner product search in case of unit vectors. For a full-attention mechanism with appropriate query and keys, this task is very easy as we can evaluate all pair-wise inner products.
Big Bird: Transformers for Longer Sequences
The impossibility for sparse-attention follows from hardness results stemming from Orthogonal Vector Conjecture(OVC) [1,2,7,97]. The OVC is a widely used assumption in fine-grained complexity. Informally, it states that one cannot determine if the minimum inner product among n boolean vectors is 0 in subquadratic time. In App. C, we show a reduction using OVC to show that if a transformer g ∈ T H=1,m=2d,q=0 D for any sparse directed graph D can evaluate the Task 1, it can solve the orthogonal vector problem. Proposition 1. There exists a single layer full self-attention g ∈ T H=1,m=2d,q=0 that can evaluate Task 1, i.e. g(u 1 , ..., u n ) = [u 1 * , . . . , u n * ], but for any sparse-attention graph D with O(n) edges (i.e. inner product evaluations), would requireΩ(n 1−o(1) ) layers.
Big Bird: Transformers for Longer Sequences
We give a formal proof of this fact in App. C.
Big Bird: Transformers for Longer Sequences
In this section our goal is to showcase benefits of modeling longer input sequence for NLP tasks, for which we select three representative tasks. We begin with basic masked language modeling (MLM; Devlin et al. 23) to check if better contextual representations can be learnt by utilizing longer contiguous sequences. Next, we consider QA with supporting evidence, for which capability to handle longer sequence would allow us to retrieve more evidence using crude systems like TF-IDF/BM25. Next, we tackle long document classification where   discriminating information may not be located in first 512 tokens. Finally, we present BigBird for generating abstractive summaries for documents longer than 512 token, which is a very realistic scenario. Below we summarize the results for BigBird using sequence length 4096, while we defer all other setup details including computational resources, batch size, step size, number of steps, etc. in App. E.
Big Bird: Transformers for Longer Sequences
Pretraining and MLM We follow [23,64] to create base and large versions of BigBird and pretrain it using MLM objective. This task involves predicting a random subset of tokens which have been masked out. We use four standard data-sets for pretraining (listed in Tab. 2), warm-starting from the public RoBERTa checkpoint 1 . We compare performance in predicting the masked out tokens in terms of bits per character, following [8]. As seen in Tab. 3, both BigBird and Longformer perform better than limited length RoBERTa, with BigBird-etc performing the best. We remark that Longformer was trained with a batch size of 16-32 with 48GB memory/chip, where as we trained on a more reasonable 16GB memory/chip with batch size of 32-64. Our increased memory efficiency is due to efficient blocking and sparsity structure of the sparse attention mechanism described in Sec. 2.
Big Bird: Transformers for Longer Sequences
For these tasks we only need to train encoder and we employ sparse attention mechanism of BigBird as described in Sec. 2.
Big Bird: Transformers for Longer Sequences
Question Answering (QA) We considered following four challenging datasets:
Big Bird: Transformers for Longer Sequences
1. Natural Questions [53]: For the given question, find a short span of answer (SA) from the given evidences as well highlight the paragraph from the given evidences containing information about the correct answer (LA). 2. HotpotQA-distractor [101]: Similar to natural questions, it requires finding the answer (Ans) as well as the supporting facts (Sup) over different documents needed for multi-hop reasoning from the given evidences. 3. TriviaQA-wiki [42]: We need to provide an answer for the given question using provided Wikipedia evidence, however, the answer might not be present in the given evidence. On a smaller verified subset of question, the given evidence is guaranteed to contain the answer. Nevertheless, we model the answer as span selection problem in this case as well     As these tasks are very competitive, multiple highly engineered systems have been designed specific each dataset confirming to respective output formats. For a fair comparison, we had to use some additional regularization for training BigBird, details of which are provided in App. E along with exact architecture description. We experiment using the base sized model and select the best configuration on the development set for each dataset (as reported in Tab. 4). We can see that BigBird-etc, with expanded global tokens consistently outperforms all other models. Thus, we chose this configuration to train a large sized model to be used for evaluation on the hidden test set.
Big Bird: Transformers for Longer Sequences
In Tab. 5, we compare BigBird-etc model to top-3 entries from the leaderboard excluding BigBird. One can clearly see the importance of using longer context as both Longformer and BigBird outperform models with smaller contexts. It is worth noting that BigBird submission is a single model, whereas the other top-3 entries for Natural Questions are ensembles, which might explain the slightly lower accuracy in exact answer phrase selection. Also, high performing Fusion-in-Decoder [40] Table 6: Classification results. We report the F1 micro-averaged score for all datasets. Experiments on smaller IMDb and Hyperpartisan datasets are repeated 5 times and the average performance is presented along with standard deviation.
Big Bird: Transformers for Longer Sequences
a sparse long sequence encoder-decoder transformer, despite being motivated differently in their paper. In particular, their sparse attention at encoder side would be block diagonal, thereby being a special case of BigBird with r = 0, g = 0.
Big Bird: Transformers for Longer Sequences
Document Classification We experiment on datasets of different lengths and contents as listed in tab. 6. We compute the fraction of the dataset that exceeds 512 tokens, i.e. the length at which the document are often truncated. Following BERT, we used one layer with cross entropy loss on top of the first [CLS] token. We see that gains of using BigBird are more significant when we have longer documents and fewer training examples. For instance, using base sized model, BigBird improves state-of-the-art for Arxiv dataset by about 5% points. On Patents dataset, there is improvement over using simple BERT/RoBERTa, but given the large size of training data the improvement over SoTA is not significant. Note that this performance gain is not seen for much smaller IMDb dataset. Along with experimental setup detail, we present results on GLUE tasks in app. E which show competitive performance.
Big Bird: Transformers for Longer Sequences
For an encoder-decoder setup, one can easily see that both suffer from quadratic complexity due to the full self attention. We focus on introducing the sparse attention mechanism of BigBird only at the encoder side. This is because, in practical generative applications, the length of output sequence is typically small as compared to the input. For example for text summarization, we see in Tab. 7 that the median output sequence length is ∼ 200 where as the input sequence's median length is > 3000. For such applications, it is more efficient to use sparse attention mechanism for the encoder and full self-attention for the decoder.    Summarization Document summarization is a task of creating a short and accurate summary of a text document. We used three long document datasets for testing our model details of which are mention in Tab. 7. In this paper we focus on abstractive summarization of long documents where using a longer contextual encoder should improve performance. The reasons are two fold: First, the salient content can be evenly distributed in the long document, not just in first 512 tokens, and this is by design in the BigPatents dataset [79]. Second, longer documents exhibit a richer discourse structure and summaries are considerably more abstractive, thereby observing more context helps. As has been pointed out recently [77,107], pretraining helps in generative tasks, we warm start from our general purpose MLM pretraining on base-sized models as well as utilizing state-of-the-art summarization specific pretraining from Pegasus [107] on large-sized models. The results of training BigBird sparse encoder along with full decoder on these long document datasets are presented in Tab. 8. We can clearly see modeling longer context brings significant improvement. Along with hyperparameters, we also present results on shorter but more widespread datasets in App. E.4, which show that using sparse attention does not hamper performance either.
Big Bird: Transformers for Longer Sequences
There has been a recent upsurge in using deep learning for genomics data [13,87,106], which has resulted in improved performance on several biologically-significant tasks such as promoter site prediction [72], methylation analysis [56], predicting functional effects of non-coding variant [109], etc. These approaches consume DNA sequence fragments as inputs,
Big Bird: Transformers for Longer Sequences
...  Thus each document had a continuous strand of 25000-100,000 bps of DNA. This process was repeated 10 times to create 10 sets of document for each chromosome of GRCH37. The resulting set of documents was then passed through Sentencepiece that created tokens of average 8bp. For pretraining we used masked language model and masked 10% of the tokens and trained on predicting the masked tokens. and therefore we believe longer input sequence handling capability of BigBird would be beneficial as many functional effects in DNA are highly non-local [12]. Furthermore, taking inspiration from NLP, we learn powerful contextual representations for DNA fragments utilizing abundant unlabeled data (e.g. human reference genome, Saccharomyces Genome Database) via MLM pretraining. Next, we showcase that our long input BigBird along with the proposed pretraining significantly improves performances in two downstream tasks. Detailed experimental setup for the two tasks are provided in App. F.  Pre-training and MLM Instead of operating directly on base pairs, we propose to first segment DNA into tokens so as to further increase the context length (Fig. 2). This idea is not entirely novel and has been previously explored in Liang [59]. In particular, we build a byte-pair encoding [51] table for the DNA sequence of size 32K, with each token representing 8.78 base pairs on average. We learn contextual representation of these token on the human reference genome (GRCh37) 2 using MLM objective.
Big Bird: Transformers for Longer Sequences
We then report the bits per character (BPC) on a held-out set in Tab. 9. We find that attention based contextual representation of DNA does improve BPC, which is further improved by using longer context in BigBird.  Promoter Region Prediction Promoter is a DNA region typically located upstream of the gene, which is the site of transcription initiation. Multiple methods have been proposed to identify the promoter regions in a given DNA sequence [100,60,11,99,72], as it is an important first step in understanding gene regulation. The corresponding machine learning task is to classify a given DNA fragment as promoter or non-promoter sequence. We use the dataset compiled by Oubounyt et al. [72] which was built from Eukaryotic Promoter Database (EPDnew) [25] 3 . We finetuned the pretrained BigBird model from above, using the training data and report F1 on test dataset. We compare our results to the previously reported best method in Tab. 10. We see that BigBird achieve nearly perfect accuracy with a 5% jump from the previous best reported accuracy.   4 and Roadmap Epigenomics projects 5 . The corresponding ML task is to predict, for a given non-coding region of DNA, these 919 chromatin-profile including 690 transcription factors (TF) binding profiles for 160 different TFs, 125 DNase I sensitivity (DHS) profiles and 104 histone-mark (HM) profiles. We jointly learn 919 binary classifiers to predict these functional effects from sequence of DNA fragments. On held-out chromosomes, we compare AUC with the baselines in Tab. 11 and see that we significantly improve on performance on the harder task HM, which is known to have longer-range correlations [28] than others.
Big Bird: Transformers for Longer Sequences
We propose BigBird: a sparse attention mechanism that is linear in the number of tokens.
Big Bird: Transformers for Longer Sequences
BigBird satisfies a number of theoretical results: it is a universal approximator of sequence to sequence functions and is also Turing complete. Theoretically, we use the power of extra global tokens preserve the expressive powers of the model. We complement these results by showing that moving to sparse attention mechanism do incur a cost. Empirically, BigBird gives state-of-the-art performance on a number of NLP tasks such as question answering and long document summarization. We further introduce attention based contextual language model for DNA and fine-tune it for down stream tasks such as promoter region prediction and predicting effects of non-coding variants.
Big Bird: Transformers for Longer Sequences
[ We begin by setting up some notations following Pérez et al. [73] to formally describe the complete architecture of Transformers. A single layer of Transformer encoder is a parametric function Enc receiving a sequence X = (x 1 , ..., x n ) of vectors in R d and returning a sequence Z = (z 1 , ..., z n ) of the same length. Each z i is a d dimensional vector as well.
Big Bird: Transformers for Longer Sequences
We interchangeably treat the sequence X as a matrix in R n×d . Enc has two components:
Big Bird: Transformers for Longer Sequences
1. An attention mechanism Attn that takes in the sequence X and returns sequence (a 1 , ..., a n ) of the same length and dimensionality; and 2. A two layer fully connected network O that takes in a vector in R d and returns a vector in R d .
Big Bird: Transformers for Longer Sequences
Then i-th output vector of Enc(X) is computed as follows:
Big Bird: Transformers for Longer Sequences
Now it remains to define Attn and O which we do next.
Big Bird: Transformers for Longer Sequences
As described in Sec. 2, an attention mechanism is parameterized by three functions: Q, K, V : R d → R m . In this paper, we assume that they are simply matrix products: 
Big Bird: Transformers for Longer Sequences
where N (i) denote the out-neighbors set of node i in D. In other words, the set of arcs (directed edges) in D represents the set of inner products that our attention mechanism will consider. Also recall that σ is a scoring function such as softmax or hardmax.
Big Bird: Transformers for Longer Sequences
Lastly, we define the output fully connected network as follows:
Big Bird: Transformers for Longer Sequences
Here W 1 ∈ R d×q , W 2 ∈ R q×d , b 1 ∈ R p , and b 2 ∈ R d are parameters of output network O.
Big Bird: Transformers for Longer Sequences
We introduce a few pieces of additional notation that will be useful.
Big Bird: Transformers for Longer Sequences
Let 
Big Bird: Transformers for Longer Sequences
In this section, we will present the full proof of theorem 1. The proof will contain three parts. The first and the third part will largely follow standard techniques. The main innovation lies is in the second part.
Big Bird: Transformers for Longer Sequences
A.2.1 Approximate F CD by piece-wise constant functions First, we consider a suitable partition of the region (0, 1) into a grid of granularity δ, which we denote by G δ . We do this using Lemma 8 from Yun et al. [105], which we restate for completeness:
Big Bird: Transformers for Longer Sequences
Lemma 1 (Lemma 8 [105]). For any given f ∈ F CD and 1 ≤ p ≤ ∞, there exists a δ > 0 such that there exists a piece-wise constant functionf with
Big Bird: Transformers for Longer Sequences
Since transformers can learn a positional embedding E, without any loss of generality, we can consider the translated function. In particular, define
Big Bird: Transformers for Longer Sequences
To do so, we will apply a suitable modification of Lemma 1, which will consider the discretized grid
Big Bird: Transformers for Longer Sequences
Therefore, it suffices to approximate a functionf : G E δ → R n×d defined as
Big Bird: Transformers for Longer Sequences
Throughout this section, we will assume that we are given a function that has an extra global token at index 0 and all vectors have an extra dimension appended to them. The latter assumption is without loss of generality as we can use the Feed-Forward Network to append sparse dimensions. In particular, we will associate X ∈ R (n+1)×(d+1) where we write X = (x 0 , x 1 , . . . , x n ). Although our function is only defined for G E δ ⊂ R n×d , we can amend the function in a natural way by making it ignore the first column. To avoid excessive clutter, we will assume that the function value is evaluated on the last n columns.
Big Bird: Transformers for Longer Sequences
The main idea in this section is the use of contextual mapping to enable Transformers to compute any discretized function. A contextual mapping is an unique encoding of each tuple (X, x i ) where X ∈ G E δ , and each column
Big Bird: Transformers for Longer Sequences
. We restate the definition adapted to our setting below Definition 2 (Defn 3.1 [105]). (Contextual Mapping) A contextual mapping is a function mapping q : G E δ → R n if it satisfies the following:
Big Bird: Transformers for Longer Sequences
1. For any P ∈ G E δ , q(P ) contains distinct entries.
Big Bird: Transformers for Longer Sequences
2. For any two P, P ∈ G E δ with P = P , all entries of q(P ) and q(P ) are distinct.
Big Bird: Transformers for Longer Sequences
The key technical novelty of the proof is computing a contextual mapping using only the sparse attention mechanism. We create a "selective shift" operator which only shifts entries of a vector that lie in a certain range. We will use this shift operator strategically to ensure that we attain a contextual mapping at the end of the process. The lemma below, which is based on parts of the proof of Lemma 6 of [105], states that we can implement a suitable "selective" shift operator using a sparse attention mechanism.
Big Bird: Transformers for Longer Sequences
Lemma 2. Given a function ψ : R (n+1)×(d+1) × R 2 → R (n+1)×1 and a vector u ∈ R d+1 and a sparse attention mechanism based on the directed graph D, we can implement a selective shift operator that receives as input a matrix X ∈ R (n+1)×(d+1) and outputs
Big Bird: Transformers for Longer Sequences
Note that e 1 ∈ R d+1 denotes (1, 0, . . . , 0).
Big Bird: Transformers for Longer Sequences
Proof. Consider the function , which can be implemented by a sparse attention mechanism :
Big Bird: Transformers for Longer Sequences
This is because the Key, Query and Value functions are simply affine transformations of X.
Big Bird: Transformers for Longer Sequences
Given any graph D, the above function will evaluate to the following:
Big Bird: Transformers for Longer Sequences
Therefore we can say thatψ(Z; b Q ) −ψ(Z; b Q ) satisfies
Big Bird: Transformers for Longer Sequences
The following lemma, which is the heart of the proof, uses the above selective shift operators to construct contextual mappings.
Big Bird: Transformers for Longer Sequences
Lemma 3. There exists a function g c : R (n+1)×(d+1) → R (n+1) and a unique vector u, such that for all P ∈ G E δ g c (P ) := u, g(P ) satisfies the property that g c is a contextual mapping of P . Furthermore, g c ∈ T 2,1,1 D using a composition of sparse attention layers as long as D contains the star graph.
Big Bird: Transformers for Longer Sequences
Proof. Define u ∈ R d+1 = [1, δ −1 , δ −2 , . . . , δ −d+1 , δ −nd ] and let X 0 = (0, . . . , 0, 1). We will assume that x i , x 0 = 0, by assuming that all the columns x 1 , . . . , x n are appended by 0.
Big Bird: Transformers for Longer Sequences
To successfully encode the entire context in each token, we will interleave the shift operator to target the original columns 1, . . . , n and to target the global column 0. After a column i is targeted, its inner product with u will encode the entire context of the first i columns. Next, we will shift the global token to take this context into account. This can be subsequently used by the remaining columns.
Big Bird: Transformers for Longer Sequences
For i ∈ {0, 1, . . . , n}, we will use l i to denote the innerproducts u, x i at the beginning. For f i = u, x i after the i th column has changed for i ∈ {1, . . . , n} and we will use f k 0 to denote u, x 0 after the k th phase. We need to distinguish the global token further as it's inner product will change in each phase. Initially, given X ∈ G E δ , the following are true:
Big Bird: Transformers for Longer Sequences
Note that all l i ordered in distinct buckets l 1 < l 2 < · · · < l n < l 0 .
Big Bird: Transformers for Longer Sequences
We do this in phases indexed from i ∈ {1, . . . , n}. Each phase consists of two distinct parts:
Big Bird: Transformers for Longer Sequences
The low shift operation: These operation will be of the form
Big Bird: Transformers for Longer Sequences
The range is chosen so that only l i will be in the range and no other l j j = i is in the range. This will shift exactly the i th column x i so that the new inner product f i = u, x i is substantially larger than l i . Furthermore, no other column of X will be affected. The high shift operation: These operation will be of the form
Big Bird: Transformers for Longer Sequences
The range [S i , T i ) δ is chosen to only affect the column x 0 (corresponding to the global token) and no other column. In particular, this will shift the global token by a further δ −nd . Letf i 0 denote the value off i 0 = u, x 0 at the end of i th high operation. Each phase interleaves a shift operation to column i and updates the global token. After each phase, the updated i th column f i = u, x i will contain a unique token encoding the values of all the l 1 , . . . , l i . After the high update,f i 0 = u, x 0 will contain information about the first i tokens.
Big Bird: Transformers for Longer Sequences
Finally, we define the following constants for all k ∈ {0, 1, . . . , n}.
Big Bird: Transformers for Longer Sequences
After each k phases, we will maintain the following invariants:
Big Bird: Transformers for Longer Sequences
The order of the inner products after k th phase is
Big Bird: Transformers for Longer Sequences
Base case The case k = 0, is trivial as we simply set S 0 = δ −(n+1)d , T 0 = δ −(n+1)·d + δ.
Big Bird: Transformers for Longer Sequences
The first nontrivial case is k = 1.
Big Bird: Transformers for Longer Sequences
Step First, in the low shift operation is performed in the range [δ −(k−1)d , δ −kd ) δ Due to the invariant, we know that there exists only one column x k that is affected by this shift. In particular, for column k, we will have max j∈N (k) u, x j = u, x 0 =f k−1 0 . The minimum is l k . Thus the update will be f k = δ −d (f k−1 0 − l k ) + l k . Observe that for small enough δ, f k ≥f k−1 0 . Hence the total ordering, after this operation is
Big Bird: Transformers for Longer Sequences
Now when we operate a higher selective shift operator in the range [S k−1 , T k−1 ) δ . Since only global token's innerproductf k−1 0 is in this range, it will be the only column affected by the shift operator. The global token operates over the entire range, we know from Eq. (2) that, f k = max i∈[n] u, x i and l k+1 = min i∈[n] u, x i . The new valuef k 0 = δ −nd ·(f k −l k+1 )+f k−1 0 . Expanding and simplifying we get,
Big Bird: Transformers for Longer Sequences
Expanding the above recursively, we get
Big Bird: Transformers for Longer Sequences
Since we know thatf 0 0 = δ −nd and each l i < δ −id , we can substitute this to get Eq. (UP) and we can get an lower-bound Eq. (LP) by using l i ≥ δ −(i−1)d .
Big Bird: Transformers for Longer Sequences
By construction, we know that S k ≤f k 0 < T k . For sufficiently small δ, observe that S k ≤f k 0 < T k all are essentially the dominant term ≈ O(δ −n(k+1)d−kd ) and all the lower order terms do not matter. As a result it is immediate to see that that f k > δ −d (f k−1 0 − l k ) > T k−1 and hence we can see that the invariant 2 is also satisfied. Since only column k and the global token are affected, we can see that invariant 3 is also satisfied.
Big Bird: Transformers for Longer Sequences
After n iterations,f n 0 contains a unique encoding for any P ∈ G E δ . To ensure that all tokens are distinct, we will add an additional layer X = X + δ −n 2 d ψ(X, v − δ/2, v + δ/2) for all v ∈ [S 1 , T n ) δ . This ensures that for all P, P ∈ G E δ , each entry of q(P ) and q(P ) are distinct.
Big Bird: Transformers for Longer Sequences
The previous lemma shows that we can compute a contextual mapping using only sparse transforms. We now use the following lemma to show that we can use a contextual mapping and feed-forward layers to accurately map to the desired output of the functionf .
Big Bird: Transformers for Longer Sequences
Lemma 4 (Lemma 7 [105]). Let g c be the function in Lemma 3, we can construct a function g v : R (n+1)×(d+1) → R (n+1)×d composed of O(nδ −nd ) feed-forward layers (with hidden dimension q = 1) with activations in Φ such that g v is defined as g v (Z) = [g tkn v (Z 1 ), . . . , g tkn v (Z n )], where for all j ∈ {1, . . . , n}, g tkn v (g c (L) j ) = f (L) j
Big Bird: Transformers for Longer Sequences
The previous section assumed we used Transformers that used hardmax operator σ H and activations functions belonging to the set Φ. This is without loss of generality as following lemma shows.
Big Bird: Transformers for Longer Sequences
Lemma 5 (Lemma 9 [105]). For each g ∈T 2,1,1 and 1 ≤ p ≤ ∞, ∃g ∈ T 2,1,4 such that d p (g,ḡ) ≤ /3
Big Bird: Transformers for Longer Sequences
Combining the above lemma with the Lemma 3, we get our main result:
Big Bird: Transformers for Longer Sequences
Theorem 2. Let 1 ≤ p ≤ ∞ and > 0, there exists a transformer network g ∈ T 2,1,4 D which achieves a ratio of d p (f, g) ≤ where D is the sparse graph.
Big Bird: Transformers for Longer Sequences
Since the sparsity graph associated with BigBird contains a star network, we know that it can express any continuous function from a compact domain.
Big Bird: Transformers for Longer Sequences
In this section, we will extend our results to the setting of Pérez et al. [73]. Our exposition will largely use their proof structure but we will make a few changes. We repeat some of the lemmas with the amendments to make the exposition self-contained.
Big Bird: Transformers for Longer Sequences
Transformer Decoder We need both an encoder and a decoder in the transformer for simulating a Turing machine. We utilize the same notation used in App. A.1 for encoders. The decoder is similar to an encoder but with additional attention to an external pair of key-value vectors (K e ∈ R n×m , V e ∈ R n×d ), which usually come from the encoder stack. A single layer of Transformer decoder is a parametric function Dec receiving a sequence Y j = (y 1 , . . . , y j ) of vectors in R d plus the external (K e , V e ) and returning a sequence of vectors Z j = (z 1 , . . . , z j ) of the same length. Each z i is a d dimensional vector as well. Dec has three components, one more than Enc:
Big Bird: Transformers for Longer Sequences
1. An attention mechanism Attn that takes in the sequence Y j and returns sequence (p 1 , ..., p j ) of the same length and dimensionality;
Big Bird: Transformers for Longer Sequences
CrossAttn that takes in the sequence (p 1 , ..., p j ) plus the external (K e , V e ) and returns sequence (a 1 , ..., a j ), with each a i ∈ R d ; and
Big Bird: Transformers for Longer Sequences
3. A two layer fully connected network O that takes in a vector in R d and returns a vector in R d .
Big Bird: Transformers for Longer Sequences
Then i-th output vector of Dec(Y j ; K e , V e ) is computed as follows:
Big Bird: Transformers for Longer Sequences
where
Big Bird: Transformers for Longer Sequences
and
Big Bird: Transformers for Longer Sequences
Attn D and O are as defined in App. A.1 and it remains to define CrossAttn. The i th output vector of multi-head cross-attention attention is given by
Big Bird: Transformers for Longer Sequences
Turning Machine We will use the same setup of Turning Machine that was used by Pérez et al. [73] (see section B.4). Given a Turing Machine M = (Q, Σ, δ, q init , F ), we use the following notation Vector representations For a symbol s ∈ Σ, s denotes its one-hot vector representation in Q |Σ| . All the transformer intermediate vectors used in our simulations have dimension d = 2|Q| + 4|Σ| + 16. Note that we use five extra dimension as compared to Pérez et al. [73]. We follow the convention used in Pérez et al. [73] and write a a vector v ∈ Q d arranged in four groups of values as follows
Big Bird: Transformers for Longer Sequences
where q i ∈ Q |Q| , s i ∈ Q |Σ| , and x i ∈ Q.
Big Bird: Transformers for Longer Sequences
In this section, we give more details on the architecture of the encoder and decoder needed to implement our simulation strategy.
Big Bird: Transformers for Longer Sequences
High Level Overview: Given the Turing machine M , we will show that a transformer with an appropriate encoder and decoder T D can simulate each step of M 's execution.
Big Bird: Transformers for Longer Sequences
Our simulation strategy will mostly follow Pérez et al. [73], except we will use a sparse attention mechanism. The main idea is to maintain the current Turing machine state q (j) and symbol under the head s (j) as part of the decoder sequence Y for all time step j so that we can always simulate the corresponding Turing machine transition δ(q (j) , s (j) ) = (q (j) , v (j) , m (j) ).
Big Bird: Transformers for Longer Sequences
The key difference will rise in Lemma B.4 of Pérez et al. [73], where full attention is used to select the appropriate symbol from tape history in one step. To accomplish the same task with sparse attention, we will exploit the associative property of max and break down the symbol selection over multiple steps. Thus, unlike Pérez et al. [73] one decoding step of our sparse transformer T D does not correspond to one step of the Turing machine M . In particular, we will have two type of steps: compute step corresponding to update of M 's state and intermediate steps corresponding to aggregating the max (which in turn is used for symbol selection). Let i denote the step of T D and g(i) denote the step of M being simulated at step i of the decoder. At each decoding step we want to maintain the current Turing   App. B.2.3). Finally, at the next compute step we will make the transition to new state q g(i)+1 , new head movement m g(i) , and new output symbol v g(i) to be written. Thereby we are able to completely simulate the given Turing machine M . As a result, we can prove the following main theorem: Theorem 3. There exists a sparse attention mechanism using O(n) inner products such that the resulting class of Transformer Networks using this sparse attention mechanism is Turing Complete.
Big Bird: Transformers for Longer Sequences
As [73], we use the same trivial single layer encoder where resulting K (e) contains position embedding and V (e) contains one-hot symbol representation.
Big Bird: Transformers for Longer Sequences
Sparse Self-Attention mechanism for Decoder In this section, we will consider a particular instance of the sparse graph D at decoder. We define its edges to be given by the following relations: ∀j ∈ N + , 1 ≤ k ≤ j + 1, j(j + 1) 2 + k, k(k + 1) 2 and
Big Bird: Transformers for Longer Sequences
This graph can be seen as a special case of BigBird where first type of edges are realizations of random and second type of edges correspond to locality. Also note that this graph satisfies the left-to-right constraint of decoder, i.e. no node attends to a node in the future. and h(i) = g(i + 1) − g(i). Note that h(i) reduces to a binary
Big Bird: Transformers for Longer Sequences
We next show how to construct the decoder layers to produce the sequence of outputs y 1 , y 2 , . . ., where y i is given by:
Big Bird: Transformers for Longer Sequences
That is, at step i of our sparse decoder y i , it will contain the information about the state of the turing machine M at time g(i), the symbol under the head of M at time g(i), and the current location of head of M at time g(i). We also have a placeholder symbol w and placeholder scalars u 1 , u 2 , u 3 , whose role will be clear from our construction.
Big Bird: Transformers for Longer Sequences
We consider as the starting vector for the decoder the vector
Big Bird: Transformers for Longer Sequences
We assume that the start head is at c (0) = 0, the initial state is q (0) = q init , and s (0) = # as we initialize from clean tape. We show the correctness of our construction by an inductive argument: we describe the architecture piece by piece and at the same time will show for every r ≥ 0 , our architecture constructs y r+1 from the previous vectors (y 0 , . . . , y r ).
Big Bird: Transformers for Longer Sequences
Thus, assume that y 1 , . . . , y r satisfy the properties stated above. Since we are using positional encodings, the actual input for the first layer of the decoder is the sequence y 1 + pos Dec (1), y 2 + pos Dec (2), . . . , y r + pos Dec (r).
Big Bird: Transformers for Longer Sequences
We denote by y i the vector y i plus its positional encoding. Thus we have ∀ 1 ≤ i ≤ r that
Big Bird: Transformers for Longer Sequences
In this layer, we use the cross-attention between encoder and decoder to access the input string and a feed-forward network to simulate the transition function of M . The first self attention in Eq. (5) is not used in this layer and we just produce the identity. This identity function is achieved by setting all queries, keys, values to be 0 everywhere plus the residual connection. Thus, we have p 1 i = y i .
Big Bird: Transformers for Longer Sequences
Since p 1 i is of the form [ , . . . , , 1, g(i) + 1, , . . . , ], we know by Lemma B.1 of Pérez et al. [73] that if we use p 1 i to attend over the encoder we obtain
Big Bird: Transformers for Longer Sequences
where α and β are as defined in Eq. (21) of [73]. Thus in Eq. (4) we finally produce the vector a 1 i given by
Big Bird: Transformers for Longer Sequences
As the final piece of the first decoder layer we use a function O 1 (·) (Eq. (3)) that satisfies the following lemma.
Big Bird: Transformers for Longer Sequences
Lemma 6 (Lemma B.2 [73]). There exists a two-layer feed-forward network O 1 : Q d → Q d such that with input vector a 1 i (Eq. (7)) produces as output That is, function O 1 (·) simulates transition δ(q g(i) , s g(i) ) to construct q g(i)+1 , v g(i) , and m g(i) besides some other linear transformations.
Big Bird: Transformers for Longer Sequences
Thus, finally the output of the first decoder layer is
Big Bird: Transformers for Longer Sequences
In this layer, we only use the feed-forward network to evaluate the next location of the head. The self-attention and cross-attention are set to be the identity function, so a 2 i = p 2 i = z 1 i . Recall that c g(i) is the cell to which M is pointing to at time g(i), and that it satisfies the following recursion c g(i)+1 = c g(i) + m g(i) , which can be expanded to see that that c g(i)+1 = m (0) + m (1) + · · · + m g(i) . Its not difficult to see that a two layer network with nonlinearity can compute c g(i)+1 /(g(i) + 1) and c g(i) /(g(i) + 1) from c g(i) , m g(i) , and 1/(g(i) + 1) using the relation c g(i)+1 = c g(i) + m g(i) . At the end of layer 2, we obtain
Big Bird: Transformers for Longer Sequences
This is an additional layer (not present in the work of [73]), where we propagate computations in our sparse graph. In particular, we will use this layer to "compute" or accumulate state in intermediate nodes. We make this clear below. The self-attention and cross-attention are all set to be the identity function, so a 3 i = p 3 i = z 2 i . In this layer, we only use the dense attention layers to select the newly computed states or to continue with previous states. Using idea similar to Lemma B.6 of [73], we can construct a dense network such that
Big Bird: Transformers for Longer Sequences
The negatives are generated to offset results from skip connection. We utilize such network to switch Turing machine state and position embedding for intermediate steps to the values received from previous time step and do nothing for compute nodes. We use h(i) as the flipping bit b. Thus, at end of layer 3, we obtain
Big Bird: Transformers for Longer Sequences
where we used h(i) for selecting old states. In particular,
Big Bird: Transformers for Longer Sequences
• We copy the input state and head position as is for intermediate nodes. We do not need to transition to next Turing machine states in these nodes.
Big Bird: Transformers for Longer Sequences
• To preserve the symbol under the head for intermediate nodes, we copy the previous symbol to α location and set β = g(i) + 1, as the symbol at α location will be copied as the symbol under head for next transformer step by the final transformation layer if β = g(i) + 1. Thus, we correctly preserve the previous symbol under head as Turing machine does not transition these nodes. For compute nodes, things happen as usual.
Big Bird: Transformers for Longer Sequences
• Finally for the intermediate nodes, we copy the position embedding corresponding to current best symbol w, which is stored in u 1 , u 2 , u 3 , u 4 . For compute node, we let the position embedding correspond to current Turing machine step.
Big Bird: Transformers for Longer Sequences
For further simplification note that g(i + 1) = g(i) if h(i) = 0 else g(i) + 1 when h(i) = 1.
Big Bird: Transformers for Longer Sequences
With this fact, we can conclude thatq (i) = q g(i+1) andĉ (i) = c g(i+1) . Thus, we can write,
Big Bird: Transformers for Longer Sequences
To find the symbol on tape under next head position c g(i)+1 , we try to find what was written last at the location c g(i)+1 . To facilitate this, following [73], we define (j) to be the last time (previous to j) in which M was pointing to position c (j) , or it is j − 1 if this is the first time that M is pointing to c (j) . Recall j is the Turing machine step counter, which is different from sparse transformer step i. [73] could utilize full attention mechanism to find v (j+1) at one go, but we have to do it over multiple steps owing to our sparse attention mechanism.
Big Bird: Transformers for Longer Sequences
We use similar query, key, value functions as used for full attention by [73] ∀i:
Big Bird: Transformers for Longer Sequences
It is clear that the three functions are linear transformations and thus they can be defined by feed-forward networks. Notice that the query vector is always formed using current time step position embedding, whereas key and value vectors are formed using copied over entries for intermediate nodes and using current entries only for compute node.
Big Bird: Transformers for Longer Sequences
Pérez et al. [73] find the desired v l(j+1) as v m(j) using full attention, where m(t) = arg min m∈{0,...,t}
Big Bird: Transformers for Longer Sequences
Note the minimization is only over Turing machine steps, i.e. over compute nodes in our case. We show below that we can estimates m(j) by parts using sparse attention mechanism. The main idea is just to notice that minimization problem min m∈{0,...,t} χ j t can be expressed as min{· · · min{min{χ j 0 , χ j 1 }, χ j 2 }, ..., χ j t } by the associativity property. By definition of our graph D, at every intermediate node i of the form j(j + 1)/2 + k, i.e. where k > 0, g(i) = j and h(i) = 0, we will attend over node k(k + 1)/2 and best till now copied from i − 1. The node k(k + 1)/2 is never an intermediate node as h(k(k + 1)/2) = 1 for all k and in fact corresponds to Turing machine's step k. This will help us select the key and value corresponding to min between node k(k + 1)/2 and i − 1. In other words, at node i of the form j(j + 1)/2 + k we would have evaluated m(k) and corresponding value selected:
Big Bird: Transformers for Longer Sequences
and similarly for u's. So after going through all the intermediate nodes, finally at the next compute node, i.e. when k = j + 1, we will obtain the minimum value over all of 0, 1, ..., j. This implies at a compute node will be able to recover (g(i) + 1) and its corresponding value as shown in Lemma B.4 of [73]. Then we have that p 4 i is given by
Big Bird: Transformers for Longer Sequences
The cross-attention and feed-forward network are set to be identity, so z 4 i = a 4 i = p 4 i .
Big Bird: Transformers for Longer Sequences
We finish our construction by using the final transformation function F (·) from the corresponding lemma from Pérez et al. [73], with a slight modification.
Big Bird: Transformers for Longer Sequences
Lemma 7 (Lemma B.5 [73]). There exists a function F : Q d → Q d defined by a feed-forward network such that
Big Bird: Transformers for Longer Sequences
The modification is to let w, u 1 , u 2 , u 3 to pass through. This yields the desired input to transformer at next time step for both intermediate and compute node, thereby concluding our induction.
Big Bird: Transformers for Longer Sequences
Step 2 Construct query, key, value functions as follows:
Big Bird: Transformers for Longer Sequences
Then
Big Bird: Transformers for Longer Sequences
Then,
Big Bird: Transformers for Longer Sequences
Step 3 Let O(a i ) = 0, then the output z i = [u i ; u i * ] as desired.
Big Bird: Transformers for Longer Sequences
To complete the argument, observe that it now only takes O(n) inner products to check if there is a pair of orthogonal vectors as we need only compare u i , u i * .
Big Bird: Transformers for Longer Sequences
Part 2: Every Sparse Attention Mechanism will needΩ(n 1− ) layers We prove by contradiction that it is impossible to solve Task 1 by any g ∈ T
Big Bird: Transformers for Longer Sequences
sparse-attention graph D withÕ(n) edges.
Big Bird: Transformers for Longer Sequences
Suppose we can solve Task 1 using a network g ∈ T
Big Bird: Transformers for Longer Sequences
that has l layers. Recall that all the computation we do in one layer is:
Big Bird: Transformers for Longer Sequences
where Attn D is defined in eq. (AT).
Big Bird: Transformers for Longer Sequences
Thus, total computation per layer isÕ(nd 3 ) and consequentlyÕ(nld 3 ) for the whole network consisting of l layers.
Big Bird: Transformers for Longer Sequences
We can use the result of Task 1 to solve the orthogonal vector (OV) problem (defined in Conjecture 1) in linear time. So in total, we will be able to solve any instance of OV iñ O(nld 3 ) time.
Big Bird: Transformers for Longer Sequences
Now if l = O(n 1− ) for any > 0 and d = Θ(log 2 n), then it appears that we are able to solve OV inÕ(n 2− ) which contradicts Conjecture 1. Therefore, we need at leastΩ(n 1−o(1) ) layers.   Figure 5: Structured block sparsity in BigBird helps in compactly packing our operations of sparse attention, thereby making our method efficient on GPU/TPU. The matrix on the left is same as that shown in Fig. 4d. The figure on the right shows how the attention is actually calculated for all blocks of query except the first. The green column is fixed, the second column can be obtained by rotating the key matrix to the left, the third column is the key matrix, the fourth column can be obtained by rotating the key matrix right, and then only for the last column we need to use the computationally inefficient gather operation. Even then we do block gather rather than gathering keys separately for each query, further speeding up operations in the GPU/TPU. well known [34,103], that such sparse multiplications cannot be efficiently implemented in GPUs. GPUs have thousands of cores performing operations in parallel. Thus, we would not be able to efficiently perform the sparse matrix multiplication mentioned in section Sec. 2.
Big Bird: Transformers for Longer Sequences
As a result we propose to first blockify the Attention matrix i.e. we pack sets of query and keys together and then define attention on these blocks. It is easier to explain this process using the example shown in Fig. 4. Suppose, there are 12 query and 12 key vectors to attend to. Using a block size of 2, we split the query matrix into 12/2 = 6 blocks and similarly the key matrix into 12/2 = 6 blocks. Then the three different building components of BigBird are defined on the block matrix. In particular the three different components are:
Big Bird: Transformers for Longer Sequences
1. Random Attention: Each query block attends to r random Key blocks. In Fig. 4a, r = 1 with block size 2. This implies that each query block of size 2 randomly attends to a key block of size 2.
Big Bird: Transformers for Longer Sequences
2. Window Attention: While creating the block, we ensure that the number of query blocks and the number of key blocks are the same. This helps us in defining the block window attention. Every query block with index j attends to key block with index j − (w − 1)/2 to j + (w − 1)/2, including key block j. In Fig. 4b, w = 3 with block size 2. It means that each query block j (size 2 queries) attends to key block j − 1, j, j + 1.
Big Bird: Transformers for Longer Sequences
Global attention remains the same as defined in Sec. 2, but we compute it in terms of blocks. In Fig. 4c, g = 1 with block size 2. For BigBird-itc this implies that one query and key block, attend to everyone.
Big Bird: Transformers for Longer Sequences
The resulting global attention matrix is shown in Fig. 4d.
Big Bird: Transformers for Longer Sequences
We use another trick to speed-up computation. For multiplying pairs of vectors with indexes in a matrix in a GPU, we need to use gather operation. Unfortunately this operation is not efficient and one would want to reduce its use as much as possible. Upon closer examination of window and global attention, we observe that we can perform them without using a gather operation. The global attention for all blocks other than Q1 in Fig. 5, can be computed if we "roll" the key-block matrix forward and backward and multiply with the query block matrix as shown in Fig. 5. While rolling, we do not include global key blocks (K1 in Fig. 5). We only use gather for the random attention, which is very small (r = 3 for all of our experiments). This improves our efficiency. As seen in Fig. 5, rolling of window attention, leads to Q2 query block attending to K6 key block and Q6 attending to K2, which were not present in the original attention. 
Big Bird: Transformers for Longer Sequences
The detailed statistics of the four datasets used are given in Tab. 13. All the hyperparameters for BigBird, used for creating Tab. 4 are shown in Tab. 14 and those submitted to get Tab. 5 are shown in Tab. 15. We use two types of regularization in training:
Big Bird: Transformers for Longer Sequences
• We used a variant of contrastive predictive coding [71] as a dual encoder model.
Big Bird: Transformers for Longer Sequences
• We use position embedding for itc and relative position encoding [80] for etc. Next, we will mention the dataset/task specific part of the model. HotpotQA The data consists of each question with multiple evidence paragraphs. We filtered 16 QA where the answer was not in the given evidences. For BigBird-itc, we use first 128 global tokens. For BigBird-etc, we have one global token for each question token, one for each evidence paragraph, and one for each sentence within the paragraph, for a maximum of 256 global token. We use a dense layer on the output corresponding to global token of the evidence paragraph to predict whether its a supporting fact with a threshold over the output logits. The answer type (yes/no/span) is predicted with a single dense layer from the global CLS token. For span based answers, the spans are predicted with dense layers on the sequence with the distance between start and end positions to be no more than 30 words. The spans are ranked by sum of start and end logits.
Big Bird: Transformers for Longer Sequences
Here also the data consists of question with supporting evidence, but in form of a single, potentially long, document and not multiple paragraphs. We largely follow the setup of [5]. For documents, that are longer than 4096, a sliding window approach is used with stride of 2048. We use CLS token at the beginning, followed by the question followed by a separator token followed by the document as input. For BigBird-itc, we make the first 128 tokens as global. For BigBird-etc, we make a global token for CLS, question, and one token for each of the paragraphs. We train four predictors at the final layer to predict long answer start, long answer end, short answer start and short answer end to a maximum of 430. Further, we linked answer tokens to their mentions using relative position label. Lastly, we use a dense layer that takes in the output vector corresponding to a candidate answer, and predicts a score for the current candidate to be the correct answer. We apply this dense layer to each candidate independently and the candidate with the best score is picked as our final answer.
Big Bird: Transformers for Longer Sequences
It is worthwhile to note that explicitly designed attention connection in etc works slightly better, the random connection based itc is pretty competative.
Big Bird: Transformers for Longer Sequences
Experimental Setup Tab. 16 gives the detail list of hyperparameters used to get results shown in Tab. 6. In all experiments we use a learning rate warmup over the first 10% steps, and linear decay of the learning rate.  
Big Bird: Transformers for Longer Sequences
we used the BigBird-itc base model pretrained on MLM task). We compare the performance of BigBird to BERT, XLNet [102] and RoBERTa in Tab. 17. We find that even on task that have a much smaller context, our performance is competitive to full attention models.
Big Bird: Transformers for Longer Sequences
As discussed in Sec. 4   Following success of several recent works [77,64], we warm start our encoder-decoder BigBird transformer model with pretrained weights and the weights between encoder and decoder are shared. In particular, the query/key/value matrix of self-attention and all the feedforward layers are shared between encoder and decoder. The only variable that is initialized randomly is the encoder-decoder attention. For base sized model, we utilize our MLM pretrained model on 4096 sequence length from App. E.1, which is in turn initialized using the public RoBERTa checkpoint. For the large size model, we lift weight from the state-of-the-art Pegasus model [107], which is pretrained using an objective designed for summarization task.
Big Bird: Transformers for Longer Sequences
To check if sparse attention causes significant degradation as compared to full attention, we further experiment on two shorter but popular datasets, where full attention can be used without significantly truncating the document. The statistics of these two datasets are in Tab. 19. We see that our performance is competitive, which shows that sparse attention can achieve similar performance to a full attention models.
Big Bird: Transformers for Longer Sequences
In this section we provide details of the experimental setup for BigBird on genomics data.
Big Bird: Transformers for Longer Sequences
We try to keep the experimental setup as close to a typical NLP pipeline. In this regard, we take human reference GRCh37 6 and convert it into documents D. Each document d ∈ D is a sequence of sentences, where each sentence is a sequence of fragments of DNA. We construct the documents as follows:
Big Bird: Transformers for Longer Sequences
1. Start with empty document set D = ∅.
Big Bird: Transformers for Longer Sequences
For each chromosome C, repeat the following procedure 10 times.
Big Bird: Transformers for Longer Sequences
(a) Pick uniformly at random a starting point q between base pairs 0 and 5000 from the 5' end.
Big Bird: Transformers for Longer Sequences
(b) Repeat until q > |C| i. Pick uniformly at random s a number between 50 and 100 to denote number of sentences per document. ii. Constructs a document d containing s sentences using consecutive base pairs (bps). The length of each sentence is chosen uniformly at random between 500-1000. Thus the resulting document has 25, 000 -100, 000 bps. iii. D = D d iv. q = q + |d| By this procedure we end-up with approximately 450K documents.
Big Bird: Transformers for Longer Sequences
Next we run sentencepiece [51] tokenization on the resulting documents. In particular, using 5 characters as the building blocks (four for bases -A, T, C, G and one for missing symbol N), we construct a byte pair encoding table of size 32k, with each token representing 8.78 base pairs on average.
Big Bird: Transformers for Longer Sequences
Using the above constructed documents, we construct a dataset for two pretraining tasks following Devlin et al. [23]:
Big Bird: Transformers for Longer Sequences
• Masked Language Model (MLM): In order to train a deep bidirectional representation, BERT training introduces the MLM task, where we simply mask out 15% of the input tokens at random, and then predict those masked tokens. We can simply replace such masked out of the tokens with a [MASK] placeholder, but it leads to a distribution mis-match for downstream tasks which will not have such placeholders.
Big Bird: Transformers for Longer Sequences
To mitigate with this issue, out of the 15% of the tokens selected for masking:
Big Bird: Transformers for Longer Sequences
-80% of the tokens are actually replaced with the token [MASK]. 6. https://www.ncbi.nlm.nih.gov/assembly/GCF_000001405.39
Big Bird: Transformers for Longer Sequences
...
Big Bird: Transformers for Longer Sequences
Masking ... ... Figure 6: Visual description of how the masked language modeling data was generated from raw DNA dataset. The raw DNA sequences of GRCh37, where split at random positions to create documents with 50-100 sentences where each sentence was 500-1000 base pairs (bps). Thus each document had a continuous strand of 25000-100,000 bps of DNA. This process was repeated 10 times to create 10 sets of document for each chromosome of GRCH37. The resulting set of documents was then passed through Sentencepiece that created tokens of average 8bp. For pretraining we used masked language model and masked 15% of the tokens and trained on predicting the masked tokens.
Big Bird: Transformers for Longer Sequences
-10% of the time tokens are replaced with a random token.
Big Bird: Transformers for Longer Sequences
-10% of the time tokens are left unchanged, but are still predicted at output.
Big Bird: Transformers for Longer Sequences
We run this entire sequence through the BigBird transformer encoder and then predict corresponding to the masked positions, based on the context provided by the other non-masked tokens in the sequence. . For 50% of the time the second sequence comes from true sequence after the first one. Remaining 50% of the time it is a a random sequence from the full dataset. The model is then required to predict this relationship using the output corresponding to the [CLS] token, which is fed into a simple binary classification layer.
Big Bird: Transformers for Longer Sequences
The sequence of steps is visually elaborated in Fig. 8. The model is trained with both MLM and NSP together. Training hyperparameter is provided in second columns of Tab. 21. In all experiments we use a learning rate warmup over the first 10,000 steps, and linear decay of the learning rate. We additionally performed a simple ablation study to validate the hypothesis, that similar to NLP, having a larger context improves performance. We use MLM task described above to test how BigBird performed with sequences of different length. Accuracy on MLM task with increasing sequence length is shown in Fig. 7. Not only longer context improves final accuracy, it also leads to faster learning, as we have now more opportunities for masking.
Big Bird: Transformers for Longer Sequences
The promoter region plays an important role in transcription initiation and thus its recognition is an important area of interest in the field of bioinformatics. Following Oubounyt et al. [72], we use datasets from Eukaryotic Promoter Database (EPDnew) [25], which contains 29,597 promoter region in the human genome. Around the transcription start site (TSS), we extract a sequence of 8000 bp (-5000 +3000 bp) from the human reference genome GRCh37. Since EPDnew uses newer GRCh38, we convert to GRCh37 coordinates using LiftOver [45].
Big Bird: Transformers for Longer Sequences
Following Oubounyt et al. [72] for each promoter region example, a negative example (nonpromoter sequences) with the same size of the positive one is constructed as follow: The positive sequence is divided into 20 subsequences. Then, 12 subsequences are picked randomly and substituted randomly. The remaining 8 subsequences are conserved. This process is illustrated in Figure 1 of [72]. Applying this process to the positive set results in new non-promoter sequences with conserved parts from promoter sequences (the unchanged Context 5000 bp  . These parameters enable generating a negative set that has 32 and 40% of its sequences containing conserved portions of promoter sequences.
Big Bird: Transformers for Longer Sequences
We prefix and append each example with [CLS] and [SEP] token respectively. The output corresponding to the [CLS] token from BigBird transformer encoder is fed to a simple binary classification layer. We fine-tune the pretrained BigBird from App. F.1 using hyperparameters described in Tab. 21. We note that high performance is not surprising due to the overlap in the nature of negative example generation and MLM pretraining.
Big Bird: Transformers for Longer Sequences
The first step of sequence-based algorithmic framework for predicting non-coding effects is to build a model to predict, large scale chromatic profile [109]. In this paper, we use the dataset provided in Zhou and Troyanskaya [109] 7 , to train BigBird to predict the chromatic profile.
Big Bird: Transformers for Longer Sequences
Each training sample consists of a 8,000-bp sequence from the human GRCh37 reference genome centered on each 200-bp bin and is paired with a label vector for 919 chromatin features. As before, we prefix and append each example with [CLS] and [SEP] token respectively. The output corresponding to the [CLS] token from BigBird transformer encoder is fed to a linear layer with 919 heads. Thus we jointly predict the 919 independent binary classification problems. We fine-tune the pretrained BigBird from App. F.1 using hyper-parameters described in Tab 
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Transformer models such as BERT , and other variants (Liu et al., 2019;Radford et al., 2019;Raffel et al., 2019a;Lewis et al., 2020) have achieved state-of-the-art results on many challenging NLP tasks. Moreover, recent work in longinput transformers (Ainslie et al., 2020;Zaheer et al., 2020b;Beltagy et al., 2020;Tay et al., 2021) has shown that increasing the input length a Transformer is able to process results in further performance gains. Additionally, it is also known that increasing model size also leads to performance gains in many tasks (Kaplan et al., 2020).
LongT5: Efficient Text-To-Text Transformer for Long Sequences
In this paper, we present a new model, called LongT5, with which we explore the effects of scaling both the input length and model size at the same time. To achieve this, we integrate long-input * Equal contributions. † Corresponding authors.   (Cohan et al., 2018) with different input length (x axis). Baseline models: HAT-BART (Rohde et al., 2021), BigBird-PEGASUS (Zaheer et al., 2020b), PRIMER (Xiao et al., 2021), LED (Beltagy et al., 2020). The size of circle roughly indicates the # of parameters for each model. transformer attention and pre-training ideas into the scalable T5 (Raffel et al., 2019a) model architecture. The resulting model, as shown in Figure 1, achieves state-of-the-art performance on several tasks which require handling long sequence inputs.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Regarding attention, we design a new attention mechanism, which we call Transient Global (TGlobal), that mimics ETC's local/global mechanism (Ainslie et al., 2020). Importantly, TGlobal attention removes the need for the additional side inputs in ETC, in order to fit within the T5 architecture. The main idea of ETC's local/global mechanism is to introduce local sparsity in the attention mechanism to reduce the quadratic cost when scaling to long inputs. Specifically, ETC only allows tokens in the input (called the long input) to attend to a local neighborhood, and adds a secondary input called the global memory, through which tokens in the long input can attend to each other indirectly. One disadvantage of this mechanism is that it requires designing this secondary global input for each new problem. In order to adapt it to T5, our new TGlobal mechanism synthesizes these global tokens on the fly (as aggregations of groups of tokens in the input), at each attention layer. Our experiments show that this mechanism results in only a small degradation in performance with respect to full attention in the same input length but allows the model to scale to much larger input lengths, resulting in significant performance gains.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Regarding pre-training, we adopt the pretraining strategy in the PEGASUS (Zhang et al., 2019a) model. This pre-training strategy was originally designed for abstractive summarization, but in our experiments, we found it also improves model performance for other tasks, such as question answering, and hence we adopted it in LongT5.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
The key idea is to mask out key (principle) sentences from a document and ask the model to reproduce them as a single string, as if it was a summary.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
We evaluate LongT5 on several summarization and question answering tasks (see Sections 4.2.1 and 4.3.1 for detailed descriptions of these datasets). Thanks to the scaling of both input length and model size, we achieve state-of-the-art results on many of them.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
The main contributions of this work are:
LongT5: Efficient Text-To-Text Transformer for Long Sequences
• A new Transformer architecture, LongT5, that allows for scaling both input length and model scale at the same time.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
• A new attention mechanism (TGlobal), which mimics ETC's local/global mechanism but is a drop-in replacement to regular attention for existing Transformer architectures like T5.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
• An analysis of model performance when varying both input length and model size of vanilla T5 and LongT5 models (pushing both models up to the maximum lengths they can handle before encountering memory issues), to understand the trade-offs in both performance and computation cost.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
• State-of-the-art results on the arXiv, PubMed, BigPatent, MediaSum, and TriviaQA datasets. For Natural Questions, we used a slightly different formulation than the original tasks, and hence we do not make state-of-the-art claims.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
• We open source our model architecture 1 and training code, as well as pre-trained model checkpoints on GitHub 2 . 
LongT5: Efficient Text-To-Text Transformer for Long Sequences
T5 (Raffel et al., 2019a) is a transformer based textto-text pre-trained language model that is gaining popularity for its unified framework that converts all text-based language problems into a text-to-text format, and its ease to scale up in number of parameters (from 60M to 11B parameters) with model parallelism. With full attention transformer, T5 has been successfully applied to many NLP tasks, but the tasks only require shorter input sequences. This is due to the limitation of quadratic computation growth with respect to input sequence length, resulting in larger memory consumption and longer training time. Recently, Press et al. (2021) explored scaling up T5 style models at inference time to longer sequences than seen during training, but how to scale up T5 style models in the input sequence length during training remains underexplored.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
We extend the original T5 encoder with globallocal attention sparsity patterns (Ainslie et al., 2020;Zaheer et al., 2020a) to handle long inputs. For the work reported in this paper, we used a standard T5 decoder since all of the tasks we considered require relatively short output sequence lengths.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Architecturally, the main difference between T5 and LongT5 lies in the attention mechanism. We experiment with two attention mechanism variations for LongT5, illustrated in Figure 2: (1) Local Attention and (2) Transient Global Attention (TGlobal). Both variations preserve several properties of T5: relative position representations, support for example packing, and compatibility with T5 checkpoints.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
For Local Attention, we simply replace the encoder self-attention operation in T5 with a sparse slidingwindow local attention operation following the implementation in ETC (Ainslie et al., 2020). Specifically, for a given local radius r, this formulation only allows each token to attend r tokens to the left and right of it (see Figure 2.a). We found r = 127 to be sufficient in practice, where r is the number of neighboring tokens to the left and to the right.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Local Attention does not introduce any new parameters and easily accommodates the attention masking required for example packing 3 . For a 3 Example packing refers to packing more than one short Attention keys Attention queries
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Attention queries
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Each global token is the result of averaging k input tokens.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Each input token can attend to its neighborhood (like in local attention), plus to all global tokens.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Each input token can attend to its neighborhood: r tokens to the left, and r tokens to the right. given choice of r, complexity is linear in input sequence length l: O(l × r).
LongT5: Efficient Text-To-Text Transformer for Long Sequences
To allow input tokens to interact with each other in each layer of the encoder at a longer range than Local Attention's local radius, we introduce Transient Global Attention as a modification of ETC's globallocal attention in a "fixed blocks" pattern. Namely, we divide the input sequence into blocks of k tokens, and for each block we compute a global token by summing (and then normalizing) the embeddings of every token in the block (see Figure 2.b). Now when computing attention, we allow each input token to attend not only to nearby tokens like in Local Attention, but also to every global token. We call these global tokens transient because in contrast to ETC-like global-local attention patterns, these tokens are dynamically constructed (and subsequently discarded) within each attention operation, removing any requirement for deciding which input tokens should be treated as "global". TGlobal attention only introduces a couple new parameters 4 : (1) T5-style relative position biases representing the distance from an input token's block to the block of each global token it's attending to, and (2) T5-style layer normalization parameters for normalizing each global token's embedding. The rest of the parameters are identical to T5, and we accommodate sequence packing by additionexample in the same input sequence to increase training efficiency. This is specially useful in LongT5, since with the large input lengths used in our model, if many examples are short, most of the input sequence would be dedicated to padding, wasting significant computation. 4 For base models, we introduced 10k additional parameters, 25k for large, and 50k for xl. ally masking attention from input tokens to global tokens of other examples. We found block size k = 16 to be sufficient in practice. Notice thus, that TGlobal attention introduces a block of l * l/k additional attention key-value pairs to calculate on top of Local Attention (l input tokens, attending to l/k global tokens; represented by the right most rectangle in Figure 2.b), hence for input sequence length l, complexity is O(l(r + l/k)).
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Generation Pre-training T5 is pre-trained with a span corruption objective, where spans of consecutive input tokens are replaced with a mask token and the model is trained to reconstruct the masked-out tokens. While it is effective, recent work on masked language modeling (MLM) (Liu et al., 2019;Zhang et al., 2019b) shows that carefully selecting the prediction objective could lead to significantly better performance. One argument is that predicting more informative tokens from the text could force the model to learn better semantics of the text. Motivated by that, we explore masking and generating the principle sentences from the text. In particular, we adopt the Gap Sentences Generation with Principle Ind-Uniq strategy from Zhang et al. (2019a), which was used for summarization pre-training. Following Zhang et al. (2019a), we select top-m scored (Principle) sentences based on ROUGE-F1 score (Lin, 2004) 
LongT5: Efficient Text-To-Text Transformer for Long Sequences
where i is the sentence index, D is the collection of sentences in the document. Each sentence is scored independently (Ind), and each n-gram is only counted once (Uniq).
LongT5: Efficient Text-To-Text Transformer for Long Sequences
LongT5 is implemented using JAX 5 and the Flaxformer 6 library. Following the same setup as T5.1.1 7 , we consider models of 3 sizes: base (∼220M), large (∼770M), and xl (∼3B), and use the same cased English SentencePiece vocab model used by T5.1.1, which contains 32000 sentence pieces. We use batch size of 128 and Adafactor as the optimizer in all experiments. We decide to use greedy decoding instead of beam search for all our experiments even with the test sets, therefore, our results reported below could potentially be improved further by using beam search, but we would like to make the setup consistent with our dev setup.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
We pre-train LongT5 models for 1M steps on 4096 input sequence length and 910 output sequence length. We use the same inverse squareroot learning rate schedule as T5, with learning rate set to 1/ max(step, warm_up steps), where warm_up steps is set to 10000. The same as T5.1.1, we pre-train LongT5 only on the C4 dataset (Raffel et al., 2019b), and we do not apply dropout during pre-training. As described in section 3.2, we use the PEGASUS Principle Sentences Generation objective as our pre-training objective. The configuration is similar to what was described by Zhang et al. (2019a) for their larger models, except for the masked sentence ratio in which we use a value of 0.2 instead of 0.45 8 . In section 5.3, we will show our ablation study between Principle Sentences Generation and Span Corruption.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
For fine-tuning, we use a constant learning rate of 0.001 and dropout rate of 0.1 for all tasks. For summarization tasks, we experiment with values of 4096, 8192, and 16384 for input lengths and 512 for output lengths. For QA tasks, we experiment with values starting at 512 and scale up to 36864 for input lengths and 128 for output lengths.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
We choose to benchmark our models on summarization tasks that cover various context lengths, because of their long context understanding and generative nature.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
LongT5 was benchmarked on the following six datasets. (Nallapati et al., 2016) News from CNN and Daily Mail are used as input and the article's summary bullets are the target summary.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
PubMed (Cohan et al., 2018) Scientific documents were collected from PubMed, with a document's content used as input and its corresponding abstract as the target summary.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
arXiv (Cohan et al., 2018) Similar to PubMed, but with documents taken from arXiv.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
BigPatent (Sharma et al., 2019) U.S. patent documents, with the patent's details used as input and the patent's abstract as the target summary.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
MediaSum (Zhu et al., 2021) Interview transcripts from CNN and NPR were used as input and their corresponding topic and overviews used as the target summary.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Multi-News (Fabbri et al., 2019) The task involves summarizing multiple news documents about a topic into a human-written summary. Table 1 provides statistics for the number of examples in train, validation, and test splits, and the average, median, max, and 90th percentile input sequence length. As can be seen, these datasets are long in input length, and would benefit from models that can model lengthier inputs. We included the CNN / Daily Mail dataset to benchmark on a common task, especially to see how using TGlobal attention impacts the model, despite the length of the inputs being smaller than the other datasets.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
We compare LongT5 with various top approaches: BigBird-PEGASUS (Zaheer et al., 2020b), HAT-BART (Rohde et al., 2021), DANCER PEGASUS (Gidiotis and Tsoumakas, 2020), PRIMER (Xiao et al., 2021), TG-MultiSum (Cui and Hu, 2021), LED (Beltagy et al., 2020), and an application of BART by Zhu et al. (2021). For these comparisons, we use common evaluation metrics of ROUGE-1, ROUGE-2, and ROUGE-L.   As can be seen in Table 2, LongT5 is able to achieve state-of-the-art rouge scores for arXiv, PubMed, BigPatent, and MediaSum. For arXiv and PubMed, which are composed of longer inputs, being able to scale up to 16k input length helps LongT5 achieve strong results.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
One dataset where LongT5 is not able to achieve state-of-the-art results is with Multi-News. LongT5 is the 2nd best model, slightly worth than PRIMER. This is understandable as the PRIMER model was pre-trained on a large corpus of documents related to news events, thus exposing the model to a similar corpus as that seen in Multi-News.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
When looking at CNN / Daily Mail, we can see that LongT5 was comparable with HAT-BART, despite not having full attention. LongT5 did at least get stronger scores in the ROUGE-2 metric.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
For the evaluation on QA tasks, we choose two popular benchmarks, Natural Questions and TriviaQA, that require long context understanding.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
NaturalQuestions (NQ) Questions are real queries issued by multiple users to Google search that retrieve a Wikipedia page in the top five search results. Answer text is drawn from the search results (Kwiatkowski et al., 2019).
LongT5: Efficient Text-To-Text Transformer for Long Sequences
The original NQ dataset asks models to predict a short answer (including no-answer or yes/no) and a long answer. We framed the task as a seq2seq task and ignored the long answer. Hence, our results focus only on short answer. Moreover, since our models predict answer texts instead of answer spans, our evaluation method differs slightly from the leader boards, and our results are not directly comparable to other existing approaches: (1) Since only the train and dev sets are publicly available, we use 90% of the official train set for training while using 10% as hold-out dev set to fine-tune the hyperparameters and training epoch, and use   the official dev set as our test set.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
(2) We benchmark LongT5 against the corresponding T5.1.1 models instead of directly comparing to the leader boards.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
TriviaQA Trivia enthusiasts authored questionanswer pairs. Answers are drawn from Wikipedia and Bing web search results, excluding trivia websites (Joshi et al., 2017). We use the official train/validation splits for training and fine-tuning the hyperparameters and training epoch, then re-train that model combining both train and validation sets to evaluate on the Wikipedia domain on the leader board 9 . Table 3 shows the dataset statistics for the number of examples in train and validation splits, and the average, median, max, and 90th percentile input sequence length. Table 4 shows a summary of the results for the NQ and TriviaQA datasets (see Appendix B for full results). For each dataset, we show two metrics: EM (Exact Match) and F1 score (evaluating precision and recall of individual words in the answer compared to the ground truth, ignoring stop words).
LongT5: Efficient Text-To-Text Transformer for Long Sequences
For NQ, we compare T5.1.1, LongT5 with Local Attention, and LongT5 with TGlobal attention. We decided to run T5.1.1 (1) with the default 512 input sequence length 10 and (2) with the largest input sequence length that can fit into device memory 11 , and use those as baselines. Since we are comparing against T5.1.1, for LongT5 experiments we report results at 512 input length for base and large, and the largest input length allowed by each model before running out of memory on the same hardware configuration used in our T5.1.1 experiments.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
As the table shows, increasing input length generally results in significant benefits in NQ, with models with larger input lengths significantly outperforming those with smaller input lengths in most cases. Some times, models with the largest input lengths underperform those with 4k length, but we believe those to be due to noise in the experiments, as results are the output of just one repetition of each experiment due to resource constraints. Moreover, while LongT5 with Local Attention often underperforms T5.1.1, LongT5 with TGlobal attention significantly outperforms T5.1.1. For example, considering the large size models, T5.1.1 was able only to scale up to an input length of 3k tokens, while the TGlobal model was able to reach 6k tokens, outperforming T5.1.1 at 4k token length (there was a dip at 6k token length, but we hypothesize this is just due to variance, as we only did one run for each configuration). For TriviaQA, we compare LongT5 with various top approaches on the leader board: BigBird-ETC (Zaheer et al., 2020a), Fusion-in-Decoder (Izacard and Grave, 2021), and ReadTwice (Zemlyanskiy et al., 2021). As shown in Table 3, TriviaQA inputs are quite long, therefore being able to scale up both in model size and to 16k input length helps LongT5 achieve state-of-the-art.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
In order to evaluate the training speed and memory consumption of LongT5, compared to T5.1.1, we performed a series of training runs in the NQ data set starting at input length 512, and increasing We can see that at shorter lengths (512), T5.1.1, LongT5 Local, LongT5 TGlobal have similar speeds, but as we increase the sequence length, LongT5 becomes significantly faster. For example at sequence length 2048, T5.1.1 base can only process 479 sequences per second, while LongT5 (base TGlobal) can process 765 and LongT5 (base Local) can process 860. The differences grow even larger as sequence length increases.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Another important fact that Figure 3 shows is that T5.1.1 models reach their out of memory point much earlier. For example, we could only scale up to 6k tokens for T5.1.1 base. On the other hand, LongT5 (base Local) can go up to 36k tokens in length, and LongT5 (base TGlobal) up to 12k. Large models show a similar picture with T5.1.1 large going only up to 3k, but the LongT5 variants going to 10k (large Local) and 6k (large TGlobal).
LongT5: Efficient Text-To-Text Transformer for Long Sequences
This section presents a similar analysis, but where we plotted model speed versus performance in NQ (F1 score). Results are shown in Figure 4 for models with large size. Each point in the curves is annotated with the corresponding sequence length.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
As Figure 4 shows, performance increases significantly as input length increases, highlighting the benefits of LongT5. Moreover, input length by itself is not enough to achieve good performance in all datasets, and in particular, in the NQ dataset (used in this figure), using Local Attention significantly hurts performance when compared with TGlobal or with T5.1.1. So, even at very long input lengths, LongT5 with Local Attention just matches T5.1.1 with input length of 3k in NQ. However, LongT5 with TGlobal attention outperforms T5.1.1. Moreover, note that although the plot shows a few irregularities (such as 8k length for LongT5 with Local Attention, or 6k length with TGlobal Attention), that is because the plot shows only the results of a single run, and hence there is some noise. However, trends can clearly be seen.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
As mentioned in section 3.2, we use PEGASUS Principle Sentences Generation instead of default Span Corruption used in T5 as our pre-training objective. Table 5 shows our ablation study for fine-tuning on NQ and arXiv from a model pretrained using the default Span Corruption objective, a model pre-trained with Principle Sentences Generation, and a model pre-trained with both objectives. The comparison is done on the dev set of the tasks, and with TGlobal base models. Both pretraining and fine-tuning on the models mentioned above are done with input sequence length 4096. The table shows, even though Principle Sentences Generation was developed by Zhang et al. (2019a) as a pre-training strategy for summarization, it benefits both summarization and QA tasks, but using both objectives together perform worse than just using PSG. Table 6 shows an additional ablation study with arXiv and PubMed, where we compare using regular T5.1.1 with Span Corruption compared to T5.1.1 pretrained with Principle Sentences Generation while using the same pre-training input sequence length of 512 (as was done in the original T5.1.1 pre-training task). As expected, Principle Sentences Generation helped the model achieve better results compared to Span Corruption when seeing the same amount of pre-training data. We  Table 5: Ablation study on dev set for different pretraining strategies using span corruption (SC) vs. principle sentences generation (PSG) and the effects on NQ and arXiv fine-tuning tasks. The models are TGlobal base, and fine-tuning is done with input sequence length 4096.  Table 6: Ablation study on arXiv and PubMed for different pre-training strategies using span corruption (SC) vs. principle sentences generation (PSG) with T5.1.1 model along with LongT5 with TGlobal attention. Fine-tuning was done on large model size, with input sequence length of 4096 except where otherwise noted. also compare this with dev scores from LongT5 with TGlobal attention at 4k and 16k input lengths, such that we can see having full attention will allow for better results, but being able to scale to longer input sequence lengths allows LongT5 to achieve its stronger results.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
This paper presented a new Transformer-based neural model called LongT5, with which we have explored the effects of scaling both input length and model size at the same time. Specifically, the main differences of LongT5 with respect to T5.1.1 are (1) a new scalable attention mechanism called Tran-sient Global attention, which is a drop-in replacement to the standard T5 attention mechanism, and hence can be used without needing additional sideinputs to the model or modifications to the model inputs; and (2) using a PEGASUS-style Principle Sentences Generation pre-training objective.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
Via experimentation in several challenging summarization and question answering datasets, we have explored the performance gains that can be achieved by scaling both input length and model size, resulting in state-of-the-art results on several datasets: arXiv, PubMed, BigPatent, MediaSum, and TriviaQA.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
As part of our future work, we would like to pursue several directions such as studying efficient attention mechanisms in the decoder and decoder-toencoder attention pieces of the model (both Local Attention and TGlobal attention are only applied to the encoder in LongT5 for now). Additionally, we would like to incorporate additional long-input transformer ideas into the LongT5 architecture, that could further improve model efficiency. Table 8 shows the full set of results on the summarization datasets used in this paper. This includes both standard T5 model (using version T5.1.1), T5 with PEGASUS Principle Sentences Generation pre-training, and LongT5 model.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
As can be seen, scaling up the input size for the models helps achieve better performance metrics. T5 models though struggle when scaling up to 4k for input, as the fine-tuning task can take many days even when using a large topology of TPUv3.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
When comparing regular T5.1.1 model with a T5.1.1 model using PEGASUS Principle Sentences Generation pre-training, the latter was able to achieve better results, with the results also improving as the input size scaled up. This helps show that both using the latter pre-training objective along with scaling up allows us to get the best results from these models.
LongT5: Efficient Text-To-Text Transformer for Long Sequences
LongT5, despite having a reduced attention from using TGlobal attention, is able to get strong performance results due to both scaling up to larger inputs and leveraging the Gap Sentences Generation pre-training strategy. Table 7 shows the full set of results comparing T5.1.1 and LongT5 models on the QA datasets used in this paper. For both NQ and TriviaQA in this comparison study, we use 90% of the official training set for training while using 10% as holdout dev set to fine-tune the hyperparameters and training epoch, and use the official dev set to report the numbers in this table. We run each model to the largest input length allowed before running out of memory on specific hardware configurationbase/large models on 4x8 TPUv3 with no model partitioning, and xl models on 8x16 TPUv3 with 8 partitions.   Table 8: Summarization results comparing T5, T5 with PEGASUS-style Principle Sentences Generation (PSG) pre-training, and LongT5 with best known approaches for the various datasets. All T5 scores are with standard T5.1.1 model. All LongT5 scores are with models using TGlobal attention. For each task, we scale up the input length depending on the statistics of the inputs, thus not all of the tasks were scaled to 16k. We do not include input length of other models because each model uses the input differently, and hence, direct comparison is not possible.