轻量化模型开发必备 | 中心化Kernel对齐实现大模型知识向小模型迁移

634次阅读

点击下方卡片，关注「AI视界引擎」公众号

知识蒸馏作为一种有效的方法，可以弥合大规模模型和轻量级模型之间的表示差异。常用的方法包括利用适当的度量标准来最小化教师模型提取的知识和学生模型学习的知识的差异或距离。中心核对齐（CKA）被广泛用于衡量表示相似性，并已在多个知识蒸馏方法中应用。然而，这些方法都很复杂，无法揭示CKA的本质，因此无法回答如何利用CKA实现简单有效的蒸馏的问题。

本文首先从一个理论角度说明CKA的有效性，将CKA分离到最大均方差异（MMD）的上界和常数项。作者提出了一种新颖的关系中心核对齐（RCKA）框架，该框架实际上将CKA与MMD建立联系。此外，作者根据每个任务的特性动态地定制CKA的应用，其计算资源较少，但性能与以前的方法相当。

在CIFAR-100，ImageNet-1k和MS-COCO上的广泛实验表明，RCKA在几乎所有教师-学生对图像分类和目标检测上实现了最先进的效果，验证了RCKA的有效性。

1 Introduction

大量工作已经完成，将大规模模型压缩成轻量级模型。代表性的方法包括网络剪枝，模型量化，神经架构搜索[23]和知识蒸馏（KD）。其中，KD最近由于其有效性和广泛的应用而成为最热门的话题之一。特别是，KD的核心思想是通过蒸馏学习表示，将大规模和高性能模型的获得表示转移到轻量级模型中，实现精确可靠的知识转移。研究行人普遍认为，从教师到学生的知识蒸馏有两种主流方法。第一种方法是基于对数预测的蒸馏，通过利用适当的度量标准最小化教师和学生的概率预测（响应）得分。第二种方法是基于特征的蒸馏，研究模型中的中间表示知识，以进一步提高蒸馏性能。

在这些方法中，度量标准的设计在知识转移中至关重要，并吸引了学者的关注。Kornblith等人[19]提出了中心核对齐（CKA）来定量理解神经网络之间的表示。CKA不仅关注模型预测，而且强调模型内的高阶特征表示，提供全面丰富的知识转移。

近年来，[20, 18]项研究引入了CKA来定量缩小教师模型和学生模型之间学习表示的差距。无疑，这些方法已经取得了显著的成功。然而，它们的设计过于复杂，需要大量的计算资源，使得实现细粒度知识转移具有挑战性，并导致低可扩展性。此外，这些方法未能揭示CKA的本质，缺乏对CKA在知识蒸馏中的深度分析。CKA有效的原因尚未被探索。因此，作者专注于CKA的理论分析，并重新思考一个既保证简单性又保证有效性的更合理架构设计，同时具有良好的泛化能力。

在本文中，作者从一个新的角度说明了CKA的有效性，其中CKA被视为最大均方差异（MMD）的上界和常数项，具体而言。从中，作者提出了一种关系中心核对齐（RCKA）框架，该框架实际上将CKA与MMD建立联系。此外，作者在实例级任务上动态定制CKA的应用，并引入基于块的中心核对齐（PCKA），与以前的方法相比，计算资源更少，但性能具有竞争力。RCKA不仅可以直接应用于基于对数预测的蒸馏，还可以应用于基于特征的蒸馏，具有卓越的可扩展性和扩展性。作者利用CKA计算类别之间的和高阶表示信息，这更好地激励了减轻教师和学生模型之间性能差距。

为了验证作者方法的有效性，作者在图像分类（CIFAR-100 和 ImageNet-1k）和目标检测（MS-COCO）任务上进行了广泛的实验。结果表明，RCKA在几乎所有的定量比较实验中实现了最先进（SOTA）的性能，并进行了公平的比较。此外，遵循作者的处理架构，在目标检测任务中，先前的蒸馏方法的表现得到了进一步提高。

作者的贡献可以总结如下：

作者从知识蒸馏的全新角度重新思考CKA，提供了CKA在知识蒸馏中有效的原因的理论依据。
作者提出了一种关系中心核对齐（RCKA）框架，用于构建CKA与MMD之间的关系，该框架具有较少的计算资源，但性能与先前的方法相当，这验证了作者对理论分析的正确性。
作者进一步动态地定制CKA在实例级任务上的应用，并提出了一种基于块的中心核对齐（PCKA）架构用于目标检测中的知识蒸馏，这进一步提高了先前蒸馏方法的表现。
作者进行了大量的消融研究来验证RCKA的有效性，该方法在各种视觉任务上实现了最先进（SOTA）的表现。此外，作者可视化了CKA的特征信息并发现了其中新的模式。

3 Methodology

在这一部分，作者首先回顾知识蒸馏的范式，然后引入中心核对齐（CKA）的公式。具体而言，作者推导出CKA与最大均方差异（MMD）之间关系的公式，其中CKA可以分离为MMD的上界和常数项。根据上述推导，作者概述了本文的方法。作者在图像分类和目标检测中应用所提出的算法，动态地定制CKA以适应每个任务。

The Paradigm of Knowledge Distillation

现有的知识蒸馏方法可以分为两组。特别是，基于对数的蒸馏方法通过将教师模型和学生的软目标对齐来缩小它们之间的差距，这可以表示为以下损失项：

其中，和分别是从学生和教师模型中得到的对数，是Softmax函数，将logits转换为类别概率，是非负温度超参数，用于缩放预测分布的平滑性。

具体而言，作者有。是一个损失函数，用于捕获对数差异分布，例如Kullback-Leibler散度。和分别表示学生在和教师中的转换函数，通常是指Vanilla KD中的身份映射。

类似地，基于特征的知识蒸馏方法，其目标是模拟教师和学生之间的特征表示，也被表示为一个损失项：

其中，和分别表示学生和教师特征图，转换模块和将它们的维度对齐。计算两个特征图之间的距离，例如-或范数。因此，知识蒸馏方法可以表示为一个通用的范式。最终的损失是交叉熵损失、对数预测损失和对特征预测损失的加权和：

其中，和是超参数，用于控制这三个损失之间的权衡。

Distilling with the Upper Bound

中心核对齐（CKA）被提出作为一种鲁棒的方法来衡量神经网络之间的表示相似性。作者首先证明CKA衡量教师和学生之间的Gram矩阵余弦相似度。

定理1：令和是矩阵。CKA相似度等于XX’和YY’的余弦相似度，分别表示和的Gram矩阵。换句话说，

在这里，操作表示将矩阵 Reshape 为向量。

然后作者推导出CKA与MMD之间关系的公式，其中CKA可以被视为MMD的上界和常数项。

定理2：最大化CKA相似性等价于最小化MMD距离的上界。

其中，不等式由Jesen不等式给出。

根据Jesen不等式，CKA可以分解为MMD上界的常数项。第一个项对应于使用RKHS核最小化MMD距离的上界。相反，后一个常数项作为权重正则化器，增强MMD的影响，促进同一批次的特征相似性，不仅是在同一类别内的实例，还包括不同类别的实例。一方面，优化MMD的上界，具有额外的强约束，允许它更快、更稳定地收敛到最优解。另一方面，后一个项作为权重缩放机制，有效地避免了由于MMD值过小而导致的优化挑战，从而产生小梯度的问题。

根据上述推导，作者成功地从最大化CKA的目标转换为最小化MMD的上界，这使得RCKA更加直观和简洁。在这些发现的基础上，作者提出了RCKA，这些方法比以前的方法更有效。

Relation-based Centered Kernel Alignment

如图1所示，作者在图像分类中提出了一种关系中心核对齐（RCKA）框架。在这个框架中，作者利用CKA作为损失函数，以确保中心相似矩阵被蒸馏，而不是强制学生模仿教师相似矩阵的不同尺度。这是非常重要的，因为模型的判别能力取决于其特征的分布，而不是其尺度，这对于分类是无关紧要的。

轻量化模型开发必备 | 中心化Kernel对齐实现大模型知识向小模型迁移

假设作者有一个大规模的教师模型和一个轻量级的学生的模型。教师模型的层的激活图可以表示为，而学生模型的层的激活图可以表示为。，和分别表示教师的通道、高度和宽度，而，和分别表示学生的。批处理大小用表示。教师和学生的对数表示为和，其中和（或）分别表示样本的数量和这个样本属于哪个类别的概率。因此，RCKA的公式类似于方程3，可以表示为：

其中，和是超参数，用于控制特征损失和对数损失之间的权衡。

，和可以表示为：

在方程6中，指的是变换模块。

与先前的方法相比，RCKA具有优越的可扩展性和可扩展性，并且可以直接应用于特征和对数蒸馏。作者计算Gram矩阵来收集高阶的类内和类间表示，鼓励学生学习更有用的知识。此外，作者在附录C.2中提供了CKA和MMD之间的关系，以更好地展示作者方法的理论支持。

由于CKA的值范围为，在训练过程的初期，比所有CKA损失更重要，用于驱动学生的优化，这有助于学生避免匹配极其复杂的表示。

Patch-based Centered Kernel Alignment

在这个子节中，作者进一步将提出的RCKA方法适应到实例级任务，如目标检测。然而，直接将RCKA应用于实例级任务可能会降低性能，因为这些任务通常使用较小的mini-batch大小（例如每个GPU训练2或4个mini-batch），导致gram矩阵无法收集足够的信息。此外，增加mini-batch大小需要大量的计算资源，这在实践中是不可行的。因此，作者动态地自定义作者的RCKA方法以适应目标检测。

近年来，[26, 19]项研究发现在目标检测中，蒸馏中间层的表示比蒸馏对数更有效。因此，作者调整RCKA只针对中间层。然而，作者在分类任务中仍然遵循作者的核心思想，通过使用CKA计算不同实例之间的相似性。因此，作者将图像特征图分成几个块，并计算不同块之间的相似性。

作者重新设计的算法在图2中得到说明。在这个框架中，作者首先用大小为的块裁剪教师和学生模型的特征图，然后将特征图转换为每个块之间的gram矩阵。最后，作者计算损失并从维度中获取平均值。在这里，和分别表示分别沿着高度和宽度切割的块的数量。因此，基于块的中心核对齐（PCKA）损失可以表示为：

其中，和分别表示学生的块数和教师的块数。通常，。表示损失权重因子。

4 Experiments

作者在图像分类和目标检测的基准测试上进行了广泛的实验。图像分类数据集包括CIFAR-100和ImageNet-1k，目标检测数据集包括MS-COCO。

此外，作者为提出的各种方法进行了各种消融和分析。有关这些数据集的更多详细信息可以在附录A中找到。作者在CIFAR-100上使用批量大小为128和初始学习率0.1的SGD优化器，并遵循[12]中ResNet34-ResNet18对和ResNet50-MobileNet对在ImageNet-1k上的设置。其他分类和检测任务的其他设置可以在附录B中找到。作者的代码将公开供可重复性。

Image Classification

在CIFAR-100分类上，作者在9组学生-教师对中比较了最先进的基于特征的和基于对数的蒸馏算法。其中，6组对教师和学生具有相同的结构，其余的对具有不同的架构。

结果如表1所示。RCKA在4组学生-教师对中超过了其他所有算法，并在其余的对上达到了可比的性能，同时需要的计算资源和时间消耗比SOTA方法DPK和ReviewKD要少得多。计算成本的比较可以在附录D中找到。

在ImageNet-1k分类上。作者还在大规模的ImageNet上进行实验，以评估RCKA。作者的RCKA在其他算法上实现了可比较的结果，甚至超过了它们，如表2所示。作者发现，随着类别和实例的增加，学生模仿教师的高阶分布变得更加困难。

此外，在附录8中探索了在ImageNet-1k上的ViT模型特征蒸馏。值得注意的是，RCKA超过了其他方法，这意味着RCKA具有良好的可扩展性和良好的性能。

轻量化模型开发必备 | 中心化Kernel对齐实现大模型知识向小模型迁移

Object Detection

在MS-COCO检测上，作者进行了三种不同检测器的比较实验，即两阶段检测器、单阶段检测器和 Anchor-Free 检测器。如表3所示，PCKA在所有三种指标上都超过了其他方法，通过将高阶的块级表示进行对齐。作者认为，在低阶特征图上模仿学生和教师特征图也可能改善PCKA的性能，通过在早期阶段模仿低阶表示，然后逐渐学习高阶和复杂的表示。

因此，作者遵循[14]的方法，通过添加辅助模仿损失，即通过卷积层将学生特征图从教师特征图中翻译过来，并使用进行监督，将其添加到检测蒸馏任务中。最终，作者从表3中得出结论，基于PCKA的模仿损失在Cascade RCNN-X101-Cascade RCNN-R50和RetinaNet-X101-RetinaNet-R50对上取得了最佳性能。作者还进行了其他四个架构的实验，如表4和表9所示。这些结果进一步验证了作者提出的方法的 effectiveness。

轻量化模型开发必备 | 中心化Kernel对齐实现大模型知识向小模型迁移

Ablations and Visualizations

作者在三个方面进行了消融研究：

超参数的影响
所提出模块的有效性
在训练过程中未探索到的现象。

超参数消融研究。如附录中表12，表11和表13所示，作者在mini-batch大小，的损失缩放因子和蒸馏的中间层数量方面进行了超参数消融研究。作者发现局部最优值是mini-batch大小，的损失缩放因子和层蒸馏。

轻量化模型开发必备 | 中心化Kernel对齐实现大模型知识向小模型迁移

MMD的上界。在定理2中，作者推导出CKA与MMD之间的关系，其中CKA是MMD的上界和常数项。为了验证这一点，作者进行了实验，如表6所示。作者注意到CKA，即MMD的上界，具有额外的更强约束。因此，与MMD相比，CKA更快、更稳定地收敛到最优解。

平均维度的选择。在PCKA框架中，作者将教师和学生模型的激活值裁剪成的形状。作者还进行了不同维度的实验，如表5所示。作者发现，在通道维上平均是优化。

划分激活图的蒸馏。作者探索将激活图划分成块的有效性。如表7所示，几种标准的蒸馏方法[13, 14, 15]都通过划分块得到了良好的效果，验证了划分块的有效性。由于块中的表示分布较小，因此可以更容易地对学生和教师进行对齐。因此，提出的PCKA架构极大地提高了以前的方法。

提供了一些可视化来展示RCKA在对数 Level 上成功地缩小了教师-学生之间的差距。特别地，作者在附录E中可视化了6组教师-学生对的对数相似性。作者发现RCKA显著地提高了对数相似性。
进一步可视化了不同检测器的训练过程以及块对RetinaNet-X101-RetinaNet-R50对的影响。结果如附录中的表4和表3所示。
首先可视化了所提出方法的混淆矩阵（见图6），然后分别可视化了训练时有无块平均不同维度的标注图像（见图7）。这些图像揭示了作者方法可以收集不同类别之间的相似性，并显示出作者在目标检测任务上方法的有效性。

轻量化模型开发必备 | 中心化Kernel对齐实现大模型知识向小模型迁移

5 Discussion

在图像分类任务中，作者应用PCKA。它也很好地超过了具有相同蒸馏类型的其他方法，如表14所示。然而，对于具有不同架构的教师-学生对，PCKA的表现非常糟糕。因为不同架构的激活包含更多不同和有害的表示，这使得向学生传递知识变得困难。

在通道上平均，提升性能。表7中的结果揭示了有趣的现象，即在将激活值裁剪成块后，通过对通道维度的损失进行平均，之前蒸馏方法的表现得到了提升。通过将表示值裁剪成块，可以使得教师和学生的对齐更容易，且块中的表示值分布更小。此外，将表示值裁剪成块符合作者提出的分类框架中的思想，因此PCKA可以计算块中的类间相似性和类内相似性。此外，由于余弦相似性相对于基于距离的损失具有优越性，以及Gram矩阵收集的高阶分布表示，PCKA超过了DIST和AT。

位置信息损失。在PCKA中，作者将教师和学生模型的激活值切割成块，然后将其 flatten 成一个向量。尽管这个操作会损害原始的位置信息，但性能不会下降。作者认为，CKA确保优化的重点在于分布的形状，而不是Gram矩阵中的原始值，这很重要，因为模型的判别能力取决于其特征的分布，而不是其规模。此外，在开始时，PCKA带来的效果较小，与CE损失相比。因此，CE损失促使学生模型的优化稳步进行，而从某个时刻开始，PCKA驱动学生模型对复杂和高阶表示进行对齐，提高泛化能力。