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

©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.

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·