MedShift: Implicit Conditional Transport for X-Ray Domain Adaptation

1. El Problema

La adaptación de dominio en imágenes médicas es un desafío crítico, especialmente al trasladar modelos entrenados con datos sintéticos a entornos clínicos reales. Existen brechas significativas (domain gaps) entre las simulaciones y las radiografías reales de la cabeza, causadas principalmente por:

• Perfiles de atenuación: Las simulaciones a menudo no replican la compleja atenuación de rayos X a través de estructuras anatómicas heterogéneas (hueso, aire, tejidos blandos).
• Características de ruido: Los sistemas reales introducen ruido estructurado, dispersión y artefactos de compresión que los simuladores no modelan.
• Representación de tejidos blandos y contraste: Las dinámicas de contraste, especialmente en los límites óseos o en características anatómicas superpuestas, suelen estar simplificadas en las simulaciones.

Los métodos existentes (como GANs o modelos de difusión) suelen requerir entrenamiento específico por par de dominios o dependen de datos emparejados (paired data), lo cual es raro en medicina. Además, muchos enfoques actuales carecen de un marco unificado para traducir entre múltiples dominios simultáneamente.

2. Metodología: MedShift

El paper propone MedShift, un modelo generativo condicional unificado diseñado para la traducción de imágenes de alta resolución entre dominios de rayos X (sintético $\leftrightarrow$ real) sin necesidad de datos emparejados.

• Arquitectura Base: MedShift combina Flow Matching (FM) y Puentes de Schrödinger (Schrödinger Bridges). Aprende un mapa de transporte implícito condicionado a etiquetas de dominio.
• Espacio Latente: Opera en un espacio latente aprendido a través de un VAE (Autoencoder Variacional) preentrenado, lo que reduce drásticamente el costo computacional mientras preserva el detalle semántico y la resolución espacial.
• Mecanismo de Transporte (Inferencia):
  1. Codificación (Atrás): Una imagen fuente ($x_1$) se codifica en una representación latente agnóstica al dominio ($z_\tau$) integrando hacia atrás desde $t=1$ hasta un tiempo intermedio $\tau \in (0,1)$ bajo la condición del dominio fuente.
  2. Traducción (Adelante): Se integra hacia adelante desde $\tau$ hasta $t=1$ bajo la condición del dominio objetivo para generar la imagen traducida ($\hat{x}_1$).
• Flexibilidad: El modelo utiliza una guía sin clasificador (Classifier-Free Guidance, CFG) y permite ajustar el parámetro $\tau$ en tiempo de inferencia. Esto permite a los usuarios priorizar la fidelidad perceptual (bajo $\tau$) o la consistencia estructural (alto $\tau$) sin reentrenar el modelo.

3. Contribuciones Clave

1. MedShift: Un nuevo modelo basado en Flow Matching y Puentes de Schrödinger para la traducción de imágenes no emparejadas de alta resolución, capaz de generalizar entre múltiples dominios en un solo marco unificado.
2. X-DigiSkull: Lanzamiento de un nuevo dataset que contiene radiografías de cráneos sintéticos y reales alineadas. Incluye:
   • Datos sintéticos generados con el simulador Mentice VIST®.
   • Datos reales adquiridos con el sistema Philips Azurion IGT.
   • Variaciones en dosis de radiación (Baja, Normal, Exposición) y múltiples ángulos de adquisición.
3. Benchmarking: Evaluación exhaustiva de varios modelos de vanguardia (GANs, Flows, Difusión) en este nuevo dataset, estableciendo una referencia para futuras investigaciones en adaptación de dominio médica.

4. Resultados

Los experimentos comparan MedShift con métodos como Hierarchy Flow, CycleGAN-Turbo, Z-STAR y SDEdit (basado en Stable Diffusion).

• Rendimiento General: MedShift logra el mejor equilibrio entre realismo (medido por CFID, Coverage, CMMD) y preservación estructural (medido por SSIM, LPIPS).
  • CycleGAN-Turbo tiene el mejor CFID (realismo) pero introduce distorsiones anatómicas graves.
  • Hierarchy Flow preserva la estructura perfectamente pero falla en la adaptación de estilo (bajo CFID).
  • MedShift supera a ambos, ofreciendo una traducción fiel que mantiene la anatomía crítica mientras adopta las características de ruido y contraste del dominio real.
• Eficiencia Computacional:
  • MedShift utiliza una arquitectura U-Net personalizada que es 6 veces más pequeña que los U-Nets de Stable Diffusion utilizados por los competidores (SDEdit, Z-STAR).
  • Ocupa significativamente menos memoria VRAM (aprox. 1.5 GB vs. >6 GB de los modelos basados en difusión), haciéndolo viable para despliegue en hardware con recursos limitados.
  • Aunque requiere múltiples pasos de integración (50 pasos), su tamaño compacto compensa la latencia en comparación con modelos de difusión masivos.

5. Significado e Impacto

• Generalización Unificada: MedShift demuestra que es posible aprender un espacio latente compartido que permite la traducción entre cualquier par de dominios vistos durante el entrenamiento, eliminando la necesidad de entrenar un modelo separado para cada combinación de dominio.
• Aplicabilidad Clínica: Al poder ajustar el equilibrio entre estructura y estilo en tiempo de inferencia, el modelo se adapta a necesidades clínicas específicas (ej. priorizar la anatomía para navegación quirúrgica o el realismo para entrenamiento).
• Escalabilidad: La combinación de alta fidelidad con un modelo pequeño y eficiente sugiere que la adaptación de dominio en imágenes médicas puede ser escalable y desplegable en entornos clínicos reales, no solo en servidores de investigación.

En resumen, MedShift representa un avance significativo hacia la creación de modelos robustos y generalizables para imágenes médicas, resolviendo el problema de la brecha de dominio mediante un enfoque de transporte óptimo implícito eficiente y flexible.