西安交大提出DeepKD | 双层解耦+自适应去噪破解知识蒸馏困境,刷新ImageNet等多项SOTA

大模型向量数据库机器学习

点击下方卡片,关注「集智书童」公众号

点击加入👉「集智书童」交流群

picture.image

picture.image

picture.image

picture.image

picture.image

picture.image

精简阅读版本

本文主要解决了什么问题

目标类和非目标类知识流之间的固有冲突 :现有方法未能有效解决目标类和非目标类知识流在优化过程中的冲突,导致次优的优化轨迹。

低置信度暗知识引入的噪声信号 :非目标类中的低置信度logits会引入噪声,阻碍有效的知识迁移。

动量机制缺乏理论基础 :现有方法对不同知识组件的梯度动量分配缺乏严格的理论支持,仅依赖经验观察。

本文的核心创新是什么

DeepKD框架 :提出了一种集成了双层解耦和自适应去噪的新型训练框架,通过深度解耦优化和动态top-

Mask机制提升知识蒸馏效果。 2. 2. 基于GSNR的动量分配策略 :通过对梯度信噪比(GSNR)的理论分析,设计了独立的动量更新器,防止面向任务和蒸馏梯度之间的相互干扰。

**动态top-

Mask机制** :逐步过滤低置信度logits,遵循课程学习原则,有效净化暗知识并保留语义相关性。

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

CIFAR-100上的显著改进

  • • 在同类架构下,DeepKD实现了+0.61%到+3.70%的准确率提升。
  • • 使用动态top-

Mask后,进一步提升了高达+1.86%,达到79.15%的准确率。

ImageNet-1K上的卓越表现

  • • 对于ResNet50和MobileNet-V1,DeepKD将top-1准确率提升了+4.15%到+74.65%。
  • • CRLD+Ours(w. top-k)在ResNet34→ResNet18上达到了74.23%(+1.86%),在RegNetY-16GF→Deit-Tiny上达到了75.75%(+1.90%)。

MS-COCO目标检测任务

  • • 动态top-

Mask机制使AP提升了+1.93%,达到32.16%。

  • • 在

上,DeepKD显著优于其他方法,最高达到36.59%。

局限性与未来工作

局限性

  • • 当前工作主要集中在基于logit的蒸馏,尚未扩展到特征蒸馏场景。
  • • 动态top-

Mask机制虽然有效,但其参数设置仍需经验调整。

未来工作

  • • 探索将DeepKD应用于更复杂的场景,如多教师蒸馏和跨模态知识迁移。
  • • 进一步优化动态top-

Mask策略,使其适应更多模型架构和数据集。

  • • 研究SNR驱动的动量解耦机制在特征蒸馏中的应用潜力。
深入阅读版本

知识蒸馏领域的最新进展强调了分离不同知识组件的重要性。现有方法利用动量机制来分离面向任务和蒸馏梯度,但它们忽略了目标类和非目标类知识流之间的固有冲突。

此外,非目标类中的低置信度暗知识引入了噪声信号,阻碍了有效的知识迁移。为解决这些局限性,作者提出了DeepKD,这是一个集成了双层解耦和自适应去噪的新型训练框架。首先,通过对面向任务和非面向任务知识蒸馏中梯度信噪比(GSNR)特性的理论分析,作者为每个组件设计了独立的动量更新器,以防止相互干扰。作者观察到,面向任务梯度(TOG)、目标类梯度(TCG)和非目标类梯度(NCG)的最佳动量系数应与其GSNR呈正相关。其次,作者引入了一种动态顶部

Mask (DTM)机制,该机制从较小的初始值逐渐增加K,以随着训练的进行纳入更多非目标类,遵循课程学习原则。DTM联合过滤教师模型和学生模型中的低置信度logits,在早期训练中有效净化暗知识。在CIFAR100、ImageNet和MS-COCO上的大量实验表明了DeepKD的有效性。

代码: https://github.com/haiduo/DeepKD

1 引言

知识蒸馏(KD)自Hinton等人[1]提出以来,已成为模型压缩领域的一种强大范式,在计算机视觉[2; 3; 4]和自然语言处理[5]领域得到广泛应用。通过将大型教师模型中的暗知识迁移到紧凑型学生网络中,KD解决了在资源受限设备上部署高性能模型的重大挑战——这对于自动驾驶[6]和具身AI系统[7; 8]等新兴应用来说是基本要求。

近年来,知识蒸馏(KD)方法的主要进展集中在三个方向:(1)多教师集成蒸馏[9; 10]以增强信息传递,(2)通过复杂的对齐机制进行中间特征蒸馏[11],以及(3)通过噪声注入[14; 15]和正则化[16; 17]对输入空间进行增强或输出空间操控。然而,这些方法未能对两个基本问题进行系统性分析:哪些知识传递组件会影响学生模型性能?在优化过程中如何最优地协调不同知识组件?尽管已有研究取得显著进展——DKD[18]通过损失重新参数化将KD损失分解为目标类知识蒸馏(TCKD)和非目标类知识蒸馏(NCKD)组件,揭示了NCKD在暗知识传递中的关键作用,以及DOT[19]引入了任务损失与蒸馏损失的梯度动量解耦——但关键局限性依然存在。首先,现有方法未能解决解耦损失的联合优化及其对应梯度动量。其次,动量分配的理论基础缺乏严格论证,仅依赖对损失空间的经验观察。

为解决这些局限性,作者提出了DeepKD,如图4所示,这是一个具有理论基础的 знаний дистилляция 框架。作者的研究始于对知识蒸馏中损失组件及其相应优化参数的全面分析。通过严格随机优化分析[20],作者发现任务导向梯度(TOG)、目标类梯度(TCG)和非目标类梯度(NCG)组件的最优动量系数应与其梯度信噪比(GSNR)[21]在具有动量的随机梯度下降优化器[22]中呈正相关。这实现了不同知识类型优化动态的深度解耦。

picture.image

如图1(a)所示,作者可视化了vanilla KD [1]、DKD [18]、DOT [19]以及作者提出的DeepKD在整个训练过程中的GSNR(梯度信号噪声比),其中梯度采样间隔为200次迭代。结果表明,DeepKD进一步将KD梯度解耦为TCG(任务相关梯度)和NCG(非任务相关梯度),实现了更高的整体GSNR。这一提升直接促进了模型泛化能力的提升[21; 24]。此外,图1(b)显示,与其他方法相比,DeepKD呈现出更平坦的损失函数曲面。这一观察结果与已有研究结论一致,即损失函数曲面中的平坦最小值通常与模型泛化能力的提升相关[25; 26]——这是有效知识蒸馏的关键因素。值得注意的是,仅通过基于GSNR的深度梯度解耦与动量机制,DeepKD在多个基准数据集上实现了当前最佳性能,具体实验结果详见实验部分。

picture.image

此外,虽然已有研究[1; 18; 19]强调了教师logits的暗知识的重要性,但它们通常对所有非目标类logits进行统一处理。作者通过两个关键见解挑战这一惯例:(1) 只有与目标类在语义上相邻的非目标类提供了有意义且真实的暗知识。(2) 低置信度logits可能引入的优化噪声会超过其信息价值。

为解决这些问题,作者引入了一种动态top.

Mask (DTM)机制,该机制逐步从教师和学生的输出中过滤低置信度logits(即潜在的噪声源),并作为课程学习过程[27]实现,如图1(c)所示。与通过可学习温度参数调节任务难度的CTKD[28]不同,DeepKD动态调整

,从5%的类别到完整类别数量,平衡早期阶段的稳定性和后期阶段的精细化。值得注意的是,虽然ReKD[29]对目标相似类应用top-k选择但保留了其他非目标类,作者仅根据教师保留top-

最大的非目标类logits,并动态丢弃其余部分。

跨多种模型架构和多个基准数据集的全面实验验证了DeepKD的有效性。该框架通过与现有的基于logit的蒸馏方法无缝集成,展现出卓越的通用性,在所有评估场景中均持续实现最先进性能。

2 相关工作

知识蒸馏范式:知识蒸馏主要沿着两个方向演进:基于特征和基于logit的方法。基于特征的方法转移中间表示,始于FitNets [30]使用回归损失对隐藏层激活进行建模。该方法通过注意力转移 [31] 和关系蒸馏 [32] 发展,以捕捉结构知识,最终演变为多级对齐技术,如Chen等人 [33] 的多阶段知识回顾和USKD [34] 的归一化特征匹配。

虽然FRS [35] 和MDR [36] 等方法通过球形归一化和自适应阶段选择解决教师-学生差异问题,但它们通常需要复杂的特征变换,且忽略了梯度层面的干扰。基于logit的知识蒸馏由Hinton等人 [1] 开创,专注于通过软化logit转移暗知识。最近的进展如DKD [18] 将KD损失解耦为目标类(TCKD)和非目标类(NCKD)分量,揭示了NCKD的关键作用。扩展方法包括NTCE-KD [37] 和MDR [36] 增强了非目标类的利用,但忽视了梯度层面的优化动态。

理论基础与方法学进步:知识蒸馏领域的近期进展深入探讨了优化策略和理论基础。在优化方面,DOT [19] 采用动量机制实现梯度解耦,CTKD [28] 使用课程温度调度,ReKD [29] 实施静态top-k过滤,但这些方法通常依赖经验启发式。通过多种策略应对了暗知识纯化问题:TLLM [38] 通过互信息分析识别不可蒸馏类别,RLD [39] 提出logit标准化,TALD-KD [14] 结合目标增强与logit扭曲。模型泛化的理论基础通过损失函数几何结构 [25; 23] 得到广泛研究,Jelassi等人 [22] 分析了动量在泛化中的作用。

作者的工作通过建立梯度信号噪声比(GSNR)与知识蒸馏中动量分配间的首个理论联系,将优化动态与知识迁移效率相结合。不同于DKD的损失层级解耦或DOT的经验动量分离,作者提供了基于GSNR的联合损失梯度优化的理论保证。相较于CTKD以温度为中心的课程调度或ReKD的静态过滤,作者提出的动态top.

Mask 在保留语义相关性的同时提供原则性噪声抑制,解决了先前方法中对非目标logit进行统一处理的局限性。

3 方法论

3.1 预备知识

Vanilla 知识蒸馏:给定教师模型

和学生模型

,知识蒸馏在保持性能的同时将知识从

迁移到

。设

为输入空间,

为标签空间。对于输入

,模型分别产生logits

。标准知识蒸馏[1]损失结合:

其中

是softmax函数,

是具有硬标签

的交叉熵损失,

是软logits

之间的KL散度,即

是控制分布软度的温度超参数,

平衡损失。

DKD:解耦知识蒸馏(DKD)[18] 将普通的KD损失分解为目标类知识蒸馏(TCKD)和非目标类知识蒸馏(NCKD)两个组成部分:

其中

是目标类别

以及所有其他非目标类别

的二元概率函数,且

转移目标类别知识,

捕捉非目标类别之间的关系,揭示 NCKD 在 KD 中的重要性。

DOT:面向蒸馏的培训器(DOT)[19]为交叉熵损失和蒸馏损失梯度维护独立的动量缓冲区。对于每个小批量,DOT分别从

计算梯度

,然后使用不同的系数更新动量缓冲区

其中

是基础动量,

是控制动量差异的超参数。通过将更大的动量应用于蒸馏损失梯度,DOT 增强了知识迁移,同时减轻了任务目标与蒸馏目标之间的优化权衡。然而,DOT 存在两个关键局限性:(1) 它未能解决目标类别和非目标类别知识流之间的固有冲突,这可能导致次优的优化轨迹;(2) 它缺乏对不同损失分量及其对应梯度动量的优化动态的系统分析,特别是在处理引入噪声信号并阻碍有效知识迁移的低置信度暗知识时。

3.2 基于GSNR的动量分配

与以往仅在一个层级上解耦任务和蒸馏梯度的DOT [19]不同,DeepKD引入了一种双层解耦策略,进一步将学生模型训练损失的梯度分解为三个分量:面向任务的梯度(TOG)、目标类梯度(TCG)和非目标类梯度(NCG)。作者将其梯度信噪比(GSNR)定义为:

其中

是采样间隔步长(默认值:200),

表示在步长

处的梯度,包括

。对于目标类别

和任意类别

,梯度可以表示为:

其中

分别表示学生模型和教师模型的softmax概率,该过程的详细数学推导见附录A.2。

在深度神经网络中的随机优化过程中,每个前向传播结束时计算的梯度用于反向传播。这些梯度形成了一个随机向量序列,其中梯度的统计期望和方差可以通过在时间窗口内使用短期样本均值来估计[20]。经验上,作者发现每隔200次迭代进行梯度采样可以获得更好的性能。这种估计为计算GSNR[40]奠定了基础。具体而言,统计期望代表了真实的梯度方向,而方差量化了随机采样引入的噪声。

为了更好地理解优化动态,作者对整个训练过程中解耦的普通知识蒸馏(vanilla KD)的GSNR曲线进行实证分析。如图2(a)所示,在相同的动量系数下,普通KD对于非目标类梯度(NCG)始终表现出高信噪比(SNR),这表明将"暗知识"迁移到学生模型中存在固有困难。这种持续的梯度发散表明蒸馏目标与目标任务之间存在未解决的冲突,最终导致梯度累积缓冲区(BSNR)中的SNR不稳定(图2(c))。

picture.image

值得注意的是,梯度发散反映了次优收敛,因为收敛良好的模型通常梯度接近于零,从而产生平滑衰减的SNR轨迹。作者假设信噪比较高的梯度分量应通过启发式加权优先处理以帮助优化。图2(b)表明,作者的DeepKD框架在保持所有梯度分量平衡的同时,实现了从NCG中加速且更好地吸收暗知识。所得SNR轨迹在缓冲区中表现出平滑的均匀性(图2(d)),验证了作者的理论命题。实证验证进一步证实,在知识蒸馏优化过程中,梯度分量的动量系数与其各自的SNR呈正相关。关键地,实验结果表明对特定系数值具有鲁棒性,与先前工作的观察结果(DOT [19])一致。

让作者首先考察带有动量的SGD的标准优化公式[41]:

在时间步

中,

分别表示动量缓冲区和模型的可训练参数,

是当前梯度,

是基础动量系数,

是学习率。通过分析图2(a)中的GSNR,作者观察到NCG和TOG相较于TCG保持了更高的GSNR。这一关键观察促使作者提出了自适应动量分配策略:

其中

是一个控制动量差异的超参数。如图2(b)和(d)所示,作者的DeepKD在不同动量系数下,在梯度缓冲区和原始梯度中均显著提升了GSNR,进一步验证了作者深度动量解耦方法对梯度分量的必要性。

picture.image

GSNR驱动方法确保每个知识分量遵循其最优优化路径,同时保持分量独立性,从而实现更有效的知识迁移。请注意,DeepKD通过仅修改Adam优化器[42]的一阶动量,同样适用于Adam优化器,这一点在DeiT[43]上得到了验证(见表3)。

3.3 动态Top-K Mask

现有的先进方法[18; 19; 44]通常通过对非目标类logits进行统一处理或加权分离[29]来处理,作者在这类传统方法中识别出两个关键局限性:(1) 教师模型对目标类表现出极强的置信度(softmax概率

超过92%的样本),而非目标类集体置信度低但包含有价值的暗知识,如图3(a)所示。(2) 非目标类中的暗知识表现出不同程度的同化性——语义上与目标类相似的类别(例如,目标类为"cat"时的"tiger")提供有益的暗知识,而语义上距离较远的类别(例如"airplane")引入噪声和学习困难。

picture.image

为解决这些局限性,作者首先开发了一种静态top-k Mask 方法,通过固定的

值 Mask 永久过滤语义差异极大的类别,如图3(b)所示,从而实现了 Baseline 改进。在此基础上,作者提出了一种更复杂的动态top-k Mask 机制,该机制实现了受课程学习[27]启发的阶段式

值调度。

该机制通过准确率曲线(图3(c))以三个不同的阶段运行:

  • • 易学阶段:K值从总类别数的5%线性增加到最优静态K值
  • • 过渡阶段:保持最优静态

  • • 硬学习阶段:将

线性扩展以涵盖全部类别数量。通过消融实验或使用

的训练数据来确定最优静态 K 值,以降低训练成本。动态顶部-

Mask 学习的完整过程如图1(c) 所示。

对于每个训练迭代

,作者计算 Mask

为:

其中rank

表示logits的升序排列的秩,

从类别总数的5%逐渐增加到类别总数。 Mask 蒸馏损失公式表述如下:

其中

表示逐元素乘法。这种机制有效地抑制了噪声,同时保留了语义相关的暗知识。

3.4 DeepKD框架

基于对GSNR的理论分析,作者提出了DeepKD,一个引入深度解耦优化和自适应去噪的知识蒸馏综合框架。如图4所示,DeepKD将知识迁移过程分解为三个并行梯度流:面向任务的梯度(TOG)、目标类梯度(TCG)和非目标类梯度(NCG)。每个梯度流独立管理,具有各自的动量缓冲区,并根据其独特的GSNR特性进行优化。这种解耦架构通过允许各组件独立优化,实现了更有效的知识迁移。DeepKD的完整损失函数(详细实现见附录中的算法1):

其中

表示标准交叉熵损失,

是平衡各损失分量贡献的固定系数,

是目标类损失,

是非目标类的动态 top-

Mask 损失。这种公式化方法使框架能够在保持计算效率的同时,有效结合任务特定学习和知识蒸馏。

4 实验

4.1 数据集和实现

作者对三个广泛使用的基准进行了全面评估:CIFAR-100 [45](100个类别,50k训练/10k验证,32×32图像),ImageNet-1K [46](1,000个类别,1.28M/50k裁剪至224×224图像),以及MS-COCO [47](80类检测,118k训练/5k验证图像)。

在实现方面,作者遵循标准做法,使用动量为0.9的SGD优化器,权重衰减为5×10^-4(CIFAR)或1×10^-4(ImageNet)。训练计划因数据集而异:CIFAR使用240个epoch,批量大小为64,初始学习率为0.01-0.05;而ImageNet使用100个epoch,批量大小为512,学习率为0.2。所有实验均在配备Nvidia RTX 4090 GPU和AMD 64核处理器的系统上进行。

4.2 图像分类

在CIFAR-100上的结果。作者的DeepKD框架在同类和异类蒸馏设置中均表现出持续的性能提升。在同类架构下(表1),

picture.image

DeepKD在不使用top-

Mask 的情况下实现了

的准确率提升,在使用 Mask 的情况下实现了

的准确率提升,其性能优于基于特征的方法

。top-

变体进一步将性能提升了高达

,MLKD

Ours达到了

的准确率。

在异构场景(表2)中,DeepKD展现出强大的泛化能力:CRLD

urs在ResNet32×4 MobileNet-V2上实现了

的准确率(

),而

在WRN-40-2ResNet8×4上达到了

的准确率(

)。在超参数变化下,性能保持稳定(方差

),证实了DeepKD在处理冲突蒸馏信号方面的有效性。

picture.image

ImageNet-1K上的结果。如表3所示,DeepKD在不同教师-学生对之间均取得了显著提升。对于ResNet50和MobileNet-V1,

(w. top-k)将top-1准确率提升了

,相较于

,这是所有配置中最大的提升。CRLD

Ours(w. top-k)建立了新的SOTA结果:

(ResNet34

Deit-Tiny,

。动态top-k Mask 始终提升性能,在top-5准确率上贡献了额外的

提升。值得注意的是,DeepKD展现出强大的可扩展性:(1) 对于轻量级学生(MobileNet-V1/Deit-Tiny),提升达到

;(2) 使用大型教师(RegNetY-16GF),top-5准确率超过

,超越了所有基于特征的方法。这些结果验证了DeepKD在大规模蒸馏场景中的有效性。

4.3 在MS-COCO上的目标检测

DeepKD在目标检测任务上表现出色(见表4)。通过动态top-k Mask

,DeepKD将 Baseline KD的AP提升了+1.93%,达到32.16%(对比30.13%),并且仅使用logit蒸馏就超过了ReviewKD的33.71% AP。

达到了34.20% AP,在所有KD变体中表现最佳,比vanilla DKD提升了+1.86%。动态top-

机制提供了额外的+0.05% - 0.34% AP提升,其中在

上的提升最为显著(对于

为+2.64%)。DeepKD在定位方面优于基于特征的LSKD,

达到了36.59%

(对比LSKD的36.34%)。这些结果证实了DeepKD在密集预测任务上的有效性。

picture.image

5 消融实验

动量系数与损失函数。基于梯度信噪比(GSNR)分析,作者将DeepKD中的梯度动量分解为目标类梯度(TCG)和非目标类梯度(NCG)。为验证这种解耦策略,作者在CIFAR-100上使用ResNet32×4(教师)和ResNet8×4(学生)对进行了全面的消融研究(见表5)。结果表明了DeepKD的有效性,并突出了每个组件的重要性。

值得注意的是,DeepKD仅引入了一个超参数Δ,该参数在不同数据集上均表现出鲁棒性。遵循DOT [19],作者在CIFAR-100上的KD+DeepKD中设置Δ=0.075,在DKD+DeepKD、MLKD+DeepKD和CRLD+DeepKD中设置Δ=0.05。对于ImageNet(教师知识更可靠),作者为所有变体使用Δ=0.05。作者的损失函数消融研究表明,每个组件都对整体性能有显著贡献,其中NCKD是主要因素——这与DKD [18]中的发现一致。这些结果验证了作者的解耦策略,并展示了其在提升模型性能方面的有效性。

picture.image

动态top-k Mask 。为了排除动态top-k Mask 的影响,作者在动量系数和损失函数的消融研究中不采用这一策略。对于动态top-k Mask 配置,作者经验性地设置参数为

值=55,相位1=60,相位2=170(见表5)。实验结果表明,即使简单的超参数调整也能进一步提升模型性能,这表明将动态top-k Mask 机制集成到KD中具有巨大潜力,值得未来工作进行进一步探索。

此外,作者对梯度解耦和动态顶部两种 Mask 策略进行了消融研究。结果表明,每个组件单独使用时均能提升性能,而它们的组合则能带来进一步的改进。需要注意的是,当单独启用动态顶部 Mask (DTM)时,该机制作用于所有logits,包括目标类别。

6 局限性与未来工作

尽管作者当前的工作专注于基于logit的蒸馏,但SNR驱动的动量解耦机制自然地扩展到特征蒸馏场景。通过将特征对齐损失视为额外的优化组件,DeepKD能够自动处理多级知识迁移,无需手动加权,补充了现有的特征增强技术,如注意力迁移[31]和对比蒸馏[53]。

未来的工作可以探索将DeepKD应用于更复杂的场景,例如多教师蒸馏和跨模态知识迁移。此外,作者的动态top-

Mask 策略在提高蒸馏性能方面显示出良好的效果,表明其具有进一步优化和适应不同模型架构和数据集的潜力。

参考

[1]. DeepKD: A Deeply Decoupled and Denoised Knowledge Distillation Trainer.

picture.image

扫码加入👉「集智书童」交流群

(备注:方向+学校/公司+昵称)

picture.image

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
大规模高性能计算集群优化实践
随着机器学习的发展,数据量和训练模型都有越来越大的趋势,这对基础设施有了更高的要求,包括硬件、网络架构等。本次分享主要介绍火山引擎支撑大规模高性能计算集群的架构和优化实践。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论