SlowGuess/NAACL_2022_OCR huggingface dataset:
Mandy Guo*, Joshua Ainslie*, David Uthus*, Santiago Ontañón*  
Jianmo Ni, Yun-Hsuan Sung, Yinfei Yang
{xyguo, jainslie, duthus, santiontanon, jianmon, yhsung, yinfoei}  $@$  google.com
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 LongT5, a new model that explores the effects of scaling both the input length and model size at the same time. Specifically, we integrate attention ideas from long-input transformers (ETC), and adopt pretraining strategies from summarization pretraining (PEGASUS) into the scalable T5 architecture. The result is a new attention mechanism we call 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 and question answering tasks, as well as outperform the original T5 models on these tasks. We have open sourced our architecture and training code, as well as our pre-trained model checkpoints.
Transformer models such as BERT (Devlin et al., 2019), 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 long-input 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).
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
Figure 1: The average ROUGE score  $(R - 1 + R - 2 + R - L) / 3$  of LongT5 and baseline models on arXiv and PubMed summarization tasks (Cohan et al., 2018) with different input length ( $x$  axis). Baseline models: HATBART (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.
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
Findings of the Association for Computational Linguistics: NAACL 2022, pages 724 - 736 July 10-15, 2022 ©2022 Association for Computational Linguistics
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.
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. 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.
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.
- A new Transformer architecture, LongT5, that allows for scaling both input length and model scale at the same time.
- 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.
- 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.
- 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.
- We open source our model architecture<sup>1</sup> and training code, as well as pre-trained model checkpoints on GitHub<sup>2</sup>.
2 T5
T5 (Raffel et al., 2019a) is a transformer based text-to-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.
3 LongT5
We extend the original T5 encoder with global-local 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.
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.
3.1.1 Local Attention
For Local Attention, we simply replace the encoder self-attention operation in T5 with a sparse sliding-window 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.
Local Attention does not introduce any new parameters and easily accommodates the attention masking required for example packing<sup>3</sup>. For a
Published under the Flaxformer GitHub https: //github.com/google/flaxformer/tree/ main/flaxformer/architectures/longt5 2https://github.com/google-research/ longt5
Example packing refers to packing more than one short
Figure 2: Illustration of the two attention mechanisms we experimented with in LongT5.
given choice of  $r$ , complexity is linear in input sequence length  $l$ :  $O(l \times r)$ .
3.1.2 Transient Global Attention (TGlobal)
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 global-local 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: (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 addition
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))$ .
3.2 PEGASUS Principle Sentences 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) using  $s_i = \text{rouge}(x_i, D \setminus \{x_i\}, \forall i)$ , 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).
example 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.
4For base models, we introduced 10k additional parameters, 25k for large, and 50k for x1.
4.1 Configurations
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 ( $\sim$ 220M), large ( $\sim$ 770M), and xl ( $\sim$ 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.
4.1.1 Pre-training
We pre-train LongT5 models for 1M steps on 4096 input sequence length and 910 output sequence length. We use the same inverse square-root learning rate schedule as T5, with learning rate set to  $1 / \sqrt{\max(\text{step}, \text{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.
4.1.2 Fine-tuning
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.
4.2 Evaluation on Summarization Tasks
We choose to benchmark our models on summarization tasks that cover various context lengths, because of their long context understanding and generative nature.
LongT5 was benchmarked on the following six datasets.
CNN / Daily Mail (Nallapati et al., 2016) News from CNN and Daily Mail are used as input and the article's summary bullets are the target summary.
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.
arXiv (Cohan et al., 2018) Similar to PubMed, but with documents taken from arXiv.
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.
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.
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.
We compare LongT5 with various top approaches: BigBird-PEGASUS (Zaheer et al., 2020b), HATBART (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.
<sup>5</sup>https://github.com/google/jax
$^{6}$ https://github.com/google/flaxformer
<sup>7</sup>https://github.com/google-research/text-to-text-transfer-transformer/blob/main/releases_checkpoints.md#t511
<sup>8</sup>We briefly experimented with other values, but found 0.2 to work best with the downstream tasks of interest.
Dataset	Example Count	Input Length
Train	Validation	Test	Average	Median	Max	90th percentile
CNN / Daily Mail	287,113	13,368	11,490	982.39	894	5268	1659
arXiv	203,037	6,436	6,440	10,720.18	8,519	378,825	20,170
PubMed	119,924	6,633	6,658	4,747.97	3,883	452,915	8,883
BigPatent	1,207,222	67,068	67,072	6,537.32	5,236	294,004	11,328
MediaSum	443,596	10,000	10,000	2,302.02	1,748	125,974	4,128
Multi-News	44,972	5,622	5,622	2,593.81	1,902.5	683,544	4,853
Table 1: Statistics for the summarization datasets. Input length measured in tokens using a SentencePiece Model.
Approach	arXiv
R-1	R-2	R-L
DANCER PEGASUS	45.01	17.6	40.56
BigBird-PEGASUS (large)	46.63	19.02	41.77
HAT-BART	46.68	19.07	42.17
LED (large)	46.63	19.62	41.83
PRIMER	47.6	20.8	42.6
LongT5 (large - 16k input)	48.28	21.63	44.11
LongT5 (x1 - 16k input)	48.35	21.92	44.27
Approach	PubMed
R-1	R-2	R-L
DANCER PEGASUS	46.34	19.97	42.42
BigBird-PEGASUS (large)	46.32	20.65	42.33
HAT-BART	48.36	21.43	37.00
LongT5 (large - 16k input)	49.98	24.69	46.46
LongT5 (x1 - 16k input)	50.23	24.76	46.67
Approach	BigPatent
R-1	R-2	R-L
BigBird-PEGASUS (large)	60.64	42.46	50.01
LongT5 (large - 16k input)	70.38	56.81	62.73
LongT5 (x1 - 16k input)	76.87	66.06	70.76
Approach	MultiNews
R-1	R-2	R-L
TG-MultiSum	47.10	17.55	20.73
PRIMER	49.9	21.1	25.9
LongT5 (large - 8k input)	47.18	18.44	24.18
LongT5 (x1 - 8k input)	48.17	19.43	24.94
Approach	MediaSum
R-1	R-2	R-L
BART (large)	35.09	18.05	31.44
LongT5 (large - 4k input)	35.54	19.04	32.20
LongT5 (x1 - 4k input)	36.15	19.66	32.80
Approach	CNN / Daily Mail
R-1	R-2	R-L
HAT-BART	44.48	21.31	41.52
LongT5 (large - 4k input)	42.49	20.51	40.18
LongT5 (x1 - 4k input)	43.94	21.40	41.28
Table 2: Summarization results comparing LongT5 with best known approaches. LongT5 scores are with models using TGlobal attention. For each task, we scale up the input length depending on the inputs' statistics, thus not all are scaled to  $16k$ . For more results, please see Section A in the Appendix.
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.
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.
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.
4.3 Evaluation on QA Tasks
For the evaluation on QA tasks, we choose two popular benchmarks, Natural Questions and TriviaQA, that require long context understanding.
4.3.1 Datasets
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).
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
Dataset	Example Count	Input Length
Train	Validation	Test	Average	Median	Max	90th percentile
NQ	307,373	7,830	6,695.92	4,486	151,519	15,290.8
TriviaQA	87,622	11,313	10,832	69,082.51	45,011	1,174,918	150,643
Table 3: Statistics for the QA datasets. Input length measured in tokens using a SentencePiece Model.
Approach	NQ
EM	F1
T5.1.1 (base - 512 input)	50.93	52.54
T5.1.1 (base - 6k input)	56.73	56.73
T5.1.1 (large - 512 input)	57.29	60.68
T5.1.1 (large - 3k input)	60.09	64.17
T5.1.1 (xl - 4k input)	60.75	64.07
Local:
LongT5 (base - 512 input)	54.39	58.24
LongT5 (base - 36k input)	55.77	59.66
LongT5 (large - 512 input)	55.19	58.00
LongT5 (large - 10k input)	60.01	64.40
TGlobal:
LongT5 (base - 512 input)	55.73	59.06
LongT5 (base - 12k input)	58.12	62.44
LongT5 (large - 512 input)	57.55	61.53
LongT5 (large - 4k input)	60.77	65.38
LongT5 (large - 6k input)	59.17	63.38
LongT5 (xl - 8k input)	62.66	66.61
Approach	TriviaQA
EM	F1
BigBird-ETC (random attn)	80.86	84.5
Fusion-in-Decoder	80.09	84.35
ReadTwice	76.86	80.85
TGlobal:
LongT5 (base - 16k input)	74.67	78.9
LongT5 (large - 16k input)	78.38	82.45
LongT5 (xl - 16k input)	81.00	84.83
Table 4: QA results: (1) NQ results comparing T5.1.1 and LongT5. Base/large models are trained on  $4 \times 8$  TPUv3 with no model partitioning. Xl models are trained on  $8 \times 16$  TPUv3 with 8 partitions. (2) TriviaQA results compared to top models on leader board. LongT5 scores using Local and TGlobal attention. Full results in Appendix B.
the official dev set as our test set. (2) We benchmark LongT5 against the corresponding T5.1.1 models instead of directly comparing to the leader boards.
TriviaQA Trivia enthusiasts authored question-answer 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<sup>9</sup>.
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).
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.
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
<sup>9</sup>https://competitions.codalab.org/competitions/17208
10For base and large models.
<sup>11</sup>For base and large models, we used 4x8 TPUv3 and no model partitioning; for xl model, we used 8x16 TPUv3 and 8 partitions.
Figure 3: Sequences per second as a function of input length for T5.1.1, LongT5 with Local Attention and LongT5 with TGlobal attention. Input lengths start at 512, and go as far as possible before running out of memory. Measurements taken with batch size 128, on 4x8 TPUv3 slices. base and large model sizes shown.
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).