[論文レビュー] Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation

Jiayi Lv, Hee-Chul Yang+1 more|arXiv (Cornell University)|Dec 11, 2024

Brain Tumor Detection and Classification被引用数 5

ひとこと要約

本論文は WKD を提案する。Wasserstein 距離に基づく知識蒸留のアプローチで、 logits のためのクロスカテゴリ間関係推論を可能にする WKD-L と、中間特徴の連続分布一致を実現する WKD-F を提供し、ImageNet、CIFAR-100、MS-COCO において KL-Div 系変種や最新の蒸留法を上回る。

ABSTRACT

Since pioneering work of Hinton et al., knowledge distillation based on Kullback-Leibler Divergence (KL-Div) has been predominant, and recently its variants have achieved compelling performance. However, KL-Div only compares probabilities of the corresponding category between the teacher and student while lacking a mechanism for cross-category comparison. Besides, KL-Div is problematic when applied to intermediate layers, as it cannot handle non-overlapping distributions and is unaware of geometry of the underlying manifold. To address these downsides, we propose a methodology of Wasserstein Distance (WD) based knowledge distillation. Specifically, we propose a logit distillation method called WKD-L based on discrete WD, which performs cross-category comparison of probabilities and thus can explicitly leverage rich interrelations among categories. Moreover, we introduce a feature distillation method called WKD-F, which uses a parametric method for modeling feature distributions and adopts continuous WD for transferring knowledge from intermediate layers. Comprehensive evaluations on image classification and object detection have shown (1) for logit distillation WKD-L outperforms very strong KL-Div variants; (2) for feature distillation WKD-F is superior to the KL-Div counterparts and state-of-the-art competitors. The source code is available at https://peihuali.org/WKD

研究の動機と目的

• カテゴリ間の相互関係 IRs を活用して、カテゴリ別 KL-Div にとどまらない知識蒸留の改善を動機づける。
• logits の蒸留のために WD ベースの蒸留手法を提案する（WKD-L）と、中間特徴の蒸留のために WD ベースの蒸留手法を提案する（WKD-F）。
• Centered Kernel Alignment (CKA) を用いてカテゴリ IRs をモデリングし、それを logits 蒸留の WD 輸送コストへ変換する。
• ガウス分布としての中間層特徴分布をモデリングし、WD をリーマン幾何学的距離として計算して特徴を蒸留する。

提案手法

• 教師特徴量上で計算される CKA によってカテゴリ間の IR を定義し、それを logit 蒸留の WD 輸送コストへ変換する。
• IR に基づく類似性から導かれたコストとエントロピー正則化付きの輸送問題として、教師と学生の logits の離散 WD を定式化する。
• ロジットにターゲットとノンターゲットの分離を組み込み、ノンターゲットに対する WD とターゲットに対するクロスエントロピーを組み合わせた2項ロスを適用する。
• 特徴については、教師と学生の分布を平均と共分散を持つ Gaussians としてモデリングし、Gaussians 間の閉形式 WD（平均項と共分散項の和）を用いる。
• 実用性の観点から空間ピラミディングを任意で適用し、実用性の観点から Gaussian Diag 対 Full 共分散を選択し、平均-共分散の寄与を gamma パラメータで調整する。

(a) Real-world categories exhibit rich interrelations (IRs) in feature space, e.g., dog is near other mammal while far from artifact like car. We quantify pairwise IRs as feature similarities among categories. Best viewed by zooming in .

実験結果

リサーチクエスチョン

• RQ1 WD ベースの蒸留はカテゴリ間相互関係を活用して logits 蒸留において KL-Div ベースの手法を上回れるか？
• RQ2 中間層の特徴をガウス分布として WD でモデリングすることは、KL-Div および非パラメトリック手法と比較して知識伝達を改善するか？
• RQ3 IR モデリング手法（CKA とさまざまなカーネル）の WKD-L の性能に与える影響は？
• RQ4 WKD-L および WKD-F は単独および組み合わせで、画像分類と物体検出タスクでどのような性能を示すか？

主な発見

• WKD-L は ImageNet および CIFAR-100 での logits 蒸留において強力な KL-Div 系変種を上回る。
• WKD-F は特徴蒸留で KL-Div 対応を上回り、ガウス分布（Diag）の方が頑健性と効率性の点で好まれることが多い。
• CKA を用いたカテゴリ間関係のモデリング（特に RBF や線形カーネル）は WD ベースの logits 蒸留を高める。
• WKD-L と WKD-F の組み合わせは、いずれか単独よりも分類と検出タスクの両方でさらなる改善をもたらす。
• MS-COCO の物体検出では、WD ベースの蒸留は KL-Div ベースの手法に対して競争力のある利得を示す。

(b) For logit distillation, discrete WD performs cross-category comparison by exploiting pairwise IRs, in contrast to KL-Div that is a category-to-category measure and lacks a mechanism to use such IRs (cf. Figure 2 ). For feature distillation, we use Gaussians for distribution modeling and continuo