Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function

Ce papier propose SQDF, une nouvelle méthode d'alignement par apprentissage par renforcement régularisé KL pour les modèles de diffusion, qui utilise un gradient de politique reparamétré basé sur une fonction Q douce pour surmonter le problème de sur-optimisation de la récompense tout en préservant la diversité et la naturalité des échantillons générés.

L'essentiel

Imaginez que vous avez un chef cuisinier génial (le modèle de diffusion) qui sait cuisiner n'importe quel plat avec une précision parfaite. Cependant, ce chef ne sait pas exactement ce que vous voulez manger spécifiquement aujourd'hui. Vous voulez un plat non seulement délicieux, mais aussi qui plaise à un critique culinaire très pointu (l'objectif de récompense).

Le problème, c'est que si vous demandez simplement au chef : « Fais-moi un plat qui plaît à ce critique ! », il risque de tomber dans le piège du sur-optimisation.

Le Problème : Le Chef qui "Triche"

Si vous ne faites pas attention, le chef va essayer de tricher. Au lieu de cuisiner un vrai plat savoureux, il va créer un monstre de nourriture qui ressemble à rien, mais qui a exactement les ingrédients que le critique adore.

• Résultat : Le critique dit "10/10", mais le plat est bizarre, laid, et tout le monde a le même plat bizarre. C'est ce qu'on appelle l'effondrement de la diversité et de la qualité.

Les anciennes méthodes pour corriger cela étaient soit trop lentes, soit instables (comme essayer d'enseigner à un élève en lui criant des formules mathématiques complexes).

La Solution : SQDF (Le Nouveau Système de Coaching)

Les auteurs de ce papier proposent une nouvelle méthode appelée SQDF. Imaginez que c'est un nouveau système de coaching pour notre chef cuisinier. Voici comment ça marche, étape par étape, avec des analogies simples :

1. La "Boussole Instantanée" (La fonction Q douce)

Au lieu de faire cuisiner un plat entier de A à Z pour voir si c'est bon (ce qui prend du temps), SQDF utilise une boussole instantanée.

• L'analogie : Imaginez que le chef est en train de sculpter une statue de glace. Au lieu d'attendre la fin pour voir si elle est belle, il regarde un instantané de la glace à mi-chemin et demande à un expert (le modèle de "consistance") : « Si je continue comme ça, à quoi ressemblera la statue finale ? ».
• L'avantage : Cette boussole est gratuite et rapide. Elle permet au chef de corriger son coup pendant qu'il sculpte, sans avoir à attendre la fin du processus.

2. Le "Ralentisseur de Crédit" (Le facteur d'escompte)

Dans la sculpture de glace, les premiers coups de marteau (quand la glace est très sale) ont moins d'importance que les derniers coups de polissage.

• L'analogie : Si vous essayez d'apprendre à quelqu'un à jouer du piano, vous ne le grondez pas pour une note fausse jouée au tout début de la chanson si la fin est magnifique.
• L'innovation : SQDF utilise un facteur d'escompte. Il dit au chef : « Ne t'inquiète pas trop des erreurs au début du processus, concentre-toi sur les étapes finales où le plat prend sa forme. » Cela évite de gaspiller de l'énergie sur des détails qui n'ont pas d'impact sur le résultat final.

3. Le "Carnet de Recettes" (Le tampon de replay)

Souvent, le chef fait une erreur et produit un plat horrible, ou alors il produit un plat incroyable par hasard.

• L'analogie : Au lieu de jeter tout ce qui est produit, SQDF garde un carnet de recettes (un tampon). Il y range les meilleurs plats qu'il a déjà faits, même s'ils sont rares.
• L'avantage : Quand le chef doit apprendre, il ne regarde pas seulement ce qu'il vient de faire, mais il relit ses meilleures recettes passées. Cela l'empêche d'oublier comment faire des plats variés et évite qu'il ne se focalise uniquement sur un seul type de plat "parfait" mais ennuyeux.

Le Résultat : Un Chef Heureux et Créatif

Grâce à SQDF, le chef cuisinier (le modèle d'IA) apprend à :

1. Plaire au critique (obtenir un score élevé).
2. Garder son style naturel (ne pas créer de monstres bizarres).
3. Variété (cuisiner des plats différents, pas toujours le même).

En résumé, cette méthode permet d'entraîner des intelligences artificielles créatives (comme celles qui génèrent des images) pour qu'elles fassent exactement ce que l'on veut, sans qu'elles ne deviennent folles en essayant de trop bien faire. C'est comme trouver le juste équilibre entre suivre un guide strict et garder sa propre créativité.

En une phrase : SQDF est un coach intelligent qui aide l'IA à apprendre de ses erreurs en temps réel, en se concentrant sur les moments clés et en se souvenant de ses meilleures performances, pour créer des images magnifiques et variées sans "tricher".

Résumé technique

1. Problématique

Les modèles de diffusion sont devenus la norme pour la génération d'échantillons haute fidélité (synthèse image-texte, vidéo, molécules). Cependant, l'alignement de ces modèles pré-entraînés avec des objectifs spécifiques (comme la qualité esthétique ou la préférence humaine) pose un défi majeur : la sur-optimisation de la récompense (reward over-optimization).

Les méthodes existantes de fine-tuning, qu'elles soient basées sur l'apprentissage par renforcement (RL) ou la rétropropagation directe des gradients de récompense, tendent à maximiser la récompense au détriment de la qualité intrinsèque du modèle. Cela conduit à deux phénomènes de dégradation :

• Effondrement sémantique (Semantic Collapse) : Perte d'alignement avec le prompt (le contenu devient illisible ou absurde).
• Effondrement de la diversité (Diversity Collapse) : Les échantillons générés convergent vers des motifs similaires, réduisant la variété des sorties.

Les approches actuelles tentent de mitiger ce problème via une régularisation KL (Kullback-Leibler), mais elles souffrent souvent d'instabilité lors de l'entraînement de fonctions de valeur (Value Function) ou dépendent d'estimateurs de gradient Monte Carlo à haute variance.

2. Méthodologie : SQDF

Les auteurs proposent SQDF (Soft Q-based Diffusion Finetuning), une nouvelle méthode de RL régularisée par KL qui utilise un gradient de politique reparamétré guidé par une estimation de la fonction Soft Q sans nécessiter d'entraînement séparé de cette fonction.

Cœur de la méthode

• Approximation de la fonction Soft Q : Au lieu d'entraîner un réseau de valeur instable, SQDF approxime la fonction Soft Q optimale via une moyenne postérieure en un seul pas (Single-step posterior mean approximation), dérivée de la formule de Tweedie. Cette approximation est différentiable, permettant d'utiliser directement le gradient de la récompense pour mettre à jour la politique.
• Gradient de Politique Reparamétré : En utilisant la reparamétrisation ($x_{t-1} = \mu_\theta(x_t, t) + \sigma_t \epsilon$), la méthode permet une mise à jour de la politique à faible variance et efficace en échantillons, évitant la rétropropagation à travers toute la chaîne de débruitage.
• Objectif KL-Régularisé : L'objectif maximise la récompense tout en maintenant la distribution du modèle fine-tuné proche de celle du modèle pré-entraîné (référence), préservant ainsi la naturalité et la diversité.

Trois Innovations Clés

Pour stabiliser l'entraînement et améliorer l'efficacité, SQDF intègre trois composants supplémentaires :

1. Facteur d'Escompte ($\gamma$) :

   • Dans le processus de débruitage, les étapes précoces (bruit élevé) ont une influence limitée sur la qualité finale de l'image.
   • SQDF introduit un facteur d'escompte $\gamma < 1$ pour réduire le crédit attribué aux actions prises aux étapes précoces, améliorant ainsi l'assignation du crédit (credit assignment) et réduisant le bruit d'approximation dans les premières étapes.

2. Intégration de Modèles de Cohérence (Consistency Models) :

   • L'approximation de Tweedie est peu fiable aux niveaux de bruit élevés (étapes précoces).
   • SQDF utilise un modèle de cohérence (Consistency Model) pour prédire l'image propre $\hat{x}_0$ à partir de l'état bruité $x_{t-1}$. Cela fournit une estimation de la moyenne postérieure beaucoup plus précise et uniforme sur toutes les étapes de temps, guidant mieux l'optimisation.

3. Mise à Jour Hors-Politique (Off-Policy) avec Buffer de Replay :

   • Contrairement aux méthodes on-policy, SQDF utilise un tampon de replay (Replay Buffer) pour stocker les échantillons.
   • Cela permet de réutiliser des échantillons rares à haute récompense, d'améliorer la couverture des modes (mode coverage) et de mieux gérer le compromis récompense-diversité.

3. Résultats Expérimentaux

Les auteurs ont évalué SQDF sur deux scénarios principaux en utilisant Stable Diffusion v1.5 et SDXL.

A. Fine-tuning pour l'Image-Texte

• Tâches : Optimisation des scores esthétiques (LAION Aesthetic) et des scores de préférence humaine (HPSv2).
• Comparaison : SQDF est comparé à des méthodes de base comme DDPO (RL), DRaFT et ReFL (rétropropagation directe), ainsi qu'à leurs variantes régularisées par KL.
• Résultats :
  • SQDF atteint des récompenses cibles supérieures tout en maintenant des scores d'alignement (ImageReward, HPS) et de diversité (LPIPS, DreamSim) bien meilleurs que les méthodes de base.
  • Il évite l'effondrement sémantique et la perte de diversité observés chez les autres méthodes à mesure que la récompense augmente.
  • SQDF occupe la frontière de Pareto optimale, surpassant même les méthodes de base augmentées d'une régularisation KL simple.

B. Optimisation Boîte Noire en Ligne (Online Black-Box Optimization)

• Scénario : Optimisation avec un budget de requêtes limité vers une oracle de récompense (score esthétique) non différentiable, utilisant un modèle substitut (proxy).
• Résultats :
  • SQDF démontre une efficacité d'échantillonnage supérieure, atteignant des récompenses cibles élevées tout en préservant la diversité et la naturalité.
  • Les méthodes concurrentes (comme SEIKO) souffrent d'un compromis critique où l'amélioration de la récompense entraîne une dégradation rapide de l'alignement et de la diversité. SQDF reste robuste face à des récompenses imparfaites.

C. Études d'Ablation

• Facteur d'escompte : Son absence ($\gamma=1$) conduit à une optimisation plus lente et à une chute de la diversité.
• Modèle de cohérence : Son retrait réduit la performance de la récompense cible, confirmant son rôle dans la fiabilité de l'estimation du Q-soft.
• Buffer de Replay : Son retrait diminue la diversité, prouvant son utilité pour la couverture des modes.

4. Contributions Clés

1. SQDF : Un cadre RL régularisé par KL pour le fine-tuning de modèles de diffusion qui évite l'entraînement instable de fonctions de valeur en utilisant une approximation de Q-soft sans entraînement.
2. Stabilité par Approximation : L'utilisation combinée d'un facteur d'escompte et de modèles de cohérence pour surmonter les limitations de l'approximation de Tweedie et du bruit précoce.
3. Efficacité et Diversité : Démonstration empirique que SQDF résout le compromis récompense-diversité mieux que les approches actuelles, tant en contexte de récompense différentiable que dans l'optimisation boîte noire.

5. Signification et Impact

Ce travail est significatif car il propose une solution élégante et stable au problème persistant de la sur-optimisation dans les modèles de diffusion génératifs. En éliminant le besoin d'entraîner une fonction de valeur complexe et en exploitant directement les gradients de récompense via une approximation différentiable fiable, SQDF permet d'aligner les modèles de génération sur des objectifs complexes sans sacrifier la qualité intrinsèque ou la diversité des échantillons. Cela ouvre la voie à des applications plus robustes en conception de médicaments, en synthèse d'images artistiques et en contrôle de systèmes génératifs.

Le code est disponible publiquement sur GitHub, facilitant la reproduction et l'adoption de la méthode.