Beyond Pairwise Preferences: Listwise Reward-Aware Alignment for Diffusion Models

Resumen Técnico: Diffusion LAIR

Enunciado del Problema
Los métodos actuales de optimización de preferencias para alinear modelos de difusión de texto a imagen, como la Optimización Directa de Preferencias (DPO) y sus variantes específicas para difusión (por ejemplo, Diffusion-DPO, DSPO), dependen en gran medida de la supervisión por pares. Estos métodos reducen los datos de entrenamiento a comparaciones binarias de "ganador-perdedor" para un prompt dado. Los autores argumentan que esta reducción es limitante por dos razones principales:

1. Estructura de los Datos: Los conjuntos de datos de preferencias prácticos (por ejemplo, Pick-a-Pic) a menudo contienen múltiples imágenes candidatas para un solo prompt, creando una estructura natural de lista que los métodos por pares ignoran.
2. Pérdida de Información: Los modelos de recompensa modernos proporcionan puntuaciones continuas. Reducir estas a una única etiqueta binaria descarta la calidad relativa de los candidatos no seleccionados e ignora la magnitud de las brechas de recompensa entre las muestras.
   Además, los enfoques existentes de aprendizaje por refuerzo (RL) en línea para la alineación de difusión son computacionalmente costosos, difíciles de paralelizar y sensibles a los cambios de distribución debido a su naturaleza en línea. Si bien los métodos fuera de línea son más eficientes, no han aprovechado plenamente la señal más rica disponible en datos de lista con puntuaciones de recompensa.

Metodología: Diffusion LAIR
El artículo propone Diffusion LAIR (Recompensa Implícita Ponderada por Ventaja de Lista), un método de optimización de preferencias fuera de línea que trata la alineación como un problema de lista consciente de la recompensa.

• Representación de Datos: En lugar de pares independientes, el método procesa un prompt $c$ con un conjunto de $N_c$ imágenes candidatas $\{x^{(i)}_0\}$ y sus puntuaciones de recompensa continuas correspondientes $\{r_i\}$.
• Ponderación por Ventaja: Las puntuaciones de recompensa se convierten en pesos de ventaja centrados ($w_i$). Primero, se calcula una probabilidad softmax $p_i$ utilizando una temperatura $\tau$. Luego, los pesos se centran contra una línea base uniforme:
  $$w_i = p_i - \frac{1}{N_c}$$
  Esto asegura que $\sum w_i = 0$, asignando pesos positivos a muestras de alta recompensa y pesos negativos a muestras de baja recompensa en relación con el promedio del grupo.
• Objetivo de Optimización: El método optimiza un objetivo de regresión ponderado por ventaja sobre la recompensa implícita. La recompensa implícita se define como la mejora en la pérdida de eliminación de ruido del modelo actual ($\theta$) sobre un modelo de referencia fijo ($ref$):
  $$s_\theta(x_0, x_t, t, c, \epsilon) = \omega(\lambda_t) (l_{ref} - l_\theta)$$
  La función de pérdida es:
  $$L_{LAIR} = \mathbb{E} \left[ -\sum_{i=1}^{N_c} w_i s^{(i)}_\theta + \frac{\lambda}{N_c} \sum_{i=1}^{N_c} (s^{(i)}_\theta)^2 \right]$$
  El primer término alienta al modelo a aumentar la recompensa implícita para muestras de alta recompensa y disminuirla para muestras de baja recompensa. El segundo término es una penalización cuadrática que regulariza la magnitud de la recompensa implícita, evitando desviaciones agresivas del modelo de referencia.

Análisis Teórico
Los autores proporcionan una caracterización teórica del objetivo:

1. Óptimo de Forma Cerrada: El objetivo admite un óptimo de forma cerrada único y acotado en el espacio de recompensa implícita: $s^*_i = \frac{N_c}{2\lambda} w_i$. Esto aclara que la fuerza de regularización $\lambda$ controla directamente la magnitud de la actualización de preferencias.
2. Límite de KL Sustituto: Bajo supuestos específicos (aproximación de log-ratio y extensión de soporte completo acotada), los autores derivan un límite sobre la divergencia de Kullback-Leibler (KL) entre la distribución inducida y la distribución de referencia: $D_{KL} \le \frac{\Delta_c}{\eta}$. Esto sugiere que $\lambda$ controla la nitidez de la inclinación de preferencias, asegurando que la actualización permanezca conservadora.
3. Ventaja sobre Pares: A diferencia de los objetivos por pares (por ejemplo, Bradley-Terry) que pueden impulsar la recompensa implícita hacia el infinito para maximizar el margen entre un ganador y un perdedor, la formulación de lista de LAIR asigna objetivos finitos y relativos basados en todo el conjunto de candidatos, preservando la estructura de clasificación y el espaciado de recompensas.

Resultados Experimentales
El método se evaluó en Stable Diffusion 1.5 (SD1.5) y Stable Diffusion XL (SDXL) utilizando el conjunto de datos Pick-a-Pic v2.

• Configuración: Las imágenes se puntuaron fuera de línea utilizando el modelo de recompensa PickScore. Los modelos se ajustaron finamente y se evaluaron en generación de texto a imagen (HPD, Parti-prompts), generación composicional (GenEval) y edición de imágenes basada en instrucciones (InstructPix2Pix).
• Rendimiento: Diffusion LAIR superó a bases de referencia sólidas (Diffusion-DPO, DSPO, Diffusion-KTO, MaPO, InPO) en la mayoría de las métricas.
  • En SD1.5, obtuvo las puntuaciones más altas en PickScore, HPS, Estética, CLIP e ImageReward para ambos conjuntos de datos Parti-prompt y HPD.
  • En SDXL, superó de manera similar a las bases de referencia, incluidos MaPO y Diffusion-DPO.
  • Generalización: Notablemente, a pesar de haber sido entrenado solo con PickScore, el modelo mostró mejoras significativas en otras métricas (por ejemplo, ImageReward), lo que sugiere que el objetivo de lista mejora la calidad general de las muestras más allá del modelo de recompensa específico utilizado.
  • Composición y Edición: El método demostró ganancias consistentes en las tasas de victoria de GenEval (razonamiento composicional) e InstructPix2Pix (edición de imágenes) en comparación con las bases de referencia por pares.
• Eficiencia: El proceso de entrenamiento fue significativamente más eficiente que las bases de referencia de RL en línea y comparable o más eficiente que otros métodos fuera de línea, requiriendo aproximadamente 24 horas de GPU para SD1.5 y 140 horas de GPU A100 para SDXL.

Significado y Afirmaciones
El artículo afirma que Diffusion LAIR ofrece una dirección simple y efectiva para alinear modelos de difusión al superar las limitaciones de la reducción por pares. Sus contribuciones clave son:

1. Aprendizaje de Lista: Es el primer método que aprovecha directamente grupos de imágenes con puntuación de recompensa para la optimización de preferencias de difusión, utilizando toda la estructura de la calidad relativa de las muestras.
2. Claridad Teórica: Proporciona un óptimo de forma cerrada acotado, ofreciendo información sobre cómo la regularización controla el cambio de distribución y la magnitud de la actualización.
3. Superioridad Empírica: Demuestra un rendimiento sólido en tareas de generación, composición y edición, demostrando que la optimización de lista consciente de la recompensa produce mejoras más amplias que los métodos por pares.

Los autores señalan limitaciones, específicamente la dependencia de la calidad del modelo de recompensa utilizado para el entrenamiento y los supuestos idealizados en el límite teórico de KL. Sin embargo, postulan que el método es robusto ante etiquetas de preferencia ruidosas comunes en conjuntos de datos anotados por humanos.