PaperlayerView exact revision on arXiv

Source-linked AI summary

Pre-trained Summarization Distillation

Sam Shleifer, Alexander M. Rush

arXiv:2010.13002v2cs.CLcs.AI

TL;DR

As pre-trained Transformers grow, researchers seek smaller, faster models that retain performance, but it is unclear which distillation approach suits pre-trained Seq2Seq summarizers. The paper compares KD, pseudo-labeling, and shrink-and-fine-tune for BART and Pegasus. SFT works best on CNN, while XSUM results vary by teacher: KD for BART and pseudo-labeling for Pegasus.

  • Problem

    As pre-trained Transformers grow, it is unclear whether direct knowledge distillation or pseudo-labeling is the better approach for pre-trained Seq2Seq summarization models.

    Source: §1 · p. 1

  • Method

    The paper compares knowledge distillation, pseudo-labeling, and shrink-and-fine-tune, which copies maximally spaced layers from a fine-tuned teacher before continuing fine-tuning.

    Source: §1 · p. 1

  • Results

    No approach dominates across datasets: SFT works best for both teachers on CNN, while KD works best for BART and pseudo-labeling for Pegasus on XSUM.

    Source: §1 · p. 1, §5 · p. 5

  • Takeaways & Limitations

    Removing carefully chosen decoder layers and continuing fine-tuning can generate high-quality student models quickly, while costlier methods can add quality in some situations.

    Source: §7 · p. 8

    Experiments were uneven across models and datasets, with fewer Pegasus runs; the authors note that 60 more Pegasus XSUM experiments might find a better approach.

    Source: §4 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

As pre-trained Transformers grow, researchers seek smaller, faster models that retain performance, but it is unclear which distillation approach best suits pre-trained Seq2Seq summarizers. The paper compares three approaches and finds their relative performance differs between CNN and XSUM.

1 INTRODUCTION

  • As pre-trained Transformers grow, compression aims to produce smaller, faster versions that retain strong performance.Source: §1 · p. 1
  • BART achieved state-of-the-art summarization results on XSUM and CNN, and Pegasus later improved performance with summarization-specific pre-training and a 25% larger model.Source: §1 · p. 1
  • SFT copies maximally spaced layers from a fine-tuned teacher into a student, then reruns the original fine-tuning procedure.The authors hypothesize that removing full layers has minimal impact because Transformer layers use residual connections.Source: §1 · p. 1
  • SFT works best for both teachers on CNN, while XSUM favors KD for BART and pseudo-labeling for Pegasus.On XSUM, no technique matches the Pegasus teacher, though pseudo-labeling comes closest.Source: §1 · p. 1
  • The paper details implementation, compares quality and speed across settings, and examines pseudo-labeling and KD extensions for XSUM.Source: §1 · p. 1

2 RELATED WORK

Related work covers knowledge distillation for pre-trained models, pseudo-labeling for sequence generation, and methods that extract smaller models from trained teachers. The paper positions its approaches by how they initialize students and transfer teacher knowledge.

2 RELATED WORK

  • Knowledge distillation transfers a teacher’s logits, allowing a language-model student to learn a distribution over possible next words rather than only the training label.Source: §2 · pp. 1–2
  • The comparison distinguishes methods by weight initialization, student pre-training, and transfer through logits, hidden states, or generations.SFT is described as similar to BERT-of-Theseus, KD to TinyBERT, and pseudo-labeling to sequence-level knowledge distillation.Source: Table 1 · p. 2, §2 · p. 3
  • BERT distillation methods target GLUE tasks, while sequence-to-sequence machine-translation distillation primarily uses pseudo-labels generated by teacher beam search.Pseudo-labeling trains a smaller model on a regenerated dataset of source documents and teacher-generated outputs.Source: §2 · p. 2
  • Sequence-level knowledge distillation trains on teacher-generated translations or summaries, unlike word-level transfer through logits.Source: §2 · p. 3
  • Liu et al. (2020a) fine-tune summarization models on their own noisy logits, a method that could be applied before or after the approaches discussed here.Source: §2 · p. 3

3 BACKGROUND AND METHODS

The section establishes Seq2Seq summarization fine-tuning, then specifies three teacher-to-student compression strategies: shrink-and-fine-tune, pseudo-labeling, and direct knowledge distillation. Students learn respectively from ground-truth summaries, teacher generations, or teacher distributions and hidden states.

3 BACKGROUND AND METHODS

  • Seq2Seq encoders map source documents to contextual embeddings, while decoders predict the next word from those embeddings and previously decoded words.BART pre-trains by reconstructing corrupted documents; Pegasus pre-trains to generate important sentences extracted from documents.Source: §3 · p. 3
  • Fine-tuning uses document-summary pairs and standard cross-entropy loss, and the distillation experiments begin with a teacher fine-tuned this way.The loss trains the model to predict the correct target word.Source: §3 · p. 3

3.1 DISTILLATION

  • SFT copies maximally spaced decoder layers from the teacher, then fine-tunes the smaller student on the summarization data using the data loss.For a one-layer student, the authors copy layer 0 and report that this works better than copying layer 11; the same initialization is used for the other methods.Source: §3.1 · p. 3
  • Pseudo-labeling replaces ground-truth summaries with the teacher’s beam-search generations and fine-tunes the student only on those generated targets.Generating pseudo-labels can dominate the cost: 16.5 of 19 GPU hours in the described setup.Source: §3.1 · pp. 3–4
  • Direct KD trains the student to match the teacher’s next-word distribution with KL divergence and adds hidden-state matching because corresponding layers have the same shape.The student layers are paired with teacher layers to emulate their outputs; a three-layer BART student, for example, copies teacher layers 0, 6, and 11 and pairs them with teacher layers 3, 7, and 11.Source: §3.1 · pp. 3–4
  • 16.5 of 19 GPU hours for a CNN BART student go to pseudo-label generation; after initialization, SFT is much cheaper because it does not use the teacher.Table 2 estimates the costs for a BART teacher and a 12-encoder-layer, 6-decoder-layer student on a Titan RTX 2080 GPU.Source: §3.1 · p. 4, Table 2 · p. 4

4 EXPERIMENTAL SETUP

The experiments define summarization and translation datasets, evaluation and cost measures, and implementation choices. They also report that experimental effort varied across model–dataset combinations.

4 EXPERIMENTAL SETUP

  • CNN summaries are roughly three sentences and resemble article openings, whereas XSUM summaries are single, more abstractive BBC-news sentences.The XSUM target sentence is removed from the article.Source: §4 · p. 4
  • Evaluation uses beam search and ROUGE; inference speed is measured in summaries per second with batch size 32, using mixed precision for BART and full precision for Pegasus.Mixed precision overflows for Pegasus.Source: §4 · p. 4
  • The comparison also includes English–Romanian WMT 2016 translation, evaluated with BLEU without post-processing, using mBART and Marian teachers.mBART is pretrained multilingually and fine-tuned on bilingual data; Marian is trained from scratch on bilingual data.Source: §4 · p. 4
  • Training stops at epoch five or after four consecutive evaluations without validation improvement; training cost includes GPU training and pseudo-label generation when applicable.The stated training-cost setup uses one Nvidia RTX 2080 GPU.Source: §4 · p. 4
  • When using a fully copied encoder, the experiments freeze its parameters and positional and token embeddings; this made fine-tuning 5.6 times faster without observed performance impact in initial experiments.The encoder-freezing comparison is qualified as an initial observation.Source: §4 · p. 4
  • The TinyBERT-inspired approach omits per-layer weights and learning rates, embedding loss, and attention loss; KD can reuse a shared encoder pass and stop backpropagation at the encoder.The paper also points to a complete list of its φ mappings.Source: §4 · p. 4
  • Table 4 counts teacher-to-student attempts lasting at least 15 minutes, but effort was unequal, with fewer CNN and Pegasus experiments.The authors note that 60 additional Pegasus-on-XSUM experiments might have found a better-performing approach.Source: Table 4 · p. 5, §4 · p. 5
  • Students use encoder-layer/decoder-layer notation; Table 5 reports ROUGE-2 for summarization and BLEU for translation, with end-to-end GPU-hour cost.Pseudo-labeling cost includes beam search over the full training set.Source: §4 · p. 5, Table 5 · p. 5

5 RESULTS

No distillation approach dominates across the evaluated tasks: the best method varies by dataset and teacher. Among the best students by size, several match or exceed teacher scores with substantial speedups, while one BART/XSUM result uses Pegasus-generated pseudo-labels.

5 RESULTS

  • No method dominates: SFT leads on CNN for both teachers, KD on BART/XSUM, PL on Pegasus/XSUM and English–Romanian translation.ROUGE-1 and ROUGE-L follow a similar pattern to ROUGE-2 on the summarization datasets.Source: §5 · p. 5
  • In 3 of 4 contexts, distillation brings relatively minor performance losses and significant speedups; BART/XSUM students beat the teacher at 93% and 43% speedups.The BART 12-3 and 12-6 students correspond to the 93% and 43% speedups, respectively; the CNN BART 12-6 student also outperforms its teacher, while the XSUM Pegasus student falls more than one ROUGE-2 point below its teacher.Source: §5 · p. 5
  • The stronger BART/XSUM 12-3 result in Table 6 used Pegasus-generated pseudo-labels, unlike Table 5’s PL result, which uses labels from the student’s teacher.The authors discuss this result as an extension of pseudo-labeling.Source: §5 · p. 5

6 ANALYSIS

The analyses show that initialization and pseudo-labeling affect student performance differently across tasks, while most tested KD changes do not improve on SFT. Inference time is more sensitive to decoder depth than to encoder depth, width, or feed-forward dimensions.

6.1 HOW DOES INITIALIZATION IMPACT DISTILLATION?

  • Copying decoder layers 0, 6, and 11 from the fine-tuned XSUM BART teacher gave dBART-12-3 faster convergence and better validation loss than other tested initializations.The comparison uses the same frozen encoder across initialization strategies.Source: §6.1 · p. 5

6.2 WHEN DOES PSEUDO-LABELING HELP PERFORMANCE?

  • For EN-RO translation, pseudo-labels can replace original training data, whereas CNN pseudo-labeling performs worse than SFT.Source: §6.2 · pp. 5–6
  • For Pegasus/XSUM, adding BART pseudo-labels worsened the result from 1.2 to 1.6 ROUGE-2 below the teacher; label quality may explain the pattern, as Pegasus pseudo-labels scored 4 points higher.The authors treat pseudo-label ROUGE-2 against training labels as a proxy for quality.Source: §6.2 · p. 6

6.3 DO CHANGES TO Lkd IMPROVE PERFORMANCE?

  • Except for BART on XSUM, KD did not improve over SFT and was always more expensive; removing hidden-state loss hurt only BART/XSUM, while added attention or cosine losses did not improve results.Source: §6.3 · p. 7

6.4 INFERENCE TIME ANALYSIS

  • Adding one decoder layer slowed inference by 8%, versus 4% for an encoder layer; changing width or feed-forward dimensions had negligible runtime impact.The decoder-layer effect is amplified during beam search because the decoder runs once per beam.Source: §6.4 · pp. 7–8

6.2 WHEN DOES PSEUDO-LABELING HELP PERFORMANCE?

  • 1.2 ROUGE-2 points below the teacher for BART/XSUM SFT became 0.8 points below with original plus teacher pseudo-labels and 0.1 points above with added Pegasus pseudo-labels.Table 10 reports these results for combinations of original data and pseudo-labels.Source: §6.2 · p. 6, Table 10 · p. 7

7 CONCLUSION

Removing carefully chosen decoder layers and continuing fine-tuning can produce high-quality student summarizers quickly. The paper also identifies further evaluation and CPU-efficiency questions for future work.

7 CONCLUSION

  • Removing carefully chosen decoder layers from a Seq2Seq transformer and continuing fine-tuning generates high-quality student models quickly.Source: §7 · p. 8
  • More expensive techniques with the same initialization can provide additional quality improvements in some situations.Source: §7 · p. 8
  • The experiments target GPU speedups; future work could evaluate other datasets, tasks, and teachers, while CPU inference may require reducing student-layer width.Source: §7 · p. 8
Loading 2010.13002v2…