@inproceedings{
lv2024wasserstein,
title={Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation},
author={Jiaming Lv and Haoyuan Yang and Peihua Li},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=1qfdCAXn6K}
}
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.
Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation
Jiaming Lv Haoyuan Yang Peihua Li