Addressing Class Imbalance in Contrastive Knowledge Distillation via Teacher-Guided Feature Augmentation for Semantic Segmentation

초록

Contrastive knowledge distillation (CKD) exploits the powerful discriminative capability of contrastive learning when transferring knowledge from a strong teacher model to a lightweight student model to compress deep neural networks. However, current CKD methods often overlook the inherent class imbalance in training datasets, which hampers effective knowledge transfer between teacher and student models. This imbalance causes the student model to focus primarily on majority classes when distilling knowledge from the teacher model, while paying less attention to the representation of minority classes. To address this limitation within the CKD framework, we propose a novel approach, called Teacher-Guided Feature Augmentation (TGFA), which selectively augments student features corresponding to minority classes for balanced, class-wise representation learning. TGFA maximizes the effect of feature augmentation by generating new anchors that emulate the important pixels of the teacher via a spatial attention map and capture its accurate feature correlations through a class-wise channel distribution. Additionally, to further diversify the oversampled features of the minority classes, we harness color variations as image-level augmentation on student inputs. When distilling DeepLabV3/PSPNet-ResNet101 teachers into ResNet18-based students, TGFA-CKD establishes state-of-the-art semantic segmentation KD results, surpassing the strongest CKD baselines by up to + 1.22, + 0.28, + 0.18, and + 0.35 mIoU points on Cityscapes, CamVid, PASCAL VOC, and ADE20K, respectively.

키워드