Source-linked AI summary
3D Self-Supervised Methods for Medical Imaging
Aiham Taleb, Winfried Loetzsch, Noel Danz, Julius Severin, Thomas Gaertner, Benjamin Bergner, Christoph Lippert
TL;DR
The paper addresses the cost of expert annotation and the mismatch between 2D representations and 3D medical-image structure. It proposes five 3D self-supervised proxy tasks and evaluates their pretrained representations across medical-imaging downstream tasks. The methods improve data-efficiency, performance, and convergence speed, including when transferred from large unlabeled 3D data to smaller task-specific datasets.
Problem
Expensive expert annotation and reliance on 2D slices limit representation learning for 3D medical images, whose full spatial context contains rich anatomical information.
Method
The paper extends five self-supervised proxy tasks to 3D images and pretrains encoders on unlabeled data for downstream fine-tuning.
Results
3D pretraining improves data-efficiency, downstream performance, and convergence speed across three tasks, outperforming 2D counterparts and scratch baselines in reported comparisons.
Takeaways & Limitations
Large unlabeled 3D corpora can provide transferable representations for smaller medical-imaging datasets, offering an alternative to ImageNet-based transfer.
Abstract
from arXiv · showhide
Self-supervised learning methods have witnessed a recent surge of interest after proving successful in multiple application fields. In this work, we leverage these techniques, and we propose 3D versions for five different self-supervised methods, in the form of proxy tasks. Our methods facilitate neural network feature learning from unlabeled 3D images, aiming to reduce the required cost for expert annotation. The developed algorithms are 3D Contrastive Predictive Coding, 3D Rotation prediction, 3D Jigsaw puzzles, Relative 3D patch location, and 3D Exemplar networks. Our experiments show that pretraining models with our 3D tasks yields more powerful semantic representations, and enables solving downstream tasks more accurately and efficiently, compared to training the models from scratch and to pretraining them on 2D slices. We demonstrate the effectiveness of our methods on three downstream tasks from the medical imaging domain: i) Brain Tumor Segmentation from 3D MRI, ii) Pancreas Tumor Segmentation from 3D CT, and iii) Diabetic Retinopathy Detection from 2D Fundus images. In each task, we assess the gains in data-efficiency, performance, and speed of convergence. Interestingly, we also find gains when transferring the learned representations, by our methods, from a large unlabeled 3D corpus to a small downstream-specific dataset. We achieve results competitive to state-of-the-art solutions at a fraction of the computational expense. We publish our implementations for the developed algorithms (both 3D and 2D versions) as an open-source library, in an effort to allow other researchers to apply and extend our methods on their datasets.
1 Introduction
Medical imaging benefits from machine learning, but expert annotation of 3D images is expensive and existing 2D transfer features may not match medical-image distributions. The paper therefore develops 3D self-supervised proxy tasks that use full spatial context to improve downstream learning.
- Motivation: Expert annotation of 3D medical images is expensive, limiting large-scale machine-learning development.Medical imaging supports prevention, detection, diagnosis, and treatment, increasing the value of effective automated analysis.
- Motivation: ImageNet features can be reused for related tasks, but natural and medical images have significantly different distributions.This domain gap limits straightforward transfer learning in supervision-starved medical settings.
- Motivation: Self-supervised proxy tasks derive supervision from the data, encouraging semantic representations that facilitate data-efficient supervised fine-tuning.This approach can reduce the burden of manual annotation.
- Motivation: Using 2D slices for 3D imaging can discard spatial context and rich anatomical information needed by downstream tasks.The paper identifies full 3D context as a better basis for representation learning in 3D domains.
- Contributions: The paper proposes five 3D self-supervised tasks and evaluates them across four datasets and three downstream tasks.The tasks are 3D-CPC, 3D Rotation prediction, 3D Jigsaw puzzles, Relative 3D patch location, and 3D Exemplar networks.
2 Related work
Self-supervised representation learning creates proxy tasks from unlabeled data, with spatial context among several possible supervision sources. Prior work spans patch-location, jigsaw, rotation, contrastive, video, and medical-imaging applications, but the paper focuses on extending these ideas to 3D medical data.
- Self-supervised learning: Self-supervised learning constructs embeddings by creating supervised proxy tasks from the input data itself.Embeddings learned for proxy tasks can support other downstream tasks.
- Spatial-context methods: Spatial context is a common supervision source, including predicting one image patch’s position relative to another.This patch-based idea was later extended to jigsaw puzzles.
- Other proxy tasks: Other self-supervised methods use colors, clustering, rotation prediction, object saliency, image reconstruction, or contrastive predictive coding.CPC methods use contrastive learning in latent space and have achieved strong results on imaging benchmarks.
- 3D representations: 3D convolutions in video models jointly capture spatial features within frames and temporal features across stacked frames.This video literature provides a related use of 3D-CNNs outside medical imaging.
- Medical imaging: Medical self-supervision has been applied to tasks including endoscopy depth estimation, registration, body-part recognition, cardiac segmentation, and instrument segmentation.These applications demonstrate varied uses of self-supervised learning in medical contexts.
3 Self-Supervised Methods
The paper formulates five proxy tasks that learn encoders from unlabeled 3D images for later downstream fine-tuning. The methods use predictive, spatial, permutation, rotation, and metric-learning objectives to encode 3D structure.
- Overview: Each proposed task learns a pretrained encoder from unlabeled 3D images without manual annotation during pretraining.The resulting encoder can subsequently be fine-tuned for downstream tasks.
- 3D-CPC: 3D-CPC splits scans into overlapping patches and predicts adjacent-patch latent representations from a context vector.The context network summarizes latent vectors from preceding patches, while InfoNCE distinguishes the correct future representation from negatives.
- 3D-CPC: 3D-CPC uses an inverted-pyramid neighborhood to balance computational cost and performance.The context includes the current patch and selected patches above it.
- 3D-RPL: 3D-RPL samples a 3D patch grid, uses a central reference patch, and predicts the query patch’s relative location.Random gaps between neighboring patches discourage shortcut solutions such as edge continuity.
- 3D-Jig: 3D-Jig shuffles an n × n × n grid of 3D patches using predefined permutations, then trains the model to predict the applied permutation.Random jitter is also used in this task.
- 3D-Rot: 3D-Rot rotates input scans by selected degrees and trains the network to predict the applied rotation.The described implementation considers multiples of 90 degrees along the three coordinate axes.
- 3D-Exe: 3D-Exe replaces prohibitively expensive surrogate-class classification with triplet loss over transformed positives and different-sample negatives.The loss brings transformed versions of the same sample closer while separating different samples.
4 Experimental Results
The methods are evaluated through fine-tuning on three downstream medical-imaging tasks, measuring data-efficiency, performance, and convergence speed. Across segmentation experiments, 3D pretraining improves results over scratch and 2D-slice baselines, while transfer learning is also assessed for diabetic retinopathy detection.
- Three downstream tasks evaluate the learned representations: brain-tumor segmentation from 3D MRI, pancreas-tumor segmentation from 3D CT, and diabetic-retinopathy detection from 2D fundus images.
- Additional implementation details cover architectures, training procedures, Exemplar augmentation, and decoder initialization for segmentation tasks.
- Brain Tumor Segmentation: BraTS experiments compare training from scratch, training on 2D slices, supervised pretraining, and BraTS challenge baselines.
- Brain Tumor Segmentation: The 3D methods outperform scratch and 2D counterparts on BraTS, with large data-efficiency gains using few labeled samples and comparable results to BraTS challenge baselines.
- Pancreas Tumor Segmentation: Pancreas experiments report better performance than baselines at 5% and 10% labeled-data settings, and much higher performance within the first 20 epochs on the full dataset.
- Diabetic Retinopathy Detection: The diabetic-retinopathy experiment pretrains 2D methods on roughly 170K UK Biobank images, then fine-tunes on roughly 5590 Kaggle fundus images and evaluates labeled-data subsets.
5 Conclusion
The proposed self-supervised methods improve downstream performance, data efficiency, and convergence, particularly with limited labeled data, while 3D methods outperform 2D counterparts in 3D domains.
- Our methods improve downstream performance, data efficiency, and speed of convergence across three medical imaging tasks.
- 3D methods outperform their 2D counterparts, supporting the use of full 3D spatial context for 3D imaging tasks.
- Our models converge faster than supervised baselines, helping achieve high results using only a few epochs.
- With fewer labeled samples, our methods consistently outperform supervised baselines in pancreas segmentation and diabetic retinopathy classification.
- Pretraining on large unlabeled corpora and fine-tuning on smaller downstream datasets can reduce manual annotation effort and support transfer beyond ImageNet features.
Broader Impact
Medical imaging supports prevention, detection, diagnosis, and treatment, but supervised machine learning is constrained by the cost and difficulty of expert annotation, especially for 3D scans.
- Medical imaging aids disease prevention, early detection, diagnosis, and treatment, making automated analysis relevant to patient healthcare.
- Machine-learning analysis of medical images is increasingly important as digitization expands and imaging can provide rapid clinical insights.
- Expert annotation is expensive, time-consuming, and difficult to generate at scale, particularly for 3D medical scans.
- These annotation constraints motivate annotation-efficient approaches such as self-supervised representation learning from unlabeled data.
A Implementation and training details for all tasks
The implementation provides flexible architectures, training configurations, preprocessing, multimodal input handling, and task-specific 3D self-supervised settings.
- The code supports multiple encoder, decoder, and classifier architectures, losses, augmentation techniques, and evaluation metrics.
- Architecture details: 3D segmentation models use a five-level residual 3D U-Net encoder, with the decoder added during downstream fine-tuning.
- Optimization details: Adam optimization uses task-dependent learning rates, with warm-up freezing of pretrained encoder weights during initial fine-tuning epochs.
- Input preprocessing: 3D scans are cropped, resized to 128 × 128 × 128, and intensity-normalized to [0, 1]; 2D images are resized to 224 × 224.
- Processing multimodal inputs: For MRI modality mismatch between UK Biobank pretraining and BraTS fine-tuning, only the pretrained input-layer weights are duplicated.
- Task settings: Patch-based tasks use 3 × 3 × 3 patches and 64-dimensional patch embeddings, while 3D-CPC and 3D-Exe use latent codes of size 1024.
- Augmentation in Exemplar: Exemplar augmentations include flipping, rotation, brightness and contrast changes, and zooming, with omission reducing downstream performance and high rates increasing time cost.
B Detailed experimental results
Detailed results show consistent low-data advantages and faster convergence for the proposed methods across pancreas segmentation and diabetic retinopathy classification, with 3D methods stronger than 2D counterparts in pancreas.
- Pancreas segmentation: Our methods consistently outperform the supervised baseline in low-data pancreas segmentation cases.
- Pancreas segmentation: Within the first 20 epochs on the full pancreas dataset, our models achieve much higher performance than the from-scratch baseline.
- Diabetic retinopathy: Pretrained models outperform the supervised baseline with fewer labels in diabetic retinopathy classification.
- Speed of convergence: Detailed retinopathy convergence results associate faster training with achieving high results using only a few epochs.