读论文：STKD :从同步教学中提炼知识以实现高效的模型压缩

标题: STKD: Distilling Knowledge From Synchronous Teaching for Efficient Model Compression
作者: Tongtong Su , Jinsong Zhang , Zhaoyang Yu , Gang Wang
原文地址：https://ieeexplore.ieee.org/document/9757208

---

Abstract

现有的高级KD方法都局限于固定的特征提取，也就是捕获教师的结构只是来指导学生训练，通常不能为学生获得全民的知识。为此，本文提出了同步教学知识蒸馏STKD，以整合在线和离线教学，从而向学生传递丰富而全面的知识。通过该方法可以结合各种不同KD方法的优点。并在几个真实数据集上取得了显著的性能提升。

Index Term——Knowledge distillation (KD), model compression, mutual learning network, prior knowledge.

---

Introduction

深度卷积神经网络在计算机视觉任务上表现优异,但其需要大量数据、参数和计算资源。而日常中的移动设备受限于计算能力、存储空间和能耗,部署模型于其上存在困难。模型压缩技术能够减小模型规模,使其部署到资源受限的移动设备成为可能,是实现高效推理的重要手段。而知识蒸馏KD方法就是模型压缩的有效方法之一：它从一个较大的教师网络中提取知识来指导一个较小的学生网络的训练，如下图所示：

而传统的KD都是以离线蒸馏的方式进行的：先训练一个高容量的复杂模型来作为教师模型，然后将其作为暗知识来训练学生模型。这种方法大多聚焦于从预训练教师中利用类概率和中间特征等知识,意味着它们只提供教师和学生网络之间的单向KD。但是却忽略了训练过程中生成的有价值知识，而这也可以互相迁移进行交互学习。

• 在线蒸馏：对教师和学生没有明确的定义，并采用对等学习策略来训练所有对等网络。但是这种方法也有一定的局限性：

1. 首先，有限的知识被提炼并传递给每个学生，因为它只利用来自两个或更多对等网络的最终预测来进行知识共享；
2. 其次，中间层中的信息交互被忽略，这导致编码在中间层中的有用特征和知识在很大程度上被遗忘；
3. 最后，这种策略缺乏高能力的教师来为学生产生指导性的指导，因此对目标网络提供无效的监督；

对此，本文就推出了一种新的同步教学知识蒸馏STKD，它将在线KT和离线KT集成到一个统一的框架中,以生成更好地指导学生网络训练的有用和全面的知识,如下图所示：

执行教师-学生优化来训练紧凑的学生网络和从两个教师中提取知识。具体来说,学生网络以在线KD的方式从后验教师中学习知识,以离线KD的方式从先验教师中学习知识。
方法通过在线离线双向知识迁移的统一框架,使学生网络在中间层与后验教师网络互学习提高软预测质量,在最终层与其对齐软概率获取高级知识特征,并与先验教师网络协同训练获取额外指导性监督,从而同时获得教师网络的多方面知识,改进知识蒸馏的训练性能。

---

贡献

1. 提出了一种创新的KD学习方法STKD，将线上与线下知识转移KT都纳入到统一框架中，向小型的学生传递全面的知识；
2. 本文将教师的多层次知识引入到学生的学习过程中，引导其学习中高级知识。还将教师模型中获得的先验知识来引导学生网络；
3. 在四个基准数据集上评估了STKD方法。并发现其具有最优的性能；

---

综述

一、相关工作

训练铜川的超参数话DNN模型往往需要占用大量的存储空间和计算资源。因此，在一些资源受限的边缘设备上部署如此复杂的网络是不可能的。

• 离线KD：就是训练一个庞大的教师网络，然后向一个小型的学生提供知识。该方法主要提供先验知识；
• 在线KD：对等网络之间相互交换分类知识，协同工作。该方法主要提供中级知识和高级知识；

本文则是结合了上面两种方式，从而为学生传递丰富而全面的知识，即采用两位教师监督学生的训练过程，也就是先用离线KD的方法来引导学生，再用在线KD的方法来和学生互相转移中高级知识。其结构如下图中(d)所示：

(a)就是离线KD方法，(b) 是在线KD方法，©是在线KD方法的一种改进。

---

二、 方法

2.1 动机与概述

现有的KD方法在训练过程中只提取了最后一层的有限知识，而忽略了中间层的特征。为此，本文倾向于提取中底层的特征，再利用两位教师来传递更全面、多层次的知识，来对学生产生较强的监督作用。

• 动机：当老师再教授某种知识时，学生不仅局限于当前的课堂知识，还要会借鉴和参考以前的知识。这也就意味着要同时应用线上KD和线下KD；

本文的框架如下图所示：

先采用后向师生对，以相互学习的方式交互，生成中级知识和对应的软化对数来作为高级知识。然后应用前向老师执行标准KD损失，以提供包含有用信息的软目标

2.2 在线蒸馏KD

本文使用WRN模型来作为骨干网络，如下图所示：

WRN是由一些特殊的块组成，而这些块被称为宽剩余块WRB；
在该工作中，中项知识被提炼出来，然后通过KL巡视提供相互监督。

• 工作原理：最小化教师和学生模型的中间特征之间的分歧，这被用作基础知识；
• 将KL损失应用于中间KD有两个优点：

1. 在KL损失的监督过程中，并行网络根据自身固有差异来学习特征图知识，可以提取更多的信息来提升学习能力；
2. 对特征映射对进行直接比较，从而使得它们的行为相似。而KL损失可以减少特征映射级别的差异，从而使学生表现更好；

此处，后验教师T和小型学生S由M个区块组成，并产生M个中间特征分布，分别为 { f i T } i = 1 M \{f_i^T\}^M_{i=1} {fiT​}i=1M​和 { f i S } i = 1 M \{f_i^S\}^M_{i=1} {fiS​}i=1M​（M是块的数量）。为了利用中间知识，本文引入了一个中间约束条件，也就是让学生的特征分布与教师的一致。采用了KL分歧来匹配两个模型的中间特征分布。第i个区块对的损失可以表示为：
$$L_{KL}^{block}(X,f_i^T,f_i^S)=KL(f_i^T,f_i^S)$$

对于以上的这个区块对，本文设计了一个块单元（也就是虚线框起来的），并强迫学生在每个阶段结束时接受老师的外部指导。这种方法可以使得学生具有与教师
同等的预测能力，它还有利于提高学生的特征提取能力。

将块单元的中间特征损失定义为个特征映射对的和：
$$L^{Inter}=\sum _{i=1}^M{L}_{KL}^{block}(X,f_i^T,f_i^S)$$
通过该函数来优化输出信息的质量。此外还采用了L2损失和交叉熵损失来训练模型：

L2损失

还通过交叉熵CE损失来调整神经网络的权值：

整体在线教学的KT损失函数如下所示：

其中α和β是平衡的权重系数。我们采用Linter损失来使并行网络之间基于中间特征对共享信息。此外,我们使用LLogits损失通过对齐其logit信息来生成有用的监督。

---

2.3 离线蒸馏KD

该方法可以提供额外的知识，这不仅可以提高学生模型的质量，还可以加快训练收敛速度。

师生模型被训练成一开始就互相分享信息。而一个优化的教师网络对于该框架来说十分重要。因此，通过先任教师将有价值的知识以离线的方式转移，进一步优化学生网络，可以提供不同类型的知识。

该方法主要是给学生模型提供了logits信息，从而加快了训练。logits信息中通常携带有暗知识到学生网络，而其中就包含了不正确分类结果的相对概率，这可以用来理解教师。本文引入了临时参数$\tau$来控制第I个软目标的重要性：

其中，$a_{T2}$为教师网络的logits值，$a_S$是学生网络的logits值。通过这个softmax函数将原始向量转换为概率向量。教师模型在之前就已经经受过了足够的训练，可以以离线KD的方式提炼和传递知识给学生

离线损失包括两部分：KD损失和CE损失。我们通过CE定义KD损失以匹配学生的软化logits$\tau (o_S)$和教师的软化logits$\tau (o_{T2})$

KD损失的优点：
首先，它可以迫使学生网络从老师的柔化输出中学习；
此外，标准KD方法可以加快数据几何因子的收敛速度；

纯KD损失可以表示为：

式中γ为LKD的平衡权重系数。同样，我们使用γ = 0.7和τ = 4进行KD训练

---

2.4 全损失训练算法

就是将前面两个损失加在一起然后进行训练。

其中：
首先优化先验老师以向学生提供更多信息，同时通过整体损失函数训练具有后向老师的学生。

1. 对于在线KT，我们采用DNN模型的逐块单元进行知识传递，每个块的输出作为中间特征分布。此外，采用L2损失来匹配logits，从而为学生提供高质量的监督；
2. 对于线下KT，考虑到高容量的教师可以对学生进行较强的监督，我们利用预训练的预测输出提供额外的指导，从而进一步提高学生的表现；

---

三、 实验

3.1 实验设置

1. 数据集：分别在CIFAR-10、 CIFAR-100、 SVHN和ImageNet ILSVRC 2012上进行实验来验证其有效性：
   (1) CIFAR-10/100：包含了从10个类别中提取的3232的RGB小图像，每个类别包含6000张图像。而CIFAR-100则是包含有100个类别（就是有20个大类，每个大类中包含有5个子类）；
   (2) SVHN：这是一个真实世界的图像数据集，它类似于MNIST数据集，具有0到9类别中的3232的RGB小裁剪数字；
   (3) ImageNet：这是一个包含有1000个不同对象类别的大规模数据集，并且其中每个图像的大小也很多样。本文将图像裁剪为224*224的大小，用于训练和测试，并且这些图像也都包含了更多不相关的噪声，更好的反应真实情况；

2. 基础方法：本文拿来对比的方法可以分为两部分：
   (1) 离线KD方法：ATKD和SPKD；
   (2) 在线KD方法：DML、RLKD、CTKD、FFL和OKDDip；
   base则是表示没有老师的普通CNN模型；

ATKD：老师将注意力映射信息传递给学生；
SPKD：更加强调内部块相似性，鼓励学生保持教师的成对激活相似性；
DML: 使用KL散度来匹配对等网络的软化预测；
RLKD:采用轻网络和助推网络进行协同训练；
CTKD:采用零工教师和专家教师来培养一个紧凑的学生，利用从零开始教师的最终逻辑和专家级教师的中级注意力特征，共同培养学生；
FFL:将特征融合和KD相结合，提高了子网和融合模块的性能；
OKDDip：它采用不同的对等网络进行在线知识发现，执行两级蒸馏，其中基础蒸馏作为多样性保持的组蒸馏与几个辅助对等体一起工作，第二级蒸馏将多样性增强的组知识转移到最终的学生模型，称为组领导者；

---

3.2 在CIFAR上的结果

CIFAR数据集是一种比较常用的用于验证KD性能的数据集。

首先，我们在CIFAR-10和CIFAR-100数据集上训练先验教师，分别达到95.06%和76.35%。第二，同时训练后验教师和学生，以生成中间特征图和最终输出。后一位>老师和学生在前一位老师的监督下，通过相互教学来学习。

表1展示了不同知识蒸馏方法在CIFAR-10和CIFAR-100上的比较结果：

前三行代表了三种经典的离线知识蒸馏方法。我们可以观察到SPKD可以比之前的在线知识蒸馏方法取得更好的准确率。STKD分别在CIFAR-10上比Base和SPKD提高了1.19%和0.21%,实现了最好的结果。此外,STKD在两个CIFAR数据集上一致优于所有在线知识蒸馏方法。具体来说,与DML、RLKD、CTKD和FFL相比,STKD在CIFAR-10上的准确率提高可达0.97%、0.43%、0.27%和0.61%,在CIFAR-100上的准确率提高可达1.88%、1.09%、0.45%和0.2%。
STKD与CTKD的主要区别在于后者以单向离线的方式传输中间知识。相比之下,我们的STKD方法在传递中间知识时进行教师教学和学生学习之间的互动学习。这意味着两个并行网络都可以根据实际需求深挖隐藏信息以互相学习,从而优化教师和学生网络的性能；
DML和FFL不能从中间层提取有价值的特征和知识。我们的STKD方法充分利用了中级知识。

下图则是展示了不同蒸馏方法的识别错误率对比：

可以观察到,所提出的STKD一致优于其他方法。此外,STKD在CIFAR-100上的改进显著,将学生网络的准确率提高了基准的2.7%。与以前的知识蒸馏算法不同,STKD方法在CIFAR-100上的性能提升远大于CIFAR-10。因为该方法旨在从两个老师模型中提取全面的知识,所以它在更具挑战性的数据集上表现出色。

下图显示了在两个数据集上的分类准确度和测试损失：

直观地说，我们提出的STKD极大地提高了学生网络的性能，因为我们传输双向知识来生成指导性的指导；
通过上面两张图也可以看出，本文所提出的STKD方法可以适用于各种复杂数据集

---

3.3 在SVHN和ImageNet上的结果

采用WRN-40-1作为教师网络，WRN-16-1作为学生网络，实验结果如下所示：

可以看出，学生网络始终比相应的基础模型更加优秀。

为了进一步研究STKD在大规模数据集上的性能，使用ILSVRC 2012数据集来进行实验：

resnet152被指定为教师网络，而resnet50被指定为学生。可以看出来，STKD方法比起其他KD方法都有着更高的性能。

---

3.4 消融实验

(1) 组件贡献分析：执行消融实验来调查STKD方法中每个组件的有效性。结果入下表所示：

案例A代表我们从零开始培养学生网络。案例B是标准的KD方法，案例C应用L2损失将学生logit回归到目标值；案例D表明我们引入了中间表征来监督学生的训练。案例E声明学生通过我们的STKD接受培训。
通过CD对比可以体现中间特征改善KD性能的优异能力；
通过DB对比可以体现组合教师被用来通过利用多级特征表示来更好地监督学生的训练；
通过DE对比可以体现使用在线KT和离线KT的培训为学生提供了全面的知识

---

3.5 效率分析

比较了本文的方法与各种KD机制的训练时间，结果如下表所示：

这表明STKD可以在提高网络分类准确率的同时减少网络训练所花费的时间成本。其次,在STKD的指导下,学生网络相比教师网络实现了高达3倍的压缩率。

---

总结

本文提出的一种新颖的知识蒸馏方法STKD。该方法通过整合在线知识迁移和离线知识迁移到一个统一的框架中,使用后期教师模型提供中间层的中间级知识和最终层的高级知识,以及使用先验教师模型提供额外的有价值知识,同时向学生网络传递不同级别的知识信息。这种方式不仅可以提高模型性能,也可以加速训练过程。他们在多个图像分类数据集上做了实验验证,结果表明该方法可以优于当前的知识蒸馏方法。总体来说,这种方法集成了在线互学习和离线优化的方式,提供了一种新颖高效的知识蒸馏方案。

---

未来展望

就是使用更先进的在线/离线蒸馏方式来改进这一框架