Source-linked AI summary

Masked Generative Distillation

Zhendong Yang, Zhe Li, Mingqi Shao, Dachuan Shi, Zehuan Yuan, Chun Yuan

arXiv:2205.01529v2cs.CV

TL;DR

Existing distillation methods commonly improve students by directly imitating teacher outputs, while task-specific designs limit reuse across tasks. MGD instead masks student features and trains them to generate the teacher’s full feature, producing improvements across classification, detection, and segmentation benchmarks.

  • Problem

    Many distillation methods are task-specific, and classification distillation often fails on dense prediction despite the potential generality of feature-based distillation.

  • Method

    MGD randomly masks student-feature pixels and uses a simple projector to generate the teacher’s complete feature rather than directly mimicking it.

  • Results

    MGD yields significant improvements across image classification, object detection, semantic segmentation, and instance segmentation, including 3.6 Boundingbox mAP and 3.1 Mask mAP gains on COCO.

  • Takeaways & Limitations

    Because MGD operates on feature maps, it can be applied to different models and tasks and combined with logit-based or head-based distillation.

Abstract

from arXiv · show

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.

1 Introduction

MGD addresses the limited cross-task applicability of specialized distillation methods by training students to recover teacher features from masked features rather than directly imitate them. Experiments report improvements across classification, detection, and segmentation tasks.

  • Motivation: Feature-based distillation is theoretically reusable across tasks, but specialized methods often transfer poorly between classification, detection, and other settings.OFD and KR bring limited improvement for detectors, while detector-specific FKD and FGD cannot be used in other tasks because they lack a neck.
  • Method: MGD masks random student-feature pixels and uses a simple block to generate the teacher’s full feature.Random masks across iterations expose all pixels during training, aiming to improve feature robustness and representation power.
  • Results: MGD produced considerable improvements across image classification, object detection, semantic segmentation, and instance segmentation.The method can also be combined with logit-based or head-based distillation methods for further gains.

2 Related work

The related work contrasts general feature-based distillation with task-specific approaches and highlights the difficulty of transferring classification distillation to dense prediction.

  • Feature-based distillation: Knowledge distillation includes logit-based methods and feature-based methods that transfer information from intermediate representations.Examples include FitNet, attention transfer, OFD, and contrastive representation distillation.
  • Dense prediction: Classification and dense prediction differ substantially, and many classification distillation methods fail on dense prediction tasks.The paper frames feature-based distillation as a potential approach for both task families.
  • Object detection: Detection distillation methods focus on where to distill because foreground-background imbalance can introduce noise.Prior methods use detector necks, heads, object-adjacent regions, or other spatial selection strategies.
  • Feature comparison: Other feature-distillation methods enforce pair-wise, holistic, latent-domain, or channel-wise probability-map consistency.These approaches differ in how they represent and compare teacher and student features.

3 Method

MGD aligns student features, randomly masks them, and uses a projector to generate teacher features; its distillation loss is combined with the task’s original loss.

  • Feature distillation: Feature-based distillation traditionally minimizes squared differences between aligned student and teacher feature maps.The adaptation layer aligns the student feature dimensions with the teacher feature dimensions.
  • Generation with Masked Feature: MGD generates the teacher’s feature from a masked student feature instead of directly mimicking the teacher.The method relies on partial feature pixels to recover the complete teacher feature map.
  • Generation with Masked Feature: The method defines teacher and student feature maps T^l and S^l for each distillation layer and applies a random mask to the student map.The mask uses a random variable and λ, which denotes the masked ratio.
  • Projector: The masked, aligned student feature is passed through a projector to produce the teacher feature.The projector contains two convolutional layers with a ReLU activation; the paper uses 3×3 convolutions for these layers.
  • Overall Loss: MGD trains the student with L_all = L_original + α · L_dis, where α balances the original task loss and distillation loss.L_original is the model’s task loss, while L_dis is the MGD distillation loss.
  • Algorithm: MGD is designed as a simple feature-based method that can be applied across multiple tasks and models.The algorithm updates the student using the original loss and MGD distillation loss after computing teacher and student features.

4 Main Experiments

MGD is evaluated across multiple computer-vision tasks and uses standardized ImageNet training and evaluation settings for classification.

  • MGD is applied to classification, object detection, semantic segmentation, and instance segmentation across different models and datasets.
  • ImageNet classification uses 1.2 million training images, 50k testing images, 1000 object categories, and accuracy evaluation.
  • Classification distillation uses the backbone’s last feature map, with α = 7 × 10−5 and λ = 0.5 across experiments.

Implementation Details.

The paper compares MGD with multiple distillation methods on ImageNet and reports gains for both homogeneous and heterogeneous classification settings.

  • Table 1 compares different distillation methods on ImageNet, with T and S denoting teacher and student, respectively.
  • MGD improves student ResNet-18 and MobileNet by 1.68 and 3.14 Top-1 accuracy, respectively.
  • The classification experiments include ResNet-34 to ResNet-18 homogeneous distillation and ResNet-50 to MobileNet heterogeneous distillation.

Classification Results.

The paper evaluates detection and instance segmentation on COCO using several detector architectures and reports comparisons with recent detector-specific distillation methods.

  • COCO2017 contains 80 object categories, with 120k training images and 5k validation images used for evaluation.
  • Detection and instance-segmentation models use neck feature maps for distillation, with different α and λ settings for one-stage and two-stage models.
  • Object detection experiments cover Faster RCNN, RetinaNet, and RepPoints, spanning two-stage, anchor-based one-stage, and anchor-free one-stage detectors.
  • Instance segmentation experiments use SOLO and Mask RCNN, and compare MGD with three recent state-of-the-art detector distillation methods.

Object Detection and Instance Segmentation Results.

MGD improves instance-segmentation and object-detection students on COCO, with substantial gains reported for RetinaNet and SOLO using ResNet-50 backbones.

  • 3.6 Boundingbox mAP improvement is reported for ResNet-50-based RetinaNet on COCO with MGD.
  • 3.1 Mask mAP improvement is reported for ResNet-50-based SOLO on COCO with MGD.
  • Table 3 reports results for different distillation methods on COCO instance segmentation, with AP denoting Mask AP and MS denoting multi-scale training.

4.3 Semantic Segmentation

MGD is evaluated for semantic segmentation on CityScapes using homogeneous and heterogeneous student–teacher settings, with consistent improvements over prior distillation methods.

  • Semantic segmentation experiments use CityScapes, which contains 5000 high-quality images, and evaluate models with mean Intersection-over-Union (mIoU).
  • MGD uses PspNet-Res101 as teacher and PspNet-Res18 or DeepLabV3-Res18 as students under 512×1024 input settings.
  • 3.78 mIoU improvement is achieved by the ResNet-18-based PspNet with MGD.
  • Adding CWD head logit distillation provides another 0.47 mIoU improvement for PspNet and 0.29 mIoU improvement for DeepLabV3.

Semantic Segmentation Results.

Table 4 reports semantic segmentation results on CityScapes, averaging results over three runs and distinguishing teacher, student, and optional CWD head distillation.

  • Table 4 reports semantic segmentation results on CityScapes, with values averaged across three runs.
  • The table uses T and S to denote teacher and student, while an asterisk marks addition of CWD head distillation loss.

5 Analysis

The analyses examine MGD’s feature-recovery mechanism, teacher and masking choices, generative-block design, distillation stages, and hyper-parameter sensitivity. Across these studies, feature recovery, compatible teachers, deeper features, and moderate masking are associated with better student performance.

  • Analysis: MGD masks the student feature and generates the teacher’s complete feature map, aiming to improve the student’s input representation.
  • Analysis: 1.01 accuracy improvement is achieved when ResNet-18 distills itself with MGD, while directly mimicking the teacher feature yields limited improvement.
  • Analysis: Direct feature mimicking reduces the student–teacher L2 difference, whereas MGD produces greater improvement despite not simply matching that feature.
  • Analysis: Masking random channels with β = 0.15 and α = 7 × 10^-5 improves image-classification performance after pooling reduces spatial dimensions.
  • Analysis: 71.8 accuracy is achieved with ResNetV1D-152 teaching ResNet-18, compared with 70.91 using ResNet-18 as teacher; heterogeneous teachers provide limited gains.
  • Analysis: The selected generative block uses two 3×3 convolutional layers and one ReLU activation; one convolution improves least, while three worsen Top-1 accuracy.
  • Analysis: Distilling deeper stages benefits students more than distilling shallower layers, because deeper features contain more semantic information.
  • Analysis: Accuracy rises as the mask ratio increases below 0.5, but at 0.8 the remaining semantic information is too poor to generate the complete teacher feature map.

6 Conclusions

MGD improves student representation power by generating the teacher’s feature from masked student features rather than directly mimicking it. Because it operates on feature maps, the method applies across multiple vision tasks and is supported by extensive experiments.

  • MGD trains students to generate the teacher’s feature from masked student features instead of directly imitating the teacher.
  • The feature-map-based method can be applied to image classification, object detection, semantic segmentation, and instance segmentation.
  • Extensive experiments across models and datasets show that MGD is simple and efficient.
Loading 2205.01529v2…