EDGE: Edge distillation and gap elimination for heterogeneous networks in 3D medical image segmentation

Compressing the cumbersome Vision Transformers (ViTs) or ConvNets into compact students can effectively facilitate the deployment of 3D medical image segmentation models on embedded devices. However, teacherstudent heterogeneity coupled with huge capacity discrepancies poses tremendous challenges in facilitating effective knowledge transfer. This paper introduces the Edge Distillation and Gap Elimination (EDGE) framework for 3D medical image segmentation while effectively mitigating the knowledge gap across heterogeneous networks. We observe that the degradation of knowledge transfer between heterogeneous networks in 3D medical image segmentation may be attributed to a lack of edge-sensitive knowledge transfer. Initially, Edge Constraint Knowledge Distillation (ECKD) is introduced to facilitate knowledge transfer of edge-sensitive logits and enhance the student's edge perception ability. Subsequently, we propose quantifying the knowledge gap as Hausdorff Distance (HD) in 3D abdominal CT images. Accordingly, Segmentation Refinement Knowledge Distillation (SRKD) is further introduced after the ECKD module to address knowledge transfer degradation in heterogeneous networks by minimizing the HD. Lastly, Scale Adaptation Knowledge Distillation (SAKD) is presented to enhance the student's performance in segmenting multi-organs with varying scales after projecting multiple stages of features into logits space. Extensive experiments are conducted on two 3D segmentation datasets, WORD and BTCV. Our proposed EDGE method consistently exhibits excellent performance across various teacher-student combinations. These results suggest that our method holds promise in bridging the knowledge gap with a significant margin compared to other competitive methods. The effectiveness of the individual modules is further confirmed via ablation experiments and visualization results.