CVPR 2023 | G2SD: 让小模型也能从自监督预训练中受益的蒸馏方法

科技

2023-03-28 04:03

©Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

研究方向 | 计算机视觉、模型压缩

论文标题：

Generic-to-Specific Distillation of Masked Autoencoders

论文链接：

https://arxiv.org/pdf/2302.14771.pdf

代码链接：

https://github.com/pengzhiliang/G2SD

动机&背景

大型视觉 Transformer（ViT）在自监督预训练机制的推动下取得了前所未有的进步。然而，由于模型容量的限制，轻量级 ViT 模型从这些预训练机制中受益较少。知识蒸馏是将大型（教师）模型的表示转移到小型（学生）模型的模型压缩方法。传统的单阶段蒸馏容易陷入特定任务知识的转移中，无法保留对模型泛化至关重要的任务无关知识。

本文提出了通用到特定蒸馏（generic-to-specific distillation, G2SD），在 MAE（Masked AutoEncoders）[1]预训练的大型模型的监督下挖掘小型 ViT 模型的潜力。在通用蒸馏中，鼓励小模型的解码器将特征预测与大模型的隐藏表示对齐，以传递任务无关的知识。在特定蒸馏中，约束小模型的预测与大模型的概率预测保持一致，以学习能提供任务性能的特定任务的特征表示。

在自然语言处理领域，TinyBERT [1] 已经证明了两阶段蒸馏范式的有效性。本文所提出的 G2SD 方法正是这种范式计算机视觉领域的实践。

G2SD：通用到特定蒸馏

与 TinyBERT [2] 类似，本文提出的蒸馏方法分为两个阶段：通用蒸馏（Generic Distillation, GD）和特定蒸馏（Specific Distillation, SD），具体示意图见图 1。通用蒸馏从大规模预训练 MAE 中获取任务无关的知识，以提升学生的泛化能力；特定蒸馏从下游任务中获取特定任务的知识，以进一步提升模型性能。

▲ 图1. G2SD 示意图。[M] 表示掩码令牌。

2.1 通用蒸馏：传递任务无关知识

为了讲清楚 GD 阶段的蒸馏损失函数，我们先对图 1（左）的各个元素进行数学定义。将 MAE 预训练获得的教师模型的编码器和解码器分别表示为和，同理，学生模型的编码器和解码器可表示为和。

MAE 方法会将输入令牌随机分为可见令牌和掩码令牌，可见令牌经过和可得到特征提取和。根据式（1），编码器的输出会和掩码令牌拼接得到解码器的输入和。

将教师和学生解码器的层数定义为和（），设教师解码器第层的输出为，学生解码器的输出为。为了能够进行特征蒸馏，需要将两者的输出进行通道对齐，学生对齐后的输出为。

根据上述定义，GD 阶段的蒸馏损失函数可定义为：

式（2）就是一个比较常规的特征蒸馏函数，是和范数的权衡函数。该损失函数的目标是让学生的编码器学习与教师编码器类似的特征提取方式。同时，由于损失函数中包含了掩码令牌的上下文信息，这也能进一步帮助学生模仿教师的建模能力，从而传递任务无关的知识。

2.2 特定蒸馏：特定任务的表示

GD 阶段能为学生模型提供较强的泛化能力，并在下游任务中获得较好的性能。然而，由于模型规模和参数量的限制，学生模型与其教师模型仍存在一定的性能差距。为了进一步缩小差距，可在下游任务（例如图像分类、目标检测和语义分割）中进行特定蒸馏，以学习更适合特定任务的特征表示。具体损失函数如下：

其中，是下游任务的损失函数，是原始蒸馏损失（即 KL 散度）。

2.3 理论基础：互信息

我们可以从互信息 [3] 的角度解释两阶段蒸馏的合理性。知识蒸馏可以从最大化模型互信息的视角理解。设教师和学生的模型参数为和，模型为和，预训练和下游数据集分别为和。单阶段特定任务蒸馏可解释为在下游数据集中最大化模型互信息，即：

而 G2SD 两阶段蒸馏则可解释为在预训练和下游两个数据集上最大化模型互信息，即：

显然，式（5）的互信息比式（4）的大，这意味着 G2SD 传递了更多的信息。

实验

图像分类：表 1 对比了 G2SD 和现有监督学习/自监督学习/知识蒸馏方法的性能，G2SD 在同等参数量和计算量下取得了超过它们所有方法的性能。值得注意的是，G2SD-Ti 在相近参数量下比精心设计的 MobileNet-v3 性能高出了 1.8%。

▲ 表1. ImageNet-1k 上 Top-1 准确率对比

目标检测&实例分割：如表 2 所示，在目标检测和实例分割任务中，G2SD 均取得了优于现有流行方法的性能。

▲ 表2. MS COCO 上目标检测和实例分割的结果对比

语义分割：如表 3 所示，G2SD 获得了最优分割性能。同时，G2SD-Ti 在只需要 11M 参数量的情况下，取得了和 Swin-T（59.9M）一样的性能，将轻量级 ViT 模型的性能推到了一个新的高度。

▲ 表3. ADE20K 数据集上的结果对比

消融实验：表 4 验证了两阶段蒸馏的有效性。从中可知，通用蒸馏就能使学生模型获得有竞争力的性能，而特定蒸馏进一步提升了模型性能，两个阶段均有很大贡献。

▲ 表4. 对单阶段和两阶段蒸馏的消融研究，G2SD w/o S.D是指仅进行通用蒸馏，MAE⚗是指在 MAE 的微调阶段进行特定蒸馏

其他实验：本文还就通用蒸馏阶段的损失函数设计进行了消融探究，包括选择哪个位置进行特征蒸馏（编码器特征、编码器特征相对关系、解码器特征等）、掩码比例、哪层特征（的选择）、学生解码器的设计。具体的实验结果可参考原文。

参考文献

[1] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Doll ́ar, and Ross Girshick. Masked autoencoders are scalable vision learners. In IEEE CVPR, 2022.

[2] Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. ArXiv, abs/1909.10351, 2020.

[3] Sungsoo Ahn, Shell Xu Hu, Andreas C. Damianou, Neil D. Lawrence, and Zhenwen Dai. Variational information distillation for knowledge transfer. IEEE CVPR, pages 9155– 9163, 2019.