[論文レビュー] An Embarrassingly Simple Approach for Knowledge Distillation

Mengya Gao, Yujun Shen+6 more|arXiv (Cornell University)|Dec 5, 2018

Advanced Neural Network Applications参考文献 40被引用数 12

ひとこと要約

本論文は、段階的知識蒸留（SSKD）を提案する。これは、知識蒸留を二段階に分離するシンプルで効果的な手法であり、第一に、教師から生徒のバックボーンへの中間特徴表現の蒸留をラベルを用いずに実行し、第二に、正解ラベルを用いてタスク固有のヘッドのみを微調整する。SSKDは損失重みのハイパラメータチューニングの必要性を排除し、ImageNet、CIFAR-100、COCOオブジェクト検出、IJB-A顔認識のベンチマークで最先端の性能を達成しており、COCOでは最大2.8%のAP上昇、ImageNetでは最大1.8%の精度上昇を達成した。

ABSTRACT

Knowledge Distillation (KD) aims at improving the performance of a low-capacity student model by inheriting knowledge from a high-capacity teacher model. Previous KD methods typically train a student by minimizing a task-related loss and the KD loss simultaneously, using a pre-defined loss weight to balance these two terms. In this work, we propose to first transfer the backbone knowledge from a teacher to the student, and then only learn the task-head of the student network. Such a decomposition of the training process circumvents the need of choosing an appropriate loss weight, which is often difficult in practice, and thus makes it easier to apply to different datasets and tasks. Importantly, the decomposition permits the core of our method, Stage-by-Stage Knowledge Distillation (SSKD), which facilitates progressive feature mimicking from teacher to student. Extensive experiments on CIFAR-100 and ImageNet suggest that SSKD significantly narrows down the performance gap between student and teacher, outperforming state-of-the-art approaches. We also demonstrate the generalization ability of SSKD on other challenging benchmarks, including face recognition on IJB-A dataset as well as object detection on COCO dataset.

研究の動機と目的

• 知識蒸留におけるハイパラメータ感受性、特にタスク損失とKD損失のバランスをとるための損失重みチューニングの難しさに対処すること。
• トレーニングプロセスを二つの明確な段階に分離することで、多様なデータセットやタスクにおいて知識蒸留の一般化性能と安定性を向上させること。
• 中間特徴の段階的・順次的蒸留が、一度に全特徴を蒸留する単一ステップの蒸留よりも優れた知識伝達をもたらすことを示すこと。
• 従来の方法が直接適用できなかった非分類タスク、例えばオブジェクト検出や顔認識に知識蒸留を拡張すること。
• COCOでFPNとRetinaNetにSSKDを適用することで、オブジェクト検出における知識蒸留の新しいベンチマークを確立すること。

提案手法

• 生徒のトレーニングを二段階に分解する：(1) 正解ラベルを一切使用せずに、教師のネットワークから生徒のバックボーンへの中間特徴表現の知識蒸留、(2) バックボーンを固定したまま、ラベル付きデータを用いてタスク固有のヘッドのみをトレーニングする。
• 段階的知識蒸留（SSKD）を実装し、ネットワークの各段階で順次的に特徴を模倣する。各サブネットワークを1つずつ別々に模倣する。
• 第一段階では、教師と生徒の対応する層間で特徴レベルの蒸留損失（例：MSEまたはKLダイバージェンス）を用いる。
• 蒸留が終了した後、生徒のバックボーンを固定し、第二段階ではクロスエントロピーまたは他のタスク固有の損失を用いて、タスクヘッドのみを微調整する。
• ResNetを用いたImageNetおよびCIFAR-100、COCOオブジェクト検出におけるFPNおよびRetinaNetへの適用を含め、さまざまなアーキテクチャとタスクにSSKDを適用する。
• 各段階でネットワークの大部分を固定することで、トレーニング時間の著しい増加を回避し、効率的なトレーニングを確保する。

実験結果

リサーチクエスチョン

• RQ1 バックボーン特徴の模倣とタスクヘッドの微調整を分離することで、損失重みのハイパラメータチューニングの必要性がなくなるか？
• RQ2 段階的・順次的な中間特徴の蒸留は、固定された損失重みを用いた一括最適化と比較して、より優れた性能を発揮するか？
• RQ3 提案手法は、画像分類、オブジェクト検出、顔認識といった多様なタスクに一般化可能か？
• RQ4 ImageNetやCOCOといった標準ベンチマークにおいて、SSKDは最先端の知識蒸留手法と比較して優れているか？
• RQ5 KDがまだ十分に検討されていない複雑な検出アーキテクチャ、例えばFPNやRetinaNetに対しても、SSKDは効果的に適用可能か？

主な発見

• ResNet-18を生徒、ResNet-34を教師として用いたImageNetでは、SSKDが2番目に優れた手法よりもトップ-1精度を1.8%向上させ、80.5%の精度を達成した。
• FPNを用いたCOCOオブジェクト検出では、ResNet-152から蒸留した場合、ResNet-50のAPが37.7%から40.5%に上昇し、2.8%の向上を達成した。
• RetinaNetでは、ResNet-152から蒸留した場合、ResNet-18のAPが32.4%から35.6%に上昇し、3.2%の改善を達成した。
• IJB-A顔認識では、SSKDが顕著な一般化性能を示し、ベースライン手法を常に上回った。
• COCOでは、すべてのオブジェクトサイズカテゴリ（AP_S、AP_M、AP_L）でSSKDが一貫した性能向上を示し、スケール変動に対する頑健性を示した。
• アブレーションスタディにより、段階的蒸留が単一ステップ蒸留を上回り、一括最適化よりも安定的かつ効果的な二段階トレーニング方式であることが確認された。