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

©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.

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·