[論文レビュー] Multi-head Knowledge Distillation for Model Compression

Huan Wang, Suhas Lohit+2 more|arXiv (Cornell University)|Dec 5, 2020

Advanced Neural Network Applications参考文献 24被引用数 5

ひとこと要約

本稿では、教師および生徒ネットワークの中間層に補助分類器ヘッドを追加することで、モデル圧縮に効果的なシンプルな手法であるマルチヘッド知識蒸留（MHKD）を提案する。これらのヘッドの出力間のKLダイバージェンスを最小化することで、多様なアーキテクチャとデータセットにおいて生徒の性能が向上し、従来の知識蒸留手法を常に上回り、対照的学習と組み合わせることで最先端の結果を達成する。

ABSTRACT

Several methods of knowledge distillation have been developed for neural network compression. While they all use the KL divergence loss to align the soft outputs of the student model more closely with that of the teacher, the various methods differ in how the intermediate features of the student are encouraged to match those of the teacher. In this paper, we propose a simple-to-implement method using auxiliary classifiers at intermediate layers for matching features, which we refer to as multi-head knowledge distillation (MHKD). We add loss terms for training the student that measure the dissimilarity between student and teacher outputs of the auxiliary classifiers. At the same time, the proposed method also provides a natural way to measure differences at the intermediate layers even though the dimensions of the internal teacher and student features may be different. Through several experiments in image classification on multiple datasets we show that the proposed method outperforms prior relevant approaches presented in the literature.

研究の動機と目的

• 中間層の監視を活用することで、モデル圧縮における知識蒸留の性能を向上させること。
• 蒸留過程における生徒ネットワークと教師ネットワークの特徴次元の不一致という課題に対処すること。
• 複雑な知識定義やアーキテクチャを導入せずに、生徒の性能を向上させるシンプルで即座に適用可能な手法を開発すること。
• 多様なネットワークアーキテクチャとデータセット、特に異種アーキテクチャのペアに対しても頑健であることを保証すること。
• 他の蒸留技術（例：メモリバンクを用いた対照的学習）と互換性を持たせ、さらなる性能向上を可能とすること。

提案手法

• 本手法は、教師および生徒ネットワークの複数の中間層に補助分類器ヘッドを導入する。
• 各中間層について、生徒と教師の補助ヘッドの出力をKLダイバージェンス損失を用いて比較する。
• 学生の学習中に、標準的な交差エントロピー損失とKLダイバージェンス損失を組み合わせることで、特徴学習をガイドする。
• 補助ヘッドが入力次元に関係なく同等の出力を学習するため、特徴次元の不一致を自然に処理できる。
• 補助ヘッドは蒸留中に訓練されるが、推論時には削除され、残るのは生徒のメイン分類器のみである。
• 本手法は、対照的学習などの他の蒸留技術と互換性があり、さらなる性能向上が可能である。

実験結果

リサーチクエスチョン

• RQ1 補助分類器による中間層監視は、モデル圧縮における知識蒸留の性能を向上させ得るか？
• RQ2 MHKDは、異種アーキテクチャペアを含む多様な教師-生徒アーキテクチャでどのように性能を発揮するか？
• RQ3 FitNet、RKD、CCなどの既存手法と比較して、MHKDは正確性と頑健性の面で優れているか？
• RQ4 MHKDは、ネガティブペアを用いた対照的学習と効果的に組み合わせられるか？
• RQ5 MHKDは補助ヘッドの配置する中間層の選択にどれほど敏感か？また、複数のヘッドを用いることで性能が向上するか？

主な発見

• MHKDは、評価されたすべてのデータセットとネットワークペアにおいて、標準的な知識蒸留（KD）を一貫して上回る。
• CIFAR-100では、ResNet32x4を教師、ShuffleNetV2を生徒とした場合、トップ1正解率が96.68%に達し、KDを最大1.2ポイント上回る。
• VGG13/MobileNetV2ペアでは、FitNet、RKD、CCと比較して正解率が1%以上向上する。
• アブレーションスタディにより、複数の補助ヘッド（例：ヘッド1〜3）を用いることで、単一ヘッドよりも優れた性能が得られ、手動でのヘッド選択の必要性がなくなることが確認された。
• CRDと組み合わせることで、最先端の性能を達成しており、対照的学習との相性と相乗効果が示された。
• 本手法はデータセットにわたる頑健性を示しており、CIFAR-100およびSVHNの両方で、すべてのベースラインを上回った。他の手法の性能トレンドがデータセット間で異なる中でも、MHKDは一貫した優位性を示した。