A Simple and Generic Framework for Feature Distillation via Channel-wise Transformation

TL;DR

This paper introduces a simple, generic framework for feature distillation that aligns teacher and student features along the channel dimension using a learnable nonlinear transformation, improving performance across various vision tasks.

Contribution

It proposes a novel channel-wise feature alignment method with a learnable transformation, enhancing distillation effectiveness and versatility in multiple computer vision applications.

Findings

Improves image classification accuracy by 3.28% on ImageNet.

Enhances object detection mAP by 3.9% on MS COCO.

Boosts semantic segmentation mIoU by 4.66% on Cityscapes.

Abstract

Knowledge distillation is a popular technique for transferring the knowledge from a large teacher model to a smaller student model by mimicking. However, distillation by directly aligning the feature maps between teacher and student may enforce overly strict constraints on the student thus degrade the performance of the student model. To alleviate the above feature misalignment issue, existing works mainly focus on spatially aligning the feature maps of the teacher and the student, with pixel-wise transformation. In this paper, we newly find that aligning the feature maps between teacher and student along the channel-wise dimension is also effective for addressing the feature misalignment issue. Specifically, we propose a learnable nonlinear channel-wise transformation to align the features of the student and the teacher model. Based on it, we further propose a simple and generic framework for feature distillation, with only one hyper-parameter to balance the distillation loss and the task specific loss. Extensive experimental results show that our method achieves significant performance improvements in various computer vision tasks including image classification (+3.28% top-1 accuracy for MobileNetV1 on ImageNet-1K), object detection (+3.9% bbox mAP for ResNet50-based Faster-RCNN on MS COCO), instance segmentation (+2.8% Mask mAP for ResNet50-based Mask-RCNN), and semantic segmentation (+4.66% mIoU for ResNet18-based PSPNet in semantic segmentation on Cityscapes), which demonstrates the effectiveness and the versatility of the proposed method. The code will be made publicly available.