SSKD:自我监督知识蒸馏

动机:希望探索一种更通用和模型不可知的方法,从预训练的教师模型中提取“更丰富的隐藏知识”。看似不同的自我监督任务可以作为一个简单但强大的解决方案。

简介:利用具有选择性传输的噪声自监督信号进行蒸馏的实用方法。自监督信号改善了传统蒸馏,在少镜头和噪声标签场景下具有显著的增益。

SSKD:自我监督知识蒸馏_第1张图片

       对正常数据和单个分类任务的模拟扩展到对转换数据的模拟,并使用额外的自我监督任务。教师的自我监督预测包含丰富的结构化知识,可以促进学生更全面的知识提炼。在本例中,对变换图像的对比学习作为自我监督的任务。它通过图像变换t(·)构造单个正对和多个负对,然后鼓励网络识别正对。教师和学生的骨干分别表示为f_{t}f_{s},而相应的输出表示为t和s,下标表示索引。

1、方法模型 

1.1先验方法

1.1.1知识蒸馏

       与硬标签(一个热矢量)相比,优化良好的教师模型预测的软目标可以提供额外的信息。分配给错误类别的相对高概率编码了不同类别之间的语义相似性。强迫学生模仿老师的预测会导致学生学习到这些无法用硬标签单独表达的次要信息。为了获得软目标,引入了温度标度,以软化峰值软最大分布:

p^{i}\left ( x;\tau \right )=Softmax\left ( s\left ( x \right );\tau \right )=\frac{e^{s_{i}\left ( x \right )/\tau }}{\sum_{k}e^{s_{k}\left ( x \right )/\tau }}

其中x是数据样本,i是类别指数,s_{i}\left ( x \right )是x在类别i上获得的得分logit,\tau是温度。通过KL发散度测量的知识蒸馏损失L_{kd}为:

L_{kd}=-\tau ^{2}\sum_{x\sim D_{x}}\sum_{i=1}^{C}p_{t}^{i}\left ( x;\tau \right )\log \left ( p_{s}^{i} \left ( x;\tau \right )\right )

其中t和s分别表示教师和学生模型,C表示类的总数,D_{x}表示数据集。学生模型的完全损失函数L是标准交叉熵损失L_{ce}和知识蒸馏损失L_{kd}的线性组合:

L=\lambda _{1}L_{ce}+\lambda _{2}L_{kd}

其中λ1和λ2是平衡参数。

1.1.2基于对比预测的自监督任务

基于对比预测方法在自我监督学习中的成功,在框架中采用了对比预测作为自我监督任务。对比预测的总体目标:通过潜在空间中的对比损失最大化数据点与其转换版本之间的一致性。

给定一个包含N个数据点\left \{ x_{i} \right \}_{i=1:N}的小批量,我们对每个数据点应用独立变换t(·)(从同一分布t中采样),并获得\left \{ \widehat{x_{i}} \right \}_{i=1:N}x_{i}\widehat{x_{i}}都被送入教师或学生网络,以提取特征\phi _{i}=f\left ( x_{i} \right ),\widehat{\phi _{i}}=f\left ( \widehat{x_{i}} \right )。在网络顶部添加投影头。投影头是一个2层多层感知器。它将表示映射到应用对比损失的潜在空间,即z_{i}=MLP\left ( \phi _{i} \right ),\widehat{z_{i}}=MLP\left ( \widehat{\phi _{i}} \right )

     将\left ( \widehat{x_{i}} ,x_{i}\right )作为正对\left ( \widehat{x_{i}},x_{k} \right )_{k\neq i}作为负对

给定一些\widehat{x_{i}},对比预测任务是从集合\left \{ x_{i} \right \}_{i=1:N}中识别相应的x_{i}。网络最大化正对之间的相似性,最小化负对之间的类似性。使用余弦相似性。将\left \{ \widehat{x_{i}} \right \}\left \{ x_{i} \right \}之间的相似性组织成矩阵形式A

A_{i,j}=\cos \left ( \widehat{z_{i}},z_{j} \right )=\frac{dot\left ( \widehat{z_{i}} ,z_{j}\right )}{\left \| \widehat{z_{i}} \right \|_{2}\left \| z_{j} \right \|_{2}}

其中A_{i,j}表示\widehat{x_{i}}x_{j}之间的相似性。对比预测的损失是

L=-\sum_{i}\log \left ( \frac{\exp \left ( \cos \left ( \widehat{z_{i}} ,z_{i}\right ) /\tau \right )}{\sum_{k}\exp \left ( \cos \left ( \widehat{z_{i}},z_{k} \right ) /\tau \right )} \right )=-\sum_{i}\log \left ( \frac{\exp \left ( A_{i,i} /\tau \right )}{\sum_{k}\exp \left ( A_{i,k}/\tau \right )} \right )

其中\tau另一个温度参数(不同于蒸馏中的\tau)。损失形式类似于softmax损失,可以理解为最大化\widehat{z_{i}}z_{i}来自正对的概率。在匹配\left \{ \widehat{x_{i}} \right \}\left \{ x_{i} \right \}的过程中,网络学习变换不变特征

       然而,在SSKD中,主要目标不是学习变换不变的表示,而是利用对比预测作为辅助任务,从教师模型中挖掘更丰富的知识

1.2SSKD学习

SSKD:自我监督知识蒸馏_第2张图片

该框架的教师和学生网络都包含三个组件,即主干f(·)、分类器p(·)和SS模块c(·,·)。教师培训分为两个阶段。

第一阶段:通过分类任务训练f_{t}\left ( \cdot \right )p_{t}\left ( \cdot \right )

第二阶段:通过自我监督任务微调c_{t}\left ( \cdot ,\cdot \right )

在学生培训中,除了标准标签丢失外,我们还强制学生在分类输出和自我监督输出上模仿老师。

对比预测被选为SS任务,因此SS模块c_{t}\left ( \cdot ,\cdot \right )c_{s}\left ( \cdot ,\cdot \right )两层MLP和相似度计算模块组成。

1.2.1教师网络训练

       输入为正常数据\left \{ x_{i} \right \}和转换版本\left \{ \widehat{x_{i}} \right \}。变换t(·)从预定义的变换分布t中采样。选择了四种变换,即颜色下降、旋转、裁剪,然后调整大小和颜色失真。我们将x\widehat{x}喂入到主干网络,并获得它们的特征\phi _{i}=f\left ( x_{i} \right ),\widehat{\phi _{i}}=f\left ( \widehat{x_{i}} \right )

       教师网络训练包括两个阶段。

第一阶段:利用分类损失对网络进行训练。仅更新主干f_{t}\left ( \cdot \right )和分类器p_{t}\left ( \cdot \right )

ps:分类损失不在变换数据\widehat{x}上计算,因为变换T比通常的数据增强更重。它的目标不是扩大训练集,而是使\widehat{x}视觉上不太类似于x

第二阶段:固定了f_{t}\left ( \cdot \right )p_{t}\left ( \cdot \right ),并且仅使用等式中的对比预测损失更新SS模块中c_{t}\left ( \cdot ,\cdot \right )的参数,如上面的对比预测损失公式所示。

第一阶段简单地是用于分类的网络的典型训练。第二阶段,旨在调整SS模块,以使用现有主干的特征进行对比预测。

这允许我们从SS模块中提取知识进行蒸馏。值得指出的是,由于MLP网络较小,第二阶段训练效率较高,因此很容易为SSKD准备教师网络。

1.2.2学生网络训练

在训练教师的SS模块后,我们将softmax(具有温度标度\tau)沿行维度应用于教师的相似性矩阵A引出概率矩阵B^{t},其中B_{i,j}^{t}表示\widehat{x_{i}}x_{j}正对的概率。类似的操作应用于学生以获得B^{s}。使用B^{t}B^{s},可以计算教师和学生的SS模块输出之间的KL散度损失

L_{ss}=-\tau ^{2}\sum_{i,j}B_{i,j}^{t}\log \left ( B_{i,j}^{s} \right )

转换后的数据点\widehat{x}是对比预测任务的副产品。虽然我们不要求学生正确分类,但我们可以鼓励学生的分类器输出p_{s}\left ( f_{s}\left ( \widehat{x} \right ) \right )接近教师的输出。损失函数为:

L_{T}=-\tau ^{2}\sum_{\widehat{x}\sim \tau \left ( D_{x} \right )}\sum_{i=1}^{C}p_{t}^{i}\left ( \widetilde{x} ;\tau \right )\log \left ( p_{s}^{i} \left ( \widetilde{x} ;\tau \right )\right )

学生网络的最终损失是上述项的组合,即方程中的交叉熵损失L_{ce}L_{kd}L_{ ss}L_{T}:

L=\lambda _{1}L_{ce}+\lambda _{2}L_{kd}+\lambda _{3}L_{ss}+\lambda _{4}L_{T}

其中\lambda _{i}是平衡参数。

1.3自我监督提升

       当执行对比预测时,教师可能会产生不准确的预测,例如,将x_{k}分配给\widetilde{x_{i}}i\neq k。因为教师的骨干网络没有与SS模块一起进行对比预测微调。与传统的知识蒸馏类似,教师为错误答案分配的相对概率包含教师丰富的知识。转移这些不准确但结构化的知识是SSKD的核心。

       根据经验发现,极不正确的预测仍然可能误导学生的学习。为了改善这些异常值的负面影响,采用启发式方法执行选择性迁移。具体而言,将预测的误差水平定义为分类任务中相应真实标签的排名。给定变换后的样本\widetilde{x_{i}}和相应的正配对指数i,我们将网络分配给每个\left \{ x_{i} \right \}_{i=1:N}的分数按降序排序x_{i}的秩表示\widetilde{x_{i}}预测的误差水平。

       秩为1表示预测完全正确。较低的秩指示较高的误差程度。在学生的训练过程中,我们根据教师预测的误差水平,以升序对小批量中的所有\widetilde{x}进行排序,并仅传输所有正确预测和排名前k%的错误预测。这种策略抑制了教师预测中的潜在噪音,并只传递有益的知识

2、总结 

1、提出了一个新的框架,称为SSKD,这是首次尝试将自我监督与知识蒸馏相结合。采用对比预测作为辅助任务,帮助从教师网络中提取更丰富的知识。

2、设计了一种选择性迁移策略来抑制教师知识中的噪声

3、通过使用各种架构在CIF AR100和ImageNet上进行彻底的实验来检验我们的方法。

4、SSKD可以使学生更接近老师,并在少量镜头和嘈杂标签场景下工作良好。

你可能感兴趣的:(知识蒸馏,迁移学习,机器学习,人工智能,计算机视觉,深度学习)