Resumen Técnico: STAR (Asignación de Recompensa Adaptativa Espacio-Temporal)

Planteamiento del Problema

Los métodos existentes de post-entrenamiento de Aprendizaje por Refuerzo (RL) para la generación de texto a imagen (T2I), como Flow-GRPO, típicamente convierten la recompensa final a nivel de imagen en una única ventaja escalar. Esta escalar se aplica luego con una fuerza uniforme a lo largo de toda la trayectoria generativa (todos los pasos de eliminación de ruido y todos los tokens latentes espaciales).

El artículo identifica un desajuste de granularidad fundamental en este enfoque:

1. Estructura Temporal: Diferentes pasos de eliminación de ruido son responsables de distintas etapas de generación (por ejemplo, diseño global frente a refinamiento de detalles).
2. Estructura Espacial: El contenido que determina la alineación con el texto (por ejemplo, objetos específicos, atributos o relaciones espaciales) a menudo aparece solo en regiones específicas de la imagen, mientras que otras regiones (por ejemplo, el fondo) contribuyen menos a los requisitos semánticos del prompt.

Cuando una recompensa escalar se aplica uniformemente, la actualización de la política tiene dificultades para distinguir entre las regiones críticas para satisfacer el prompt y aquellas que no lo son. Esto dificulta que el modelo se concentre en el aprendizaje de los componentes generativos específicos que realmente impulsan la recompasa, lo que conduce a un rendimiento subóptimo en la generalización compositiva, el conteo de objetos y la renderización de texto.

Metodología: STAR

Los autores proponen la Asignación de Recompensa Adaptativa Espacio-Temporal (STAR), un marco diseñado para reasignar recompensas a nivel de imagen a través de la trayectoria de eliminación de ruido estructurada de los modelos de difusión y flujo. STAR opera sin modificar la fuente de recompensa externa, aprovechando en su lugar la mecánica interna del modelo generativo.

1. Asignación Guiada por Atención

STAR construye mapas de asignación espacial dinámica que varían a través de los pasos de eliminación de ruido ($t$) y los rollouts ($i$). El proceso implica:

• Extracción de Unidades de Texto: El prompt $c$ se descompone en un conjunto de unidades de texto centrales $K(c)$ (por ejemplo, sujetos, estilos, atributos).
• Mapeo de Atención: El modelo utiliza las capas de atención texto-imagen (específicamente desde el backbone de atención conjunta de Stable Diffusion 3.5) para extraer señales de enrutamiento. Para cada posición espacial latente $u$ y cada token de texto $j$, se computa un peso de atención $\alpha(u, j)$.
• Construcción de Mapas de Calor: Para cada unidad de texto, se genera un mapa espacial agregando los pesos de atención sobre los tokens de texto correspondientes. Estos mapas se agregan en un mapa de calor específico de la trayectoria $H_t^i$ que resalta las regiones latentes más activas con respecto a los componentes clave del prompt.
• Normalización: El mapa de calor se normaliza y escala para producir un mapa de coeficientes $C_t^i(u)$. Este mapa asigna pesos más altos a las regiones relevantes del prompt y pesos más bajos (pero no nulos) a las regiones del fondo.

2. Objetivo de Política con Resolución Espacial

STAR convierte la ventaja escalar relativa al grupo $A_i$ en una señal de entrenamiento con resolución espacial $A_t^i(u) = A_i \cdot C_t^i(u)$.

• El objetivo de la política se modifica para operar sobre contribuciones de probabilidad logarítmica espacial en lugar de una probabilidad de trayectoria escalar.
• Se computa una razón espacial normalizada $\rho$ para cada posición latente.
• La función de objetivo final (basada en el recorte de tipo PPO) suma las razones recortadas ponderadas por el mapa de ventaja espacial $A_t^i(u)$ a través de todos los pasos de tiempo y ubicaciones espaciales. Esto asegura que las actualizaciones de la política sean más fuertes en las regiones que corresponden a los componentes clave del prompt.

Contribuciones Clave

1. Marco STAR: Un método para el post-entrenamiento de RL en modelos de difusión/flujo T2I que convierte las recompensas a nivel de imagen en señales de entrenamiento distribuidas sobre el tiempo de eliminación de ruido y el espacio latente, permitiendo que las actualizaciones se centren en las regiones generativas correspondientes a objetos, atributos y relaciones clave.
2. Interfaz Unificada: El uso de interacciones internas texto-imagen para construir mapas de asignación espacial específicos de cada paso de tiempo, extendiendo los objetivos de RL existentes hacia un objetivo de política con resolución espacial sin requerir nuevos modelos de recompensa.
3. Validación Empírica: Demostración de que STAR mejora la alineación semántica compositiva, la renderización de texto y la optimización de preferencias a través de tres tareas distintas (GenEval, OCR, PickScore) utilizando Stable Diffusion 3.5 Medium como base.

Resultados Experimentales

Los autores evaluaron STAR en Stable Diffusion 3.5 Medium (SD3.5-M) comparándolo contra la línea base Flow-GRPO y otros modelos de vanguardia.

• GenEval (Generación Compositiva): STAR logró una puntuación general de 0.9759, superando a Flow-GRPO (0.95) y al SD3.5-M base (0.63). Se observaron mejoras notables en conteo (0.98), colores (0.96) y vinculación de atributos (0.94).
• Renderización de Texto OCR: STAR logró una precisión de 0.9757, superando a la línea base Flow-GRPO (0.92).
• PickScore (Optimización de Preferencias): STAR logró una puntuación de 23.60, comparado con 23.31 para Flow-GRPO.
• Eficiencia: El método introduce una sobrecarga computacional mínima, añadiendo menos del 2% al tiempo de entrenamiento y aumentando el uso de memoria pico en solo un 1.56%, ya que reutiliza los mapas de atención del pase hacia adelante (forward pass).

Significado y Reivindicaciones

El artículo afirma que STAR aborda el problema de la "asignación de crédito" en el RL de T2I al reconocer la naturaleza espacio-temporal de la generación por difusión. Al enrutar la misma señal de recompensa escalar hacia las regiones latentes específicas que determinan la alineación con el texto, STAR permite actualizaciones de política más selectivas y efectivas.

Los autores enfatizan que STAR logra estas ganancias sin cambiar la fuente de recompensa externa ni requerir modelos de recompensa adicionales. Proporciona una "interfaz simple, ligera y efectiva" para asignar el feedback a nivel de imagen a trayectorias estructuradas. Los resultados sugieren que la asignación adaptativa espacio-temporal es un factor crítico para optimizar la comprensión semántica fina, el control de contenido local y la alineación de preferencias en modelos generativos.