[论文解读] Pre-trained Summarization Distillation

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

Topic Modeling参考文献 25被引用 57

一句话总结

本文将缩减并微调（SFT）、知识蒸馏（KD）和伪标签（PL）用于把大型预训练摘要模型（BART、Pegasus）蒸馏成更小、更快的学生模型，并显示在 CNN/DailyMail 上 SFT 往往获胜，而在 XSUM 上 PL 获胜或并列，权衡取舍各不相同。

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.

研究动机与目标

• 鼓励将大型预训练摘要模型压缩以便实际使用。
• 系统比较三大蒸馏范式（SFT、KD、PL）在标准数据集 CNN/Daily Mail 与 XSUM 上的表现。
• 提供实现细节和基于任务与模型规模的经验性指南，帮助选择蒸馏策略。
• 量化质量（ROUGE）与速度（推理与训练时间）之间的权衡。

提出的方法

• 定义三种蒸馏范式：缩减并微调（SFT）、伪标签（PL）、直接知识蒸馏（KD）。
• 通过从教师模型中拷贝最大间距的层并在任务数据上进行微调来初始化学生模型。
• 对于 KD，通过最小化 logits、数据和隐藏状态损失的加权组合来实现（L_KD = alpha_logits L_logits + alpha_data L_data + alpha_hidden L_hidden）。
• 通过使用束搜索生成教师摘要并在这些伪目标上训练学生来使用伪标签（L_pseudo）。
• 在 CNN/Daily Mail（较少抽象）和 XSUM（较为抽象）上进行评估，并报告多组教师-学生对的 ROUGE-1/2/L。
• 报告每种蒸馏方法的训练时间成本和推理速度提升。

实验结果

研究问题

• RQ1 哪种蒸馏方法（SFT、KD、PL）在 BART 与 Pegasus 的速度与摘要质量之间提供最佳折衷？
• RQ2 通过拷贝教师中的最大间距层进行初始化是否能够在不进行预训练的情况下实现有效压缩？
• RQ3 在 CNN/Daily Mail 与 XSUM（抽象性不同）上蒸馏方法的表现有何差异？
• RQ4 伪标签质量和跨任务迁移（翻译基准）对蒸馏性能有何影响？

主要发现

• 在 CNN 上，SFT 在所有教师（BART、Pegasus）中都优于更昂贵的 KD 和 PL 方法。
• 在 XSUM 上，KD（对 BART 有利）和 PL（对 Pegasus 有利）可以超越 SFT，在某些设置中 PL 最接近教师性能。
• 某些蒸馏模型可实现显著的加速（例如 43-93% 的提速），相对于教师仅有较小的 ROUGE 损失，且取决于数据集与模型对。
• 当伪标签质量高且与原始微调数据结合时（Orig+PL、Orig+PL+PL），伪标签可以显著提升性能。
• KD 在 XSUM 上通常对 SFT 的提升有限且由于需要教师前向传播而更具计算成本；PL 可以成为强有力的替代，尤其是对 Pegasus。