Investigating Efficiently Extending Transformers for Long Input Summarization (2023)

Abstract: While large pretrained Transformer models have proven highly capable at tackling natural language tasks, handling long sequence inputs still poses a significant challenge. One such task is long input summarization, where inputs are longer than the maximum input context of most models. Through an extensive set of experiments, we investigate what model architectural changes and pretraining paradigms most efficiently adapt a pretrained Transformer for long input summarization. We find that a staggered, block-local Transformer with global encoder tokens strikes a good balance of performance and efficiency, and that an additional pretraining phase on long sequences meaningfully improves downstream summarization performance. Based on our findings, we introduce PEGASUS-X, an extension of the PEGASUS model with additional long input pretraining to handle inputs of up to 16K tokens, which achieves strong performance on long input summarization tasks comparable with much larger models.

## introduction large pretrained transformer models have proven to be extremely capable at tackling natural language tasks @xcite . however, handling long textual sequences continues to be a significant challenge for these models. training models to handle long sequences is expensive in both computation and memory, and moreover requires training and evaluating on long sequence data, which is rarer and more costly to collect. given the broad success of transformer models on short-sequence language tasks, our goal is to investigate the best way to extend these models to handle longer input sequences. in this work, we focus on the task of long input summarization: summarizing long input documents into shorter text sequences. the inputs of such tasks are often significantly longer than the maximum input lengths of most standard transformer models, and hence warrant both architecture modifications as well as new training regimes. for instance, to avoid the quadratic growth in memory consumption of attention in transformers, many memory-efficient transformer variants have been proposed @xcite . however, the manner in which these changes are incorporated into models has been inconsistent and ad-hoc, and there are few established best practices. for instance, some works directly fine-tune on long-input summarization tasks @xcite , while others first perform additional pretraining @xcite . because of the high cost of training these models, there has yet to be a systematic study of how best to adapt models for long input sequences. hence, it has been difficult to establish which model and training changes are necessary or complementary. to answer these questions, we conduct an extensive empirical investigation into the architectural changes, model configurations and pretraining schemes to identify better approaches to training transformer models for long input summarization. we evaluate a set of efficient transformer variants, and propose a simple block-wise local transformer architecture with staggered blocks and global tokens that strikes a good balance of performance and memory efficiency. we show that given a fixed token budget, pretraining on short sequences and then pre-adapting the model to an efficient transformer architecture by training on longer sequences leads to better performance than only long input pretraining or no adaptation at all. we also investigate model design choices such as position encoding schemes, encoder-decoder layer distributions, and the impact of discrepancies between pretraining and fine-tuning architecture hyperparameters. based on the findings from our empirical investigation, we adapt the pretrained pegasus large model @xcite to tackle long input summarization on up to 16k input tokens. the resulting model, which we call pegasus-x, attains top scores on long summarization tasks, outperforming much larger models like longt5 @xcite . moreover, impact on short input summarization performance is minimal. a smaller version which we call pegasus-x base attains similar scores with much fewer parameters. beyond summarization, we believe that many of our findings will be useful to the community for efficiently adapting transformer models to handle ever longer input sequences for other tasks. in summary, our contributions are: ## experimental setup similar to @xcite , we perform the majority of our experiments with a pegasus basesized model, before applying our findings to pegasus large -sized model. ## pegasus-x based on our findings, we settle on the following recipe for adapting pegasus models @xcite to long sequence summarization. we experiment with two model sizes: pegasus-x (pegasus extended) based on pegasus large , and pegasus-x base based on a newly trained pegasus base model which we call pegasus base+ . 4 we initialize the weights of pegasus-x and pegasus-x base with the pretrained weights of pegasus large and pegasus base+ respectively. only two new sets of parameters are introduced: global token embeddings, and a new layernorm for the global input representations in each transformer layer. this is ∼1m more parameters for pegasus-x base and 2m more for pegasus-x. we initialize the global token embeddings by randomly sampling tokens from the input embeddings, and we initialize the layernorm weights with the regular input layernorm weights. the task-and model-specific hyperparameters for fine-tuning can be found in appendix 15. for this section, we report rouge-lsum foot_4 rather than rouge-l for consistency with the metrics reported in other papers and leaderboards. pegasus-xbase pegasus-x # parameters 272m 568m # global tokens 128 128 block size 512 512 batch size 512 1024 additional pretraining 300k steps 300k steps we also compare to reported results of big bird-pegasus 6 @xcite , led (beltagy et al., 2020), top-down transformer @xcite with both average-pool (avgp) and adaptive-pool (adap) variants, bart-ls @xcite , longt5-large and xl, and sled @xcite . led, top-down and sled are initialized with bart large weights with no additional pretraining on long input sequences. bart-ls is concurrent work that also incorporates staggered block-local attention and addition long-sequence pretraining, in addition to pooling layers and different pretraining data. pegasus-x outperforms big bird-pegasus on all tasks, and top-down-avgp on both compared tasks. although top-down-adap outperforms pegasus-x, it uses a much more complex fine-tuning setup, using an importance tagger on reference summaries to construct token pooling weights, whereas pegasus-x only uses standard fine-tuning. even so, pegasus-x still outperforms top-down-adap on pubmed. pegasus-x outperforms bart-ls on pubmed and slightly underperforms on arxiv; as mentioned above, google-research/blob/master/rouge/readme. md#two-flavors-of-rouge-l 6 big bird-pegasus only has a context of 3072 tokens, likely due to the larger memory consumption of big bird. pegasus-x and bart-ls share many similarities, and we see the strong performance of bart-ls as confirmation of the efficacy of parts of our recipe for longer sequence models. pegasus-x also outperforms longt5 on both arxiv and pubmed, despite both compared longt5 models having more parameters. however, we find that longt5 performs much better on bigpatent, which is a largely extractive summarization task. we hypothesize that a larger hidden size may improve extraction over very long sequences. ## conclusion in this work, we investigate a range of proposed improvements to transformer models to effectively and economically handle long inputs in summarization tasks. through extensive ablation experiments, we find a simple but effective recipe for extending short-input models to tackle long-input summarization. based on our findings, we introduce pegasus-x, an extended version of pegasus with a modified architecture and additional longsequence pretraining. we show that pegasus-x sets the state of the art on two long input summarization tasks (govreport and pubmed) and performs competitively on many others, even despite being much smaller than some compared models. our findings can be extended to models in other domains beyond summarization, both for pretraining long input models from scratch as well as extending already pretrained short sequence models.