Enhancing pretraining efficiency for medical image segmentation via transferability metrics

TL;DR

This paper introduces a new transferability metric based on contrastive learning to identify optimal pretraining points, reducing training time and enhancing medical image segmentation performance.

Contribution

We propose a novel contrastive learning-based transferability metric that predicts the best pretraining stage for medical image segmentation tasks, improving efficiency and outcomes.

Findings

Shorter pretraining often yields better segmentation results.

ImageNet accuracy is a poor predictor of downstream performance.

Our metric effectively indicates optimal pretraining timing.

Abstract

In medical image segmentation tasks, the scarcity of labeled training data poses a significant challenge when training deep neural networks. When using U-Net-style architectures, it is common practice to address this problem by pretraining the encoder part on a large general-purpose dataset like ImageNet. However, these methods are resource-intensive and do not guarantee improved performance on the downstream task. In this paper we investigate a variety of training setups on medical image segmentation datasets, using ImageNet-pretrained models. By examining over 300 combinations of models, datasets, and training methods, we find that shorter pretraining often leads to better results on the downstream task, providing additional proof to the well-known fact that the accuracy of the model on ImageNet is a poor indicator for downstream performance. As our main contribution, we introduce a novel transferability metric, based on contrastive learning, that measures how robustly a pretrained model is able to represent the target data. In contrast to other transferability scores, our method is applicable to the case of transferring from ImageNet classification to medical image segmentation. We apply our robustness score by measuring it throughout the pretraining phase to indicate when the model weights are optimal for downstream transfer. This reduces pretraining time and improves results on the target task.