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

Sebo Diaz; Polina Golland; Elfar Adalsteinsson; Neel Dey

eess.IVApr 2, 2026

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

Authors: Sebo Diaz, Polina Golland, Elfar Adalsteinsson, Neel Dey

Organizations: MIT · MIT, MGH, HMS

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.

Figures & tables

Figure 1: Training on Stable Representations. When standard ERM models are trained on one domain (A, top) but tested on another (A, bottom), their representations are unstable under domain shifts (B). While standard training is performant in-domain (E, top), this instability significantly degrades performance on out-of-distribution images (E, bottom). MaskGen provides a theoretically grounded and easy-to-implement framework that exploits both in-domain image intensities (A, top) and stable, domain-invariant representations extracted by foundation models (C, top) to train robust models that effectively generalize to new domains automatically without any adaptation (F, bottom). Representations visualized by mapping 3 arbitrary channels to RGB.

Panel of Figure 2 — Figure 2: Method. Left : MaskGen is exceedingly simple to implement. Given a standard PyTorch training loop, the green lines are the only additions required. Right : The probabilistic graphical model we assume for domain generalization. Label $Y$ generates both stable $X_{s}$ and unstable $X_{u}$ variables and the environment $E$ influences only $X_{u}$ .

Figure 3: Qualitative all-data segmentation results. Rows correspond to different datasets and domain shift types. Cols. 1 and 2 visualize in-domain training and out-of-domain testing samples, respectively. Cols. 3–7 illustrate qualitative results on out-of-domain test set examples.

Method	AMOS Ji et al. (2022) Modality	BraTS Antonelli et al. (2022) Sequence	CoW Yang et al. (2024) Modality	HVSMR Pace et al. (2024) Disease	PanDG Zhang et al. (2025) Phase	Prostate Liu et al. (2021) Site	Avg. Rank
ERM	57.4 (1.7)	38.9 (1.8)	24.6 (1.5)	63.6 (3.0)	25.3 (2.5)	83.0 (0.7)	8.0
GIN Ouyang et al. (2022)	65.1 (0.7)	29.9 (1.6)	17.9 (1.5)	64.1 (3.0)	41.1 (1.8)	82.1 (0.8)	7.2
GIN+IPA Ouyang et al. (2022)	$\underline{66.5(0.8)}$	$\underline{44.7(1.6)}$	$\underline{37.6(1.9)}$	64.7 (2.6)	50.3 (1.4)	$\underline{84.1(0.6)}$	$\underline{2.5}$
MixUp Zhang et al. (2017)	61.5 (0.9)	41.2 (1.7)	36.3 (1.5)	64.5 (3.2)	26.7 (2.5)	82.5 (0.7)	5.5
CutOut DeVries and Taylor (2017)	55.5 (1.7)	41.8 (1.6)	26.5 (1.9)	64.9 (3.5)	30.5 (2.4)	84.0 (0.6)	5.0
DSBN Zhou et al. (2022)	58.0 (1.4)	40.5 (1.7)	29.6 (1.9)	$66.1(3.2)$	26.3 (2.6)	83.4 (0.6)	5.2

Table 1: All-data domain generalization results. Mean Dice score (w/ std. error) across datasets against domain generalization baselines, with avg. ranking on the far-right.

Method	AMOS Modality	CoW Modality	HVSMR Disease	Prostate Site
ERM	5.4(1.0)	18.2(1.3)	40.4(3.8)	27.8(5.0)
GIN	6.0(0.7)	13.4(1.1)	$44.9(4.0)$	20.3(4.3)
GINIPA	$\underline{37.8(1.5)}$	$23.8$ (1.7)	$44.5$ (3.8)	$27.9$ (5.1)
MixUp	$15.6$ (1.6)	$\underline{33.7(1.2)}$	$43.7$ (3.9)	$\underline{53.2(3.4)}$
CutOut	8.4(1.1)	18.2 (1.4)	41.3(3.6)	40.8(4.5)
DSBN	$19.3$ (1.8)	$22.2$ (1.5)	$44.2$ (3.7)	$44.1$ (5.2)

Table 2: Few-shot domain generalization results. Left : Mean Dice score (w/ std. error) for few-shot experiments. Right : Arbitrarily selected few-shot results.

Figure 4: Ablating Regularization via View Masking . Top: We train models with only the image (“Image Only"), only the representations (“Reps. Only"), and both combined with increasing view masking rates (MR=

0\%,25\%

50\%

75\%

). MaskGen increases validation performance over training on the image alone, the representations alone, and over combination without masking (MR=

0\%

). Bottom: We perform a one-channel removal analysis to measure the sensitivity in validation (

\mathrm{\Delta Dice}

) w.r.t. the full input (image and representations). The drop in performance indicates that, without view masking, the model does not balance information sources and heavily relies on the image input.

Dataset	Model	0.0	0.25	0.50	0.75	RepOnly
AMOS Ji et al. (2022)	nnInteract. Isensee et al. (2025)	64.3	65.0	65.9	66.4	63.9
AMOS Ji et al. (2022)	anatomix Dey et al. (2025)	65.1	65.8	66.6	67.4	64.4
BraTS Antonelli et al. (2022)	nnInteract. Isensee et al. (2025)	41.5	43.4	44.2	42.9	42.7
BraTS Antonelli et al. (2022)	anatomix Dey et al. (2025)	41.7	42.5	47.2	47.2	41.3
CoW Yang et al. (2024)	nnInteract. Isensee et al. (2025)	48.7	45.9	48.9	50.8	50.7
CoW Yang et al. (2024)	anatomix Dey et al. (2025)	49.0	49.5	51.0	52.2	51.7

Table 3: Varying the representation extractor. Validation Dice scores for mask rates of 0.0, 0.25, 0.50, and 0.75, as well as training on reps. alone.

Method	Adaptation Needed?	AMOS Modality	BraTS Sequence	CoW Modality	HVSMR Disease	PanDG Phase	Prostate Site	Avg. Rank
EntMin (ERM) Wang et al. (2020)	Yes	64.7 (1.2)	38.9 (1.8)	24.4 (1.5)	64.4 (2.9)	27.7 (2.4)	84.5 (0.5)	4.2
EntMin (MR=0%) Wang et al. (2020)	Yes	$\underline{65.8(1.2)}$	$\underline{39.3(1.6)}$	45.1 (1.2)	65.0 (2.8)	$\underline{35.7(2.0)}$	83.6 (0.5)	$\underline{3.0}$
SFB Eastwood et al. (2023)	Yes	63.6 (1.0)	38.8 (1.6)	41.0 (1.3)	65.2 (2.6)	30.4 (2.5)	85.4 (0.5)	3.3
Alpha Tuning Jiang and Veitch (2022)	Yes	64.1 (1.3)	41.7 (1.8)	$\textbf{49.0}(1.4)$	64.8 (3.3)	28.5 (2.5)	84.3 (0.5)	3.2
Ours	No	67.5 (0.9)	46.2 (1.6)	$\underline{47.8(1.3)}$	65.8(3.4)	46.0 (1.9)	$\underline{84.7(0.5)}$	1.3

Table 4: Comparisons with methods that perform adaptation. Mean Dice score for test-time adaptation (TTA) and joint-predictor settings. TTA methods observe the test distribution at inference, whereas MaskGen uses no adaptation and shares the same model weights across all evaluations.

Dataset	All-data Setting				Few-shot Setting
	$0.0$	$0.25$	$0.50$	$0.75$	$0.0$	$0.25$	$0.50$	$0.75$
AMOS Ji et al. (2022)	65.1	65.8	66.6	67.4	23.5	41.9	50.8	53.6
BraTS Antonelli et al. (2022)	41.7	42.5	47.2	47.2	–	–	–	–
CoW Yang et al. (2024)	49.0	49.5	51.0	52.2	39.0	43.5	37.7	43.0
HVSMR Pace et al. (2024)	61.7	62.6	63.5	62.0	43.4	44.4	45.3	44.5
PanDG Zhang et al. (2025)	51.8	60.2	60.6	60.1	–	–	–	–

Table 6: Effect of Mask Rate. Validation Dice score as a function of mask rate used in training. Left: all-data results; right: few-shot results.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Comparing cross-modality representations extracted by foundation models. Rows 1–4: Given a subject from BraTS, we visualize its representations across MRI sequence shifts produced by nnInteractive [ 19 ] (rows 1, 2) and anatomix [ 9 ] (rows 3, 4), finding that anatomix representations are more stable and suitable for domain generalization. Rows 5–8: We now produce a similar visualization for two unpaired subjects with cross-modality domain shift from the AMOS dataset and find the same stability trend.

Figure 6: CKA Representation Similarity of Foundation Models Under Sequence Shifts. A: CKA [ 28 ] between each 3D foundation model’s chosen embeddings across MRI sequence shifts given subjects with multiple MRI sequences from BraTS. B: Given a fixed sequence, we now perturb image contrast (via gamma adjustments) and visualize CKA as a measure of representational stability under contrast shifts. C: A per-subject visualization of the plot in panel A .

Figure 7: Choice of layer from which to use representations. We compare anatomix’s final block representations (rows 1–2, 5-6) to its penultimate block representations (rows 3–4, 7-8) for two different datasets with two different kinds of domain shifts. Rows 1–4: FLAIR (rows 1 and 3) and T2 (rows 2 and 4) from BraTS, representing intra-subject MRI sequence shifts. Rows 5–8: CT (rows 5 and 7) and MRI (rows 6 and 8) from AMOS, representing inter-subject modality shifts. Representation channels (columns 2-7) were arbitrarily selected. Qualitatively, the penultimate output better preserves high-frequency details while retaining inter-domain stability.

Figure 8: Histogram of Cosine Similarities. Cosine similarities for the two masked regimes of (1,0) and (1,1) for our method sampled every 100 steps during an "All-data" training experiment on BraTS [ 34 , 1 ] to empirically support Condition 1 . A KDE curve is fitted to the histogram for clarity.

Method	AMOS [ 22 ]	BraTS [ 34 , 1 ]	CoW [ 53 ]	HVSMR [ 36 ]	PanDG [ 57 ]	Prostate [ 32 ]
F.T. Extractor	56.8 (1.5)	37.7 (1.7)	17.3 (1.7)	64.6 (3.2)	23.7 (2.6)	84.2 (0.5)
F.T. Decoder	57.2 (1.5)	39.5 (1.7)	18.5 (1.7)	63.2 (2.9)	27.7 (2.6)	83.8 (0.7)
Ours	67.5 (0.9)	46.2 (1.6)	47.8 (1.3)	65.8 (3.4)	46.0 (1.9)	84.7 (0.5)

Appendix

Table 7: Finetuning domain-stable foundation models does not yield domain generalization . Ablation comparing finetuning all of the foundation model’s weights (F.T. Extractor), finetuning its decoder-alone (F.T. Decoder), and our proposed approach in the All-data setting.

Ablations	AMOS [ 22 ]	BraTS [ 34 , 1 ]
Squeeze+Excite [ 15 ]	64.3 (1.1)	38.8 (1.7)
Ours	67.5 (0.9)	46.2 (1.6)

Appendix

Table 8: Alternative approaches for feature combination. A comparison of alternative feature combination methods: squeeze+excite channel attention vs. our input view masking.

Table 16

Method	AMOS [ 22 ] Modality	BraTS [ 1 ] Sequence	CoW [ 53 ] Modality	HVSMR [ 36 ] Disease	PanDG [ 57 ] Phase	Prostate [ 32 ] Site	Avg. Rank
ERM	57.4 (1.7)	38.9 (1.8)	24.6 (1.5)	63.6 (3.0)	25.3 (2.5)	83.0 (0.7)	8.0
GIN [ 35 ]	65.1 (0.7)	29.9 (1.6)	17.9 (1.5)	64.1 (3.0)	41.1 (1.8)	82.1 (0.8)	7.2
GIN+IPA [ 35 ]	$\underline{66.5(0.8)}$	$\underline{44.7(1.6)}$	$\underline{37.6(1.9)}$	64.7 (2.6)	50.3 (1.4)	$\underline{84.1(0.6)}$	$\underline{2.5}$
MixUp [ 55 ]	61.5 (0.9)	41.2 (1.7)	36.3 (1.5)	64.5 (3.2)	26.7 (2.5)	82.5 (0.7)	5.5
CutOut [ 8 ]	55.5 (1.7)	41.8 (1.6)	26.5 (1.9)	64.9 (3.5)	30.5 (2.4)	84.0 (0.6)	5.0
DSBN [ 62 ]	58.0 (1.4)	40.5 (1.7)	29.6 (1.9)	$66.1(3.2)$	26.3 (2.6)	83.4 (0.6)	5.2

Appendix

Table 11: All-data domain generalization results for untuned masking rate. Mean Dice score (w/ std. error) across datasets against domain generalization baselines, with avg. ranking on the far right. Companion to Table 1 with a default untuned masking rate of

r=0.5

Method	AMOS Modality	CoW Modality	HVSMR Disease	Prostate Site	Avg. Rank
ERM	5.4 (1.0)	18.2 (1.3)	40.4 (3.8)	27.8 (5.0)	8.0
GIN	6.0 (0.7)	13.4 (1.1)	$44.9(4.0)$	20.3(4.3)	6.5
GINIPA	$\underline{37.8(1.5)}$	$23.8$ (1.7)	$44.5$ (3.8)	$27.9$ (5.1)	3.5
MixUp	$15.6$ (1.6)	$\underline{33.7(1.2)}$	$43.7$ (3.9)	$\underline{53.2(3.4)}$	3.5
CutOut	8.4 (1.1)	18.2 (1.4)	41.3 (3.6)	40.8 (4.5)	5.5
DSBN	$19.3$ (1.8)	$22.2$ (1.5)	$44.2$ (3.7)	$44.1$ (5.2)	3.5

Appendix

Table 12: Few-shot domain generalization results for untuned masking rate. Mean Dice score (w/ std. error) across datasets against domain generalization baselines, with avg. ranking on the far right. Companion to Table 2 with a default untuned masking rate of

r=0.5

Figure 9: Corruption Analysis. We apply simulated domain shifts with increasing strengths, via bias ( top ) and contrast/gamma ( middle ) corruptions and observe that MaskGen maintains high robustness. In the bottom row, we perturb the trained model weights with additive Gaussian noise, and find that MaskGen is stable under these corruptions, indicating a flatter and more generalizable solution. Dashed vertical line in all plots signifies the "identity" or no corruption level.

Table 20

Block	Contents
0	Conv-BN-ReLU( $n_{in},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	Conv-BN-ReLU( $2n_{c}$ )

Appendix

Table 15: U-Net architectural details. Conv-BN-ReLU refers to a sequence of 3D convolutional with

3\times 3\times 3

kernels, batch normalization, and pointwise ReLU non-linearities.

n_{in}

is the number if input channels.

n_{c}

is the channel width multiplier and

n

is the number of output channels.

Dataset	Crop Size	Spacing (mm)	# Train	# Val.	# Test
AMOS	$(192,192,128)$	$(1.0,1.0,1.5)$	$300$	$30$	$30$
BraTS	$(128,128,128)$	$(1.0,1.0,1.0)$	$1020$	$48$	$96$
CoW	$(192,160,128)$	$(0.3,0.3,0.6)$	$75$	$20$	$30$
HVSMR	$(128,160,128)$	$(0.74,0.74,0.8)$	$23$	$18$	$19$
PancreasDG	$(256,224,64)$	$(1.09,1.09,4.4)$	$53$	$100$	$100$
Prostate	$(224,256,64)$	$(0.52,0.51,1.25)$	$49$	$25$	$42$

Appendix

Table 16: Dataset details. Pre-processing details such as crop size and spacing, alongside splits used for each dataset in the experiments.

Explore similar work

Frequency Adapter with SAM for Generalized Medical Image Segmentation
May 11, 2026Phuoc-Nguyen Bui, Van-Nguyen Pham, Duc-Tai Le +2Frequency-Domain Feature Learning Domain Generalization
Sungkyunkwan University, Korea
Multiple Domain Generalization Using Category Information Independent of Domain Differences
Apr 8, 2026Reiji Saito, Kazuhiro HottaDomain Generalization Cross-Domain Generalization
Meijo University, 1-501 Shiogamaguchi, Tempaku-ku, Nagoya 468-8502, Japan
Simple, Safe, and Overlooked: Reclaiming Sustainable Domain Generalization with Statistical Color Matching
Aug 19, 2026Sebastian Doerrich, Francesco Di Salvo, Shyam Nandan Rai +2Data Augmentation Domain Generalization
xAILab Bamberg, University of Bamberg, Bamberg, Germany

May 11, 2026J/K move · Enter open · S save

Frequency Adapter with SAM for Generalized Medical Image Segmentation

Phuoc-Nguyen Bui, Van-Nguyen Pham, Duc-Tai Le, Junghyun Bum+1

Sungkyunkwan University, Korea

Frequency-Domain Feature Learning Domain Generalization Vision Foundation Model Adaptation Medical Image Segmentation

Open paper AskPDF arXiv

Medical image segmentation is a critical task in computer-aided diagnosis and treatment planning. However, deep learning models often struggle to generalize across datasets due to domain shifts arising from variations in imaging protocols, scanner types, and patient populations. Traditional domain generalization (DG) methods utilize causal feature learning, adversarial consistency, and style augmentation to improve segmentation robustness. While effective, these approaches rely on explicit feature alignment, adversarial objectives, or handcrafted augmentations, which may not fully exploit the capabilities of foundation models. Recently, the Segment Anything Model (SAM) has demonstrated strong generalization capabilities in segmentation tasks. SAM-based DG methods attempt to improve medical image segmentation. However, these approaches primarily operate in the spatial domain and overlook frequency-based discrepancies that significantly affect model robustness. In this work, we propose Frequency-based Domain Generalization with SAM (FSAM), a novel framework that integrates Low-Rank Adaptation (LoRA) for efficient fine-tuning and a frequency adapter to incorporate frequency-domain representations for single-source domain generalization. FSAM enhances SAM's segmentation robustness by extracting domain-invariant high-frequency features, mitigating frequency-related domain shifts. Experimental results on fundus and prostate datasets demonstrate that FSAM outperforms existing traditional DG and SAM-based DG approaches in domain generalization. Codes and pre-trained models will be made available on GitHub.

Abstract

Explore similar work

Frequency Adapter with SAM for Generalized Medical Image Segmentation

Multiple Domain Generalization Using Category Information Independent of Domain Differences

Simple, Safe, and Overlooked: Reclaiming Sustainable Domain Generalization with Statistical Color Matching