Learning General-purpose Biomedical Volume Representations using Randomized Synthesis

Thank you for the insightful comments and questions. We are happy to see that the reviewer found the work interesting and novel.

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

Thank you for raising this, we agree that it could be clearer and have edited the text accordingly. These steps are also described in detail in Appendix B.1.

To summarize, many datasets in radiology consist of a foreground (body, head, etc.) and an empty background (e.g., air). Our data engine synthetically simulates these configurations by masking the synthetic volume with a random foreground mask (the deformed sphere). Further, many biomedical domains have a layer-like structure near the foreground (e.g., skin/fat in abdominal CT/MRI, the cortex in brain MRI, etc.). To simulate structures like these, we synthetically create a layer around the foreground as well.

Quantitatively, we ablate this data engine choice in Table 3 (Ours w/o FG mask) and find a consistent drop in performance across both tasks and all four radiological datasets without it, highlighting the benefits of this step.

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

We apologize for the potential confusion and clarify this point in the revision. The generated images are not intended to look like real medical images. Instead, they provide biomedically-informed generic training data that is densely labeled but not necessarily realistic. It is not because the deformation is too large.

As intuition: Networks often overfit to both the intensity and shape characteristics of the training domain, limiting cross-domain generalization (e.g., CT to MRI segmentation). While current domain randomized methods like SynthSeg show that training on synthetic brain images with random intensities yields contrast-invariant brain segmentation, these methods still overfit to the domain (i.e., brain label maps), need specific training frameworks for specific tasks (e.g., segmentation, registration), and need dense ground truth labels for any new application domain.

Our work avoids overfitting to specific biomedical tasks and datasets by randomly sampling and deforming biomedical shape templates/primitives to construct label maps that are broadly generalizable across various biomedical contexts and tasks. Combined with our contrastive training mechanism that learns intensity invariance and a bias towards shape, we observe broad gains across datasets and tasks.

Our ablations (Table 3) show that training on only real labels (brains) does not generalize to new biomedical domains. Further, training on purely synthetic labels (smshapes) does not work either, suggesting that biomedical shape priors are important for learning a robust representation. Our work balances between these extremes by instead independently deforming real biomedical shapes to achieve robust generalization across domains.

Thank you for your valuable and relevant questions and constructive comments. We are happy to see that the reviewer found the approach straightforward, motivated by classic computer vision, and well-presented, all leading to improved results across multiple datasets and tasks.

We believe that there may be some miscommunications on our part and hope to address them below.

”The core idea is … to learn invariance to pixel values through paired synthetic shapes and different pixel values, and to learn semantic-independent shape representation through random geometric shapes – both key ideas come from SynthMorph and SynthSeg … Despite leveraging more anatomical shapes beyond brain structures and applied to a contrastive framework, the essence remains unchanged.”

Thank you for raising this point. To clarify, SynthMorph and SynthSeg are highly specialized frameworks designed exclusively for brain registration and brain segmentation, respectively. Each network is trained in a fully supervised regime and is strongly biased towards the task and the anatomical region (the brain) it was trained to handle, as suggested by our experimental results (Table 3 rows 2–3 and Figure 5).

Building networks that extract robust representations of biomedical datasets that are task- and application-independent requires a new approach, which we tackle with two contributions:

• To ensure generalization to a wide range of shapes, we introduce a data engine that synthesizes random configurations of biomedical shape templates that yield robust off-the-shelf representations for any unseen anatomical region and imaging modality.
• To provide generalization to a wide range of downstream tasks, we employ a contrastive loss that is supervised at every voxel and that utilizes paired samples from the data engine to learn invariance to common imaging variation and supports generalization across tasks.

To our knowledge, both contributions are novel (as also noted by Reviewers 8tqR and voNv) and essential for generalist performance.

To empirically emphasize the distinctions from SynthMorph-shapes and SynthSeg:

• Without our proposed data engine, contrastive training on only SynthSeg’s synthetic brain labels (Table 3, row 3, brains) or on only SynthMorph’s random shapes (Table 3, row 2, smshapes) drops performance consistently across all tasks and non-Brain datasets.
• Further, in Figure 5, specifically for registration, our method outperforms the SynthMorph-shapes model trained on random shapes (mentioned by the reviewer) by large margins. On the abdominal dataset, there is a 72-point median Dice gap between SynthMorph-shapes and our method. On the cardiac dataset, there is an 11-point median Dice difference.

Moreover, without our proposed contrastive training framework, training on entirely randomized shape configurations does not scale with other losses. In Table 3, “ablating pretraining losses”, we show that other pretraining losses such as denoising or unsupervised representation learning (NeurIPS22) result in worse performance.

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

To clarify, we do not aim to address tasks that require subtle texture-based differentiation such as lesion subtyping (this is listed in our limitations section).

Importantly, this is an intentional choice as it has been established that networks that are biased towards shape instead of texture are more robust and more broadly applicable across domain shifts. For examples, see:

• ICLR’19 oral
• NeurIPS’21 spotlight

Our work is similarly inspired. By focusing on learning biomedical shape representations, we achieve broad robustness to new biomedical domains and flexibility for multiple tasks, as illustrated by

• SOTA multimodal registration across two datasets (abdominal MR/CT and cardiac MR/CT).
• Consistent few-shot segmentation improvements across several completely distinct datasets featuring contexts such as fetuses, hearts, prostates, etc.

As such, texture differentiation is out of scope for our paper which makes gains on several other fronts. We have now expanded the limitations regarding texture for clarity, thank you for bringing this up.

Thank you for the relevant feedback and questions and for emphasizing the work’s technical contributions and experimental thoroughness.

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

We agree. However, this is true of all few-shot segmentation evaluation frameworks which, by definition, cannot match the performance of fully supervised networks finetuned on the entire training set. As we use only a few annotated volumes, we assess all pretrained initializations for data efficiency. We find that in the data-constrained setting, our work provides a robust dataset- and modality-agnostic improvement over previous work.

This data efficiency in few-shot segmentation has several downstream benefits such as accelerating pseudo-labeling pipelines widely used to annotate vast datasets (e.g., TotalSegmentator), to provide a warm-start to active learning pipelines, among others. Further, although diminishing in extent, our method does yield benefits at higher levels of supervision as well, please see Figure 8 in the main paper.

”Can Authors elaborate more on the intuition for why their randomly deformed template shapes are so effective?”

Thank you for raising this important question.

Networks often overfit to both the intensity and shape characteristics of the training domain, limiting cross-domain generalization (e.g., CT to MRI segmentation). While current methods like SynthSeg show that training on synthetic brain images with random intensities yields contrast-invariant brain segmentation, these methods still overfit to the domain (i.e., brain label maps), need specific training frameworks for specific tasks (e.g., segmentation, registration), and need dense ground truth labels for any new domain they are intended to be used in.

Our work avoids overfitting to specific biomedical tasks and datasets by randomly sampling and deforming biomedical shape templates/primitives to construct label maps that are broadly generalizable across various biomedical contexts and tasks. Combined with our contrastive training mechanism that learns intensity invariance and a bias towards shape, we find broad gains across datasets and tasks.

Our ablations (Table 3) show that training on only real labels (brains) does not generalize to new biomedical domains. Further, training on purely synthetic labels (smshapes) does not work either, suggesting that biomedical shape priors are important for learning a robust representation. Our work balances between these extremes by instead independently deforming real biomedical shapes to achieve robust generalization across domains.

”Is there some point at which the extent of the deformation causes the representations to be less useful?”

The extent of deformation can hypothetically hamper downstream performance if the upper bound of the scale range of the affine deformations in the data engine is set too large. If it is too large for the last template sampled, the template can fill the entire volume, overwriting all others. In our work, the global scale range is set to [0.5, 1.5] and we did not observe this effect in practice.

Thank you for the valuable feedback and for highlighting the method’s innovativeness, excellent downstream task performance, sustainability, reproducibility, and the paper’s directness with any limitations.

“Considered datasets. If I am not missing any details the authors mostly evaluate their method on CT and MRI images.”

We apologize for the miscommunication on our part. To clarify, MRI is not a single modality and is a broad grouping of various imaging sequences that all highlight vastly different tissue properties.

Our MRI datasets vary in sequences, using BOLD, HASTE, T2-weighted, bSSFP, and SPIR. For example, WUFetal uses BOLD MRI which measures blood oxygenation whereas SPIR is used for fat-suppressed structural imaging. Further, each dataset contains disparate anatomical regions and contexts, from pregnant uteruses to cardiac images. These differences make each dataset a unique domain for assessing generalizability.

In our revision, we now explicitly clarify the sequence differences between datasets. Thank you for catching this.

“They even consider an MRI image as out of distribution.”

If the reviewer is referring to L405, we meant that WUFetal (whole uterus BOLD images of fetuses, placentae, and surrounding organs) is not a publicly available dataset type, and thus it is out-of-distribution for our baselines pretrained on multiple public datasets. We have edited this sentence for clarity.

’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”’

We agree. As suggested, we have now added two new few-shot segmentation experiments on 3D Ultrasound and microscopy datasets into Appendix A.6, providing further details therein.

• For 3D ultrasound, we use ultrasound volumes from the SegThy dataset which images the thyroid and provides labels for the thyroid, carotid artery, and jugular vein.
• For 3D microscopy, we use the NucMM-M dataset that images cortical nuclei in the mouse visual cortex. Instance-level annotations were converted to binary semantic labels for consistency with the semantic segmentation approach used throughout the paper.

For ease of reference, we summarize the Dice results below, with the best method in bold and runner-up in code formatting:

For 3D ultrasound (SegThy), we improve upon all baseline pretrained models by a clear margin, demonstrating our model’s usefulness to this new domain. W.r.t. 3D microscopy (NucMM-M), we achieve 2nd place performance to SMIT. However, as in Tables 1 and 2 in our paper (plus SegThy), our method consistently outperforms SMIT across all other seven segmentation tasks and both registration tasks, highlighting its generalizability.

To add to the discussion of claims, please see below.

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

We agree. We have now added clarifications that our model is specifically suited for radiology throughout the paper.

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

These details are provided in the appendices, our apologies for the confusion.

For the registration experiments, all baseline grid search design choices and results are in Appendix A.2. All registration baseline implementation details are in Appendix B.4.1–4.3. For the few-shot segmentation experiments, implementation, tuning, and hyperparameter details are provided in Appendix B.5. Lastly, our ablation implementation details are provided in Appendix B.6 and B.7. For further clarity, we have now added additional details where relevant.