Learning General-Purpose Biomedical Volume Representations using Randomized Synthesis (2411.02372v2)

Abstract: Current volumetric biomedical foundation models struggle to generalize as public 3D datasets are small and do not cover the broad diversity of medical procedures, conditions, anatomical regions, and imaging protocols. We address this by creating a representation learning method that instead anticipates strong domain shifts at training time itself. We first propose a data engine that synthesizes highly variable training samples that would enable generalization to new biomedical contexts. To then train a single 3D network for any voxel-level task, we develop a contrastive learning method that pretrains the network to be stable against nuisance imaging variation simulated by the data engine, a key inductive bias for generalization. This network's features can be used as robust representations of input images for downstream tasks and its weights provide a strong, dataset-agnostic initialization for finetuning on new datasets. As a result, we set new standards across both multimodality registration and few-shot segmentation, a first for any 3D biomedical vision model, all without (pre-)training on any existing dataset of real images.

• The paper introduces a synthetic data engine that generates diverse 3D volumes to overcome limited real-world biomedical datasets.
• The method employs contrastive pretraining on semantically aligned volume pairs to learn robust, appearance-invariant features.
• The approach significantly enhances multi-modality registration and few-shot segmentation by improving Dice overlap and reducing deformation errors.

This paper addresses the significant challenge of poor generalization in volumetric biomedical vision models, which stems from the limited size and high variability (across imaging protocols, anatomical regions, conditions, etc.) of publicly available 3D medical datasets. The authors propose a novel representation learning framework that side-steps the need for large, diverse real-world datasets by training exclusively on synthetically generated data designed to anticipate and encompass wide domain shifts.

The core of their approach is a data engine that synthesizes highly variable 3D training samples. This engine operates in two stages:

1. Label Ensemble Model: It creates synthetic 3D label maps by randomly sampling and composing binary shape templates from a large database (specifically, segmentations from the TotalSegmentator dataset [wasserthal2023totalsegmentator]). These templates are randomly deformed and assigned sequential labels. To simulate realistic anatomy within a field-of-view, a foreground mask (randomly deformed sphere) is applied, and an optional "envelope" layer (dilated/eroded foreground mask) can be added around the foreground. This process results in diverse 3D label maps with randomized spatial configurations and shapes.
2. Appearance Model: Given a synthesized label map with $K$ labels, the engine generates two distinct volumes ($V_1$ and $V_2$) from it. This is done by sampling intensities for each label from two independent $K$-component Gaussian Mixture Models (GMMs) with randomized parameters. Spatial texture is added using Perlin noise [perlin1985image], and further realistic imaging variations are simulated through a comprehensive set of augmentations, including random bias fields, Fourier spikes, Gamma shifts, blurring, Gibbs ringing, resolution degradation, noise, motion, flips, and affine warps. Geometric augmentations are shared between $V_1$ and $V_2$ to preserve semantic layout, while intensity augmentations are independent, creating pairs of volumes that share semantic content but differ in appearance and imaging artifacts.

The data engine generates pairs of volumes ($V_1$, $V_2$) from a shared label map ($L$) that are semantically aligned but appearance-wise distinct. This property is crucial for the subsequent contrastive pretraining strategy. A 3D convolutional UNet [ronneberger2015unet] is trained to learn representations that are stable across these nuisance imaging variations while preserving semantic identity. The pretraining objective is a spatial extension of multi-positive supervised contrastive learning [khosla2020supervised]. For any given voxel (anchor) in $V_1$ or $V_2$ with label $k$, its features should be similar to features of all other voxels with label $k$ in both $V_1$ and $V_2$ (positives), and dissimilar to features of voxels with different labels (negatives). The contrastive loss, applied at multiple decoder layers of the UNet, enforces this inductive bias of appearance-invariant, semantics-aware spatial representations. A small MLP projection head is used during training and discarded afterward.

The resulting pretrained UNet provides general-purpose representations and weights that can be applied to diverse voxel-level tasks in 3D biomedical imaging without requiring pretraining on real datasets.

Practical Applications and Implementation:

The paper demonstrates the utility of the learned representations and weights on two key downstream tasks:

1. Multi-modality Deformable 3D Registration: Traditional image similarity metrics like Mean Squared Error (MSE) or Mutual Information (MI) often struggle with large intensity differences between modalities (e.g., MRI and CT). The proposed method uses the pretrained UNet's output features as a dense, multi-channel input representation for existing, high-performance registration solvers like ANTs [tustison2020antsx] and ConvexAdam [siebert2021fast]. Instead of minimizing image intensity differences, these solvers minimize the differences between the learned feature maps of the fixed and moving images.
   • Implementation: The 16 output channels of the pretrained UNet are treated as multi-channel images. For ANTs, a sum of MSE losses is used across all feature channels. For ConvexAdam, the learned features are concatenated with its default hand-crafted features. Hyperparameters for the registration solver (e.g., regularization weights) are tuned on a small validation set.
   • Results: Using the proposed features significantly improves registration accuracy (measured by Dice overlap of anatomical structures) on challenging intra-subject MRI-CT (L2RAb [hering2021learn2reg]) and inter-subject cardiac MRI-CT (MM-WHS [zhuang2018multivariate]) datasets, outperforming state-of-the-art methods, including those specifically designed for registration but trained on real data. The method achieves high Dice scores while maintaining low deformation folding percentages ($<0.5\%$).
2. Few-shot 3D Multi-label Semantic Segmentation: In scenarios with very limited annotated data (1-3 volumes), training a segmentation model from scratch is challenging. The pretrained UNet weights serve as a strong, dataset-agnostic initialization.
   • Implementation: The pretrained UNet is finetuned on small annotated datasets for specific segmentation tasks. A new convolutional layer with softmax is added to the UNet's output to predict per-voxel labels for $N$ classes. Standard segmentation losses (Dice and cross-entropy) are used. Extensive data augmentation is applied during finetuning to maximize performance from limited data.
   • Results: The finetuned model consistently improves performance over random initialization across a diverse set of 3D segmentation datasets (cardiac MRI, abdominal CT, prostate MRI, abdominal MRI, fetal brain MRI). It achieves competitive or better results than several large, task-specific foundation models pretrained on collections of real datasets, demonstrating strong generalization even to datasets significantly out-of-distribution from common training sources (e.g., fetal MRI). The performance gains over random initialization are more pronounced in the extremely few-shot setting (1-3 volumes) but persist even with more data.

Implementation Considerations & Trade-offs:

• Computational Requirements: Generating synthetic data can be computationally intensive. The authors mitigate this by generating a large dataset (120,000 label volumes, 240,000 volume pairs) offline, supplemented by lighter online augmentations during training. Pretraining the 3D UNet requires significant GPU resources and time (600,000 iterations with batch size 1, $128^3$ volumes).
• Data Engine Design: The specific parameters and components of the data engine (template sources, deformation ranges, GMM parameters, augmentation types and probabilities) are critical. Ablation studies show that using real brain labels or simpler synthetic shapes as templates, or removing augmentations, degrades performance. The choice of contrastive temperature ($\tau$) also impacts the learned representations and downstream task performance, indicating a trade-off between optimal features for registration vs. segmentation.
• Inductive Bias: The chosen inductive bias (stability to appearance variation within semantic labels) is powerful for many tasks but might be suboptimal for niche applications relying on precise relative intensity values (e.g., certain quantitative MRI analyses).
• Downstream Task Adaptation: While the pretrained features are general-purpose, specific tasks like segmentation still require finetuning with labeled data. While this is standard practice, future work could explore how to make segmentation more "promptable" directly from the synthetic pretraining.
• Architecture Choice: The UNet architecture was chosen as a widely used standard, but the framework is designed to be compatible with other volume-to-volume networks.

Overall Practical Impact:

The research offers a practical pathway to develop robust 3D biomedical vision models without relying on the costly and logistically challenging collection of large, diverse real-world annotated datasets. By leveraging structured synthetic data and a tailored contrastive learning objective, the method learns representations that generalize well to unseen domains and tasks. This has direct implications for accelerating the development and deployment of AI models in various clinical applications, particularly in resource-constrained settings or for rare conditions where real data is scarce. The provided code and model weights at \url{https://github.com/neel-dey/anatomix} enable practitioners to readily apply this approach.