Uncertainty Driven Adaptive Self-Knowledge Distillation for Medical Image Segmentation

Deep learning has recently significantly improved the precision of medical image segmentation. However, due to the commonly limited dataset scale and reliance on hard labels (one-hot vectors) in medical image segmentation, deep learning models often overfit, which reduces segmentation performance. To mitigate the problem, we propose an uncertainty driven adaptive self-knowledge distillation (UAKD) model for medical image segmentation that regularizes the training process through self-generated soft labels. The innovation of UAKD is to integrate uncertainty estimation into soft label generation and student network training, ensuring accurate supervision and effective regularization. In detail, UAKD introduce teacher network ensembling to reduce semantic bias in soft labels caused by the teacher networks' fitting biases. An adaptive knowledge distillation mechanism is also proposed, which utilizes uncertainty to generate adaptive weights for soft labels to compute the loss function, thereby efficiently transferring reliable knowledge from the teacher network to the student network while suppressing unreliable information. Finally, we introduce a gradient ascent based cyclic ensemble method to reduce teacher network overfitting on the training data, further enhancing the aforementioned teacher ensembling and uncertainty estimation. Experiments on three medical image segmentation tasks show that UAKD outperforms existing regularization methods and demonstrates the effectiveness of uncertainty estimation for assessing soft label reliability.