【论文笔记_知识蒸馏_2022】It’s All in the Head:Representation Knowledge Distillation through Classifier Sharin

代码地址:https://github.com/Alibaba-MIIL/HeadSharingKD

摘要

表示知识提取旨在将丰富的信息从一个模型传递到另一个模型。目前的表示提取方法主要集中于模型嵌入向量之间距离度量的直接最小化。这种直接方法可能局限于传输嵌入在表示向量中的高阶依赖，或者处理教师和学生模型之间的能力差距。在本文中，我们介绍了两种通过教师和学生之间的分类器共享来增强表示升华的方法。具体地说，我们首先表明，将教师的分类器连接到学生主干并冻结其参数有利于表征提取过程，从而产生一致的改进。然后，我们提出了另一种方法，要求根据能力有限的学生定制教师模型。这种方法与第一种方法竞争，有时甚至超过第一种方法。通过大量的实验和分析，我们证明了所提出的方法在各种数据集和任务上的有效性，包括图像分类、细粒度分类和人脸验证。例如，对于MobileFaceNet模型，我们在IJB-C数据集上实现了最先进的人脸验证性能：TAR@（FAR=1e-5）=93.7%。

1.介绍

…
这种旨在最小化嵌入向量之间的距离度量的直接方法可能在将代表性知识从教师转移到学生时受到限制[Shen等人，2021]，因为辨别力可能存在于奇异维度中，或者隐藏在嵌入维度之间的复杂相关性中。此外，由于教师的复杂性可能明显高于学生的复杂性，因此学生可能没有能力模仿教师的表现空间。这就是所谓的产能缺口问题[Mirzadeh等人，2019年]。因此，我们要求通过弥合教师和学生之间的差距来支持特征蒸馏过程的学习策略。特别是，我们研究了使用模型分类器辅助训练过程的能力。

模型的分类器捕获有关模型的表示空间结构和识别能力的基本信息。例如，在[Kang等人，2020]中，他们仅通过调整分类器权重来解决长尾识别中的不平衡问题。之前的工作在训练新模型以实现向后表示兼容性时使用了预训练模型的分类器权重[Shen等人，2021]， 或用于无监督的域适应[Liu等人，2021b]。受这些方法的启发，在本文中，我们建议通过在教师和学生模型之间共享分类器来增强表征知识的蒸馏。

具体来说，我们探索了两种在教师和学生之间分配分类器权重的方法。在第一种方法中，教师分类器通过将教师分类器连接到学生主干作为辅助头部（带有冻结参数）来约束学生表征学习。我们将这种方法命名为Teacher Head-KD，用TH-KD表示（图1（a））。在学生优化过程中，共享教师的分类边界可能有助于形成与教师相似的表示空间。在第二种方案中，使用一个临时学生分类器，通过用学生的头部训练教师来规范教师的表征学习。我们将这种方法命名为学生负责人KD，用SH-KD表示（图1（b））。通过这种方式，教师被迫学习可能更适合学生模型有限能力的功能。

图1：基于教师和学生之间分类器共享的两个方案的说明。（a） 第一种方法使用教师的分类器作为辅助头部，使用冻结权重，以帮助学生学习。（b） 在第二种方法中，我们将预先训练好的学生的分类器权重复制到教师的分类器（并冻结它们），以规范教师的表征学习。然后，使用知识提炼对一名新生进行培训。

通过广泛的分析，我们研究了这两种方法的性能，并将其与其他KD方法进行了比较。例如，我们表明TH-KD方案改进了表示蒸馏，使得嵌入空间中的类内变化更小，类间可分性更高。TH-KD和SH-KD方法在多个数据集上进行了各种任务的测试：CIFAR-100[Krizhevsky等人，2009]，斯坦福汽车[Krause等人，2013]FoodX-251[Kaur等人，2019]，以及在IJB-C数据集上进行人脸验证[Maze等人，2018]。具体来说，使用SHKD方案，我们在使用MobileFaceNet模型时，在IJB-C数据集上获得了最先进的结果：TAR@（FAR=1e5）=93.7%。

论文的主要贡献如下:
1.我们介绍了两种基于教师和学生模型之间分类器共享的表示提取新技术：TH-KD和SH-KD。这些技术易于实现，并与其他知识提取方法相辅相成。
2.我们分析了学生表征的质量，测量了通过TH-KD和SH-KD方法获得的与教师表征的相似性，并展示了它们在增强表征升华方面的有效性.
3.我们的方法在不同的设置、不同的数据集和不同的体系结构上实现了一致的准确性改进，包括在IJB-C数据集上获得最先进的人脸验证结果。

2.基于分类器共享的表征知识提取

在本节中，我们将介绍两种基于教师和学生模型之间的分类器共享的方法，以促进表示提取过程。在第一种方案中，教师的分类器用于约束学生的表征学习。在第二种方案中，使用学生分类器来规范教师的表征学习。
2.1问题表述
给定一个教师模型ft，我们的目标是在教师的指导下培养一个更小的学生模型fs。对于给定的输入样本x，我们分别用zt=ft（x；φt）和zs=fs（x；φs）表示教师和学生模型的表示（嵌入）向量，其中φt和φs分别是教师和学生模型的参数。
教师的分类器由gt（z）=Wtz+bt定义，学生的分类器由gs（z）=Wsz+bs定义。最后的预测是通过分别对教师和学生应用softmax激活和ps=h（gs（zs））。为了简单起见，我们分别用θt={Wt，bt}和θs={Ws，bs}表示教师和学生的分类器权重和偏差项。

对于给定的训练样本x和相应的真实标签向量y，用于训练学生模型的损失函数的一般形式可以写成：

其中，L CE（·）是交叉熵损失，H（·）是教师和学生模型的概率输出（logit）之间的KD距离函数，例如KL散度[Hinton et al.，2015]。术语D（·）是指应用于教师和学生模型表示向量的距离度量，如L2损失、余弦距离或对比损失（CRD loss）[Tian等人，2020]，其中α和β是控制每个损失项贡献的恒定超参数。

特别是，研究发现，表征蒸馏的L2损失对于人脸识别[Liu等人，2021a]和其他一般细粒度分类任务非常有用。L2损失由归一化嵌入向量的欧氏距离计算。我们称这种损失为L2E。

请注意，如果教师和学生的嵌入维度不同，我们会向架构的头部添加一个线性变换，以匹配它们的维度。虽然KD损失H（·）能够转移封装在教师软预测中的有价值知识，最小化表征损失D（·）来强制学生的嵌入空间与教师的嵌入空间对齐。因此，这些损失项是互补的，它们共同促进了教师向学生强有力的知识转移。
2.2 教师头分享TH-KD
我们的目标是利用由教师分类决策边界表示的判别信息来指导学生模型的优化过程。在该方案中，我们建议使用教师分类器作为辅助头部来训练学生模型。假设是教师分类器对给定学生嵌入输入的预测向量输出，我们将两个分类器计算的KD损失组合如下：
其中，α^TH是一个恒定的超参数，用于平衡两个分类头的损失。同样，分类损失由下式给出：

在推理时间，可以通过组合头部输出得到最终预测：p′s=（1）− α^TH）ps+α^TH p^THs，该方法名为TH-KD，如图1（a）所示。注意，对于α^TH=1，学生的头部只是教师的头部，其权重在训练期间是固定的。

考虑到教师的头部损失，鼓励学生模仿教师的特征空间，同时解决其高维依赖性。在第3.2节中，我们展示了TH-KD方案提高了表示质量和准确性。

2.3学生头分享 SH-KD
第二种方法旨在解决学生在蒸馏过程中能力有限的问题。在常规的知识提炼过程中，教师模型独立于学生提炼过程进行训练。通常情况下，教师模型的能力高于学生模型的能力，因此教师学习的特征可能不适用于学生培训。

为此，我们建议在考虑到学生能力有限的情况下，通过使用临时学生头部的权重初始化教师的分类器，并在培训期间固定它们来培训教师。这一过程可以被视为一种规范化机制，强制教师学习适合学生局限性的有用特征。
该方法可描述为三步训练程序：
步骤1：在有KD或没有KD的情况下训练一个学生网络，提供bockbone网络和带有参数{φ0s，θs}（同步3中的写法）的分类器头部。
步骤2：通过初始化并用固定了θs的分类器头来训练教师模型，以获得最终的教师模型的参数{φt，θs}。
步骤3：使用方程式（1）中的损失训练学生模型，并使用步骤2中获得的教师模型进行知识蒸馏，以产生最终的学生参数：

该方法名为SH-KD，如图1（b）所示。SH-KD方法提供了一个简单而有效的方案，可以在教师培训期间调整教师模型，以符合学生的能力，而代价是额外的培训迭代。在第3.2节中，我们表明，使用SH-KD进行培训可以提高学生和教师表达的相似性，并提高准确性。

表1:CIFAR100数据集上的测试精度（%）。RN缩写为ResNet[He等人，2015年]。我们遵循与CRD工作相同的协议[Tian等人，2020年]。附录中提供了这些方法的引文。

3.实验

在本节中，我们将报告我们在三个领域的主要结果：图像分类、细粒度分类和人脸验证。具体来说，我们报告了CIFAR100、斯坦福汽车、FoodX-251和IJB-C数据集的结果。

此外，我们还研究了所提出的方案THKD和SH-KD对表征蒸馏质量的影响。培训详情见附录。
3.1基准结果

3.2分析

4.结论

在本文中，我们探索了两种基于教师和学生之间分类器共享的表征提取技术。TH-KD方法与学生共享教师的分类器，以约束表征蒸馏过程。SH-KD允许在教师培训中共享学生的分类器，而代价是另一次培训迭代。大量的实验和分析证明了所提方案在不同领域和数据集上的有效性。我们发现TH-KD和SH-KD都加速了表征蒸馏过程。此外，TH-KD技术有助于提高由学生主干提取的嵌入的识别能力。最后，使用SH-KD和TH-KD进行训练可以增加教师和学生嵌入的相似性，从而提高学生的准确性。