【CVPR 2021】基于Wasserstein Distance对比表示蒸馏方法：Wasserstein Contrastive Representation Distillation

论文地址：

https://arxiv.org/abs/2012.08674

主要问题：

目前大部分知识蒸馏（例如使用KL散度的知识蒸馏方法）可能无法在教师网络中捕获重要的结构性知识，并且往往缺乏特征泛化的能力，特别是在教师和学生被用来解决不同分类任务的情况下。

主要思路：

作者提出了一个基于Wasserstein Distance对比表示蒸馏方法，称之为Wasserstein Contrastive Representation Distillation。在该算法中，作者同时使用基础形式和对偶形式的Wasserstein Distance。

其中：

对偶形式用于度量global的知识迁移，即产生一个相反的学习目标，最大化教师和学生网络之间相互信息的下界（跟第一篇思路很类似，只不过第一篇里面距离仍用的是KL距离）；

原始形式用于mini-batch内的局部对比知识迁移，有效地匹配了教师和学生网络之间的特征分布。最终结果能够达到state-of-the-art的效果。

Wasserstein Distance：

基本内容：

Wasserstein Distance是最近在基于对比的知识蒸馏方法中提出的一种距离度量，使用它的目标往往是将相似的样本移动得更近，同时将特征空间中不同的样本分开。

定义：

考虑两个概率分布：$x_1\sim p_1$和$x_2\sim p_2$，那么$p_1,p_2$的Wasserstein-1距离就可以写作：

$W(p_1,p_2)=in{f}_{\pi \in \prod (p_1,p_2)}{\int }_{M\times M}c(x_1,x_2)d\pi (x_1,x_2)$

其中$c(\cdot )$是一个点对点的用来评估距离的损失函数，$\prod$是$p_1(x_1)$和$p_2(x_2)$所有可能的联合概率分布，$M$是$x_1,x_2$所在的特征空间，$\pi (x_1,x_2)$是满足${\int }_M\pi (x_1,x_2)d{x}_2=p_1(x_1)$和${\int }_M\pi (x_1,x_2)d{x}_1=p_2(x_2)$的联合概率分布

基于Kantorovich-Rubenstein二元性，WD可以写作对偶的形式：

$W(p_1,p_2)=su{p}_{||g|{|}_L\le 1}{\mathbb{E}}_{x_1\sim p_1}[g(x_1)]-{\mathbb{E}}_{x_2\sim p_2}[g(x_2)]$

其中$g$是一个满足1-Lipschitz约束的函数（往往是个神经网络）

具体实现：

Global Contrastive Knowledge Transfer：

对于全局对比知识迁移，作者考虑在logits之前的层最大化两个特征表示$h^S,h^T$的相关信息（$MI$），即试图通过KL散度将联合分布$p(h^T,h^S)$与边缘分布$\mu (h^T)$和$\nu (h^S)$的乘积相匹配：

$I(h^S,h^T)=KL(p(h^S,h^T)||\mu (h^T)\nu (h^S))$

由于联合分布和边缘分布都是隐式的（即我们没法直接计算），因此我们可以用NCE（Noise Contrastive Estimation）的方法来近似估计$MI$。

具体地说，我们将来自联合分布的对表示为同余对（congruent pair），独立的边缘分布的乘积的对表示为不余对。换句话说就是同余对是指将相同的数据输入提供给教师和学生网络，而不同余对由不同的数据输入组成。

跟 Complementary Relation Contrastive Distillation 论文中的做法类似，作者也引入了带有隐变量 $\eta$ 的分布 $q$：

$q(h^T,h^S|\eta =1)=p(h^T,h^S)$

$q(h^T,h^S|\eta =0)=\mu (h^T)\nu (h^S)$

这里我们假设1个相关关系对带有$1$个不相关关系对，那么$q(\eta =1)=q(\eta =0)=1/2$

基于Complementary Relation Contrastive Distillation我们同样可以推导出：

$I(h^T,h^S)\ge {\mathbb{E}}_{q(h^T,h^S|\eta =1)}\log q(\eta =1|h^T,h^S)$

同样使用一个函数$g$来评估一个关系对是否来自联合分布，并且可以通过NCE loss来学习：

${\mathcal{L}}_{NCE}={\mathbb{E}}_{q(h^T,h^S|\eta =1))}\log g(h^T,h^S)+{\mathbb{E}}_{q(h^T,h^S|\eta =0)}\log [1-g(h^T,h^S)]$

并且${\mathcal{L}}_{NCE}$可同时优化函数$g$的参数和网络$S$的参数

不同于Complementary Relation Contrastive Distillation使用神经网络作为函数$g$，因为这样有两个缺陷：

• g可能对输入中的小数值变化很敏感，从而产生比较差的性能，尤其是当网络架构或学生和教师网络的训练数据集不同的时候
• g可能会出现模式坍塌的问题（参考Wasserstein GAN）

为了解决这个问题，作者使用了 spectral normalization。即对于一个任意矩阵$A$，它的spectral normalization定义为：

$\sigma (A)=ma{x}_{||\beta |{|}_2\le 1}||A\beta |{|}_2$

它相当于A的最大奇异值

通过将此正则化器应用于$\hat{g}$中每个层的权重，就可以满足1-Lipschitz约束了，因此最终将其损失函数改写为：

${\mathcal{L}}_{GCKT}={\mathbb{E}}_{q(h^T,h^S|\eta =1))}\log \hat{g}(h^T,h^S)-M{\mathbb{E}}_{\hat{g}(h^T,h^S|\eta =0)}\log \hat{g}(h^T,h^S)$

训练中采样方法则与Complementary Relation Contrastive Distillation类似

Local Contrastive Knowledge Transfer

对比学习也可以应用于一个mini-batch，以进一步提高性能

具体来说，在小批量处理中，在训练学生网络时，从教师网络中提取的特征 { h i T } i = 1 n \{h^T_i\}^n_{i=1} {hiT​}i=1n​可以被视为一个固定的集合。理想情况下，分类信息被封装在特征空间中，因此来自学生网络中的每个元素 { h j S } j = 1 n \{h^S_j\}^n_{j=1} {hjS​}j=1n​都应该能够在这一个固定的集合中找到nerighbor

而且我们可以推测，在特征空间临近的样本可能共享相同的类。因此，我们可以尝试鼓励模型将$h_j^S$同时推送到几个邻居 { h i T } i = 1 n \{h^T_i\}^n_{i=1} {hiT​}i=1n​，而不是仅仅从教师网络中的一个以更好地泛化模型性能

这可以用WD的原始形式有效地实现。即当只使用有限的训练样本时，原始形式可以被解释为将概率质量从$\mu (h^T)$转移到$\nu (h^S)$的一种简单有效的方法

具体解释为我当我们有$\mu (h^T)={\sum }_{i=1}^n{u}_i{\delta }_{h_i^T}$和$\nu (h^S)={\sum }_{j=1}^n{v}_j{\delta }_{h_j^S}$，其中${\delta }_x$是以x为中心的狄拉克函数

那么我们可以把WD的一般形式写作：
$W(\mu ,\nu )={\min }_{\pi }{\sum }_{i=1}^n{\sum }_{j=1}^n{\pi }_{ij}c(h_i^T,h_j^S)={\min }_{\pi }〈\pi ,C〉$

其中${\sum }_{j=1}^n{\pi }_{ij}=u_i$，${\sum }_{i=1}^n{\pi }_{ij}={\nu }_j$，$\pi$是$h^T$和$h^S$中的离散联合概率，$C$是由$C_{ij}=c(h_i^T,h_j^S)$给出的损失矩阵，$〈\pi ,C〉=Tr〈{\pi }^{T}C〉$表示Frobenius 点积，$C(\cdot )$
表示一个用来度量两个特征向量不相关性的损失函数（例如cosine距离）

理想情况下，可以使用线性规划获得上式的全局最优值。但是该方法是不可微的，使它与现有的深度学习框架不兼容。作为一种替代方案，作者应用了Sinkhorn算法，通过添加一个凸正则化项来求解上式，即：

${\mathcal{L}}_{LCKT}={\min }_{\pi }{\sum }_{i,j}{\pi }_{ij}c(h_i^T,h_j^S)+ϵH(\pi )$

其中$H(\pi )={\sum }_{i,j}{\pi }_{ij}\log {\pi }_{ij}$，$ϵ$是一个超参数

更具体的算法可以看论文给出的伪代码：

Unifying Global and Local Knowledge Transfer：

虽然GCKT和LCKT是为不同的目标而设计的，但它们是互补的。通过优化LCKT，我们的目标是最小化边缘分布之间的差异，这相当于减少两个特征空间之间的差异，以便LCKT可以为GCKT提供一个更受约束的特征空间。另一方面，通过优化GCKT，学习到的表示也可以形成一个更好的特征空间，这反过来又帮助LCKT匹配边缘分布。

因此最终的Loss就可以写作：

$L_{WCoRD}({\theta }_S,\varphi )=L_{CE}({\theta }_S)-{\lambda }_1{L}_{GCKT}({\theta }_S,\varphi )+{\lambda }_2{L}_{LCKT}({\theta }_S)$

实验结果：

关注我的公众号：

感兴趣的同学关注我的公众号——可达鸭的深度学习教程：

联系作者：

B站：https://space.bilibili.com/470550823

CSDN：https://blog.csdn.net/weixin_44936889

AI Studio：https://aistudio.baidu.com/aistudio/personalcenter/thirdview/67156

Github：https://github.com/Sharpiless