[论文解读] LightPAFF: A Two-Stage Distillation Framework for Pre-training and Fine-tuning

Kaitao Song, Hao Sun+5 more|arXiv (Cornell University)|Apr 27, 2020

Topic Modeling参考文献 17被引用 9

一句话总结

LightPAFF 是一种两阶段知识蒸馏框架，在预训练和微调阶段均将大型预训练模型（如 BERT、GPT-2、MASS）的知识迁移至更小的学生模型，实现与原始大模型相当的性能，同时模型尺寸缩小至原来的 1/5，推理速度提升 5–7 倍，从而实现在资源受限系统中的高效部署。

ABSTRACT

While pre-training and fine-tuning, e.g., BERT~\citep{devlin2018bert}, GPT-2~\citep{radford2019language}, have achieved great success in language understanding and generation tasks, the pre-trained models are usually too big for online deployment in terms of both memory cost and inference speed, which hinders them from practical online usage. In this paper, we propose LightPAFF, a Lightweight Pre-training And Fine-tuning Framework that leverages two-stage knowledge distillation to transfer knowledge from a big teacher model to a lightweight student model in both pre-training and fine-tuning stages. In this way the lightweight model can achieve similar accuracy as the big teacher model, but with much fewer parameters and thus faster online inference speed. LightPAFF can support different pre-training methods (such as BERT, GPT-2 and MASS~\citep{song2019mass}) and be applied to many downstream tasks. Experiments on three language understanding tasks, three language modeling tasks and three sequence to sequence generation tasks demonstrate that while achieving similar accuracy with the big BERT, GPT-2 and MASS models, LightPAFF reduces the model size by nearly 5x and improves online inference speed by 5x-7x.

研究动机与目标

• 解决在在线或移动环境中部署大型预训练模型（如 BERT、GPT-2）时因高内存占用和延迟成本带来的挑战。
• 克服先前知识蒸馏方法仅在微调阶段压缩模型的局限性，避免导致显著的性能下降。
• 开发一种统一框架，适用于多种预训练目标——掩码语言建模（BERT）、因果语言建模（GPT-2）以及掩码序列到序列建模（MASS），以确保广泛适用性。
• 通过利用预训练和微调后的教师模型的知识，保留轻量化学生模型的一般语言理解与生成能力。
• 在多个 NLP 任务（包括分类、语言建模和序列到序列生成）中实现高效、低延迟的推理，且不牺牲性能。

提出的方法

• 在两个不同阶段应用知识蒸馏：第一阶段在预训练期间使用预训练数据集，第二阶段在微调期间使用特定任务数据。
• 使用参数量少于教师模型的学生模型，训练其在两个阶段均模仿教师模型的输出 logits 和隐藏表示。
• 引入超参数 λ 以平衡知识蒸馏损失与来自真实标签的标准交叉熵损失。
• 通过最小化学生模型与教师模型参数之间的 L2 距离实现逐参数知识蒸馏，并添加高斯噪声以提升鲁棒性。
• 统一 formulate 蒸馏目标，使其适用于 BERT、GPT-2 和 MASS，从而在不同预训练目标和模型架构间保持一致的应用效果。
• 采用两阶段流程：(1) 预训练教师模型，(2) 在下游任务上微调教师模型，(3) 将预训练的教师模型蒸馏以生成轻量化预训练学生模型，(4) 在下游任务上使用蒸馏方法微调学生模型。

实验结果

研究问题

• RQ1 在预训练和微调两个阶段均应用知识蒸馏，是否能显著减小模型尺寸并降低推理延迟，同时保持性能？
• RQ2 与仅在微调阶段进行蒸馏的单阶段方法相比，两阶段蒸馏在轻量化模型的准确率和鲁棒性方面表现如何？
• RQ3 预训练任务的难度（如掩码语言建模与因果语言建模）是否会影响最优蒸馏超参数设置及模型性能？
• RQ4 所提出的 LightPAFF 框架能否在具有不同预训练目标的多种预训练模型（如 BERT、GPT-2、MASS）上实现统一应用？
• RQ5 与仅在微调阶段进行蒸馏相比，两个阶段均进行蒸馏在泛化能力和鲁棒性方面提升程度如何？

主要发现

• LightPAFF 将 BERT、GPT-2 和 MASS 的模型尺寸缩小近 5 倍，同时保持与原始大模型相当的性能。
• 该框架在所有评估模型上将在线推理速度提升 5 倍至 7 倍，支持在边缘设备上实际部署。
• 在 9 项下游任务上的实验——每类 3 项，分别对应语言理解（SST-2、QQP、多音字消歧）、语言建模（WikiText-2、PTB、WikiText-103）以及序列到序列生成（WMT Zh-En、En-De、En-Fr）——表明压缩后的模型在性能上保持一致提升。
• 两阶段蒸馏方法展现出更好的泛化能力，表现为对参数扰动的更强鲁棒性，以及在预训练和微调阶段性能的稳定性。
• 最优蒸馏超参数 λ 因模型而异：BERT 为 λ = 1.0，MASS 为 λ = 0.7，GPT-2 为 λ = 0.4，反映出不同预训练任务难度和信息保留能力的差异。
• 预训练任务难度与性能呈正相关：BERT（平均 85% 输入 token 可见）最简单，其次是 MASS（75%），GPT-2（50%）最难，这与掩码预测任务中观察到的准确率下降趋势一致。