ICLR 2023 | HomoDistil：蒸馏和剪枝在知识传递上的有机结合

©Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

研究方向 | 计算机视觉、模型压缩

论文标题：

HomoDistil: Homotopic Task-Agnostic Distillation of Pre-trained Transformers

论文链接：

https://arxiv.org/pdf/2302.09632.pdf

动机&背景

随着预训练大模型规模的不断增加，任务特定蒸馏（下游微调蒸馏）的成本越来越高，任务无关蒸馏变得越来越重要。然而，由于教师模型的模型容量和表示能力远超学生模型，因此学生很难在大量开放域训练数据上模仿教师的预测。本文提出了同源蒸馏（Homotopic Distillation, HomoDistil）来缓解这一问题，该方法充分利用了蒸馏和剪枝的优势，将两者有机结合在了一起。

具体来说，本文用教师模型初始化学生模型，以缓解两者在蒸馏过程中的容量和能力差异，并通过基于蒸馏损失的重要性得分的迭代剪枝，来逐步将学生模型修剪至最终想要的目标结构。在整个蒸馏+剪枝的过程中，教师和学生一直保持着较小的预测差异，这有助于知识更有效的传递。其核心动机如图 1 所示。

▲ 图1. HomoDistil 动机说明（用剪枝给知识蒸馏做初始化，并迭代式地获得最终的学生模型结构）

HomoDistil：同源任务无关蒸馏

如图 2 所示，本文所提出的 HomoDistil 先用教师模型初始化学生，并以类似 TinyBERT [1] 的蒸馏损失函数作为修剪的目标函数，在每次迭代中，根据重要性得分从学生中删除最不重要的神经元并用蒸馏损失指导学生的训练。在整个训练过程中不断重复此过程，直至学生达到目标规模。该方法可从「蒸馏损失函数」和「迭代剪枝细节」两部分进行介绍。

▲ 图2. HomoDistil 方法的示意图，矩形的宽度表示层的宽度，颜色的深度反映训练的充分性。

2.1 蒸馏损失函数