CVPR 2023 | DisWOT：给老师找个好学生？KD和NAS的高效组合

2023-04-18 05:04

©Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

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

动机&背景

知识蒸馏（Knowledge Distillation, KD）可以用大模型指导小模型学习，但大小模型之间的容量差距阻碍了知识的传递，过去的研究集中在训练范式（早停[1]）和架构调整（助教[2]）等，他们要么性能提升有限，要么训练成本过高。本文从无需训练的神经网络架构搜索（NAS）的角度出发，为教师寻找最佳的学生架构，从而缓解模型容量差距。无需训练的 NAS 也被称为 Zero-cost proxy，这一表述会在后文出现。

论文标题：

DisWOT: Student Architecture Search for Distillation WithOut Training

论文链接：

https://arxiv.org/pdf/2303.15678.pdf

代码链接：

https://github.com/lilujunai/DisWOT-CVPR2023

方法：DisWOT

DisWOT 的目标是在给定的教师模型和约束条件下，通过无需训练的 NAS 方法找到最适合教师的学生模型，在此学生模型上，用 KD 范式以更高效地传递知识，提升学生模型的性能。

2.1 搜索最优学生网络

无需训练的 NAS 的三大关键要素是「度量指标」、「搜索空间」和「搜索策略」，图 1 是该环节的示意图，图1（a）反映了本文的两大度量指标，图1（b）反映了本文所选择的搜索空间和搜索策略，我们将在这一小节对具体细节依次进行介绍。

▲ 图1. 基于 DisWOT 的 NAS 示意图

语义相似度度量。受 [3] 的结论启发，KD 范式中教师网络往往具有更大的感受野和更丰富的语义特征，语义信息上的差距会阻碍知识的有效传递。反过来讲，最优的学生模型应该与教师的语义信息差距尽可能小。与此同时，[4] 的研究结论表明随机初始化的模型也具备一定的语义定位能力，因此本文采用 Gram-CAM [5] 生成的定位热图来反映度量语义信息。

设教师和学生模型全局平均池化（GAP）层之前的特征图分别为和，，其中第个通道的空间图可以表示为和，Grad-CAM 的具体计算公式为：

其中和是模型最后一个全连接层的权重。最终的语义相似度度量可定义为教师和学生的累积 Grad-CAM 图的相互关系，即：

关系相似度度量。类似 RKD [6] 的思想，输入样本之间的关系对于知识蒸馏也很重要，最优学生模型应该与教师的关系相似度尽可能大。小批量样本的相关矩阵公式如下：

其中，是的变形。基于式（3），关系相似度度量可定义为：

搜索空间。参考过往的基于 ResNet 的工作，本文选择了三种不同的搜索空间。待搜索参数包括块重复次数、连接关系、通道数和内核大小等。不同的搜索空间有着不同的空间大小。

无需训练的进化搜索。适应度可定义为，具体算法见图 2。

该搜索算法先初始化一组满足约束条件的学生模型架构作为初始种群，并计算他们对应的适应度，再后续的每轮迭代中，算法先从种群中随机挑选部分网络，从中选出适应度前 top-k 的网络，再从这 k 个网络中随机选出 1 个网络作为父网络，进行变异，变异后的网络如果符合约束条件，则加入种群，并计算适应度，每次有新模型加入种群后都将从种群中移除一个适应度最小的模型架构。

在多轮迭代中，整个种群的平均适应度不断增加，最终选择适应度最大的模型架构作为后续 KD 时的学生模型。

▲ 图2. 基于 DisWOT 的进化搜索算法伪代码

2.2 高阶知识蒸馏

为了证明搜索得到的学生模型架构的优越性，本文最初采用了原始的蒸馏损失（例如 KL 散度）。但不难想象，NAS 阶段的度量指标（语义相似度和关系相似度）也可以作为 KD 阶段知识转移的优化目标，这将进一步提升学生模型的蒸馏性能，即：

实验

1. 图像分类 - CIFAR100：如表 1 所示，DisWOT 在不同 KD 方法上都取得了超过其他 Zero-cost proxy 方法的性能表现，这说明了本文所提出的基于 DisWOT 的搜索方法的有效性。

▲ 表1. 在 1M 参数量约束下，不同 KD 方法和不同 Zero-cost proxies 方法的性能对比

2. 图像分类 - ImageNet：如表 2 所示，DisWOT 在不同教师和不同蒸馏策略下都取得了超过对照组的性能表现。

▲ 表2. ImageNet 下不同教师变体不同蒸馏方法的性能对比

3. NAS-Bench-201：如表 3 所示，DisWOT 以最低的训练成本实现了近乎 SOTA 的性能表现。尤其是在 CIFAR-10 数据集上，比需要训练的 NAS 方法提速了 180×。

▲ 表3. 不同 NAS 方法在 CIFAR-10/CIFAR-100/ImageNet-16-120 数据集上的精度和训练成本对比（表中部分时间数据应该有误）

4. 消融实验 - 度量指标：表 4 说明本文提出的语义相似度优于 FitNets，关系相似度优于 RKD，且两者存在互补性，共同作用能获得更好的相关性结果。

▲ 表4. S0 搜索空间不同度量指标的 Spearman 相关性

5. 其他：作者还针对“初始化方式对方法的影响”“不同教师对方法的影响”“DisWOT 与最终 KD 的性能相关性”等进行了实验分析，细节可见原文。

总结

贡献

本文为“缓解教师和学生模型差距”提供了一种新的解决方案——利用 NAS 寻找更适合教师传递知识的学生模型架构；

本文证明了在 ResNet 网络中，随机初始化的师生网络之间的特征语义和样本关系的相似性与最终蒸馏性能具有良好的相关性。这种知识蕴藏在结构而非权重中。

局限性

该方法基于两点假设“大小模型的感受野不对齐”“随机初始化模型已能很好地定位物体，具备语义信息”，但这两点假设都是建立在 CNN 网络结构上的，Transformer 结构是否满足这种假设还需验证；

只在图像分类任务上进行评估，下游任务的扩展性有待验证。

参考文献

[1] Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. In ICCV, 2019.

[2] Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. In AAAI, 2020.

[3] Sihao Lin, Hongwei Xie, Bing Wang, Kaicheng Yu, Xiaojun Chang, Xiaodan Liang, and Gang Wang. Knowledge distillation via the target-aware transformer. In CVPR, 2022.

[4] Yun-Hao Cao and Jianxin Wu. A random cnn sees objects: One inductive bias of cnn and its applications. In AAAI, 2022.

[5] Bolei Zhou, Aditya Khosla, Agata Lapedriza, Aude Oliva, and Antonio Torralba. Learning deep features for discriminative localization. In CVPR, 2016.

[6] Wonpyo Park, Dongju Kim, Yan Lu, and Minsu Cho. Relational knowledge distillation. In CVPR, 2019.

更多阅读