AssistDistil for Medical Image Segmentation

Abstract

Deep learning models have demonstrated significant effectiveness in addressing intricate object segmentation and image classification tasks. Nevertheless, their widespread use is impeded by high computational demands, limiting their applicability on resource-constrained devices and in contexts like medical image segmentation. This paper proposes AssistDistil, a semi-knowledge distillation technique designed to facilitate the transfer of knowledge from a larger teacher network to a more compact student model. During the inference process, the student model works in conjunction with the teacher model by condensing the teacher model’s latent information into its own latent representation, thereby boosting its representational capacity. The effectiveness of the proposed approach is demonstrated for multiple case studies in medical image segmentation task of eye segmentation, skin lesion segmentation, and chest X-ray segmentation. Experimental results on the IIITD Cataract Surgery, HAM10000, PH2, Shenzhen and Montgomery chest X-ray datasets demonstrate the efficacy of the proposed approach both in terms of accuracy and computational cost. For example, in comparison to the AUNet-based teacher model, the proposed approach achieves a similar mIOU with only 0.5% of the model size. In the future, we plan to explore knowledge distillation approaches to improve the distillation process in case of large model capacity gap between teacher and student networks. With fewer parameters, we intend for the student model to attain performance comparable to that of the teacher model without additional assistance.