🤖 machine learning

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

This paper identifies that entropy minimization in test-time adaptation can lead to model collapse by amplifying prediction bias caused by distribution shifts, and proposes Distribution Shift Bias Reduction (DSBR) to mitigate this issue by equalizing class contributions, thereby stabilizing adaptation and preventing collapse across medical imaging and ImageNet-C benchmarks.

Original authors: Tim Nielen, Sameer Ambekar, Johannes Kiechle, Daniel M. Lang, Julia A. Schnabel

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tim Nielen, Sameer Ambekar, Johannes Kiechle, Daniel M. Lang, Julia A. Schnabel

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Confident Fool"

Imagine you have a very smart student (the AI model) who studied hard for a math test using a specific textbook (the training data). They learned the rules perfectly.

But on the day of the final exam, the teacher hands them a new version of the test where the font is slightly different, the paper is a different color, and the lighting in the room is dimmer (this is Distribution Shift). The questions are still the same, but the look of the test has changed.

The student gets confused. They start guessing. Because they are unsure, they lean heavily on the answers they feel most "comfortable" with, even if those answers are wrong.

Entropy Minimization (EM) is a technique used to help the student adapt. The idea is: "If you are unsure, try to become more confident." The computer tries to reduce its own confusion (entropy) by making its predictions sharper and more certain.

The Catch: In this paper, the authors show that this "just be more confident" strategy can backfire spectacularly. If the student starts leaning toward the wrong answer because of the weird lighting, telling them to "be more confident" just makes them more confidently wrong. Eventually, they stop looking at the questions entirely and just shout the same answer for every single question. The paper calls this Model Collapse.

The Root Cause: The "Merging Clusters"

Why does this happen? The authors discovered a specific mechanism they call Prediction Bias.

Imagine the student's brain has two neat piles of notes: one pile for "Normal" cases and one pile for "Tumors."

The Shift: When the test conditions change (e.g., a different medical scanner), the "Tumor" notes get messy and start sliding over into the "Normal" pile. The piles start to merge.
The Bias: Because the piles are merging, the student starts seeing "Tumors" where there are none, or "Normal" where there are tumors. They are biased toward one side.
The Amplification: Standard Entropy Minimization sees this mess and tries to tidy it up. But instead of separating the piles back out, it just pushes the "Normal" pile harder, absorbing the messy "Tumor" notes into it. It tightens the grip on the wrong grouping.
The Collapse: The "Tumor" pile gets crushed and disappears. The student now thinks everything is "Normal." They have collapsed into a single, useless answer.

The Solution: DSBR (The "Fairness Coach")

The authors propose a new method called Distribution Shift Bias Reduction (DSBR).

Think of DSBR as a fairness coach standing next to the student during the exam.

The Problem: The student is shouting "Normal! Normal! Normal!" because they are biased.
The Fix: The coach says, "Wait a minute. You are shouting 'Normal' way too much. Let's make sure you give equal attention to 'Tumor' as well."

DSBR works by reweighting the student's learning process:

If the student is predicting "Normal" 90% of the time, the coach tells the computer: "Don't listen to those 'Normal' answers as much. They are too loud."
If the student is predicting "Tumor" only 10% of the time, the coach says: "Listen closely to these rare 'Tumor' answers. They are valuable."

By forcing the model to treat all classes equally, DSBR stops the "Normal" pile from swallowing the "Tumor" pile. It keeps the two piles separate, preventing the student from collapsing into a single, wrong answer.

Why This Matters for Medicine

The paper tested this on medical imaging (like X-rays and MRI scans).

The Stakes: In medicine, if a model collapses and says "Everything is healthy" when there is actually a tumor, it could be life-threatening.
The Results: The authors tested their "Fairness Coach" (DSBR) on four different medical datasets and a standard image dataset (ImageNet).
- Old methods (like Tent) often failed, causing the model to collapse and give wrong answers.
- DSBR consistently stopped the collapse. It kept the model balanced and accurate, even when the images looked weird or came from different hospitals.

Summary

The Villain: Entropy Minimization, which tries to make AI more confident but accidentally makes it stubbornly wrong when the data changes.
The Crime: "Prediction Bias," where different types of data merge together, tricking the AI into thinking one answer is always right.
The Hero: DSBR, a method that acts like a fairness coach, ensuring the AI doesn't ignore rare or difficult cases just because it's trying to be confident.
The Outcome: A more reliable AI that doesn't give up and say "everything is fine" when it's actually in trouble.

Technical Summary: Entropy Minimization without Model Collapse

Problem Statement
Test-Time Adaptation (TTA) is essential for deploying machine learning models in real-world scenarios where data distributions shift due to changes in sensors, environments, or acquisition protocols. In medical imaging, these shifts are often subtle and heterogeneous (e.g., different scanners or staining protocols), yet they can cause significant performance degradation. Entropy Minimization (EM) has become the dominant objective for TTA, aiming to reduce prediction uncertainty on unlabeled test data. However, EM suffers from a critical failure mode known as model collapse, where the model converges to predicting a single class regardless of the input.

The authors identify that the root cause of this collapse is not merely the adaptation process itself, but a phenomenon they term prediction bias. Under distribution shifts, feature clusters corresponding to distinct classes in the representation space can drift and partially merge, while the classifier's decision boundary remains fixed. This induces a skew in the predicted class distribution where some classes become overrepresented and others suppressed. Standard EM amplifies this existing bias by tightening the merged clusters, creating a positive feedback loop that drives the model toward a trivial, single-class solution.

Methodology: Distribution Shift Bias Reduction (DSBR)
To mitigate this failure mode, the authors propose Distribution Shift Bias Reduction (DSBR), a bias-correcting objective designed to operate solely at test time without access to source data or labels.

Core Mechanism: DSBR addresses the prediction bias by equalizing the contribution of each predicted class to the unsupervised entropy minimization loss. Instead of allowing dominant clusters to dictate the gradient updates, DSBR reweights the loss to prevent any single class from overwhelming the optimization signal.
Mathematical Formulation: The method defines a balanced target distribution $p^*_X$ that assigns equal probability mass ( $1/K$ ) to each predicted cluster. Using importance reweighting, the loss for a sample $x$ is scaled by $1 / (K \cdot \hat{p}_{\hat{Y}}(\hat{y}(x)))$ , where $\hat{p}_{\hat{Y}}$ is the estimated frequency of the predicted class. This ensures that samples in underrepresented clusters contribute more to the gradient, while those in overrepresented clusters contribute less, actively encouraging the separation of merged clusters.
Implementation: Since the true predicted class distribution is unknown at test time, DSBR approximates it using an Exponential Moving Average (EMA) of cluster frequencies across batches. This is particularly crucial for high-class-count datasets (e.g., ImageNet with $K=1000$ ) where small batch sizes would otherwise lead to degenerate per-batch estimates.

Key Contributions

Identification of Prediction Bias: The paper establishes that prediction bias—a skew in the predicted class distribution present before adaptation begins due to feature cluster merging—is the root cause of model collapse in EM. It demonstrates that EM does not create this bias but amplifies it until collapse occurs.
DSBR Algorithm: The authors introduce DSBR, a method that specifically targets this failure mode by reweighting the entropy loss to balance class contributions. Unlike other approaches, it does not rely on sample filtering, domain-specific augmentations, or complex regularizers.
Comprehensive Evaluation: The study designs challenging adaptation scenarios across four medical imaging datasets (WILDSCamelyon, Histopantum, MammoBench, GliomaMRI) and the natural vision benchmark ImageNet-C.

Results

Prevention of Collapse: On medical imaging benchmarks, standard EM methods (like Tent) frequently collapse to chance-level accuracy on specific source-target domain pairs. DSBR consistently prevents this collapse, achieving near-optimal accuracy across all configurations.
Performance Gains: DSBR outperforms state-of-the-art baselines (including Tent, SAR, ROID, and DeYO). On medical datasets, it achieved an average improvement of +4.2% (ResNet) and +2.6% (ViT) over non-adapted baselines, significantly surpassing the next-best methods.
Robustness: On ImageNet-C, DSBR demonstrated superior stability under severe corruptions, mixed corruptions, and extreme label shifts. It effectively handled batch sizes as small as 1, provided the EMA decay was tuned appropriately, whereas other methods often failed or collapsed.
Mechanism Validation: Visualizations of feature space embeddings confirm that while EM causes clusters to merge and collapse, DSBR maintains well-separated clusters throughout the adaptation process.

Significance
The paper argues that stabilizing entropy minimization requires a principled, mechanistic understanding of why it fails. By isolating prediction bias as the primary driver of collapse, the authors show that correcting this single quantity is sufficient to stabilize TTA. The significance of this work lies in its demonstration that robust test-time adaptation in critical domains like medical imaging does not require complex heuristics or data filtering, but rather a direct correction of the optimization dynamics that lead to model collapse. DSBR offers a simple, hyperparameter-light solution that operates entirely at test time, making it highly applicable for clinical deployment where reliability is paramount.