[论文解读] Masked Generative Distillation

Zhendong Yang, Zhe Li+4 more|arXiv (Cornell University)|May 3, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文提出了一种新型基于特征的知识蒸馏方法——掩码生成蒸馏（Masked Generative Distillation, MGD），通过在训练过程中让学生模型从随机掩码的学生特征中重建教师的完整特征图，从而提升学生模型的表征能力。MGD在图像分类、目标检测、语义分割和实例分割等任务上均取得了当前最优的性能提升，其中在ImageNet上的top-1准确率最高提升了1.79%，在检测与分割任务上的mAP提升超过3.5个百分点。

ABSTRACT

Knowledge distillation has been applied to various tasks successfully. The current distillation algorithm usually improves students' performance by imitating the output of the teacher. This paper shows that teachers can also improve students' representation power by guiding students' feature recovery. From this point of view, we propose Masked Generative Distillation (MGD), which is simple: we mask random pixels of the student's feature and force it to generate the teacher's full feature through a simple block. MGD is a truly general feature-based distillation method, which can be utilized on various tasks, including image classification, object detection, semantic segmentation and instance segmentation. We experiment on different models with extensive datasets and the results show that all the students achieve excellent improvements. Notably, we boost ResNet-18 from 69.90% to 71.69% ImageNet top-1 accuracy, RetinaNet with ResNet-50 backbone from 37.4 to 41.0 Boundingbox mAP, SOLO based on ResNet-50 from 33.1 to 36.2 Mask mAP and DeepLabV3 based on ResNet-18 from 73.20 to 76.02 mIoU. Our codes are available at https://github.com/yzd-v/MGD.

研究动机与目标

• 为解决现有基于特征的蒸馏方法依赖于直接模仿教师特征所带来的局限性，这些方法可能无法最优地提升学生模型的表征能力。
• 探究从掩码输入中重建特征是否能比直接模仿更有效地增强学生特征的表征能力。
• 开发一种简单、通用的蒸馏方法，适用于包括分类、检测和分割在内的多种视觉任务。
• 验证生成式重建目标是否能带来更强的鲁棒性与性能提升，即使学生与教师的特征图在视觉上仍存在明显差异。

提出的方法

• 在训练过程中对学生的特征图随机掩码像素，形成用于重建的掩码输入。
• 训练一个轻量级生成模块（由两层3×3卷积层与ReLU激活函数构成），从掩码后的学生特征中重建教师的完整特征图。
• 使用重建损失，最小化生成特征图与真实教师特征图之间的L2差异。
• 将MGD损失与标准蒸馏目标（如交叉熵或检测头损失）联合优化。
• 在神经网络的不同阶段应用该方法，发现在深层、更具语义性的特征层上性能最佳。
• 调节两个超参数：损失平衡系数α与掩码比率λ，在消融实验中表现出较低的敏感度。

实验结果

研究问题

• RQ1 学生模型是否能通过从掩码输入中重建教师的完整特征，而非直接模仿教师特征，实现更优的表征学习？
• RQ2 即使学生与教师的特征图在视觉上仍明显不同，生成式重建目标是否仍能带来性能提升？
• RQ3 MGD在包括图像分类、目标检测和密集预测任务在内的多种视觉任务中表现如何？
• RQ4 生成模块与掩码比率的最优架构与超参数配置是什么？
• RQ5 在深层网络阶段进行蒸馏是否比在浅层阶段带来更大的性能增益？

主要发现

• MGD将ResNet-18在ImageNet上的top-1准确率从69.90%提升至71.69%，提升1.79%。
• 在目标检测任务中，使用ResNet-50骨干的RetinaNet mAP从37.4提升至41.0，提升3.6个百分点。
• SOLO模型使用ResNet-50骨干时，mask mAP从33.1提升至36.2，提升3.1个百分点。
• DeepLabV3使用ResNet-18骨干时，在ADE20K数据集上的mIoU从73.20%提升至76.02%。
• 即使学生与教师特征在视觉上仍不相似，MGD仍优于当前最优的基于特征的蒸馏方法（如FGD与KR）。
• 该方法对超参数选择具有鲁棒性，尤其对掩码比率λ具有较强容忍度，最优性能出现在中等值（λ < 0.5），且对损失平衡系数α的敏感度极低。