[Paper Review] Masked Generative Distillation

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

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Masked Generative Distillation (MGD), a novel feature-based knowledge distillation method that improves student model representation by training it to reconstruct the teacher's full feature map from randomly masked student features using a simple generative block. MGD achieves state-of-the-art performance gains across image classification, object detection, semantic segmentation, and instance segmentation, with top-1 accuracy improvements of up to 1.79% on ImageNet and mAP gains exceeding 3.5 points on detection and segmentation tasks.

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.

Motivation & Objective

• To address the limitation of existing feature-based distillation methods that rely on direct imitation of teacher features, which may not optimally improve student representation.
• To explore whether feature reconstruction from masked inputs can enhance student feature representation power more effectively than direct imitation.
• To develop a simple, general-purpose distillation method applicable across diverse vision tasks, including classification, detection, and segmentation.
• To validate that the generative reconstruction objective leads to improved robustness and performance, even when student and teacher features remain visually dissimilar.

Proposed method

• Mask random pixels in the student's feature map during training, creating a masked input for reconstruction.
• Train a lightweight generative block (two 3×3 conv layers with ReLU activation) to reconstruct the teacher’s full feature map from the masked student feature.
• Use a reconstruction loss that minimizes the L2 difference between the generated and ground-truth teacher feature maps.
• Integrate the MGD loss with standard distillation objectives (e.g., cross-entropy or detection heads) for joint optimization.
• Apply the method at different stages of the network, with optimal performance observed at deeper, more semantic layers.
• Tune two hyperparameters: the loss balance coefficient α and the mask ratio λ, with minimal sensitivity observed in ablation studies.

Experimental results

Research questions

• RQ1 Can a student model achieve better representation learning by reconstructing the teacher’s full feature from masked inputs, rather than directly imitating the teacher’s features?
• RQ2 Does the generative reconstruction objective lead to improved performance even when the student and teacher feature maps remain visually distinct?
• RQ3 How does MGD perform across diverse vision tasks, including image classification, object detection, and dense prediction tasks?
• RQ4 What is the optimal architecture and hyperparameter configuration for the generative block and mask ratio in MGD?
• RQ5 Does distillation at deeper network stages yield greater performance gains than at shallower stages?

Key findings

• MGD improves ResNet-18’s ImageNet top-1 accuracy from 69.90% to 71.69%, a gain of 1.79%.
• On object detection, RetinaNet with ResNet-50 backbone improves from 37.4 to 41.0 mAP, a 3.6-point gain.
• SOLO with ResNet-50 backbone achieves a 3.1-point increase in mask mAP, from 33.1 to 36.2.
• DeepLabV3 with ResNet-18 improves mIoU from 73.20% to 76.02% on ADE20K.
• MGD outperforms state-of-the-art feature-based distillation methods like FGD and KR, even when student and teacher features remain visually dissimilar.
• The method is robust to hyperparameter choices, especially the mask ratio λ, with optimal performance at moderate values (λ < 0.5), and shows minimal sensitivity to the loss balance coefficient α.