基于Sinkhorn距离的增量蒸馏方法及其在目标检测中的应用

Sinkhorn Distance-Based Incremental Distillation and Its Application in Object Detection

摘要: 目标检测模型通常需要不断学习新目标知识以满足业务场景的变化。为了避免增量学习中的“灾难性遗忘”问题，提高目标检测模型在增量学习中的性能，本论文提出了一种基于Sinkhorn距离的增量蒸馏方法。该方法采用Sinkhorn距离作为蒸馏损失函数，引导两个YOLOv11模型(学生模型和老师模型)间的知识传递。和常用于知识蒸馏的Kullback-Leibler散度相比，Sinkhorn距离能够有效避免学生模型输出平滑问题，从而提高增量蒸馏的效果。该论文提出的增量蒸馏方法在coco 2017数据集上进行验证，结果表明该方法能够有效解决“灾难性遗忘”问题，从而提高目标模型的性能。

Abstract: Object detection model continually learns new object knowledges for the changes in the business. To avoid the “catastrophic forgetting” problem, this paper proposed an incremental distillation method based on Sinkhorn distance. In the method, the Sinkhorn distance is used as loss to trans form the knowledge from the teacher to the student in distillation process. Sinkhorn distance can avoid the output smooth problem of the student caused by the Killback-Leibler (KL) divergence in knowledge distillation process. The proposed method is evaluated on coco 2017 dataset, and the results shown that the proposed method has ability to effectively solve the “catastrophic forgetting” problem and increase the performance of object detection.

1. 引言

目标检测是计算机视觉的重要任务，用于在图像中找出目标并识别目标的类型。目标检测的输入通常为一张带有目标的图像，而输出为目标的类型以及目标的边界框(bounding box, BBox)以标定目标在图像中的位置。目前，目标检测在交通控制、飞行安全、环境检测、军事侦察等领域有着广泛应用[1]。在实际应用中，目标检测模型通常需要不断学习新目标的知识以满足业务变化[2]。微调现有模型以适应新任务需求是一个不错的选择，但是仅使用新目标数据去微调模型，会导致模型能够非常准确的检测出新目标，而不能准确检测出旧目标，及出现“灾难性遗忘”(catastrophic forgetting)问题[3]。知识回放(knowledge replay) [4]和增量蒸馏(incremental distillation) [5]是最常用的两种解决“灾难性遗忘”问题的增量学习方法。知识回放将所有目标(旧目标和新类型)数据一起用于微调模型，以确保模型学到新目标知识且避免遗忘旧目标知识。由于引入了旧目标数据，知识回放在微调模型时需要大量的标注数据，从而大大增加微调的人力、计算和时间成本。增量蒸馏是另一种解决“灾难性遗忘”问题的方法。它是一种基于知识蒸馏的增量学习方法，将已有的具备旧目标知识的模型作为老师，将学习新目标知识的模型作为学生，利用蒸馏过程，将老师的知识传递给学生，同时让学生学会新知识。和知识回放不同，增量蒸馏仅将新类型数据用于模型微调，减少了标注数据量，节约了资源开销。目前，增量蒸馏已经广泛的应用于安全、医疗、工业生产等领域[6]。为了引导学生学习老师的知识，增量蒸馏和知识蒸馏一样，采用蒸馏损失函数量化学生与老师的差距。因此，蒸馏损失函数的性能，直接决定了增量蒸馏的性能。Kullback-Leibler (KL)散度是最常用的蒸馏损失函数。它能够准确地量化老师和学生输出分布的差异，从而引导学生学会老师的知识。目前，KL散度已经广泛地用于自然语言、图像、语音知识蒸馏中[7]。然而，研究表明当老师和学生输出差异较大时，KL测度会导致学生输出过于平滑或者过于简单，从而导致学生模型的检测性能下滑[8]。为了解决上面问题，我们提出一种基于Sinkhorn距离的增量蒸馏方法。该方法采用Sinkhorn距离代替KL散度作为蒸馏损失函数引导学生学习老师知识。Sinkhorn距离[9]是Wasserstein距离的近似，以减少Wasserstein距离的计算消耗。与KL散度相比，Sinkhorn距离量化学生和老师输出分布的转化所代价，而不是直接量化老师和学生输出分布差异。因此，即使老师和学生输出存在较大差异，通过最优化传输过程，Sinkhorn距离也能用引导学生学习到老师的知识，从而提高增量蒸馏的性能。

2. 相关工作

2.1. 目标检测模型

本模基于卷积神经网络(Convolutional Neural Network, CNN)的深度模型是具有最佳检测性能的一类模型，这类模型可以分为两类：二阶段模型和一阶段模型。二阶段模型，通常先常用一个骨干模型对图像进行多尺度特征提取，形成特征金字塔，并为每个尺度的特征图设置备选外包边框，框出图上目标，最后根据备选外包框识别出目标类型。二阶段模型的代表为RCNN系列模型，如RCNN、fast-RCNN和faster RCNN等。和二阶段模型相比，一阶段模型合并多尺度特征提取、备选框选取和目标类型识别过程，形成端到端模型。一阶段模型的代表为YOLO系列模型[10]。到目前为止，二阶段模型无论从应用量还是检测性能和一阶段模型的差距不大，但是由于一阶段模型采用端到端的方式，屏蔽了特征提取骨干模型选择、备选框模型选择等细节，所以降低了模型在实际落地应用的难度。

2.2. 增量蒸馏

增量蒸馏是一种基于知识蒸馏技术的增量学习方法。和知识蒸馏一样，需要两种模型参与，一种是老师模型(简称老师)，另一种是学生模型(简称学生)，老师和学生的数量可以为一个或者多个。增量蒸馏将具备旧知识的模型作为老师，而将学习新知识的模型作为学生，通过知识蒸馏过程，将老师的知识传递给学生，同时让学生学习到新知识。因此，学生在学会新知识时，不会遗忘掉旧知识，从而解决了增量学习中的“灾难性遗忘”问题。另外，和知识回放相比，增量蒸馏仅用新知识数据微调模型，降低了微调过程的资源消耗。在增量蒸馏中，损失函数可以表示为，

$L=L_m+\alpha L_d$ (1)

其中，$\alpha$ 为权重，L为总体损失函数，表示为$L_m$ 和$L_d$ 的加权和。$L_m$ 为检测损失函数，量化学生输出和真实数据标签的差距。基于softmax的交叉熵是最常用的检测损失函数。$L_d$ 为蒸馏损失函数，量化学生和老师的差距，代表学生学习老师知识的情况。在目标检测中，由于模型通常不仅需要输出目标类别，还需要输出目标的边界框以定位目标在图像中的位置，所以对于目标检测来说，增量蒸馏的损失函数由三部分组成即，

$L=L_m+\alpha L_{cd}+\beta L_{rd}$ (2)

其中，$\alpha$ 和$\beta$ 为权重，$L_{cd}$ 和$L_{rd}$ 均为蒸馏损失函数，$L_{cd}$ 量化学生和老师输出目标类型的差距，而$L_{rd}$ 量化学生和老师输出目标边界框的差距。由于学生和老师针对目标类型的输出均为sotfmax概率分布，因此KL散度成为目前最常用的$L_{cd}$ 。对于YOLO v8以上版本的模型来说，因为其针对目标边界检测采用DFL策略，将边界框坐标表示为softmax概率分布，所以基于YOLO的增量蒸馏通常也使用KL散度来实现$L_{rd}$ 。

然而，研究表明，当老师和学生输出差异较大时，KL散度要么导致学生的输出过于平滑，无法区分目标类别；要么导致学习的输出过于简单，无法利用老师知检测目标。因此，我们提出了基于Skinborn距离的增量蒸馏方法。

3. 基于Sinkhorn距离的增量蒸馏方法

文本基于YOLO模型结构设计增量蒸馏流程，具体过程如图1所示。

图中，YOLO模型可以分类三个部分，分别为骨干网络(backbone)，颈部网络(neck)和检测头。骨干网络用于从输入图像中提取多尺度特征，颈部网络用于对抽取出的多尺度特征进行融合，形成特征图。检测头分为类型检测和边界框回归，用于根据特征图输出目标类型和目标边界框。由于YOLO对边界框回归采用DFL策略将目标边界框坐标表示为softmax概率分布，YOLO的类型检测和边界框回归都视为分布学习问题处理，能够更好的处理检测类别不平衡且提高检测精度。增量蒸馏需要两个YOLO模型参与，一个是老师模型，另一个是学生模型。老师具备检测旧目标的能力，例如图1中的老师具备检测人的能力。增量蒸馏的输入为仅标注新目标的数据，如图1中的输入仅标注了马的信息，而没有标注人信息。增量蒸馏的目的就是通过微调过程让学生能够同时检测新旧目标(例如同时检测人和马)的类型和位置。在增量蒸馏过程中，老师被冻结，用于生成旧目标的特征图和计算蒸馏损失函数。学生通过它的骨干网络和颈部网络生成新目标的特征图。将老师生成的旧目标的特征图和学生生成的新目标的特征图一起输入到学生的检测头，输出最终的检测结果。最后，为了让学生同时学会老师知识和新增知识，基于损失函数最优化的反向传播过程用于调节学生模型参数。增量蒸馏的损失函数极为重要，直接影响蒸馏效果。本论文采用Sinkhorn距离代替KL散度作为蒸馏损失函数，提升蒸馏性能。假设$f_s$ 和$f_t$ 为老师和学生模型，对于输入$x_i$ ，他们的softmax输出可以表示为，

Figure 1. Flow chart of incremental distillation

图1. 增量蒸馏流程图

$\begin{array}{l}{t}_i={\text{softmax}}_{\theta }\left(f_t\left(x_i\right)\right)\\ s_i={\text{softmax}}_{\theta }\left(f_s\left(x_i\right)\right)\end{array}$ (3)

其中，$\theta$ 为sofmax的温度，$t_i$ 和$s_i$ 表示为老师和学生模型对输入$x_i$ 的输出向量，那么蒸馏损失函数通常可以定义为，

$L=J\left(t_i,s_i\right)$ (4)

其中，$J$ 为距离函数，量化老师输出和学生输出的差距。

Sinkhorn距离是Wasserstein距离的近似，以减少Wasserstein距离的复杂度，用于量化老师和学生分布的转化代价。Sinkhorn距离表示为：

$J_{SD}\left(t_i,s_i\right)=\langle P^{\lambda },D\rangle$ (5)

其中，D为代价矩阵，其元素表示为，

$D_{mn}=\left|t_i^m-s_i^n\right|$ (6)

其中，$t_i^m$ 和$s_i^n$ 分别为$t_i$ 和$s_i$ 的第m个和n个元素。$P^{\lambda }$ 为最优传递矩阵，定义为

$P^{\lambda }=\arg \min \left(\langle P,D\rangle \right)-\lambda h\left(P\right)$ (7)

其中，P为结果矩阵，满足$t_i=P\ast 1_d,s_i=P^{\text{T}}\ast 1_d$ ，$1_d$ 为全1矩阵。$\ast$ 为矩阵乘法，$h$ 为熵。$P^{\lambda }$ 通过下面公式迭代生成。

$\begin{array}{l}{K}^0=\exp \left(-\frac{1}{2}D\right),\\ K^t=G^{t}diag\left({\left(G^t\right)}^{\text{T}}{1}_d÷t_i\right){}^{-1},\\ G^t=diag{\left(K^{t-1}{1}_d÷s_i\right)}^{-1}{K}^{t-1},\\ P^{\lambda }=K^A\end{array}$ (8)

其中，$t=1,2,3,\cdots ,\lambda$ ；$÷$ 表示向量的元素除法。因为YOLO模型的类型检测和边界框回归都是分布学习问题，所以基于Sinkhorn距离的损失函数不仅可以量化老师和学生输出目标类型的差异，同时可以用于量化老师和学生输出边界框的差异。因此，增量蒸馏的损失函数表示为：

$L=L_m+J_{SD}\left(C_i^t,C_i^s\right)+J_{SD}\left(B_i^t,B_i^s\right)$ (9)

其中，$C_i^t$ 和$C_i^s$ 分别为老师模型和学生模型对输入图像$x_i$ 的目标类型输出，$B_i^t$ 和$B_i^s$ 分别为老师模型和学生模型对输入图像$x_i$ 的目标边界框输出。通常，在计算蒸馏损失函数时，老师和学生模型的所有输出值都要参与，这样会让蒸馏损失函数的计算变得复杂。Elastic Response Distillation (ERD)是一种蒸馏损失函数计算加速的技术，他通过过滤仅让部分对蒸馏过程影响力较大的模型输出参与损失函数的计算，从而较少了计算量。因此，我们也在计算Sinkhorn距离时，引入了ERD技术(即图1中绿色的ERS模块)。因此，Sinkhorn距离改写为：

$J_{SD}\left(t_i,s_i\right)=\langle P_M^{\lambda },D_M\rangle$ (10)

其中，$P_M^{\lambda }$ 和$D_M$ 为在输出集合M上计算的最有传递矩阵和代价矩阵。N为老师和学生模型输出的全集，有$M\subset N$ 。在本文中集合M由Tao提出的方法[11]从集合N中筛选得到。

4. 实验

4.1. 实验数据集及实验步骤

本论文采用COCO数据集[12]进行实验。该数据集包含了80种不同大小，不同类别的图像目标。采用训练数据集蒸馏模型，采用验证集得到实验结果。在实验中，参与对比的方法有数据回放方法(Data Replay, DR)，KL距离蒸馏方法(KL-based Distillation, KL-D) [13]，RKL距离蒸馏方法(RKL-based Distillation, RKL-D) [14]和Skinborn距离[15]蒸馏方法(Sinkhorn-based Distillation, S-D)。实验将两个YOLOv8模型分别作为老师和学生模型，实现基于不同距离的增量蒸馏方法，仅用一个YOLOv8模型实现数据回放方法。所有实验均在4张Nvidia Geforce 4090 GPU上进行，每张GPU有40G显存。实验中的模型均在显存中进行计算，不使用内存或者CPU计算。采用Transformers框架运行模型，不对模型进行额外加速。

4.2. 老师模型对蒸馏效果的影响

该实验用于测试老师模型对蒸馏性能的影响。一般而言，一个识别模型的精准率和识别类型数量成反比，模型识别的类型越多，那么模型识别精准率随之下降。在该实验中，从COCO数据集中选出固定40类目标训练老师模型，然后分别采用新的10、20、30、40种新类型目标蒸馏学生模型，最后通过验证集，测试学生模型的识别精准率。在训练和蒸馏时，每类目标均为100个样本。采用AP50 (简写为AP)作为量化学生模型识别精准率的标准，其中AP50表示学生模型输出的类型和真实类型一致，且输出的目标边界框和真实边界框有50%的重合，则判为识别正确。实验结果如下图2所示。

为了简洁的表示老师训练集(即旧类型数据集)和学生训练集(即新类型数据集)，采用老师训练集 + 学生训练集的表示方法，例如40 + 40表示采用40个类型的数据训练老师，采用另外40个类型的数据蒸馏学生。图2(a)中，随着新目标类型的增加，学生的识别精准率在逐渐下降。这是因为，随着识别的目标增多，模型需要形成的分类平面就会越复杂。当模型规模一定时，模型拟合复杂的分类平面就越困难，所以导致模型的识别性能下降。DR模型始终保持较高的识别精准率，这是因为其将所有新旧类型的目标数据合在一起训练模型，因此模型不仅能记住旧目标知识，同时能够记住新目标知识。我们提出的S-D方法，获得比KL-D和RKL-D更好的识别性能。这是因为S-D采用Skinborn距离对学生模型进行蒸馏，而KL-D和RKL-D分别采用KL散度和RKL散度蒸馏学生模型。在学生模型和老师模型输出分布较大时，如增量蒸馏情况，Skinborn距离能够更好的引导学生模型学习老师的知识。图2(b)中展示了四种方法的训练时间消耗，采用40 + 10情况时的DR训练时间为基准时间(及时耗百分比为100)。可以看出DR的训练用时远远高于S-D、KL-D和RKL-D方法。这是因为DR方法需要将新旧目标数据一起训练模型，而基于增量蒸馏的方式(S-D、KL-D、RKL-D)仅需要新类型目标蒸馏学生模型。因此，增量蒸馏大大减少了训练时间消耗。另外，S-D的训练时间仅为DR的28%左右，需要的标注数据仅为DR方法的50%左右，但获得DR方法近95%的识别性能。这说明，我们提出的S-D方法是一种有效的增量学习方法。

(a)

(b)

Figure 2. The distillation results for different teachers. (a) Precision rate, (b) Time cost

图2. 不同老师模型的蒸馏效果。(a) 精准率，(b) 时间消耗

4.3. 对抗“灾难性遗忘”的能力

该实验用于测试增量学习方法对抗“灾难性遗忘”的能力。在实验中，我们采用10个、20个、30个和40个类型的目标数据训练老师模型，然后固定用40个类型的目标数据蒸馏学习模型。每类目标均为100个样本。同样，采用DR、KL-D、RKL-D和我们提出的S-D方法进行比较。比较结果如下图3所示。

图3中，DR在所有实验中都取得了最好的识别精准率，这是因为DR采用将新旧类型数据一起训练模型的策略，对抗“灾难性遗忘”问题，得到最稳定的增量学习效果，而代价是较高的人力标注成本和训练时间、算力成本。同时，可以看到，我们提出的S-D模型在三种增量蒸馏方法中取得了最好的识别性能。这是因为S-D采用Sinkhorn距离作为损失函数，引导将老师模型的知识传递给学生模型。相比KL散度和RKL散度，Sinkhorn距离能够更好的量化老师和学生知识分布的差距，从而有效抑制了“灾难性遗忘”问题。

Figure 3. The performance of different models in resisting the “catastrophic forgetting”

图3. 不同模型对抗“灾难性遗忘”的性能

4.4. 对小样本目标的学习能力

本实验用于测试上增量蒸馏方法在识别稀有目标上的能力。在实验中对KL-D、RKL-D和我们提出的S-D进行比较。为了测试上述三种增量蒸馏方法在小样本目标上的识别性能，我们选出60种目标数据作为训练集，将这60种目标分类三份。第一份20种目标仅包含100个样本/目标，作为小样本目标数据，记作S。第二份20种目标包含400个样本/目标，作为大样本目标数据，记作L1。最后一份20种目标包含400个样本/目标，作为大样本目标数据，记作L2。我们实验分类两部分，第一部分采用S + L1训练老师模型，而用L2蒸馏学生模型。第二部分采用L1 + L2训练老师模型，而用S蒸馏学生模型。在实验验证阶段，仅用小样本目标的验证集去验证学生模型的识别精准率。实验结果如图4所示。

图4(a)展示采用S + L1训练老师，采用L2蒸馏学生模型的情况。S-D方法蒸馏的学生模型的精准率非常接近老师模型，而用KL-D方法和RKL-D方法蒸馏的学生模型的精准率要远远低于老师模型，这表明S-D方法能够在小样本条件下，将老师的知识传递给学生。这是因为，Sinkhorn距离，不仅能够引导学生学习到老师的大样本目标输出，同时可以引导学生学习老师小样本目标的输出。相对地，基于KL散度的距离，仅能让学生关注到老师大样本目标的输出。图4(b)展示采用L1 + L2训练老师模型，而采用S蒸馏学生模型的情况。可以看出，三种增量蒸馏方法获得相似的识别性能。这是因为，增量蒸馏分的损失函数分为两部分，识别损失函数用于引导学生模型学习新知识，而蒸馏损失函数引导学生学习老师知识。由于三种蒸馏方式均采用经典的交叉熵作为识别损失函数，所以它们在引导学生模型学习新知识的性能相差不大。

(a)

(b)

Figure 4. The performance of different models in flew-shot learning. (a) Using S + L1 to train teacher and L2 to distill students; (b) Using L1 + L2 to train teacher and S to distill students

图4. 不同模型的小样本学习性能。(a) 采用S + L1训练老师，L2蒸馏学生的结果；(b) 采用L1 + L2训练老师，S蒸馏学生的结果

5. 总结

增量蒸馏是一种增量学习方法。它利用蒸馏过程，让学生能够学习新知识，同时从老师那里继承旧知识，从而克服在模型学习新知识时出现的“灾难性遗忘”问题。在传统的蒸馏方法中，采用基于KL散度的损失函数引导学生模型学习老师模型的知识。然而当老师模型和学生模型的输出分布差距较大时(如增量蒸馏)，KL散度会导致学生模型不能准确地学习到老师的知识。为了克服这个问题，我们提出一种基于Sinkhorn距离的增量学习方法，该方法采用Sinkborn距离作为蒸馏损失函数。和KL散度直接量化老师模型和学生模型输出分布不同，Sinkhorn量化从学生模型输出分布最优化转移到老师模型输出分布的代价。因此，该距离在老师和学生输出分布差异较大时，也能较准确的量化两者的差距，能够准率的引导学生模型学习到老师的知识。换句话说，Skinhorn距离更适合增量蒸馏。在实验中，我们提出的基于Skinhorn距离的增量蒸馏方法，和基于KL散度和基于RKL散度的方法相比，能够得到识别精准率更高的学生模型。