剪辑：LRST

【新智元导读】大连理工大学的扣问东谈主员提倡了一种基于Wasserstein距离的常识蒸馏要道，克服了传统KL散度在Logit和Feature常识移动中的局限性，在图像分类和观点检测任务上阐扬更好。

自Hinton等东谈主的首创性使命以来，基于Kullback-Leibler散度（KL-Div）的常识蒸馏一直占主导地位。

然而，KL-Div仅比较老师和学生在相应类别上的概率，空泛跨类别比较的机制，应用于中间层蒸馏时存在问题，其无法处理不重复的散布且无法感知底层流形的几何结构。

为了惩处这些问题，大连理工大学的扣问东谈主员提倡了一种基于Wasserstein距离（WD）的常识蒸馏要道。所提倡要道在图像分类和观点检测任务上均得到了现时最好的性能，论文已被NeurIPS 2024接纳为Poster

论文地址：https://arxiv.org/abs/2412.08139

名堂地址：https://peihuali.org/WKD/

代码地址：https://github.com/JiamingLv/WKD

配景与动机先容

常识蒸馏（KD）旨在将具有大容量的高性能老师模子中的常识移动到轻量级的学生模子中。比年来，常识蒸馏在深度学习中受到了越来越多的热心，并得到了权臣进展，在视觉识别、观点检测等多个领域得到了平庸应用。

在其首创性使命中，Hinton等东谈主引入了Kullback-Leibler散度（KL-Div）用于常识蒸馏，不停学生模子的类别概率瞻望与老师模子通常。

从其时起，KL-Div在Logit蒸馏中占据主导地位，而且其变体要道DKD、NKD等也得到了令东谈主慎重的性能。此外，这些Logit蒸馏要道还不错与将常识从中间层传递的好多先进要道互相补充。

尽管KL-Div得到了浩大的奏效，但它存在的两个时弊穷苦了老师模子常识的移动。

当先，KL-Div仅比较老师和学生在相应类别上的概率，空泛延长跨类别比较的机制。

然而，践诺全国中的类别呈现不同进度的视觉通常性，举例，哺乳动物物种如狗和狼相互间的通常度较高，而与汽车和自行车等东谈主工成品则有很大的视觉各异，如图1所示。

倒霉的是，由于KL-Div是类别对类别的比较，KD和其变体要道无法显式地哄骗这种丰富的跨类别常识。

图1 左图使用t-SNE展示了100个类别的镶嵌散布。不错看出，这些类别在特征空间中阐扬出丰富的互关研究 (IR)。然而，右图中的KL散度无法显式地哄骗这些互关研究

其次，KL-Div在用于从中间层特征进行常识蒸馏时存在局限性。图像的深度特征频频是高维的且空间尺寸较小，因此其在特征空间中相配稀薄，不仅使得KL-Div在处理深度神经网罗特征的散布时存在困难。

KL-Div无法处理不重复的破坏散布，而且由于其不是一个度量，在处理团结散布时才气有限，无法感知底层流形的几何结构。

图2 基于Wasserstein距离（WD）的常识蒸馏要道的总览图

为了惩处这些问题，扣问东谈主员提倡了一种基于Wasserstein距离的常识蒸馏要道，称为WKD，同期适用于Logit蒸馏（WKD-L）和Feature蒸馏（WKD-F），如图2所示。

在WKD-L中，通过破坏WD最小化老师和学生之间瞻望概率的各异，从而进行常识转动。

通过这种形势，延长跨类别的比较，大要有用地哄骗类别间的互关研究（IRs），与KL-Div中的类别间比较变成通晓对比。

关于WKD-F，扣问东谈主员哄骗WD从中间层特征中蒸馏常识，选拔参数化要道来建模特征的散布，并让学生径直匹配老师的特征散布。

具体来说，哄骗一种最平庸使用的团结散布（高斯散布），该散布在给定特征的1阶和2阶矩的情况下具有最大熵。

论文的主要孝顺不错总结如下：

提倡了一种基于破坏WD的Logit蒸馏要道（WKD-L），不错通过老师和学生瞻望概率之间的跨类别比较，哄骗类别间丰富的互关研究，克服KL-Div无法进行类别间比较的时弊。

将团结WD引入中间层进行Feature蒸馏（WKD-F），不错有用地哄骗高斯散布的Riemann空间几何结构，优于无法感知几何结构的KL-Div。

在图像分类和观点检测任务中，WKD-L优于相配强的基于KL-Div的Logit蒸馏要道，而WKD-F在特征蒸馏中优于KL-Div的对比要道和起初进的要道。WKD-L和WKD-F的结合不错进一步提高性能。

用于常识移动的WD距离

用于Logit蒸馏的破坏WD距离

类别之间的互关研究（IRs）

如图1所示，践诺全国中的类别在特征空间中阐扬出复杂的拓扑关系。相易类别的特征会连归拢变成一个散布，而相邻类别的特征有重复且不成皆备分离。

因此，扣问东谈主员提倡基于CKA量化类别间的互关研究（IRs），CKA是一种归一化的Hilbert-Schmidt孤独性准则（HSIC），通过将两个特搜集映射到再生核希尔伯特空间（RKHS）来建模统计关系。

当先将每个类别中整个检察样本的特征组成一个特征矩阵，之后通过狡计任性两个类别特征矩阵之间的CKA得到类间互关研究（IR）。狡计IR的资本不错忽略，因为在检察前仅需狡计一次。

由于老师模子频频包含更丰富的常识，因此使用老师模子来狡计类别间的互关研究

。

死亡函数

用

和

分手示意老师模子和学生模子的瞻望类别概率，其通过softmax函数和温度对Logit狡计得到。将破坏的WD示意为一种熵正则化的线性贪图：

其中

和

分手示意每单元质地的运输资本和在将概率质地从

出动到

时的运输量；

是正则化参数。

界说运输资本

与通常度度量

成负关联。

因此，WKD-L的死亡函数不错界说为：

用于Feature蒸馏的团结WD距离

特征散布建模

将模子某个中间层输出的特征图重塑为一个矩阵，其中第i列

示意一个空间特征。

之后，揣度这些特征的一阶矩

和二阶矩

，并将二者当作高斯散布的参数来建模输入图像特征的散布。

死亡函数

设老师的特征散布为高斯散布

。访佛地，学生的散布记为

。

两者之间的团结Wasserstein距离（WD）界说为：

其中，

和

是高斯变量，q示意息争散布。最小化上式不错得到闭集体式的WD距离。此外，为了均衡均值和协方差的作用，引入了一个均值-协方差比率γ，临了死亡界说为：

实验分析和比较

扣问东谈主员在ImageNet和CIFAR-100上评估了WKD在图像分类任务中的阐扬，还评估了WKD在自蒸馏（Self-KD）中的有用，而且将WKD扩展到观点检测，并在MS-COCO上进行了实验。

在ImageNet上的图像分类

扣问东谈主员在ImageNet的在两种开导下与现存使命进行了比较。开导（a）触及同质架构，其中老师和学生网罗分手为ResNet34和ResNet18；开导（b）触及异质架构，在该开导中，老师网罗为ResNet50，学生网罗为MobileNetV1。

关于Logit蒸馏，WKD-L在两种开导下均优于经典的KD偏激整个变体。关于特征蒸馏，WKD-F也迥殊现时的最好要道ReviewKD；临了，WKD-L和WKD-F的结合进一步进步了性能，超越了强有劲的竞争要道。

表1 在ImageNet上的图像分类后果

在CIFAR-100上的图像分类

扣问东谈主员在老师模子为CNN、学生为Transformer或反之的开导下评估了WKD要道，使用的CNN模子包括ResNet（RN）、MobileNetV2（MNV2）和ConvNeXt；Transformer模子包括ViT、DeiT和Swin Transformer。

关于Logit蒸馏，WKD-L在从Transformer到CNN移动常识或反之的开导下永恒优于最新的OFA要道。关于特征蒸馏，WKD-F在整个实验开导中名循序一；

扣问东谈主员以为，关于跨CNN和Transformer的常识转动，斟酌到两者特征各异较大，WKD-F比像FitNet和CRD这么径直对原始特征进行对皆的要道更为适合。

表2 CIFAR-100上跨CNN和Transformer的图像分类后果（Top-1准确率）

在ImageNet上的自蒸馏

扣问东谈主员在Born-Again Network（BAN）框架中将WKD要道用于自蒸馏任务（Self-KD）。

使用ResNet18在ImageNet上进行实验，后果如表3所示，WKD-L得到了最好后果，比BAN的Top-1准确率迥殊约0.9%，比第二高的USKD要道迥殊0.6%。这一比较标明，WKD要道不错很好地延长到自蒸馏任务中。

表3 在ImageNet上自蒸馏（Self-KD）的后果

在MS-COCO上的观点检测

扣问东谈主员将WKD扩展到Faster-RCNN框架下的观点检测中。关于WKD-L，使用检测头中的分类分支进行Logit蒸馏。关于WKD-F，径直从输入到分类分支的特征中进行常识移动，即从RoIAlign层输出的特征来狡计高斯散布。

关于Logit蒸馏，WKD-L权臣优于经典的KD，并略略优于DKD。关于特征蒸馏，WKD-F在两个开导中均权臣迥殊之前的最好特征蒸馏要道ReviewKD。临了，通过结合WKD-L和WKD-F，阐扬迥殊了DKD+ReviewKD。当使用非常的边框回来进行常识瞬息间，WKD-L+WKD-F进一步提高并超越了之前的起初进要道FCFD。

表4 在MS-COCO上的观点检测后果。†：使用了非常的边框回来进行常识移动

论断

Wasserstein距离（WD）在生成模子等多个领域中已展现出相较于KL散度（KL-Div）的权臣上风。

然而，在常识蒸馏领域，KL散度仍然占据主导地位，现在尚不了了Wasserstein距离能否放置更优的阐扬。

扣问东谈主员以为，早期基于Wasserstein距离的常识蒸馏扣问未能充分施展该度量的后劲。

因此，文中提倡了一种基于Wasserstein距离的全新常识蒸馏要道，大要从Logit和Feature两个方面进行常识移动。

多数的实验标明，破坏体式的Wasserstein距离在Logit蒸馏中是现时主流KL散度的极具后劲的替代决策，而团结体式的Wasserstein距离在中间层特征移动中也得到了令东谈主折服的性能阐扬。

尽管如斯，该要道仍存在一定局限性：WKD-L比较基于KL散度的Logit蒸馏要道狡计支出更高彩票游戏app平台，而WKD-F假定特征遵照高斯散布。