用沃瑟斯坦距离替代KL散度，提升知识蒸馏效果
Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation
- 用离散沃瑟斯坦距离实现类别间概率对比，捕捉类别关联
- 在图像分类和目标检测上，对蒸馏效果优于主流KL散度方法
- 适合需要深层特征传递的模型压缩场景
自Hinton等人开创性工作以来，基于KL散度的知识蒸馏占据主导地位，其变体也取得了优异性能。然而，KL散度仅比较教师与学生模型对应类别的概率，缺乏类别间对比机制；且在中间层应用时，无法处理非重叠分布，也不考虑底层流形几何结构。为此，本文提出基于沃瑟斯坦距离（WD）的知识蒸馏方法：针对输出层提出基于离散WD的逻辑值蒸馏方法WKD-L，可显式利用类别间的丰富关系；针对中间层提出参数化建模特征分布并采用连续WD的特征蒸馏方法WKD-F。在图像分类与目标检测任务上的全面评估表明：(1) WKD-L在逻辑值蒸馏中超越强基线的KL-Div变体；(2) WKD-F在特征蒸馏中优于同类KL-Div方法及当前最优方法。代码已开源。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。