Resumen Técnico: FullFlow

Enunciado del Problema

Los modelos modernos de difusión y flujo rectificado de texto a imagen (T2I) poseen priores visuales ricos y capacidades generativas sólidas, pero operan estrictamente de manera unidireccional: texto de entrada, imagen de salida. Aunque existen modelos unificados de visión y lenguaje (VLM), típicamente se construyen condicionando un modelo de lenguaje preentrenado con un codificador visual o realizando un preentrenamiento conjunto costoso desde cero. Estos enfoques a menudo descartan los fuertes priores de imagen ya codificados en las bases de T2I o requieren un reentrenamiento sustancial que es computacionalmente prohibitivo. La pregunta central abordada por este trabajo es si un modelo de flujo T2I preentrenado puede actualizarse a un generador bidireccional de visión y lenguaje sin un preentrenamiento multimodal completo ni sacrificar el prior de imagen original.

Metodología

Los autores proponen FullFlow, una receta eficiente en parámetros que actualiza un modelo T2I de flujo rectificado preentrenado (específicamente Stable Diffusion 3 y FLUX.1-dev) a un generador bidireccional. La innovación central radica en tratar las dos modalidades de manera asimétrica pero compatible dentro de una única base compartida.

1. Extensión Multimodal en Espacio Nativo

• Modalidad de Imagen: Las imágenes permanecen en su espacio latente continuo nativo de flujo rectificado. El modelo preserva el proceso de imagen continuo preentrenado.
• Modalidad de Texto: El texto se maneja mediante un proceso de inserción discreto basado en Flujos de Edición. A diferencia de la difusión discreta basada en máscaras (que introduce tokens fuera del vocabulario incompatibles con codificadores congelados), FullFlow utiliza un proceso hacia adelante basado en eliminación. Esto asegura que, en cualquier nivel de ruido, el estado de texto corrupto sea una subsecuencia de tokens válida que pueda decodificarse y re-tokenizarse para codificadores auxiliares (por ejemplo, CLIP).
• Arquitectura: Una única base compartida (MM-DiT) procesa ambas modalidades. El modelo se condiciona en dos tiempos desacoplados: $t$ para el tiempo de imagen y $\tau$ para el tiempo de texto.
  • Componentes Congelados: Los pesos de la base y todos los codificadores de texto auxiliares (T5, CLIP) permanecen congelados.
  • Componentes Entrenables: Solo se entrenan adaptadores LoRA (aplicados a las capas lineales del transformador), cabezas de texto ligeras (para conteos de inserción e identidades de tokens) y nuevas vías de condicionamiento de tiempo. Esto equivale a entrenar solo aproximadamente el 5% de los parámetros para SD3 y aproximadamente el 2.5% para FLUX.1.

2. Objetivo de Entrenamiento de Doble Tiempo

El objetivo de entrenamiento implica predecir conjuntamente la velocidad de la imagen y las inserciones de texto faltantes basándose en el par de tiempos $(t, \tau)$.

• Programa de Corrupción: Los autores comparan un programa "alternado-limpio" (entrenando solo en tareas condicionadas por el punto final) con un programa "corrupción-mixta" donde $(t, \tau) \sim \text{Unif}([0, 1]^2)$. Descubren que un enfoque híbrido: preentrenar con corrupción mixta para aprender la correspondencia cruzada de modalidades, seguido de un refinamiento con alternado-limpio, produce los mejores resultados.
• Estabilización:
  • Equilibrio Adaptativo de Gradientes: Para abordar la discrepancia de escala entre la pérdida de imagen basada en MSE y la pérdida de texto basada en entropía cruzada/recuento, los autores ajustan dinámicamente el peso de la pérdida de texto $\lambda_{txt}$ utilizando una media móvil exponencial (EMA) de la relación de gradientes.
  • Coincidencia con Profesor: Para prevenir la degradación del prior de imagen preentrenado durante el ajuste fino en nuevas distribuciones de datos, la regresión de velocidad de imagen se reemplaza con un objetivo de "coincidencia con profesor al mismo ruido". El profesor es la misma base con las capas LoRA deshabilitadas, lo que no requiere memoria adicional.

3. Modos de Inferencia

Desacoplar $t$ y $\tau$ transforma la inferencia en la selección de trayectorias en un espacio generativo 2D:

• Texto $\to$ Imagen: Fijar $\tau=1$ (texto limpio), evolucionar $t$ de 0 a 1.
• Imagen $\to$ Texto: Fijar $t=1$ (imagen limpia), evolucionar $\tau$ de 0 a 1.
• Muestreo Conjunto: Evolucionar simultáneamente tanto $t$ como $\tau$.
• Generación Parcial de Texto (VQA): Fijar la imagen y los tokens de la pregunta, restringiendo el proceso de inserción al segmento de la respuesta. Esto permite responder preguntas visuales sin un decodificador separado.

Contribuciones Clave

1. C1: Convertir un modelo de flujo rectificado T2I preentrenado en un generador bidireccional entrenando solo adaptadores LoRA y cabezas de texto ligeras, preservando con éxito el prior de imagen preentrenado.
2. C2: Introducir un marco de doble tiempo donde las imágenes permanecen en flujo continuo y el texto utiliza un proceso de inserción discreto, permitiendo la generación condicional, conjunta y de texto parcial como trayectorias en un espacio unificado $(t, \tau)$.
3. C3: Demostrar un rendimiento superior frente a una línea base Dual Diffusion-LoRA coincidente en SD3 y FLUX.1-dev tanto en la retención de texto a imagen como en la calidad de imagen a texto, mientras se soporta VQA aguas abajo con recursos computacionales mínimos.

Resultados Experimentales

Los autores evalúan FullFlow en SD3 y FLUX.1-dev utilizando un conteo de parámetros coincidente y rango LoRA frente a una línea base Dual Diffusion reimplementada.

• Eficiencia: FullFlow entrena en dos GPUs RTX A5000 en menos de 24 horas, reduciendo la VRAM pico de ~84 GB (Dual Diffusion) a ~38 GB y aumentando el rendimiento en ~8$\times$.
• Texto a Imagen (SD3): Con tiempo de reloj coincidente, FullFlow mejora el FID de 62.7 a 31.6 (y 29.3 con pasos coincidentes) y el CMMD de 0.870 a 0.121, demostrando una fuerte retención del prior de imagen.
• Imagen a Texto (SD3): FullFlow logra una puntuación CIDEr de 99.4 (vs. 2.0 para la línea base) con tiempo de reloj coincidente, y 13.4 con pasos coincidentes.
• Tareas Aguas Abajo: En los benchmarks de VQA (VQAv2, VizWiz, OKVQA), FullFlow (con solo 130M de parámetros entrenables) logra una precisión competitiva, superando notablemente a la línea base Dual Diffusion en VizWiz (50.7 vs 29.9).
• Transferibilidad: La receta se transfiere con éxito a FLUX.1-dev, logrando un CIDEr de 73.6 después de 36 horas de entrenamiento.

Significado y Afirmaciones

El artículo afirma que FullFlow demuestra que se pueden desbloquear fuertes capacidades bidireccionales de visión y lenguaje a partir de modelos de flujo T2I preentrenados sin necesidad de un preentrenamiento multimodal completo. Los resultados sugieren que los priores T2I preentrenados codifican una estructura visual-semántica más rica de lo que comúnmente se asume, capaz de soportar tareas inversas (imagen a texto) y generación conjunta mediante una adaptación ligera.

Los autores posicionan este trabajo como una "prueba de concepto factible en términos de cómputo" en lugar de una afirmación de rendimiento óptimo de extremo a extremo. Enfatizan que su método hace que los modelos de flujo rectificado bidireccionales sean accesibles para la comunidad de investigación más amplia al requerir solo GPUs comerciales y una pequeña fracción de los parámetros típicamente necesarios para modelos unificados. Reconocen limitaciones, incluido el enfoque actual en una resolución de 256$\times$256 y la falta de escalado a conjuntos de datos más grandes o ajuste de instrucciones para diálogos complejos, señalando estos como direcciones futuras.