Why Invariance is Not Enough for Biomedical Domain Generalization and How to Fix It

TL;DR

MaskGen is a simple, theoretically grounded method for improving biomedical image segmentation robustness across diverse clinical settings by leveraging source intensities and foundation model features.

Contribution

It introduces MaskGen, a novel approach that enhances domain generalization in biomedical segmentation without complex modifications or extensive augmentations.

Findings

Significant performance improvements in biomedical segmentation across clinical shifts.

Effective in both fully supervised and few-shot learning scenarios.

Architecture- and loss-agnostic, compatible with standard pipelines.

Abstract

We present MaskGen, a theoretically grounded and deliberately simple approach for domain generalization in 3D biomedical image segmentation. Modern segmentation models degrade sharply under shifts in modality, disease severity, clinical sites, and more, limiting their reliable adoption. Existing generalization methods address this using extreme augmentations, hand-engineered domain statistics mixing, or architectural redesigns that add significant implementation overhead while yielding inconsistent performance across biomedical settings. MaskGen instead presents a principled learning strategy with marginal overhead that utilizes both source-domain image intensities and domain-stable foundation model representations to train robust segmentation models. As a result, MaskGen achieves strong gains in both fully supervised and few-shot segmentation across broad clinical shifts in biomedical studies. Unlike prior approaches, MaskGen is architecture- and loss-agnostic, compatible with standard augmentation pipelines, easy to implement, and tackles arbitrary anatomical regions. Its implementation is freely available at https://github.com/sebodiaz/MaskGen.