随着深度学习的发展,神经网络结构的设计逐渐由手工设计转变为算法自动设计。在近期的神经网络设计(Neural Architecture Search, NAS)研究中,现有的方法存在一定缺陷,结果往往不能真正体现出众多子网络正确的排序关系。为解决这一问题,微软亚洲研究院的研究员们提出了基于优先路径蒸馏的网络结构搜索方法。采用这一方式选出的网络在 ImageNet 上的分类准确率达到了80.0%, 超越了现有的 EfficientNet-B0/B1 和 MobileNetV3。该论文已被 NeurIPS 2020 接收。
从数亿数量级的神经网络结构中搜索出高性能的网络,是一个充满挑战但又令人向往的研究任务。正如深度学习变革了传统的手工图像特征设计,神经网络结构的设计也正在逐渐由手工设计转变为算法自动设计。
面对数以亿级的网络结构,将每一个可能的结构都训练收敛,并选择其中最好的结构是不现实的。在近期的神经网络设计研究中,一个被广泛使用的解决方法是先训练一个包含了所有可能结构的超网(hypernetwork),当测试某一个网络结构的性能时,直接继承超网训练后的参数。这样的做法省去了重新训练的时间,大大加快了网络搜索的速度。然而,虽然预训练超网的方法能够大幅度加速网络搜索,但因为很难对所有的路径(子模型)进行充分训练,所以其给出的结果往往不能真正体现出众多子网络正确的排序关系。
为了解决这一问题,微软亚洲研究院的研究员们提出维护一个优先路径组(prioritized path board)。也就是说,在训练超网的某一条路径时,使用元网络(meta-network)从组中选出一条性能较好的子网对其进行网络蒸馏(distillation),从而提升超网的收敛程度与性能。采用这种方式选出的网络在 ImageNet 上的分类准确率达到了80.0%, 超越了现有的 EfficientNet-B0/B1 和 MobileNetV3。该论文已被 NeurIPS 2020 接收。
基于优先路径的网络蒸馏
现有的超网训练方式多种多样,研究员们采用了简单有效的均匀采样单条路径(single-path uniform sampling strategy)作为基础,即每一次随机从超网中采一条路径进行训练。与之不同的是,在训练单条路径时,会从优先路径组中选出一条路径对其进行蒸馏,从而提升训练效果。
图1:方法示意图,左侧为常规的蒸馏方法,即采用一个预训练的模型进行蒸馏,右侧为提出的基于优先路径的蒸馏方法。
优先路径组
优先路径组是由少量性能优异的路径构成的。超网训练会对优先路径组进行动态的维护,如果采样出来的网络在效果和复杂度上均优于有限路径组中的网络,那么就会将其替换到优先路径组中。不仅如此,维护这样一个优先路径组还使得超网训练结束后可以直接从中选取最优路径,从而节省以往方法在进行网络结构搜索时运用强化学习方法或进化算法(Evolution Algorithm)的时间。在选取优先路径时,可根据公式
进行选择,其中M代表元网络,N(*, *, *)代表网络最后一层输出的特征(logits),ρ代表元网络所预测的两条路径的匹配程度。
知识蒸馏
知识蒸馏是一种被广泛应用的模型压缩方法,通过让小型网络来模仿预训练大型网络的最后一层输出特征,可以使小型网络达到接近于大型网络的表现。研究员们通过优先路径来进行蒸馏,从而无需提前训练一个大型神经网络。对超网进行更新的具体公式如下:
其中L_CE为正常的训练损失, L_KD为蒸馏损失,ρ仍为前面所提到的匹配程度。
元网络
由于不同的子网结构差异可能非常之大,因此研究员们希望能够从优先路径组中选出最有助于子网训练的优先路径对其进行知识蒸馏。经过尝试,研究员们采用了元网络的技术,将采样路径和优先路径最后一层输出的特征差输入到元网络中,由元网络判断优先路径与采样路径的匹配程度。当训练元网络时,研究员们会在验证集上计算损失,通过匹配程度ρ进行梯度回传并更新元网络M:
实验结果
对基于优先路径蒸馏的网络结构搜索算法的测试是在 ImageNet 上进行的。实验结果如图2和表1所示。可以看出,在各种模型大小下,该方法的搜索结果均超越了此前的 EfficientNet-B0/B1 和 MobileNetV3,实现了优越的性能。不仅如此,该方法搜索所需要的时长也是各种网络结构搜索算法中最短的。
图2:基于优先路径蒸馏的网络结构搜索得到的模型在 ImageNet 上的性能
表1:基于优先路径蒸馏的网络结构搜索得到的模型在 ImageNet 上的性能
除了图像分类任务外,研究员们还在物体检测任务上对算法进行了测试,结果如表2所示。可以看到,该模型同样能够泛化到物体检测任务上,在各种指标下都比已有模型有近1%的提升。
表2:基于优先路径蒸馏的网络结构搜索模型在物体检测任务上的表现
基于NNI工具接口的源码实现
NNI (Neural Network Intelligence) 是当下最热门的开源自动机器学习(AutoML)工具之一,由微软亚洲研究院与微软(亚洲)互联网工程院领衔开发。NNI 对机器学习生命周期的各个环节都做了较为全面的支持,包括特征工程、神经网络架构搜索、超参调优和模型压缩。
目前,微软亚洲研究院的研究员们已将此次提出的基于优先路径蒸馏的网络结构搜索算法通过 NNI 平台的 One-Shot 算法接口进行了实现,提供了完整的搜索、重训练和测试代码以及模型。由于 NNI 提供了统一的接口表达网络搜索空间,所以有对比此算法与其他神经网络架构搜索结果需求的用户可选择这份代码实现做参考。代码以及更多技术细节,请参见:https://github.com/microsoft/nni。
结语
本篇 NeurIPS 2020 论文针对网络结构搜索中超网训练不充分的问题,提出了使用优先路径进行蒸馏的方法,加快了超网训练的收敛和搜索速度,且搜索得到的网络性能超越了大多数现有算法。尽管如此,该方法目前仍处于初步的探索阶段,未来可能还会将模型延迟考虑到优先路径选择中,并对优先路径蒸馏的理论依据进行进一步的探索。
更多技术细节,详见论文:Cream of the Crop: Distilling Prioritized Paths For One-Shot Neural Architecture Search
论文链接:https://arxiv.org/abs/2010.15821
代码链接:https://github.com/microsoft/Cream
NNI实现链接:https://github.com/microsoft/nni