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

TL;DR

This study shows that wild-pretrained Swin transformers are more robust to CT imaging differences in lung tumor segmentation than self-pretrained models, highlighting the importance of pretraining data source and architecture choice.

Contribution

It compares the robustness of wild versus self-pretrained transformer models for lung tumor segmentation in CT images, revealing the advantages of wild pretraining and Swin architecture.

Findings

Wild-pretrained Swin outperformed self-pretrained Swin in robustness.

ViT models showed similar accuracy regardless of pretraining type.

Masked image prediction improved accuracy over contrastive learning.

Abstract

Self-supervised learning (SSL) is an approach to extract useful feature representations from unlabeled data, and enable fine-tuning on downstream tasks with limited labeled examples. Self-pretraining is a SSL approach that uses the curated task dataset for both pretraining the networks and fine-tuning them. Availability of large, diverse, and uncurated public medical image sets provides the opportunity to apply SSL in the "wild" and potentially extract features robust to imaging variations. However, the benefit of wild- vs self-pretraining has not been studied for medical image analysis. In this paper, we compare robustness of wild versus self-pretrained transformer (vision transformer [ViT] and hierarchical shifted window [Swin]) models to computed tomography (CT) imaging differences for non-small cell lung cancer (NSCLC) segmentation. Wild-pretrained Swin models outperformed self-pretrained Swin for the various imaging acquisitions. ViT resulted in similar accuracy for both wild- and self-pretrained models. Masked image prediction pretext task that forces networks to learn the local structure resulted in higher accuracy compared to contrastive task that models global image information. Wild-pretrained models resulted in higher feature reuse at the lower level layers and feature differentiation close to output layer after fine-tuning. Hence, we conclude: Wild-pretrained networks were more robust to analyzed CT imaging differences for lung tumor segmentation than self-pretrained methods. Swin architecture benefited from such pretraining more than ViT.