Rethinking Preference Alignment for Diffusion Models with Classifier-Free Guidance

Cet article propose une méthode d'alignement des préférences pour les modèles de diffusion qui, en découpant l'apprentissage en modules positifs et négatifs pour générer un vecteur de « guidance contrastive » lors de l'inférence, améliore la généralisation et le contrôle sans nécessiter de réentraînement du modèle de base.

L'essentiel

Le Problème : L'Artiste qui a perdu le fil

Imaginez un grand artiste numérique (c'est le modèle de diffusion, comme Stable Diffusion) qui a appris à peindre en regardant des milliards d'images sur Internet. Il est très talentueux, mais comme il a tout vu, il ne sait pas toujours ce que les humains aiment vraiment. Il peut peindre un chat, mais il risque de lui mettre trois pattes ou de ne pas respecter la consigne "un chat mignon".

Pour l'entraîner à mieux faire, les chercheurs utilisent une méthode appelée DPO (Optimisation Directe des Préférences). C'est un peu comme si on lui montrait deux dessins : "Celui-ci est bien, celui-ci est nul". L'artiste apprend de ses erreurs.

Le hic ? Si on force trop l'artiste à apprendre de ces exemples, il devient un mimique obsessionnel. Il finit par copier exactement les exemples qu'on lui a montrés, perdant sa créativité et sa capacité à peindre des choses nouvelles. C'est ce qu'on appelle le "surapprentissage" (overfitting). Il devient rigide et ne sait plus s'adapter à de nouvelles demandes.

La Solution : Le Chef et le Critique (PGD et cPGD)

Les auteurs de ce papier ont eu une idée brillante : au lieu de forcer l'artiste à changer sa façon de peindre (ce qui le rend rigide), pourquoi ne pas lui donner un guide pendant qu'il peint ?

Ils s'inspirent d'une technique existante appelée Guidage sans Classifieur (CFG). Imaginez que vous peignez, et qu'un ami vous dit : "Non, non, le ciel devrait être plus bleu". Vous ajustez votre coup de pinceau en temps réel.

Voici leurs deux nouvelles méthodes :

1. PGD : Le Chef et le Critique

Imaginez que vous avez deux personnes dans la pièce :

• Le Chef (le modèle de base) : Il sait peindre n'importe quoi, mais il est un peu moyen.
• Le Critique (le modèle affiné) : C'est un expert qui a regardé des milliers de dessins "parfaits" et "ratés". Il ne peint pas lui-même, il donne juste des conseils.

Au lieu de remplacer le Chef par le Critique, on garde le Chef et on écoute le Critique.

• Le Chef commence à peindre.
• Le Critique dit : "Hé, ce n'est pas assez comme les dessins que j'aime !".
• On ajuste le coup de pinceau du Chef en fonction de la différence entre ce qu'il fait et ce que le Critique veut.

L'analogie : C'est comme cuisiner. Le Chef prépare le plat de base. Le Critique goûte et dit : "Ajoute un peu de sel". Vous ne jetez pas le plat pour le refaire de zéro, vous l'ajustez juste au moment de servir. Cela évite que le plat devienne trop salé (surapprentissage).

2. cPGD : La Balance de la Cuisine

Pour aller encore plus loin, les auteurs ont une idée encore plus fine. Ils disent : "Pour bien savoir ce qu'on aime, il faut aussi savoir ce qu'on n'aime pas".

Ils entraînent deux experts :

• L'Expert "Oui" : Il a appris uniquement avec les dessins que les gens ont adorés.
• L'Expert "Non" : Il a appris uniquement avec les dessins que les gens ont détestés.

Au moment de peindre, on fait une soustraction :

Ce que l'Expert "Oui" veut MOINS Ce que l'Expert "Non" veut.

L'analogie : Imaginez que vous essayez de trouver le point d'équilibre parfait sur une balance. D'un côté, vous mettez les ingrédients que vous aimez (le "Oui"). De l'autre, vous enlevez ceux que vous détestez (le "Non"). Le résultat est un plat parfaitement équilibré, sans avoir besoin de changer toute la recette de base.

Pourquoi c'est génial ?

1. Pas de réapprentissage coûteux : On n'a pas besoin de réentraîner l'artiste géant (ce qui prend des semaines et coûte cher). On utilise juste ces "guides" pendant la création.
2. Moins de rigidité : Comme on ne force pas l'artiste à changer sa personnalité, il reste créatif et ne fait pas de "catastrophes" (comme dessiner des mains avec 7 doigts).
3. Plug-and-Play : Ces guides sont comme des plugins ou des filtres. Une fois créés, on peut les utiliser avec n'importe quel modèle de base, même ceux qui sont différents. C'est comme ajouter un filtre Instagram à n'importe quelle photo.

En résumé

Au lieu de forcer un artiste à changer sa façon de voir le monde (ce qui le rend rigide et le fait oublier ses bases), cette méthode lui donne un compas pendant qu'il travaille.

• Le PGD utilise un seul compas (ce qu'on aime).
• Le cPGD utilise deux compas (ce qu'on aime moins ce qu'on n'aime pas) pour trouver la direction exacte.

Le résultat ? Des images plus belles, qui respectent mieux les consignes, et qui restent naturelles et variées, sans avoir besoin de réapprendre tout le métier à l'artiste.

Résumé technique

1. Problématique

L'alignement des modèles de diffusion text-to-image (T2I) à grande échelle avec les préférences humaines subtiles reste un défi majeur. Bien que les modèles pré-entraînés sur des données internet massifs soient capables de générer des images de haute fidélité, ils dévient souvent des préférences humaines réelles en termes d'esthétique et de suivi des instructions.

La méthode dominante pour résoudre ce problème est l'Optimisation Directe des Préférences (DPO), adaptée aux modèles de diffusion (Diffusion-DPO). Cependant, le DPO souffre de plusieurs limitations critiques :

• Surapprentissage (Overfitting) : Les modèles finement ajustés avec DPO ont tendance à mémoriser les données d'entraînement plutôt qu'à généraliser.
• Effondrement de mode (Mode Collapse) : La distribution générée peut devenir instable, perdant la diversité des échantillons.
• Généralisation médiocre : Les performances chutent sur des prompts hors distribution (OOD).
• Oubli catastrophique : Le modèle peut perdre ses capacités de génération de base.

Les auteurs illustrent ce problème via une expérience 2D simplifiée montrant que le DPO pousse le modèle vers des solutions dégénérées, tandis que les approches basées sur le renforcement (RLHF) sont complexes et coûteuses en hyperparamètres.

2. Méthodologie

Les auteurs proposent de repenser l'alignement non pas comme un problème de réentraînement complet, mais comme un problème d'inférence guidée, s'inspirant du Classifier-Free Guidance (CFG).

Concept Central : PGD (Preference-Guided Diffusion)

L'idée fondamentale est de traiter le modèle finement ajusté (sur les préférences) comme un signal de contrôle conditionnel, et le modèle de base (non ajusté) comme une distribution a priori (inconditionnelle).

• Au lieu de réentraîner le modèle avec une fonction de perte DPO complexe, on utilise un modèle finement ajusté (avec peu d'itérations) pour guider l'échantillonnage du modèle de base.
• La fonction de score pour l'inférence est construite comme suit :
  $$\nabla \log \pi_{PGD}(x) = \nabla \log \pi_{ref}(x) + w \cdot (\nabla \log \pi_{DPO}(x) - \nabla \log \pi_{ref}(x))$$
  Où $w$ est un poids de guidage. Cela permet d'amplifier la différence entre la préférence et l'a priori sans nécessiter un modèle DPO parfaitement convergé, évitant ainsi le surapprentissage.

Extension : cPGD (Contrastive PGD)

Pour améliorer encore la généralisation, les auteurs découpent l'apprentissage des préférences en deux modules distincts, entraînés séparément :

1. Un modèle $\theta_+$ entraîné uniquement sur les données positives (préférées).
2. Un modèle $\theta_-$ entraîné uniquement sur les données négatives (non préférées).

L'alignement est ensuite réalisé à l'inférence en soustrayant les prédictions du modèle négatif de celles du modèle positif, créant un vecteur de guidage contrastif :
$$\nabla \log \pi_{cPGD}(x, t) = \nabla \log \pi_{ref}(x, t) + w \cdot (\nabla \log \pi(x, t; \theta_+) - \nabla \log \pi(x, t; \theta_-))$$

Cette approche est interprétée comme une ré-pondération dynamique des gradients de la perte DPO. Elle permet de séparer les forces d'attraction (positif) et de répulsion (négatif), offrant un signal d'alignement plus net et contrôlable.

3. Contributions Clés

1. Reformulation théorique : Traitement de l'alignement des modèles de diffusion comme un cas particulier d'inférence de type CFG, transformant le problème de réentraînement en un problème de guidage à l'inférence.
2. Proposition de PGD : Une méthode simple qui aligne la distribution générée avec les préférences humaines en utilisant un modèle finement ajusté comme signal de guidage, sans nécessiter de réentraînement complet du modèle de base.
3. Proposition de cPGD : Une variante contrastive qui entraîne deux modèles indépendants (positif et négatif) et les combine à l'inférence, améliorant la robustesse et la généralisation.
4. Modules Plug-and-Play : La méthode produit des modules qui peuvent être réutilisés pour aligner d'autres modèles de base, une fois entraînés.

4. Résultats Expérimentaux

Les méthodes ont été évaluées sur Stable Diffusion 1.5 (SD1.5) et Stable Diffusion XL (SDXL) en utilisant les jeux de données Pick-a-Pic v2 et HPDv3.

• Performance Quantitative : PGD et cPGD surpassent systématiquement les baselines (DPO, MaPO, NPO, KTO) sur les métriques de récompense (PickScore, HPSv2, HPSv3, ImageReward).
  • Sur SDXL, cPGD atteint un taux de victoire moyen (win rate) de ~70-82% contre le modèle de base, contre ~66% pour DPO.
  • Les méthodes proposées obtiennent des améliorations Pareto : elles augmentent la récompense tout en maintenant (voire en améliorant) la diversité des échantillons et la préservation de l'a priori (mesurée par le FID).
• Robustesse : Les méthodes montrent une meilleure généralisation sur des prompts de test variés (Parti-Prompts) et résistent mieux aux variations de qualité des données d'entraînement (sous-ensembles de haute qualité vs données brutes).
• Étude Humaine : Une évaluation humaine sur 55 prompts montre que PGD est sélectionné dans 45,5% des cas, surpassant largement DPO (29,5%) et les autres méthodes.
• Efficacité : Bien que l'inférence nécessite de calculer deux modèles (base + guidage), les auteurs montrent qu'il est possible de distiller ces modules en un seul checkpoint (merged checkpoint) avec une perte minime de performance, réduisant ainsi le coût d'inférence.

5. Signification et Impact

Ce travail apporte un changement de paradigme dans l'alignement des modèles de diffusion :

• Simplicité et Efficacité : Il démontre qu'un réentraînement complexe avec DPO n'est pas nécessaire pour obtenir un alignement de haute qualité. Un guidage léger à l'inférence suffit.
• Stabilité : En évitant l'optimisation directe de la perte DPO sur le modèle de génération, on élimine les risques d'effondrement de mode et de surapprentissage.
• Flexibilité : L'approche "Plug-and-Play" permet d'adapter des modules d'alignement à différents modèles de base (même avec des architectures différentes mais un espace latent partagé, comme démontré avec KOALA), facilitant le déploiement et la mise à jour des préférences sans réentraîner l'ensemble du modèle.

En résumé, l'article propose une solution élégante et robuste pour aligner les modèles de diffusion sur les préférences humaines, en exploitant la puissance du guidage sans classifieur plutôt que celle de l'optimisation directe des paramètres.