UNIC:
Universal Classification Models via Multi-teacher Distillation
Relative gains using our UNIC encoder distilled from four teachers (DINO, DeiT-III, iBOT, dBOT-ft), over the respective best teacher for each task using a single encoder and no task-specific parameters. All models (UNIC and teachers) are trained on ImageNet-1K.
Introduction
背景与动机:近年来,预训练模型在多种任务上取得了显著成果,成为了一种普遍使用的工具。这些模型通常共享相同的架构,有时甚至使用相同的训练数据,但根据学习策略的不同,它们在特定的视觉识别场景中表现出色。例如,自监督学习模型在迁移学习(即对新类别的泛化)方面表现出色,而使用掩码建模技术训练的模型更适合于像素级任务。同时,在有标注数据的预训练情况下,监督学习在特定分类任务上仍然表现最佳。
研究目标:本文旨在学习一个通用的编码器,该编码器能够在广泛的分类任务中实现强大的泛化能力。具体来说,除了ImageNet分类任务(所有教师模型均在此数据集上训练,学生模型在此数据集上进行蒸馏)之外,本文还关注于新类别的分类、新领域的分类,以及诸如语义分割或深度估计等密集预测任务。目标是学习一个单一的编码器,可以直接应用于所有这些任务,而无需除每个分类任务的线性分类器之外的任务特定参数。
方法概述:本文提出使用多教师蒸馏的方法,通过结合多种具有互补优势的专门化教师模型来训练一个编码器,该编码器旨在匹配或超越每个任务的最佳教师模型。通过从多个教师中彻底分析蒸馏过程,并在各种任务上评估模型,包括ImageNet-1K及其15个迁移数据集上的图像级分类任务,以及语义分割和深度估计等像素级分类任务,本文逐步开发了一种方法,该方法在多个任务和维度上展示了改进的泛化能力。
主要改进:
可扩展的projector(Expendable Projectors):
引入了教师特定的可扩展projector,这些projector 位于编码器输出和损失计算空间之间,增加了中间特征在蒸馏过程中的影响力。这些projector也被设计成信息高速公路,使信号从中间层更直接地传播到蒸馏损失中。
教师丢弃(Teacher Dropping):
提出了一种正则化机制,即教师丢弃,用于在多教师蒸馏中更好地平衡各教师的影响,从而在蒸馏模型原本表现不佳的任务上取得显著增益。
**实验结果:**本文展示了通过从四个强大的ViT-Base教师模型(DINO、DeiT-III、iBOT、dBOT-ft)中蒸馏得到的通用编码器,在所有考虑的任务上均表现出色。
主要贡献:
对ViT编码器的多教师蒸馏进行了全面分析,并利用这些发现来改进蒸馏过程和学生的泛化能力。
提出了像可扩展projector和教师丢弃正则化等简单但关键的改进,使学生能够保持或改进多个不同任务上最佳教师的性能。
进行了广泛的评估,以验证所提出的模型在多个泛化轴上的表现,并研究了蒸馏编码器如何利用其权重和特征空间。
Method
1. 多教师蒸馏的基本框架
教师网络:T = {T1, . . . , T𝑀 },每个教师网络都是ViT的编码器;输入x,输出𝑑-维特征向量 𝒚𝑡,𝑖 = 𝑓𝑡 (𝒙; 𝑖);
学生网络:学生编码器输出添加一个针对每位教师的projector head ℎ𝑡,该projector head将每个标记转换为特定于教师的表示 ℎ𝑡(𝒛𝑖)。
然后,基于对应projector head的输出 ℎ𝑡(𝒛𝑖) 计算每位教师的损失。
这些projector head是消耗品,即在知识蒸馏后会被移除,并不属于学生编码器的一部分。
projector head的目标是辅助学习过程。
将projector head设置为具有两个线性层、GeLU 非线性和隐藏维度 𝑑ℎ = 4𝑑 的多层感知机(MLP),其中 𝑑 是特征维度;将在后续章节中进一步分析这些projector head。
使用两种常见的蒸馏损失的组合:余弦损失和平滑-ℓ1损失;来自教师𝑡的标记𝑖的损失由以下公式给出:
这个损失是分别针对CLS和每个patch token P计算的。为了得到最终的损失,将所有教师的损失相加,同时也包括CLS标记𝑐和所有patch token 的损失:
| P | 是patch tokens 的个数。
2. 协议总结
首先,概述了本节分析中使用的实验协议。
**数据集和主干网络。**为了更好地隔离不同蒸馏组件的影响,对所有教师和学生使用相同的训练数据和架构,即分别使用ImageNet-1K数据集和ViT-Base。在蒸馏过程中,丢弃了ImageNet的标签,仅使用图像;没有将监督损失与上述蒸馏损失相结合。
**教师模型。**考虑了使用自监督学习(SSL)预训练的模型,如DINO或iBOT,以及针对ImageNet-1K分类任务优化的监督学习模型,如DeiTIII[62]或微调后的dBoT。针对𝑀 = 2个教师模型进行了分析,特别是DINO和Deit-III。第4节和补充材料中探讨了更多教师模型和组合。
**任务。**在多个任务上测量性能,这些任务按以下轴线划分:
1)在ImageNet-1K验证集(IN-val)上训练集类别的Top-1准确性;
2)在未见类别上的迁移学习性能;报告了在15个不同的图像分类数据集上平均的Top-1准确性;
3)语义分割和4)深度估计的密集预测性能;报告了ADE-20k上的mIoU和NYUD上的RMSE,这些指标是使用本质上属于密集分类的协议测量的,即使用如[39]中的线性探针。
为所有任务直接在编码器输出𝒛上学习线性探针。
3. 分析ViT标记的多教师蒸馏
跨标记和教师的特征统计均等化
首先分析了从两位教师的CLS和块标记中提取的特征的统计信息,并表明这应在多教师蒸馏中予以考虑。计算了这些统计信息,并注意到给定教师的CLS和块标记之间以及不同教师之间的一阶和二阶矩值存在许多差异。
例如,DINO的CLS标记特征的范数和标准差是同模型块标记的两倍,而相同的统计信息在DeiT-III和DINO标记之间也存在差异。
为了探究这些特征之间的统计不一致性是否影响蒸馏,在计算损失之前,对每个教师的输出进行了特征标准化,即将教师特征归一化为零均值和单位方差。这不仅消除了CLS和块标记之间的差异,还消除了不同教师之间标记的差异。
为了方便和通用性,建议在蒸馏过程中使用指数移动平均法实时学习这种归一化统计信息。从表1中可以看到,无论是图像级还是块级任务,使用特征标准化的蒸馏学习模型的性能始终更高(第4行与第5行相比)。
CLS和块标记的projector head
除了统计差异外,CLS和块标记在概念上也不同:CLS是一个全局标记,预期会编码图像级语义,而块标记则编码局部信息。为了更好地从CLS和块标记中捕获这些特定信息,本文尝试为每种类型的标记使用专用的、特定于教师的projector head。这在实际操作中不会产生额外成本,因为在蒸馏后会丢弃这些projector head。关于可丢弃的projector he,将在第4节中进一步讨论。比较表1中的第5行和第6行,可以看到,将特定于教师的projector head专门化为CLS或块标记可以带来进一步的收益。
ImageNet和新类别的分类
表1中的结果表明,与针对该特定任务高度优化的模型(如DeiT-III,准确率为83.6%)相比,通过多教师蒸馏学习的模型在ImageNet-1K性能上有所欠缺(准确率为82.2%)。有人可能会认为这是由于在蒸馏过程中没有使用标签。为了验证这一点,本文还使用仅DeiT-III模型作为教师进行了蒸馏。
在这种情况下,能够在ImageNet上达到83.1%的top-1准确率。这比从多个教师联合蒸馏得到的82.2%要高得多,因此本文认为蒸馏过程本身仍有改进空间。
从表1中还可以看到,通过多教师蒸馏学习的模型在迁移学习和新类别分类方面大大优于DINO。这对于最近的iBOT模型也是如此,该模型也达到了最先进的top-1准确率,即在本文的设置下迁移学习的平均准确率为72.4%。
密集预测的多教师蒸馏
为了评估单个块标记的判别能力,在线性探测后考虑了两个密集预测任务:语义分割和深度预测。
表1显示,即使是基本的多教师蒸馏设置也优于最佳教师(第4行)。
更重要的是,使用标准化和针对CLS和块标记的专用projector head后,性能进一步提高(第6行)。
对于学生编码器,分割任务的mIoU比最佳教师高出4.6%。与那些旨在改进密集预测性能的模型相比,这一结果更加令人印象深刻。
本文的模型是通过监督学习和对比学习训练的教师蒸馏而来的,其密集预测性能与在密集任务上表现出色的模型相当,即通过掩码块预测训练的模型,如iBOT:iBOT在ADE-20K上实现了36.6%的mIoU,而本文的学生模型达到了36.9%。
保留教师的互补优势
从表1的结果中,可以看到,通过多教师蒸馏设置和简单的修改(如特征标准化和针对CLS/块标记的专用projector head)学习的模型开始在许多方面显示出强大的泛化性能。
这些模型似乎保留了教师的互补优势:它们已经在迁移学习和密集预测任务上超越了最佳教师,同时在ImageNet任务上也取得了不错的性能。
然而,正如上面所讨论的,仍有改进的空间;理想中的模型应该能够在所有任务上匹配或超越最佳教师。
4. 用于蒸馏的projector阶梯
上述基本设置使用可丢弃的projector head作为在蒸馏过程中注入教师特定参数的一种方式。这些模块附加在编码器的末端,并作为编码器输出与损失函数考虑的特征空间之间的小型"缓冲区"。
本文提议以互补的方式使用更多这些可丢弃的模块:作为信息高速公路,以更直接的方式从中间层传播信息到损失函数。中间层已被用于改进蒸馏,通常是通过在这些层之上添加额外的损失。然而,这导致了更具挑战性的优化。此外,具有许多额外损失的超参数调整是组合性的,并且变得繁琐。在多个教师的情况下,这些问题更为突出。
本文提出不在中间表示上添加损失,而是扩展现有的可丢弃的教师特定projector head以接收来自中间层的输入,并在损失之前附加模块,这些模块将所有中间层标记直接连接到教师特定的projector head。本文将这种扩展后的projecto称为projector阶梯。这种架构与通常用于将模型适应到新任务的适配器架构相似。然而,在本文的情况下,蒸馏过程中附加的类似适配器的模块是可丢弃的。
具体来说,本文将多层感知机(MLP)projector附加到中间层,并扩展教师特定projector ℎ𝑡的输入,该projector投影到目前为止仅在学生编码器的最后一层上运行。设𝒛𝑙表示学生编码器第𝑙层的输出,其中𝑙 = 1, ..., 𝐿。投projecto梯的头变为:
其中ℎ^𝑙_𝑡表示附加在层𝑙 ∈ 𝐿之后的MLPprojector head。ℎ^𝑙_𝑡的架构与ℎ_𝑡相同,但是,由于添加了多个这样的projector head,显著减少了隐藏维度𝑑^𝑙_ℎ,并在𝑙 < 𝐿时设置𝑑^𝑙_ℎ = 𝑑。
从表1中可以看到,这种projector阶梯总体上提高了性能(第8行),特别是对于密集预测。似乎密集连接导致了更好的主块标记。在有监督分类方面,增益也是显著的:ImageNet-1K准确率提高了+0.5%。
5. 平等地学习所有教师
基本设置假设蒸馏编码器的最终目标是能够平等地表示每位教师。
当蒸馏在所有教师之间使用特征标准化以及诸如余弦和平滑ℓ1之类的简单损失时,存在一种直接的方法来比较每位不同教师的学习程度:可以简单地比较损失的大小,这指示了近似每位教师特征空间的好坏程度。
图3a展示了使用第3.3节中介绍的设置对UNIC模型进行多教师蒸馏时的损失曲线(虚线)。可以看到,DINO教师似乎比DeiT-III更快且更好地被学习。
因此,本文的学生在ImageNet-1K(即DeiT-III擅长的任务)上表现不佳也就不足为奇了。但如果DINO根本不是蒸馏过程的一部分呢?在图3b中,展示了在使用DINO和DeiT-III作为教师以及仅从DeiT-III进行蒸馏的情况下,ImageNet-1K准确率在蒸馏过程中的变化。
可以看到,使用多位教师时,本文的模型学习得更快,但最终收敛到较低的准确率:学生似乎利用了来自额外教师的特征来更快地提升性能,但未能达到仅蒸馏DeiT-III时的准确率(83.1%)。
图3表明,某种形式的损失平衡可能是有益的。在多任务设置中,损失平衡是很常见的:在大多数情况下,这是通过添加控制每个损失的超参数来手动完成的。
然而,对于像本文的案例这样拥有多位教师和多种损失的情况,这种方法相当繁琐,这一点也在[46]中有所讨论。重要的是要避免手动调参的组合性质。
另一种方法是使用为多任务学习提出的现有损失平衡方法,例如Adaloss[24]。
本文认为,与多任务学习相比,在标准化特征和简单回归损失上进行多教师蒸馏时,平衡损失要简单得多:损失的大小是可比较的,并且可以用于平衡和调整蒸馏过程的速度。
教师正则化丢弃
本文引入了一种名为"教师丢弃"的简单损失平衡方案。本文没有设计一些软损失加权算法,而是从随机丢弃[58]和路径丢弃[25]等方法中汲取灵感,提出对一部分教师进行"丢弃",即将其损失置为零。
然而,随机丢弃教师并不会鼓励教师之间的损失均衡。相反,本文提出在选择要丢弃哪些教师时直接使用损失的绝对大小,即保留损失大小最大的教师,并以一定概率丢弃其他教师。
这与自适应丢弃[4]在概念上有相似之处,但本文的方法是非参数的,只是简单地利用了受限表示上特征空间损失的可比性。
本文在图像级别执行基于损失的教师丢弃。在每次迭代和每张图像上,本文为每个教师𝑡定义一个二进制系数𝛼𝑡 = {0, 1},该系数与相应的损失L𝑡相乘。这决定了对于该图像,教师𝑡是否以概率𝑝被丢弃。
为了确保始终有一些信号可以学习,选择从不丢弃损失绝对值最大的教师,即当前模型近似效果最差的教师。其他教师可以以概率𝑝被丢弃。具体来说,对于每张图像,教师𝑡 ∈ T的系数由以下公式给出:
在所有情况下,当前迭代中近似效果最差的教师始终会被使用。本文还尝试了补丁级别的教师丢弃,但并未发现明显的增益。
蒸馏过程中教师丢弃的效果
本文在图3a中研究了蒸馏过程中教师丢弃的影响:随着训练的进行,教师丢弃使得教师之间的损失大小变得更加相似(实线)。在图4中,本文绘制了蒸馏过程中教师系数𝛼𝑡的变化情况;教师利用率变得更加均衡,并在几个周期后趋于稳定。
教师丢弃如何影响性能?
本文将教师丢弃正则化与手动平衡教师损失、随机丢弃[15]以及最近的Adaloss[24]损失平衡方法进行了比较。从表1的第6行结果开始,本文发现这些策略中没有一种能够显著改进,更不用说超越使用教师丢弃(第8行)的结果了。具体来说,Adaloss在四个任务上分别达到了80.1/73.6/34.3/0.565(详细情况见补充材料)。除了性能之外,本文认为所提出的教师丢弃的有效性和简单性是无可比拟的。
本文研究了教师丢弃概率𝑝的影响,并发现不同值的性能是稳定的。然而,较高的𝑝有利于ImageNet的性能,而在学生已经超越最佳教师的任务上则略有下降(见补充材料)。
从表1(第8行)中可以看到,教师丢弃提升了ImageNet-1K的性能,即改进了本文在蒸馏模型中最缺乏的任务上的蒸馏效果。当将教师丢弃与一系列projector结合使用时,能够在该任务上达到83.2%的顶级性能。
这一性能仅比高度优化的DeiT-III(第3行)低0.4%。更重要的是,本文还缩小了仅使用DeiT-III进行多教师蒸馏和专用蒸馏之间观察到的差距。教师丢弃对此做出了重大贡献,将具有一系列projector的最佳模型的性能提高了0.5%(第7行与第8行相比)。
3.6. 迈向通用分类模型
使用一系列projector和教师丢弃正则化的多教师蒸馏使能够达到与高度优化的DeiT-III相当的ImageNet分类性能,同时在包括长尾类别在内的大多数新类别上的15个数据集上的迁移学习性能,以及像语义分割和深度估计这样的补丁级分类任务上,都超越了最佳教师的表现。本文认为这些证据表明,本文的蒸馏模型作为更通用的分类模型在运作。本文将使用增强的多教师蒸馏设置学习的模型称为UNIC模型(代表"通用分类",发音为"unique")。