Wasserstein 距离在知识蒸馏中的优势：对抗 Kullback-Leibler 散度
计算机视觉与模式识别
2024-12-12 v1 人工智能
机器学习
摘要
自 Hinton 等人开创性工作以来，基于 Kullback-Leibler 散度（KL-Div）的知识蒸馏始终占据主导地位，最近其变体也取得了令人瞩目的性能。然而，KL-Div 仅比较教师与学生对应类别之间的概率，缺乏跨类别比较的机制。此外，KL-Div 在应用于中间层时存在问题，无法处理非重叠分布且对底层流形的几何结构毫无所知。为此，我们提出了基于 Wasserstein 距离（WD）的知识蒸馏方法。具体而言，我们提出了基于离散 WD 的 logit 蒸馏方法 WKD-L，它进行跨类别概率比较，能够显式利用类别之间丰富的相互关系。此外，我们引入基于特征蒸馏的方法 WKD-F，使用参数化方法建模特征分布，并采用连续 WD 将知识从中间层传递。全面评估表明：(1) 在 logit 蒸馏方面，WKD-L 优于非常强大的 KL-Div 变体；(2) 在特征蒸馏方面，WKD-F 优于 KL-Div 对手以及最先进的竞争方法。源代码已公开于 https://peihuali.org/WKD。
引用
@article{arxiv.2412.08139,
title = {Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation},
author = {Jiaming Lv and Haoyuan Yang and Peihua Li},
journal= {arXiv preprint arXiv:2412.08139},
year = {2024}
}
备注
Accepted to NeurIPS 2024. Equal contribution from first two authors