A3-TTA: Adaptive Anchor Alignment Test-Time Adaptation for Image Segmentation

Resumen Técnico de A3-TTA: Adaptación en Tiempo de Prueba de Alineación de Anclajes Adaptativos para la Segmentación de Imágenes

Planteamiento del Problema
Los modelos de segmentación de imágenes basados en aprendizaje profundo suelen sufrir una degradación del rendimiento cuando se despliegan en dominios de destino que difieren de sus dominios de entrenamiento (fuente) debido a variaciones en los protocolos de imagen, escáneres o condiciones ambientales. Aunque la Adaptación de Dominio No Supervisada (UDA) y la Adaptación de Dominio No Supervisada Libre de Fuente (SFUDA) abordan esto, a menudo requieren un entrenamiento de múltiples épocas o acceso a datos de origen, lo que limita su utilidad en entornos clínicos en tiempo real donde los datos llegan en flujos. La Adaptación en Tiempo de Prueba (TTA) ofrece una solución al adaptar los modelos de forma online utilizando únicamente imágenes de prueba no etiquetadas sin datos de origen. Sin embargo, las estrategias de TTA existentes enfrentan desafíos significativos:

1. La adaptación de la Normalización por Lotes (BN) a menudo falla ante grandes desplazamientos de dominio debido a distribuciones de características desajustadas.
2. La optimización no supervisada (por ejemplo, la minimización de la entropía) carece de supervisión a nivel de píxel y puede provocar un olvido catastrófico.
3. Los métodos basados en pseudo-etiquetas suelen depender de heurísticas de perturbación por ensamble (por ejemplo, dropout, aumentación, inyección de ruido) para generar auto-etiquetas. Estos métodos carecen de un fundamento distributivo explícito, lo que resulta en pseudo-etiquetas ruidosas e inestables que desencadenan la acumulación de errores y degradan el rendimiento.

Metodología: Marco de Trabajo A3-TTA
Los autores proponen A3-TTA (Adaptación en Tiempo de Prueba de Alineación de Anclajes Adaptativos), un marco diseñado para construir pseudo-etiquetas fiables mediante supervisión guiada por anclajes sin requerir datos de origen. El método opera de manera de paso único y online, y consta de cuatro componentes principales:

1. Identificación de Imagen de Anclaje-Objetivo (ATI) mediante Densidad de Compacidad de Clase (CCD):
   En lugar de depender de perturbaciones estocásticas, A3-TTA identifica un subconjunto de imágenes del dominio de destino que son "bien predichas" y están distribucionalmente cerca del dominio de origen. Estas se denominan Imágenes de Anclaje-Objetivo (ATI). La selección se basa en una nueva métrica llamada Densidad de Compacidad de Clase (CCD), que calcula una puntuación de compacidad ($S_i$) basada en la distribución de probabilidad softmax de las predicciones de clase. Una puntuación más baja indica una mayor compacidad de predicción y una menor ambigüedad interclase. A diferencia de la entropía a nivel de píxel o del MC Dropout, la CCD es una señal global a nivel de clase computable en un solo paso hacia adelante.

2. Banco de Características Dinámico y Alineación de Características:
   Las ATI seleccionadas se almacenan en un Banco de Características de capacidad fija. Para cada imagen de prueba entrante, el modelo recupera la característica de anclaje más similar de este banco. Se aplica entonces una estrategia de Refinamiento de Alineación de Características (FAR):

   • Fusión de Características Sensible a la Similitud: La característica latente de la imagen de prueba se fusiona con la característica de anclaje recuperada utilizando un coeficiente ponderado por similitud.
   • Normalización de Características Basada en Anclajes: La característica fusionada se normaliza utilizando la media y la desviación estándar de la característica del ancla.
     Este proceso refina la representación de la imagen de prueba, alineándola más cerca de la distribución similar a la fuente para generar pseudo-etiquetas más fiables.

3. Supervisión Densa Guiada por Bordes:
   Las pseudo-etiquetas refinadas se utilizan para supervisar el modelo mediante una estrategia de pérdida dual:

   • Consistencia Semántica ($L_{sem}$): Una pérdida de entropía cruzada normalizada que asegura la alineación entre la predicción original y la pseudo-etiqueta refinada.
   • Minimización de Entropía Sensible a los Bordes ($L_{be}$): Una pérdida que se dirige específicamente a las regiones de alta entropía (típicamente los bordes de los objetos) para reducir la incertidumbre y mejorar la claridad de los bordes, evitando la inestabilidad de la minimización de la entropía global.

4. Profesor Medio Robusto Auto-Adaptativo:
   Para estabilizar las actualizaciones en un entorno de una sola época, A3-TTA emplea un marco de Profesor Medio (Mean Teacher). Crucialmente, introduce una tasa de Promedio Móvil Exponencial (EMA) Auto-Adaptativa. La tasa de actualización se regula dinámicamente mediante la divergencia de entropía cruzada normalizada entre las predicciones del profesor y el estudiante. Si la divergencia es grande (indicando un desplazamiento de distribución significativo), el profesor se actualiza de forma más agresiva; si la divergencia es pequeña, el profesor retiene más de sus parámetros previos para mantener la estabilidad.

Contribuciones Clave

• Imágenes de Anclaje-Objetivo (ATI) y Métrica CCD: Introducción de las ATI como referencias estables y una nueva métrica de Densidad de Compacidad de Clase para su identificación eficiente en un solo paso basada en la confianza de la softmax.
• Refinamiento de Alineación de Características: Una estrategia novedosa que utiliza un banco de características dinámico y operaciones de fusión/normalización para refinar las representaciones de las imágenes de prueba para la generación de pseudo-etiquetas.
• Supervisión Mejorada: Integración de consistencia semántica y minimización de la entropía sensible a los bordes para mejorar la calidad de la predicción a nivel de píxel.
• EMA Auto-Adaptativo: Una estrategia de actualización dinámica para el marco de profesor-estudiante que equilibra la adaptación rápida con la estabilidad basada en la consistencia de la predicción.
• Agnosticismo del Modelo: El marco no requiere datos de origen, tareas proxy adicionales ni modificaciones arquitectónicas en los modelos de segmentación existentes.

Resultados Experimentales
Los autores evaluaron A3-TTA en tres conjuntos de datos multidominio:

1. Resonancia Magnética Cardíaca (M&MS): Segmentación del ventrículo izquierdo, miocardio y ventrículo derecho a través de cuatro proveedores de escáneres.
2. Resonancia Magnética de Próstata: Segmentación a través de seis sitios diferentes.
3. Cityscapes (Condiciones Adversas): Segmentación bajo niebla, noche, lluvia y nieve utilizando DeepLabV3+.

Rendimiento Cuantitativo:

• Imágenes Médicas: A3-TTA mejoró los coeficientes Dice promedio en un 10.40% (Corazón) y un 17.68% (Próstata) respecto a la línea base de solo origen (source-only). Superó consistentemente a los métodos de TTA de vanguardia (por ejemplo, TENT, CoTTA, SAR) en todos los dominios de destino, logrando los coeficientes Dice más altos y la menor Distancia de Superficie Simétrica Promedio (ASSD).
• Imágenes Naturales: En el conjunto de datos de condiciones adversas de Cityscapes, A3-TTA mejoró el Dice en un 16.90% respecto al modelo de origen, superando al mejor método TTA existente (TENT) por 2.14–2.93 puntos porcentuales.
• TTA Continua: En escenarios de adaptación secuencial (adaptación a través de múltiples dominios de destino), A3-TTA demostró una fuerte capacidad anti-olvido, manteniendo un alto rendimiento en rondas subsecuentes donde otros métodos sufrieron acumulación de errores.
• Robustez: Bajo ruido Rician simulado y desenfoque de movimiento, A3-TTA mantuvo un rendimiento significativamente superior en comparación con otros métodos de TTA y el modelo de origen.
• Eficiencia: El método es computacionalmente eficiente, procesando cortes en aproximadamente ~0.044 segundos con un menor consumo de memoria GPU (4.99 GB) en comparación con CoTTA (7.23 GB) bajo configuraciones similares.

Significancia y Reivindicaciones
El artículo afirma que A3-TTA aborda los problemas críticos de inestabilidad y acumulación de errores inherentes a los métodos actuales de TTA basados en pseudo-etiquetas. Al reemplazar las perturbaciones heurísticas con una alineación guiada por anclajes, el método proporciona un mecanismo con fundamento distributivo para generar señales de supervisión fiables. Los autores enfatizan que A3-TTA es particularmente significativo para el despliegue clínico en tiempo real, donde los datos llegan en flujos, los datos de origen no están disponibles y el entrenamiento de múltiples épocas es inviable. La capacidad del marco para mantener el rendimiento a través de dominios secuenciales sin olvido catastrófico lo convierte en una solución robusta para entornos dinámicos. Los autores señalan que, si bien el método es efectivo para desplazamientos de distribución típicos, las brechas extremas de trans-modalidad (por ejemplo, de CT a MRI) siguen siendo un desafío, y el trabajo futuro explorará la extensión del enfoque a modelos 3D y modelos fundacionales.