Self-supervised learning improves robustness of deep learning lung tumor segmentation to CT imaging differences

Med Phys 2025 AI 6 Explanations View Original
Original Paper (PDF)

Unable to display PDF. Download it here or view on PMC.

Plain-English Explanations
Pages 1-2
Making AI Tumor Segmentation Robust Across CT Scan Variations

The CT Variation Problem: Deep learning models for lung tumor segmentation are trained on specific CT scan datasets but must perform reliably across diverse real-world CT acquisitions - varying contrast agents, scanner manufacturers, reconstruction kernels, and slice thicknesses. This distribution shift between training and deployment settings is a major barrier to clinical adoption.

Self-Supervised Learning for Robustness: Self-supervised learning (SSL) trains neural networks using unlabeled data by creating pretext tasks where the image itself serves as ground truth. Unlike supervised learning that requires expensive manual annotations, SSL can leverage large pools of unannotated CT scans to pre-train models that generalize better across imaging variations.

Wild vs Self-Pretraining: This study from Memorial Sloan Kettering compared two SSL pretraining strategies. Wild-pretraining uses a large, diverse, uncurated dataset of over 10,000 CT scans covering many body regions and diseases. Self-pretraining uses only the task-specific downstream dataset (lung cancer CTs) without labels. The key question is which approach yields more robust segmentation models.

Three Architectures Compared: CNN-based (PRCLv2), Vision Transformer (ViT), and hierarchical Swin Transformer networks were each tested under both pretraining strategies, providing a comprehensive comparison of how architecture interacts with pretraining approach to determine robustness.

TL;DR: This Memorial Sloan Kettering study systematically compared wild-pretraining (10,412 diverse CTs) versus self-pretraining strategies across CNN, ViT, and Swin architectures for robust NSCLC segmentation, finding wild-pretrained Swin most robust to CT imaging variations.
Pages 2-5
Datasets, Pretext Tasks, and Experimental Design

Upstream Dataset for Wild-Pretraining: A total of 10,412 uncurated 3D CT scans were assembled from public cancer imaging archives and internal institutional datasets (head and neck, lung, and esophageal cancer patients undergoing radiotherapy), covering diverse body regions, scanner types, and imaging protocols without any curation.

Downstream Dataset: 377 labeled NSCLC CT scans from the TCIA NSCLC-Radiomics dataset (stage I-IV, single institution) served as both the self-pretraining data and the fine-tuning training set. Two separate test sets evaluated generalization: LRad (156 public early-stage patients post-surgery) and LC (196 internal advanced-stage patients).

Pretext Tasks - SMIT: Transformer models used three pretext tasks from the Self-Distilled Masked Image Transformer (SMIT) framework: Masked Image Prediction (MIP) - predicting masked image patches from unmasked context; Image Token Distillation (ITD) - matching global feature tokens between teacher and student networks; and Masked Patch Token Distillation (MPD) - matching features for masked patches. A teacher-student architecture with exponential moving average updates was used.

Robustness Evaluation: Segmentation was evaluated on CT scans varying by: (1) contrast vs non-contrast, (2) three convolution reconstruction kernels (soft to sharp), and (3) two slice thicknesses (2.5 mm and 5 mm). Statistical comparisons between wild- and self-pretrained model instances used paired Wilcoxon signed rank tests. Feature reuse was analyzed using Centered Kernel Alignment (CKA).

TL;DR: 10,412 uncurated CTs for wild-pretraining, 377 labeled NSCLC CTs for self-pretraining and fine-tuning, with SMIT pretext tasks and systematic robustness testing across contrast, reconstruction kernels, and slice thicknesses.
Pages 4-7
Wild-Pretrained Swin Dominates - Architecture Determines Benefit

Swin Benefits Most from Wild-Pretraining: Wild-pretrained Swin produced significantly higher segmentation accuracy than self-pretrained Swin on the LRad early-stage dataset (DSC p less than 0.0001, HD95 p less than 0.0001), and demonstrated significantly better robustness to reconstruction kernel variations. Wild-pretrained Swin also achieved the highest tumor detection rate on the phantom dataset.

ViT and CNN Show No Clear Benefit: Neither wild-pretrained ViT nor CNN showed statistically significant improvement over their self-pretrained counterparts across the main test sets. CNN models showed similar performance across all pretraining strategies, suggesting that CNN's strong inductive bias (spatial locality) makes it less dependent on large pretraining datasets.

Architecture-Specific Explanation: The results can be explained by inductive bias differences. ViT treats images as token sequences, destroying local spatial relationships - making it data-hungry and benefiting from diverse pretraining. Swin adds hierarchical windowed attention that restores some spatial inductive bias but still benefits from large-scale pretraining. CNN's full spatial inductive bias makes it most self-sufficient.

Two-Stage Pretraining Fails: Sequential wild-pretraining followed by self-pretraining (two-stage) actually decreased accuracy compared to wild-pretraining alone. CKA analysis showed that self-pretraining stage discarded wild-pretrained features (very low CKA values in middle layers), explaining why the two-stage approach did not accumulate benefits from both stages.

TL;DR: Wild-pretrained Swin significantly outperformed self-pretrained Swin on early-stage and diverse-kernel test sets, while ViT and CNN showed no clear benefit from wild-pretraining - and two-stage pretraining paradoxically reduced performance by discarding pretrained features.
Pages 7-8
Feature Reuse Analysis Explains Performance Differences

CKA as Interpretability Tool: Centered Kernel Alignment (CKA) measured similarity between feature representations before and after fine-tuning. High CKA means features were preserved from pretraining into the fine-tuned model (feature reuse); low CKA means features changed substantially (feature adaptation). The optimal pattern for medical image segmentation appears to be high reuse in early layers (edge detection) and low reuse in late layers (task-specific specialization).

Wild-Pretraining Produces Optimal Feature Reuse: Wild-pretrained Swin showed the desired pattern: high CKA in early layers (1-4) indicating reuse of low-level features, but low CKA in late layers (13-14) indicating task-specific specialization near the output. This pattern - retained generic features combined with specialized high-level features - is associated with superior downstream performance.

Self-Pretraining Shows Feature Redundancy: Self-pretrained models showed uniformly high CKA across all layers including late layers and off-diagonal (cross-layer) comparisons, indicating high feature redundancy with limited adaptation to the segmentation task. The model changed little from pretraining to fine-tuning, suggesting pretraining on the same data created features too specific to the training set.

Pretext Task Impact: Among pretext tasks, the SMIT combination (MIP + ITD + MPD) produced the best performance with wild-pretraining. Contrastive learning produced the worst results and highest feature redundancy - indicating that learning global image similarity (not local structure) is insufficient for precise segmentation requiring local boundary detection.

TL;DR: CKA analysis revealed that wild-pretrained Swin achieves optimal feature reuse - retaining generic low-level features while adapting high-level task-specific features during fine-tuning - explaining its superior robustness compared to self-pretraining's feature redundancy.
Pages 8-9
Building Clinically Deployable Lung Tumor Segmentation Models

Real-World CT Variation is the Critical Challenge: In clinical practice, lung cancer CT scans come from many different scanners, protocols, and acquisition settings. A segmentation model that works well only on CTs similar to its training data will fail frequently in real deployment. Wild-pretrained Swin's demonstrated robustness to contrast, reconstruction kernels, and slice thickness makes it substantially more suitable for clinical translation.

Selecting the Right Architecture Matters: This study provides actionable guidance for teams developing clinical segmentation tools: Swin transformers should be preferred when large unlabeled pretraining datasets are available, as they extract maximal benefit. CNN models may be more appropriate when only small task-specific labeled datasets are available and pretraining resources are limited.

Implications for Radiotherapy Planning: Robust automated GTV segmentation across diverse CT protocols is essential for radiotherapy planning centers that receive patients imaged at different facilities. Wild-pretrained Swin models could enable reliable AI-assisted contouring even when the incoming CTs differ substantially from training data, reducing the need for site-specific model retraining.

Foundation Model Development Path: The demonstrated benefits of wild-pretraining with 10,000 CT scans suggest that scaling to hundreds of thousands of scans - as done in computer vision foundation models - could further improve robustness and performance. This study charts a clear path toward lung cancer segmentation foundation models.

TL;DR: Wild-pretrained Swin transformers are the recommended architecture for clinical lung tumor segmentation when diverse unlabeled CTs are available, providing significantly better robustness across real-world CT protocol variations essential for clinical deployment.
Page 9
Scaling and Refinement for Clinical Deployment

Scaling Pretraining Data: Wild-pretraining with 10,412 CTs improved Swin robustness substantially. Future work should systematically study whether scaling to 100,000+ CTs (easily achievable by incorporating national imaging registries or public datasets) continues to improve robustness, potentially approaching foundation model performance.

Layer Freezing for Efficient Fine-Tuning: The CKA analysis showed that early layers change little during fine-tuning (suggesting they could be frozen) while late layers need significant adaptation. Future experiments with partial layer freezing during fine-tuning could reduce the labeled data requirements, enabling robust models with fewer expensive manual annotations.

Multi-Site Validation: All testing was performed on datasets from a limited number of institutions. Prospective multicenter evaluation with CT scans from diverse hospitals, scanner manufacturers (Siemens, GE, Philips, Canon), and clinical workflows is needed to confirm that wild-pretrained Swin's robustness generalizes to the full range of clinical environments.

Extension to Lymph Node and Metastasis Segmentation: Current work focused on primary tumor GTV segmentation. Extending the wild-pretraining approach to lymph node delineation and metastatic lesion detection would address additional critical clinical needs in lung cancer staging and radiotherapy planning, where accurate nodal contouring is equally important.

TL;DR: Scaling pretraining to 100,000+ CTs, developing partial layer freezing for efficient fine-tuning with fewer labels, multicenter validation, and extension to lymph node segmentation are the priority directions for translating these findings into clinical deployment.
Citation: Open Access, 2025. Available at: PMC12044619.