[論文レビュー] 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発散度に代えてα発散度を採用し、勾配クリッピングによる訓練の安定化を実施することで、444M FLOPsでImageNetにおけるトップ1精度80.0%を達成し、多様なFLOP制約下で先行手法を上回るSOTA性能を実現した。

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におけるトップ1精度80.0%を達成し、200–800 MFLOPs範囲の効率的モデルにおいて新たなSOTAを樹立した。
• BigNAS、Once-for-Allネットワーク、AttentiveNASなど、先行SOTAモデルを多様なFLOP制約下で上回った。
• 適応的α発散度を用いることで、教師モデルの不確実性の過小・過大評価を同時にペナルティ化し、サブネットワークの性能が向上した。
• 固定α値（例：α = -1 または α = 1）を用いてもKLベースのKDを上回る性能を示し、適応的バージョンが最良の結果を達成した。
• 本手法はNASにとどまらず一般化可能である。単一モデル蒸留においても、MobileNetV3、ShuffleNet、ビジョントランスフォーマーを含む複数のアーキテクチャで、ベースラインKDよりも高いImageNetトップ1精度を達成した。
• 勾配クリッピングにより訓練が安定化したが、発散度の有効性は損なわれず、適応的α発散度目的関数の信頼性ある最適化が可能となった。