The Mean is the Mirage: Entropy-Adaptive Model Merging under Heterogeneous Domain Shifts in Medical Imaging

Sameer Ambekar; Reza Nasirigerdeh; Peter J. Schuffler; Lina Felsner; Daniel M. Lang; Julia A. Schnabel

cs.LGFeb 24, 2026

The Mean is the Mirage: Entropy-Adaptive Model Merging under Heterogeneous Domain Shifts in Medical Imaging

Authors: Sameer Ambekar, Reza Nasirigerdeh, Peter J. Schuffler, Lina Felsner, Daniel M. Lang, Julia A. Schnabel

Organizations: School of Computation, Information and Technology, Technical University of Munich, Germany · Institute of Machine Learning in Biomedical Imaging, Helmholtz Munich, Germany · relAI – Konrad Zuse School of Excellence in Reliable AI · Munich Center for Machine Learning (MCML) · Institute of Pathology, Technical University of Munich, Germany · School of Biomedical Engineering and Imaging Sciences, King’s College London, UK

Abstract

Model merging under unseen test-time distribution shifts often renders naive strategies, such as mean averaging unreliable. This challenge is especially acute in medical imaging, where models are fine-tuned locally at clinics on private data, producing domain-specific models that differ by scanner, protocol, and population. When deployed at an unseen clinical site, test cases arrive in unlabeled, non-i.i.d. batches, and the model must adapt immediately without labels. In this work, we introduce an entropy-adaptive, fully online model-merging method that yields a batch-specific merged model via only forward passes, effectively leveraging target information. We further demonstrate why mean merging is prone to failure and misaligned under heterogeneous domain shifts. Next, we mitigate encoder classifier mismatch by decoupling the encoder and classification head, merging with separate merging coefficients. We extensively evaluate our method with state-of-the-art baselines using two backbones across nine medical and natural-domain generalization image classification datasets, showing consistent gains across standard evaluation and challenging scenarios. These performance gains are achieved while retaining single-model inference at test-time, thereby demonstrating the effectiveness of our method.

Figures & tables

Figure 1: Mean merging and Entropy-Adaptive merging (Ours). (a) Mean merging averages independently trained hospital models into a static model that doesn’t consider the target information into account, thus can fail on mixed target batches. (b) Our method, adaptively, uses entropy for each unlabeled target batch to calculate per-batch merging coefficients along linear mode connectivity directions, producing batch-specific merged models while leveraging target information. Additional illustration with loss heatmaps in Fig. 3 .

Figure 2: Why Mean Merging fails for the PACS dataset using ViT-B32 domain-specific models trained on Photo, Cartoon, and Sketch domains and evaluated on Art domain. We first compute the angles as provided in ( Jang et al., 2024 ) between weights and bias terms of domain-specific trained models, then compute directional and angular drift to obtain deeper insights. (A) Data diversity: Model parameters drift in different directions and scales across domains, with the classification head showing the strongest misalignment. (B) Encoder-classification head mismatch: Mean angular layerwise drift. Misalignment increases with network depth and becomes most pronounced at the classification head. (C) Fixed mean merging is prone to failure: Mean merging leads to substantial signal loss in the classification head compared with individual models.

Panel of Figure 3 — Figure 3: Illustration of the loss landscape. (a) Mean merging: The linear path between models crosses a high-loss barrier because straightforward mean connectivity assumes compatible representations across all the layers. (b) Our method learns adaptive merging coefficients that follow a lower-loss path, resolving feature incompatibility near the unseen target optimum.

	5 Medical Datasets
	MiDog		Organs		Histo		ISIC Skin		Messidor
Method	B32	B16	B32	B16	B32	B16	B32	B16	B32	B16
Offline Model merging
Ensemble	89.24	89.22	69.87	73.71	92.16	89.87	86.89	86.24	49.75	54.58
DiWA ( Rame et al., 2022 )	89.65	89.07	66.93	71.26	89.48	89.48	85.60	85.82	49.50	52.81
Online Model merging

Table 1: State-of-the-art comparisons across medical datasets . We report the mean accuracy for ViT-B/32 (B32) and ViT-B/16 (B16) under the standard leave-one-domain-out ( Gulrajani and Lopez-Paz, 2020 ) setting, where domain-specific models are trained on respective domains and evaluated on an unseen target domain (adaptively). We re-implement the baselines as either offline or online model-merging methods and detail them in the Appendix. Our entropy-adaptive method achieves the best performance ( bold ).

	4 Natural Vision Datasets
	PACS		VLCS		Office-Home		Terra
Method	B32	B16	B32	B16	B32	B16	B32	B16
Offline Model merging
Ensemble	80.46	84.00	78.51	77.21	74.20	77.83	30.42	41.22
DiWA ( Rame et al., 2022 )	83.91	84.72	77.54	77.67	75.46	78.44	27.64	40.43
Online Model merging

Table 2: State-of-the-art comparisons across natural vision datasets . We report the mean accuracy for ViT-B/32 (B32) and ViT-B/16 (B16) and follow the same settings as in Table 1 . We reimplement the baselines as either offline or online model-merging methods and detail them in the Appendix. Our entropy-adaptive method achieves the best performance ( bold ).

Figure 4: Consistent improvements under Dirichlet and Temporal sampling. Mean accuracy (%) on PACS and Organs under Dirichlet splits (0.05, 0.50) and Temporally correlated data. Our Entropy-Adaptive (green) consistently outperforms the baselines with larger gains under a stronger skew factor (

\alpha=0.05

	Coefficients strategy		Mean accuracy
Method	Entropy only (Eq. 4 )	Decoupled head coef (Eq. 9 )	PACS	Organs
Entropy-Adaptive	✓		85.23 $\pm$ 0.2	69.81 $\pm$ 0.2
		✓	84.27 $\pm$ 0.2	67.80 $\pm$ 0.3
	✓	✓	87.33 $\pm$ 0.2	70.70 $\pm$ 0.2

Table 3: Benefits of Entropy only and decoupled head coefficients for ViT-B/32 on PACS and Organs. Both components improve performance; together, they deliver the strongest gains, highlighting the need for distinct coefficients for the encoder and the classification head.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Layer-wise domain-specific models misalignment. Pairwise parameter angles between domain-specific models across ViT and classification head components show that encoder layers are largely aligned, while the classification head (C) is strongly misaligned, motivating head-aware merging.

Explore similar work

Entropy Minimization without Model Collapse: Mitigating Prediction Bias in Medical Imaging
Jun 1, 2026Tim Nielen, Sameer Ambekar, Johannes Kiechle +2Distribution Shift Robustness Test-Time Adaptation
School of Computation, Information and Technology, Technical University of Munich, Germany · Institute of Machine Learning in Biomedical Imaging, Helmholtz Munich, Germany · Munich Center for Machine Learning (MCML) +3
When To Adapt? Adapting the Model or Data in Federated Medical Imaging
Apr 27, 2026Chamani Shiranthika, Parvaneh SaeediDomain Adaptation Medical Image Classification
School of Engineering Science, Simon Fraser University, Burnaby, BC, V5A 1S6
Model Merging for Medical LVLMs: A Benchmark and a Winner-Take-All Approach
Jul 17, 2026Lichao Mou, Shilan Zhang, Chunlei Li +7Low-Rank Adaptation Medical VLMs
MedAI Technology (Wuxi) Co. Ltd. · Wuhan University of Technology · Technical University of Munich +1

Jun 1, 2026J/K move · Enter open · S save

Entropy Minimization without Model Collapse: Mitigating Prediction Bias in Medical Imaging

Tim Nielen, Sameer Ambekar, Johannes Kiechle, Daniel M. Lang+1

School of Computation, Information and Technology, Technical University of Munich, Germany · Institute of Machine Learning in Biomedical Imaging, Helmholtz Munich, Germany · Munich Center for Machine Learning (MCML)+3

Distribution Shift Robustness Test-Time Adaptation Model Collapse Entropy Regularization

Open paper AskPDF arXiv

Entropy minimization (EM) is the dominant objective for test-time adaptation, yet its failure mode, model collapse, remains poorly understood. In this work, we show that distribution shifts can cause feature clusters corresponding to distinct classes in the model's representation space to merge, while the decision boundary remains fixed. This induces a systematic skew in the predicted class distribution, referred to as prediction bias. Prediction bias refers to a shift in the predicted class distribution, with some classes overrepresented and others suppressed. We show that entropy minimization amplifies this prediction bias by tightening the existing clusters, reinforcing the incorrect groupings until all predictions collapse to a trivial solution. Next, to demonstrate the significance of prediction bias and mitigate it, we further propose Distribution Shift Bias Reduction (DSBR), a bias-correcting objective that specifically targets this failure mode by equalizing the contribution of each predicted class to the unsupervised entropy minimization loss. To study this failure mode, we design suitable adaptation settings using four medical-imaging datasets and additionally evaluate on ImageNet-C. We find that DSBR consistently stabilizes test-time adaptation, prevents model collapse, and matches or outperforms state-of-the-art methods. Moreover, DSBR operates solely at test-time.

Abstract

Explore similar work

Entropy Minimization without Model Collapse: Mitigating Prediction Bias in Medical Imaging

When To Adapt? Adapting the Model or Data in Federated Medical Imaging

Model Merging for Medical LVLMs: A Benchmark and a Winner-Take-All Approach