Resumen Técnico: Diff-Instruct con Recompensa Difusa (Didr)

Enunciado del Problema

Los avances recientes en generadores de texto a imagen (T2I) de un solo paso han permitido la síntesis en tiempo real al mapear el ruido latente directamente a imágenes en una sola pasada hacia adelante. Sin embargo, alinear estos generadores con las preferencias humanas mediante Aprendizaje por Refuerzo (RL) presenta un desafío estructural fundamental.

Los métodos anteriores para RLHF de un solo paso (por ejemplo, Diff-Instruct*, Diff-Instruct++) combinan la optimización de recompensas en el espacio de imágenes con la coincidencia de distribuciones en el espacio latente. Estos enfoques suelen aplicar señales de preferencia (recompensas) exclusivamente en el punto final de imagen limpia ($x_0$), mientras imponen una regularización Kullback-Leibler (KL) contra la distribución de referencia a lo largo de toda la trayectoria de difusión ruidosa. Los autores identifican una disparidad estructural crítica en este paradigma: la recompensa actúa solo en el estado terminal, mientras que el regularizador abarca todos los niveles de ruido.

Esta disparidad conduce a un modo de fallo denominado dominación de recompensa terminal. Dado que el ruido hacia adelante enmascara naturalmente los detalles de la imagen, la penalización de regularización KL contra la divergencia de modos se debilita significativamente en niveles altos de ruido. En consecuencia, los optimizadores explotan esta vulnerabilidad de la trayectoria, abandonando la distribución de referencia para colapsar sobre regiones de alta recompensa sin incurrir en penalizaciones KL suficientes. Esto resulta en puntuaciones de recompensa mejoradas, pero a menudo a expensas de la fidelidad de la imagen y la estabilidad distribucional.

Metodología: Diff-Instruct con Recompensa Difusa (Didr)

Para resolver esta desalineación, los autores proponen Didr, un marco de alineación a nivel de trayectoria libre de datos derivado de la minimización de la KL Integral (IKL). La idea central es propagar sistemáticamente la distribución de imagen limpia inclinada por la recompensa óptima de RLHF a través de todos los niveles de ruido de la trayectoria de difusión.

1. El Objetivo de Trayectoria Fundamentado

El método comienza estableciendo el objetivo de alineación correcto para imágenes limpias bajo RLHF regularizado por KL. La distribución de imagen limpia óptima $q^*(x_0|c)$ es la densidad inclinada por la recompensa:
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
donde $q_0$ es la distribución de referencia, $r$ es la función de recompensa y $\tau$ es un parámetro de temperatura.

En lugar de aplicar recompensas solo en $x_0$, Didr "eleva" este objetivo óptimo a toda la trayectoria difuminando $q^*$ a través del proceso hacia adelante de referencia. Esto produce una familia de marginales inclinadas por recompensa $\{q^*_t\}_{t \in [0, T]}$ definidas en cada nivel de ruido:
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
Los autores demuestran que minimizar la divergencia KL Integral entre las marginales del generador y estas marginales inclinadas por recompensa comparte el mismo minimizador que el objetivo estándar de RLHF regularizado por KL. Esto asegura que la recompensa y la regularización estén equilibradas en cada nivel de ruido, previniendo la dominación de recompensa terminal.

2. La Puntuación de Recompensa Difusa (DRS)

Mediante la diferenciación del objetivo IKL, los autores derivan la función de puntuación objetivo. La puntuación óptima se descompone en dos componentes:
$$s_{target} = s_{ref} + s_r$$
donde $s_{ref}$ es la función de puntuación de referencia y $s_r$ es la Puntuación de Recompensa Difusa (DRS). La DRS actúa como un término de corrección impulsado por recompensa definido en cada nivel de ruido:
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
A diferencia de la guía por clasificador (que es una aproximación de primer orden), la DRS retiene la inclinación log-exponencial completa, propagando información de preferencia solo donde la imagen limpia es identificable desde la observación ruidosa.

3. El Proxy de Recompensa Difusa (DRP)

Dado que la DRS requiere una expectativa posterior intratable $q(x_0|x_t, c)$, los autores introducen el Proxy de Recompensa Difusa (DRP) para hacer práctico el método. El DRP aproxima la expectativa posterior utilizando cadenas de eliminación de ruido de pasos cortos diferenciables desde el modelo de referencia congelado.

• Mecanismo: Desde un estado ruidoso $x_t$, el método ejecuta $K$ cadenas independientes de eliminación de ruido de $S$ pasos utilizando el modelo de referencia congelado para generar muestras limpias aproximadas $\{\hat{x}^{(k)}_0\}$.
• Estimación de Gradientes: La DRS se aproxima como un promedio ponderado por softmax de gradientes de recompensa por trayectoria:
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  donde los pesos $\omega^{(k)}$ son proporcionales a $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$.
• Eficiencia: Este enfoque permite que los gradientes de recompensa se propaguen de manera estable hacia atrás a través de la trayectoria sin almacenar activaciones del modelo de referencia ni requerir datos de entrenamiento de imágenes.

4. Marco de Entrenamiento

El proceso de entrenamiento alterna entre dos etapas:

1. Actualización del Asistente de Enseñanza (TA): Un modelo de puntuación TA ($s_\psi$) se entrena mediante Coincidencia de Puntuación de Eliminación de Ruido (DSM) para rastrear las marginales evolutivas del generador, suplantando la puntuación implícita del generador.
2. Actualización del Generador: El generador ($g_\theta$) se actualiza utilizando el gradiente del objetivo IKL, que compara la puntuación TA contra la suma de la puntuación de referencia corregida por CFG y la estimación DRP.

Contribuciones Clave

1. Identificación de la Dominación de Recompensa Terminal: El artículo identifica y analiza formalmente un modo de fallo estructural en métodos anteriores de RLHF de un solo paso donde los optimizadores explotan niveles de ruido para colapsar sobre regiones de alta recompensa, violando el equilibrio RLHF previsto.
2. Alineación a Nivel de Trayectoria Fundamentada: Los autores proponen Didr, un marco que deriva un objetivo a nivel de trayectoria matemáticamente sólido al propagar hacia adelante la densidad de imagen limpia óptima de RLHF. Esto asegura que la recompensa y la regularización estén equilibradas en cada nivel de ruido.
3. Derivación Analítica de la DRS: El artículo deriva analíticamente la Puntuación de Recompensa Difusa, mostrando que descompone exactamente el objetivo de trayectoria en una puntuación de referencia más una corrección inducida por recompensa.
4. Estimador Práctico (DRP): Los autores introducen el Proxy de Recompensa Difusa, un estimador libre de datos que calcula la DRS mediante eliminación de ruido de pasos cortos diferenciables, permitiendo un entrenamiento estable sin datos de imágenes.

Resultados Empíricos

Experimentos extensivos demuestran que Didr supera consistentemente a las líneas base existentes de un solo paso:

• Dominancia de Pareto: En la arquitectura base SDXL (1024×1024), la variante estándar de Didr domina en el sentido de Pareto a todas las líneas base de alineación de un solo paso existentes en el compromiso PickScore–FID, superando simultáneamente a modelos de múltiples pasos como FLUX-dev y SDXL-DPO en ambos ejes.
• Métricas de Preferencia: La variante de entrenamiento extendido (Didr_longer) logra la puntuación PickScore más alta (23.9) entre los modelos de un solo paso evaluados y mejora ImageReward en un 8,9% sobre la línea base anterior más fuerte (Diff-Instruct*). También lidera en HPSv2.1 y Puntuación Estética.
• Generalización: Didr se transfiere con éxito a la arquitectura base Z-Image de 6 mil millones de parámetros. En un solo paso de generación, supera a su modelo maestro de 50 pasos en métricas de preferencia, demostrando escalabilidad a través de arquitecturas.
• Mejoras Cualitativas: Las comparaciones visuales muestran que Didr produce detalles finos más nítidos y colores más naturales en comparación con métodos anteriores de un solo paso, evitando la sobresaturación y los artefactos pictóricos que a menudo se observan en las líneas base optimizadas por recompensa.

Significado y Afirmaciones

El artículo afirma que la propagación de recompensas a nivel de trayectoria fundamentada es una estrategia de alineación efectiva que resuelve la desalineación inherente de los métodos de recompensa solo en el punto final. Al propagar rigurosamente la distribución óptima de RLHF a través del proceso de difusión, Didr logra el mejor compromiso entre la alineación de preferencias y la fidelidad de la imagen entre los generadores de un solo paso.

Los autores enfatizan que su método requiere ningún dato de entrenamiento de imágenes, confiando únicamente en indicaciones de texto y un modelo de referencia congelado. Posicionan a Didr como una solución que permite a los generadores de un solo paso igualar o superar la alineación de preferencias de los maestros de múltiples pasos mientras mantienen velocidades de inferencia en tiempo real. El trabajo sugiere que la disparidad estructural identificada en métodos anteriores es una limitación fundamental que solo puede resolverse alineando toda la trayectoria, no solo el estado terminal.