NeurIPS 2022 | DIST: 少学点反而更好！商汤开源松弛知识蒸馏方法

科技

2023-03-01 04:03

©Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

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

论文标题：

Knowledge Distillation from A Stronger Teacher

论文链接：

https://arxiv.org/pdf/2205.10536.pdf

代码链接：

https://github.com/hunto/DIST_KD

引言

深度神经网络在计算机视觉领域取得了显著成功，在追求更好性能的道路上，现有的深度学习模型通常变得更深更宽。然而，由于计算和内存资源的限制，这些模型难以实际应用部署。因此，针对深度学习模型的压缩方法得到了广泛的研究，知识蒸馏（Knowledge Distillation, KD）就是其中非常重要的一类方法。

知识蒸馏是指在训练期间从更大模型（教师）中提取知识来提高小模型（学生）性能的方法，其本质在于如何将知识从教师传递给学生。最直观有效的方法是通过 Kullback-Leibler（KL）散度 [1] 来匹配师生的概率预测得分，其他常见的知识蒸馏方法还包括特征蒸馏 [2] 和关系蒸馏 [3]。

随着模型规模和模型性能的不断增加，实验发现，性能更好的模型作为教师，往往并不能带来学生蒸馏性能的提升，甚至比从头开始训练的表现更糟糕。这一现象可被概况为“更强的模型不一定是好教师”。过往的研究主要集中在解决模型规模显著不同时的知识蒸馏问题，常见的解决方案是引入中等规模的模型作为中间过渡，缓解教师和学生在模型大小上的巨大差异。

然而，模型大小的增加只是更强大教师的一种范式，现有的方法缺乏对更先进的训练策略对知识蒸馏影响的深入分析。本文旨在探究和提出一个足够通用的方案来解决“更强的模型不一定是好教师”的问题，既包括更大的模型规模，也包括更强的训练策略。

作者认为将知识从教师传递到学生时，我们真正关心的是教师的概率预测的相对顺序，而不需要准确地模仿其绝对值。本文提出了一种利用皮尔逊（Pearson）相关系数 [4] 替换 KL 散度的知识蒸馏方法。除了学习概率预测的类间关系（inter-class relations），作者还提出可以学习不同实例相对于每个类的类内关系（intra-class relations）。该方法被称为 DIST（Knowledge Distillation from A Stronger Teacher）。

如图 1 所示，不同于过往的方法，学生只是被适当地引导以提取那些真正有用的关系。DIST 的训练成本和原始 KD 几乎一致，在多个任务和数据集上取得了 SOTA 的效果。

▲ 图1. DIST 和现有 KD 方法的不同

重新审视KD的概率匹配

2.1 原始KD介绍

原始 KD [1] 利用最小化教师和学生概率预测得分的差异，将知识从预先训练好的教师模型传递到学生模型。具体公式如下：

过往研究表明与 ground-truth 一同训练有利于提高学生性能，整体训练损失由原始分类损失和蒸馏损失组成：

2.2 更强教师带来的灾难性分歧

通过系统地研究设计和训练深度神经网络的流行策略，作者发现——如图 2 所示，当教师和学生采用更先进的训练策略时，教师的性能会更好，但它和学生之间的差异也会变得相当大。从图中可以看出，当采用更先进的训练策略时（B2），与 ResNet-50 相比，ResNet-18 的输出并未有太大变化。

由此可以推测，当教师过于强大时，利用 KL 散度准确模仿教师的概率预测得分可能十分具有挑战性，这可能是原始 KD 的失败的原因。同时，蒸馏损失和分类损失的不一致也会更严重，这将导致对学生训练的干扰。因此，作者认为可以采用一种更松弛的方式匹配教师和学生的概率预测。

方法：DIST

3.1 对关系的松弛匹配

概率预测得分表示教师对所有类别的置信度（或偏好）。根据上文的实验现象和推测，我们真正关心的是教师预测的相对关系，而不是绝对值。因此，我们可以选择一种的度量，在精确匹配的基础上引入的额外映射和，即：

因此，并不一定要求和完全相同。当和 是保序映射时，我们就可以保证 和 具有相同的预测排序，而不影响预测向量的语义信息和推断结果。

一种常见的保序映射为正线性变换，即：

这种匹配方式具有尺度和平移不变性，可以满足对关系的松弛匹配。为了满足式（5）的性质，作者采用了广泛使用的皮尔逊距离作为度量，即：

通过这种方式，原始 KD 的精确匹配得到松弛，取而代之的是最大线性相关，以保留教师和学生每个实例在概率预测的关系，即类间关系（inter-class），具体损失函数可表示为：

3.2 类内关系的补充

除了类间关系，每个类在多个实例的概率预测的相对关系也蕴含着丰富的知识。该知识反映了教师对不同实例在同一类的相似度判断。例如，标签为“猫”“狗”和“飞机”的三张图像，在“猫”类上的预测得分分别是、、。因为飞机是无生命的，的相对关系蕴含着语义相似度的信息，这种类内关系值得传递。其具体损失函数可表示为：

整体的训练损失函数可由分类损失、类间 KD 损失和类内 KD 损失组成，即：

实验

1. 图像分类：如表 1 所示，DIST 方法明显优于过往的 KD 方法，在相同和不同模型结构间都取得了最好的分类性能。

▲ 表1. ImageNet 上不同蒸馏方法的对比

2. 目标检测：如表 2 所示，通过简单地替换损失函数，DIST 就显著优于原始 KD。此外，在 DIST 基础上加上最简单的 FPN 特征均方误差最小化，就能取得超越目标检测中最先进的 KD 方法。这充分说明了 DIST 在下游任务中的有效性。

参考文献

[1] G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015. 1, 3, 5, 6, 7, 8, 14

[2] B. Heo, J. Kim, S. Yun, H. Park, N. Kwak, and J. Y. Choi. A comprehensive overhaul of feature distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1921–1930, 2019. 1, 6

[3] Y. Tian, D. Krishnan, and P. Isola. Contrastive representation distillation. In International Conference on Learning Representations, 2019. 1, 6, 8, 14

[4] K. Pearson. Vii. mathematical contributions to the theory of evolution.—iii. regression, heredity, and panmixia. Philosophical Transactions of the Royal Society of London. Series A, containing papers of a mathematical or physical character, (187):253–318, 1896. 2

更多阅读