[论文解读] MixKD: Towards Efficient Distillation of Large-scale Language Models

Kevin J Liang, Weituo Hao+5 more|arXiv (Cornell University)|Nov 1, 2020

Topic Modeling参考文献 52被引用 30

一句话总结

MixKD通过对数据进行混合并在插值样本上向教师查询，提升大型语言模型的知识蒸馏效果，从而在GLUE任务上为紧凑学生提供更强的泛化能力和更好表现。

ABSTRACT

Large-scale language models have recently demonstrated impressive empirical performance. Nevertheless, the improved results are attained at the price of bigger models, more power consumption, and slower inference, which hinder their applicability to low-resource (both memory and computation) platforms. Knowledge distillation (KD) has been demonstrated as an effective framework for compressing such big models. However, large-scale neural network systems are prone to memorize training instances, and thus tend to make inconsistent predictions when the data distribution is altered slightly. Moreover, the student model has few opportunities to request useful information from the teacher model when there is limited task-specific data available. To address these issues, we propose MixKD, a data-agnostic distillation framework that leverages mixup, a simple yet efficient data augmentation approach, to endow the resulting model with stronger generalization ability. Concretely, in addition to the original training examples, the student model is encouraged to mimic the teacher's behavior on the linear interpolation of example pairs as well. We prove from a theoretical perspective that under reasonable conditions MixKD gives rise to a smaller gap between the generalization error and the empirical error. To verify its effectiveness, we conduct experiments on the GLUE benchmark, where MixKD consistently leads to significant gains over the standard KD training, and outperforms several competitive baselines. Experiments under a limited-data setting and ablation studies further demonstrate the advantages of the proposed approach.

研究动机与目标

• 在不牺牲准确性的前提下，通过减小模型规模和推理成本来实现大规模语言模型的高效部署。
• 在任务数据有限时通过扩充训练数据进行增强来解决过拟合与记忆化问题。
• 提出一种与数据无关的蒸馏框架（MixKD），利用混合（mixup）生成用于师生学习的增强样本。
• 为通过增强KD实现的泛化改进提供理论依据。
• 在GLUE任务上展示经验性提升，尤其在数据较少的情况下，并与基线KD及相关方法进行比较。

提出的方法

• 将大型教师模型（BERT）对小型学生模型（BERT-3或BERT-6）进行知识蒸馏。
• 对词嵌入应用混合（mixup）以为学生和教师监督生成增强的输入（x' = λx_i + (1-λ)x_j; y' = λy_i + (1-λ)y_j）。
• 以联合目标进行训练，组合如下：对原始数据的标准交叉熵、对混合数据的学生端交叉熵，以及对混合数据的教师-学生蒸馏损失（L = L_MLE + α_SM L_SM + α_TMKD L_TMKD）。
• 在混合样本上对教师进行查询，并使蒸馏损失对抗于学生在这些样本上的预测最小化（L_TMKD）。
• 给出理论结果，显示在何种条件下增强数据可缩小经验误差与泛化误差之间的差距，从而改善泛化。

实验结果

研究问题

• RQ1 在KD中基于混合的数据增强（MixKD）是否相较于标准KD提高了小型学生模型的泛化能力并缩小泛化差距？
• RQ2 在不同数据可用性下，MixKD在GLUE任务中的表现如何，与DistilBERT、PKD及其他基线相比如何？
• RQ3 在KD中将混合增强与回译和其他数据增强技术结合的影响是什么？
• RQ4 在有限数据情形下，MixKD对超参数和混合比的敏感性如何？

主要发现

• MixKD在GLUE任务上持续优于标准KD及相关基线，在数据有限的设置中尤有显著提升。
• 一个6层的MixKD学生模型，结合SM+TMKD+BT，通常能捕捉到教师性能的很大一部分，同时保持显著更快的推理速度和更少的参数量。
• 在SST-2任务中，结合SMTKD和回译的MixKD相对于从头训练或标准KD，显著缩小了与教师之间的差距。
• 理论结果表明，在几种数据生成情景下，利用混合数据增强可以减小泛化差距，从而支持经验上的提升。
• 嵌入向量可视化表明，MixKD引导学生将插值样本与数据流形对齐，改善潜在空间的组织。