[Paper Review] Imbalance-Aware Self-Supervised Learning for 3D Radiomic Representations

Hongwei Li, Fei-Fei Xue+6 more|arXiv (Cornell University)|Mar 6, 2021

Radiomics and Machine Learning in Medical Imaging24 references4 citations

TL;DR

This paper proposes a 3D Siamese self-supervised learning framework that enhances radiomic representations for medical image analysis by addressing data imbalance through unsupervised re-weighting and batch composition balancing. It achieves state-of-the-art performance in brain tumor classification and lung cancer staging, outperforming both traditional radiomics and supervised learning, especially under label scarcity.

ABSTRACT

Radiomic representations can quantify properties of regions of interest in medical image data. Classically, they account for pre-defined statistics of shape, texture, and other low-level image features. Alternatively, deep learning-based representations are derived from supervised learning but require expensive annotations from experts and often suffer from overfitting and data imbalance issues. In this work, we address the challenge of learning representations of 3D medical images for an effective quantification under data imbalance. We propose a \emph{self-supervised} representation learning framework to learn high-level features of 3D volumes as a complement to existing radiomics features. Specifically, we demonstrate how to learn image representations in a self-supervised fashion using a 3D Siamese network. More importantly, we deal with data imbalance by exploiting two unsupervised strategies: a) sample re-weighting, and b) balancing the composition of training batches. When combining our learned self-supervised feature with traditional radiomics, we show significant improvement in brain tumor classification and lung cancer staging tasks covering MRI and CT imaging modalities.

Motivation & Objective

• To address data imbalance in 3D medical imaging self-supervised representation learning, which is underexplored despite its prevalence in clinical datasets.
• To develop a self-supervised 3D Siamese network that learns high-level, discriminative features without requiring annotated data.
• To integrate self-supervised representations with traditional radiomics to improve downstream diagnostic performance in neuroimaging and oncology.
• To evaluate the effectiveness of unsupervised strategies—sample re-weighting and balanced batch composition—for mitigating class imbalance in self-supervised training.
• To demonstrate that self-supervised features can outperform or complement supervised learning, particularly when labeled data is limited.

Proposed method

• A 3D Siamese network with shared weights is used to learn invariant representations by comparing two augmented views of the same 3D medical volume.
• Extensive data augmentation is applied across four categories: affine transforms, appearance enhancement, contrast changes, and noise injection to generate positive pairs.
• An unsupervised re-weighting strategy (RE) is introduced to adjust loss weighting based on class frequency, reducing bias toward majority classes.
• A clustering-based sampling strategy (SE) constructs balanced batches by grouping samples into clusters and sampling from each to ensure class representation.
• The SE module uses k-means clustering to form new batches with equal class distribution, improving training stability and representation quality.
• The final model combines self-supervised features with classical radiomics, using a late fusion strategy for downstream classification.

Experimental results

Research questions

• RQ1 Can self-supervised 3D representation learning improve radiomic feature quality in imbalanced medical imaging datasets?
• RQ2 How do unsupervised re-weighting and balanced batch sampling mitigate data imbalance in self-supervised 3D medical image learning?
• RQ3 Does combining self-supervised features with traditional radiomics outperform supervised learning in low-data regimes?
• RQ4 How does the proposed method compare to proxy-task-based self-supervised learning (e.g., Rubik’s cube) in terms of discriminative power?
• RQ5 What is the optimal number of clusters (k) in the SE module for balancing representation quality and generalization?

Key findings

• The combination of traditional radiomics and self-supervised 3D representations (3DSiam+SE) achieved 91.6% sensitivity and 71.1% specificity in brain tumor classification, outperforming supervised learning (88.8%/71.1%) under full-label conditions.
• With only 50% of training labels, the hybrid approach achieved 84.8% sensitivity and 76.3% specificity, surpassing supervised learning (80.4%/56.6%) in the same setting.
• The SE module reduced feature redundancy, as shown by lower correlation coefficients in the covariance analysis, indicating more compact and discriminative representations.
• The optimal number of clusters (k) in the SE module was found to be 3, as higher values (e.g., k=5) led to decreased AUC due to biased sampling.
• The proposed method outperformed proxy-task-based self-supervised learning (e.g., Rubik’s cube) in both brain tumor and lung cancer tasks, indicating superior discriminativeness of direct self-supervised representation learning.
• In lung cancer staging, the hybrid model achieved 53.8% overall accuracy and 37.5% minor-class accuracy with full labels, exceeding supervised learning (52.6%/37.5%) and traditional radiomics (49.0%/37.5%).