PaperlayerView exact revision on arXiv

Source-linked AI summary

A generalizable 3D framework and model for self-supervised learning in medical imaging

Tony Xu, Sepehr Hosseini, Chris Anderson, Anthony Rinaldi, Rahul G. Krishnan, Anne L. Martel, Maged Goubran

arXiv:2501.11755v2eess.IVcs.CV

TL;DR

Existing 3D medical SSL often uses simple pretext tasks and datasets similar to downstream applications, limiting generalizability and scalability. 3DINO adapts DINOv2 to 3D medical inputs and pretrains a general-purpose ViT on nearly 100,000 multimodal, multi-organ volumes; the resulting weights improve data efficiency and generalizability across segmentation and classification tasks.

  • Problem

    Many 3D medical SSL methods use simple pretext tasks and datasets similar to downstream applications, limiting generalizability and scalability.

    Source: Introduction · p. 3

  • Method

    3DINO adapts DINOv2 to 3D medical imaging and pretrains a general-purpose ViT on a large, multimodal, multi-organ dataset.

    Source: Introduction · pp. 3–4

  • Results

    3DINO improved data efficiency and generalizability across segmentation and classification, with 3DINO-ViT consistently outperforming compared pretrained models on the primary Dice segmentation metric.

    Source: Discussion · p. 6

  • Takeaways & Limitations

    3DINO-ViT can reduce labeled-data needs across diverse downstream medical imaging tasks without expensive retraining on in-domain unlabeled datasets.

    Source: Discussion · p. 7

    Pretraining primarily uses MRI and CT and does not cover a balanced list of organs; the full segmentation pipeline has relatively high runtime complexity.

    Source: Discussion · p. 7

Abstract

from arXiv · show

Current self-supervised learning methods for 3D medical imaging rely on simple pretext formulations and organ- or modality-specific datasets, limiting their generalizability and scalability. We present 3DINO, a cutting-edge SSL method adapted to 3D datasets, and use it to pretrain 3DINO-ViT: a general-purpose medical imaging model, on an exceptionally large, multimodal, and multi-organ dataset of ~100,000 3D medical imaging scans from over 10 organs. We validate 3DINO-ViT using extensive experiments on numerous medical imaging segmentation and classification tasks. Our results demonstrate that 3DINO-ViT generalizes across modalities and organs, including out-of-distribution tasks and datasets, outperforming state-of-the-art methods on the majority of evaluation metrics and labeled dataset sizes. Our 3DINO framework and 3DINO-ViT will be made available to enable research on 3D foundation models or further finetuning for a wide range of medical imaging applications.

Introduction

Existing 3D medical SSL often uses simple pretext tasks and datasets similar to downstream applications, limiting generalizability and scalability. 3DINO adapts SSL to 3D medical inputs and introduces a general-purpose model pretrained on a large, multimodal, multi-organ dataset.

Introduction

  • Similar-distribution pretraining limits generalizability and scalability, while detailed 3D labels are costly and separate task-specific pretraining adds overhead.The introduction notes that label scarcity can be especially acute for rare diseases and difficult-to-acquire high-resolution or specialized-modality data.Source: Introduction · p. 3
  • Scaling diverse SSL pipelines to 3D is computationally prohibitive, while slice-wise training omits full 3D anatomical context.DINOv2 is identified as a candidate for genuinely 3D representations because of its accuracy and computational-efficiency improvements [28].Source: Introduction · p. 3
  • 3DINO adapts DINOv2 for 3D medical inputs and 3DINO-ViT is pretrained on nearly 100,000 unlabeled volumes from 35 studies.The dataset includes MRI (N = 70,434), CT (N = 27,815), and brain PET (N = 566).Source: Introduction · p. 3
  • 3DINO converts image-specific operations to 3D and adds a 3D ViT-Adapter with 3D multi-scale deformable attention for segmentation.The authors describe the model as generalizing across multiple modalities simultaneously and make the framework, adapter, attention module, and weights available.Source: Introduction · p. 4

Results

Across segmentation and classification benchmarks, 3DINO-ViT improved performance, including with limited labeled data. The evaluation also covers an unseen organ and modality, using comparisons against several pretrained and randomly initialized baselines.

Results

  • The comparisons included random ViT and Swin initializations, Swin Transfer, and MONAI-ViT pretrained on the same unlabeled dataset.MONAI-ViT was adapted to a vanilla ViT for a more direct comparison of pretraining methods.Source: Results · p. 4

Within and out-of-distribution downstream tasks

  • Evaluation covered BraTS and BTCV segmentation, unseen left atrium MRI and breast ultrasound segmentation, and brain-age and COVID-CT-MD classification.The classification experiments also varied the percentage of labeled training data.Source: Within and out-of-distribution… · p. 4

Performance on segmentation and classification tasks

  • 3DINO-ViT significantly improved segmentation over SOTA techniques in most comparisons (p < 0.05) and matched full-label baselines with under 50% of labeled data.At 10% BraTS labels, Dice was 0.90 (0.88, 0.91) versus 0.87 (0.85, 0.89) for Random; at 25% BTCV, it was 0.77 (0.72, 0.81) versus 0.59 (0.53, 0.65).Source: Performance on segmentation and… · p. 5
  • With all labeled data, segmentation gains over the next-best baseline narrowed and were not significant: 0.8% Dice on BraTS and 0.9% on BTCV.Source: Performance on segmentation and… · p. 5
  • Across dataset sizes, 3DINO-ViT achieved 18.9% higher AUC on COVID-CT-MD and 5.3% higher AUC on ICBM than the next-best baselines.The reported class-specific AUC improvements over the next-best baseline were 23% for COVID-19 classification and 22% for individuals aged [40, 50) years.Source: Performance on segmentation and… · p. 5

3DINO patch-level representations for image segmentation

A lightweight decoder evaluation found strong segmentation results from 3DINO patch-level representations. The model also retained an advantage on segmentation tasks involving an unseen organ and modality, including with the full labeled datasets.

3DINO patch-level representations for image segmentation

  • 61% Dice improvement on BTCV and 15% on BraTS over the next-best baseline came from frozen 3DINO-ViT weights with a two-layer linear decoder.Both improvements were statistically significant: p < 0.05 on BTCV and p < 0.001 on BraTS.Source: 3DINO patch-level representations for… · p. 5

Performance on unseen organs and modalities

  • On out-of-distribution tasks, 3DINO-ViT improved Dice by 1.9% for left atrium segmentation and 24% for 3D ultrasound tumor segmentation over the next-best baseline with 25% labeled data; advantages were 1.1% and 0.8%, respectively, at 100%.Source: Performance on unseen organs and… · pp. 5–6

Discussion

3DINO-ViT improved performance across segmentation and classification tasks, including out-of-distribution data, while supporting both patch-level and image-level tasks. The study also notes that pretraining coverage is mainly MRI and CT and the organ mix is unbalanced.

Discussion

  • Across tasks and labeled-data sizes, frozen 3DINO-ViT outperformed state-of-the-art methods on most metrics; Dice gains over Random reached 13.0% on BraTS and 55.1% on BTCV.The Dice improvements are maximum gains over the Random encoder.Source: Discussion · p. 6
  • 3DINO-ViT outperformed other pretrained models on classification tasks in all metrics and improved class-wise pseudo-probabilities, including brain-age cases between classification bins.For some ICBM brain-age cases, the predictions better captured uncertainty.Source: Discussion · p. 6
  • The authors report 3DINO as the first 3D medical-image pretraining method to jointly improve patch-level and image-level tasks.They describe this claim as based on their extensive validation.Source: Discussion · p. 7
  • On unseen heart MRI and breast ultrasound segmentation, 3DINO significantly improved over other state-of-the-art methods.For left-atrium segmentation, the gain was 1.9% with 25% of labels, 9.3% with 10%, and 23% with 5%.Source: Discussion · p. 7
  • PCA visualizations showed foreground-background, organ-surface, and anatomical-axis variation; tumor regions were often distinct, while unseen ultrasound visualizations were less salient.The comparisons are reported for representations visualized across datasets and modalities.Source: Discussion · p. 7
  • Pretraining was mainly MRI and CT with an unbalanced organ list, and the full segmentation pipeline had relatively high runtime complexity.The authors suggest 3DINO-ViT could reduce labeled-data needs without retraining on in-domain unlabeled datasets.Source: Discussion · pp. 7–8

Figures

The figures depict 3DINO’s dual-objective approach, its pretraining data, and evaluations across segmentation and classification tasks. They also show example predictions and out-of-distribution evaluations.

Figures

  • Figure 1 combines image-level and patch-level objectives, while its dataset panel summarizes pretraining data across modalities and organs.The final loss sums the two objectives; the patch-level objective compares masked student patches with corresponding unmasked EMA teacher representations.Source: Figure 1 · p. 9, Figures · p. 10
  • Figure 2 reports BraTS and BTCV Dice and HD95, ICBM and COVID-CT-MD AUC and F1, across labeled-data sizes.Its caption specifies training-set percentages and sample counts; the figure also includes confidence intervals and per-class comparisons.Source: Figure 2 · p. 11, Figures · p. 12
  • Figures 3–5 visualize segmentation outputs and class-specific Dice alongside classification probabilities for selected cases and labeled-data sizes.Figure 3 identifies BraTS tumor-region and BTCV organ colors; Figure 4 separates BraTS scores by class, and Figure 5 shows predicted class probabilities.Source: Figure 3 · p. 12, Figure 4 · p. 13, Figure 5 · p. 14
  • Figure 6 evaluates left atrium segmentation on an unseen organ and breast-ultrasound tumor segmentation on an unseen modality.It includes Dice and HD95 scores, example predictions at different labeled-data sizes, and unsupervised visualizations.Source: Figure 6 · p. 15, Figures · p. 16

Pretraining unlabeled dataset

  • The pretraining dataset contains 98,815 unlabeled volumes after slice-count filtering and subset-based deduplication; its high-resolution adaptation subset contains 53,758 volumes.The original dataset includes publicly available data and one internal dataset.Source: Pretraining unlabeled dataset · p. 17

Image-level objective

  • DINO’s image-level objective uses self-distillation between student and teacher networks, with the teacher obtained as an exponential moving average of the student.The method forms global and local crops from an unlabeled image and trains the student to predict teacher representations.Source: Image-level objective · p. 17

DINOv2

  • DINOv2 adds masked patch prediction to the image-level objective, with the combined loss supporting patch representations for dense tasks such as segmentation.The cited work associates this objective with improved patch-level representation quality and robustness to corruption; the method also uses a short high-resolution adaptation stage.Source: DINOv2 · p. 18

3DINO

The 3D adaptation modifies the ViT encoder and positional encodings for volumetric inputs while retaining most of DINOv2’s low-level modules. It also adapts cropping and masking to 3D data.

3DINO

  • The ViT encoder flattens and projects 3D patches, and learned position encodings are interpolated to permit variable-sized inputs.Masking locations are sampled uniformly across patches rather than blockwise.Source: 3DINO · pp. 18–19
  • Medical imaging augmentations include random contrast adjustment, additive noise, Gibbs noise, and histogram shift instead of color-related operations for one-channel inputs.Intensity normalization maps the 0.05th and 99.95th percentiles to -1 and 1.Source: 3DINO · p. 19
  • Cropping applies 2D RandomResizedCrop in-plane and samples the cross-slice crop size independently to accommodate variable volume dimensions.This can stretch or squash the depth axis while preserving local-to-global crop correspondences; the final pretrained model is named 3DINO-ViT.Source: 3DINO · p. 19

3DINO pretraining implementation details

The implementation describes the 3DINO pretraining schedule and a separate ViT reimplementation of Swin pretraining for comparison. That comparison uses the same pretraining dataset as 3DINO-ViT.

3DINO pretraining implementation details

  • Standard 3DINO pretraining used ViT-Large with 16×16×16 patches, batch size 512, and 125,000 iterations over approximately nine days.Global and local crops were 96×96×96 and 48×48×48, respectively; high-resolution adaptation ran 12,500 iterations over approximately two days.Source: 3DINO pretraining implementation details · pp. 19–20

SOTA comparisons

  • MONAI-ViT reimplements Swin pretraining tasks on a vanilla ViT to compare pretraining algorithms while holding the dataset constant with 3DINO-ViT.The tasks are inpainting, contrastive coding, and rotation prediction; the ViT was pretrained for 100,000 iterations over approximately 10 days.Source: SOTA comparisons · p. 20, MONAI-ViT pretraining implementation… · p. 20

MONAI-ViT pretraining implementation details

  • The MONAI-ViT implementation used ViT-Large, 16×16×16 patches, 96×96×96 images, and batch size 64 on four A100 GPUs.Intensity scaling was adjusted for non-quantitative scans in the multimodal dataset.Source: MONAI-ViT pretraining implementation… · p. 20

Finetuning Methods

The experiments use six segmentation and classification datasets, with patient-level splits or randomly assigned train, validation, and test sets. They include an unseen organ and an unseen imaging modality.

Brain Tumor Segmentation (BraTS) dataset

  • BraTS 2021 provides four MRI sequences and expert labels for whole tumor, tumor core, and enhancing tumor; 1,084 patients remained after removing pretraining-overlapping scans.The remaining patients were split into 758 training, 108 validation, and 218 test cases.Source: Brain Tumor Segmentation (BraTS) dataset · p. 21

Beyond the Cranial Vault (BTCV) dataset

International Consortium for Brain Mapping (ICBM) dataset

  • ICBM supplies T1-weighted MRI scans for brain-age classification, with scans outside ages 20–60 discarded and splits made at the patient level.The resulting sets contain 756 training, 151 validation, and 233 test scans.Source: International Consortium for Brain… · p. 21

COVID-CT-MD dataset

  • COVID-CT-MD contains lung CT scans for classifying COVID-19, community-acquired pneumonia, and healthy patients.Its 305 patients were split into 214 training, 30 validation, and 61 test scans.Source: COVID-CT-MD dataset · p. 21

Left Atrium Segmentation Challenge (LA-SEG) dataset

  • LA-SEG is a left-atrial-cavity MRI segmentation benchmark used to assess generalization to an out-of-domain organ.The dataset contains 154 scans from 60 patients; the training set was divided into 80 training and 20 validation scans, with 54 test scans.Source: Left Atrium Segmentation Challenge… · pp. 21–22

Tumor Detection, Segmentation and Classification Challenge on Automated 3D Breast Ultrasound (TDSC-ABUS) dataset

  • TDSC-ABUS provides breast-ultrasound lesion segmentations for evaluating generalization to a modality absent from pretraining.Its 100 scans were split into 70 training, 10 validation, and 20 test scans.Source: Tumor Detection, Segmentation and… · p. 22

3D ViT-Adapter

The 3D ViT-Adapter adds spatial information to a vanilla ViT for dense tasks, while the multi-channel strategy uses pretrained single-channel weights without changing the patch embedding layer. (p07, p08, p15, p16)

3D ViT-Adapter

  • The 3D ViT-Adapter combines multi-scale spatial features with pretrained ViT features to produce a feature pyramid for dense segmentation.Its Spatial Prior Module extracts local multi-scale features, which the Injector and Extractor integrate and adapt using pretrained ViT features.Source: 3D ViT-Adapter · p. 22
  • The Spatial Prior Module uses 3D convolutions and scales of 1/4, 1/8, and 1/16 to avoid resampling errors for 112×112×112 inputs.Source: 3D ViT-Adapter · p. 22
  • 3D MSDA adds a depth-axis deformable offset and uses trilinear interpolation to sample 3D feature maps at learnable offsets.Each query attends to a fixed, small number of keys rather than all spatial locations.Source: 3D ViT-Adapter · p. 23
  • With eight attention heads, 3D MSDA initializes offsets toward the eight 3D octants to provide more even feature-map coverage.The 2D initialization instead uses equal angular separation between heads.Source: 3D ViT-Adapter · p. 23

Multi-channel inputs

  • For multi-channel inputs, each channel passes through the single-channel pretrained network separately, and its features are concatenated and linearly projected with GELU.The patch embedding layer is left unchanged to retain the pretrained weights.Source: Multi-channel inputs · pp. 23–24
  • For the ViT-Adapter, frozen Transformer features are averaged across channels, while Injector spatial information is copied across channels before addition.The resulting features are concatenated and linearly projected; the authors report benefits on multi-channel segmentation tasks.Source: Multi-channel inputs · p. 24

3DINO segmentation

The section describes frozen-weight 3DINO-ViT segmentation and classification evaluations, comparison setups, feature visualizations, and statistical analyses. (p20, p24, p28, p35, p37)

3DINO segmentation

  • Segmentation experiments froze high-resolution-adapted 3DINO-ViT weights and trained the ViT-Adapter and UNet-like decoder on random subsets of training labels.Inputs were 112×112×112, with a four-level feature pyramid and a decoder using transposed convolutions and encoder-decoder connections.Source: 3DINO segmentation · p. 24

SOTA segmentation comparisons

  • SOTA segmentation comparisons used end-to-end SwinUNETR baselines and a frozen MONAI-ViT with trained adapter and decoder.The Swin and MONAI-ViT comparison inputs were 96×96×96; BraTS used the same multi-channel adaptation strategy.Source: SOTA segmentation comparisons · p. 25

Linear decoder segmentation

  • The linear segmentation decoder maps concatenated patch representations from the final four ViT layers to class logits, then upsamples them with trilinear interpolation.For 96×96×96 inputs with 16×16×16 patches, the example output is a 6×6×6 logit map.Source: Linear decoder segmentation · p. 25

Linear classification probing

  • Classification used frozen pretrained models with linear probes, selecting probe settings by validation performance and testing label efficiency on randomized training subsets.The probe settings included learning rate and the number of final ViT layer outputs to concatenate.Source: Linear classification probing · pp. 25–26

Visualization Methods

Principal component analysis (PCA)

  • PCA visualizations color the next three principal components as RGB after the first component is thresholded to separate foreground and background.PCA was applied to 3DINO-ViT features at each 16×16×16 patch location.Source: Principal component analysis (PCA) · p. 26

Multi-head self-attention (MHSA)

Statistical analysis

  • Segmentation confidence intervals used 1,000 scan-wise bootstrap samples, while classification intervals used five randomized runs; significance tests differed by task.Segmentation used paired Wilcoxon tests on Dice and HD95; classification used unpaired Welch’s t-tests on AUC and F1.Source: Statistical analysis · p. 27

Related papers

Loading 2501.11755v2…