TinyMIM: An Empirical Study of Distilling MIM Pre-trained Models
Abstract
- 掩模图像建模(MIM)在大型(ViTs)的预训练中具有很强的性能。然而,对于现实世界的应用程序至关重要的小型模型不能或只能从这种预训练方法中获得少量好处。在本文中,我们探索了蒸馏技术,以将基于大型MIM的预训练模型的成功转移到较小的模型。我们系统地研究了提取框架中的不同选项,包括提取目标、损失、输入、网络正则化、顺序提取等,揭示了:1)提取令牌关系比基于CLS令牌和特征的提取更有效;2) 当学生的深度与教师的深度不匹配时,作为目标的教师网络的中间层比使用最后一层的表现更好;3) 弱正则化是优选的等。有了这些发现,我们使用所有的ViT Tiny、ViT Small和ViTbase模型,在ImageNet-1K分类上实现了与scratch MIM预训练相比的显著微调精度改进,分别获得了+4.2%/+2.4%/+1.4%的增益。我们的TinyMIM基本大小模型在AE20K语义分割中实现了52.2mIoU,比MAE基线高+4.1。我们的微小尺寸TinyMIM模型在ImageNet-1K图像分类上实现了79.6%的前1级准确率,创下了相同尺寸和计算预算的小视觉模型的新纪录。这种强大的性能为开发小型视觉 Transformer 模型提供了一种替代方法,即探索更好的训练方法,而不是像大多数以前的工作那样在架构中引入归纳偏差。代码位于https://github.com/OliverRensu/TinyMIM.
- 论文地址:[2301.01296] TinyMIM: An Empirical Study of Distilling MIM Pre-trained Models (arxiv.org)
- TinyMIM,一种对小型视觉Transformer(ViT)掩码图像建模(MIM)预训练进行知识蒸馏的成功方法;蒸馏token关系、以中间层为目标、弱正则化和顺序蒸馏对于增强性能很重要。为开发小型视觉Transformer模型探索了改进训练方法的新路,而不是在架构中引入更多归纳偏差。
Introduction
-
掩蔽图像建模(MIM)掩盖了图像区域的大部分,并训练网络来恢复掩蔽区域的原始信号,已被证明是一种非常有效的预训练视觉 Transformer 的自监督方法。由于其强大的微调性能,MIM现在已经成为视觉 Transformer 的主流预训练方法,并在这一研究领域进行了大量的跟进,例如研究如何设置解码架构、重建目标等,以及揭示其特性。
-
然而,如下表所示,MIM预训练主要影响相对较大的模型。当模型大小小到ViT Tiny(500万个参数)(这对现实世界的应用至关重要)时,MIM预训练甚至会影响ImageNet-1K分类的微调精度。事实上,与从头开始训练的对手相比,准确度下降了-0.6。这就提出了一个问题:小型模型也能从MIM预训练中受益吗?如何实现这一点?
-
-
在ImageNet1K上使用ViT-T、-S、-B和-L对MAE预训练的ViT和从头开始训练的ViTs进行比较。我们在从头开始训练时采用DeiT。我们报告的准确性排名前 1。随着模型尺寸的缩小,MAE的优势逐渐消失。MAE甚至损害了ViT-T的性能。
-
-
此外,现有关于小型视觉 Transformer 的研究主要集中在将一定的 certain inductive bias 引入架构设计中。额外的体系结构归纳偏差促进了优化,但限制了表达能力。人们很自然地会问,我们是否可以推动普通的小视觉 Transformer 发挥同样的作用。
-
在这项工作中,我们提出了TinyMIM,它回答了上述问题。TinyMIM没有使用MIM pretext 任务直接训练小型ViT模型,而是使用蒸馏技术将较大MIM预训练模型的知识转移到较小模型。蒸馏将较大的MIM预训练模型的良好特性赋予了较小的模型,同时避免了解决“太”困难的MIM任务。注意到知识提炼已经得到了很好的发展,特别是对于监督模型,我们的主要工作是首次系统地研究当使用MIM预先训练的模型作为教师时,在提炼框架中不同设计选项的影响。具体而言,我们考虑了蒸馏目标、数据扩充、网络正则化、辅助损失、宏观蒸馏策略等,并得出了几个有用的发现:
- 蒸馏目标。与提取目标相关的主要发现有两个:1)提取标记关系比提取 CLS 标记和特征图更有效。2) 使用中间层作为目标可以比使用最后一层执行得更好,并且用于不同下游任务(例如,分类和分割)的最佳目标层可以是不同的。
- 数据和网络正则化。弱增强和正则化是优选的:1)使用掩蔽图像的性能比使用原始图像差;2) 相对较小的下降路径率(教师为0,学生为0.1)表现最好。
- 辅助损失。我们发现辅助MIM损失并不能提高微调精度。
- 宏观蒸馏策略。我们发现,使用顺序蒸馏策略,即“ViT-B➡ViT-S➡ViT-T”,比直接从ViT-B蒸馏到ViT-T更好。
-
通过选择最佳的框架选项,我们使用不同大小的ViT模型,在ImageNet-1K分类上实现了与直接MIM预训练相比的显著微调精度改进,如下图所示。具体而言,TinyMIM在ViT Tiny、ViT Small和ViT基础模型上的增益分别为+4.2%/+2.4%/+1.4%。
-
-
TinyMIM(我们的)、MAE和在ImageNet-1K上使用ViT-T、-S和-B从头开始训练之间的比较。我们报告的准确性排名前1。我们在从头开始训练时采用DeiT。我们首次成功地对较小的ViT进行了掩蔽图像建模预训练。
-
-
特别是我们的TinyMIM-在微调过程中进行知识提取的T模型在ImageNet-1K分类中实现了79.6%的前1级精度(见表),这比之前所有通过引入架构归纳偏差或较小的特征分辨率来开发小型视觉 Transformer 模型的工作都表现得更好。它使用相似的模型大小和计算预算创造了新的精度记录。在ADE20K语义分割上,TinyMIM-T实现了45.0mIoU,比第二好的方法MobileViTv3-S高+1.9。TinyMIM强大的微调精度-T提出了一种开发小型视觉转换器模型的替代方法,即通过探索更好的训练方法,而不是像以前的大多数工作那样在架构中引入归纳偏差。
-
-
与最先进的微型 Transformer 和架构变体的比较。这些参数指示主干参数,该主干参数不包括分类中的最后一个分类层或分割中的解码器的参数。我们报告了ImageNet-1K分类的前1精度和ADE20K分割的mIoU。
-
Related Works
Masked Image Modeling
- 用于自监督 Transformer 预训练的掩蔽语言建模(MLM)在自然语言处理(NLP)领域取得了令人难以置信的成功。受掩蔽和重建相同思想的启发,BEiT是通过对掩蔽图像进行编码和预测DALL-E生成的掩蔽标记,将这种成功带到计算机视觉领域的先驱。SimMIM和MAE发现,重建RGB像素会产生有利的表示。MAE采用非对称编解码器结构。编码器仅对可见令牌进行编码,并丢弃高部分屏蔽令牌以减少计算负担。然后,轻量级解码器生成重构的补丁。与具有丰富语义的自然语言处理中的标记不同,计算机视觉中的像素是低级信息,因此,最近的许多工作旨在寻找更好的监督。MaskFeat将手工制作的HOG描述符产生的局部梯度特征作为监督。PeCo通过增强感知相似性来训练一种新的标记器。iBot和data2vec将指数移动平均(EMA)更新模型作为标记器。MILAN 采用了一个经过训练的CLIP作为老师。类似地,BeiTv2也使用CLIP进行标记化训练。与这些使用各种 tokenizers /教师的作品不同,我们采用了一个经过预训练的蒙面图像建模模型作为我们的教师。
- MIM预训练在从基本尺寸到巨型尺寸的相对较大的模型上表现非常好。然而,当模型小到很小的尺寸时,这将损害微调,可能是因为有限的captheMIM任务对小模型来说“太”困难了。本文通过对蒸馏技术的系统研究,探讨了如何使小型视觉 Transformer 模型也能从MIM训练中受益
Knowledge Distillation
-
知识蒸馏是一种将知识从繁琐的模型转移到小型模型的经典方法,由[Distilling the knowledge in a neural network]开创。原始知识提取框架采用教师退火的分类逻辑作为学生的提取目标。从那时起,已经进行了广泛的变体来提高提取效率,包括将提取目标更改为中间特征和关系,教师和学生的数据增强,正则化,提取策略等。
-
虽然几乎所有的研究都是针对监督环境下的CNN架构进行的,但最近,有一些工作针对视觉转换器和基于对比学习的方法。在DeiT中,教师被设置为CNN架构,以便将CNN中涉及的感应偏差转移到视觉 Transformer 中。它还建议使用硬蒸馏,该方法使用教师网络的硬伪类标签作为蒸馏目标,其性能优于朴素的知识蒸馏。在[Seed: Self-supervised distillation for visual representation]中,应用了一种关于实例之间相似性的提取方法,将对比预训练的大型CNN模型的能力转移到小型CNN。
-
在[Contrastive learning rivals masked image modeling in fine-tuning via feature distillation]中,提出了一种基于特征图提取的方法,通过不同的预训练方法来改进视觉 Transformer,包括图像分类、基于实例对比的自视觉学习和CLIP预训练。然而,对于MIM预训练模型,它没有显示出任何增益。本文首次研究了MIM预训练视觉 Transformer 的提取框架。通过系统的研究,得出了一些有用的发现和最佳选择,在这些发现和选择下,各种尺寸的视觉 Transformer 都取得了显著的效果。
Small Vision Transformers
- 近年来,设计高效的CNN模型已得到广泛研究。随着(ViT)的出现,已经有几项工作在研究如何开发高效的视觉 Transformer,其中大多数工作集中在将感应偏置引入架构中[Levit,Efficientformer,Mobilevit,Edgevits]。与这些通过在架构中引入复杂组件来开发小型视觉 Transformer 的工作不同,我们证明了小规模的普通视觉 Transformer 可以表现得同样好,甚至更好。我们的主要见解是,MIM预训练可以隐含地包含必要的归纳偏差,从而避免了显式架构偏差的需要。我们的微型平面视觉 Transformer 使用相似的模型大小和计算预算,在ImageNet-1K图像分类和ADE20K语义分割方面实现了最先进的精度。
TinyMIM
- 我们采用一个较大的MIM预训练模型作为教师,采用一个较小的ViT作为学生。TinyMIM的目标是通过模仿教师以知识提炼的方式产生的目标来训练随机初始化的学生。经过预训练后,TinyMIM预训练的模型可以转移到各种下游任务中。在这项工作中,由于其流行性和简单性,我们采用MAE作为MIM模型。
Factors
Distillation Target
-
块特征和输出特征。给定一个输入图像 x,我们首先将其划分为N个不重叠的块,并使用线性投影层将N个块映射为块嵌入 F 0 ∈ R N × D F_0 \in \R^{N×D} F0∈RN×D,其中 D 是隐藏特征的维度。假设我们有一个包含 L 个 Transformer 块的ViT。每个Transformer块以最后一个Transformer块的输出 F i − 1 F_{i−1} Fi−1 作为输入,并生成当前块的特征 F i F_i Fi,可以公式化为:
-
F i = T r a n s f o r m e r ( F i − 1 ) ; i ∈ [ 1 ; L ] F_i = Transformer(F_{i−1}); i \in[1; L] Fi=Transformer(Fi−1);i∈[1;L]
-
我们将 F i F_i Fi 称为第 i 个Transformer块的块特征。特别地,我们将最后一个Transformer块中的特征 F L F_L FL 命名为输出特征。
-
-
注意力特征和FFN特征。每个Transformer块由一个自关注层和一个前馈层组成,可以定义为:
-
H i = A t t e n t i o n ( L N ( F i − 1 ) ) ; H ^ i = H i + F i − 1 ; H ~ i = F F N ( L N ( H ^ i ) ) ; F ‾ i = H ^ i + H ~ i H_i = Attention(LN(F_{i−1}));\\ \hat H_i = H_i + F_{i−1};\\ \widetilde H_i = FFN(LN(\hat H_i));\\ \overline F_i =\hat H_i +\widetilde H_i Hi=Attention(LN(Fi−1));H^i=Hi+Fi−1;H i=FFN(LN(H^i));Fi=H^i+H i
-
其中Attention(·)、FFN(·)和LN(·)分别表示自注意层、前馈层和层范数。我们将 H ^ i \hat H_i H^i 和 H ~ i \widetilde H_i H i 称为第 i 个Transformer块的注意力特征和FFN特征
-
-
Query/Key/Value Features. 每个自关注层由M个头部网络组成,每个头部网络将输入特征 F i − 1 F_{i−1} Fi−1 映射到查询(Q)、关键字(K)和值(V):
-
Q i m = L N ( F i − 1 ) W i Q ; K i m = L N ( F i − 1 ) W i K ; V i m = L N ( F i − 1 ) W i V ; Q ^m _i = LN(F_{i−1})W ^Q _i ;\\ K ^m _i = LN(F_{i−1})W ^K _i ;\\ V ^m _i = LN(F_{i−1})W ^V _i ; Qim=LN(Fi−1)WiQ;Kim=LN(Fi−1)WiK;Vim=LN(Fi−1)WiV;
-
其中 Q i ; K i ; V i ∈ R N × D M Q_i;K_i;V_i \in \R^{N×\frac DM} Qi;Ki;Vi∈RN×MD 表示第M个头网络的 query, key and value 。 query, key and value 特征( Q i ; K i ; V i ∈ R N × D Q_i;K_i;V_i \in \R^{N×D} Qi;Ki;Vi∈RN×D)分别是M个 Q i m / K i m / V i m Q^m_i/K^m_i/V^m _i Qim/Kim/Vim 的级联。
-
-
Relations. 对于来自第 i 个 Transformer 块的第 m 个头部网络,我们可以计算其Q-Q、K-K、V-V和Q-K关系( R i ; m Q Q ; R i ; m K K , R i ; m V V ; R i ; m Q K ∈ R N × N R^{QQ}_{i;m};R^{KK} _{i;m},R^{V V}_{i;m};R^{QK}_{i;m}\in \R ^{N×N} Ri;mQQ;Ri;mKK,Ri;mVV;Ri;mQK∈RN×N),这些关系被实现为标度积关系:
- R i ; m Q Q = S o f t m a x ( Q i m Q i m T D / M ; R i ; m K K = S o f t m a x ( K i m K i m T D / M ; R i ; m V V = S o f t m a x ( V i m V i m T D / M ; R i ; m Q K = S o f t m a x ( Q i m K i m T D / M ; R^{QQ} _{i;m} = Softmax(\frac {Q^m_i {Q^m_i}^T}{\sqrt{D/M}} ;\\ R^{KK} _{i;m} = Softmax(\frac {K^m_i {K^m_i}^T}{\sqrt{D/M}} ;\\ R^{VV} _{i;m} = Softmax(\frac {V^m_i {V^m_i}^T}{\sqrt{D/M}} ;\\ R^{QK} _{i;m} = Softmax(\frac {Q^m_i {K^m_i}^T}{\sqrt{D/M}} ;\\ Ri;mQQ=Softmax(D/MQimQimT;Ri;mKK=Softmax(D/MKimKimT;Ri;mVV=Softmax(D/MVimVimT;Ri;mQK=Softmax(D/MQimKimT;
-
第 i 个 Transformer 组的 Q-Q/K-K/V-V/Q-K关系( R i Q Q , R i K K , R i V V , R i Q K ∈ R M × N × N R^{QQ}_ i,R^{KK}_ i,R^{V V}_ i,R^{QK}_i \in \R^{M×N×N} RiQQ,RiKK,RiVV,RiQK∈RM×N×N)是 M 个 R i , m Q Q / R i ; m K K / R i ; m V V / R i ; m Q K R^{QQ}_ {i,m}/R^{KK}_{i;m}/R^{V V}_{i;m}/R^{QK}_{i;m} Ri,mQQ/Ri;mKK/Ri;mVV/Ri;mQK的级联。
Input
- MIM模型在输入图像 x 上随机掩蔽高比例的图像块,产生用于预训练的掩蔽图像 x ~ \widetilde x x 。我们还研究了TinyMIM在执行知识提取时的输入——可以是原始图像 x 或掩码图像 x ~ \widetilde x x 。
Target Block
- 考虑这样一种情况,即我们倾向于使用包含24个块的MAE预训练的 ViT-L(教师)来提取包含12个块的ViTB(学生)。在这种情况下,学生的块号与老师的块号不匹配。我们调查教师的哪一块可以提供最合适的目标。所选择的块被称为目标块。
Knowledge Distillation as MIM Pre-training
-
在第3.1.1节中,我们描述了各种候选蒸馏目标。在本节中,我们将介绍不同蒸馏目标的不同知识蒸馏损失。设 x 表示输入图像, f t f_t ft 和 f s f_s fs 分别表示教师模型和学生模型。知识提取的目标是通过在冻结 f t f_t ft 的同时优化 f s f_s fs,将知识从 f t f_t ft 转移到 f s f_s fs。通常,训练由KL散度监督,KL散度定义为:
-
L K L ( p ; t ) = t l o g t p ; \mathcal L_{KL}(p; t) = tlog\frac t p ; LKL(p;t)=tlogpt;
-
其中 t 表示由 f t ( x ) f_t(x) ft(x) 生成的目标,p是由 f s ( x ) f_s(x) fs(x) 产生的预测。
-
-
类别标记蒸馏。我们使用 c t c_t ct 和 c s c_s cs 分别表示 f t f_t ft 和 f s f_s fs 的类标记特征。类令牌蒸馏的损失公式为:
- L = L K L ( c s ; c t ) \mathcal L=\mathcal L_{KL}(c_s;c_t) L=LKL(cs;ct)
-
Feature Distillation. 一般来说,教师网络和学生网络的特征维度不匹配。为了解决这个问题,我们在学生网络的输出上采用了一个额外的线性层来匹配教师目标的特征维度。设 F t F_t Ft 和 F s F_s Fs 分别表示目标特征和学生产生的预测,然后是线性投影层。我们可以将特征蒸馏的损失公式化如下:
-
L = L 1 ( F s ; N o r m ( F t ) ) \mathcal L =\mathcal L_1(F_s; Norm(F_t)) L=L1(Fs;Norm(Ft))
-
其中,Norm(·)是由无隶属的层范数实现的白化运算,L1是平滑的L1损失,
-
-
其中β设置为2.0。
-
-
关系蒸馏。这是我们默认的知识提取策略,如下图所示。为了清晰起见,我们使用 R t ➡ m Q K R ^{QK}_ {t➡m} Rt➡mQK 来表示来自教师网络的第m个头生成的Q-K关系目标(参见等式4),以及 R s ➡ m Q K R^{QK}_ {s➡m} Rs➡mQK 来表示来自学生网络的相应的Q-K关系预测。
-
-
TinyMIM默认的知识提取策略。对学生(例如ViT-B)进行优化,以模拟由MIM预训练教师(例如ViT-L)的中间块以原始图像作为输入生成的关系。我们用自适应块替换学生的最后一个块,以匹配老师的数(没有额外的计算成本)。经过预训练(知识提炼)后,学生模型可以转移到各种下游任务中。
-
-
我们定义 R t ➡ m V V R^{V V}_{t➡m} Rt➡mVV 和 R s ➡ m V V R^{V V}_{s➡m} Rs➡mVV 以类似的方式。关系损失蒸馏公式如下:
-
Head Alignment for Relation Distillation. 通常情况下,学生网络的头数低于教师网络。例如,ViT-L(教师)每个块包含16个头,而ViT-B(学生)每个块仅包含12个头。回想一下,关系蒸馏损失是逐头计算的,因此在执行关系蒸馏之前,我们必须解决头错位问题。为此,我们将学生的最后一个区块替换为自适应区块,该区块保留了原始的隐藏维度,但根据老师调整了head数。具体地,给定每个块具有 M t M_t Mt 个头的教师网络和每个块具有 M s M_s Ms 个头、隐藏维度 D s D_s Ds 和头维度Ds=Ms的学生网络,自适应块被设计为每个块具有 M t M_t Mt 个头、隐蔽维度 D s D_s Ds 以及头维度Ds=8的Transformer块。
Sequential Distillation
- 当训练像ViT-S这样的小模型时,教师有两种选择:预先训练的ViT-B和预先训练的ViTL。直观地说,经过预训练的ViT-L是一位好老师,因为它具有更高的表现能力。然而,ViT-L和ViT-S之间存在巨大的容量差距,导致蒸馏结果不佳。在[On the efficacy of knowledge distillation,Born again neural networks]之后,我们采用顺序蒸馏策略来改进预训练。例如,当预训练ViT-S时,教师被选择为TinyMIM预训练的ViT-B,该ViT-B已经由TinyMIM以ViT-L作为教师进行训练。
Experiments
Implementation Details
-
Pre-training. 所有模型都在ImageNet-1K训练集上的100epoch时间表下进行预训练。我们使用4096的批量大小和lr=1.5e-4×batchsize/256的学习率。我们采用了余弦衰减时间表,其中包括5个时期的预热。我们采用了AdamW优化器,其权重衰减为0.05。我们使用随机调整大小的裁剪随机水平翻转,颜色抖动仅适用于学生。输入大小设置为224×224。
-
Fine-tuning. 我们将TinyMIM预训练模型转移到ImageNet图像分类和ADE20K语义分割中。对于ImageNet,我们使用权重衰减为0.05的AdamW优化器。对于数据扩充,我们遵循MAE中的设置。我们对100个时期的ViT-B进行微调,批量大小为1024,学习率为2e-3,下降路径率为0.1。我们对200个时期的ViT-S和ViT-T进行微调,批量大小为2048,学习率为5e-3, drop path rate 为0.1。对于ADE20K,我们遵循MAE中的相同设置,并采用UperNet作为我们的框架,使用TinyMIM预训练的主干。输入图像分辨率为512×512,用于训练和评估。我们使用mIoU作为评估指标。
-
此外,我们评估了TinyMIM在各种域外ImageNet数据集上的稳健性,这些数据集是通过在ImageNet上应用不同的扰动生成的,例如自然对抗性示例(ImageNet-A)、语义转移(ImageNet-R)、常见图像损坏(ImageNetC)。我们报告了ImageNet-A/R的前1精度和ImageNet-C的mCE误差(越低越好)。
-
默认设置。默认情况下,我们采用上等式中公式化的关系蒸馏。头部对齐,原始图像作为输入,顺序蒸馏和MAE预训练的ViT-L的第18个块作为TinyMIM-ViT-B预训练的目标块。
Main Results
- 如下表所示,我们将我们的TinyMIM与以前使用不同ViT的ImageNet图像分类和ADE20K语义分割方法进行了比较。特别是,TinyMIM预训练的ViT-T实现了75.8%的top1准确率,比MAE基线高出+4.2。一个名为TinyMIM的增强模型-T、 保留了ViT-T的简单架构和计算预算,进一步实现了79.6%的前1级精度。有关TinyMIM的详细信息,请参阅附录-T.此外,TinyMIM预训练的ViT-S实现了83.0%的前1级精度,分别比MAE基线和先前的最佳方法CIM高+2.4和+1.4。通过转移MAE预训练的ViT-L的知识,TinyMIM预训练的ViT-B在ImageNet-1K上实现了85.0%的前1级准确率。
-
-
ImageNet-1K和ADE20K上的微调结果。所有模型都在ImageNet-1K上进行了预训练。“标记器/教师数据”:教师和标记器的训练数据。+: 使用官方代码复制结果。*:该模型使用DeiT风格的知识提取对1000个时期进行了微调。++: 该模型在ADE20K分割微调之前对ImageNet-1K分类进行中间微调。
-
- 在语义分割方面,TinyMIM预训练的ViTB分别比MAE基线和最先进的CAE高出+4.1和+2.0。在ADE20K分割微调之前对ImageNet-1K分类进行中间微调进一步提高了性能。我们还在下表中的域外数据集上评估了我们的模型。我们的TinyMIM预训练模型比MAE预训练模型更稳健。具体而言,TinyMIM-ViT-B在ImageNet-A和ImageNet-R上分别比MAE-ViT-B高+6.4和+4.6,并将mCE降低-5.1。
-
-
域外数据集的稳健性评估。
-
Ablation Study
-
除非另有规定,所有消融研究均在TinyMIM-ViT-B上进行,教师为MAE预训练的ViT-L,关系提取策略,原始图像作为输入,ViT-L的第18个块作为目标块,在100个epoch的预训练计划下。我们在ImageNet-1K上报告了排名前1的准确性。
-
类别标记蒸馏。对于这种蒸馏策略,我们研究了两种变体:1)如等式所示的类标记蒸馏;2) 类令牌蒸馏具有额外的MAE重建损失。结果如下表所示。这两种变体的表现都比MAE基线差,表明 cls 令牌不适合作为提炼目标,因为在教师的预训练期间没有对课堂令牌进行明确的监督。
-
-
研究公式中的 cls 标记蒸馏。
-
-
特征提取。如第3.1.1节所述,有四种类型的特征可以作为公式的特征提取的目标:输出特征、FFN特征、注意力特征和Q/K/V特征。下表比较了使用不同特征作为蒸馏目标的结果。我们还报告了残差连接之前的FFN特征和注意力特征的结果。一个有趣的发现是,在残差连接之后提取FFN特征和注意力特征会显著降低性能。
-
-
特征蒸馏的研究。有关不同特征的定义,请参见第3.1.1节和等式2。
-
-
关系蒸馏。Eq.9公式我们的默认关系蒸馏,它联合蒸馏Q-K关系和V-V关系(见Eq.4)。在这里,我们通过将目标关系从Q-K/V-V更改为Q-K/K-K/V-V来研究一个变体。我们还研究了是否在每个关系上应用Softmax算子。结果如下表所示。
-
-
研究方程中公式化的关系蒸馏。关于不同关系的定义,请参见第3.1.1节和等式4。
-
-
不同蒸馏策略的比较。在这项研究中,所有模型都是在300 epoch 时间表下进行预训练的。我们比较了ImageNet图像分类和ADE20K语义分割上的三种提取策略。对于每种策略,我们都使用产生最佳结果的目标。我们还强调了MAE基线的改进。
-
-
ImageNet-1K图像分类中三种提取策略的比较。模型在300 epoch 时间表下进行预训练。
-
-
三种提取策略在ADE20K语义分割中的比较。模型是在300 epoch 时间表下进行预训练的。
-
-
Target Block. 如第3.1.3节所述,我们考虑学生的块号与老师的块号不匹配的情况。在这里,我们使用包含24个块的MAE预训练的ViT-L来提取包含12个块的ViT-B。在这里,我们检查使用ViT-L的第12、第15、第18、第21和第24(最后)个块作为目标块的效果。比较结果如下表所示。我们通过实验发现,使用第18块产生了最好的结果。
-
-
ImageNet-1K和ADE20K上Target Block 的研究。
-
-
Sequential Distillation. 在第3.3节中,我们主张采用顺序蒸馏策略,以实现从较大模型(如ViT-L)到较小模型(如ViT-S)。下表比较了采用或不采用顺序蒸馏的不同教师的结果。我们得出两个结论:1)使用较大的教师(MAE-ViT-L)来提取较小的学生(ViT-S)会降低成绩;2) 顺序蒸馏显著提高了ViT-T的性能(MAE-ViT-B➡TinyMIM-ViT-S作为老师,ViT-T作为学生)。
-
-
顺序蒸馏的研究。
-
-
将MAE集成到TinyMIM中。MAE是一种简单但有效的自我监督预训练范式,通过要求模型预测掩蔽输入来训练模型。相比之下,TinyMIM以知识蒸馏的方式预训练较小的ViT。在这里,我们将MAE集成到TinyMIM中,生成一个集成模型。该模型在两种损失下进行了优化:TinyMIM的知识提取损失和MAE的重构损失。为了实现MAE预训练,我们随机屏蔽75%的图像补丁,并将可见补丁馈送到网络中,以启动集成模型的预训练。下表显示了TinyMIM-ViT-B与集成模型之间的比较。从表中,我们可以得出一个结论——将MAE集成到我们的TinyMIM中并不能提高性能。此外,我们还研究了TinyMIM-ViT-B的输入,它可以是原始图像,也可以是掩码图像,如表所示——将原始图像作为输入会产生更好的结果。
-
-
TinyMIM-ViT-B(第一排)和集成模型(第三排)之间的比较。我们还研究了TinyMIM-ViT-B的输入,它可以是原始图像(第一行)或掩蔽图像(第二行)。
-
-
Drop Path. Drop Path是训练transformer最关键的技术之一。使用适当的下降路径速率可以显著地缓解过拟合问题。然而,MAE在其实现中禁用了这项技术。在这里,我们验证了将Drop Path 应用于TinyMIM的效果。结果如下表所示。对于学生模型,最佳Drop Path率为0.1。对于教师模型,禁用Drop Path产生最佳结果。
-
-
教师和学生使用下降路径率(DPR)的消融研究。
-
Conclusion
- 在本文中,我们提出了TinyMIM,它是第一个成功地对较小的ViT模型执行掩蔽图像建模(MIM)预训练的系统。我们不是采用掩码和预测 pretext 任务,而是通过以知识蒸馏的方式模拟大ViT的关系来预训练小ViT。TinyMIM的成功可归因于对可能影响TinyMIM预训练的各种因素的全面研究,包括蒸馏靶、蒸馏输入和target block。通过大量的实验,我们得出了一系列的结论。例如,关系蒸馏优于特征蒸馏和类表征蒸馏;以原始图像作为输入是最优的;顺序蒸馏对于训练较小的ViT是必要的等。凭借其简单性和强大的性能,我们希望我们的方法可以作为未来研究的坚实基础。
A. Hyper-parameters
- ImageNet-1K预训练的超参数。见表。
-
-
ImageNet-1K预训练的超参数。
-
- ImageNet-1K图像分类微调的超参数。见表。
-
-
ImageNet-1K图像分类微调的超参数。*表明我们分别对100个epoch和300个epoch的预训练模型使用0.65和0.6。
-
- TinyMIM-T保留了ViT-T的简单架构和计算预算。我们微调TinyMIM?在ImageNet-1K上使用DeiT风格的知识蒸馏进行1000个epoch。在MobileNetV3之后,在分类层之前放置一个额外的完全连接层,以将特征维度从192增加到1280。头编号设置为12,而不是默认的3。ADE20K语义分割微调的超参数。见表。
-
-
ADE20K语义分割的超参数微调。
-
知识蒸馏
-
大规模机器学习和深度学习模型越来越常见。例如,GPT-3 是在 570 GB 的文本上训练的,包含 1750 亿个参数。然而,虽然训练大型模型有助于提高最先进的性能,但要在边缘设备上部署如此繁琐的模型并非易事。此外,大多数数据科学建模工作都侧重于训练一个大型模型或不同模型的集合,以便在通常不代表真实世界数据的暂留验证集上表现良好。训练目标和测试目标之间的这种不一致导致开发出的机器学习模型在经过策划的验证数据集上产生了良好的准确性,但在实际测试数据上进行推理时,往往无法达到性能、延迟和吞吐量基准。知识蒸馏通过捕捉和 “蒸馏 ”复杂机器学习模型或模型集合中的知识,将其转化为更小的单个模型,从而有助于克服这些挑战,而且更易于部署,同时又不会明显降低性能。
-
知识蒸馏指的是将知识从一个或一组庞大而笨重的模型转移到一个更小的模型的过程,而这个模型可以在现实世界的限制条件下实际部署。从本质上讲,这是一种模型压缩形式,Bucilua 及其合作者于 2006 年首次成功地展示了这一形式。知识蒸馏通常在与复杂架构(包括多个层和模型参数)相关的神经网络模型上进行。因此,随着近十年来深度学习的出现,以及它在语音识别、图像识别和自然语言处理等多个领域的成功应用,知识蒸馏技术在实际应用中的重要性日益凸显。
-
对于内存和计算能力有限的边缘设备来说,部署大型深度神经网络模型的挑战尤为严峻。为了应对这一挑战,人们首先提出了一种模型压缩方法,将大型模型中的知识转移到训练较小的模型中,而不会对性能造成任何重大损失。Hinton 及其同事将这种从大型模型学习小型模型的过程形式化为 “知识蒸馏 ”框架。在知识蒸馏过程中,小型 “学生 ”模型学会模仿大型 “教师 ”模型,并利用教师的知识获得类似或更高的准确度。下一节,将深入探讨知识蒸馏框架及其底层架构和机制。知识提炼系统由三个主要部分组成:知识、提炼算法和师生架构。
-
在神经网络中,知识通常是指学习到的权重和偏差。与此同时,大型深度神经网络中的知识来源也多种多样。典型的知识蒸馏使用对数作为教师知识的来源,而其他知识则侧重于中间层的权重或激活。其他类型的相关知识包括不同类型的激活与神经元之间的关系或教师模型本身的参数。
-
不同形式的知识可分为三种不同类型: 基于反应的知识、基于特征的知识和基于关系的知识。下图 展示了教师模型中这三种不同类型的知识。我将在下一节中详细讨论这些不同的知识来源。
-
-
Response-based knowledge
- 基于反应的知识侧重于教师模型的最终输出层。假设是学生模型将学会模仿教师模型的预测。如下图所示,这可以通过使用损失函数来实现,该函数被称为 “蒸馏损失”,分别捕捉学生模型和教师模型的对数之差。随着训练过程中损失的最小化,学生模型就能更好地做出与教师模型相同的预测。
- 在图像分类等计算机视觉任务中,软目标包括基于响应的知识。软目标代表了输出类别的概率分布,通常使用 softmax 函数进行估算。每个软目标对知识的贡献都使用一个称为温度的参数来调节。基于软目标的响应式知识蒸馏通常用于监督学习。
-
Feature-based knowledge
-
经过训练的教师模型还能捕捉到中间层的数据知识,这对深度神经网络尤为重要。中间层可以学习辨别特定特征,这些知识可用于训练学生模型。如下图所示,目标是训练学生模型学习与教师模型相同的特征激活。蒸馏损失函数通过最小化教师模型和学生模型的特征激活度之间的差异来实现这一目标。
-
-
除了神经网络输出层和中间层所代表的知识外,捕捉特征图之间关系的知识也可用于训练学生模型。这种知识形式被称为基于关系的知识,如下图所示。这种关系可以建模为特征图之间的相关性、图、相似性矩阵、特征嵌入或基于特征表示的概率分布。
-
-
-
-
训练学生和教师模型的方法主要有三种,即离线、在线和自我蒸馏。蒸馏训练方法的分类取决于教师模型是否与学生模型同时修改,如图所示。
-
-
Offline distillation
- 离线提炼是最常见的方法,即使用预先训练好的教师模型来指导学生模型。在这种方法中,教师模型首先在训练数据集上进行预训练,然后从教师模型中提炼知识来训练学生模型。鉴于深度学习的最新进展,各种预先训练好的神经网络模型都可以公开获得,它们可以根据使用情况充当教师模型。离线蒸馏是深度学习中一种成熟的技术,也更容易实现。
-
Online distillation
- 在离线蒸馏中,预训练的教师模型通常是大容量深度神经网络。对于一些使用案例,离线蒸馏可能无法使用预先训练好的模型。为了解决这一限制,可以使用在线蒸馏法,在单个端到端训练过程中同时更新教师模型和学生模型。在线蒸馏可以使用并行计算,因此是一种高效的方法。
-
Self-distillation
- 在自蒸馏过程中,教师模型和学生模型使用的是同一个模型。例如,深度神经网络深层的知识可以用来训练浅层。它可以被视为在线蒸馏的一种特殊情况,并以多种方式实例化。教师模型早期 epoch 的知识可以转移到其后期 epoch,用于训练学生模型。
-
-
师生网络结构的设计对于高效获取和提炼知识至关重要。通常情况下,较为复杂的教师模型与较为简单的学生模型之间存在模型容量差距。通过高效的师生架构优化知识传递,可以缩小这种结构性差距。由于深度神经网络的深度和广度,从深度神经网络中转移知识并不简单。最常见的知识转移架构包括一个学生模型,该模型是:
- 较浅版本的教师模型,层数较少,每层神经元数量较少、
- 教师模型的量化版本、
- 具有高效基本操作的小型网络、
- 具有优化全局网络结构的较小网络、
- 与教师模型相同的模型。
-
除上述方法外,神经结构搜索等最新技术也可用于在给定特定教师模型的情况下设计最佳学生模型结构。
-
Adversarial distillation
- 最近在生成式对抗网络中提出的对抗学习概念,用于训练生成器模型和判别器模型,生成器模型学习生成尽可能接近真实数据分布的合成数据样本,判别器模型学习判别真实数据样本和合成数据样本。这一概念已被应用于知识蒸馏,使学生和教师模型能够更好地学习真实数据分布的表示方法。
- 为了实现学习真实数据分布的目标,对抗学习可用于训练生成器模型,以获得合成训练数据,并将其作为真实数据使用或增强原始训练数据集。第二种基于对抗学习的提炼方法主要是基于对数或特征图的判别模型来区分学生和教师模型的样本。这种方法有助于学生很好地模仿教师。第三种基于对抗学习的蒸馏技术侧重于在线蒸馏,在这种方法中,学生模型和教师模型被联合优化。
-
Multi-Teacher distillation
-
如下图所示,在多教师提炼法中,学生模型从多个不同的教师模型中获取知识。使用教师模型组合可以为学生模型提供不同种类的知识,这些知识可能比从单一教师模型获取的知识更有益。
-
-
来自多个教师的知识可以合并为所有模型的平均响应。教师通常传授的知识类型基于对数和特征表征。多名教师可以传授不同类型的知识。
-
-
Cross-modal distillation
- 下图显示了跨模态提炼训练方案。在这里,教师接受一种模态的训练,然后将其知识提炼给需要不同模态知识的学生。在训练或测试过程中,如果没有特定模式的数据或标签,就会出现这种情况,因此需要跨模式传授知识。
-
基于图的提炼法利用图形来捕捉数据内部的关系,而不是从教师到学生的单个实例知识。图有两种用途–作为知识传递的一种手段,以及控制教师知识的传递。在基于图的蒸馏中,图的每个顶点代表一个自我监督的教师,它可以分别基于基于反应或基于特征的知识,如对数和特征图。
-
基于注意力的蒸馏法是利用注意力图从特征嵌入中转移知识。
-
由于隐私、安全或保密原因,在没有训练数据集的情况下,无数据提炼基于合成数据。合成数据通常由预训练教师模型的特征表示生成。在其他应用中,GAN 也用于生成合成训练数据。
-
量化蒸馏用于将知识从高精度教师模型(如 32 位浮点)转移到低精度学生网络(如 8 位)。
-
终身蒸馏是基于持续学习、终身学习和元学习的学习机制,即先前学习的知识被积累并转移到未来的学习中。
-
基于神经架构搜索的蒸馏方法可用于识别合适的学生模型架构,从而优化对教师模型的学习。
-
知识蒸馏在计算机视觉领域的应用非常广泛。最先进的计算机视觉模型越来越多地基于深度神经网络,这些模型可以从模型压缩部署中获益。知识蒸馏技术已成功应用于以下领域:image classification, face recognition, image segmentation, action recognition, object detection, lane detection, pedestrian detection, facial landmark detection, pose estimation, video captioning, image retrieval, shadow detection, text-to-image synthesis, video classification, visual question answering
- 知识蒸馏还可用于跨分辨率人脸识别等细分用例,在这些用例中,基于高分辨率人脸教师模型和低分辨率人脸学生模型的架构可提高模型性能和延迟。由于知识蒸馏可以利用不同类型的知识,包括跨模态数据、多领域、多任务和低分辨率数据,因此可以针对特定的视觉识别用例训练各种蒸馏学生模型。
-
鉴于大容量深度神经网络(如语言模型或翻译模型)的盛行,知识提炼在 NLP 应用中的应用尤为重要。最先进的语言模型包含数十亿个参数,例如,GPT-3 包含 1750 亿个参数。这比之前最先进的语言模型 BERT(基础版本包含 1.1 亿个参数)高出几个数量级。因此,知识蒸馏在 NLP 中非常流行,它可以获得快速、轻量级的模型,训练起来更容易,计算成本也更低。除语言建模外,知识蒸馏还可用于 NLP 用例,例如:neural machine translation, text generation, question answering, document retrieval, text recognition
- 通过知识蒸馏,可以获得高效、轻量级的 NLP 模型,这些模型在部署时对内存和计算要求较低。学生-教师训练还可用于解决多语言 NLP 问题,多语言模型中的知识可以相互转移和共享。
- DistilBERT 是 Hugging Face 开发的一种更小、更快、更便宜、更轻便的 BERT 模型 。在这里,作者预训练了一个较小的 BERT 模型,该模型可以在各种 NLP 任务上进行微调,并具有相当高的准确性。在预训练阶段应用了知识蒸馏法,以获得一个蒸馏版的 BERT 模型,该模型的体积小了 40%(6600 万个参数对 1.1 亿个参数),速度快了 60%(GLUE 情感分析任务的推理时间为 410s 对 668s),同时保留了相当于原始 BERT 模型 97% 准确率的模型性能。在 DistilBERT 中,学生采用了与 BERT 相同的架构,并使用了新颖的三重损失(triplet loss),该损失结合了语言建模、蒸馏和余弦距离损失。
-
最先进的语音识别模型也基于深度神经网络。现代 ASR 模型采用端到端训练,其架构包括卷积层、具有注意力的序列到序列模型以及最近的 Transformer。对于设备上的实时语音识别,最重要的是获得更小、更快的模型,以获得有效的性能。语音知识蒸馏有多种用途:speech recognition, spoken language identification, audio classification, speaker recognition, acoustic event detection, speech synthesis, speech enhancement, noise-robust ASR, multilingual ASR, accent detection
- Parthasarathi 和 Strom(2019)利用师生训练为 100 万小时的无标注语音数据生成软目标,而训练数据集仅包括 7000 小时的标注语音。教师模型生成了所有输出类别的概率分布。学生模型在相同的特征向量下也会产生输出类别的概率分布,目标函数优化了这两种分布之间的交叉熵损失。在这里,知识蒸馏有助于简化在大量语音数据语料库中生成目标标签的过程。