[论文解读] ViTKD: Practical Guidelines for ViT feature knowledge distillation

Zhendong Yang, Zhe Li+4 more|arXiv (Cornell University)|Sep 6, 2022

Advanced Neural Network Applications被引用 23

一句话总结

本论文研究了基于特征的知识蒸馏在 Vision Transformers (ViT) 上，推导出三个实用准则，提出 ViTKD，并在 ImageNet-1k 上显示出与基于 logit 的 KD 相辅相成的一致改进。

ABSTRACT

Knowledge Distillation (KD) for Convolutional Neural Network (CNN) is extensively studied as a way to boost the performance of a small model. Recently, Vision Transformer (ViT) has achieved great success on many computer vision tasks and KD for ViT is also desired. However, besides the output logit-based KD, other feature-based KD methods for CNNs cannot be directly applied to ViT due to the huge structure gap. In this paper, we explore the way of feature-based distillation for ViT. Based on the nature of feature maps in ViT, we design a series of controlled experiments and derive three practical guidelines for ViT's feature distillation. Some of our findings are even opposite to the practices in the CNN era. Based on the three guidelines, we propose our feature-based method ViTKD which brings consistent and considerable improvement to the student. On ImageNet-1k, we boost DeiT-Tiny from 74.42% to 76.06%, DeiT-Small from 80.55% to 81.95%, and DeiT-Base from 81.76% to 83.46%. Moreover, ViTKD and the logit-based KD method are complementary and can be applied together directly. This combination can further improve the performance of the student. Specifically, the student DeiT-Tiny, Small, and Base achieve 77.78%, 83.59%, and 85.41%, respectively. The code is available at https://github.com/yzd-v/cls_KD.

研究动机与目标

• 激励并理解针对 ViT 模型的基于特征的知识蒸馏，其与 CNN 不同，因为 ViT 具有基于注意力的结构。
• 确定在不同层和模块中蒸馏 ViT 特征的有效策略。
• 开发一种针对 ViT 的特定蒸馏方法，在 ImageNet-1k 上取得持续改进。
• 证明 ViTKD 可以与基于 logit 的 KD 方法互为补充，并有益于下游任务。

提出的方法

• 分析 ViT 在不同层的特征图和注意力行为，以设计蒸馏指南。
• 研究在浅层使用 mimicking（线性层对齐和相关性矩阵）。
• 研究基于生成的蒸馏（遮罩标记并使用生成块，如跨注意、自注意或卷积投影器）用于深层。
• 将 ViTKD 定义为结合浅层 mimicking 与深层 generation，并使用总损失：L = L_ori + alpha L_lr + beta L_gen。
• 在特征和生成目标上使用基于 L2 的蒸馏损失，并在需要时使用自适应层。
• 提供实现细节，包括掩码比率 lambda = 0.5，以及 ImageNet-1k 实验中的超参数 alpha = 3e-5, beta = 3e-6。

实验结果

研究问题

• RQ1 在将知识迁移到更小的 ViT 学生模型时，ViT 特定的基于特征的蒸馏能否超过基于 CNN 的特征蒸馏？
• RQ2 哪些层（浅层 vs 深层）以及哪些蒸馏机制（mimicking vs generation）对 ViT 特征蒸馏收益最大？
• RQ3 ViTKD 是否与基于 logit 的 KD 方法互补？它们的结合是否能进一步提升性能？
• RQ4 蒸馏策略如何迁移到图像分类以外的下游任务（如目标检测）？

主要发现

• 出现三条实用准则：浅层使用 mimicking，深层使用 generation；聚焦 FFN-out 或 MHA-out 特征有利于 FFN-out 的蒸馏；浅层知识对 ViT 蒸馏尤为有益。
• ViTKD 将 DeiT-Tiny 的 Top-1 从 74.42% 提升到 76.06%，DeiT-Small 从 80.55% 提升到 81.95%，DeiT-Base 从 81.76% 提升到 83.46%。
• 与基于 logit 的 KD（NKD）结合时，ViTKD 进一步提升到 Tiny/Small/Base 的 Top-1 分别为 77.78%、83.59%、85.41%。
• ViTKD 训练的模型也提升了下游任务；例如，将 ViTKD 与 Mask-RCNN 一起使用可提升 COCO 的 AP box 和 AP mask 指标。
• 与学生同架构的教师在 ViTKD 中提供更好的指导；跨架构的教师可能降低性能。
• ViTKD 对超参数 alpha 和 beta 显示出鲁棒性，并在不同的师生对中与 NKD 取得互补增益。