[論文レビュー] Efficient Knowledge Distillation from Model Checkpoints

Chaofei Wang, Qisen Yang+3 more|arXiv (Cornell University)|Oct 12, 2022

Machine Learning and Data Classification被引用数 7

ひとこと要約

本稿では、完全に収束したモデルではなく、中間のモデルチェックポイントを知識蒸留における教師モデルとして使用することを提案する。情報ボトルネック理論を用いて、中間モデルが入力特徴間の相互情報量が高いために「暗黙の知識」をより多く保持しており、これが学生モデルの性能向上に寄与することを特定した。実験では、最適な中間教師モデルが最終モデルや強力なアンサンブルを上回ることを示し、提案された選択アルゴリズムによりCIFAR-100で最大0.8%の精度向上を達成した。

ABSTRACT

Knowledge distillation is an effective approach to learn compact models (students) with the supervision of large and strong models (teachers). As empirically there exists a strong correlation between the performance of teacher and student models, it is commonly believed that a high performing teacher is preferred. Consequently, practitioners tend to use a well trained network or an ensemble of them as the teacher. In this paper, we make an intriguing observation that an intermediate model, i.e., a checkpoint in the middle of the training procedure, often serves as a better teacher compared to the fully converged model, although the former has much lower accuracy. More surprisingly, a weak snapshot ensemble of several intermediate models from a same training trajectory can outperform a strong ensemble of independently trained and fully converged models, when they are used as teachers. We show that this phenomenon can be partially explained by the information bottleneck principle: the feature representations of intermediate models can have higher mutual information regarding the input, and thus contain more "dark knowledge" for effective distillation. We further propose an optimal intermediate teacher selection algorithm based on maximizing the total task-related mutual information. Experiments verify its effectiveness and applicability.

研究の動機と目的

• 知識蒸留における学生モデルの最良化に寄与する高精度な、完全に訓練された教師モデルが最適であるという従来の仮定に挑戦すること。
• 低い精度であるにもかかわらず、中間モデルが完全に収束したモデルよりも優れた教師として機能する理由を解明すること。
• 情報理論に基づいた理論的裏付けと効率性を兼ね備えた、最適な中間教師モデルの選択手法を開発すること。
• 単一の訓練経路からの弱いスナップショットアンサンブルが、独立に訓練された強力なアンサンブルを上回ることを実証すること。
• 複数のフルトレーニングモデルやアンサンブルの訓練を回避することで、トレーニングコストを削減しながら蒸留性能を向上させること。

提案手法

• 非パラメトリックなエントロピー推定器を用いて、各訓練チェックポイントにおける入力と特徴間の相互情報量（I(X;F)）およびターゲットと特徴間の相互情報量（I(Y;F)）を推定する。
• I(X;F)とI(Y;F)の和を最大化する基準を定式化し、知識の豊かさとタスク関連性の両立を図る中間教師モデルの最適選択基準を提示する。
• 複数の独立した教師モデルのトレーニングを必要とせず、1つの訓練経路に沿った最適なチェックポイントを特定するアルゴリズムを構築する。
• 標準的なKD損失関数と互換性があり、任意の学生-教師アーキテクチャに適用可能である。
• チェックポイント同士を比較するために正規化された相互情報量スコアを用い、タスク関連情報の合計を最大化するモデルを選択する。
• 高価なフルトレーニング済みモデルやアンサンブル教師の代わりに、最適なタイミングで選択された1つの中間チェックポイントを用いることで、効率的な蒸留を実現する。

実験結果

リサーチクエスチョン

• RQ1 中間モデルチェックポイントは、精度が低いにもかかわらず、なぜ完全に収束したモデルよりも知識蒸留の教師として優れているのか？
• RQ2 1つの訓練経路からの弱いスナップショットアンサンブルが、独立に訓練された強力なアンサンブルを上回ることができるか？
• RQ3 情報ボトルネック原理は、なぜ中間モデルが優れた知識転送能力を示すのかをどのように説明できるか？
• RQ4 蒸留性能を最大化する最適な中間教師モデルの選択戦略は何か？
• RQ5 複数の高価なアンサンブルを排除し、1つの最適タイミングのチェックポイントを用いることで、学生モデルの精度を維持または向上させることができるか？

主な発見

• 相互情報量最大化により選択された最適な中間教師モデルは、CIFAR-100におけるWRN-40-1学生モデルでKD精度73.26%を達成し、完全に収束した教師モデル（72.68%）を0.58%上回った。
• ResNet-110では、最適な中間教師モデルが72.63%のKD精度を達成し、フルモデルの72.48%と0.7xチェックポイントの72.46%を上回った。
• 中間地点（0.5x）のチェックポイントは、すべての評価アーキテクチャで完全に収束したモデルを上回り、CIFAR-100で平均0.31%の向上を示した。
• 1つの訓練走行からの弱いスナップショットアンサンブルが、独立に訓練された強力なフルアンサンブルを上回ることを実証し、提案された選択戦略の有効性を示した。
• 複数のフルトレーニングモデルやアンサンブルの訓練を排除することでトレーニングコストを削減しながら、依然として優れた蒸留性能を達成した。
• 過剰訓練により分類相関情報が抑制され、蒸留に利用可能な「暗黙の知識」の質が低下することを確認した。