SegWithU: Uncertainty as Perturbation Energy for Single-Forward-Pass Risk-Aware Medical Image Segmentation

Resumen Técnico: SegWithU

1. El Problema

La segmentación de imágenes médicas es una herramienta fundamental para la cuantificación anatómica y la toma de decisiones clínicas. Sin embargo, en la práctica clínica, una segmentación precisa no es suficiente; es crucial saber cuándo y dónde una predicción es poco fiable.

• Limitaciones de los métodos actuales:
  • Los métodos de alta calidad (como Deep Ensembles, Monte Carlo Dropout o Test-Time Augmentation) requieren múltiples pasadas de inferencia, lo que incrementa drásticamente el costo computacional y la latencia, haciéndolos difíciles de desplegar en entornos clínicos reales.
  • Las alternativas de una sola pasada (single-forward-pass) a menudo ofrecen un rendimiento inferior en la detección de errores o dependen de suposiciones restrictivas sobre la geometría del espacio de características (requiriendo reentrenamiento del modelo base).
  • En muchos flujos de trabajo médicos, el modelo de segmentación ya está validado e integrado; reentrenarlo o reemplazarlo para añadir incertidumbre es costoso e inviable.

Objetivo: Desarrollar un marco de trabajo post-hoc (después del entrenamiento) que pueda añadir estimación de incertidumbre a un modelo de segmentación preentrenado y congelado, sin modificar su arquitectura base ni requerir múltiples inferencias.

2. Metodología: SegWithU

SegWithU es un marco ligero que "envuelve" un modelo de segmentación preentrenado (backbone) y aprende una "cabeza de incertidumbre" (uncertainty head) sobre las características intermedias extraídas.

Principio Central:
La intuición fundamental es que las delineaciones anatómicas poco fiables son aquellas que son inestables bajo pequeñas perturbaciones latentes de la representación de características. En lugar de estimar densidades complejas, SegWithU mide la energía de perturbación en el espacio de características.

Componentes Clave:

1. Extracción de Características (Feature Tapping):
   • Se extraen mapas de características intermedias del backbone congelado (antes del bloque de salida).
   • Soporta una sola extracción o una fusión multi-escala de características.
2. Sondas Posteriores de Rango 1 (Rank-1 Posterior Probes):
   • Es el núcleo del método. Proyecta las características en un espacio de sondas compacto y de baja dimensión.
   • Utiliza convoluciones $1\times1$ para generar respuestas de sondas ($v$) que se combinan linealmente para simular perturbaciones en los logits de salida ($\Delta z$).
   • Calcula la varianza de las probabilidades de clase resultantes bajo estas perturbaciones para generar un mapa de incertidumbre epistémica ($U_{epi}$).
3. Descomposición de Mapas de Incertidumbre:
   SegWithU genera dos mapas de incertidumbre distintos con propósitos diferentes:
   • Mapa Orientado a la Calibración ($U_{cal}$): Se utiliza para ajustar (temperar) las probabilidades predictivas. Mejora la consistencia numérica entre la confianza y el resultado real (métrica Brier).
   • Mapa Orientado a la Clasificación/Ranking ($U_{rnk}$): Optimizado para detectar errores y priorizar qué píxeles deben ser revisados por un experto. No necesita estar calibrado probabilísticamente, sino ordenar los píxeles por severidad de error.
4. Función de Pérdida (Training Objective):
   Se entrena solo la cabeza de incertidumbre (el backbone permanece congelado) utilizando una combinación de pérdidas:
   • NLL: Para calibración de probabilidades.
   • Pérdida de Correlación de Error y Ranking Pares: Para asegurar que los píxeles erróneos reciban puntuaciones de incertidumbre más altas que los correctos (crucial para AUROC y AURC).
   • Pérdida de Cola (Tail Loss): Penaliza la baja incertidumbre en píxeles erróneos graves.
   • Pérdida de Confianza (Trust Loss): Regula la magnitud de las perturbaciones para evitar cambios inestables.

3. Contribuciones Clave

1. Marco Post-Hoc Eficiente: Introduce SegWithU, que añade una cabeza de incertidumbre ligera a un backbone preentrenado y congelado, evitando el reentrenamiento costoso y la inferencia múltiple.
2. Modelado de Energía de Perturbación: Propone una nueva visión de la incertidumbre basada en la sensibilidad a perturbaciones latentes mediante sondas de rango 1, en lugar de estimación de densidad o muestreo estocástico.
3. Separación de Roles: Demuestra que la calibración probabilística y la detección de fallos (ranking) son tareas distintas que deben modelarse con mapas separados para maximizar el rendimiento en ambos frentes.
4. Rendimiento Superior en una Sola Pasada: Logra el mejor rendimiento global entre los métodos de una sola pasada, superando a baselines como DUQ, DDU-Seg y DUE, y siendo competitivo con métodos de múltiples pasadas (como Ensembles) sin su costo computacional.

4. Resultados Experimentales

El método se evaluó en tres conjuntos de datos médicos estándar: ACDC (corazón), BraTS2024 (tumores cerebrales) y LiTS (tumores hepáticos).

• Métricas de Segmentación: SegWithU mantiene la calidad de segmentación original (puntuación Dice), demostrando que no degrada el modelo base.
• Calibración (Brier Score): Obtiene puntuaciones competitivas, a menudo las mejores entre los métodos de una sola pasada.
• Detección de Errores (AUROC y AURC):
  • AUROC: Logra los valores más altos en BraTS2024 (0.9946) y LiTS (0.9925), superando a todos los baselines de una sola pasada y compitiendo con los de múltiples pasadas.
  • AURC (Área bajo la curva de riesgo-cobertura): Consistentemente el mejor o cercano al mejor en todos los conjuntos de datos (ej. 2.4885 en ACDC, 0.2660 en BraTS2024). Esto indica que el modelo puede filtrar eficazmente los píxeles erróneos al abstenerse de las predicciones más inciertas.
• Análisis Cualitativo: Los mapas de incertidumbre de SegWithU están más concentrados en las fronteras ambiguas y regiones sospechosas, evitando las respuestas difusas o desconectadas típicas de otros métodos como Deep Ensembles o DUQ.

5. Significado e Impacto

• Viabilidad Clínica: SegWithU ofrece una ruta práctica para la segmentación médica consciente del riesgo. Permite desplegar sistemas de "auto-auditoría" en modelos existentes sin alterar su comportamiento ni requerir infraestructura de inferencia pesada.
• Cambio de Paradigma: El trabajo sugiere que la estimación de incertidumbre no debe tratarse como un subproducto de la confianza de clasificación, sino como una tarea estructurada de aprendizaje posterior.
• Eficiencia vs. Rendimiento: Demuestra que es posible obtener una calidad de incertidumbre de nivel superior (comparable a los ensembles) utilizando únicamente una pasada de inferencia, lo cual es crítico para aplicaciones en tiempo real o con recursos limitados.
• Robustez: La separación entre calibración y ranking permite optimizar específicamente para la detección de fallos, que es el uso más valioso en entornos clínicos (triaje y revisión experta).

En conclusión, SegWithU establece un nuevo estado del arte para la estimación de incertidumbre en segmentación médica de una sola pasada, demostrando que el modelado de perturbaciones basado en sondas es una estrategia efectiva, eficiente y clínicamente relevante.