Resumen Técnico: CoCA – Recompensa a Nivel de Paso Gratuita en Modelos de Difusión T2I Ajustados mediante RL

1. Planteamiento del Problema

Los modelos de difusión de texto a imagen (T2I) de vanguardia dependen cada vez más del Aprendizaje por Refuerzo (RL) para el ajuste fino (fine-tuning) con el fin de alinearse con las preferencias humanas (p. ej., atractivo estético, corrección semántica). Sin embargo, estos métodos enfrentan una limitación crítica: la dispersión de la recompensa (reward sparsity).

En las formulaciones estándar de RL para la difusión, el proceso se modela como un Proceso de Decisión de Markov (MDP) donde el agente recibe una única recompensa terminal solo después del paso final de denoisado ($t=0$). Los métodos de política de gradiente existentes (p. ej., REINFORCE, DDPO) distribuyen esta recompensa escalar de manera uniforme a través de los $T$ pasos de denoisado. Esto crea un desajuste:

• Impacto Desigual: Los pasos temporales iniciales son decisivos para establecer la estructura global de la imagen, mientras que los pasos posteriores sirven principalmente para refinar detalles finos.
• Asignación de Crédito Ineficiente: La ponderación uniforme no logra identificar qué pasos específicos contribuyeron significativamente a la calidad final, lo que conduce a una eficiencia de muestreo subóptima y una precisión limitada al atribuir recompensas a las acciones generativas.
• Limitaciones de las Soluciones Existentes: Los intentos previos para abordar esto incluyen el descuento temporal fijo (que depende de suposiciones a priori que no logran generalizarse entre diversos prompts) o el entrenamiento de redes auxiliares (p. ej., críticos paso a paso) para predecir recompensas intermedias, lo que introduce una carga computacional significativa y requiere entrenamiento adicional.

2. Metodología: CoCA (Asignación de Crédito Basada en la Contribución)

Los autores proponen CoCA, un marco que redistribuye la recompensa terminal a lo largo de la trayectoria de denoisado para generar señales densas a nivel de paso sin entrenar redes auxiliares.

Mecanismo Central

CoCA estima la contribución de cada paso de denoisado midiendo la progresión latente hacia la imagen limpia realizada.

1. Proxy de Progreso Latente: En lugar de asumir una importancia fija, CoCA rastrea la similitud entre el estado latente intermedio ($x_{T-t}$) y el latente terminal realizado ($x_0$) en el espacio de embeddings.
   • Calcula la similitud de coseno: $Sim_t = \frac{\langle x_{T-t}, x_0 \rangle}{\|x_{T-t}\| \cdot \|x_0\|}$.
   • La contribución de un paso se deriva del cambio en esta similitud ($\Delta Sim_t$).
2. Suavizado por Ventanas (Window Smoothing): Para reducir las fluctuaciones en la similitud por paso, la trayectoria se segmenta en ventanas no superpuestas de tamaño fijo $W$. Se calcula la similitud promedio dentro de una ventana, y el puntaje de progreso es la diferencia entre ventanas consecutivas ($\Delta \bar{Sim}_i$).
3. Redistribución de la Recompensa: La recompensa terminal $r(x_0, c)$ se redistribuye a cada paso $t$ basándose en los puntajes de progreso normalizados:
   $$\hat{R}(s_t, a_t) = w_t \cdot r(x_0, c)$$
   donde $w_t$ es proporcional a la contribución estimada de ese paso.

Estabilización y Garantías Teóricas

• Normalización de la Recompensa: Para estabilizar el entrenamiento, CoCA aplica la normalización dos veces:
  1. Normalización por prompt antes de la redistribución para mantener la consistencia de ranking entre prompts.
  2. Normalización por prompt y por paso temporal después de la redistribución para asegurar escalas de recompensa consistentes a través de los pasos temporales y reducir la varianza.
• Consistencia del Gradiente de Política: Los autores demuestran que si las recompensas de paso redistribuidas suman la recompensa terminal original para cada trayectoria ($\sum \hat{R}_t = r(x_0, c)$), el gradiente de política esperado permanece sin sesgo. Por lo tanto, CoCA preserva los puntos estacionarios del objetivo original mientras proporciona señales de aprendizaje más densas.

3. Contribuciones Clave

1. Marco Novel: Introducción de CoCA, un mecanismo de asignación de crédito basado en la contribución que aborda el desajuste entre las actualizaciones de política iguales y la importancia variable de los pasos en modelos de difusión ajustados con RL.
2. Sin Entrenamiento Auxiliar: A diferencia de otros métodos de nivel de paso, CoCA no requiere redes adicionales (críticos) ni anotaciones de recompensa; opera puramente como un mecanismo de redistribución de recompensa condicionado en las trayectorias realizadas.
3. Observación Empírica: El artículo demuestra que los patrones de progreso latente a nivel de paso temporal varían sustancialmente entre diferentes prompts y trayectorias, lo que explica por qué los esquemas de ponderación temporal fijos fallan al generalizarse y motiva la necesidad de una asignación adaptativa.
4. Aplicabilidad General: El marco es agnóstico al algoritmo, integrándose con éxito tanto en los pipelines de optimización basados en REINFORCE (DDPO) como en los basados en GRPO (DanceGRPO).

4. Resultados Experimentales

Los experimentos se realizaron a través de cuatro funciones de recompensa de preferencia humana (Aesthetic, ImageReward, PickScore, HPSv2) y dos configuraciones de optimización (DDPO y DanceGRPO).

• Eficiencia de Muestreo: CoCA mejora la eficiencia de la optimización entre un 25% y un 100% en comparación con las líneas base de nivel de trayectoria (DDPO) y las líneas base existentes de nivel de paso (TDPO, UCA). Logra curvas de aprendizaje más pronunciadas y una convergencia más rápida.
• Rendimiento: En la configuración REINFORCE, CoCA alcanzó los puntajes más altos en todas las métricas, superando tanto a las líneas base de nivel de trayectoria como a las de nivel de paso. En la configuración GRPO, mostró resultados superiores, particularmente en escenarios de optimización de múltiples recompensas.
• Generalización: CoCA demostró una mejor generalización en prompts no vistos y funciones de recompensa no vistas (por ejemplo, mejorando PickScore cuando se entrena con HPSv2), lo que sugiere que las mejoras provienen de mejoras genuinas en la calidad de la imagen y no de un "hackeo" de la recompensa (reward hacking).
• Estudios de Ablación:
  • Tamaño de Ventana: Un tamaño de ventana moderado ($W=5$) produjo el mejor rendimiento, equilibrando estabilidad y granularidad.
  • Métrica de Distancia: La similitud de coseno superó a la distancia $\ell_2$ para medir el progreso latente.
  • Normalización: La combinación de la normalización previa y posterior a la redistribución proporcionó el entrenamiento más estable.

5. Significado y Reivindicaciones

El artículo afirma que CoCA ofrece una mejora "gratuita" en el ajuste fino por RL al resolver el problema de la asignación de crédito inherente a los entornos de recompensa dispersa. Su importancia radica en:

• Eficiencia: Reducción drástica del número de consultas de recompensa (y, por tanto, de cómputo/datos) requeridas para alinear modelos generativos con las preferencias humanas.
• Simplicidad: Proporcionar una solución general y "plug-and-play" que no aumenta la complejidad del modelo ni requiere entrenamiento auxiliar.
• Solidez Teórica: Asegurar que la redistribución de las recompensas no sesga el objetivo de optimización, manteniendo la integridad del gradiente de política original.

Los autores concluyen que, al proporcionar señales de aprendizaje más densas y alineadas con el resultado, CoCA reduce la barrera para aplicaciones beneficiosas en la creación de contenido y el diseño asistido, reconociendo al mismo tiempo que una mayor eficiencia también podría facilitar el uso indebido o amplificar los sesgos si no se gestiona cuidadosamente.