TopKD革新蒸馏：Top-K缩放+解耦余弦损失，ResNet/ViT通杀，性能全面超越DKD/CRD

本文主要解决了什么问题

1. 1. 现有知识蒸馏方法主要关注特征层面的知识迁移，忽视了教师模型logit分布中嵌入的关键信息。
2. 2. 传统基于logit的蒸馏方法通常使用KL散度强制教师和学生模型logit之间的严格对齐，限制了知识迁移的丰富性，可能阻碍学生模型的泛化能力。
3. 3. 基于特征的蒸馏方法虽然能传递丰富的语义表示，但通常计算成本高、训练流程复杂且架构依赖性强，难以部署到不同模型中。

本文的核心创新是什么

1. 1. 识别并强调了教师logits中Top-K知识的重要性，发现这些Top-K logit包含更丰富和更具语义意义的监督信号，即使教师模型的top预测错误时也能提供有价值的语义关系。
2. 2. 提出了Top-scaled Knowledge Distillation（TopKD）框架，包含两个核心组件：

• • Top-K缩放模块（TSM）：自适应地放大最具信息量的logit，当教师Top-1预测错误时，对真实标签的logits应用更大的缩放因子以纠正教师的输出偏差。
• • Top-K解耦损失（TDL）：基于教师logits的符号和幅度将全局余弦相似度解耦为三个不同分量（Positive Top-K、Negative Top-K和Non-Top-K），提供有针对性和有效的监督。

• 3. 采用对比损失而非KL散度作为主要损失函数，更好地捕捉Top-K知识的结构信息，鼓励语义上对应的教师-学生对之间的对齐，同时抑制不匹配的对。
• 4. TopKD与架构无关，可以无缝集成到现有的KD方法中，无需引入额外模块或要求架构变更。

结果相较于以前的方法有哪些提升

1. 1. 在CIFAR-100、ImageNet、STL-10和Tiny-ImageNet等多个数据集上的大量实验表明，TopKD始终超越最先进的知识蒸馏方法，无论是基于logit还是基于特征的蒸馏方法。
2. 2. TopKD在蒸馏Vision Transformers时表现出显著的有效性，证明了其在不同网络架构中的通用性和鲁棒性。
3. 3. TopKD的模块化设计使其可以作为通用增强模块集成到现有蒸馏方法中，每次集成都能带来一致且显著的性能提升，展现出强大的兼容性、可迁移性和鲁棒性。
4. 4. TopKD提高了迁移表征的质量，在下游任务上展现出跨数据集和任务的强泛化能力，如在STL-10和Tiny-ImageNet上的迁移学习实验中表现优异。
5. 5. TopKD能够缓解"大模型并不总是更好的教师"的问题，即使教师-学生能力差距显著，也能根据教师能力表现出持续改进，更可靠地从能力更强的教师那里获益。

局限性总结

1. 1. TopKD需要选择合适的K值（Top-K的数量），论文实验表明K=3,5,10效果较好，但可能需要针对不同任务进行调整，无限增大K值并不会带来更多知识。
2. 2. TopKD需要设置合适的权重参数α和β（论文中默认设置为α=3和β=1），这些超参数可能需要针对不同场景进行调整，增加了调参的复杂性。
3. 3. 虽然TopKD与架构无关，但论文主要在图像分类任务上进行了验证，在其他类型的任务（如目标检测以外的任务或其他领域如NLP）上的效果尚未充分验证。

导读

知识蒸馏（KD）的最新进展主要侧重于特征层面的知识迁移，经常忽视教师模型logit分布中嵌入的关键信息。本文重新审视基于logit的知识蒸馏，揭示了一个尚未充分探索但至关重要的要素：Top-K知识。基于这一见解，作者提出了Top-scaled Knowledge Distillation（TopKD），这是一个简单、高效且与架构无关的框架，显著增强了基于logit的知识蒸馏。TopKD由两个主要组件构成：（1）Top-K缩放模块（TSM），自适应地放大最具有信息量的logit；（2）Top-K解耦损失（TDL），提供有针对性和有效的监督。值得注意的是，TopKD可以无缝集成到现有的KD方法中，无需引入额外模块或要求架构变更。在CIFAR-100、ImageNet、STL-10和Tiny-ImageNet上的大量实验表明，TopKD始终超越最先进的知识蒸馏方法。此外，TopKD在蒸馏Vision Transformers时表现出显著的有效性，证明了其在不同网络架构中的通用性。这些发现突显了logit在推动知识蒸馏方面的重要潜力。

1 引言

知识蒸馏（KD）由Hinton等人[16]提出，旨在通过校准其预测结果，将强大教师模型的泛化能力迁移到轻量级学生模型。自那以后，由于其实用性、高效性和通用性，KD得到了广泛应用。它已证明适用于各种网络架构，并且可以与其他压缩技术（如剪枝[11, 22]和量化[18, 47, 7]）无缝集成，以进一步减小模型尺寸。

尽管知识蒸馏（KD）取得了广泛的成功，目前最先进的方法[41, 14, 20, 39, 4, 2, 28]主要集中于从教师模型的中间层蒸馏特征表示。基于特征的知识蒸馏能够有效传递丰富的语义表示，但通常会产生高昂的计算成本、复杂的训练流程以及强烈的架构依赖性。相比之下，基于logit的知识蒸馏[16, 52, 51, 25]通过监督学生模型使用教师模型的输出分布，提供了一种更轻量级和灵活的替代方案，使其更容易部署到不同的模型中。

然而，基于logit的蒸馏方法通常表现出比基于特征的蒸馏方法更差的性能，这主要是因为在最终输出层存在信息 Bottleneck ，其中蒸馏监督通常仅限于类别概率。这一局限性因Kullback-Leibler散度（KL-Div）的广泛使用而进一步加剧，KL-Div强制要求学生模型和教师模型的logit之间严格对齐。尽管在某些情况下这种方法有效，但这种僵化的对齐往往会限制知识迁移的丰富性，并可能阻碍学生模型的泛化能力。认识到这一局限性，近期研究[52, 51, 53, 29]通过放宽或完全移除KL-Div约束，探索了替代性公式，证明重新思考基于logit的监督可以取得具有竞争力甚至更优的结果。这些进展突显了基于logit蒸馏的未开发潜力，激励了对其结构和学习动态的更深入研究。

基于这一见解，作者重新审视了教师logits的分布，并识别出未被充分探索的Top-K知识，这种知识能够捕捉到特别具有信息量的监督信号。为了利用这一特性，作者提出了Top-scaled Knowledge Distillation（TopKD），这是一个在蒸馏过程中明确强调教师Top-K知识的框架。

作者的贡献可以总结为：

• • 作者通过分析传统基于logit的蒸馏方法的局限性，并强调教师logits中Top-K知识先前被忽视的重要性，引入了一种新的视角。
• • 为有效利用这一见解，作者开发了Top-scaled Knowledge Distillation（TopKD），该模型包含一个轻量级的Top-K Scaling Module（TSM）和一个Top-K Decoupled Loss（TDL）。这些组件与架构无关，可以无缝集成到现有的蒸馏方法中，并持续提升其性能。此外，作者避免使用KL-Div作为主要损失函数，而是采用对比损失来更好地捕捉Top-K知识结构信息。
• • 在多个数据集上的实验结果表明，TopKD始终优于基于logit和基于特征的蒸馏方法，实现了当前最佳性能。

2 相关工作

2.1 知识蒸馏

知识蒸馏（KD）将知识从大型教师模型迁移到紧凑型学生模型，已成为模型压缩和性能提升的广泛采用方法。当前的KD方法主要分为基于logit的蒸馏和基于特征的蒸馏。基于logit的蒸馏[16, 30, 52, 51, 53]通过使教师模型和学生模型的输出logit对齐来迁移知识，通常使用Kullback-Leibler散度作为目标。尽管简单且计算高效，但这些方法往往无法充分捕捉教师模型的内部表征结构。基于特征的蒸馏[41, 14, 20, 4, 2, 28]则指导学生模型模仿教师模型的中间特征表征，提供更丰富的监督，但代价是架构约束和计算开销的增加。此外，一些方法将对比学习融入蒸馏过程[39, 3, 33, 1]，旨在通过鼓励学生模型将表征拉近教师模型（正样本对）而推远无关样本（负样本对）来增强其判别能力。

2.2 对比学习

对比学习被广泛认可为一种有效的表征学习方法，尤其在自监督和半监督场景中。其核心目标是通过对比正样本对（例如同一图像的不同增强视图）与负样本对（例如不同类别的实例）来学习特征嵌入。SimCLR [5]、MoCo [12] 和 InfoNCE [32] 等重要框架表明，强制实例级区分能够获得鲁棒且可迁移的特征表示。在知识蒸馏的背景下，对比损失为编码关系知识提供了一种有效机制，补充了传统的实例级对齐，并进一步增强了学生模型的泛化能力。

3 方法论

在本节中，作者介绍了作者的Top-scaled Knowledge Distillation（TopKD）方法，该方法旨在通过显式利用教师模型的输出分布来增强基于logit的知识蒸馏。通过对logit分布的深入分析，作者发现Top-K logit包含了更丰富和更具语义意义的监督信号。然而，传统方法通常依赖KL散度来强制整个输出分布的对齐，从而忽略了Top-K logit中固有的语义相关性。为了解决这一局限性，作者采用了一种对比损失，它能更有效地捕捉类别预测之间的语义关系信息。此外，作者提出了两个核心组件：（1）Top-K缩放模块（TSM），该模块在训练过程中自适应地强调最具信息的logit；（2）Top-K解耦损失（TDL），该损失在受Top-K logit影响的嵌入空间中指导对齐。

3.1 回顾教师日志

作者重新审视了教师模型在训练过程中的logits分布，并观察到其top-1预测并不总是准确的。如图2所示，对于一个被 Token 为海豹的输入，教师模型预测为水獭，这是一个视觉上相似但错误的类别。这些错误分类并非完全随机，而是往往反映了教师模型学习到的语义层次结构。进一步检查Top-K预测（例如，K=10）可以发现，许多高排名的类别与真实标签在语义上相关，表明Top-K logits编码了丰富的关系知识。然而，正如[29]所指出的，基于KL散度的传统蒸馏方法存在根本性限制：它仅鼓励教师模型和学生模型输出概率的对齐，而未能捕捉到类间的潜在结构。为了有效利用嵌入在教师模型logits中的跨类别信息，作者引入了一种对比学习方法，以对齐教师模型和学生模型表示之间的关系模式。即使教师模型的top预测是错误的，如图2所示，这种方法也能让学生模型提取有意义的语义关系，从而增强其泛化能力。

为实现此目标，作者采用一种对比损失函数，该函数鼓励语义上对应的教师-学生对之间的对齐，同时抑制不匹配的对。其目标是拉近相同样本的logits，推远不同样本的logits。令

和

分别表示每个批次样本的学生logits和教师logits，对比损失函数定义如下：

其中

是交叉熵（CE）损失，

表示批次大小，

是类别数，

是一个控制相似度分数尖锐度的温度参数，遵循先前工作中的常见做法 [5, 12]。这里，

，包含从 0 到

的索引。这种公式鼓励最大化匹配对在

对角线上的相似度，同时最小化非匹配对在非对角线上的相似度。

3.2 Top-K缩放模块（TSM）

为了更好地利用对比学习中嵌入在教师logits中的跨类别关系，作者引入了一种专门针对Top-K logits的缩放策略。具体而言，作者增加Top-K预测的值以强调其影响。此外，当教师的Top-1预测错误时，作者对真实标签的logits应用更大的缩放因子以纠正教师的输出偏差。形式上，作者对每个样本的logits进行如下重缩放：

其中

表示 Top-K 预测的索引，

是真实标签，

是一个与排名相关的缩放因子，而

是一个与 Top-K 和 Non-Top-K logit 之间的平均差异成比例的偏差项。TSM 增强了语义相关类别的相对重要性，同时降低噪声或误导性 logit 的权重，从而更有效地保留教师学习到的语义结构。由于篇幅限制，计算

和

的详细步骤在附录中的算法??中提供。

3.3 Top-K 解耦损失 (TDL)

方程1中定义的对比损失通过匹配每个学生的logits与其对应的教师的logits，沿图3中维度D i m = 0的批次维度，强制实例级对齐。然而，仅此目标无法捕捉单个预测中更细粒度的结构差异。为克服这一局限性，作者集成了余弦相似度损失以保持实例内语义一致性（图3中维度D i m = 1）。对于来自同一样本的两个向量z s 和z t（沿特征维度j），余弦相似度计算如下：

如图3(c)所示，公式3作为公式1的自然补充，优化嵌入空间中的优化方向，以更好地支持实例级对齐。这两个目标共同确保语义一致性和方向指导，从而产生更具判别性的表示。

复合余弦相似度虽然提供了学生和教师预测之间总体的一致性度量，但它并未衡量 logits 中不同类别的语义重要性。在实践中，一小部分高置信度条目（通常是 Top-K logits）编码了最自信和最有信息量的内容，而其余条目可能存在噪声或模糊性。将所有 logits 统一处理可能会稀释有意义的语义信号，阻碍有效的知识迁移，正如 [52] 中所观察到的。与采用 KL-Div 变体的 [52] 不同，TopKD利用余弦相似度。这两种方法反映了根本不同的对齐原则：KL-Div 强制精确值匹配，而余弦相似度强调嵌入空间中的方向一致性。

为了更好地利用logits，作者基于教师logits的符号和幅度将全局余弦相似度解耦为三个不同的分量：(1) Positive Top-K，对应具有最高正值的方向；(2) Negative Top-K，涵盖最负值的方向；以及 (3) Non-Top-K，代表剩余的、置信度较低的方向。Top-K解耦损失如下：

其中

表示批次大小，

表示方程 3 中定义的余弦相似度，

表示应用方程 2 中的缩放后教师模型的 logits。仅对教师模型的 logits 进行缩放，学生模型的 logits 保持不变。作者在表6 中对

和

进行了详细的消融实验。

TopKD损失。最终的TopKD损失函数定义为：

在TopKD中，整体蒸馏目标由两个互补的组成部分构成：一个对比损失，用于强制学生和教师预测在实例 Level 上保持一致，以及Top-

解耦损失（TDL），用于捕获logits中的细粒度结构一致性。

4 实验

数据集。作者的实验主要集中于图像分类。CIFAR-100 [21] 是一个广泛使用的基准数据集，包含 60,000 张彩色图像

，分为 50,000 张训练图像和 10,000 张验证图像，涵盖 100 个类别。ImageNet [35] 是一个大规模数据集，包含 1,280 万张训练图像和 50,000 张验证图像，分为 1,000 个类别。STL-10 [9] 和 Tiny-ImageNet [31] 作为中等规模的基准数据集，用于评估泛化能力。STL-10 提供 10 个高分辨率类别

和一个大型无标签集，用于无监督学习，而 Tiny-ImageNet 包含 200 个低分辨率类别

，旨在测试每个类别数据量有限时的泛化能力。

Baseline 方法。近年来，大多数最先进的方法都依赖于基于特征的蒸馏。因此，作者首先将TopKD与具有代表性的基于特征的方法进行比较，包括CRD [39]、ReviewKD [4]、SimKD [2]、CAT-KD [10]和FCFD [28]。由于TopKD是一种基于logit的方法，作者也将其与经典KD [16]以及最近的基于logit的蒸馏方法进行比较，包括CTKD [25]、DKD [52]、DOT [51]、LSKD [37]和WTTM [53]。作者报告了LSKD [37]，并选择

[52, 37]作为代表性设置，因为MLKD

[19, 37]使用了更多轮次和不同的初始学习率。作者选择了ResNets [13]、WRNs [49]、VGGs [36]、MobileNets [17]和ShuffleNets [50]作为作者在CIFAR-100 [21]和ImageNet [35]数据集上的分类实验的 Baseline 模型。对于视觉Transformer Baseline ，包括DeiT-Ti [6]、T2T-ViT7 [48]、PiT-Ti [15]和PVT-Ti [45]。

训练细节。作者采用了先前工作中的实验设置[39, 52, 4]。所有实验均使用SGD[38]优化器，批大小为64，在CIFAR-100上训练240个epoch，在ImageNet上训练100个epoch。对于STL-10[9]和Tiny-ImageNet[31]，作者通过冻结中间层并仅训练分类器进行迁移学习，以评估表征的可迁移性。更多实现细节请参见附录。

4.1 主要结果

CIFAR-100图像分类。作者在CIFAR-100上评估了各种知识蒸馏方法，使用了异构的教师-学生架构（表1）和同构架构（表2）。方法按蒸馏类型分类：基于特征和基于logit。TopKD始终表现最佳或第二好，证明了其有效性。此外，在采用MLKD+LS[19, 37]相同设置的情况下，TopKD实现了更高的准确率。详细结果见附录（表??）。

ImageNet图像分类。作者在表3中报告了Top-1和Top-5准确率。在同构设置（ResNet-34 / ResNet-18）中，教师模型和学生模型表现出极小的性能差异。在异构设置（ResNet-50 / MobileNetV1）中，教师模型显著优于学生模型。在两种设置下，TopKD始终在大规模ImageNet数据集上优于大多数最先进的蒸馏方法。

4.2 消融实验

为评估TopKD核心组件的贡献，作者对Top-K缩放模块（TSM）和Top-K解耦损失（TDL）进行了消融研究。如表4所示，这两个模块单独使用均能提升 Baseline 性能，而它们的组合则带来了额外的增益，表明它们为整体性能提供了互补的优势。作者进一步研究了TopKD对Top-K选择中K值选择的敏感性。具体而言，作者在表5中评估了K值变化对整体性能的影响。在该实验中，作者使用ResNet

[13]作为教师网络，

[13]作为学生网络。

结果表明，当

，5，10时，能够捕获更丰富的知识。无限增大K的值并不会带来更多知识，反而只会使Top知识相对于其他知识变得不那么突出。在表6中，作者观察到在两种设置下，

，5分别取得了最佳结果，且在

值上性能差距相似，这证实了Top-K知识最为关键。因此，作者采用

和

作为所有实验的默认设置。

4.3 扩展

为评估TopKD，作者从五个关键视角进行分析。首先，作者将即插即用的TSM和TDL模块集成到各种蒸馏框架中，持续提升性能，展示了TopKD的泛化能力和兼容性。其次，作者测试蒸馏模型在下游数据集上的可迁移性，评估TopKD超越原始领域的泛化能力。第三，作者提供可视化结果，清晰验证TopKD的优越性。第四，作者探讨大模型并非总是更好的教师的问题，并表明TopKD能够缓解这一问题，即使教师-学生能力差距显著。最后，作者将评估扩展到视觉Transformer模型，确认TopKD的性能提升。

即插即用能力。TopKD的一个关键优势在于其模块化设计，这使得它可以无缝集成到现有的知识蒸馏方法中。Top-K缩放模块（TSM）和Top-K解耦损失（TDL）都是轻量级且与架构无关的，无需对 Backbone 模型进行修改。为了评估这一能力，作者将TSM和TDL集成到几种具有代表性的蒸馏方法中。如图4所示，每次集成都导致原始方法性能得到一致且显著的提升。这些结果表明，TSM和TDL可以作为通用增强模块，展现出在蒸馏方法之间强大的兼容性、可迁移性和鲁棒性。

表征的可迁移性。知识蒸馏可以促进有用表征向与原始训练域不同的下游任务或数据集的迁移。为了评估通过TopKD学习到的特征的可迁移性，作者采用ResNet8x4和MobileNetV2作为学生模型，分别从ResNet32x4和ResNet50蒸馏而来，或用于在CIFAR-100 [21]上从头训练以进行比较。对于STL-10 [9]和Tiny-ImageNet [31]基准（均缩放到32×32），作者冻结模型编码器直到倒数第二层，并在顶部训练一个线性分类器以分别执行10类和200类分类。如表7所示，TopKD始终提高了迁移表征的质量，展示了跨数据集和任务的强泛化能力。

可视化。作者从两个角度对CIFAR-100 [21]提供可视化结果，以进一步说明TopKD的有效性。这些可视化评估了学生模型中特征的可区分性以及类别结构关系的保留情况。首先，图5中的t-SNE [42]图，使用ResNet32x4作为教师模型，ResNet8x4作为学生模型，揭示通过TopKD学习到的特征表示比其他蒸馏方法产生的特征表示更加紧凑和具有类别区分性。其次，作者在图6中可视化了教师模型和学生模型的logits相关矩阵之间的归一化绝对差，其中使用WRN-40-2作为教师模型，ResNet8x4作为学生模型。较小的差异表明TopKD促进了教师模型和学生模型输出之间的结构对齐。

更大规模的模型并不总是更好的教师。尽管更强的教师模型被期望提供更好的监督，但它们往往难以有效地将知识迁移给较小的学生，甚至与较弱的教师相比表现更差。先前工作将此问题归因于教师模型和学生模型之间较大的容量差距[8, 43, 52, 37]。作者认为大多数现有方法依赖于直接匹配教师的logits，这忽略了其丰富的语义结构，导致知识迁移效率低下。由于教师logits包含High-Level类别关系，直接模仿它们会阻碍学生提取有意义信息的能力。作者提出的TopKD通过关注Top-K知识中的信息结构来缓解这一限制，如表8中的结果所示。

值得注意的是，其他方法在教师能力更强时显示出学生表现波动甚至下降的情况，而TopKD则根据教师能力表现出持续改进。这一观察表明TopKD对教师能力更为敏感，并且能够更可靠地从能力更强的教师那里获益。

视觉Transformer与检测。为了进一步验证作者TopKD的鲁棒性和通用性，作者在一系列视觉Transformer架构[40, 23, 24, 46, 6, 48, 15, 45]上进行了额外的实验，并将作者的评估扩展到使用MS-COCO数据集[27]的目标检测任务。TopKD在不同基于Transformer的 Backbone 网络和任务中始终表现出色，进一步证明了其广泛的适用性和有效性（详情见附录）。

参考

[1]. TopKD: Top-scaled Knowledge Distillation