ACL 2022 | MetaDistil：基于元学习的模型蒸馏

2022-06-15 06:06

©PaperWeekly 原创 · 作者 | BNDSBilly

研究方向 | 自然语言处理

Abstract

模型蒸馏是当下较为主流的模型压缩方法，通过将大规模教师模型的知识传递给小规模学生模型，从而既能降低模型大小，提升推理速度，又能得到与原模型基本相当的表现。之前分享的几篇论文大多是对蒸馏目标函数及学生模型架构的探究，但本文的作者转而对教师模型的知识传达能力进行了探究。

最终，作者提出了基于 Meta Learning 的 MetaDistil 方法，同时考虑了提升教师模型的知识传达能力和学生模型的知识获取能力。实验表明，MetaDistil 压缩的模型在多个 benchmark 上都优于传统蒸馏方法，同时对学生模型的参数敏感性更低，可以更加灵活的应用在不同的下游任务和模型上。

论文标题：

BERT Learns to Teach: Knowledge Distillation with Meta Learning

论文链接：

https://arxiv.org/abs/2106.04570

代码链接：

https://github.com/JetRunner/MetaDistil

Methods

2.1 MetaDistil

在传统模型蒸馏过程中，一般首先训练一个大模型作为教师模型，然后训练一个小的学生模型来模拟教师模型的行为，以便获取教师模型的知识。在整个过程中，教师模型是固定的。但是这一范式有两个缺点：首先，教师模型不知道学生模型的能力。有教育学研究表明，以学生为中心的学习（考虑到学生的特点和学习能力）可以显著提升学生表现。然而，传统的知识蒸馏无法考虑到学生模型的学习能力和表现。其次，教师模型并没有针对蒸馏而最优化。一般而言，选取的教师模型都最优化以提升自的推理性能，然而并不一定是传递知识的最优化状态。打个比方，一个博士有足够的知识来自己解决问题，但需要额外的教学培训才能获得教授资格。

针对这两个问题，作者提出了 MetaDistil 蒸馏方法，这是基于 meta learning [1] 的新蒸馏范式：教师模型将知识蒸馏给学生，并通过学生模型的反馈进一步提升自己的知识传达能力。MetaDistil 的工作流如下图所示：

2.2 Pilot Update

meta learning 的核心思想是 learning to learn，也即同时考虑优化学习算法过程和算法本身。其通常涉及一个双层优化过程，外层称为元学习器，对应蒸馏过程中的教师模型；内层称为内学习器，对应学生模型。最近的一些研究使用了这种双层优化框架，例如 meta pseudo labels 工作 [2] 使用元学习来优化伪标签生成器，以实现更好的半监督学习；meta back-translation 工作 [3] 使用该框架训练反向翻译模型，从而间接地训练机器翻译模型。在每轮迭代中，内学习器通过蒸馏从元学习器学习知识并更新，然后元学习器再根据其表现进行更新。

但是，该方法的主要思路是让内学习器更加适配于元学习器，从而获得一个最优化的元学习器；但在本文中，目标是让内学习器（学生模型）达到最优，而让元学习器（教师模型）来适配内学习器的表现。

所以，本文作者提出了 pilot update 方法来解决这一问题：首先仍然按照 meta learning 的方式进行每一轮的内学习器和元学习器的更新，只不过内学习器的更新要在元学习器更新后撤销，再根据更新的元学习器进行本轮的更新，这样可以使两个学习器的更新进行同步和匹配。在实现撤销这一步时，为了方便起见，作者复制了一个内学习器的副本，在 meta learning 结束后，直接将该副本删除即可。

对于学生模型，其损失函数包括在下游任务上的损失以及蒸馏损失：

而对于教师模型，作者希望其能根据学生模型的表现反馈进行调整，而学生模型在下游任务上的表现也包含教师模型参数，所以作者采用下游任务损失作为教师模型的损失函数（）。特别地，为了防止过拟合，作者从训练集中分离出了一部分样本组成 quiz set，将学生模型在该数据集上的表现计算损失函数。

最终的算法伪代码如下：

Experiments

本文实验选取了 GLUE benchmark，将 BERT-base 压缩成了层，隐藏层维度为的学生模型。特别地，本文采用的是 fine-tune 阶段的蒸馏，其中蒸馏损失函数采用了模型 logits 输出的 MSE loss。

作为对比实验，作者选取了深度互学习 DML（Zhang et al., 2018）、教师助理知识蒸馏 TAKD（Mirzadeh et al., 2020）、路线约束优化 RCO（Jin et al., 2020）al., 2019）、邻近知识教学 ProKT（Shi et al., 2021）、学生友好型教师网络 SFTN（Park et al., 2021）、DistilBERT（Sanh et al., 2019）、TinyBERT（Jiao et al., 2019）、MiniLM v1 和 v2（Wang et al., 2020b,a）方法。为了确保公平，作者均根据各类方法压缩得到了层学生模型，并在 baseline 上进行了实验。结果如下表所示：

由表可知，MetaDistil 在大部分下游任务上均优于蒸馏 baseline。说明在元学习的帮助下，MetaDistil 能够直接优化教师的教学能力，从而进一步提高学生的准确性。此外，作者进行了消融实验，发现在不使用 pilot update 的情况下，模型性能下降，说明了该更新方法的有效性。此外，MetaDistil 在图像分类方面也取得了非常有竞争力的结果。

为了进一步证实 MetaDistil 的有效性，作者对 pilot update 算法中第两行更新后学生模型的 validation loss 进行了对比。对比发现，后者的 loss 更小，这说使用了pilot update方法更新的学生模型表现比不使用要更好。除此之外，作者还观察到，在算法第行分别更新教师模型和学生模型后，教师模型的 logits 与学生模型的相似度下降，作者推断这是由于教师模型需要强化自己的能力，以便于下一步更加好的将知识传递给学生模型。

最后，作者还探究了 MetaDistil 框架对于模型超参数的敏感性。在实验中，作者利用 MetaDistil 对 BERT-base 进行了蒸馏，得到了不同参数量、不同学习率、以及不同蒸馏温度的学生模型，并对比了其在 MNLI、SST-2 和 MRPC 数据集上的性能，结果如下所示。实验结果表明，MetaDistil 始终优于传统的蒸馏方法，且对不同的学习率和温度均不太敏感，显示出了更好的鲁棒性。

Conclusion

本文中，作者提出了 MetaDistil 框架，这是一种基于 meta learning 的知识蒸馏算法，显式地优化教师模型，以将其知识更好的传达给学生模型。大量实验表明，MetaDistil 相比于传统蒸馏算法表现更好，且具有更好的鲁棒性。

参考文献

[1] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh (eds.), ICML, 2017.

[2] Hieu Pham, Zihang Dai, Qizhe Xie, Minh-Thang Luong, and Quoc V Le. Meta pseudo labels. arXiv preprint arXiv:2003.10580, 2020.

[3] Hieu Pham, Xinyi Wang, Yiming Yang, and Graham Neubig. Meta back-translation. arXiv preprint arXiv:2102.07847, 2021.

更多阅读