[Paper Review] Embedding Task Knowledge into 3D Neural Networks via Self-supervised Learning

Jiuwen Zhu, Yuexiang Li+2 more|arXiv (Cornell University)|Jun 10, 2020

Domain Adaptation and Few-Shot Learning13 references4 citations

TL;DR

This paper proposes Task-related Contrastive Predictive Coding (TCPC), a self-supervised learning method that enhances 3D medical image classification by embedding lesion-related prior knowledge into 3D CNNs. Using supervoxel clustering via SLIC to localize candidate lesions, TCPC applies a calibrated 'U'-shaped contrastive predictive coding scheme on sub-volumes around these regions, achieving state-of-the-art performance on brain hemorrhage and lung cancer datasets with minimal labeled data.

ABSTRACT

Deep learning highly relies on the amount of annotated data. However, annotating medical images is extremely laborious and expensive. To this end, self-supervised learning (SSL), as a potential solution for deficient annotated data, attracts increasing attentions from the community. However, SSL approaches often design a proxy task that is not necessarily related to target task. In this paper, we propose a novel SSL approach for 3D medical image classification, namely Task-related Contrastive Prediction Coding (TCPC), which embeds task knowledge into training 3D neural networks. The proposed TCPC first locates the initial candidate lesions via supervoxel estimation using simple linear iterative clustering. Then, we extract features from the sub-volume cropped around potential lesion areas, and construct a calibrated contrastive predictive coding scheme for self-supervised learning. Extensive experiments are conducted on public and private datasets. The experimental results demonstrate the effectiveness of embedding lesion-related prior-knowledge into neural networks for 3D medical image classification.

Motivation & Objective

• To address the challenge of limited annotated 3D medical image data in deep learning.
• To improve 3D medical image classification by incorporating task-specific prior knowledge into self-supervised pre-training.
• To develop a proxy task in self-supervised learning that is more relevant to clinical diagnosis than generic image augmentation tasks.
• To demonstrate that lesion-aware self-supervised pre-training enhances downstream classification performance, especially with small labeled datasets.

Proposed method

• Uses Simple Linear Iterative Clustering (SLIC) to generate supervoxels and identify potential lesion regions in 3D medical volumes.
• Crops sub-volumes centered on supervoxel regions to focus on diagnostically relevant areas.
• Proposes a calibrated 'U'-shaped path for support feature generation in contrastive predictive coding, improving contextual feature learning.
• Applies contrastive predictive coding (CPC) on sub-volumes to predict features of a target cube using features from a spatially structured support set.
• Employs a 3D CNN backbone for feature extraction and pre-training on unlabeled data using the TCPC proxy task.
• Combines transfer learning from UCF101 with TCPC pre-training to further improve performance on medical datasets.

Experimental results

Research questions

• RQ1 Can self-supervised learning methods that incorporate lesion-related prior knowledge improve 3D medical image classification?
• RQ2 How does a task-specific proxy task based on lesion localization compare to generic proxy tasks in self-supervised 3D learning?
• RQ3 To what extent does TCPC improve performance when only a small amount of labeled data is available?
• RQ4 Can combining transfer learning from natural images with task-aware self-supervised pre-training further boost classification accuracy?

Key findings

• TCPC achieved 78.38% average accuracy and 0.921 AUC on the brain hemorrhage dataset, outperforming all baselines including 3D CPC and Rubik’s cube.
• On the lung cancer dataset, TCPC achieved 95.86% accuracy with 3D VGG and 98.88% with 3D ResNet-18, surpassing 3D Jigsaw and Rubik’s cube methods.
• The UCF101 + TCPC cascade method improved accuracy by 1% on average, reaching 97.85% and 99.40% for 3D VGG and ResNet-18 on lung nodules.
• TCPC demonstrated robustness to hyper-parameters, with performance increasing as cube size increased due to richer contextual information.
• With minimal labeled data, TCPC consistently improved accuracy and AUC over training from scratch, especially in low-data regimes.