[Paper Review] 3D Self-Supervised Methods for Medical Imaging

Aiham Taleb, Winfried Loetzsch+5 more|arXiv (Cornell University)|Jun 6, 2020

Medical Image Segmentation Techniques67 references128 citations

TL;DR

Introduces five 3D self-supervised proxy tasks (3D-CPC, 3D-Rot, 3D-Jig, 3D-RPL, 3D-Exe) to learn from unlabeled 3D medical images, improving data efficiency and downstream performance across brain tumor, pancreas tumor, and diabetic retinopathy tasks, with open-source code.

ABSTRACT

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.

Motivation & Objective

• Reduce reliance on expensive expert annotations in 3D medical imaging.
• Learn rich 3D representations from unlabeled data using self-supervised proxy tasks.
• Enable data-efficient fine-tuning on downstream tasks like segmentation and detection.
• Demonstrate transfer of learned 3D representations from large unlabeled corpora to small downstream datasets.

Proposed method

• Propose five 3D self-supervised tasks: 3D Contrastive Predictive Coding (3D-CPC), 3D Rotation prediction (3D-Rot), 3D Jigsaw puzzles (3D-Jig), Relative 3D patch location (3D-RPL), and 3D Exemplar networks (3D-Exe).
• 3D-CPC uses an encoder and context network to predict latent representations of future patches in an inverted-pyramid 3D context using InfoNCE loss.
• 3D-RPL trains to classify the relative location of a query patch within a 3D grid, with jitter to avoid shortcuts.
• 3D-Jig solves permutations of a 3D patch grid as a P-way classification task.
• 3D-Rot predicts discretized 3D rotations (10 classes) to force semantic understanding.
• 3D-Exe uses triplet loss to pull transformed positives closer and push negatives away in embedding space.

Experimental results

Research questions

• RQ1 Do 3D self-supervised proxy tasks improve learned representations from unlabeled 3D medical images compared to 2D or scratch training?
• RQ2 Are 3D representations transferable to downstream tasks and to smaller datasets via transfer learning from large unlabeled corpora?
• RQ3 Do the 3D proxy tasks offer data-efficiency, performance gains, and faster convergence in brain tumor segmentation, pancreas tumor segmentation, and retinopathy detection?
• RQ4 How do 3D methods compare with 2D counterparts and supervised baselines on realistic medical imaging benchmarks?

Key findings

• 3D self-supervised pretraining yields superior data-efficiency, outperforming scratch and 2D-slice baselines especially in low-data regimes.
• Across tasks, 3D methods achieve competitive or superior Dice scores compared with state-of-the-art BraTS baselines, and often require fewer downstream epochs.
• 3D representations transfer effectively from large unlabeled 3D corpora (e.g., UK Biobank) to smaller, downstream datasets.
• 3D methods outperform their 2D counterparts on the BraTS brain tumor segmentation task, highlighting the value of full 3D context.
• Open-source implementations for both 3D and 2D variants are released to facilitate adoption and extension.