SANA I2I: A Text Free Flow Matching Framework for Paired Image to Image Translation with a Case Study in Fetal MRI Artifact Reduction

Resumen Técnico: SANA-I2I

1. Problema y Contexto

El trabajo aborda el desafío de la traducción de imagen a imagen (I2I) supervisada en el contexto de la reducción de artefactos de movimiento en resonancia magnética (RM) fetal.

• Limitaciones de los enfoques actuales: Los modelos generativos de vanguardia, como SanaControlNet, dependen de condicionamiento textual (prompts) junto con imágenes de control. Sin embargo, en tareas de I2I como la eliminación de ruido o artefactos, el texto es innecesario, ambiguo o inexistente. La inclusión de condicionamiento textual puede introducir estocasticidad no deseada y desalineación con la tarea específica de transformación de imagen.
• Falta de datos reales emparejados: En la RM fetal, es imposible obtener pares reales de imágenes "limpias" y "con artefactos" del mismo feto en el mismo instante debido al movimiento. Esto dificulta el entrenamiento supervisado tradicional.

2. Metodología

Los autores proponen SANA-I2I, un marco de trabajo libre de texto basado en Flow Matching (Ajuste de Flujo) que extiende la familia de modelos SANA.

• Arquitectura y Eliminación de Texto:

  • Se elimina por completo el condicionamiento textual. El modelo se redefine para tomar como entrada el par $(x, t, y)$, donde $x$ es el ruido/estado latente, $t$ es el tiempo y $y$ es la imagen de entrada (fuente) que actúa como única señal de condicionamiento.
  • Se utiliza un Autoencoder de Compresión Profunda (DC-AE) para mapear las imágenes a un espacio latente compacto, permitiendo generación de alta resolución con menor costo computacional.
  • Se adapta la rama de control (ControlNet) para que la imagen de entrada se procese a través de la misma red de incrustación (PatchEmbed) que el ruido, unificando el espacio latente.

• Mecanismos de Entrenamiento y Guía:

  • Flow Matching: El modelo aprende un campo de velocidad condicional que mapea la distribución de ruido a la distribución de la imagen objetivo, guiado exclusivamente por la imagen fuente.
  • Classifier-Free Guidance (CFG): Se mantiene el mecanismo de CFG. Durante el entrenamiento, la variable de condicionamiento $y$ se "apaga" (se establece en cero) con una probabilidad $p_{drop}$, permitiendo que el modelo aprenda tanto ramas condicionadas como no condicionadas. Esto permite controlar la fuerza de la transformación en la inferencia sin reentrenar.
  • Variantes: Se exploran dos enfoques:
    1. Principal: Solo se apaga el condicionamiento $y$ (la señal de control permanece activa). Es el más efectivo para la corrección de artefactos.
    2. Bis (Exploratorio): Se apagan tanto $y$ como la señal de control, permitiendo generación no condicionada pura, aunque mostró resultados inferiores para la tarea específica de RM fetal.

• Generación de Datos Sintéticos:

  • Para superar la falta de datos reales emparejados, se utiliza una estrategia de generación sintética basada en el método de Duffy et al.
  • Se toman imágenes limpias (dominio B) y se les inyectan artificialmente artefactos de movimiento realistas (dominio C) en el espacio k. Esto crea pares $(C \to B)$ para el entrenamiento supervisado.
  • Se seleccionó una configuración de artefactos severa (SSIM entre 0.6 y 0.9) para asegurar que el modelo aprenda a corregir distorsiones significativas.

3. Contribuciones Clave

1. Extensión sin texto de SanaControlNet: Propone un marco específico para traducción de imagen a imagen supervisada que elimina la dependencia de prompts de texto, alineando mejor el modelo con tareas de transformación puramente visuales.
2. Eficiencia y Reducción de Pasos: Demuestra que la formulación basada en flujo (Flow Matching) permite obtener resultados de alta calidad con un número muy reducido de pasos de inferencia (tan solo 2-5 pasos), superando la ineficiencia de los modelos de difusión tradicionales.
3. Adaptabilidad a Imagen Médica: Valida el enfoque en un escenario clínico difícil (RM fetal), mostrando que es posible entrenar modelos supervisados robustos utilizando datos sintéticos generados de manera controlada.
4. Mecanismo de Control Flexible: Mantiene la capacidad de ajustar la agresividad de la corrección mediante la escala de guía (guidance scale) sin necesidad de reentrenamiento.

4. Resultados

El modelo se evaluó en la tarea de reducción de artefactos de movimiento en RM fetal, comparándolo con un método previo basado en GANs (trabajo anterior de los mismos autores) y utilizando métricas sintéticas y reales.

• Evaluación Cuantitativa (Datos Sintéticos):

  • En métricas de referencia completa (SSIM y MAE), el método anterior obtuvo puntuaciones ligeramente mejores.
  • Análisis Crítico: Los autores explican que esto se debe a que SANA-I2I corrige inconsistencias globales de luminancia (haciendo la imagen más brillante/realista), lo que penaliza las métricas pixel a pixel frente a la referencia original oscura. Sin embargo, visualmente, SANA-I2I elimina los artefactos residuales que el método anterior deja.

• Evaluación en Datos Reales (Sin Referencia):

  • Se utilizaron métricas basadas en distribución: FID (Fréchet Inception Distance) y KID (Kernel Inception Distance).
  • Resultado: SANA-I2I superó claramente al método anterior (FID: 51.74 vs 60.14). Esto indica que las imágenes generadas se alinean estadísticamente mejor con la distribución de imágenes fetales limpias y de alta calidad.

• Análisis Cualitativo:

  • SANA-I2I suprime los artefactos de movimiento de manera más efectiva, preservando los límites anatómicos y los detalles finos sin introducir estructuras visualmente implausibles.
  • El método anterior tendía a dejar patrones de artefactos residuales y a veces causaba un suavizado excesivo.

• Estudio de Ablación:

  • Pasos de Inferencia: Se encontró que 5 pasos (o incluso 2) son suficientes para lograr una reconstrucción de alta calidad.
  • Escala de Guía: Un valor entre 1.0 y 1.1 ofrece el mejor equilibrio entre la eliminación de artefactos y la preservación de detalles anatómicos. Valores más altos eliminan artefactos pero borran detalles diagnósticos.

5. Significado e Impacto

El trabajo demuestra que los modelos generativos basados en flujo pueden adaptarse exitosamente a tareas de traducción de imagen a imagen en medicina sin depender de condicionamiento textual.

• Eficiencia Computacional: La capacidad de generar resultados de alta fidelidad en pocos pasos hace que el modelo sea viable para aplicaciones clínicas donde el tiempo de inferencia es crítico.
• Generalización: Aunque se prueba en RM fetal, el marco es general y aplicable a cualquier tarea de traducción de imagen emparejada (denoising, super-resolución, cambio de modalidad).
• Validación de Datos Sintéticos: Confirma que el entrenamiento con datos sintéticos generados mediante simulación de artefactos (Duffy) es una estrategia robusta para resolver la escasez de datos emparejados en medicina, permitiendo el aprendizaje supervisado donde antes solo era posible el no supervisado o semi-supervisado.

En conclusión, SANA-I2I representa un avance significativo hacia modelos de generación médica más eficientes, precisos y libres de la complejidad innecesaria del condicionamiento por texto.