Enhancing pretraining efficiency for medical image segmentation via transferability metrics

📝 Paper Summary

Medical Image Segmentation Transfer Learning

A novel contrastive transferability metric optimizes ImageNet pretraining length for medical image segmentation, reducing training time while outperforming standard full pretraining and from-scratch initialization.

Core Problem

Pretraining U-Net encoders on ImageNet is computationally expensive, and overlong pretraining can degrade downstream medical segmentation performance due to source-domain overfitting.

Why it matters:

Medical image segmentation suffers from scarce labeled data, making pretraining essential
Full ImageNet pretraining is computationally expensive, and ImageNet accuracy poorly predicts downstream segmentation performance
Existing transferability metrics require matching source and target task types, failing for classification-to-segmentation transfer

Concrete Example: When pretraining a Swin U-Net encoder for the IDRiD retina dataset, fully pretraining for 300 epochs yields worse downstream segmentation performance than pretraining for only a fraction of the time, but ImageNet accuracy cannot detect this optimal stopping point.

Key Novelty

Contrastive Robustness Score for Transferability

Measures how robustly a pretrained model represents target medical data by comparing augmented positive image pairs against random negative pairs
Applies across different task types, enabling evaluation when transferring from an ImageNet classification source to a medical segmentation target
Acts as an early stopping indicator during pretraining, identifying the optimal checkpoint before the model overfits to ImageNet-specific features

Evaluation Highlights

Achieves over 98% average Transferability Indicator Score (TIS) offline, outperforming both full pretraining and no pretraining baselines
Online robustness (early stopping) reduces pretraining time by up to a factor of 300 while maintaining >97% TIS on average
Improves Swin U-Net decoder-only transferability by +8.7% absolute TIS compared to full ImageNet pretraining across datasets

Breakthrough Assessment

7/10

Provides a practical, theoretically sound metric for a common but under-explored problem (classification-to-segmentation transfer), significantly saving compute while slightly boosting performance.

⚙️ Technical Details

Problem Definition

Setting: Transfer learning from a source classification task (ImageNet) to a target medical image segmentation task.

Inputs: Medical images (e.g., MRI, X-ray, Retina scans)

Outputs: Pixel-wise segmentation masks

Pipeline Flow

ImageNet Pretraining (Encoder)
Robustness Evaluation (Checkpoint Selection)
Target Task Fine-tuning (Encoder-Decoder)

System Modules

Encoder

Extracts hierarchical features from input images; pretrained on ImageNet and evaluated for transferability

Model or implementation: Basic U-Net encoder, ResNet50, or Swin Transformer

Decoder

Reconstructs spatial resolution to output a segmentation mask; randomly initialized before downstream training

Model or implementation: Basic U-Net decoder, Convolutional decoder with Attention Gates, or Swin Transformer decoder

Modeling

Base Model: Basic U-Net, ResNet50 Attention U-Net, Swin U-Net

Training Method: Supervised classification pretraining on ImageNet followed by segmentation fine-tuning

Objective Functions:

Purpose: Measure the robustness of the pretrained representation on the target dataset using a contrastive triplet formulation.

Formally: \varrho_{S,T}(f_\vartheta) = 1 - \frac{1}{|T|} \sum_{x_i \in T} \max(0, d(\hat{f}_\vartheta(\tilde{q}_i), \hat{f}_\vartheta(k^+_i)) - d(\hat{f}_\vartheta(\tilde{q}_i), \hat{f}_\vartheta(k^-_i)) + \varepsilon)

Adaptation: Full fine-tuning or Decoder-only training (frozen encoder)

Training Data:

Source: ImageNet
Target: ACDC (1902 train / 1076 test), COVID-QU-Ex (2330 train / 583 test), IDRiD (6426 train / 1350 test slices)

Key Hyperparameters:

pretraining_epochs: 300
pretraining_optimizer: Adam
simple_pretraining_lr: 0.004
+ 7 more
simple_pretraining_batch_size: 4096
advanced_pretraining_lr_swin: 0.001
advanced_pretraining_lr_conv: 0.004
downstream_epochs: 150
downstream_optimizer: SGD with momentum 0.9
downstream_weight_decay: 1e-4
robustness_margin_epsilon: 0.5

Comparison to Prior Work

vs. LEEP: Robustness metric works for classification-to-segmentation transfer, whereas LEEP requires matching task types (regression/classification)
vs. LogME: Does not require target labels to estimate transferability, enabling unsupervised evaluation of target data representation
vs. TransRate: Uses mutual information between features and labels [not cited in paper]

Limitations

Randomly initialized encoders sometimes produce higher robustness scores than models pretrained for a few epochs, complicating online early stopping
Online method (early stopping) is prone to finding local maxima and generally performs worse than the offline method
Spearman correlation between the robustness score and downstream performance is relatively low, even if their optimal checkpoints coincide
Robustness is a better indicator when encoder weights are frozen during downstream training, but this setup yields worse absolute performance than full fine-tuning

Reproducibility

Code: https://github.com/aielte-research/MedSegPretrainImageNet/

Code and results are publicly available via GitHub. Checkpoints evaluated at epochs 1, 5, 20, 50, 100, 150, 200, 250, and 300. Preprocessing details, learning rate schedules, and augmentation strategies are fully documented in the appendix.

📊 Experiments & Results

Evaluation Setup

Transfer learning from ImageNet classification to medical image segmentation across three domains (MRI, X-ray, Retina).

Benchmarks:

ACDC (Cardiac MRI multi-class segmentation)
COVID-QU-Ex (Chest X-ray binary segmentation)
IDRiD (Retina scan binary segmentation (microaneurysms))

Metrics:

Dice index
Transferability Indicator Score (TIS)
Statistical methodology: Mean performance obtained from 5 independent downstream training runs per configuration.

Key Results

Benchmark	Metric	Baseline	This Paper	Δ
Main comparison results demonstrate that the offline robustness score consistently selects better pretraining checkpoints than relying on ImageNet accuracy.
ACDC	TIS	0.982	0.989	+0.007
COVID-QU-Ex	TIS	0.950	0.999	+0.049
Model breakdown results show that the robustness metric is particularly effective for transformer-based architectures like Swin U-Net.
Average across datasets	TIS	0.900	0.987	+0.087
Average across datasets	TIS	0.809	0.990	+0.181
Ablation experiments confirm that cosine distance provides the strongest transferability signal among tested metrics.
IDRiD	TIS	0.748	0.993	+0.245

Main Takeaways

Shorter ImageNet pretraining often leads to better medical segmentation performance than full convergence, confirming ImageNet accuracy is a poor transferability indicator
The proposed contrastive robustness score successfully identifies near-optimal pretraining checkpoints across different U-Net architectures and medical domains
Using robustness as an online early stopping criterion prevents unnecessary pretraining steps (often stopping before epoch 150 out of 300) while exceeding the performance of full pretraining
Cosine distance is consistently the best similarity metric for the robustness score compared to L2 distance and Pearson correlation
Extracting features from the second-to-last encoder level yields better transferability prediction than the final level, likely because the final level overfits to ImageNet-specific features

📚 Prerequisite Knowledge

Prerequisites

Transfer learning and fine-tuning
U-Net architectures
Contrastive learning

Key Terms

TIS: Transferability Indicator Score—a metric measuring the relative downstream performance of a model chosen by an indicator compared to the best possible model in a set

Robustness score: The proposed metric measuring how well a pretrained model represents target data by distinguishing augmented versions of the same image from other images

U-Net: An encoder-decoder neural network architecture with lateral skip connections, widely used for medical image segmentation

ImageNet pretraining: Training a model on a large general dataset (ImageNet) before fine-tuning on a specialized dataset to overcome data scarcity

Contrastive learning: A training paradigm where a model learns to pull representations of similar (positive) pairs together and push dissimilar (negative) pairs apart

Swin U-Net: A U-Net-style model that uses Swin Transformers instead of convolutional layers for both the encoder and decoder

Offline robustness: Evaluating the robustness score on all saved checkpoints after full pretraining finishes to select the best one

Online robustness: Evaluating the robustness score continuously during pretraining and using it as an early stopping condition

Dice index: A common evaluation metric for image segmentation that measures the overlap between the predicted mask and the ground truth

ACDC: Automated Cardiac Diagnosis Challenge—a dataset of cardiac MRI images

COVID-QU-Ex: A dataset of chest X-ray images for COVID-19 infection segmentation

IDRiD: Indian Diabetic Retinopathy Image Dataset—a dataset of retina scans