Adapting Pretrained Text-to-Text Models for Long Text Sequences

Our model is based on a standard transformer with block-sparse self-attentions (Zaheer et al., 2020) on the encoder side. While various new architectures (Wang et al., 2020; Choromanski et al., 2021; Lei, 2021; Gu et al., 2021) have been proposed, we stick to the simple architecture for the following reasons: 1) it makes it easy to reuse existing pretraining pipelines, which are often highly optimized specifically for vanilla transformers, e.g., learning rate schedules, normalization layers, optimizers; 2) using local attentions, where each token attends to only tokens in the local context, allows our model to reuse all the model parameters from existing PLMs, while other attention variants use different parameterizations that prohibit inheriting the weights of an existing pretrained model.

In addition to block attention, we investigate three mechanisms that enable long-range connections in the encoder:

1) Global-token mechanism: Previous work (Guo et al., 2022; Zaheer et al., 2020; Beltagy et al., 2020) has proposed augmenting block-sparse attention with a small set of “global tokens” that attend to the entire sequence and hence enable long-range interactions in the encoder. Specifically, we mark the first 64 tokens in each attention block as global tokens and share the projection matrices for both the global and regular tokens. This mechanism has proven effective in encoder-only models, especially for question answering tasks as shown by the aforementioned methods.

2) Overlapping (strided) attention windows: Sliding-attention with overlap is a straightforward way to introduce long-range connections in local attention models. As we stack the layers in the encoder, the receptive field of each token would increase exponentially. For example, Beltagy et al. (2020) use the stride of one token and each token attends to an equal number of tokens from both sides. We develop a simpler and faster block-wise version which makes the parallelization easier to implement; namely, tokens in each block will attend to all the tokens inside the block, and half of the tokens from its immediate left and right blocks.

3) Pooling layers: Recent work (Zhang et al., 2021; Pang et al., 2022a) has explored using pooling operations to reduce the number of key and value states in transformers. We implement a simpler version that only requires standard average pooling operations. All illustration of the pooling-augmented attention layer is shown in Figure 1. Specifically, in the top n layers of the transformer encoder, we add a second attention module which takes as input the hidden states output by the $i$th block self-attention layer $\mathbf{X_{i}}\in\mathbb{R}^{L\times h}$, where $L$ is the sequence length and $h$ is the size of the hidden states. As in the vanilla attention layers, $\mathbf{X_{i}}$ is first projected to create the key, query, value matrices $\mathbf{Q}_{i}^{p},\mathbf{K}_{i}^{p},\mathbf{V}_{i}^{p}\in\mathbb{R}^{L\times h}$.²²2The projection layers to create these matrices are not used in existing pretrained models and will be randomly initialized before further pretraining We first average pool the $\mathbf{K}_{i}^{p}$ and $\mathbf{V}_{i}^{p}$ sequences, with a fixed kernel/stride size, into smaller lengths $\tilde{\mathbf{V}}_{i}^{p}$, $\tilde{\mathbf{K}}_{i}^{p}$ $\in\mathbb{R}^{\tilde{L}\times h}$, where $\tilde{L}\ll L$. We then apply standard attention using $\mathbf{Q}_{i}^{p}$, $\tilde{\mathbf{K}}_{i}^{p}$ and $\tilde{\mathbf{V}}_{i}^{p}$ resulting in $O(L\times\tilde{L})$ complexity. The output of the pooling layers is added with $\mathbf{X_{i}}$ to form a residual connection.

We compare these variants via the performance on downstream long-sequence tasks in Sec 3.2.

The choice of the corpus has a significant impact on the downstream results. We consider long documents from formal text domains, including Books3 (Gao et al., 2021), STORIES (Trinh and Le, 2018), RealNews Zellers et al. (2019); and long dialogues including MediaSum (Zhu et al., 2021) and OpenSubtitles (Tiedemann, 2016). While collecting a long-document corpus seems to be a natural choice for long-sequence downstream tasks, as they are more likely to include long-range dependencies than common short texts on the internet, pretraining only on these datasets also brings the risk of overfitting to specific domains, instead of achieving consistent gains on a range of tasks. Thus, we also consider a general-domain corpus – C4 as used by T5 (Raffel et al., 2020). Additionally, instead of using randomly concatenated sequences, we also tried to concatenate semantically similar C4 documents (using similarity metric learned by dense retrieval models) with the hope that the model can learn to capture more long-range dependencies across relevant documents. We discuss the effects of these corpus variants in Sec 3.3.

A variety of self-supervised pre-training objectives have been proposed for sequence-to-sequence models (Lewis et al., 2020; Raffel et al., 2020; Guo et al., 2022). In the long document setting, we ideally seek an objective that promotes long-range reasoning ability in the model. We investigate the following different pretraining objectives and the effect of input length during pretraining.

1) T5 Span Denoising: Applying BART’s denoising objective to long sequences is computationally expensive as it requires reconstructing the entire input and incurs significant computation overhead on the decoder-side attention. Moreover, reconstructing the entire input would be at odds with most downstream tasks such as question-answering and summarization, which require generating shorter text. Thus, we adopt T5-style denoising for pretraining our model, i.e., we randomly pick a set of spans in the input sequence as the decoding target and mark them with special sentinel tokens. The model is then trained to generate the uncorrupted spans. This objective is readily applicable to long documents as we can control both the length and the number of spans. We experiment with both fixed span lengths as in (Raffel et al., 2020), and also mixed span lengths with both short and long spans, with which we hope the model is able to perform well on a range of tasks requiring differing output lengths.

2) Pegasus – Primary Sentence Prediction: Originally proposed for summarization pretraining in (Zhang et al., 2020) and recently used for long documents by Guo et al. (2022), this objective identifies and masks out a set of principle sentences, i.e., sentences with a high ROUGE score with the rest of the document. The model is then trained to generate these principle sentences. The output length can be controlled by choosing the number of principle sentences to mask.

3) Model-based Denoising: Apart from randomly selecting the decoding targets, we also explore a novel model-based objective. Here we use a separate encoder-only model (with local attention) to select decoding targets for the sequence-to-sequence model. This approach is inspired by ELECTRA (Clark et al., 2020) and we hope the prediction loss of the encoder-only model can be a good proxy to select spans that require long-range dependencies to predict. Specifically, we first mask a larger number of tokens (5,120 tokens instead of 1,024) in the input sequence. We then apply an encoder-only masked language model to recover the masked spans. Based on the losses of the masked language model, we only keep the top 20% hard spans to train the text-to-text model. The encoder-only model can either be frozen or jointly trained with the sequence-to-sequence model.

We evaluate the models on six summarization datasets and four QA datasets. The summarization datasets are from formal domains, including GovReport (Huang et al., 2021), ArXiv & PubMed (Cohan et al., 2018) and BookSum Chapters (Kryściński et al., 2021); or informal conversational domains , such as TVMegaSite & ForeverDreaming (Chen et al., 2022). For QA, we consider Qasper (Dasigi et al., 2021), which contains questions over NLP papers; QMSum³³3QMSum is proposed as a ”query-based summarization” dataset. We consider it as a special case of QA as our model uses the same input format for QMSum and other QA datasets. (Zhong et al., 2021), longform QA over meeting scripts, and two QA datasets on books: QuALITY (Pang et al., 2022b) and NarrativeQA (Kočiský et al., 2018).

We finetune the pretrained model with a maximum of 16,384 tokens. For long-sequence QA tasks, we adopt the input format as used by the state-of-the-art open-domain QA system (Izacard and Grave, 2021) and the query-based summarization model (Vig et al., 2022). Specifically, we repeat the question/query at the start of each attention block. For finetuning, we utilize the robust finetuning technique proposed by Aghajanyan et al. (2021). We also conduct a grid search over finetuning hyperparameters, such as learning rate and dropout rate, the details of which are presented in Table 8 in Appendix A. We report ROUGE⁴⁴4https://github.com/pltrdy/files2rouge scores for summarization datasets, and for QA we use Exact Match (EM) scores for datasets with short answers and F1 scores for datasets with long answers.

We study the effectiveness of different model choices before launching the resource-consuming pretraining. We first initialize a base-size block-attention model using BART’s weights. We augment the model with three additional long-range mechanisms, as described in Sec 2.1. Note that only the pooling layers introduce additional parameters that will be randomly initialized. Table 1 shows the results on both QA and summarization tasks. For the global-token mechanism, we mark the first $64$ tokens of each block as global tokens. We see that pooling layers produce the most consistent improvements even for GovReport, where the baseline already achieves strong numbers. Consistent with a prior study on encoder-only models (Xiong et al., 2022), attention window overlaps fail to produce further improvements over the disjoint block-attention layers. Adding global tokens consistently helps on QA tasks but not on summarization tasks. We hypothesize that in encoder-decoder models, the cross-attention can offset the effect of global tokens, as each decoding position has access to all input tokens’ representations. When finetuning our final pretrained model, we also try to combine global tokens with pooling layers for QA tasks, but we do not observe further improvements.

With the assumption that models should be exposed to as many long dependencies as possible at pretraining time, we initially tried to only pretrain the model with natural long documents that are collected from sources like books, news, and TV dialogues. However, we did not achieve consistent improvements with this corpus alone. Instead, we found it is important to include sufficient documents from diverse domains, even if those documents are mostly short sequences. We present our ablation analysis in Table 2. Here we reported results on small summarization datasets where the gaps are more visible. Note that the sizes of long-document corpora are usually smaller than open-domain corpus. To remove the size factor that affects model performance, we limit the pretraining steps such that the model does not see repeated examples from each corpus. In Figure 2, we show the length statistics of each document source.

From the table, we see that pretraining on corpora that only have long documents, which are often from specific domains, hurts the downstream performance for most of the datasets, except for NarrativeQA, which is from a very close domain. On the other hand, pretraining on randomly concatenated C4 documents brings visible gains for most of the tasks. In addition to directly using concatenations of random C4 documents, we tried to assemble long sequences using semantically similar C4 documents, with the hope of creating more long-range connections in the pretraining sequences. For each document, we use a dense retrieval model (Izacard et al., 2021) to find similar documents and concatenate them as long pretraining sequences. We denote this corpus as “C4-linked". However, this new corpus is either similar or worse compared to directly using C4. We conjecture that it is because the retrieved documents may contain redundant information, making some of the masked spans trivial to predict — the training perplexity after 100k updates on “C4-linked" is significantly lower than that on the original C4 corpus (10.5 vs 12.2). Example sequences and details on how to build this corpus can be found in the Appendix.

We compare the effects of different pretraining objectives in Table 3. The generation targets are usually paragraph-length for QMSum, while Qasper expects the model to predict spans or single sentences most of the time. All the models are pretrained for $100k$ updates on the C4 corpus. To investigate the effect of pretraining sequence length, we compare the $16k$ model with a model pretrained with $8k$ sequence length. We double the batch size for the $8k$ length pretraining such that the input tokens in each batch stays the same. We also increase the masking ratio for the $8k$ model to $1/8$ so that the decoding sequence length remains 1,024. Note that under this setting, pretraining with $8k$-length batches is a bit slower compared to the $16k$ batches due to the decoder-side self-attention.

To further investigate the performance gains of our proposed model, we compare the performance of the proposed model against the base model as a function of source document length for two summarization datasets, namely SummScreen and TVMegaSite. To conduct our analysis, we divide the validation split of both the datasets into short and long documents. The cutoff length to separate the two groups is chosen such that approximately 75% of the documents are classified as short documents. Figure 3 presents the results of this comparison. For both the datasets: (a) there’s a performance drop for both the best and the base model for longer documents, and (b) the best model is better than the base model on all data splits. For SummScreen the performance gap between the best and the base model is bigger for long documents than for short documents – relative ROUGE-L increase of 0.80% and 3.96% for short and long documents respectively. This suggests that the performance gains for the best model can be attributed to better long-context modeling. For TVMegaSite this trend of increasing performance gap between the best and the base model with an increase in document length still holds true, though the increase in performance gap is modest in comparison to the increase observed for SummScreen – relative ROUGE-L increase of 2.43% and 2.75% for short and long documents respectively.

A long list of works propose to reduce the complexity of the attention layers of transformers. The simplest paradigm to achieve efficiency is to restrict each token’s attending context to a subset of the whole sequences, e.g., Reformer (Kitaev et al., 2020) and the Routing transformer (Roy et al., 2021) proposes hashing or clustering based attention, where each token only attends to tokens of a single bucket/cluster. Our model architecture is mostly influenced by previous work like Longformer (Beltagy et al., 2020), BigBird (Zaheer et al., 2020) and ETC (Ainslie et al., 2020) that have demonstrated strong downstream performance. These models assume strong locality bias in language data and restrict each token’s attending context to nearby tokens. In contrast to these works, we augment the block attention with pooling layers and study the effect of additional pretraining on long sequences. Other popular approach to tackle the efficiency bottleneck includes kernel-based (Choromanski et al., 2021; Peng et al., 2021) and low-rank approximation (Wang et al., 2020) of the N$\times$N attention matrix. However, in contrast to local attention transformers, the effectiveness of these approximation approaches is yet to be validated in large models and downstream tasks.

Apart from methods that solely modify transformer’s attention calculations, several recent works propose alternative architectures that are free of the quadratic complexity over input length by design. For instance. Perceiver IO (Jaegle et al., 2021) proposes to maintain a limited-length latent sequence instead of all tokens’ hidden states, and only conduct attention computation on the latent space. In another notable work, Gu et al. (2021) propose to use structured state-space models rooted in control theory to model the dependencies in long sequences. The model can be trained as a CNN with large kernel size, which is implemented efficiently with fast fourier transforms, and tested as a recurrent network that has $O(1)$ complexity. For these models to achieve reasonable performance on the downstream tasks we studied, it is necessary to pretrain them from scratch, which will require significant computation compared to our adapting approach. However, testing these emerging models on scaled settings and real-world long-sequence problems is indeed an important research problem and we hope to explore these alternative architectures in future work.

To apply pretrained models to long-sequence tasks, early work (Zaheer et al., 2020; Beltagy et al., 2020) simply reuses parameters from models pretrained on short sequences and replaces the encoder full attention with sparse local attentions. While the models are not exposed to long sequences at pretraining time, they demonstrates consistent improvements over previous models that can only take truncated inputs. Complementary to local attentions, Zhang et al. (2021) show that pooling layers can be inserted into a pretrained transformer at finetuning time and bring additional performance gains on summarization. Instead of relying on a single model that directly processes the whole input, Mao et al. (2022) proposes a two-stage extract-and-generate approach, where the extractor can leverage the supervision signal learned by the generator. However, despite the complicated training recipe, it does not bring consistent gains and underperforms our non-pretrain baselines. The most relevant work to ours is LongT5 (Guo et al., 2022), which adopts both global tokens as well as local attention, and pretrains the model with 4k text sequences from C4. Compared to LongT5, we augment local attentions with pooling layers and present a more comprehensive study on pretraining strategies. Without pretraining from scratch, we achieve stronger summarization performance. Concurrent to our work,  Phang et al. (2022) also present an empirical study on adapting short-text models for long document summarization. While their study mostly focuses on the architecture aspect, we present additional analysis on the choices of pretraining corpus and learning objectives.