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

Résumé Technique : Diffusion LAIR

Énoncé du Problème
Les méthodes actuelles d'optimisation des préférences pour l'alignement des modèles de diffusion texte-à-image, telles que l'Optimisation Directe des Préférences (DPO) et ses variantes spécifiques aux diffusion (par exemple, Diffusion-DPO, DSPO), reposent largement sur une supervision par paires. Ces méthodes réduisent les données d'entraînement à des comparaisons binaires « gagnant-perdant » pour un prompt donné. Les auteurs soutiennent que cette réduction est limitante pour deux raisons principales :

1. Structure des Données : Les ensembles de données de préférences pratiques (par exemple, Pick-a-Pic) contiennent souvent plusieurs images candidates pour un seul prompt, créant une structure naturelle de liste que les méthodes par paires ignorent.
2. Perte d'Information : Les modèles de récompense modernes fournissent des scores continus. Réduire ces scores à une seule étiquette binaire élimine la qualité relative des candidats non sélectionnés et ignore l'ampleur des écarts de récompense entre les échantillons.
   De plus, les approches d'apprentissage par renforcement (RL) en ligne existantes pour l'alignement des diffusion sont coûteuses en calcul, difficiles à paralléliser et sensibles aux décalages de distribution en raison de leur nature on-policy. Bien que les méthodes hors ligne soient plus efficaces, elles n'ont pas pleinement exploité le signal plus riche disponible dans les données de type liste et notées par récompense.

Méthodologie : Diffusion LAIR
L'article propose Diffusion LAIR (Listwise Advantage-weighted Implicit Reward), une méthode d'optimisation des préférences hors ligne qui traite l'alignement comme un problème de liste conscient de la récompense.

• Représentation des Données : Au lieu de paires indépendantes, la méthode traite un prompt $c$ avec un ensemble de $N_c$ images candidates $\{x^{(i)}_0\}$ et leurs scores de récompense continus correspondants $\{r_i\}$.
• Pondération par l'Avantage : Les scores de récompense sont convertis en poids d'avantage centrés ($w_i$). D'abord, une probabilité softmax $p_i$ est calculée en utilisant une température $\tau$. Ensuite, les poids sont centrés par rapport à une ligne de base uniforme :
  $$w_i = p_i - \frac{1}{N_c}$$
  Cela garantit que $\sum w_i = 0$, attribuant des poids positifs aux échantillons à haute récompense et des poids négatifs aux échantillons à faible récompense par rapport à la moyenne du groupe.
• Objectif d'Optimisation : La méthode optimise un objectif de régression pondéré par l'avantage sur la récompense implicite. La récompense implicite est définie comme l'amélioration de la perte de débruitage du modèle actuel ($\theta$) par rapport à un modèle de référence fixe ($ref$) :
  $$s_\theta(x_0, x_t, t, c, \epsilon) = \omega(\lambda_t) (l_{ref} - l_\theta)$$
  La fonction de perte est :
  $$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]$$
  Le premier terme encourage le modèle à augmenter la récompense implicite pour les échantillons à haute récompense et à la diminuer pour les échantillons à faible récompense. Le deuxième terme est une pénalité quadratique qui régularise l'ampleur de la récompense implicite, empêchant une déviation agressive par rapport au modèle de référence.

Analyse Théorique
Les auteurs fournissent une caractérisation théorique de l'objectif :

1. Optimum sous Forme Close : L'objectif admet un optimum unique et borné sous forme close dans l'espace de la récompense implicite : $s^*_i = \frac{N_c}{2\lambda} w_i$. Cela clarifie que la force de régularisation $\lambda$ contrôle directement l'ampleur de la mise à jour des préférences.
2. Limite de KL de Surrogate : Sous des hypothèses spécifiques (approximation du rapport logarithmique et extension à support complet bornée), les auteurs dérivent une borne sur la divergence de Kullback-Leibler (KL) entre la distribution induite et la distribution de référence : $D_{KL} \le \frac{\Delta_c}{\eta}$. Cela suggère que $\lambda$ contrôle la netteté de l'inclinaison des préférences, garantissant que la mise à jour reste conservatrice.
3. Avantage par rapport aux Paires : Contrairement aux objectifs par paires (par exemple, Bradley-Terry) qui peuvent conduire la récompense implicite à l'infini pour maximiser la marge entre un gagnant et un perdant, la formulation de liste de LAIR attribue des cibles finies et relatives basées sur l'ensemble complet des candidats, préservant ainsi la structure de classement et l'espacement des récompenses.

Résultats Expérimentaux
La méthode a été évaluée sur Stable Diffusion 1.5 (SD1.5) et Stable Diffusion XL (SDXL) en utilisant l'ensemble de données Pick-a-Pic v2.

• Configuration : Les images ont été notées hors ligne à l'aide du modèle de récompense PickScore. Les modèles ont été affinés et évalués sur la génération texte-à-image (HPD, Parti-prompts), la génération compositionnelle (GenEval) et l'édition d'images basée sur des instructions (InstructPix2Pix).
• Performance : Diffusion LAIR a surpassé des bases solides (Diffusion-DPO, DSPO, Diffusion-KTO, MaPO, InPO) sur la plupart des métriques.
  • Sur SD1.5, il a obtenu les scores les plus élevés sur PickScore, HPS, Aesthetics, CLIP et ImageReward pour les ensembles de données Parti-prompt et HPD.
  • Sur SDXL, il a également surpassé les bases, y compris MaPO et Diffusion-DPO.
  • Généralisation : Notamment, bien qu'entraîné uniquement avec PickScore, le modèle a montré des améliorations significatives sur d'autres métriques (par exemple, ImageReward), suggérant que l'objectif de liste améliore la qualité globale des échantillons au-delà du modèle de récompense spécifique utilisé.
  • Composition et Édition : La méthode a démontré des gains constants dans les taux de victoire de GenEval (raisonnement compositionnel) et d'InstructPix2Pix (édition d'images) par rapport aux bases par paires.
• Efficacité : Le processus d'entraînement a été significativement plus efficace que les bases de RL en ligne et comparable ou plus efficace que d'autres méthodes hors ligne, nécessitant environ 24 heures GPU pour SD1.5 et 140 heures GPU A100 pour SDXL.

Signification et Revendications
L'article revendique que Diffusion LAIR offre une direction simple et efficace pour l'alignement des modèles de diffusion en dépassant les limites de la réduction par paires. Ses contributions clés sont :

1. Apprentissage de Liste : C'est la première méthode à exploiter directement des groupes d'images notées par récompense pour l'optimisation des préférences de diffusion, en utilisant la structure complète de la qualité relative des échantillons.
2. Clarté Théorique : Elle fournit un optimum sous forme close borné, offrant un aperçu de la manière dont la régularisation contrôle le décalage de distribution et l'ampleur de la mise à jour.
3. Supériorité Empirique : Elle démontre une forte performance sur les tâches de génération, de composition et d'édition, prouvant que l'optimisation de liste consciente de la récompense produit des améliorations plus larges que les méthodes par paires.

Les auteurs notent des limitations, spécifiquement la dépendance à la qualité du modèle de récompense utilisé pour l'entraînement et les hypothèses idéalisées dans la borne théorique de KL. Cependant, ils posent que la méthode est robuste aux étiquettes de préférences bruyantes courantes dans les ensembles de données annotés par des humains.