Source-linked AI summary
Self-supervised Feature Learning for 3D Medical Images by Playing a Rubik's Cube
Xinrui Zhuang, Yuexiang Li, Yifan Hu, Kai Ma, Yujiu Yang, Yefeng Zheng
TL;DR
Limited and costly annotations constrain supervised learning for 3D medical images. The paper pre-trains 3D networks by recovering rearranged and rotated cubes, then fine-tunes them for medical-image tasks; this improves classification and segmentation accuracy without extra data.
Problem
Difficult-to-acquire annotations leave too few labeled 3D medical images to train deep-learning networks effectively.
Method
Rubik’s cube recovery pre-trains 3D networks to restore cube arrangement and orientation, with Siamese-Octad predicting permutation and rotation tasks.
Results
Fine-tuning Rubik’s cube pre-trained weights substantially improves 3D CNN accuracy on brain hemorrhage classification and brain tumor segmentation versus train-from-scratch or UCF101 transfer.
Takeaways & Limitations
Cube rearrangement and rotation provide transferable translationally and rotationally invariant features from raw volumetric medical data without extra data.
Takeaways & Limitations
For segmentation, Rubik’s cube recovery pre-trains only downsampling layers, so the U-Net decoder requires random initialization.
Abstract
from arXiv · showhide
Witnessed the development of deep learning, increasing number of studies try to build computer aided diagnosis systems for 3D volumetric medical data. However, as the annotations of 3D medical data are difficult to acquire, the number of annotated 3D medical images is often not enough to well train the deep learning networks. The self-supervised learning deeply exploiting the information of raw data is one of the potential solutions to loose the requirement of training data. In this paper, we propose a self-supervised learning framework for the volumetric medical images. A novel proxy task, i.e., Rubik's cube recovery, is formulated to pre-train 3D neural networks. The proxy task involves two operations, i.e., cube rearrangement and cube rotation, which enforce networks to learn translational and rotational invariant features from raw 3D data. Compared to the train-from-scratch strategy, fine-tuning from the pre-trained network leads to a better accuracy on various tasks, e.g., brain hemorrhage classification and brain tumor segmentation. We show that our self-supervised learning approach can substantially boost the accuracies of 3D deep learning networks on the volumetric medical datasets without using extra data. To our best knowledge, this is the first work focusing on the self-supervised learning of 3D neural networks.
1 Introduction
3D medical images are difficult to annotate at scale, limiting supervised deep-learning performance. The paper proposes Rubik’s cube recovery as a self-supervised proxy task for pre-training 3D networks and improving downstream analysis.
- 3D CT and MRI volumes require laborious physician review, with diagnosis affected by inter-observer variation.
- Difficult and time-consuming annotation of each 3D volume limits the amount of labeled data available for training 3D deep-learning models.
- Self-supervised learning pre-trains a CNN on a proxy task using non-annotated data, then fine-tunes it with limited annotations for a target task.
- Rubik’s cube recovery uses cube rearrangement and rotation to pre-train 3D networks on translationally and rotationally invariant features.
- Fine-tuning the pre-trained network improves brain hemorrhage classification and brain tumor segmentation accuracy without extra data.
2 Method
The method corrupts 3D volumes by rearranging and rotating cubes, then trains a Siamese 3D CNN to recover their configuration. Its combined ordering and orientation losses produce transferable encoder features for classification and segmentation.
- The pipeline pre-trains a 3D CNN on Rubik’s cube recovery and fine-tunes its weights on annotated target tasks.
- Rubik’s cube recovery: A volume is partitioned into cubes, randomly permuted and rotated, and the network recovers their original order and orientation.
- Rubik’s cube recovery: Compared with a 2D jigsaw puzzle, cube rotation adds difficulty and encourages networks to use more spatial information in 3D data.
- Network architecture.: Siamese-Octad uses M weight-sharing 3D CNN branches whose concatenated features feed separate permutation and rotation prediction tasks.
- Cube ordering.: Cube ordering selects K widely separated permutations and trains the network to classify which permutation was applied.
- Cube orientation.: Random cube rotation trains multi-label predictions for horizontal and vertical rotations, encouraging rotationally invariant features.
- The full objective combines permutation and rotation losses as L = αLP + βLR, with empirically best weights α = β = 0.5.
- For segmentation, pre-trained weights transfer to the FCN encoder, while the decoder remains randomly initialized.
3 Experiment
The brain hemorrhage experiment transfers Rubik’s cube pre-trained weights to classify four pathological causes using 3D CT volumes. It evaluates models with an 80:20 train-test split and classification accuracy.
- The experiment transfers Rubik’s cube pre-trained weights to pathological-cause classification and uses an 80:20 training-test split.
- 1486 brain CT scan images support classification of hemorrhage volumes into aneurysm, arteriovenous malformation, moyamoya disease, or hypertension.
- Each CT volume measures 230 × 270 × 30 voxels, and average classification accuracy is the evaluation metric.
BraTS-2018.
BraTS-2018 provides 285 preprocessed brain-tumor MR volumes with four imaging modalities for evaluating brain-tumor segmentation.
- The BraTS-2018 training set contains 285 brain-tumor MR volumes with native T1, post-contrast T1Gd, T2, and FLAIR modalities.
- The MR images are co-registered, resampled to 1 mm3 resolution, skull-stripped, and sized 240 × 240 × 155 voxels.
3.2 Performance on Solving Rubik’s Cube
The Siamese-Octad network successfully solves the Rubik’s cube proxy task, with cube rotation testing structural representation beyond cube ordering alone.
- 93.1% and 82.1% test accuracy were achieved for cube orientation on the brain hemorrhage and BraTS-2018 datasets, respectively.
- Cube-ordering accuracy decreased by 7.7% and 6.6% as random cube rotation increased difficulty on the two datasets.
- Cube rotation enabled the network to model rotated content and extract more brain structural information than rearrangement alone.
3.3 Fine-tuning Models on Target Tasks
The study fine-tunes Rubik’s cube-pre-trained networks on brain hemorrhage classification and brain tumor segmentation, comparing them with alternative initialization strategies.
- Pre-trained Rubik’s cube weights were fine-tuned for brain hemorrhage classification and brain tumor segmentation.
- Train-from-scratch and UCF101-pre-trained models were included as comparison strategies.
- Test results were reported in Table 2.
Baselines.
The experiments use train-from-scratch as the baseline and UCF101 video pre-training as a natural-data comparison, while Rubik’s cube pre-training yields the strongest reported hemorrhage-classification result.
- Train-from-scratch served as the baseline, while UCF101 provided 3D CNN pre-training from 13,320 videos across 101 action categories.
- 11.2% higher accuracy than train-from-scratch, Rubik’s cube pre-training reached 83.8% for brain hemorrhage classification.
- UCF101 pre-training improved classification by only 2.7%, reflecting the gap between natural video and volumetric medical data.
Brain tumor segmentation.
Rubik’s cube pre-training improves brain-tumor segmentation, while decoder initialization constrains transfer because only downsampling layers are pre-trained.
- Brain tumor segmentation: 76.2% and 77.3% mIoUs were achieved by Rubik’s cube fine-tuning with U-Net and 3D DUC, respectively.
- Brain tumor segmentation: Rubik’s cube pre-training produced more accurate tumor segmentations than UCF101 pre-training.
- Brain tumor segmentation: Because Rubik’s cube recovery pre-trains only downsampling layers, U-Net’s decoder must be randomly initialized and may degrade transferred representations.
- Brain tumor segmentation: 3D DUC outperformed 3D U-Net under every pre-training protocol, with gains of 1.6% for UCF101 and 1.1% for Rubik’s cube weights.
Comparison of solving different Rubik’s cubes.
Adding cube rotation to cube ordering improved downstream transfer, despite lowering cube-ordering accuracy, suggesting more robust translationally and rotationally invariant representations.
- +2.7% for brain hemorrhage classification and +2.3% for brain tumor segmentation were achieved by Rubik’s cube pre-training over cube ordering alone.The comparison uses models fine-tuned from Rubik’s cube recovery without cube rotation.
- Adding cube rotation lowered cube-ordering accuracy but improved downstream generalization on both evaluated tasks.The authors associate the improvement with a more robust feature representation.
- The combined cube-ordering and orientation tasks encouraged translationally and rotationally invariant features from raw 3D data.
4 Conclusion
The paper presents a self-supervised framework for volumetric medical images that pre-trains 3D neural networks with Rubik’s cube recovery. Its two operations are cube rearrangement and cube rotation, designed to enforce translationally and rotationally invariant features from raw 3D data.
- The framework formulates Rubik’s cube recovery as a proxy task for pre-training 3D neural networks on volumetric medical images.
- Cube rearrangement and cube rotation enforce translationally and rotationally invariant features from raw 3D data.