CrossFlow: One-Step Generation Across Latent and Pixel Spaces

Resumen Técnico: CrossFlow: Generación de un solo paso a través de espacios latentes y de píxeles

1. Planteamiento del problema

Los sistemas actuales de generación de imágenes de vanguardia, particularmente los Modelos de Difusión Latente (LDMs) y los métodos de concordancia de flujo (flow-matching), operan definiendo el prior, la trayectoria de probabilidad y el objetivo de predicción dentro de un único espacio de representación. Si bien los LDMs mejoran la eficiencia computacional al realizar la trayectoria generativa en un espacio latente de autoencoder comprimido, dependen de un proceso de dos etapas: un de-noising iterativo en el espacio latente seguido de un paso de decodificación separado para mapear los latentes de vuelta a los píxeles.

Esta arquitectura introduce un desajuste fundamental:

• Desajuste de optimización: El generador está optimizado para predecir desplazamientos latentes, mientras que la calidad final de la imagen depende de un decoder entrenado por separado.
• Desplazamiento de distribución: Durante la inferencia, el decoder recibe latentes generados que pueden diferir distribucionalmente de los latentes limpios producidos por el encoder durante el entrenamiento. Este desplazamiento puede amplificar los errores.
• Limitación de supervisión: El generador no puede ser entrenado directamente con pérdidas perceptuales o adversariales a nivel de imagen a menos que su salida pase a través del decoder separado, lo que debilita la supervisión directa en el espacio de píxeles.

El artículo plantea si el espacio de ruido y el espacio de salida deben ser idénticos, proponiendo un marco que preserva la eficiencia de las representaciones latentes mientras predice directamente la imagen misma.

2. Metodología: CrossFlow

CrossFlow es una formulación de flujo de espacio cruzado que mapea entradas latentes ruidosas directamente a imágenes en el espacio de píxeles en una sola evaluación de función. Resuelve el desajuste de espacio derivando objetivos donde la velocidad en el espacio latente se elimina del objetivo supervisado.

Formulación técnica central:

• Marco unificado: Los autores definen una función objetivo unificada $Z$ para paradigmas de generación de un solo paso. Los métodos estándar dependen de campos de velocidad en el espacio latente, lo que impide la generación de espacio cruzado.
• Criterios de diseño: Para permitir la supervisión de espacio cruzado, los autores establecen dos criterios:
  1. Eliminación de dependencias de la velocidad latente: El objetivo debe estar estructuralmente desacoplado de las velocidades en el espacio latente. La salida de la red no debe ser un vector de desplazamiento latente.
  2. Condiciones de contorno del punto final: El marco debe producir una imagen válida en el límite de destino ($r=0$).
• El objetivo CrossFlow: Al satisfacer estos criterios, los autores derivan una función de pérdida general (Ec. 10) donde la red $F_\theta$ predice una imagen en el espacio de píxeles $x_0$ en lugar de una velocidad latente.
  • La trayectoria latente (definida por $\alpha(t)$ y $\beta(t)$) sigue dictando cómo se corrompen los latentes limpios para crear las entradas.
  • El objetivo supervisado es la imagen de píxeles limpia.
  • El objetivo específico utilizado en los experimentos (Ec. 11) involucra un término de derivada temporal computado mediante el Producto Jacobiano-Vector (JVP) a lo largo de la trayectoria latente, asegurando que el modelo aprenda la dinámica de flujo correcta sin predecir velocidades latentes.

Detalles de implementación:

• Arquitectura: El modelo utiliza un codificador de imagen preentrenado y fijo (VA-VAE) para mapear imágenes a latentes y un generador de Transformer de Visión (ViT) ($F_\theta$) para mapear latentes ruidosos directamente a píxeles.
• Pipeline de entrenamiento: Se utiliza un enfoque de dos etapas para asegurar la estabilidad:
  1. Entrenar el codificador $E$ utilizando objetivos estándar de autoencoder (reconstrucción + adversaridad).
  2. Congelar $E$ y entrenar el generador $F_\theta$ desde cero utilizando el objetivo CrossFlow.
• Función de pérdida: El objetivo total combina la pérdida analítica de CrossFlow con pérdidas auxiliares en el espacio de píxeles:
  • Pérdida de reconstrucción L1: Alineación píxel a píxel.
  • Pérdida Perceptual: Similitud estructural.
  • Pérdida Adversarial (GAN): Fidelidad de alta frecuencia.
  • Nota: El factor de ponderación analítico del objetivo de flujo se omite durante el cómputo de la pérdida perceptual para evitar la inestabilidad del entrenamiento.

3. Contribuciones clave

1. Marco de espacio cruzado: Introducción de CrossFlow, un marco de generación de un solo paso donde el prior de entrada (latente) y la imagen de salida (píxel) residen en diferentes espacios de representación.
2. Objetivo de un solo paso libre de velocidad: Una derivación novedosa que permite que un modelo actúe tanto como un generador de un solo paso de latente-a-píxel como un reemplazo del decoder, eliminando la necesidad de un decoder separado en la inferencia.
3. Supervisión directa en el espacio de píxeles: La capacidad de entrenar el generador directamente con pérdidas de fidelidad perceptual y adversaridad sobre la imagen generada, saltándose la etapa intermedia del decoder.
4. Validación empírica: Demostración de que un solo modelo puede lograr puntuaciones FID competitivas en ImageNet-1k (256×256) con una sola evaluación de función (1-NFE) mientras supera a los métodos latentes de un solo paso estándar y iguala a las líneas base de difusión latente de múltiples pasos.

4. Resultados experimentales

Los experimentos se realizaron en ImageNet-1k a una resolución de 256×256 utilizando generación condicionada por clase.

• Rendimiento principal:

  • CrossFlow-XL (461M parámetros) logra un FID de 1.62 con 1-NFE.
  • Esto es competitivo con las líneas base de difusión latente de múltiples pasos como DiT-XL/2 (2.27 FID con 250 pasos) y SiT-XL/2 (2.06 FID con 250 pasos).
  • Supera significativamente a otros métodos latentes de un solo paso (por ejemplo, iMF-XL/2 con 1.72 FID) y evita la necesidad de un decoder separado.
  • CrossFlow-B (95M) y CrossFlow-L (315M) logran FIDs de 2.54 y 1.69, respectivamente, demostrando escalabilidad.

• Estudios de ablación:

  • Elección del codificador: El codificador VA-VAE (32 canales) produjo los mejores resultados (2.54 FID), superando a SD-VAE (4.08 FID) y FLUX.2 (3.81 FID), lo que sugiere que la estructura de la representación latente es crítica.
  • Componentes de la pérdida: Eliminar la pérdida perceptual causó una caída severa del FID a 30.21, y eliminar la pérdida GAN la redujo a 8.86. El objetivo de flujo de CrossFlow proporciona la señal estructural, mientras que las pérdidas perceptual y adversaridad son vitales para la fidelidad de alta frecuencia.
  • Reemplazo del Decoder: Cuando se utilizó como decoder para latentes generados por un LightningDiT-B/1, CrossFlow mejoró consistentemente el FID sobre el decoder VA-VAE estándar durante todo el entrenamiento.

5. Significado y afirmaciones

El artículo afirma que CrossFlow logra cerrar la brecha entre la eficiencia de las representaciones latentes y la fidelidad de la supervisión directa en el espacio de píxeles. Al unificar el generador y el decoder en un solo modelo entrenado con entradas latentes ruidosas y objetivos de píxeles, CrossFlow:

• Elimina la etapa explícita del decoder durante la inferencia.
• Expone al modelo durante el entrenamiento a los latentes ruidosos o imperfectos que ocurren durante la generación, en lugar de solo salidas de codificadores limpios.
• Demuestra que los objetivos de flujo de espacio cruzado pueden combinar los beneficios computacionales de la difusión latente con las mejoras de calidad de la supervisión adversaridad y perceptual directa a nivel de píxel.

Los autores mantienen un tono modesto, señalando que aunque la representación latente y las pérdidas auxiliares son importantes, la contribución central es el cambio arquitectónico y de objetivo que permite la generación de un solo paso en espacio cruzado sin requerir módulos de decoder separados.