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

Resumen Técnico: Minimización de la Entropía sin Colapso del Modelo

Planteamiento del Problema
La Adaptación en Tiempo de Prueba (TTA, por sus siglas en inglés) es esencial para desplegar modelos de aprendizaje automático en escenarios del mundo real donde las distribuciones de datos cambian debido a variaciones en sensores, entornos o protocolos de adquisición. En el ámbito de la imagen médica, estos cambios suelen ser sutiles y heterogéneos (por ejemplo, diferentes escáneres o protocolos de tinción), pero pueden causar una degradación significativa del rendimiento. La Minimización de la Entropía (EM) se ha convertido en el objetivo dominante para la TTA, con el fin de reducir la incertidumbre de la predicción en datos de prueba no etiquetados. Sin embargo, la EM sufre un modo de fallo crítico conocido como colapso del modelo, donde el modelo converge a predecir una única clase independientemente de la entrada.

Los autores identifican que la causa raíz de este colapso no es meramente el proceso de adaptación en sí, sino un fenómeno que denominan sesgo de predicción. Bajo cambios de distribución, los grupos de características (feature clusters) correspondientes a distintas clases en el espacio de representación pueden desplazarse y fusionarse parcialmente, mientras que la frontera de decisión del clasificador permanece fija. Esto induce un sesgo en la distribución de clases predichas donde algunas clases se sobre-representan y otras se suprimen. La EM estándar amplifica este sesio existente al estrechar los grupos fusionados, creando un bucle de retroalimentación positiva que impulsa al modelo hacia una solución trivial de una sola clase.

Metodología: Reducción del Sesgo por Cambio de Distribución (DSBR)
Para mitigar este modo de fallo, los autores proponen la Reducción del Sesgo por Cambio de Distribución (DSBR), un objetivo de corrección de sesgo diseñado para operar únicamente en el tiempo de prueba sin acceso a datos o etiquetas de origen.

1. Mecanismo Central: DSBR aborda el sesgo de predicción igualando la contribución de cada clase predicha a la pérdida de entropía no supervisada. En lugar de permitir que los grupos dominantes dicten las actualizaciones del gradiente, DSBR repondera la pérdida para evitar que una sola clase abrume la señal de optimización.
2. Formulación Matemática: El método define una distribución objetivo equilibrada $p^*_X$ que asigna una masa de probabilidad igual ($1/K$) a cada grupo predicho. Utilizando el reponderado por importancia, la pérdida para una muestra $x$ se escala por $1 / (K \cdot \hat{p}_{\hat{Y}}(\hat{y}(x)))$, donde $\hat{p}_{\hat{Y}}$ es la frecuencia estimada de la clase predicha. Esto asegura que las muestras en grupos subrepresentados contribuyan más al gradiente, mientras que aquellas en grupos sobre-representados contribuyan menos, fomentando activamente la separación de los grupos fusionados.
3. Implementación: Dado que la verdadera distribución de clases predichas es desconocida en el tiempo de prueba, DSBR la aproxima utilizando un Promedio Móvil Exponencial (EMA) de las frecuencias de los grupos a través de los lotes (batches). Esto es particularmente crucial para conjuntos de datos con un alto número de clases (por ejemplo, ImageNet con $K=1000$) donde los tamaños de lote pequeños de otro modo conducirían a estimaciones degeneradas por lote.

Contribuciones Clave

• Identificación del Sesgo de Predicción: El artículo establece que el sesgo de predicción —un desequilibrio en la distribución de clases predichas presente antes de que comience la adaptación debido a la fusión de los grupos de características— es la causa raíz del colapso del modelo en la EM. Demuestra que la EM no crea este sesgo, sino que lo amplifica hasta que ocurre el colapso.
• Algoritmo DSBR: Los autores introducen DSBR, un método que aborda específicamente este modo de fallo mediante el reponderado de la pérdida de entropía para equilibrar las contribuciones de las clases. A diferencia de otros enfoques, no depende del filtrado de muestras, aumentos específicos del dominio o regularizadores complejos.
• Evaluación Exhaustiva: El estudio diseña escenarios de adaptación desafiantes a través de cuatro conjuntos de datos de imagen médica (WILDSCamelyon, Histopantum, MammoBench, GliomaMRI) y el benchmark de visión natural ImageNet-C.

Result los

• Prevención del Colapso: En los benchmarks de imagen médica, los métodos de EM estándar (como Tent) frecuentemente colapsan hacia una precisión de nivel de azar en pares específicos de dominio origen-objetivo. DSBR previene consistentemente este colapso, logrando una precisión casi óptima en todas las configuraciones.
• Ganancias de Rendimiento: DSBR supera a los estados del arte de referencia (incluyendo Tent, SAR, ROID y DeYO). En los conjuntos de datos médicos, logró una mejora promedio de +4.2% (ResNet) y +2.6% (ViT) sobre las bases no adaptadas, superando significativamente a los siguientes mejores métodos.
• Robustez: En ImageNet-C, DSBR demostró una estabilidad superior bajo corrupciones severas, corrupciones mixtas y cambios extremos de etiquetas. Manejó eficazmente tamaños de lote tan pequeños como 1, siempre que el decaimiento del EMA se ajustara adecuadamente, mientras que otros métodos a menudo fallaban o colapsaban.
• Validación del Mecanismo: Las visualizaciones de los incrustamientos (embeddings) del espacio de características confirman que, mientras la EM causa que los grupos se fusionen y colapsen, DSBR mantiene los grupos bien separados durante todo el proceso de adaptación.

Significancia
El artículo argumenta que estabilizar la minimización de la entropía requiere una comprensión mecánica y fundamentada de por qué falla. Al aislar el sesgo de predicción como el principal motor del colapso, los autores demuestran que corregir esta única cantidad es suficiente para estabilizar la TTA. La significancia de este trabajo radica en su demostración de que la adaptación en tiempo de prueba robusta en dominios críticos como la imagen médica no requiere heurísticas complejas o filtrado de datos, sino una corrección directa de la dinámica de optimización que conduce al colapso del modelo. DSBR ofrece una solución simple, con pocos hiperparámetros, que opera enteramente en el tiempo de prueba, lo que lo hace altamente aplicable para el despliegue clínico donde la fiabilidad es primordial.