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

Abstract

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.

Method

DSBR reweights the entropy loss so that every predicted class contributes equally. Samples in overrepresented clusters count for less and samples in underrepresented clusters count for more, which works against cluster merging. It needs no sample filtering or domain-specific augmentations, and its only hyperparameter is the EMA decay \(\alpha\).

Each step keeps an exponential moving average of the predicted class frequencies, starting from \(\hat{P}_{0}^{(k)} = 1/K\):

Here \(B^{(k)}_{t}\) is the number of samples in the batch predicted as class \(k\). The DSBR loss over batch \(\mathcal{B}_{t}\) then weights each sample’s entropy loss \(l(X)\) by the inverse of its predicted class frequency:

Results

Collapse on ImageNet-C

Predicted classes over one ImageNet-C Fog episode (severity 5). Both methods start with the same bias. Tent amplifies it until the model collapses, while DSBR dampens it until it vanishes by about batch 400.

Medical imaging benchmarks

Balanced accuracy (%), averaged over all source-target configurations. Δ Avg. is the average improvement over the non-adapted model. The best result in each column is in bold.

* Consistency component disabled, because it relies on 2D augmentations. † Not available, because the method relies on 2D augmentations.

ImageNet-C

Accuracy (%) at severity 5, in the default setting and three wild settings. The best result in each column is in bold.

Citation

@inproceedings{nielen2026entropy,
  title         = {Entropy Minimization without Model Collapse: Mitigating Prediction Bias in Medical Imaging},
  author        = {Nielen, Tim and Ambekar, Sameer and Kiechle, Johannes and Lang, Daniel M. and Schnabel, Julia A.},
  booktitle     = {Advances in Neural Information Processing Systems},
  year          = {2026},
  eprint        = {2606.02339},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}

Acknowledgments

This work is supported by DAAD programme Konrad Zuse Schools of Excellence in Artificial Intelligence, Munich Center for Machine Learning, both sponsored by the Federal Ministry of Research, Technology, Bavaria California Technology Center (BaCaTeC) under BaCaTeC-Förderprojekt Nr. 20 [2026-1], Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) – 515279324 / SPP 2177, and the German Federal Ministry of Research, Technology and Space (DECIPHER-M, 01KD2420G). DML and JAS received funding from HELMHOLTZ IMAGING, a platform of the Helmholtz Information and Data Science Incubator.