ImageRAGTurbo: Towards One-step Text-to-Image Generation with Retrieval-Augmented Diffusion Models

Resumen Técnico: ImageRAGTurbo

1. El Problema

Los modelos de difusión han establecido el estado del arte en la generación de imágenes texto-a-imagen. Sin embargo, enfrentan dos limitaciones críticas:

• Alta Latencia: El proceso de muestreo iterativo (típicamente 25-50 pasos) genera una latencia significativa, lo que limita su uso en aplicaciones en tiempo real.
• Compromiso Calidad-Velocidad en Modelos de Pocos Pasos: Los modelos recientes de "pocos pasos" (1-4 pasos) o "un paso" (one-step), como los modelos de difusión adversarialmente destilados, reducen drásticamente el tiempo de inferencia pero a menudo sacrifican la fidelidad del prompt (la capacidad de la imagen generada para reflejar con precisión los conceptos visuales y las relaciones espaciales del texto) y la calidad de la imagen. Además, el entrenamiento de estos modelos destilados suele ser computacionalmente costoso.

Las soluciones actuales de generación aumentada por recuperación (RAG) en imágenes no están optimizadas para modelos de pocos pasos, o introducen una sobrecarga computacional excesiva (por ejemplo, mediante el uso repetido de Modelos de Lenguaje Multimodal).

2. Metodología: ImageRAGTurbo

El artículo propone ImageRAGTurbo, un marco novedoso que combina la generación de difusión de pocos pasos con la Recuperación Aumentada por Generación (RAG). La intuición central es que inyectar información recuperada semánticamente relevante simplifica la tarea del modelo de mapear ruido aleatorio a la distribución objetivo.

El enfoque se divide en tres componentes principales:

• Recuperación de Pares Texto-Imagen: Dado un prompt de entrada ($p_{tgt}$), el sistema consulta una base de datos (ej. OpenImage) para recuperar pares texto-imagen relevantes ($p_{retr}, x_{retr}$) utilizando un codificador de texto preentrenado (CLIP) y búsqueda ScaNN.
• Inyección Directa en el Espacio H (Investigación Preliminar):
  • Se identificó que el espacio de características de la capa más profunda del UNet (denominado H-space) contiene representaciones semánticas de alto nivel.
  • En una fase inicial "sin entrenamiento", los embeddings recuperados se proyectan al espacio H del denoiser y se mezclan con las características del prompt objetivo mediante interpolación esférica normalizada.
  • Resultado: Esto mejora la fidelidad, pero requiere una búsqueda exhaustiva de hiperparámetros (fuerza de mezcla) en tiempo de inferencia, lo cual es impráctico.
• Sintonización Eficiente con Adaptador H (H-Adapter):
  • Para evitar la búsqueda de hiperparámetros, se introduce un adaptador entrenable en el espacio H del denoiser.
  • Este adaptador utiliza un mecanismo de atención cruzada (cross-attention) para fusionar automáticamente las características recuperadas ($h_{retr}$) con las características del objetivo ($h_{tgt}$).
  • La fórmula de fusión es: $h'_{t} = h_{tgt} + \lambda \cdot g_{\phi}(h_{tgt}, h_{retr})$, donde $\lambda$ se basa en la similitud coseno entre los embeddings de texto.
• Procedimiento de Entrenamiento:
  • Se utiliza entrenamiento adversarial latente (Latent Adversarial Training). El modelo estudiante (pocos pasos) intenta engañar a un discriminador (basado en el codificador congelado del modelo maestro) para que sus latentes parezcan reales.
  • Se añaden pérdidas auxiliares: Score Distillation Loss y Latent LPIPS Loss para estabilizar el entrenamiento y mejorar la calidad visual.
  • Solo se ajustan los parámetros del adaptador y el decodificador del estudiante (mediante LoRA), manteniendo congelado el resto del denoiser, lo que hace el proceso eficiente.

3. Contribuciones Clave

1. Nuevo Marco RAG para Pocos Pasos: Introducen el primer enfoque que integra la recuperación de información en modelos de difusión de un solo paso, mejorando significativamente la alineación texto-imagen sin sacrificar la velocidad.
2. Adaptador H-Eficiente: Desarrollan un método de ajuste fino (finetuning) ligero que utiliza un adaptador en el espacio H con atención cruzada, evitando la necesidad de búsquedas exhaustivas de hiperparámetros en tiempo de ejecución.
3. Rendimiento Superior: Demuestran experimentalmente que su método supera a los enfoques de pocos pasos existentes (como Latent Consistency Models y Stable Diffusion Turbo) en calidad de imagen y fidelidad del prompt, manteniendo una velocidad de inferencia comparable.

4. Resultados Experimentales

Los resultados se evaluaron en los benchmarks MS-COCO y TIFA:

• Calidad y Fidelidad (TIFA y CLIP):
  • ImageRAGTurbo (1 paso) logra un puntaje TIFA de 0.801, superando al Stable Diffusion Turbo (sin RAG, 1 paso) que obtiene 0.779, y acercándose notablemente al modelo maestro de 50 pasos (0.811).
  • En la métrica CLIP (alineación texto-imagen), supera al Stable Diffusion Turbo en un 1.37% y al modelo de consistencia latente (4 pasos) en un 3.7%.
  • En categorías específicas como "objeto", "actividad" y "material", el modelo de 1 paso con RAG iguala o supera ligeramente al modelo de 50 pasos.
• Calidad Visual (FID y Aesthetics):
  • Logra un FID de 25.59 en MS-COCO, comparable al modelo de 50 pasos (25.33) y superior al Turbo base (26.04).
  • En el benchmark TIFA, la puntuación estética (AES) es de 5.88, superando a todos los modelos de pocos pasos y al modelo de 50 pasos (6.04 es el máximo, pero ImageRAGTurbo es muy cercano).
• Eficiencia:
  • El tiempo de inferencia total es de 116.7 ms por imagen, casi idéntico al Stable Diffusion Turbo (113.8 ms).
  • La sobrecarga introducida por el mecanismo de recuperación y el adaptador es mínima (aprox. 2.5% de latencia adicional), mientras que la ganancia en fidelidad (TIFA) es del 3%.
  • Es aproximadamente 25 veces más rápido que el modelo maestro de 50 pasos, con una caída mínima en la calidad.

5. Significado e Impacto

El trabajo de ImageRAGTurbo es significativo porque resuelve el dilema fundamental de la generación de imágenes rápida: velocidad vs. precisión.

• Viabilidad en Tiempo Real: Al lograr una generación de alta fidelidad en un solo paso, hace viable el uso de modelos de difusión en aplicaciones interactivas y en tiempo real donde la latencia es crítica.
• Eficiencia Computacional: El método de ajuste fino con adaptadores ligeros reduce la barrera de entrada para entrenar modelos de pocos pasos, evitando el costo masivo de la destilación adversarial tradicional.
• Nueva Dirección: Establece un precedente para el uso de bases de datos de recuperación en modelos de difusión de pocos pasos, sugiriendo que la información contextual externa puede compensar la falta de iteraciones internas del modelo.

En conclusión, ImageRAGTurbo demuestra que no es necesario sacrificar la calidad del prompt para obtener velocidad, siempre que se utilice la información recuperada de manera inteligente dentro de la arquitectura del modelo.