[论文解读] Towards Oracle Knowledge Distillation with Neural Architecture Search

Minsoo Kang, Jonghwan Mun+1 more|arXiv (Cornell University)|Nov 29, 2019

Advanced Neural Network Applications参考文献 39被引用 5

一句话总结

本文提出了一种新颖的知识蒸馏框架，将神经架构搜索（NAS）与最优知识蒸馏（OD）相结合，从集成教师模型中训练出高精度、高效率的学生模型。通过搜索最优学生架构并使用模拟最佳集成预测的最优损失函数，该方法缩小了模型容量差距，在准确率和内存效率方面均超越了教师模型及标准蒸馏基线。

ABSTRACT

We present a novel framework of knowledge distillation that is capable of learning powerful and efficient student models from ensemble teacher networks. Our approach addresses the inherent model capacity issue between teacher and student and aims to maximize benefit from teacher models during distillation by reducing their capacity gap. Specifically, we employ a neural architecture search technique to augment useful structures and operations, where the searched network is appropriate for knowledge distillation towards student models and free from sacrificing its performance by fixing the network capacity. We also introduce an oracle knowledge distillation loss to facilitate model search and distillation using an ensemble-based teacher model, where a student network is learned to imitate oracle performance of the teacher. We perform extensive experiments on the image classification datasets---CIFAR-100 and TinyImageNet---using various networks. We also show that searching for a new student model is effective in both accuracy and memory size and that the searched models often outperform their teacher models thanks to neural architecture search with oracle knowledge distillation.

研究动机与目标

• 为了解决知识蒸馏中因模型容量不匹配而导致的学生与教师模型之间持续存在的性能差距。
• 通过使学生模型能够学习到教师预测的最佳组合（即最优性能），提升从集成教师模型进行知识蒸馏的效果。
• 开发一种基于搜索的方法，识别出专为蒸馏量身定制的高性能、紧凑型学生架构，而非预先固定学生结构。
• 证明结合架构搜索与最优蒸馏可生成在准确率和效率方面均优于标准蒸馏和教师集成的模型。

提出的方法

• 提出一种最优知识蒸馏（OD）损失，使学生模型学习在每个输入上模仿所有教师模型中表现最佳的预测结果，而非采用模型平均。
• 采用神经架构搜索（NAS）来发现具有自适应容量的最优学生网络架构，从基础学生模型出发，扩展至原始大小的 4 倍。
• 使用可微分的搜索空间，包含操作和跳跃连接，支持架构与权重的端到端联合优化。
• 在 NAS 过程中应用 OD 损失，引导搜索过程朝向能够最好模仿教师集成最优性能的架构。
• 在内存约束下执行搜索，搜索空间设计兼容标准卷积神经网络（CNN），并适用于图像分类任务。
• 使用搜索到的架构和 OD 损失训练最终的学生模型，实现更高的准确率与更优的参数效率。

实验结果

研究问题

• RQ1 神经架构搜索能否识别出在知识蒸馏中优于固定架构学生模型的架构？
• RQ2 最优知识蒸馏（即对每个样本模仿最佳教师预测）是否能带来高于标准蒸馏的学生准确率？
• RQ3 NAS 与 OD 的结合能否缩小学生与教师集成之间的模型容量差距，从而实现更优性能？
• RQ4 搜索得到的学生模型在性能上与教师集成相比如何，特别是在准确率和内存占用方面？

主要发现

• 在使用 5 个 ResNet-32 模型的集成教师网络时，CIFAR-100 上参数量超过 0.89M（为基线模型的 2 倍）的 KDAS 模型在使用最优蒸馏损失时，准确率达到 76.87%，与教师模型持平，但性能更优。
• KDAS 模型在参数量相近的情况下，始终优于标准知识蒸馏基线（MMNs），证明了架构搜索的有效性。
• 在 TinyImageNet 上，OD 相较于标准 KD 的优势更为显著，其中 50.4% 的训练样本在不同教师模型间存在不同的最优预测，从而实现了显著的准确率提升。
• 对于多种主干网络架构（ResNet-110、WideResNets），KDAS 检测到的学生模型在准确率上优于标准蒸馏和教师集成，同时保持了具有竞争力的内存使用量。
• 即使不使用最优损失，NAS 检测到的学生模型也优于其他 KD 方法（如 KD、DML、BSS、TAKD），表明搜索本身已具备显著价值。
• 该方法具有模型无关性，如在 CIFAR-100 和 TinyImageNet 上对多种学生架构均表现出一致的性能提升。