Learning General-purpose Biomedical Volume Representations using Randomized Synthesis

## Abstract

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 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.

## 1 Introduction

Figure: Figure 1: Representations produced by our framework, trained only on synthetic data, are approximately stable across imaging modalities, field-of-views, and poses on real unseen volumes from various datasets. For each anatomical region (rows), we show two example volumes with substantial variation (col. 1) and six arbitrarily selected network output channels (cols. 2–7) that illustrate this stability. These features and network weights can be used for several voxel-level tasks.
Refer to caption: extracted/5976921/figures/highlight_collage_v2.png

Biomedical vision models trained on imaging studies with fixed protocols rarely generalize to new populations, medical procedures, and imaging devices. These domain shifts then necessitate clinically infeasible reannotation and retraining cycles, especially for adaptation to new tasks. Further, volumetric annotated biomedical datasets are especially limited in sample size and focused on specific medical procedures, diseases, or scales of anatomy, leading current networks to overfit to a small subset of biomedical tasks. To overcome this data scarcity and heterogeneity, we present a representation learning framework driven by a synthetic data engine. Our approach yields a generalist 3D network that performs well on diverse voxel-level tasks across a range of unseen biomedical contexts.

Current biomedical foundation models are trained by aggregating publicly available datasets to cover multiple domains (Butoi et al., ; Chen et al., ; Liu et al., ; Ma & Wang, ; MH Nguyen et al., ; Pachitariu & Stringer, ; Xie et al., ). However, persistent gaps hinder their widespread adoption. For example, some methods operate only on specific modalities and regions that can be tractably scaled up in sample size, such as chest X-ray (Chen et al., ), and thus cannot learn general representations for most other domains, such as in utero fetal MRI.
Others treat 2D slices of volumetric images as independent data points (Butoi et al., ; Ma & Wang, ), often constructing training sets with high inter-sample correlation yielding models that fail to produce consistent 3D results. Further, existing foundation models almost exclusively focus on segmentation and classification and neglect other key vision tasks such as registration. To our knowledge, no biomedical vision foundation model has been demonstrated for multiple disparate 3D tasks yet.

Contributions. This paper makes advances on two fronts. To gain robustness to large domain shifts in downstream deployment, we first propose a biomedically informed data engine whose samples encompass a wide range of appearances and semantics. This engine uses randomly sampled spatial configurations of biomedical shape templates to synthesize images with arbitrary resolutions, appearances, imaging physics, and crucially, minimal influence from any existing biomedical dataset. Unlike training on samples from GANs or diffusion models, which are limited to reproducing only their original training distribution, our engine synthesizes highly diverse samples useful for arbitrarily new biomedical contexts that we do not have training data for. We then develop a contrastive learning framework that uses paired samples from the data engine to pretrain a network for general voxel-level tasks using an inductive bias of approximate stability to nuisance imaging variation that does not change image semantics, a key property for generalization across datasets (Gruver et al., ).

We experimentally demonstrate that the resulting features and weights enable broad generalization on the key biomedical tasks of 3D registration and segmentation across several diverse datasets. We achieve state-of-the-art unsupervised multimodality image registration by simply using the network’s approximately appearance invariant and pose equivariant representations (Fig. 1) to drive existing registration solvers. The proposed network can also be used as an off-the-shelf dataset-agnostic initialization for finetuning on any voxel-level task. Specifically, we demonstrate strong few-shot segmentation performance in a few-shot setting across highly diverse downstream datasets, thereby removing the need for cumbersome dataset-specific self-supervised pretraining. Our code and trained model weights are available at [https://github.com/neel-dey/anatomix](https://github.com/neel-dey/anatomix).

## 2 Related work

Figure: Figure 2: Data engine. A. We randomly sample binary labels as templates from a large database of segmentations to create 3D shape ensembles of randomly deformed templates. B. Given a synthesized shape ensemble and an appearance model, we synthesize two volumes to pretrain a network with a dense contrastive objective. C. Example synthetic training volumes.
Refer to caption: extracted/5976921/figures/data_generation_v2.png

Generative image models. Learning-based generative models (Brock et al., ; Goodfellow et al., ; Karras et al., ; Luo, ; Rombach et al., ; Song et al., ) trained on internet-scale natural vision sets (Schuhmann et al., ) can now synthesize photorealistic samples for pretraining general networks (Donahue & Simonyan, ; Fan et al., ; Li et al., ; Tian et al., ). However, such generative models trained instead on the few thousand publicly available anatomy- and modality-specific annotated volumes in biomedical datasets (Baid et al., ; LaMontagne et al., ; Qu et al., ; Wasserthal et al., ) generally do not learn representations that can generalize to new biomedical domains. In contrast, the label synthesis component of our data engine draws loose inspiration from the Dead Leaves model (Baradad Jurjo et al., ; Lee et al., ). This hand-crafted generative model considers images to be compositions of randomly deformed shape templates (such as cubes, ellipsoids, etc.) with static intensities to capture the statistics of natural images. We compose biomedical shape templates similarly but develop several further extensions and propose a distinct appearance model, as explained below.

Domain randomization. To learn robustness to domain shifts at deployment, domain randomized generative models (Tobin et al., ) trade realism for diversity when generating training data for downstream models. For example, domain randomized methods for brain segmentation (Billot et al., ; ; Gopinath et al., ; Hoopes et al., ) and registration (Hoffmann et al., ; Iglesias, ) train on synthetic brains simulated from label maps and require large collections of expert brain annotations. Recent work in dataset-agnostic instance segmentation (Chollet et al., ; Dey et al., ) generalizes beyond brains by simulating both annotations and images using a pre-specified shape prior. We build on these concepts to train a task-agnostic network to simulate images with highly variable appearances and physics from compositions of biomedical shape templates.

Invariant imaging features. Given the heterogeneous nature of biomedical imaging protocols, several existing strategies aim to extract features robust to nuisance variation.
When registering images across modalities, aligning modality-invariant hand-crafted local descriptors (Heinrich et al., ; ) and/or edges (Haber & Modersitzki, ) is common. With deep learning-based multimodality registration, this inter-modality invariance can be learned (Dey et al., ; Mok et al., ; Pielawski et al., ), leading to improved performance at the cost of dataset-specific training. Beyond registration, brain-specific invariant features have been learned by exploiting large repositories of annotated brains (Chua & Dalca, ; Liu et al., ; ).
Our work obviates the need for dataset-specific training, anatomical region-specific modeling, and large-scale annotation collection by extracting modality- and appearance-invariant features in an amortized manner.

Volumetric pretraining. Many methods pretrain on unannotated images from a dataset before supervised finetuning on a small labeled subset. These methods often employ self-supervised reconstructive (Chen et al., ; Tang et al., ; Valanarasu et al., ; Zhou et al., ; ) and/or discriminative (Chaitanya et al., ; ; Ren et al., ; You et al., ) losses. However, these pretraining strategies exploit heuristics about their target datasets that often do not broadly generalize (Dong et al., ; Ren et al., ). Our approach instead yields a network that generalizes to arbitrary datasets and does not require bespoke pretraining frameworks for each project. Lastly, recent biomedical foundation models trained on pooled datasets of 2D slices (Butoi et al., ; Ma & Wang, ; MH Nguyen et al., ; Wong et al., ) generally require interaction (via bounding boxes, scribbles, etc.), struggle with 3D consistency, and are restricted to segmentation. We instead directly train for general tasks using synthetic 3D volumes and do not work on interactive tasks.

## 3 Methods

Figure: Figure 3: Representation learning. Given a 3D label map and two synthetic volumes generated from it, we process them with a single UNet (with shared weights). The UNet is pretrained contrastively at each layer of the decoder. For a randomly sampled anchor, features sampled from corresponding labels in both volumes are treated as positives and features from other labels are considered negatives.
Refer to caption: extracted/5976921/figures/loss-methodology.png

This section first details the proposed data engine (Fig. 2), then describes the representation learning strategy (Fig. 3), and concludes with applications towards 3D registration and segmentation.

Data engine: label ensemble model (Fig. 2A). We create synthetic 3D label ensemble volumes by sampling from a repository of biomedical shape templates. As templates, we use the freely available $\sim$45,000 binary volumes from the TotalSegmentator dataset of 104 annotated organs in 1,204 CT volumes (Wasserthal et al., ). For each label ensemble volume, we iteratively populate a 3D volume with a random number of randomly sampled templates that are each then deformed and assigned the label corresponding to the sampling iteration. To simulate anatomy being surrounded by empty space (as is common in radiology), we apply a foreground mask to two-thirds of the synthesized label ensembles by multiplying them with a randomly deformed binary sphere with a random radius and center. Finally, we randomly encase half of the foreground-masked volumes within envelope labels of random widths, to emulate layer-like structures common to some biomedical applications (e.g., fat).

Data engine: appearance model (Fig. 2B). Given a 3D label ensemble $L$ with $K$ labels, we sample the intensities of two volumes $V_{1}$ and $V_{2}$ from two independent $K$-component Gaussian mixture models (GMMs) each with parameters $\{\mu_{k1},\sigma_{k1}^{2}\}_{k=1}^{K}$ and $\{\mu_{i2},\sigma_{i2}^{2}\}_{k=1}^{K}$, respectively, all of which are randomly drawn from uniform distributions. For each spatial index in $L$ with label $k$, we sample the initial intensities in $V_{1}$ and $V_{2}$ from $\mathcal{N}(\mu_{k1},\sigma_{k1}^{2})$ and $\mathcal{N}(\mu_{k2},\sigma_{k2}^{2})$, respectively. We then pointwise multiply them with Perlin noise (Perlin, ) to simulate spatial texture and augment using transformations relevant to biomedical volumes such as random bias fields, Fourier-spikes, Gamma shifts, blurring, Gibbs ringing, resolution degradations, noise, motion, flips, and affine warps. All intensity augmentations are sampled independently but the geometric augmentations are shared.

In summary, we synthesize a 3D label ensemble $L$ and draw volumes $V_{1}$ and $V_{2}$ from it that differ in appearance but share 3D semantic layouts. This is repeated with randomized hyperparameters for each sample to generate a synthetic dataset. Low-level modeling details are provided in App. B.1.

Contrastive pretraining (Fig. 3). As our data engine provides exact label supervision, we develop a representation learning loss that is a spatial extension of multi-positive supervised contrastive learning (Khosla et al., ). To pretrain network $F:\mathbb{R}^{H\times W\times D}\rightarrow\mathbb{R}^{H\times W\times D\times
C}$ where $H,W,$ and $D$ are spatial dimensions and $C$ is the number of output features, we use an inductive bias of voxels within a 3D shape having similar spatial representations, regardless of appearance. We assume that an anchor spatial index $i\in I$ (where $I=\{1,\dots,2HWD\}$, i.e., voxels pooled from $V_{1}$ and $V_{2}$) with features $f_{i}\in\mathbb{R}^{C}$ in label $k$ should have similar representations to all other indices in label $k$ in both $F(V_{1})$ and $F(V_{2})$ and dissimilar representations to indices from other labels. As in (Chen et al., ), we use a non-linear projection $Z:\mathbb{R}^{H\times W\times D\times C}\rightarrow\mathbb{R}^{H\times W\times
D
\times C_{Z}}$ on $F$’s outputs followed by an $L_{2}$-normalization when computing the contrastive loss,

$$ $\mathcal{L}=\sum_{i\in I}\frac{-1}{|P(i)|}\sum_{p\in P(i)}\log\frac{\exp(z_{i} \cdot z_{p}/\tau)}{\sum_{q\in Q(i)}\exp(z_{i}\cdot z_{q}/\tau)},$ (1) $$

where $z\in\mathbb{R}^{C_{Z}}$, $\tau$ is a hyperparameter, $Q(i)=I\backslash\{i\}$ (i.e., all non-anchor spatial indices), and $P(i)$ is the set of all positives for anchor $i$, s.t. $P(i)=\{p\in Q(i):k_{p}=k_{i}\}$, where $k_{x}$ is the label of spatial index $x$. Lastly, we use this loss on multiple decoder layers of $F$ during training to leverage multiscale (self-)supervision as in Dey et al. (); Park et al. (); Ren et al. ().

Figure: Figure 4: Multi-modality 3D registration. The ConvexAdam registration solver (Siebert et al., ) driven by our network features (“Ours”) accurately aligns challenging intra-subject abdominal MRI-CT (top) and inter-subject cardiac MRI-CT (bottom) 3D multi-modality pairs with large deformations.
Refer to caption: extracted/5976921/figures/qualitative-registration-v2.png

Pretraining implementation details. We implement $F$ as a four-level 3D convolutional UNet (Ronneberger et al., ) following the architecture from (Ren et al., ) and construct $Z$ as a 3-layer 128-node-wide MLP. While $F$ can be any volume-to-volume network, we use a U-Net as it is the standard architecture for biomedical imaging tasks and performs well across datasets (Isensee et al., ; Stringer & Pachitariu, ).
$F$ and $Z$ are pretrained jointly for 600,000 iterations with a batch size of one $128^{3}$ label map, each generating two $128^{3}$ volumes. We compute the contrastive loss (with temperature $\tau=0.33$) on 512 randomly sampled indices at each iteration for each decoder layer due to memory limitations. Lastly, $Z$ is only used during pretraining and is discarded for all downstream tasks. All other pretraining details are described in App. B.3.

Downstream tasks: multi-modality registration. Gradient-based deformable registration objectives typically take the form of $\mathcal{L}_{reg}=d(V_{\mathrm{fixed}},V_{\mathrm{moving}}\circ\varphi)+
\lambda\mathrm{Reg}(\varphi)$, where image $V_{\mathrm{moving}}$ is to be aligned to $V_{\mathrm{fixed}}$ by deformation $\varphi$ (subject to regularization $\mathrm{Reg}(\cdot)$) and $d(\cdot)$ is an image dissimilarity score. To align images across imaging modalities, we simply replace the input volumes with our pretrained network’s features as in $\mathcal{L}_{reg}=d(F(V_{\mathrm{fixed}}),F(V_{\mathrm{moving}})\circ\varphi)+
\lambda\mathrm{Reg}(\varphi)$. This approach is compatible with any existing high-performance registration solver, such as the ConvexAdam (Siebert et al., ) and ANTs (Tustison et al., ) frameworks used in our experiments below.

Downstream tasks: few-shot segmentation. For $N$-label segmentation, we use a small set of annotated volumes to finetune the pretrained network $F$, with an additional convolutional layer with softmax activation with $N$ output channels. We optimize the network using an equally weighted sum of the soft Dice and cross-entropy losses (Isensee et al., ; Taghanaki et al., ). To extract strong performance for all baselines in the challenging setting of finetuning on only one or few annotated volumes, we use extensively-tuned augmentation pipelines and finetune all layers of each network for a high number of iterations (37,500) with cosine learning rate decay.

Figure: Figure 5: Multi-modality 3D registration results. (a) Dice boxplots for each method for L2RAb (left group) and MM-WHS (right group), with corresponding medians reported on top of each box and the mean percentages of voxels with folds produced by each method reported at the bottom; (b) Using our features leads to consistent registration improvements at the subject-level.
Refer to caption: extracted/5976921/figures/registration-quantitative-v3.png

## 4 Experiments

The pretrained network’s output features for inter and intra-subject volume pairs with semantically similar content are visualized in Fig. 1. Below, we present experiments that investigate the utility of our learned representations for multi-modality registration, the network weights as a pretrained initialization for few-shot segmentation, and analyze and ablate our framework’s modeling decisions.

### 4.1 Unsupervised multi-modality deformable 3D registration

Data and setup. We use the Learn2Reg AbdomenMRCT (Hering et al., ) (L2RAb) and MM-WHS (Gao et al., ; Zhuang, ; Zhuang et al., ) datasets to benchmark MRI to CT volume registration. L2RAb is an abdominal registration benchmarking dataset, whose publicly available portion provides eight affine-aligned intra-subject MRI and CT pairs of size $192\times 160\times 192$ at $2\times 2\times 2\mathrm{mm}^{3}$ resolution, with labels for four organs. MM-WHS, originally a heart segmentation dataset, contains 20 annotated MRIs and CTs (from distinct subjects) and we affine-align all volumes to a common space of size $160\times 160\times 128$ at $1.142\times 1.142\times 1.283\mathrm{mm}^{3}$ resolution (see App. B.4.4 for further details). The registration experiments are unsupervised and we therefore split L2RAb and MM-WHS into 1/7 and 5/15 validation/testing pairs, respectively, where the validation pair(s) are only used for tuning registration hyperparameters.

Baselines and evaluation. We compare unsupervised and training-free multimodality registration frameworks. Our iterative baselines include the widely-used ANTs library (Avants et al., ; Tustison et al., ) with mutual information loss (Mattes et al., ; Wells III et al., ) (ANTs-MI) and the state-of-the-art (Hering et al., ) multimodality method, ConvexAdam (Siebert et al., ). We further use two dataset-agnostic registration networks, SynthMorph-shapes (Hoffmann et al., ) and uniGradICON (Tian et al., ), with the latter also using optional instance optimization (uniGradICON+IO). For evaluation, we report post-registration volume overlap (Dice) of anatomical structures. We also assess deformation inverse consistency using the percentage of folding voxels, where $\mathrm{det}(J_{\varphi})<0$, for Jacobian $J_{\varphi}$ of the estimated $\varphi$. Folding percentages below 0.5% of all voxels are generally considered negligible and methods producing higher Dice values while staying under this threshold are preferred (Dey et al., ; Qiu et al., ).

Adaptation to use network features. We modify ANTs and ConvexAdam to use our pretrained network’s 16 extracted features. We use ANTs’ multichannel mode with the MSE loss and tune its hyperparameters heuristically on validation pairs. For ConvexAdam, we concatenate our features with its default handcrafted features and perform a grid search for both the original implementation and our variant over four hyperparameters. Lastly, the deep learning baselines assume single-channel input volumes and thus cannot directly use our multichannel network features.

Results. Figs. 4 and 5 present results on held-out testing pairs. ANTS-Ours strongly improves upon the typically-used ANTS-MI, with 26 and 5 points of median Dice improvement on L2RAb and MM-WHS, respectively, while maintaining nearly identical low folding characteristics. Further, driven by our network features, ConvexAdam-Ours outperforms all methods in terms of volume overlap and improves on its base ConvexAdam method by 11 and 6 Dice points, with the same folding ratios.

In contrast, the SynthMorph-shapes and uniGradICON+IO methods perform well on MM-WHS (where all hearts are roughly centered) but cannot handle the larger deformations in L2RAb. They do, however, yield nearly diffeomorphic transformations, producing almost zero folds. Lastly, without iterative optimization, uniGradICON demonstrates limited generalization across large intensity-based domain gaps. We note that all methods produce folding percentages that are under the threshold of 0.5% folding voxels. Additional grid search results on the validation sets are in App. A.2.

### 4.2 Few-shot 3D multi-label semantic segmentation

Figure: Figure 6: Few-shot 3D segmentation qualitative results. All methods (columns 2–7) were fine-tuned on 3, 3, and 1 multi-label annotated volume(s) for each respective dataset (rows 1–3).
Refer to caption: extracted/5976921/figures/qualitative-segmentation.png

**Table 1: Few-shot 3D segmentation Dice means and their bootstrapped std. deviations. Bolding and underlining represent the best and second-best Dice, respectively.**
|  | Params. | MSD-Heart | PROMISE12 | L2RAb-MRI | FeTA | AMOS-CT | WUFetal |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Fine-tuning vols. |  | 1 | 2 | 3 | 3 | 3 | 3 |
| Number of classes |  | 1 | 1 | 4 | 7 | 15 | 4 |
| Rand. Init. UNet | 5.9M | .85(.01) | .80(.02) | .85(.06) | .78(.03) | .56(.01) | .73(.02) |
| Transfer Learning | 5.9M | .87(.02) | .82(.01) | .82(.06) | .78(.03) | .52(.01) | .74(.02) |
| Models Genesis | 19.1M | .84(.04) | .73(.03) | .81(.06) | .79(.03) | .55(.01) | .66(.03) |
| MedicalNet | 17.3M | .86(.02) | .53(.04) | .73(.07) | .74(.04) | .44(.02) | .50(.02) |
| PrimGeoSeg | 67.2M | .87(.01) | .79(.02) | .84(.05) | .79(.03) | .63(.01) | .76(.02) |
| SMIT | 67.2M | .88(.02) | .72(.03) | .84(.06) | .76(.04) | .58(.01) | .73(.02) |
| Disruptive AE | 67.2M | .82(.02) | .64(.03) | .77(.07) | .74(.04) | .50(.02) | .70(.02) |
| Ours | 5.9M | .89(.01) | .85(.01) | .86(.06) | .80(.03) | .61(.01) | .76(.02) |
| Full supervision | 5.9M | .91(.01) | .90(.00) | .89(.05) | .83(.01) | .85(.00) | .88(.01) |

Data and setup. We evaluate few-shot segmentation performance on a diverse collection of datasets:
cardiac MRI from MSD-Heart (Antonelli et al., ), abdominal CT from AMOS (Ji et al., ), prostate MRI from PROMISE12 (Litjens et al., ), abdominal MRI (Akin et al., ; Clark et al., ; Kavur et al., ; Linehan et al., ) from Learn2Reg-Abdomen (Hering et al., ), fetal brain MRI from FeTA (Payette et al., ), and an in-house dataset of whole uterus fetal BOLD MRI (WUFetal). WUFetal provides labels for the placenta, amniotic fluid, and the fetal brain and body. It is included as an out-of-distribution test dataset for the baselines below that are trained on multi-dataset collections of commonly imaged regions. Lastly, we operate in the few-shot regime, where only 1–3 annotated volumes per dataset are used for finetuning. The dataset splits are in App. B.5.

Baselines. We use 3D foundation models specifically pretrained for multi-label segmentation on multiple datasets. These include the masked autoencoding-based Models Genesis (Zhou et al., ), SMIT (Jiang et al., ), and Disruptive AE (Valanarasu et al., ).
Other transfer-learning baselines include MedicalNet (Chen et al., ) and PrimGeoSeg, the latter being pretrained to segment synthetic binary volumes with simplistic shapes. To explicitly test transfer learning with a matched architecture (TransferLearning), we further train a fully supervised UNet with the same architecture as Ours on a large-scale neuroimage segmentation dataset (Hoopes et al., ; LaMontagne et al., ). We also test a randomly initialized UNet (with matched architecture to Ours) trained on few (RandInitUNet) or all (FullSupervision) volumes in the training split. Current 2D interactive binary segmentation foundation models (Butoi et al., ; Ma & Wang, ) were excluded from the experiments as they require user prompts and do not apply to 3D multi-label data. All methods were finetuned with extensive data augmentation, as described in App. B.5.

Results. Table 2 and Fig. 6 present few-shot segmentation results. Our learning framework produces pretrained weights that consistently improve upon random initialization. Crucially, this improvement is achieved using only our dataset-agnostic pretrained weights and without any pretraining on unlabeled real volumes or data from similar domains, as used in previous work. Compared to larger foundation models specifically trained for segmentation, our pretrained general-purpose network achieves the first or second rank consistently and has the best average rank. Importantly, the second-best model changes dataset-to-dataset, indicating that the baseline methods do not generalize consistently to new biomedical contexts. We lastly emphasize that we achieve these gains with fewer parameters, without access to any real images, and while being applicable to other tasks as well, as described below.

Figure: Figure 7: Features for varying pretraining configurations. When our framework is trained with different $\tau$ values (cols. 2–3) or on synthetic data generated from other label sources (cols. 5–6), the network features for real biomedical volumes (col. 1) are degenerate and/or sensitive to nuisance imaging variation. We visualize an arbitrary channel for each model, with more in App. Fig. 11.
Refer to caption: extracted/5976921/figures/main_paper_temperature_and_template_variation.png

**Table 3: Effect of pretraining configurations on downstream tasks via Dice means and their bootstrapped std. deviations. Row 1 corresponds to the configuration used in our previous experiments. Registration experiments are all performed using the ConvexAdam (Siebert et al., ) solver.**
| Pretraining config. | Registration (L2RAb) | Few-shot segmentation Dice ($\uparrow$) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Pretraining loss | $\tau$ | Labels | Dice($\uparrow$) | Folds%($\downarrow$) | WUFetal | MSD-Heart | AMOS-CT |
| Ours | 0.33 | Ours | .74(.10) | 0.22% | .76(.02) | .89(.01) | .61(.01) |
| Ablating pretraining labels |  |  |  |  |  |  |  |
| Ours | 0.33 | smshapes | .68(.10) | 0.20% | .73(.02) | .88(.01) | .60(.01) |
| Ours | 0.33 | Brains | .57(.10) | 0.16% | .74(.02) | .88(.02) | .60(.01) |
| Ours | 0.33 | Ours+Brains | .71(.08) | 0.29% | .76(.02) | .89(.01) | .61(.01) |
| Temperature variation |  |  |  |  |  |  |  |
| Ours | 0.07 | Ours | .58(.10) | 0.17% | .72(.02) | .90(.01) | .60(.01) |
| Ours | 0.20 | Ours | .64(.09) | 0.19% | .78(.02) | .91(.01) | .62(.01) |
| Ablating pretraining loss |  |  |  |  |  |  |  |
| Denoising | - | Ours | .51(.10) | 0.19% | .58(.02) | .83(.03) | .46(.01) |
| Remove labels | - | Ours | .50(.11) | 0.24% | .66(.02) | .86(.02) | .58(.01) |
| Ablating pretraining augmentations |  |  |  |  |  |  |  |
| Ours (Row 1) w/o FG mask | .63(.09) | 0.27% | .73(.02) | .86(.03) | .60(.01) |  |  |
| Ours w/o FG mask, w/o offline augm. | .63(.09) | 0.22% | .74(.02) | .89(.01) | .56(.01) |  |  |
| Ours w/o FG mask, w/o all augm. | .59(.09) | 0.63% | .70(.02) | .84(.02) | .51(.01) |  |  |

### 4.3 Ablations and other analyses

Multitask capabilities. Current 3D biomedical vision foundation models are evaluated primarily using segmentation, raising the question of how well their features generalize to other tasks. To test this, we employ the few-shot segmentation baselines (that provide both pretrained encoders and decoders) as general-purpose feature extractors and use their features to drive a generic registration solver (ANTs), with implementation details provided in App. B.4.3. As reported in Table 2, all current 3D biomedical foundation models are unable to extract features that improve upon the baseline setting used by the solver (MutualInfo). In contrast, our network (Ours) outperforms it by wide margins and is the only method to yield multitask general-purpose features for the highly disparate tasks of multi-modality registration (Fig. 5) and few-shot segmentation (Table 2).

Label generation. To evaluate our label ensemble model, we replace it with other training labels while keeping the appearance model fixed. We test pretraining using: (a) synthetically generated labels that have no biomedical priors (Butoi et al., ; Hoffmann et al., ) (smshapes),
(b) 1,573 label maps of real brain MRI using the FreeSurfer protocol (Fischl, ) (Brains) that represent real anatomical structures with dense per-voxel annotations,
and (c) a combination of Brains and our model to mix real and synthetic sources of label maps. Details regarding these models are provided in App. B.6. Table 3 rows 1–4 show that both registration and segmentation performance decline with other choices of pretraining labels. Further, combining our labels with Brains does not affect segmentation but worsens abdominal registration.
Lastly, pretraining on these alternative label models leads to unstable network features on real data (Fig. 7), likely explaining the performance degradation.

Temperature ($\tau$). The temperature hyperparameter $\tau$ in the contrastive loss (eq. 1) controls the penalty weight on negative pairs (Wang & Liu, ). In natural vision, smaller $\tau$ values (such as $\tau=0.07$ (Chen et al., )) are used to upweigh difficult negative pairs.
However, as we train our network to learn similar representations within each label despite highly disparate appearances, the negative pairs are all difficult and require a relaxed $\tau$ of $0.33$ for stable training. We find that using lower $\tau$ leads to degenerate aliased features (Fig. 7) and worse
registration results (Table 3). Interestingly, segmentation benefits slightly from an intermediate setting of $\tau=0.20$, indicating a tradeoff between optimal representations for downstream registration versus segmentation.

Pretraining objectives. We now retain our data engine, but pretrain using other frameworks. We compare against matched architectures that are (a) pretrained to denoise the augmentations used in our data engine (Denoising) as in Iglesias et al. () and (b) pretrained using self-supervision on intensities alone without using label information (RemoveLabels) as in Chua & Dalca (); Ren et al. (). In Table 3, rows 1, 7, and 8, we find that our label-supervised multi-positive contrastive strategy yields the highest results for both registration and segmentation.

Figure: Figure 8: Fine-tuning performance as a function of annotation budget.
Refer to caption: extracted/5976921/figures/amosct-percentage-of-data.png

Data engine augmentations.
We now ablate the augmentations used in the proposed data engine used during pretraining by cumulatively removing the foreground masking, the augmentations used during offline image synthesis (App. Fig. 12), and all augmentations, such that the training images are simply the Perlin-corrupted Gaussian mixture model outputs. In Table 3 rows 9–11, we observe a substantial drop in both registration and segmentation performance without the proposed augmentations.

Finetuning with more annotations. Finally, while our segmentation experiments focus on the few-shot setting, our learned initialization also benefits scenarios with more supervision. Fig. 8 quantifies how finetuning the proposed network with more annotated volumes leads to improved segmentation on AMOS-CT (Ji et al., ) relative to random initialization across settings, albeit with smaller improvements given more annotations.

## 5 Discussion

Limitations and future work. Our approach does have limitations. We pretrained our network to be stable against intensity variations (among other variables) and demonstrated its utility for registration and segmentation. However, a small set of biomedical tasks rely on relative intensity values (Nakamura et al., ; Thomalla et al., ), making intensity invariance a suboptimal inductive bias for them. Further, while we operate on general volumetric tasks, certain biomedical inverse problems like MRI reconstruction take sensor-domain non-Cartesian (Schlemper et al., ) measurements as multichannel (Singh et al., ) inputs, requiring domain-specific architectural changes.
Lastly, our segmentation experiments finetune our pretrained network on specific datasets, potentially introducing complexity for some clinical users. However, future extensions could directly use our proposed data engine to train promptable 3D segmentation models that require no such finetuning.

Conclusions. When combined with the right inductive biases, synthetic data models informed by biomedical templates enable the training of powerful general-purpose networks. This is important for 3D radiology, where existing annotated datasets are limited in sample size, often acquiring only dozens to at most a few thousand volumes. This leads to inflexible models that deteriorate under domain shifts.
Trained only on synthetic volumes with our proposed framework, the resulting network provides substantial benefits on a variety of tasks. For example, its representations yield substantial improvements over the state-of-the-art in training-free multimodality deformable registration, a key area in biomedical vision. Further, the network can also serve as a downstream dataset-agnostic initialization for few-shot segmentation tasks and lead to improvements across multiple datasets.

## Acknowledgements

We gratefully acknowledge support from NIH NICHD R01HD100009, NIH NIBIB 1R01EB032708, NIH NIBIB R01EB033773, MIT-IBM Watson AI Lab, MIT Jameel Clinic, and the MIT CSAIL-Wistron Program.

## Appendix A Appendix: Additional results

### A.1 Additional synthetic data visualizations

Figure: Figure 9: Randomly selected synthetic volume (center slice) visualizations sampled from our proposed data engine (top). In the Brains and smshapes rows, we present samples corresponding to our ablations where we replace our proposed 3D label ensemble generator with real 3D brain labels (middle) or synthetically generated labels with no biomedical priors (bottom), respectively, but keep the appearance model unchanged.
Refer to caption: extracted/5976921/figures/appendix-synth-viz.png

### A.2 Registration regularization weight grid searches

Deformable registration optimization faces a trade-off between accuracy and warp field regularity, which is often tackled by using various regularizers and hyperparameters (Hoopes et al., ).
For a fair comparison with ConvexAdam, we separately tune both the original framework and our extension (ConvexAdam-Ours) via a grid search over four hyperparameters on the validation splits of both datasets (i.e., L2R-Abdomen MRCT and MM-WHS). These hyperparameters include: Adam (Kingma & Ba, ) optimization grid spacing: $\{1,2\}$, the warp smoothness penalty $\lambda$: $\{0.25,0.5,0.75,\dots,2.5\}$, grid spacing: $\{2,3,\dots,6\}$ and disp_hw: $\{1,2,\dots,5\}$. All hyperparameters are tuned such that post-registration volume overlap (Dice) is maximized while maintaining deformation folds below $0.5\%$.

In Table 4, we summarize the Dice and folding statistics over the validation set grid searches. The reported means are averaged across subjects and hyperparameter configurations while standard deviations indicate inter-configuration spread. Using our network features (ConvexAdam-Ours) leads to substantially better performance and lower sensitivity to hyperparameter settings. This is corroborated in Fig. 10, where for different settings of $\lambda$ along the x-axis, we visualize Dice and folding voxel percentages for each hyperparameter configuration. We again find better performance at lower folding percentages while maintaining a lower sensitivity to hyperparameter settings.

**Table 4: Registration validation set grid search summary statistics (mean $\pm$ std.). Here, the means are computed over all subjects and all hyperparameter configurations and the standard deviations correspond to the spread over all hyperparameter configurations.**
|  | L2R-Abdomen MRCT | MM-WHS |  |  |
| --- | --- | --- | --- | --- |
| Method | Dice ($\uparrow$) | Folds% ($\downarrow$) | Dice ($\uparrow$) | Folds% ($\downarrow$) |
| ConvexAdam-Ours | $\mathbf{0.863\pm 0.016}$ | $\mathbf{0.693\pm 1.423}$ | $\mathbf{0.661\pm 0.030}$ | $\mathbf{0.496\pm 1.249}$ |
| ConvexAdam | $0.806\pm 0.038$ | $2.269\pm 3.292$ | $0.652\pm 0.023$ | $1.715\pm 3.572$ |

Figure: Figure 10: Registration validation set grid search sweep statistics. Here, each point contained in a boxplot is the average Dice for a particular hyperparameter configuration, given a fixed $\lambda$.
Refer to caption: extracted/5976921/figures/appendix-grid-search-convexadam.png

### A.3 Additional representation visualizations

Figure: Figure 11: Companion figure to Fig. 7: Feature visualizations for varying pretraining configurations. Here, we visualize channel-wise output representations from our pretrained network for the two volumes at the top left for the first six channels. First, varying the temperature hyperparameter in the contrastive loss can lead to substantial aliasing (grouped rows 1–3). Then, changing our label ensemble synthesis model for other label sources (grouped rows 3–5) shows a loss in stability to nuisance variation. Our proposed model in grouped row 3 achieves both interpretable and stable representations on highly challenging volume pairs.
Refer to caption: extracted/5976921/figures/appendix_ablation_features_variation.png

### A.4 Network parameter count effects

Our experiments in Section 4.2 compare our pretrained U-Net with publicly released foundation models that have significantly higher parameter counts. To investigate the impact of network size, we use the same U-Net architecture as in our proposed model and retrain the second-highest average ranking method, PrimGeoSeg (Tadokoro et al., ), reducing its parameter count from 67.2M to 5.9M but matching all other training details to their code repository(^1^11[https://github.com/SUPER-TADORY/PrimGeoSeg](https://github.com/SUPER-TADORY/PrimGeoSeg)). We then use it as a feature extractor for multi-modality registration with the ANTs solver (Tustison et al., ) and also finetune it for few-shot segmentation, as in our experiments in the main text.

As shown in Table 5, matching the parameters of PrimGeoSeg results in performance drops across 3 out of 5 few-shot segmentation datasets, as well as in both registration datasets. Notably, while PrimGeoSeg had originally outperformed our method on the AMOS-CT few-shot segmentation experiment in the main paper, its performance drops below ours by 5 mean Dice points when matched in parameter count, suggesting that the performance gap is a function of network size on that dataset.

**Table 5: Experiments comparing the multitask capabilities of the second-best segmentation model in Table 2 (PrimGeoSeg) when matched for parameters and network architecture, as measured by subject-averaged Dice coefficients and their corresponding bootstrapped std. deviations.**
|  |  |  | Few-shot Segmentation | Registration w/ ANTs |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | Architecture | Params. | MSD-Heart | L2RAb-MRI | FeTA | AMOS-CT | WUFetal | L2RAb | MMWHS |
| PrimGeoSeg | SwinUNETR | 67.2M | .87(.01) | .84(.05) | .79(.03) | .63(.01) | .76(.02) | .46(.09) | .51(.02) |
| PrimGeoSeg | UNet | 5.9M | .82(.02) | .84(.06) | .79(.03) | .56(.01) | .71(.02) | .36(.06) | .48(.02) |
| Ours | UNet | 5.9M | .89(.01) | .86(.05) | .80(.03) | .61(.01) | .76(.02) | .70(.09) | .63(.02) |

### A.5 Negative results

While our proposed model consistently achieves state-of-the-art performance across several segmentation and registration benchmarks, it shows interesting trends on the preprocessed neurite-oasis (Hoopes et al., ) T1w MRI neuroimage segmentation dataset, which is derived from the larger OASIS (LaMontagne et al., ) dataset. For this few-shot segmentation experiment, we train on $160^{3}$ crops, with a batch size of 3, and finetune on one annotated subject to segment the 35 classes provided by the dataset. In Table 6, we find that our proposed pretrained network does not improve over random initialization for this particular dataset.

**Table 6: Few-shot 3D segmentation results on the neurite-oasis dataset (Hoopes et al., ) reported as the mean Dice coefficient and bootstrapped std. deviation.**
| RandInit | ModelsGenesis | MedicalNet | PrimGeoSeg | SMIT | DisruptiveAE | Ours |
| --- | --- | --- | --- | --- | --- | --- |
| .82(.01) | .84(.01) | .75(.01) | .84(.01) | .83(.01) | .80(.01) | .82(.01) |

This result is consistent with the literature on the limited benefits of representation learning for few-shot adult neuroimage segmentation on OASIS in terms of Dice coefficient gains, as also reported in Ren et al. (). These trends may stem from the relatively high resolution, tissue contrast, and low inter-subject variability in OASIS. Further supporting this, prior work (Lee et al., ) has also shown that training adult neuroimage segmentation models from scratch on very small datasets can yield competitive results.
Of the 8 remaining segmentation and registration tasks, our model either achieves the best (7 of 8 tasks) or second-best (1 of 8 tasks) performance, demonstrating its overall robustness.

## Appendix B Appendix: Additional implementation details

### B.1 Data engine details

Our data engine in Fig. 2 A & B has several components. Here, we describe low-level implementation details. We note that we make extensive use of the MONAI (Cardoso et al., ), TorchIO (Pérez-García et al., ), and scikit-image (van der Walt et al., ) libraries for both label and volume synthesis.

#### B.1.1 Label synthesis

The pseudocode in Algorithm 1 summarizes the synthesis process for a single label volume depicted in Fig.2A. This process assumes the availability of a set $Y$ of binary segmentation labels for different organs, which serve as templates. Specifically, these segmentations are taken from version 1 of the publicly available TotalSegmentator dataset (Wasserthal et al., ) (CC BY 4.0 license). To incorporate binary labels with multiple connected components, we merge individual rib labels into a single binary label for all ribs and also similarly pool the individual vertebral labels.

In Algorithm 1, $p_{\mathrm{fg}}$ and $p_{\mathrm{envelope}}$ refer to the probability of foreground masking and creating an envelope around the foreground, respectively, $w$ is the kernel width of the ball kernel used for morphological dilation and erosion, $\circ$ refers to a spatial deformation operator, and $\ast$ denotes the element-wise multiplication. The Perlin deformations $P_{\sigma}$ used are taken from Hoffmann et al. ().

Figure: Algorithm 1 3D synthetic label map $L$ generation

#### B.1.2 Volume synthesis

Given a label map $L$, we use it to conditionally sample two volumes/contrastive views $V_{1}$ and $V_{2}$ using an appearance model that is summarized in Fig. 12. Specifically, the two intensity volumes are generated by sampling from two independent Gaussian mixture models conditioned on the label map $L$. These preliminary 3D images are then independently transformed by a biomedical augmentation pipeline to form a contrastive pair of volumes. The term Zero background in Fig. 12 refers to setting intensities in the volumes spatially coinciding with the background label to 0.

As the combined pipeline of our label and volume synthesizer is computationally intensive, generating data on the fly could potentially bottleneck training. To mitigate this, we sample 120,000 3D label volumes and their corresponding 240,000 contrastive volume pairs offline using the proposed model. During training, we further apply additional online augmentations using a smaller pipeline. The offline and online augmentations are indicated in Fig. 12 by the blue and red boxes, respectively.

Figure: Figure 12: Post-Gaussian mixture model volume augmentation pipeline used to generate synthetic data for pretraining. Blue and red boxes refer to offline and online augmentations, respectively. All augmentations are applied to $128\times 128\times 128$ volumes. $p$ refers to the pre-defined probability of applying an individual augmentation and all other hyperparameters mentioned are consistent with MONAI (Cardoso et al., ) conventions.
Refer to caption: extracted/5976921/figures/appendix_appearance_model.png

**Table 7: U-Net architectural details. We use the architecture from (Billot et al., ; Ren et al., ). Conv-BN-ReLU refers to a sequence of 3D convolution with $3\times 3\times 3$ kernels, batch normalization, and pointwise ReLU activations. $n_{c}$ is the channel width multiplier and $n$ is the number of output channels. In our experiments, both $n_{c}$ and $n$ are set to 16.**
| Layer index | Layer contents |
| --- | --- |
| 0 | Conv-BN-ReLU($n_{c}$) |
| 1 | Conv-BN-ReLU($n_{c}$) |
| 2 | Conv-BN-ReLU($n_{c}$) |
| 3 | MaxPool(2), Conv-BN-ReLU($2n_{c}$) |
| 4 | Conv-BN-ReLU($2n_{c}$) |
| 5 | MaxPool(2), Conv-BN-ReLU($4n_{c}$) |
| 6 | Conv-BN-ReLU($4n_{c}$) |
| 7 | MaxPool(2), Conv-BN-ReLU($8n_{c}$) |
| 8 | Conv-BN-ReLU($8n_{c}$) |
| 9 | MaxPool(2), Conv-BN-ReLU($16n_{c}$) |
| 10 | Conv-BN-ReLU($16n_{c}$) |
| 11 | Upsample $2\times$, Concatenate with layer 8 |
| 12 | Conv-BN-ReLU($16n_{c}$) |
| 13 | Conv-BN-ReLU($16n_{c}$) |
| 14 | Upsample $2\times$, Concatenate with layer 6 |
| 15 | Conv-BN-ReLU($4n_{c}$) |
| 16 | Conv-BN-ReLU($4n_{c}$) |
| 17 | Upsample $2\times$, Concatenate with layer 4 |
| 18 | Conv-BN-ReLU($2n_{c}$) |
| 19 | Conv-BN-ReLU($2n_{c}$) |
| 20 | Upsample $2\times$, Concatenate with layer 2 |
| 21 | Conv-BN-ReLU($n_{c}$) |
| 22 | Conv-BN-ReLU($n_{c}$) |
| 23 | Conv-BN-ReLU($n$) |

### B.2 Network architectures

We employ a widely used (Billot et al., ; Ren et al., ) UNet architecture for pretraining, fine-tuning, and for use as a feature extractor in our baseline comparisons and ablations. The architecture is described in Table 7. Each layer that contributes to the contrastive loss also includes an individual MLP network that projects sampled spatial indices onto an embedding space. Specifically, the architecture of each MLP consists of two FC(128)-BN-ReLU blocks (where FC(w) is a fully connected layer of width w), followed by an FC(128) layer and an $L_{2}$-normalization sequence at the final layer.

### B.3 Pretraining implementation

We compute the contrastive loss in Eq. 1 on pre-activation convolutional features extracted from layers 7, 9, 12, 15, 18, and 23 in Table 7. Model selection is performed by tracking the validation contrastive loss to pick the best-performing checkpoint.
We train using Adam (Kingma & Ba, ) with a starting learning rate of $2\times 10^{-4}$ and a step decay towards 0 every 120,000 iterations.

### B.4 Registration experiments implementation

Below, we first provide details regarding how the baselines were implemented in our registration experiments and then describe how our network features were integrated with existing solvers. For volumes of arbitrary grid sizes, we use sliding window inference with a window size of $128^{3}$ and an overlap ratio of 0.8. Further, we use region-of-interest masks for fixed and moving volumes whenever a registration method can use them.

#### B.4.1 Baseline implementation

SynthMorph-shapes. SynthMorph-shapes is a domain-randomized diffeomorphic registration UNet trained on synthetic volume pairs generated from label maps. It is optimized using a Dice loss subject to diffusion regularization. We download its pretrained weights(^2^22[https://surfer.nmr.mgh.harvard.edu/ftp/data/voxelmorph/synthmorph/shapes-dice-vel-3-res-8-16-32-256f.h5](https://surfer.nmr.mgh.harvard.edu/ftp/data/voxelmorph/synthmorph/shapes-dice-vel-3-res-8-16-32-256f.h5)) from the VoxelMorph library and use their Tensorflow-based registration framework. Lastly, as their network is fully convolutional, we use the input volumes at their native resolution without resizing.

uniGradICON/uniGradICON+IO. uniGradICON is an approximately diffeomorphic registration foundation model trained on a variety of datasets. We use the binaries available on their repository(^3^33[https://github.com/uncbiag/uniGradICON](https://github.com/uncbiag/uniGradICON)) for our experiments. The off-the-shelf model does not have any hyperparameters at test-time and we use the default hyperparameters for the iterative variant (uniGradICON+IO).

ConvexAdam. ConvexAdam is a high-performance multimodality registration solver and we use the b2671f8 commit of the repository(^4^44[https://github.com/multimodallearning/convexAdam/tree/b2671f86902390dec8dde702d0b583b451d84e98](https://github.com/multimodallearning/convexAdam/tree/b2671f86902390dec8dde702d0b583b451d84e98)). Here we use the masked variant with default number of instance optimization iterations (80) and default hyperparameters of the MIND-SSC loss so as to use the same implementation of MIND-SSC across experiments. All the remaining parameters are fine-tuned on the validation data.

ANTs-MI. The ANTs-MI baseline was run with the following command using the ANTs library:

where fixed_volume.nii.gz and moving_volume.nii.gz are the input volumes to register and fixed_volume_mask.nii.gz and moving_volume_mask.nii.gz are binary masks indicating non-zero / non-background regions.

These settings correspond to using a three-level registration pyramid with the SyN algorithm and the mutual information loss for 200, 200, and 100 iterations at each level with corresponding level-specific smoothing kernels.

#### B.4.2 Modifications to existing registration methods to use our network features

As our network produces 16 output channels, we modify existing registration solvers as below.

ANTs. To solve for a warp between a fixed and moving volume pair, we define 16 different MSE-based loss functions with ANTs. Specifically, each loss estimates the dissimilarity between corresponding channel volumes produced by our network for the fixed and moving inputs. We also downscale each individual loss by a tenth to trade off multiple data fidelity terms and regularization. The remaining modeling decisions and hyperparameters are identical to the baseline and use the following command:

ConvexAdam. ConvexAdam already operates on multichannel inputs by using handcrafted MIND-SSC features. We therefore concatenate our network features with their original features and additionally multiply the network features by 0.1 for stable optimization. We perform a grid search over the same hyperparameters as the baseline ConvexAdam and set the remaining modeling decisions to be consistent with it.

#### B.4.3 Using existing 3D biomedical segmentation foundation models for registration

In Section 4.3 and Table 2 of the main text, we demonstrate that current 3D biomedical segmentation models do not produce features that are directly usable by registration solvers such as ANTs.
To elaborate, MedicalNet was excluded from analysis as it only provides a pretrained encoder without a decoder.
ModelsGenesis produces one output feature map for an image that we use as is.
The SwinUNETR-architecture based baselines (PrimGeoSeg, SMIT, and DAE) produce 48 output channels that would necessitate prohibitively high compute costs at full-resolution. However, based on validation set performance, we find no significant difference between using ANTs with all 48 features or a subset of 16 (every fourth channel), and thus use the latter for all results reported for the test set.
The ANTs hyperparameters for all experiments are the same as in App. B.4.2, varying only in the number of input channels corresponding to the output features produced by each method.

#### B.4.4 MM-WHS data preparation

Prealignment. The public proportion of the MM-WHS dataset consists of 20 unpaired and annotated 3D CTs and MRIs of the heart, all from different subjects. The CTs are high-resolution CT angiograms with tight fields of view around the heart, whereas the MRIs often include the subject’s trunk. As our deformable registration baselines all assume affine pre-alignment, we align all the volumes to a common space. For accurate groupwise registration to this space, we formulate this as an affine atlas construction and registration problem (Avants et al., ).

All CT volumes are first clipped to [-450, 450] HU. We then arbitrarily select the first CT volume within the dataset (by subject ID) as an initial reference. We first resample it to a grid size of 160 × 160 × 128 at $1.142\times 1.142\times 1.283mm^{3}$ resolution to define an initial coordinate system. This resampled volume is then used as an initial target for affine atlas construction with the remaining CT volumes. We use the following ANTs command(^5^55[https://github.com/ANTsX/ANTs/blob/master/Scripts/antsMultivariateTemplateConstruction2.sh](https://github.com/ANTsX/ANTs/blob/master/Scripts/antsMultivariateTemplateConstruction2.sh)) for groupwise affine alignment:

Once an affine CT atlas is estimated, we then similarly register all MRI volumes to this CT atlas. In particular, due to the difficulty of intensity-based affine registration between the MRI and CT collections due to domain and FOV shifts, we estimate these affine transformations on the segmentations provided by the dataset and not the volumes themselves. Once all subject-to-atlas affine transformations are estimated on the segmentation volumes, they are used to warp all of the intensity volumes into the desired common space. All deformable registration experiments in our paper use only the intensity volumes and use the segmentations only for evaluation.

Additional labels. MM-WHS provides manual annotations for seven structures including heart chambers and portions of arteries for its original use in segmentation benchmarking. However, there are several anatomical structures in these volumes beyond the original labels such as the spine. Therefore, to better repurpose this data for holistic and non-local multi-modality registration evaluation, we annotate additional labels for the descending aorta and the spine. We use TotalSegmentator (Wasserthal et al., ) to segment these labels on the CT volumes and manually verify the results. For the MRI volumes, these new structures are annotated by a domain expert.

### B.5 Segmentation experiments implementation

All image characteristics for each dataset are summarized in Table 8. The datasets and modeling decisions were chosen to maximize diversity between the various segmentation settings.

For the publicly available datasets, we (re)split whatever annotated data is publicly available from the respective datasets. In particular, we resplit MSD-Heart, FeTA, and L2RAb-MRI’s publicly available training sets to obtain training, validation, and held-out testing data. For AMOS-CT, we resplit the 200 volumes in the training set to obtain new training and validation splits and use their original and public 100 validation volumes as held-out testing data. For PROMISE12, we use the public splits, considering the test and livechallengetest splits to be its testing and validation splits, respectively.

PROMISE12 and AMOS-CT are highly anisotropic in resolution and are thus resampled to $0.625\times 0.625\times 0.625\ mm^{3}$ and $1.5\times 1.5\times 2.0\ mm^{3}$ resolution, respectively, with the latter resolution for AMOS-CT chosen to match the CT finetuning setting of the SwinUNETR baselines (Valanarasu et al., ). The CT intensities of AMOS-CT were clipped to [-450, 450] HU and the spatial grid extents of MSD-Heart and (post-resampling) PROMISE12 were padded to have a minimum of 96 slices.

Our UNet baselines use the crop sizes listed in Table 8 and our SwinUNETR-based pretrained baselines use the crop sizes used for pretraining and finetuning in their respective original papers for both consistency and due to their pretrained transformer backbones. All augmentations are performed online and the MRI and CT dataset experiments in Table 8 use the augmentations listed in Table 9 top and bottom, respectively. When finetuning, we train for 37,500 iterations with a batch size of four 3D crops using Adam with a starting step size of $2\times 10^{-4}$ cosine decayed to 0. Finally, we exclude the background label when reporting Dice statistics.

**Table 8: Segmentation experimental dataset statistics. All MRI modalities and sequences significantly differ from dataset to dataset.**
|  | WUFetal | PROMISE12 | MSD-Heart | L2RAb-MRI | AMOS-CT | FeTA |
| --- | --- | --- | --- | --- | --- | --- |
| Grid size | (112, 112, 80) | (320, 320, 24) | (320, 320, 115) | (192, 160, 192) | (512, 512, 115) | (256, 256, 256) |
| Original res. (mm^3) | (3.0, 3.0, 3.0) | (0.625, 0.625, 3.6) | (1.25, 1.25, 1.37) | (2.0, 2.0, 2.0) | (0.68, 0.68, 5.0) | (0.5, 0.5, 0.5) |
| Training res. (mm^3) | (3.0, 3.0, 3.0) | (0.625, 0.625, 0.625) | (1.25, 1.25, 1.37) | (2.0, 2.0, 2.0) | (1.5, 1.5, 2.0) | (0.5, 0.5, 0.5) |
| Training crop size | 80^3 | 96^3 | 96^3 | 128^3 | 96^3 | 128^3 |
| Training batch size | 4 | 4 | 4 | 4 | 4 | 4 |
| Num. of labels | 4 | 1 | 1 | 4 | 15 | 7 |
| Modality | MRI | MRI | MRI | MRI | CT | MRI |
| Finetuning vols. | 3 | 2 | 1 | 3 | 1 | 3 |
| Full supervision vols. | 60 | 50 | 6 | 24 | 180 | 40 |
| Validation vols. | 15 | 20 | 10 | 12 | 20 | 20 |
| Testing vols. | 24 | 30 | 4 | 12 | 100 | 20 |

**Table 9: Augmentations for segmentation fine-tuning experiments for MRI datasets (top table) and CT datasets (bottom table). Hyperparameters correspond to MONAI conventions (Cardoso et al., ).**
| Prob. | MRI augmentation | Hyperparameters |
| --- | --- | --- |
| 1.0 | Spatial crop | Crop size |
| 0.33 | Gaussian Noise | $\mu=0.0$, $\sigma=0.1$ |
| 0.33 | Bias field | coefficients $\sim\mathcal{U}[0,0.075]$ |
| 0.33 | Gibbs ringing | $\alpha\sim\mathcal{U}[0,0.33]$ |
| 0.33 | Gamma transform | $\gamma\sim\mathcal{U}[0,4.5]$ |
| 0.33 | Gaussian blur | per-axis $\sigma\sim\mathcal{U}[0,0.1]$ |
| 0.33 | Gaussian sharpen | $\alpha\sim\mathcal{U}[1,30]$, $\sigma_{1}\sim\mathcal{U}[0,3.0]$, $\sigma_{2}\sim\mathcal{U}[0,1.0]$ |
| 1.0 | Affine warp | rotation$\sim\mathcal{U}[-\pi/4,\pi/4]$, scale$\sim\mathcal{U}[0.8,1.2]$, shear$\sim\mathcal{U}[-0.2,0.2]$ |
|  |  | (all per axis) |
| Prob. | CT augmentation | Hyperparameters |
| 1.0 | Spatial foreground crop | Crop size, foreground label weight of 0.5 |
| 0.33 | Gaussian Noise | $\mu=0.0$, $\sigma=0.1$ |
| 0.33 | Gamma transform | $\gamma\sim\mathcal{U}[0,4.5]$ |
| 0.33 | Gaussian blur | per-axis $\sigma\sim\mathcal{U}[0,0.1]$ |
| 0.33 | Gaussian sharpen | $\alpha\sim\mathcal{U}[1,30]$, $\sigma_{1}\sim\mathcal{U}[0,3.0]$, $\sigma_{2}\sim\mathcal{U}[0,1.0]$ |
| 1.0 | Affine warp | rotation$\sim\mathcal{U}[-\pi/4,\pi/4]$, scale$\sim\mathcal{U}[0.8,1.2]$, shear$\sim\mathcal{U}[-0.2,0.2]$ |
|  |  | translation$\sim\mathcal{U}[-32,32]$ (all per axis) |

### B.6 Alternative label synthesis models

As in Sec. 4.3/label generation, we study the effect of our label ensemble synthesis model by replacing it with other label models used in biomedical image analysis. We detail our implementation of these alternatives below and visualize them in App. A.1.

smshapes. For smshapes, we follow the implementation of Hoffmann et al. () for label generation with two key exceptions. First, for the label synthesis model, they fix the number of labels to synthesize to 26. For fair comparison with our framework which varies the number of labels in each volume, we modify smshapes to sample the same variable number of shapes. Second, for the appearance model, we match the hyperparameters of their GMM implementation to ours and also use multiplicative Perlin noise such that the appearance models are now matched and only the source of labels varies.

Brains. As opposed to synthesizing labels, the Brains experiment uses real brain labels in a manner similar to (Billot et al., ). This is done to study the effect of synthesized label ensembles with randomized positions versus real biomedical anatomical configurations. We use 492, 500, and 581 T1-weighted brain scans from ADNI, HCP, and IXI, respectively, and segment them with SynthSeg (Billot et al., ) to obtain training label maps. Using these labels, we sample 120,000 label volumes with 240,000 contrastive views as in our proposed model and match all other hyperparameters.

### B.7 Alternative pretraining losses

Denoising pretraining. For denoising pretraining, we maintain our data engine but pretrain the UNet to instead invert the intensity augmentations applied to the output of the initial Gaussian mixture model per sample, which is inspired by Iglesias et al. ().
The UNet has a matched architecture to ours with an additional single convolutional layer mapping to a single-channel output volume. We train using the $L_{1}$ loss for denoising, match the optimization hyperparameters to our method, and select the checkpoint with the best validation $L_{1}$ loss.

Removing label supervision. As our data engine provides exact label supervision for each voxel, we use multi-positive label-supervised contrastive learning. However, several large-scale models are pretrained with self-supervised objectives not using any label supervision. To benchmark against these approaches, we use the self-supervised positive pair-only non-contrastive framework of Ren et al. () with its losses applied to the same network layers as ours. Its variance and covariance loss weights are set to 0.01, the orthogonality weight is set to 100, and we halve the initial learning rate for stable training on our data as opposed to real brains used in their work.

### B.8 WUFetal dataset details

The in-house Whole Uterus Fetal (WUFetal) BOLD MRI dataset consists of 99 whole uterus volumes covering various pathologies, gestational ages, imaging artifacts, and the presence of twins. Due to this variability, this is a highly challenging dataset for few-shot segmentation. These scans were acquired on a 3T Skyra Siemens scanner using multi-slice gradient echo EPI sequences at 3mm isotropic resolution (TR = [5-8] ms, TE = [32-38] ms, $\alpha=\pi/2$). All analyses were performed retrospectively on anonymized data and are IRB-approved.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see weaknesses section.
soundness: 3
strengths: -	Innovative approach with a highly powerful general synthetic data engine, I think this work adds a lot to the discussion on “foundational” models in medicine. 

-	The method requires a comparatively small number of trainable hyperparameters compared to existing models while achieving a higher accuracy. This in my opinion is an important contribution towards sustainable models. 

-	The authors provide the source code which is a major plus for reproducibility and discuss reproducibility aspects. 

-	The authors evaluate the effect of different pretraining strategies in reasonable ablation studies. 

-	The authors honestly show negative results in the appendix.
summary: The work proposes a new method for general (few-shot) medical image segmentation and registration. The method is based on generating realistic and varying synthetic data and training a generalist 3D segmentation and registration models on this data. The goal is to develop a method for generalizing to different imaging devices, medical procedures and conditions as well as different populations. 

The data synthesis engine is capable to generate highly diverse samples and uses the totalsegmentattor dataset for shape templates. The “foundational” model is then pretrained using contrastive pretraining and finetuned in multi-modality registration and few-shot segmentation. 

In experiments the authors demonstrate excellent performance in both downstream tasks.
weaknesses: -	Considered datasets. If I am not missing any details the authors mostly evaluate their method on CT and MRI images. They even consider an MRI image as out of distribution. I think experimentation and ablation on more difficult and diverse datasets, such as 3D microscopy or Ultrasound, even if the results are not all positive would add to the discussion and validate the claim of a “general volumetric biomedical foundation model”

-	Framing of contribution. I would prefer if the authors tone down their wording about the model and clearly point out that it is a model for radiology or CT and MRI dataset and not a “general volumetric biomedical foundation model”.

-	Hyperparameter selection: What range of hyperparameters was tested, and how much time or resources were spent on tuning? How were the hyperparameters for the four baseline methods chosen? Especially for fine-tuning the baselines on your datasets, which I assume is done? Clearly describing the hyperparameter search is important for reproducibility. Please correct me if I was missing such details from the main manuscript.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: * Authors compare their randomized shape template-based synthetic data engine to one that uses data with no biomedical priors and one using brain regions. Can Authors elaborate more on the intuiton for why their randomly deformed template shapes are so effective? Is there some point at which the extent of the deformation causes the representations to be less useful?
soundness: 3
strengths: * The paper is very well written - it lays out the prior work and puts the contirbute in context.
* The approach yields representations that are both modality-agnostic and task-agnostic while removing the need for dataset-specific and anatomy-specific pre-training.
* Authors present results of several downstream tasks using their features including multi-modality registration image registration and few-shot segmentation on which their method outperform the others compared.
* Authors perform ablation studies on the various components of their pipeilne.
* The Authors present extensive visualization and quantitative results in their main text, and supplementary material. Algorithms and parameters are clearly presented too which allows for further scrutiny and improved reproducability.
* Authors are aware of the limitations of their approach and include these in the paper.
summary: Authors present a method to generate highly variable synthetic imaging data which is then used to pre-train a 3D network using contrastive learning. The data generation method consists of drawing samples from a semantically labelled repository of biomedical shape templates to randomly populate an empty 3D volume. The volume is then deformed. Empty space and organ 'envelopes' are also simulated. To simulate different modalities and protocols, random intensity transformation is applied to the deformed 3D volume to yield 2 images. Typical imaging artifacts such a sbias field and blurring are simulated through random augmentations.The two images are fed into a UNet, and contrastive pre-training is performed on features at each decoder layer. An anchor point is chosen in one of the images, and all voxels of that label in both images are considered positive, and everything else negative pairs. The network yields features that can be used to finetune on other modalities and tasks. Importantly, the representations are modality-agnostic and anatomy-agnostic.
weaknesses: * The segmentation task performed using the Authors features may yield better results than the other methods that are compared, however the result still misses significant portions of the anatomical regions they aim to segment. The features require further adjustment and extensive fine-tuning to be useful in diagnosis and treatment.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: The proposed workflow involves many hyper-parameters (Figure 12) controlling the properties of generated synthetic volumes -- what is the rule of thumb for choosing them?
soundness: 3
strengths: The overall methodology is straightforward and easy to understand. It echoes with the classical computer vision concept of invariance learning in the deep neural network era (although learned through a data-driven approach).

Improved image registration results on two public image registration benchmarks and image segmentation performance on six image segmentation datasets are shown, compared to those of some existing works. 

The paper is well-written with sufficient clarity. The illustrations are self-explanatory. Readers will enjoy reading it.
summary: This work proposes a pre-training approach for downstream tasks related to fine-grained volumetric medical data: image registration and semantic segmentation. The authors propose to learn appearance invariance and shapes of human anatomy through synthetic dense pixel volumes. In this process, 3D volumes are synthesized by randomly recomposing 3D anatomical shapes and assigning multiple sets of random pixel values, in together with synthetic noises and deformations. Pairs of synthetic volumes are used for multi-scale contrastive learning. The proposed approach demonstrates improved image registration and low-shot image segmentation results compared to some previous works. Detailed ablation studies on the pre-training configurations toward downstream performances are conducted.
weaknesses: Technical novelty: The core idea behind the approach is to leverage data-driven approach to learn invariance to pixel values through paired synthetic shapes and different pixel values, and to learn semantic-independent shape representation through random geometric (real-world and pure synthetic) shapes – both key ideas come from the well-established SynthMorph (random Ising shapes with random intensity + synthetic deformation for training registration networks. Hoffmann et al.) and SynthSeg (GMM-like pixel value model for repainting brain structures. Billot et al.) Despite leveraging more anatomical shapes beyond brain structures and applied to a contrastive framework, the essence remains unchanged.  

Many medical decisions are made not only on shape but also on subtle textures, for example, differentiating subtypes of tumors/lesions – toward which the proposed over-simplified appearance model by nature falls short. More sophisticated texture models need to be carefully studied beyond this manuscript.

For the same reason, high-level global semantic information such as relative locations between anatomical structures cannot be learned due to the nature of this approach. 

Real-world value: Given the increasing number of large-scale publicly accessible volumetric image datasets such as CT-RATE (Hamamci et al.), Totalsegmenter (Wasserthal et al.), and AbdomenAtlas (Li et al.), and the derived 3D foundation models, the real-world application of the proposed framework is unclear. Some of these large-scale public datasets come with fine-grain pixel-wise labels and associated radiological reports which provide additional supervision signals and text alignment potentials. The claimed generalization capability can be learned from multi-site large real-world datasets as well, through the intrinsic heterogeneity of big data and possibly through intense data augmentation.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1- What are the exact steps of the "label ensemble model" described in Section 3? Please write elaborate description of these steps.
2- Why do the generated images not look like real medical images? Is it because the deformation is too large? Why such "unrealistic" looking images are preferred rather than more realistic ones obtained with smaller deformation?
3- How does the quality of the representations obtained by the proposed backbone compares with SoTA foundational models such as DinoV2 or SAM2?
soundness: 3
strengths: - Foundational models are showing promising performance lately; however, we lack of a 3D model that work across different modalities in medical imaging. The paper proposes a solution to this important problem using domain randomisation and contrastive learning.
- The paper contain experiments on multiple datasets both for registration and few shot segmentation, and the results demonstrate the potential of the method.
- The idea of combining the ideas of domain randomisation and local contrastive learning to train a generic 3D backbone is quite interesting and, to my knowledge, is novel.
summary: The paper proposes training a backbone that generalizes across different datasets using synthetically generated dataset. The proposed pre-training strategy has 3 main steps: 1) Given a large datasets of 104 annotated organs, randomly sample anatomies, deform them and create a volume by ensembling these anatomies, 2) add noise and other augmentations to the volumes that are sampled in the previous step to simulate realistic looking synthetic medical images from labels, 3) train a U-Net with a contrastive objective by sampling two volumes that share the same 3D semantic layout but differ in appearance, treating corresponding features at different encoder levels as positives and all others as negatives. The pre-trained backbone is validated on two different tasks: 3D registration and 3D few-shot segmentation; using multiple datasets.  The results show the effectiveness of the proposed backbone in the experiments compared to existing methods.
weaknesses: - One issue I see in the paper is the convoluted description of the data engine step, especially the part creating the label ensemble model in section 3. I understand that this step is mainly based on the domain randomisation idea proposed in the literature. However, it is not really clear to me the steps between 201-207, especially the parts multiplying the masks with a randomly deformed sphere and randomly encasing half of the foreground-masked volumes.

- The images generated in the data engine step do not seem like as real medical images. Do they look like this because the deformation is too large? It is not clear why one would prefer training the model using such unrealistic images.

- The paper does not discuss the recent foundational models that show better generalization performance on many medical image datasets [1]. The downstream task performance of the representations obtained from the proposed backbone should be compared with the those obtained by the representations of a foundational model (e.g. DinoV2 [2]). For example, [3] is a recent paper that uses DinoV2 features for registration; but the same applies for the segmentation experiments. One can use the DinoV2 features for segmentation.

[1] Cekmeceli et al. "Do Vision Foundation Models Enhance Domain Generalization in Medical Image Segmentation?"
[2] Oquab et al. "DINOv2: Learning Robust Visual Features without Supervision"
[3] Song et al. "DINO-Reg: General Purpose Image Encoder for Training-Free Multi-modal Deformable Medical Image Registration"

Meta-review: The submission presents a method for pre-training a generalist 3D backbone for medical image segmentation and registration using synthetic data, combining domain randomization and local contrastive learning. Reviewers praised the method's innovation, sustainability (fewer trainable parameters), and strong experimental validation across multiple datasets, with comprehensive ablations and reproducibility ensured by publicly available code. All reviewers agreed on the strong potential and significance of the work.

Decision: Accept (Poster)