PaperlayerView exact revision on arXiv

Source-linked AI summary

VoCo: A Simple-yet-Effective Volume Contrastive Learning Framework for 3D Medical Image Analysis

Linshan Wu, Jiaxin Zhuang, Hao Chen

arXiv:2402.17300v2eess.IV

TL;DR

3D medical image pretraining can lack high-level semantics, a reported concern for downstream tasks. VoCo uses contextual position priors to form base-crop prototypes and predict sub-volume positions by contrast; it outperforms existing state-of-the-art 3D medical image SSL methods.

  • Problem

    3D medical image SSL methods often reconstruct raw information, while high-level semantics remain lacking in pretraining and can hinder downstream performance.

    Source: §2 · p. 3

  • Method

    VoCo forms prototypes from base crops at different positions and predicts randomly cropped sub-volumes' contextual positions by contrasting their similarities to the bases.

    Source: §1 · p. 2, §3.2 · p. 3

  • Results

    VoCo clearly outperforms existing state-of-the-art 3D medical image SSL methods across experiments on six downstream tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper concludes that contextual position priors can be encoded into model representations for downstream tasks requiring high-level semantics.

    Source: §1 · p. 2, §5 · p. 8

Abstract

from arXiv · show

Self-Supervised Learning (SSL) has demonstrated promising results in 3D medical image analysis. However, the lack of high-level semantics in pre-training still heavily hinders the performance of downstream tasks. We observe that 3D medical images contain relatively consistent contextual position information, i.e., consistent geometric relations between different organs, which leads to a potential way for us to learn consistent semantic representations in pre-training. In this paper, we propose a simple-yet-effective Volume Contrast (VoCo) framework to leverage the contextual position priors for pre-training. Specifically, we first generate a group of base crops from different regions while enforcing feature discrepancy among them, where we employ them as class assignments of different regions. Then, we randomly crop sub-volumes and predict them belonging to which class (located at which region) by contrasting their similarity to different base crops, which can be seen as predicting contextual positions of different sub-volumes. Through this pretext task, VoCo implicitly encodes the contextual position priors into model representations without the guidance of annotations, enabling us to effectively improve the performance of downstream tasks that require high-level semantics. Extensive experimental results on six downstream tasks demonstrate the superior effectiveness of VoCo. Code will be available at https://github.com/Luffy03/VoCo.

1. Introduction

VoCo uses relatively consistent geometric relations between organs to train contextual-position predictions for 3D medical image pretraining. The paper reports that VoCo outperforms existing state-of-the-art 3D medical image SSL methods.

1. Introduction

  • Prior work reports that a lack of high-level semantics in pretraining can hinder downstream performance, while expert annotations are expensive.Existing 3D medical image SSL methods commonly reconstruct raw information from augmented images.Source: §1 · pp. 1–2
  • Relatively consistent geometric relations between organs offer a potential source of consistent semantic representations for 3D medical image pretraining.The paper proposes contextual-position prediction to encode these priors for downstream tasks requiring high-level semantics.Source: §1 · p. 2
  • VoCo represents non-overlapping volumes from different positions as bases, then predicts randomly cropped sub-volumes' positions by contrasting their similarities to those bases.Feature discrepancy is enforced among the base volumes, which serve as class assignments for different positions.Source: §1 · p. 2
  • VoCo clearly outperforms existing state-of-the-art 3D medical image SSL methods.The introduction reports extensive experiments across six downstream tasks.Source: §1 · p. 2

2. Related Works

Prior SSL methods use instance- or prototype-level contrast, while medical-image approaches often reconstruct low-level information and may leave high-level semantics underrepresented. VoCo instead uses contextual position priors to form prototypes and predict positions through volume contrast.

2. Related Works

  • Instance-level contrastive learning compares augmented views or model perturbations, whereas prototype-level methods contrast inputs against prototypes generated by clustering or online updates.The cited prototype approaches include online clustering [5] and randomly initialized prototypes updated during training [6] [57] [15] [16].Source: §2 · p. 2
  • VoCo uses contextual position priors to form base-crop prototypes instead of randomly initializing and updating prototypes or clustering the dataset.The paper presents this as avoiding time-consuming clustering while using the relatively consistent geometric relations between organs.Source: §2 · pp. 2–3
  • Medical-image SSL commonly reconstructs low-level information from augmented images, but such methods can leave high-level semantics underrepresented for downstream tasks.Examples include rotation [52] [76] [53], multi-view crops [71] [32], and masking [74] [56].Source: §2 · p. 3
  • Natural-image position methods predict shuffled-patch order or patch locations, but inconsistent object geometry and direct linear outputs limit their approach.The cited methods include shuffled-patch prediction [43] and location prediction [65] [8].Source: §2 · p. 3
  • VoCo introduces contextual-position prediction for 3D medical images and predicts positions through volume contrast rather than a direct linear output layer.The paper describes this prediction paradigm as more intuitive and effective.Source: §2 · p. 3

3. Methodology

VoCo uses a prediction branch to relate cropped volumes to contextual positions and a regularization branch to make base features more discriminative. Its method predicts positions through volume contrast and overlap-based labels, while regularizing relationships among bases.

3.1. Overall Framework

  • The overall framework pairs contextual-position prediction between cropped volumes with regularization that increases feature discrepancy among bases.Source: §3.1 · p. 3

3.2. Contextual Position Prediction

  • Random sub-volumes are assigned contextual positions by comparing their features with base features and training against labels proportional to overlap with each base.Source: §3.2 · pp. 3–4

3.3. Volume Contrast for Regularization

  • The regularizer pushes cosine similarities between distinct bases toward zero, promoting linearly independent, discriminative assignments; the total objective combines this loss with prediction loss, using λ=1.0 in experiments.Source: §3.3 · p. 5

4. Experiments

Experiments evaluate VoCo on six downstream datasets spanning segmentation and classification, including datasets unseen during pre-training. VoCo outperforms compared methods overall, while ablations show gains from position prediction and regularized, appropriately sized bases.

4.1. Datasets

  • Pre-training uses about 1.6k CT scans from BTCV, TCIA Covid19, and LUNA; BTCV and TCIA Covid-19 alone are used for BTCV downstream experiments.The other downstream tasks use all three pre-training datasets, and five of the six downstream datasets are unseen during pre-training.Source: §4.1 · p. 5

4.2. Implementation details

  • Experiments use SwinUNETR, AdamW, cosine learning-rate scheduling, and 100K pre-training steps, with no foundation models or post-processing.The comparison includes general and medical SSL methods, with consistent settings for MAE and MoCo v3.Source: §4.2 · p. 5

4.3. Experiments on downstream tasks

  • 83.85% Dice Score on BTCV with VoCo exceeds GL-MAE’s 82.01%, while LiTs results show 96.52% Dice Score with Swin-UNETR after VoCo pre-training.For LiTs, VoCo pre-training improves Swin-UNETR by 3.10%.Source: Table 1 · p. 6, Table 2 · p. 6, §4.3 · p. 7
  • VoCo outperforms state-of-the-art methods across six downstream datasets, as summarized in the overall comparison figure.The evaluation spans segmentation and COVID-19 classification.Source: Figure 5 · p. 7, §4.3 · p. 8
  • On MSD Spleen [1] and MM-WHS, VoCo achieves Dice scores of 96.34% and 90.54%, respectively, surpassing prior-best GVSL scores of 95.47% and 88.27%.Source: §4.3 · p. 7

4.4. Ablation Study

  • Position-prediction loss raises Dice from 80.53% to 82.96% on BTCV and 86.11% to 88.82% on MM-WHS; regularization further improves performance.With 4×4×1 bases, VoCo reaches 83.85% and 90.54% Dice on BTCV and MM-WHS; larger or Z-direction base settings do not improve results, and 4×4×1 balances performance and efficiency.Source: Table 6 · p. 7, Table 7 · p. 7, §4.4 · p. 8

5. Conclusion and Future Directions

The conclusion presents VoCo as an SSL framework that uses contextual positions between organs to learn semantic representations through position prediction. It reports superior downstream performance and identifies several directions for future evaluation and extension.

5. Conclusion and Future Directions

  • VoCo represents crops from different positions as bases and predicts a random crop’s contextual position by contrasting its similarity to those bases.The authors say this encodes contextual position priors into representations for downstream tasks requiring high-level semantics.Source: §5 · p. 8
  • The authors report superior performance in extensive experiments, while proposing larger pre-training datasets and more downstream datasets as future evaluations.They also name label-efficient evaluation, such as semi-supervised learning, and mining inter-volume relationships as future directions.Source: §5 · p. 8

6. Appendix

  • The appendix describes pre-training and downstream datasets, implementation details, and additional ablations and experiments on a 2D medical dataset.Source: §6 · p. 9

A. Datasets

The experiments use six downstream datasets spanning CT and MRI segmentation and COVID-19 classification. Pre-training uses about 0.8k CT scans for BTCV and 1.6k for the other four downstream tasks.

A. Datasets

  • Pre-training uses BTCV and TCIA Covid19 (about 0.8k CT scans) for BTCV evaluation, and adds LUNA to reach 1.6k scans for the other four downstream tasks.BTCV is also used in pre-training.Source: A. Datasets · p. 9
  • BTCV has 24 training scans and 6 validation scans, with one background and thirteen organ classes.The dataset is used in pre-training.Source: A. Datasets · p. 9
  • LiTS provides 131 abdominal CT volumes, and evaluation uses liver masks rather than tumor labels.The dataset releases annotations for training and validation.Source: A. Datasets · p. 9
  • MSD Spleen uses 32 scans for training and 9 for validation to support comparisons with prior methods.Source: A. Datasets · p. 9
  • MM-WHS is unseen during pre-training and contains seven cardiac and vascular classes.Its data splits are listed in Table 8.Source: A. Datasets · p. 9
  • BraTS evaluates generalization from CT-only pre-training to MRI, using 387 scans for experiments and 97 for evaluation across WT, TC, and ET classes.Source: A. Datasets · p. 9
  • CC-CCII is a COVID-19 classification dataset with 2,516 training and 1,644 validation scans across NCP, CP, and Normal classes.Source: A. Datasets · p. 9

B. Implementation Details

The implementation follows comparatively consistent prior-work settings and uses one A100 GPU. Supplementary experiments include qualitative segmentation visualizations and additional ablation and 2D analyses.

B. Implementation Details

  • Figure 7 shows qualitative segmentation results for LiTS and MSD Spleen, with Scratch denoting training from scratch.Blue and red dashed boxes mark the highlighted differences.Source: Figure 7 · p. 10
  • Figure 8 shows qualitative segmentation results for MM-WHS and BraTS, with Scratch denoting training from scratch.Blue dashed boxes mark the highlighted differences.Source: Figure 8 · p. 10
  • The implementation uses Monai and PyTorch on one NVIDIA A100 GPU, with settings intended to remain consistent with prior methods.Source: B. Implementation Details · pp. 9–10
  • Fine-tuning runs for 3,000 epochs on BTCV, 1,000 on LiTS, MSD Spleen, and MM-WHS, 500 on BraTS, and 100 on CC-CCII.Source: B. Implementation Details · p. 10

C. Experiments

  • The supplementary experiments cover ablations, 2D medical image analysis, and other experiments omitted from the main paper because of page limits.Source: C. Experiments · p. 10

C.1 Ablation Studies

The ablation tests the loss-balance parameter on BTCV and MM-WHS and selects equal weighting for VoCo.

C.1 Ablation Studies

  • The loss-balance parameter λ is tested at 0.5, 1.0, and 1.5 using Dice Score on BTCV and MM-WHS.Source: §C.1 · p. 10
  • The tested λ settings do not matter a lot for Dice Score on BTCV and MM-WHS.Source: §C.1 · p. 10
  • VoCo sets λ to 1, treating the loss functions as equally important.Source: §C.1 · p. 10

C.2 2D Medical Image Analysis

The 2D evaluation pre-trains and fine-tunes on NIH ChestX-ray, using the same data split as prior work. VoCo achieves competitive results, and the authors report that position priors remain useful in 2D images.

C.2 2D Medical Image Analysis

  • The NIH ChestX-ray experiment pre-trains and fine-tunes on the same dataset, using 60% of the data for pre-training and the remainder for fine-tuning.The experiments use 3D-UNet and settings following prior works.Source: §C.2 · pp. 10–11
  • VoCo achieves competitive results on NIH ChestX-ray, with results reported in Table 11.Source: §C.2 · p. 11, Table 11 · p. 11
  • The authors conclude that contextual position priors persist in 2D images and benefit VoCo training, despite 2D images containing less information than 3D scans.Source: §C.2 · p. 11

C.3 Dice Scores of MM-WHS dataset

The Dice scores for seven organs on the MM-WHS dataset are reported in Table 12.

C.3 Dice Scores of MM-WHS dataset

  • Table 12 reports Dice scores for seven organs on MM-WHS.The supplied passage does not provide the score values.Source: §C.3 · p. 11

C.4 Validation results on the leaderboard

The authors report leaderboard evaluations and visualizations across several datasets. For MSD, they use an offline validation set because its leaderboard had not been updated.

C.4 Validation results on the leaderboard

  • BTCV test results were verified, and Flare23 and Amos22 public leaderboard test sets were evaluated without ensembling, extra data, or other tricks.The comparison uses SwinUNETR with the same network and settings; results are referenced in Table 13.Source: §C.4 · p. 11, Table 13 · p. 11
  • Because the MSD leaderboard had not been updated, the authors report offline validation results and say more results will appear in the revision.They state that they strictly follow SwinUNETR settings; results are referenced in Table 14.Source: §C.4 · p. 11, Table 14 · p. 11

C.5 More Visualization Results

  • Visualization results are shown for LiTs, MSD Spleen, MM-WHS, and BraTs in Figures 7 and 8.Source: §C.5 · p. 11

Related papers

Loading 2402.17300v2…