[论文解读] Masked Distillation with Receptive Tokens

Tao Huang, Yuan Zhang+5 more|arXiv (Cornell University)|May 29, 2022

Advanced Neural Network Applications被引用 7

一句话总结

该论文提出MasKD，一种用于密集预测任务的新颖知识蒸馏框架，通过可学习的感受野标记（receptive tokens）动态识别并优先处理特征图中的感兴趣像素（PoIs）。通过生成自适应的、基于注意力的蒸馏掩码，并利用多种多样化的掩码结合Dice损失，MasKD在不依赖任务特定先验知识的前提下，在目标检测与语义分割基准上实现了最先进性能。

ABSTRACT

Distilling from the feature maps can be fairly effective for dense prediction tasks since both the feature discriminability and localization priors can be well transferred. However, not every pixel contributes equally to the performance, and a good student should learn from what really matters to the teacher. In this paper, we introduce a learnable embedding dubbed receptive token to localize those pixels of interests (PoIs) in the feature map, with a distillation mask generated via pixel-wise attention. Then the distillation will be performed on the mask via pixel-wise reconstruction. In this way, a distillation mask actually indicates a pattern of pixel dependencies within feature maps of teacher. We thus adopt multiple receptive tokens to investigate more sophisticated and informative pixel dependencies to further enhance the distillation. To obtain a group of masks, the receptive tokens are learned via the regular task loss but with teacher fixed, and we also leverage a Dice loss to enrich the diversity of learned masks. Our method dubbed MasKD is simple and practical, and needs no priors of tasks in application. Experiments show that our MasKD can achieve state-of-the-art performance consistently on object detection and semantic segmentation benchmarks. Code is available at: https://github.com/hunto/MasKD .

研究动机与目标

• 为解决标准特征蒸馏方法对所有像素一视同仁的问题，尽管其对模型性能的贡献各不相同。
• 消除在密集预测任务中选择蒸馏区域时对真实边界框或任务特定先验知识的依赖。
• 通过基于注意力的感受野标记，学习特征图中细粒度的、动态的像素级重要性模式。
• 通过捕捉多种多样化掩码中的复杂像素依赖关系，同时保持训练稳定性，从而增强蒸馏效果。
• 通过掩码加权机制提升学生模型性能，该机制可动态考虑不同掩码在蒸馏过程中的重要性差异。

提出的方法

• 引入可学习的感受野标记，通过计算像素级注意力来识别教师网络特征图中的感兴趣像素（PoIs）。
• 通过感受野标记的像素级注意力生成蒸馏掩码，指示在特征重建中应优先处理的像素。
• 采用多个感受野标记以捕捉特征图内部多样化且复杂的像素依赖关系。
• 使用主任务损失对感受野标记进行端到端训练，同时固定教师网络参数，实现端到端优化。
• 应用Dice损失以鼓励多个生成掩码之间的多样性，防止冗余。
• 引入掩码加权模块，根据各掩码对蒸馏的贡献程度，自适应地分配其重要性权重。

实验结果

研究问题

• RQ1 在特征图中通过动态、基于注意力的方式选择像素，是否能提升密集预测任务中的知识蒸馏性能？
• RQ2 如何在不依赖真实标注或任务特定先验知识的前提下，学习到信息丰富且多样的蒸馏掩码？
• RQ3 相较于均匀采样或启发式采样，多个感受野标记与自适应掩码加权机制在多大程度上能提升蒸馏性能？
• RQ4 所提出方法在不同网络架构与任务（如目标检测与语义分割）之间是否具备良好的泛化能力？
• RQ5 模型能否学习到与教师和学生特征表示相一致的有意义像素依赖关系？

主要发现

• 在COCO目标检测任务中，MasKD相较Faster R-CNN-R50实现2.4 AP的性能提升，显著优于标准蒸馏基线方法。
• 在语义分割任务中，MasKD使DeepLabV3-R18在COCO数据集上mIoU提升2.79，展现出强大的跨任务泛化能力。
• 该方法在无需任务特定先验知识的前提下，持续提升多种主干网络架构与数据集上的性能表现。
• 消融实验证实，多掩码组合、Dice损失带来的多样性，以及掩码加权机制的协同作用，显著提升了蒸馏效率。
• 可视化结果表明，学习到的掩码聚焦于物体内部的判别性、局部化区域，而非整个边界框，验证了其细粒度注意力机制的有效性。
• 该方法在未见过的网络架构与数据集上也表现出良好的泛化能力，表明其在真实部署场景中具备鲁棒性与实用性。