[Paper Review] Investigating Efficiently Extending Transformers for Long Input Summarization

Jason Phang, Yao Zhao+1 more|arXiv (Cornell University)|Aug 8, 2022

Topic Modeling4 citations

TL;DR

This paper proposes PEGASUS-X, a 568M-parameter extension of the PEGASUS model that efficiently handles long input summarization up to 16K tokens using a staggered, block-local attention mechanism with learned global tokens. It achieves state-of-the-art performance on GovReport and PubMed while maintaining strong short-input performance, outperforming much larger models like LongT5 with minimal parameter increase and no model parallelism.

ABSTRACT

While large pretrained Transformer models have proven highly capable at tackling natural language tasks, handling long sequence inputs continues to be a significant challenge. One such task is long input summarization, where inputs are longer than the maximum input context of most pretrained models. Through an extensive set of experiments, we investigate what model architectural changes and pretraining paradigms can 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. PEGASUS-X achieves strong performance on long input summarization tasks comparable with much larger models while adding few additional parameters and not requiring model parallelism to train.

Motivation & Objective

• To identify the most effective architectural and pretraining modifications for adapting short-context pretrained Transformers to long input summarization.
• To evaluate trade-offs between model efficiency, memory usage, and performance in long sequence modeling.
• To develop a practical, low-cost recipe for extending existing small-to-medium models to handle long inputs without training from scratch.
• To minimize performance degradation on short inputs while maximizing gains on long input summarization tasks.
• To release a high-performing, parameter-efficient model for long input summarization that avoids the need for model parallelism.

Proposed method

• Adopt a block-local attention mechanism with staggered blocks to reduce memory and compute costs while maintaining long-range modeling capability.
• Introduce learned global token embeddings to capture global context, improving performance without increasing sequence length.
• Apply an additional pretraining phase on long sequences (up to 16K tokens) after initializing from a short-sequence pretrained PEGASUS model.
• Use sinusoidal position embeddings instead of T5 relative position biases to reduce computational overhead while maintaining performance.
• Optimize the encoder-decoder layer distribution and attention hyperparameters through ablation studies.
• Fine-tune the resulting model on long input summarization datasets using the same objective as the original PEGASUS model.

Experimental results

Research questions

• RQ1 Which efficient Transformer architecture variants offer the best balance of performance and computational efficiency for long input summarization?
• RQ2 Does additional pretraining on long sequences significantly improve downstream summarization performance compared to fine-tuning alone?
• RQ3 How do architectural choices such as position encoding schemes and global token integration affect performance and efficiency?
• RQ4 Can a small, efficient model achieve state-of-the-art results on long input summarization without requiring model parallelism or large parameter counts?
• RQ5 To what extent does pre-adapting a short-context model to a long-context architecture improve performance compared to direct fine-tuning?

Key findings

• A staggered, block-local Transformer with learned global tokens achieves the best trade-off between performance and efficiency for long input summarization.
• Additional pretraining on long sequences after initializing from a short-sequence pretrained model leads to significantly better downstream performance than either long-only pretraining or no adaptation.
• PEGASUS-X achieves state-of-the-art performance on the GovReport and PubMed long input summarization benchmarks, outperforming models like LongT5 despite having fewer parameters.
• The model maintains minimal performance degradation on short input summarization tasks, indicating strong generalization across input lengths.
• Sinusoidal position embeddings outperform T5 relative position biases in terms of speed with only a slight drop in performance, making them more suitable for efficient inference.
• The proposed recipe—initializing from a short-context model, applying architectural modifications, and adding long-sequence pretraining—proves highly effective and scalable for extending existing models to long contexts.