Uncertainty Estimation in Pathology Foundation Models via Deep Mutual Learning

Resumen Técnico: Estimación de la Incertidumbre en Modelos Fundacionales de Patología mediante el Aprendizaje Mutuo Profundo

Planteamiento del Problema
Los Modelos Fundacionales de Patología (PFM, por sus siglas en inglés) han demostrado capacidades de generalización sólidas para la extracción de representaciones de imágenes de portaobjetos completos (WSI) en tareas como el análisis de supervivencia, el descubrimiento de biomarcadores y la subtipificación de tumores. Sin embargo, su traslación a la práctica clínica rutinaria se ve obstaculizada por dos limitaciones críticas:

1. Falta de Estimaciones de Confianza: Los flujos de trabajo actuales de los PFM suelen generar predicciones sin medidas explícitas de confianza o incertidumbre. En entornos clínicos, los modelos no solo deben ser precisos, sino también capaces de indicar cuándo sus predicciones son inciertas para apoyar una toma de decisiones segura.
2. Rendimiento Dependiente de la Tarea: Ningún PFM es universalmente superior en todas las tareas downstream, conjuntos de datos y contextos clínicos. La selección del modelo óptimo para el despliegue es un desafío, y confiar en un solo modelo conlleva el riesgo de un rendimiento subóptimo o de una confianza excesiva en casos propensos al fallo.

Los métodos de fusión actuales suelen mejorar el rendimiento, pero no proporcionan estimaciones de incertidumbre explícitas y agnósticas a la tarea. Por el contrario, los métodos de cuantificación de la incertidumbre, como los ensambles profundos (deep ensembles) o el dropout de Monte Carlo, no están diseñados específicamente para ensambles de PFM o son específicos para una tarea (por ejemplo, limitados al diagnóstico de glioblastoma).

Metodología: Marco DICE
Los autores proponen DICE (Disagreement-Informed Coordination of Experts - Coordinación de Expertos Informada por el Desacuerdo), un marco de trabajo plug-and-play que ensambla $K$ PFMs congelados para estimar la incertidumbre y mejorar el rendimiento downstream. La intuición central es que el desacuerdo entre PFMs heterogéneos sirve como una señal informativa de incertidumbre, siempre que este desacuerdo sea estructurado y no esté dominado por el ruido de optimización.

El marco opera de la siguiente manera:

1. Arquitectura: El sistema utiliza $K$ PFMs congelados (por ejemplo, Virchow2, UNI2-h, H-optimus-1, CONCHv1.5, Hibou-L). Cada PFM actúa como un "experto" $E^{(k)}$, emparejado con un agregador de Aprendizaje de Múltiples Instancias (MIL) independiente (específicamente ABMIL) y una cabeza de clasificación.

2. Objetivos de Entrenamiento: Para asegurar que el desacuerdo del ensamble refleje una incertidumbre significativa en lugar de ruido, los expertos se acoplan durante el entrenamiento mediante tres objetivos:

   • Pérdida Supervisada ($L_{sup}$): Pérdida de entropía cruzada estándar para garantizar una clasificación precisa a nivel de portaobjetos.
   • Aprendizaje Mutuo Profundo ($L_{DML}$): Un objetivo de alineación en el espacio de predicción donde cada experto minimiza la divergencia de Kullback–Leibler (KL) entre su distribución predictiva y las de sus $K-1$ pares. Esto fomenta el consenso mientras mantiene la alineación con la verdad fundamental (ground truth).
   • Regularización de Gramian ($L_{Gram}$): Un objetivo de alineación en el espacio de representación basado en la medida de Gramian. Este fomenta que las incrustaciones (embeddings) de los portaobjetos de los expertos abarquen un subespacio de bajo volumen, evitando la desalineación geométrica que podría hacerse pasar por incertidumbre.
   • Pérdida Total: El objetivo de entrenamiento combina estos términos: $L = L_{sup} + L_{DML} + \lambda_{Gram}(t) L_{Gram}$, donde $\lambda_{Gram}$ sigue un esquema de calentamiento (warm-up).

3. Inferencia y Estimación de la Incertidumbre:

   • Proxy de Incertidumbre: En la inferencia, la predicción del ensamble es el promedio de las distribuciones de los $K$ expertos. La incertidumbre predictiva total se descompone en componentes aleatorios (datos) y epistémicos (modelo). Los autores demuestran teóricamente que la incertidumbre epistémica (el desacuerdo residual) está acotada superiormente por la pérdida de agregación DML. Por lo tanto, minimizar la pérdida DML durante el entrenamiento contrae el límite superior de la incertidumbre del modelo.
   • Localización: La localización de lesiones a nivel de parche se logra de forma no supervisada mediante el cálculo del consenso de los pesos de atención de todos los expertos.

Contribuciones Clave

1. Primer Marco Agnóstico a la Tarea: DICE es el primer marco que aprovecha los ensambles de PFM específicamente para la estimación de la incertidumbre agnóstica a la tarea en patología computacional.
2. Fundamento Teórico: El artículo proporciona una prueba teórica de que la minimización del objetivo de alineación DML durante el entrenamiento contrae un límite superior de la incertidumbre epistémica del ensamble, validando el uso del desacuerdo residual como un proxy de incertidumbre fiable.
3. Rendimiento Empírico: Experimentos extensos demuestran que DICE no solo identifica con precisión las muestras propensas al fallo (tanto in-distribution como out-of-distribution), sino que también iguala o supera a los PFM individuales de vanguardia (SOTA) y a las bases de comparación de fusión en clasificación, calibración y localización de lesiones no supervisada.

Resultados Experimentales
El marco fue evaluado en tres benchmarks de aprendizaje débilmente supervisado: PANDA (gradación de cáncer de próstata), CAMELYON16 y CAMELYON17 (detección de metástasis en ganglios linfáticos mamarios).

• Clasificación y Calibración: DICE alcanzó los puntajes F1 más altos y la menor Log-Verosimilitud Negativa (NLL) en la mayoría de los conjuntos de datos. En PANDA, mejoró el F1 en 3.8 puntos porcentuales respecto al mejor PFM individual y en 2.8 puntos respecto a la mejor base de comparación de fusión.
• Estimación de la Incertidumbre: En experimentos de predicción selectiva, DICE postergó eficazmente los portaobjetos de alta incertidumbre, resultando en las tasas de error retenido más bajas. Demostró una separación superior entre las predicciones correctas e incorrectas (medida por el $\Delta$ de Cliff) en comparación con el Dropout de Monte Carlo y otros métodos de fusión. Esta capacidad se mantuvo incluso bajo cambios de cohorte (por ejemplo, entrenar en CAMELYON16 y probar en CAMELYON17).
• Localización: Mediante la agregación de mapas de atención, DICE proporcionó una localización a nivel de parche no supervisada que fue competitiva o superior a los modelos individuales, particularmente al concentrar la atención dentro de las regiones tumorales anotadas en lugar del estroma circundante.
• Estudios de Ablación: El marco resultó robusto a la elección del agregador MIL (DSMIL, TransMIL) y mostró que aumentar el número de expertos produce rendimientos decrecientes una vez que el espacio de características está suficientemente cubierto, resaltando la importancia de la complementariedad de los expertos sobre el mero tamaño del ensamble.

Significancia y Reivindicaciones
El artículo posiciona a DICE como un paso crucial hacia la traducción de los PFM en sistemas de apoyo a la decisión conscientes de la incertidumbre. Al establecer un enfoque teóricamente fundamentado para fusionar PFMs, DICE aborda la naturaleza de "caja negra" de los modelos actuales, ofreciendo un mecanismo para identificar casos propensos al error y delegarlos de forma segura a los clínicos. Los autores afirman que su método hace que los PFM sean más confiables en entornos clínicos al proporcionar estimaciones de incertidumbre fiables sin requerir supervisión explícita para la localización o la calibración de la confianza.

Los autores reconocen las limitaciones, incluyendo el costo computacional de extraer características de $K$ modelos congelados (aunque los costos de entrenamiento siguen siendo menores que el ajuste fino de extremo a extremo) y la necesidad de equilibrar cuidadosamente la alineación y la diversidad para evitar la supresión del mismo desacuerdo requerido para la estimación de la incertidumbre. Concluyen que, si bien el número de expertos es actualmente fijo, el marco allana el camino para modelos de patología más fiables y conscientes de la incertidumbre integrados en los flujos de trabajo clínicos.