Adversarial class-wise self-knowledge distillation for medical image segmentation

Self-knowledge distillation enables knowledge transfer by dynamically constructing the next-stage learning objectives, thus providing more effective path cues to optimize the compact student. The challenge lies in formulating effective learning objectives for the upcoming stage that mitigate the interference of inter-class similarity in medical image segmentation. This paper presents an Adversarial Class-Wise Self-Knowledge Distillation (ACW-SKD). ACW-SKD leverages an auxiliary head to generate coarse segmentation results, which are then utilized as prediction masks to refine class-wise features, followed by mitigating inter-class similarity via class-wise feature distillation. A feature reconstruction module (FRM) is introduced in the penultimate feature layer and class-wise feature layer to avoid plugging in multiple intermediate branches for constructing the next-stage learning objectives. Furthermore, adversarial temperature loss is incorporated to further recognize inter-class similarity by integrating a learnable temperature module. Extensive experiments are conducted on three benchmark datasets: Synapse, FLARE2022, and M2caiSeg. The results indicate that ACW-SKD surpasses several offline knowledge distillation methods, self-knowledge distillation methods, and U-Net networks. Ablation studies and visual analysis further validate the efficacy of ACW-SKD. This method notably enhances segmentation accuracy for challenging classes and mitigates the influence of inter-class similarities in medical image segmentation. Moreover, ACW-SKD delivers comparable results to U-Net with a reduced computational demand, positioning it as a viable option for deploying efficient medical image segmentation models on mobile devices. Our codes are available at https://github.com/shenjq77/ACW-SKD.