Multi-Task Multi-Scale Contrastive Knowledge Distillation for Efficient Medical Image Segmentation

TL;DR

This paper proposes a multi-task, multi-scale contrastive knowledge distillation method to improve medical image segmentation models, especially when training data and computational resources are limited.

Contribution

It introduces a novel architecture combining multi-scale feature distillation and supervised contrastive learning for effective knowledge transfer from large to small models in medical imaging.

Findings

Multi-scale feature distillation improves student model performance.

Supervised contrastive learning enhances knowledge transfer effectiveness.

Distilling features at multiple scales is crucial for optimal results.

Abstract

This thesis aims to investigate the feasibility of knowledge transfer between neural networks for medical image segmentation tasks, specifically focusing on the transfer from a larger multi-task "Teacher" network to a smaller "Student" network. In the context of medical imaging, where the data volumes are often limited, leveraging knowledge from a larger pre-trained network could be useful. The primary objective is to enhance the performance of a smaller student model by incorporating knowledge representations acquired by a teacher model that adopts a multi-task pre-trained architecture trained on CT images, to a more resource-efficient student network, which can essentially be a smaller version of the same, trained on a mere 50% of the data than that of the teacher model. To facilitate knowledge transfer between the two models, we devised an architecture incorporating multi-scale feature distillation and supervised contrastive learning. Our study aims to improve the student model's performance by integrating knowledge representations from the teacher model. We investigate whether this approach is particularly effective in scenarios with limited computational resources and limited training data availability. To assess the impact of multi-scale feature distillation, we conducted extensive experiments. We also conducted a detailed ablation study to determine whether it is essential to distil knowledge at various scales, including low-level features from encoder layers, for effective knowledge transfer. In addition, we examine different losses in the knowledge distillation process to gain insights into their effects on overall performance.