Banner Banner

Active Learning with the nnUNet and Sample Selection with Uncertainty-Aware Submodular Mutual Information Measure

Bernhard Föllmer
Kenrick Schulze
Christian Wald
Sebastian Stober
Wojciech Samek
Marc Dewey

February 13, 2024

Annotating medical images for segmentation tasks is a time-consuming process that requires expert knowledge. Active learning can reduce this annotation cost and achieve optimal model performance by selecting only the most informative samples for annotation. However, the effectiveness of active learning sample selection strategies depends on the model architecture and training procedure used. The nnUNet has achieved impressive results in various automated medical image segmentation tasks due to its self-configuring pipeline for automated model design and training. This raises the question of whether the nnUNet is applicable in an active learning setting to avoid cumbersome manual configuration of the training process and improve accessibility for non-experts in deep learning-based segmentation. This paper compares various sample selection strategies in an active learning setting in which the self-configuring nnUNet is used as the segmentation model. Additionally, we propose USIM - Uncertainty-Aware Submodular Mutual Information Measure, a sample selection strategy for UNet-like architectures. The method combines uncertainty and submodular mutual information to select batches of uncertain, diverse, and representative samples. We evaluate the performance gain and labeling costs on three medical image segmentation tasks with different segmentation challenges. Our findings demonstrate that utilizing nnUNet as the segmentation model in an active learning setting is feasible, and most sampling strategies outperform random sampling. Furthermore, we demonstrate that our proposed method yields a significant improvement compared to existing baseline methods.