FOCUS: Optimal Control for Multi-Entity World Modeling in Text-to-Image Generation

Resumen Técnico: FOCUS

1. El Problema

Los modelos de generación de imágenes a partir de texto (T2I) modernos, basados en Flow Matching (FM) como Stable Diffusion 3.5 y FLUX.1, han demostrado un alto rendimiento en prompts de una sola entidad. Sin embargo, sufren fallos críticos al generar escenas con múltiples sujetos:

• Fuga de atributos (Attribute Leakage): Los atributos de un objeto (ej. color, forma) se transfieren incorrectamente a otro.
• Entrelazamiento de identidad (Identity Entanglement): Los objetos se fusionan o pierden su identidad individual.
• Omisión de sujetos: El modelo ignora uno o más elementos del prompt.

Estos errores se deben a representaciones débiles centradas en objetos dentro del modelo, donde las restricciones de "qué objeto tiene qué atributo" no son locales. Los métodos existentes para mitigar esto son mayoritariamente heurísticos (basados en reglas de atención), carecen de un objetivo de optimización unificado y a menudo no son portables entre diferentes arquitecturas de modelos.

2. Metodología

Los autores proponen un marco teórico que reformula la generación de imágenes como un problema de Control Óptimo Estocástico (SOC).

• Marco Teórico (Flow Matching como SOC):
  Se formula la generación como un proceso de control donde se introduce una pequeña perturbación de control $u$ en la dinámica del flujo (drift) para reducir el entrelazamiento de sujetos mientras se mantiene la fidelidad al modelo base. El objetivo es minimizar una función de costo que penaliza la dispersión de los sujetos y recompensa su separación.

• FOCUS (Flow Optimal Control for Unentangled Subjects):
  Se introduce una nueva función de pérdida probabilística basada en la Divergencia de Jensen-Shannon (JSD). A diferencia de métodos anteriores que usan heurísticas simples en los mapas de atención, FOCUS trata los mapas de atención cruzada como distribuciones de probabilidad espaciales.

  • Objetivo: Minimizar la dispersión dentro de los mapas de atención de un mismo sujeto (consistencia de unión) y maximizar la divergencia entre los mapas de atención de diferentes sujetos (separación).
  • La pérdida se normaliza para ser comparable independientemente del número de sujetos.

• Dos Implementaciones Prácticas:

  1. Controlador en Tiempo de Inferencia (Test-Time Control): Un método sin entrenamiento que perturba la velocidad base del modelo en tiempo real. Utiliza una aproximación de "paso único" para estimar el gradiente de control basándose en el estado actual, evitando la necesidad de resolver ecuaciones hacia atrás complejas.
  2. Ajuste Fino (Fine-Tuning) con Adjoint Matching: Un método para aprender una red de control ligera que "amortiza" la señal de guía. Utiliza la técnica de Adjoint Matching para entrenar una red que regresa a una señal de adjunto simplificada calculada sobre trayectorias congeladas, permitiendo inferencia rápida sin sobrecarga computacional adicional.

3. Contribuciones Clave

1. FOCUS: Un objetivo de pérdida basado en divergencia probabilística diseñado específicamente para el desentrelazamiento de múltiples sujetos, compatible con cualquier arquitectura de T2I moderna.
2. Dos Instantiaciones:
   • Un controlador sin entrenamiento que funciona en tiempo de prueba.
   • Un procedimiento de ajuste fino eficiente que aprende un controlador amortizado, generalizando a prompts no vistos.
3. Perspectiva de Control Óptimo: Proporcionan una interpretación unificada de las heurísticas de atención anteriores (como Attend&Excite o Divide&Bind) como casos particulares de control óptimo, ofreciendo un "knob" (parámetro) explícito para el equilibrio entre fidelidad y desentrelazamiento.
4. Portabilidad: El método es agnóstico a la arquitectura y funciona tanto en modelos basados en Flow Matching (SD 3.5, FLUX.1) como en modelos de difusión clásicos (SD 1.5).

4. Resultados Experimentales

Los experimentos se realizaron en Stable Diffusion 3.5 y FLUX.1, utilizando un corpus de 150 prompts con 2-4 sujetos.

• Métricas Automáticas:

  • Tanto el control en tiempo de prueba como el ajuste fino mejoraron consistentemente la alineación imagen-texto y la fidelidad del texto en comparación con el modelo base.
  • FOCUS obtuvo los puntajes compuestos más altos en la mayoría de las métricas (CLIP, SigLIP, BLIP, PickScore, ImageReward) en comparación con heurísticas existentes como Attend&Excite, CONFORM y Divide&Bind.
  • El ajuste fino superó al control en tiempo de prueba en términos de puntaje compuesto, logrando mejoras sin sacrificar el estilo base ni la diversidad.

• Estudio Humano:

  • En una evaluación de preferencia pareada con 50 participantes, FOCUS obtuvo las tasas de victoria y las calificaciones Elo más altas tanto para SD 3.5 como para FLUX.1.
  • Los evaluadores humanos prefirieron consistentemente las imágenes generadas con FOCUS por una mejor separación de sujetos y una unión correcta de atributos.

• Eficiencia:

  • El control en tiempo de prueba añade una latencia moderada (~90-138% más de tiempo de inferencia) pero funciona en GPUs comerciales.
  • Los modelos ajustados finamente no tienen sobrecarga en tiempo de inferencia y generalizan bien a prompts no vistos durante el entrenamiento.

5. Significado e Impacto

El trabajo FOCUS representa un avance significativo en la generación de imágenes multi-objeto al pasar de soluciones heurísticas ad-hoc a un marco teórico riguroso basado en el control óptimo.

• Unificación: Ofrece un lenguaje común para entender y mejorar la fidelidad de múltiples sujetos en diversas arquitecturas de modelos generativos.
• Practicidad: Proporciona soluciones tanto para usuarios que no pueden reentrenar modelos (control en tiempo de prueba) como para desarrolladores que buscan modelos optimizados (ajuste fino ligero).
• Calidad: Resuelve problemas fundamentales de "mundo" en la generación de imágenes, permitiendo aplicaciones más complejas como ilustración de historias, composiciones de múltiples paneles y comunicación científica, donde la identidad y la unión de atributos son críticas.

En resumen, FOCUS demuestra que tratar la generación de imágenes como un problema de control óptimo permite un manejo superior de la complejidad de escenas con múltiples entidades, superando a los métodos actuales en fidelidad, consistencia y eficiencia.