《Search to Distill: Pearls are Everywhere but not the Eyes》论文阅读

https://arxiv.org/abs/1911.09074 CVPR Oral

Motivation

知识蒸馏一般是由 teacher，student 两个网络组成，teacher 一般是 ResNet 152 这样的大模型，student 一般是 Res50 这样的小模型。为了让小模型能学的更好，student 网络预测的结果不仅和 ground truth 算 loss，还和 teacher 网络的预测值算 loss。学者一直在研究如何让 teacher 教给 student 更多的内容，从而使 student 涨更多的点。

本文的作者提出一个观点：对于不同结构的 teacher 网络来说，其对应的最优 student 结构是不一样的。如下表所示，如果拿 EfficientNet-B7 做 teacher，student 2 表现比 student 1要好；但是把 teacher 换成 Inception-ResNet-v2 之后，反而是 student 1 表现比 student 2 要好。

之前人们研究 KD 为什么能 work 的时候，主要是从 teacher 提供的 soft label 比 ground truth 的 one-hot label 要好的角度来分析。但是却没有从网络结构的角度来考虑。虽然 KD 可以用在任意两对 teacher-student 网络上，但是作者发问： 所有学生网络是否同样有能力接受不同老师的知识 (are all student networks equally capable of receiving knowledge from different teachers)？

为了验证这个假设，作者选了 8 个不同架构的 teacher 网络，从 ResNet 152 到 EfficientNet-B7，如下表所示：

然后作者又选了 5 个不同的 student 模型。作者说这 5 个模型是从 MNAS 的搜索空间搜出来的，只用 ground truth 训练的话结果相近，从 65.4 到 65.6，只差 0.2。不过作者并没有提供这 5 个模型的结构图、参数量之类的细节。

这样 8 个 teacher，5 个 student 就可以得到 5x8=40 组实验结果，选其中的 20 组如下表所示。可以看出对于相同的 teacher 模型，不同的 student 学的效果能差 1 个点。

下图 (a) 分析了 teacher 之间的相似度：先根据模型输出的过了 softmax 之后的结果算 KL 散度，然后得到 Confusion-Matrix。由图 (a)可以看出 Teacher - a 和 Teacher - b 是分布比较接近的，但是其对应的学的好的 student 模型却完全不一样。

上述的分析都是说明 structural knowledge 对 KD 的重要性， student 的结构对于 KD 能够达到的 upper bound 有影响。作者提出了一种 Architecture-aware Knowledge Distillation (AKD) 的概念：在给定了一个 teacher 模型之后，可以用 NAS 搜一个能在上面学的最好的 student 模型。在比较小的计算量下，作者搜出来的网络最后在 ImageNet、MegaFace 等数据集上都取得了非常好的效果。

由下图可以看出，传统的 NAS 用的是 ground truth 的 one-hot label，搜索空间更大。而本文提出的 AKD 还用到了 teacher 模型提供的 soft label，搜的是对某个 teacher 模型最友好的 student 结构，两者的搜索空间是不一样的。

Method

作者用了基于强化学习的 NAS 方法来搜一个 student 模型，同时也限制了latency 让这个模型不能太大。训练的时候 teacher model 也会提供 reward。如果抛开 teacher model 的监督，其实是一个很常规的 RNN 做 controller 的方法，搜索空间和 Mnasnet 相似。由于要等 teacher 模型的反馈，作者也说明了速度比传统的 NAS 慢 3~4 倍，200 个 TPU 要跑 5 天。整个模型的 pipeline 如下图所示：

作者把基于 KD-guided reward 搜出来的网络结构统称为 AKDNet，而 普通的 classification-guided reward 搜出来的网络统称为 NASNet。 下图展示了 AKDNet 搜索过程中的一些结果，逐一介绍：(A)：target latency 是 15 ms，可以看出采样比较多的地方都集中在 15 ms附近； (B)：latency-accuracy 之间的权衡，x轴越大，latency 越高，准确率越高。因此搜索的最优解就是在画出来的那条红线附近； (C，Ｄ)：随着采样时间的增加，latency 越来越接近 15 ms，Acc 也越来越高； (E)：模型的 latency 和 FLOPS 之间的关系，大致是成线性相关的。

为了更好地展示 AKDNet 和 NASNet 的区别，作者把他们两个采样的点随着采样时间的增加的变化可视化的出来。可以看出越到后面的时候，两种方法采样的点完全分开了。

作者认为 teacher 模型的 reward 带来了“结构化信息”(structural knowledge)，而 structural knowledge 也是 teacher 教给 student 的一部分，如果 teacher 和 student 的结构差异很大，就会妨碍 student 学到的效果。 为了这个验证假设，作者继续提出了两个问题： 问题１： 如果两个相同的 RL Agent 用不用的 teacher model 提供 award，他们会收敛到搜索空间的相同点上吗？ 答案是这两个 Agent 会收敛到完全不同的点，如下图所示

问题２： 如果两个不同的 RL Agent 用相用的 teacher model 提供 award，他们会收敛到搜索空间的相同点上吗？ 结论是会收敛到的位置会非常相似，如下图。作者认为这样就验证了 KD 过程中存在 structural knowledge 的假设。

接下来作者可视化了 AKDNet and NASNet 搜出来的结构有什么区别，如下图。可以看出来他们偏爱的运算是不一样的。

实验

在 ImageNet 上的详细结果如下表。在 Latency 限制为 75 ms 的时，从74.7->75.5，也就是说搜索模型的时候加入 teacher reward 最终能涨将近一个点。而对于 Latency 限制为 15 ms 的小模型，能从 63.9 涨到 66.5。

下表说明了试了不同的 KD Policy，加在 NAS 上面总是能涨的：

在不同 Latency 限制下，和 SOTA 方法比较的结果如下表，

作者尝试了把 Latency 的约束换成 FLOPS，发现这两个指标之间有一定的线性关系，搜出来的结果差别不大，如下表所示：

换一个数据集，在ＭegaFace 上验证的结果如下表：

此外，作者也证明了涨点之后并没有过拟合。当 teacher 是多个 model ensemble 的结果的时候，本文提出的方法仍然能涨。

结论

作者在结论中总结说，student model 和 teacher model 结构比较相似的话，学的会比较好。未来一个潜在的研究方向是如何定义一个 metric 指标，衡量不同网络之间的结构相似性。

简单概况的话，本文其实就是在指定 teacher 结构的前提下，搜一个对其学的最好的 student 结构。具体做法是常规的 NAS 里面加一个 teacher reward。