[Paper Review] AlphaNet: Improved Training of Supernets with Alpha-Divergence

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

Advanced Neural Network Applications50 references4 citations

TL;DR

This paper proposes AlphaNet, a novel knowledge distillation method for training supernets using adaptive α-divergence to simultaneously reduce over- and under-estimation of teacher model uncertainty. By replacing standard KL divergence with α-divergence and stabilizing training via gradient clipping, AlphaNet achieves SOTA accuracy of 80.0% top-1 on ImageNet at only 444M FLOPs, outperforming prior methods across diverse FLOP regimes.

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.

Motivation & Objective

• To address the issue of uncertainty misestimation in knowledge distillation during supernet training, where KL divergence causes student models to over- or under-estimate teacher model uncertainty.
• To improve the performance of sub-networks in weight-sharing neural architecture search (NAS) by using a more generalized divergence metric.
• To develop a stable training procedure for α-divergence-based distillation that avoids high gradient variance.
• To demonstrate the effectiveness of adaptive α-divergence across multiple supernet architectures, including slimmable networks and weight-sharing NAS.
• To establish new state-of-the-art accuracy-FLOP trade-offs on ImageNet using the proposed AlphaNet framework.

Proposed method

• Replace standard KL divergence in knowledge distillation with α-divergence, a generalized divergence that allows adaptive control over over- and under-estimation of uncertainty.
• Introduce an adaptive α-divergence mechanism that dynamically selects α values to penalize both over- and under-estimation of the teacher model’s predictive uncertainty.
• Apply gradient clipping to the α-divergence loss to stabilize training, ensuring the optimization remains valid despite high gradient variance.
• Use the largest sub-network in the supernet as the teacher model to provide soft labels for all smaller sub-networks (students) during joint training.
• Train the supernet with a combined objective: cross-entropy loss on real labels and α-divergence loss on soft labels from the teacher.
• Apply the method to both slimmable networks and weight-sharing NAS, enabling efficient architecture search across multiple FLOP regimes.

Experimental results

Research questions

• RQ1 Can α-divergence improve knowledge distillation in supernets by reducing both over- and under-estimation of teacher model uncertainty compared to KL divergence?
• RQ2 How does adaptive α-divergence affect the performance of sub-networks in weight-sharing NAS across different FLOP regimes?
• RQ3 Does gradient clipping preserve the validity of α-divergence as an optimization objective while stabilizing training?
• RQ4 Can the proposed method generalize beyond NAS to single-model distillation, such as knowledge distillation from a large MobileNet to smaller variants?
• RQ5 What is the impact of different α values on sub-network accuracy, and does an adaptive selection strategy outperform fixed α settings?

Key findings

• AlphaNet achieves 80.0% top-1 accuracy on ImageNet with only 444M FLOPs, setting a new state-of-the-art for efficient models in the 200–800 MFLOPs range.
• The method outperforms prior SOTA models including BigNAS, Once-for-All networks, and AttentiveNAS across diverse FLOP regimes.
• Using adaptive α-divergence, AlphaNet improves sub-network performance by simultaneously penalizing over- and under-estimation of teacher model uncertainty.
• Even with fixed α values (e.g., α = -1 or α = 1), the method improves upon KL-based KD, and the adaptive version yields the best results.
• The method generalizes beyond NAS: in single-model distillation, AlphaNet-trained student models achieve higher ImageNet top-1 accuracy than baseline KD across multiple architectures, including MobileNetV3, ShuffleNet, and vision transformers.
• Gradient clipping stabilizes training without compromising the divergence’s validity, enabling reliable optimization of the adaptive α-divergence objective.