NAACL 2022 | 简单且高效！随机中间层映射指导的知识蒸馏方法

2022-08-10 05:08

©PaperWeekly 原创 · 作者 | werge

研究方向 | 自然语言处理

前言

模型蒸馏旨在将教师模型的知识传递给学生模型。一般而言，蒸馏损失函数设置为两个模型输出层的 MSE 或 KL 散度，通过加入内层知识蒸馏（教师模型的中间层向学生模型中间层传递知识）可以在此基础上进一步提升学生模型的表现。然而，内层知识蒸馏存在 skip and search problem ，不容易找到合适的内层映射方式。

本文提出了 RAIL-KD (RAndom Intermediate Layer Knowledge Distillation) 方法，在每一轮训练时，随机选取教师模型的中间层与学生模型中间层建立映射，进行知识蒸馏。可以有效地减少额外计算开销，并且可以有效提升学生模型的泛化能力。

论文标题：

RAIL-KD: RAndom Intermediate Layer Mapping for Knowledge Distillation

论文链接：

https://arxiv.org/abs/2109.10164

Introduction

预训练的语言模型在许多自然语言理解任务上都表现出了非凡的能力，然而由于模型过大，推理时间过长，在现实世界应用程序中的部署面临巨大挑战。近年来，模型压缩技术在减少模型大小和延迟方面取得了较大的进展，包括模型量化、模型剪枝、知识蒸馏等。

ILD（Intermediate layer distillation）可以在常规蒸馏的 logits 匹配以外，进一步地增强知识转移的能力，从而显著提升学生模型性能。ILD 的主要思路是将教师模型的部分中间层和全部学生模型的中间层进行一一对应，将 layer representation 映射到同一空间，并计算损失，来指导蒸馏过程。

但是 ILD 存在skip and search problem，由于教师模型的层数比学生模型多，所以教师模型的部分中间层会被跳过，但这些层可能是包含了重要信息的层，这就是skip problem；如果想寻找更有用的教师模型中间层，则需要进行大量实验，这会带来额外的搜索开销，这就是search problem。

现有的一些方法都对该问题进行了探究，但大多都是为了解决 skip problem ，而忽略了search problem。本文作者提出了 RAIL-KD 方法，在每个 epoch开始前，从教师模型的个中间层中随机选取个中间层，并分别与学生模型的个中间层进行对应，以指导蒸馏。由于层选择是随机进行的，所以教师模型的所有中间层都有机会被选择，这样不会漏掉重要的中间层，解决了skip problem ，而随机选取并不会增加训练时的计算成本，同时也就解决了search problem。

Method

下表展示了当下各类针对skip and search problem研究的层映射方式、算法复杂度以及局限性。

在 PKD [2] 中，作者使用了固定的逐层映射，例如 PKD-LAST 是指将教师模型最高的层与学生模型的个中间层对应；PKD-SKIP 则是教师模型每两层选取一层与学生模型的一层对应，然而这会导致 skip problem；CKD [3] 则在此基础上，将教师模型的中间层分为组，每一组的 representation 进行简单拼接，与学生模型的一层对应，这解决了 skip 问题，但选择分组需要额外计算，存在 search problem。ALP-KD [1] 直接将教师模型的所有中间层作为一组，与每一个学生层计算加权注意力，来将知识传递给学生，解决了 search 问题。CODIR 方法则使用了对比学习解决 search 问题，但这两种方法都需要额外计算开销。

本文提出的 RAIL-KD 方法则不会在蒸馏过程中增加任何计算成本，同时在表现上优于以前的方法。

设学生模型对于数据的第个中间层的表示为，作者采用了 mean-pooling representation 表示所有个样本的平均表示：

，教师模型的表示相同。分别将和投影到相同的向量空间，就可以计算 layerwise loss 并求和得到最终的损失函数：

本文作者还提出了另一种损失函数计算方式，直接将各层的输出进行拼接，称为了 Concatenated RAIL-KD：

最终训练的损失函数为下游任务损失、输出层 logit 损失函数以及 RAIL-KD 的损失函数的加权和：

两种计算方式的示意图如下所示：

实验

作者分别对不同的教师-学生模型对在 GLUE benchmark 上进行了实验，实验结果如下：

如上表所示，在层蒸馏中，RAIL-KD 在 dev 和 test 上的平均性能均优于当下最优的 ILD 方法，且提出的两种 RAIL-KD 方法的表现非常相似，这表明本文提出的方法对级联蒸馏和逐层蒸馏都是有效的。

同样地，根据上表所示，作者在层的模型压缩实验中也看到了类似的趋势；有趣的是，PKD 方法在测试集上的表现 ALP 和 RAIL-KD。这可能是因为 PKD 跳过了 RoBERTa24 上的大量中间层。

作者还对模型的泛化能力进行了实验：利用在 MLI, QQP 和 SST-2 上 finetune 好的模型，在 SciTail, PAWS 和 IMDB 上进行了测试，如下表所示。作者注意到，与域内结果相比，不同模型之间的表现高低在域外结果有所不同，但 RAIL-KD(l/c) 两种方法仍然明显优于其他所有 baseline。

有趣的是，作者注意到 layerwise RAIL-KD 在 Roberta-24 上的表现均优于 concatenated RAIL-KD，但在 BERT-12 上的表现则全面落后。这些结果表明，在教师模型和学生模型容量（层数）差距很大时，layerwise 蒸馏方法表现优于 concatenated 蒸馏方法，反之亦然。

总结

本文提出了一种新颖、简单且高效的中间层知识蒸馏方法，该方法在性能提升和有效训练时间方面优于传统方法。RAIL-KD 从教师中随机选择与学生模型中间层数量相同的个中间层，与学生模型一一对应，并进行蒸馏。作者认为，该方法可以帮助学生模型更好的接近教师模型的上层表示，且可以在更广泛的 NLU 任务上提高鲁棒性和泛化能力。

参考文献

[1] Peyman Passban, Yimeng Wu, Mehdi Rezagholizadeh, and Qun Liu. 2021. ALP-KD: attention-based layer projection for knowledge distillation. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, pages 13657–13665. AAAI Press

[2]Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019. Patient knowledge distillation for bert model compression. https://arxiv.org/abs/1908.09355.

[3]Yimeng Wu, Peyman Passban, Mehdi Rezagholizadeh, and Qun Liu. 2020a. Why skip if you can combine: A simple knowledge distillation technique for intermediate layers. https://arxiv.org/abs/2010.03034.