[论文解读] The Role of Masking for Efficient Supervised Knowledge Distillation of Vision Transformers

Seung‐Woo Son, Ryu, Jegwang+2 more|arXiv (Cornell University)|Feb 21, 2023

Advanced Neural Network Applications被引用 5

一句话总结

MaskedKD 是一种通过在教师模型推理过程中对输入图像块进行遮蔽，以降低视觉变换器（Vision Transformer）知识蒸馏计算成本的方法。该方法基于学生模型注意力机制计算出的显著性分数，对 25%–50% 的图像块进行遮蔽。该方法可将知识蒸馏的 FLOPs 降低高达 50%，且不损失学生模型的准确率，从而实现总训练 FLOPs 最多 28% 的减少，以及 2.0×–3.7× 的吞吐量提升。

ABSTRACT

Knowledge distillation is an effective method for training lightweight vision models. However, acquiring teacher supervision for training samples is often costly, especially from large-scale models like vision transformers (ViTs). In this paper, we develop a simple framework to reduce the supervision cost of ViT distillation: masking out a fraction of input tokens given to the teacher. By masking input tokens, one can skip the computations associated with the masked tokens without requiring any change to teacher parameters or architecture. We find that masking patches with the lowest student attention scores is highly effective, saving up to 50% of teacher FLOPs without any drop in student accuracy, while other masking criterion leads to suboptimal efficiency gains. Through in-depth analyses, we reveal that the student-guided masking provides a good curriculum to the student, making teacher supervision easier to follow during the early stage and challenging in the later stage.

研究动机与目标

• 为解决在知识蒸馏过程中，使用大型视觉变换器（ViT）教师模型带来的高计算成本问题。
• 在不降低学生模型性能的前提下，降低知识蒸馏成本——即运行教师模型所需的 FLOPs。
• 开发一种避免预计算并存储教师模型预测结果的方法，以减少内存占用，并在数据增强下避免准确率下降。
• 实现在自监督学习和基于音频的变换器设置中的高效蒸馏，扩展至监督 ViT 训练之外的应用场景。

提出的方法

• MaskedKD 基于学生模型最后一层多头注意力机制计算出的块显著性度量来选择遮蔽的块，复用现有注意力分数，计算开销可忽略不计。
• 显著性度量通过平均块间注意力分数（排除分类标记）来识别并保留对学生预测最相关的图像特征。
• 将显著性分数最低的块进行遮蔽，确保核心特征得以保留，同时减少输入教师模型的 token 数量。
• 该方法在知识蒸馏前向传播过程中执行，无需修改网络架构或重新训练教师模型。
• 仅在教师模型的前向传播中应用遮蔽，而学生模型仍处理完整图像，从而保持特征学习的完整性。
• 该方法与多种 ViT 蒸馏方法兼容，包括原始 KD、DINO 和基于块的蒸馏。

实验结果

研究问题

• RQ1 我们能否在不降低学生模型准确率的前提下，降低 ViT 蒸馏的计算成本？
• RQ2 仅利用学生模型注意力机制的信息，通过遮蔽输入教师模型的图像块，能否实现显著的 FLOP 减少？
• RQ3 基于注意力的块显著性与随机遮蔽或基于 DINO 的显著性等其他遮蔽策略相比，在保持蒸馏性能方面表现如何？
• RQ4 MaskedKD 能否扩展至自监督学习和非视觉变换器任务（如音频处理）？
• RQ5 为实现最大效率并维持性能，应遮蔽多少比例的图像块为最优？

主要发现

• MaskedKD 在不损失学生模型 top-1 准确率的前提下，将输入教师模型的图像块数量减少了 25%–50%，适用于多种 ViT 架构和蒸馏方法。
• 该方法实现了 50% 的蒸馏 FLOPs 减少（即教师模型推理成本降低），在将 DeiT-S 蒸馏至 DeiT-Ti 时，总训练 FLOPs 最多减少 28%。
• 由于教师模型计算减少，吞吐量提升 2.0× 至 3.7×，具体取决于设置。
• 使用块间注意力分数进行遮蔽会降低性能，证实了所提出的显著性度量更具有效性。
• 尽管计算成本更高，基于 DINO 注意力分数的遮蔽仍无法超越 MaskedKD 的性能，表明基于学生注意力的遮蔽更具优势。
• 保留 top-k 显著性块（即最重要块）的性能优于随机遮蔽或 bottom-k 遮蔽，验证了显著性与性能之间的相关性。