[論文レビュー] Does Knowledge Distillation Really Work?

Samuel Stanton, Pavel Izmailov+3 more|arXiv (Cornell University)|Jun 10, 2021

Domain Adaptation and Few-Shot Learning参考文献 43被引用数 11

ひとこと要約

この論文は、知識蒸留が大規模な教師モデルから小規模な学生ネットワークへ本当に知識を転送しているかを調査し、学生の一般化性能が向上しているにもかかわらず、予測分布の整合性（fidelity）は予想に反して低く抑えられていることを発見した。主な要因は最適化の難易度であり、モデル容量やデータ品質の問題ではなく、学生が教師の重みに近い初期化をしても最適解に収束しないことが原因である。

ABSTRACT

Knowledge distillation is a popular technique for training a small student network to emulate a larger teacher model, such as an ensemble of networks. We show that while knowledge distillation can improve student generalization, it does not typically work as it is commonly understood: there often remains a surprisingly large discrepancy between the predictive distributions of the teacher and the student, even in cases when the student has the capacity to perfectly match the teacher. We identify difficulties in optimization as a key reason for why the student is unable to match the teacher. We also show how the details of the dataset used for distillation play a role in how closely the student matches the teacher -- and that more closely matching the teacher paradoxically does not always lead to better student generalization.

研究の動機と目的

• 知識蒸留が強い一般化性能の向上を示すにもかかわらず、教師モデルと学生モデルの間で予測の忠実度が高くならない理由を調査すること。
• 教師と学生モデルの予測に差異が生じる原因が最適化の難易度、データ品質、あるいは同定可能性の問題のいずれに起因するかを特定すること。
• 蒸留の忠実度と一般化性能を分離し、高い精度が必ずしも高い忠実度を意味しないことを示すこと。
• データ拡張、延長された訓練、共有初期化といった介入策が忠実度を向上させるかどうかを評価し、その条件を同定すること。

提案手法

• 自己蒸留およびアンサンブル蒸留の両方の設定において、CIFAR-100でResNet-20およびResNet-56アーキテクチャを用いて実験を実施した。
• 生成対抗ネットワーク（GAN）で生成されたサンプルを実データに追加することで、蒸留データセットのサイズと品質を段階的に変化させ、その影響を評価した。
• ランダム初期化と教師の初期重みを用いた初期化の異なるスケームを比較した。
• 同意率、予測分布間のKLダイバージェンス、およびネットワークの前活性化の中心化カーネル整合性（CKA）を用いて忠実度を測定した。
• 訓練期間、データ拡張（例：MixUp）、最適化アルゴリズム（例：Adam）の選択が忠実度と一般化性能に与える影響を評価した。
• すべての指標について10回の異なる乱数シードを用い、平均値と標準偏差を推定することで、結果の堅牢性を確保した。

実験結果

リサーチクエスチョン

• RQ1 現代のディープラーニング環境下で、知識蒸留は教師と学生モデルの間で予測の忠実度をどれほど達成できるか？
• RQ2 学生モデルの容量が十分にある場合でも、蒸留データセットの品質やサイズを向上させることで、学生の忠実度が向上するのか？
• RQ3 教師と学生の重みを共有して初期化することで、蒸留の忠実度を顕著に向上させられるか、それとも根本的に最適化の問題に起因するのか？
• RQ4 教師の予測を再現できない原因が同定可能性の問題にあるのか、それとも訓練中の非最適収束に起因するのか？
• RQ5 学生の一般化性能（精度）と蒸留の忠実度（分布整合性）の関係は何か？両者は常に同時に向上するのか？

主な発見

• 十分なモデル容量があるにもかかわらず、知識蒸留はしばしば高い忠実度に到達できない。教師と学生の予測分布の間に顕著な乖離が見られることから明らかである。
• 自己蒸留の文脈では、GANで生成されたサンプルを追加することでデータセットサイズを拡大したが、これにより忠実度は向上したものの、テスト精度が低下した。これは最適化の複雑さとデータ品質の間でトレードオフが生じていることを示している。
• 学生を教師の最終重みに近い初期状態から開始すると、同じ損失の谷（loss basin）に収束し、予測の一致率がほぼ100％に達する。これは最適化の問題が忠実度の低さの根本的原因であることを示している。
• 教師と学生の重みを共有して初期化することでCKAは向上（77.174 → 77.098）したが、機能的整合性や予測KLダイバージェンスに顕著な改善は見られず、初期化だけでは不十分であることが示された。
• MixUpや延長された訓練、注意深く設計されたデータ拡張といった介入を行っても、学生は依然として教師の予測を再現できず、最適化ダイナミクスが主な障壁であることが確認された。
• 高い学生の精度は、必ずしも高い忠実度を意味しない。実際、最も精度の高いモデルが、常に教師の予測分布に最も忠実であるとは限らない。