Feature Subspace Projection Knowledge Distillation for Medical Image Segmentation

Abstract

Feature-based knowledge distillation facilitates feature knowledge transfer by aligning intermediate features of students and high-performance teachers such as TranUnet and MISSFormer in medical image segmentation. However, the bias-variance coupling resulting from redundancy or noise within high-dimensional features presents a significant challenge for effective knowledge transfer. To tackle this issue, we propose a feature subspace projection knowledge distillation (FSP-KD) method to decouple bias and variance in the high-dimensional feature space. This method decomposes the feature space into two components: the variance-dependent distribution and the bias-dependent distribution. The bias-dependent distribution is modeled as a weighted post-projection feature distribution using the feature subspace projection (FSP) module. Likewise, the variance-dependent distribution is represented by a weighted pre-projection feature distribution. Additionally, a conditional adversarial mechanism (CADV) module is integrated at the logits layer to prompt the student to identify higher-order discrepancies from the teacher. This approach leverages conditional generative adversarial learning to improve the holistic alignment between student and teacher distributions. Extensive experiments are carried out on three benchmark datasets for medical image segmentation: Synapse, Flare2022, and m2caiSeg. The experimental results show that our proposed FSP-KD method has achieved state-of-the-art performance. Notably, FSP-KD has outperformed the teacher MISSFormer when used in a teacher-student setup with ResNet18. Ablation experiments and visualization results provide additional confirmation of the effectiveness of each module.