[论文解读] AlphaNet: Improved Training of Supernets with Alpha-Divergence

Dilin Wang, Chengyue Gong+3 more|arXiv (Cornell University)|Feb 16, 2021

Advanced Neural Network Applications参考文献 50被引用 4

一句话总结

本文提出 AlphaNet，一种新颖的知识蒸馏方法，通过使用自适应 α-散度，在训练超网络时同时减少教师模型不确定性估计的过度与不足。通过用 α-散度替代标准 KL 散度，并借助梯度裁剪稳定训练，AlphaNet 在仅 444M FLOPs 的条件下实现了 ImageNet 上 80.0% 的 top-1 准确率，创下新 SOTA 记录，且在多种 FLOP 范围内均优于先前方法。

ABSTRACT

Weight-sharing neural architecture search (NAS) is an effective technique for automating efficient neural architecture design. Weight-sharing NAS builds a supernet that assembles all the architectures as its sub-networks and jointly trains the supernet with the sub-networks. The success of weight-sharing NAS heavily relies on distilling the knowledge of the supernet to the sub-networks. However, we find that the widely used distillation divergence, i.e., KL divergence, may lead to student sub-networks that over-estimate or under-estimate the uncertainty of the teacher supernet, leading to inferior performance of the sub-networks. In this work, we propose to improve the supernet training with a more generalized alpha-divergence. By adaptively selecting the alpha-divergence, we simultaneously prevent the over-estimation or under-estimation of the uncertainty of the teacher model. We apply the proposed alpha-divergence based supernets training to both slimmable neural networks and weight-sharing NAS, and demonstrate significant improvements. Specifically, our discovered model family, AlphaNet, outperforms prior-art models on a wide range of FLOPs regimes, including BigNAS, Once-for-All networks, and AttentiveNAS. We achieve ImageNet top-1 accuracy of 80.0% with only 444M FLOPs. Our code and pretrained models are available at https://github.com/facebookresearch/AlphaNet.

研究动机与目标

• 为解决在超网络训练中知识蒸馏时因 KL 散度导致学生模型对教师模型不确定性估计过度或不足的问题。
• 通过使用更通用的散度度量，提升权重共享神经架构搜索（NAS）中子网络的性能。
• 开发基于 α-散度蒸馏的稳定训练流程，避免梯度方差过高的问题。
• 在多种超网络架构（包括可裁剪网络与权重共享 NAS）中验证自适应 α-散度的有效性。
• 通过所提出的 AlphaNet 框架，在 ImageNet 上建立新的准确率-FLOP 权衡 SOTA 记录。

提出的方法

• 在知识蒸馏中用 α-散度替代标准 KL 散度，α-散度是一种可调的广义散度，能自适应控制对教师模型不确定性估计的过度与不足。
• 引入一种自适应 α-散度机制，动态选择 α 值，以同时惩罚教师模型预测不确定性估计的过度与不足。
• 对 α-散度损失应用梯度裁剪，以稳定训练，确保尽管存在高梯度方差，优化过程仍保持有效。
• 在联合训练中，使用超网络中最大子网络作为教师模型，为所有更小的子网络（学生）提供软标签。
• 以联合目标训练超网络：在真实标签上使用交叉熵损失，同时在教师提供的软标签上使用 α-散度损失。
• 将该方法应用于可裁剪网络与权重共享 NAS，实现在多种 FLOP 范围内的高效架构搜索。

实验结果

研究问题

• RQ1 与 KL 散度相比，α-散度是否能通过减少教师模型不确定性估计的过度与不足，改进超网络中的知识蒸馏？
• RQ2 自适应 α-散度在不同 FLOP 范围内对权重共享 NAS 中子网络性能有何影响？
• RQ3 梯度裁剪是否能在保持 α-散度作为优化目标有效性的前提下，稳定训练过程？
• RQ4 所提出的方法是否可泛化至 NAS 之外的单模型蒸馏任务，例如从大型 MobileNet 向小型变体进行知识蒸馏？
• RQ5 不同 α 值对子网络准确率有何影响？自适应选择策略是否优于固定 α 设置？

主要发现

• AlphaNet 在仅 444M FLOPs 的条件下，实现了 ImageNet 上 80.0% 的 top-1 准确率，创下 200–800 MFLOPs 范围内高效模型的新 SOTA 记录。
• 该方法在多种 FLOP 范围内均优于先前 SOTA 模型，包括 BigNAS、Once-for-All 网络与 AttentiveNAS。
• 通过使用自适应 α-散度，AlphaNet 同时惩罚教师模型不确定性估计的过度与不足，从而提升子网络性能。
• 即使采用固定 α 值（如 α = -1 或 α = 1），该方法仍优于基于 KL 的 KD，而自适应版本效果最佳。
• 该方法具有泛化能力：在单模型蒸馏中，AlphaNet 训练的学生模型在多种架构（包括 MobileNetV3、ShuffleNet 与视觉 Transformer）上均实现了高于基线 KD 的 ImageNet top-1 准确率。
• 梯度裁剪在不损害散度有效性的前提下稳定了训练，从而实现了对自适应 α-散度目标的可靠优化。