[Paper Review] Pre-trained Summarization Distillation

Sam Shleifer, Alexander M. Rush|arXiv (Cornell University)|Oct 24, 2020

Topic Modeling25 references57 citations

TL;DR

This paper compares shrink-and-fine-tune (SFT), knowledge distillation (KD), and pseudo-labeling (PL) to distill large pre-trained summarization models (BART, Pegasus) into smaller, faster students, showing SFT often wins on CNN/DailyMail while PL wins or ties on XSUM, with varying tradeoffs.

ABSTRACT

Recent state-of-the-art approaches to summarization utilize large pre-trained Transformer models. Distilling these models to smaller student models has become critically important for practical use; however there are many different distillation methods proposed by the NLP literature. Recent work on distilling BERT for classification and regression tasks shows strong performance using direct knowledge distillation. Alternatively, machine translation practitioners distill using pseudo-labeling, where a small model is trained on the translations of a larger model. A third, simpler approach is to 'shrink and fine-tune' (SFT), which avoids any explicit distillation by copying parameters to a smaller student model and then fine-tuning. We compare these three approaches for distillation of Pegasus and BART, the current and former state of the art, pre-trained summarization models, and find that SFT outperforms knowledge distillation and pseudo-labeling on the CNN/DailyMail dataset, but under-performs pseudo-labeling on the more abstractive XSUM dataset. PyTorch Code and checkpoints of different sizes are available through Hugging Face transformers here http://tiny.cc/4iy0tz.

Motivation & Objective

• Motivate compressing large pre-trained summarization models for practical use.
• Systematically compare three distillation families (SFT, KD, PL) on standard datasets CNN/Daily Mail and XSUM.
• Provide implementation details and empirical guidance for selecting distillation strategy based on task and model size.
• Quantify trade-offs between quality (ROUGE) and speed (inference and training time).

Proposed method

• Define three distillation paradigms: Shrink and Fine-Tune (SFT), Pseudo-Labeling (PL), Direct Knowledge Distillation (KD).
• Initialize student models by copying maximally spaced layers from the teacher and fine-tuning on the task data.
• For KD, minimize a weighted combination of logits, data, and hidden-state losses (L_KD = alpha_logits L_logits + alpha_data L_data + alpha_hidden L_hidden).
• Use pseudo-labels by generating teacher summaries with beam search and training the student on these pseudo-targets (L_pseudo).
• Evaluate on CNN/Daily Mail (less abstractive) and XSUM (more abstractive) and report ROUGE-1/2/L across multiple teacher-student pairs.
• Report training-time costs and inference-speed gains for each distillation method.

Experimental results

Research questions

• RQ1 Which distillation method (SFT, KD, PL) yields best trade-off between speed and summarization quality for BART and Pegasus?
• RQ2 Does initialization by copying maximally spaced layers from the teacher enable effective compression without pre-training?
• RQ3 How do distillation methods perform differently on CNN/Daily Mail versus XSUM (abstractive) summarization tasks?
• RQ4 What is the impact of pseudo-label quality and cross-task transfer (translation benchmarks) on distillation performance?

Key findings

• On CNN, SFT outperforms the more expensive KD and PL methods across teachers (BART, Pegasus).
• On XSUM, KD (for BART) and PL (for Pegasus) can outperform SFT, with PL closest to teacher performance in some settings.
• Some distilled models achieve substantial speedups (e.g., 43-93% faster) with only small ROUGE losses compared to teachers, depending on dataset and model pair.
• Pseudo-labeling can significantly improve performance when the pseudo-labels are high quality and combined with original fine-tuning data (Orig+PL, Orig+PL+PL).
• KD often offers limited gains over SFT for XSUM and is computationally more expensive due to requiring teacher forward passes; PL can be a strong alternative, particularly for Pegasus.