超越KL！大连理工发布Wasserstein距离知识蒸馏新方法｜NeurIPS 2024

原标题：超越KL！大连理工发布Wasserstein距离知识蒸馏新方法｜NeurIPS 2024
文章来源：新智元
内容字数：7110字

基于Wasserstein距离的知识蒸馏方法：超越KL散度的知识迁移

本文总结了大连理工大学研究人员发表在NeurIPS 2024上的论文，该论文提出了一种基于Wasserstein距离 (WD) 的知识蒸馏 (KD) 方法，在图像分类和目标检测任务上取得了显著成果，并超越了传统的基于Kullback-Leibler散度 (KL-Div) 的方法。

1. 传统KL-Div方法的局限性

传统的基于KL-Div的知识蒸馏方法，尽管在Logit蒸馏中取得了成功，但存在两个主要缺陷：首先，KL-Div只比较教师和学生模型在相同类别上的概率，忽略了类别间的相互关系；其次，在Feature蒸馏中，KL-Div难以处理高维稀疏的深度特征，无法有效处理不重叠的分布，也无法捕捉底层流形的几何结构。

2. 基于Wasserstein距离的知识蒸馏方法 (WKD)

为了克服KL-Div的局限性，研究人员提出了WKD方法，该方法包含两个部分：WKD-L (Logit蒸馏) 和 WKD-F (Feature蒸馏)。

2.1 WKD-L: 基于离散WD的Logit蒸馏

WKD-L利用离散WD最小化教师和学生模型预测概率的差异。通过这种方式，WKD-L能够进行跨类别比较，有效利用类别间的相互关系 (IRs)，这与KL-Div的类别内比较形成对比。 研究人员使用CKA来量化类别间的相互关系，并将其融入损失函数中。

2.2 WKD-F: 基于连续WD的Feature蒸馏

WKD-F利用连续WD进行中间层特征的知识蒸馏。它通过高斯分布建模特征的分布，并最小化教师和学生模型特征分布之间的WD距离。这种方法能够有效利用高斯分布的Riemann空间几何结构，克服了KL-Div无法感知几何结构的缺点。

3. 实验结果与分析

研究人员在ImageNet、CIFAR-100和MS-COCO数据集上进行了大量的实验，结果表明：

1. 在ImageNet图像分类任务中，WKD-L优于各种基于KL-Div的Logit蒸馏方法；WKD-F优于现有最先进的Feature蒸馏方法；WKD-L和WKD-F结合后，性能进一步提升。
2. 在CIFAR-100图像分类任务中，WKD在跨CNN和Transformer的知识迁移中表现出色，尤其是在Feature蒸馏方面。
3. 在ImageNet自蒸馏任务中，WKD-L取得了最佳结果。
4. 在MS-COCO目标检测任务中，WKD-L和WKD-F均显著优于现有方法，结合后性能进一步提升，甚至超越了最先进的方法。

4. 结论

这项研究证明了Wasserstein距离在知识蒸馏领域的有效性。WKD方法在Logit和Feature蒸馏方面均取得了优异的性能，为知识蒸馏研究提供了新的方向。虽然WKD-L的计算成本略高于基于KL-Div的方法，且WKD-F假设特征服从高斯分布，但这项工作仍然为知识蒸馏领域带来了重要的贡献。

联系作者

文章来源：新智元
作者微信：
作者简介：智能+中国主平台，致力于推动中国从互联网+迈向智能+新纪元。重点关注人工智能、机器人等前沿领域发展，关注人机融合、人工智能和机器人对人类社会与文明进化的影响，领航中国新智能时代。