Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Ce papier propose un cadre de fine-tuning pour les modèles de diffusion qui, en façonnant les distributions intermédiaires via P-GRAFT et en corrigeant le bruit inverse, améliore la génération d'images et d'autres contenus tout en offrant une meilleure efficacité théorique et empirique que les méthodes existantes.

L'essentiel

Imaginez que vous avez un chef cuisinier très talentueux (le modèle de diffusion pré-entraîné). Ce chef sait cuisiner des plats délicieux (générer des images, des molécules, etc.) à partir de rien. Mais parfois, il fait des erreurs : il met trop de sel, il oublie un ingrédient, ou il ne respecte pas exactement la recette que vous lui avez donnée.

L'objectif de ce papier est d'apprendre à ce chef à faire mieux, sans avoir à le rééduquer de zéro. Les auteurs proposent deux nouvelles méthodes magiques pour y parvenir.

1. Le problème : Pourquoi est-ce difficile de corriger le chef ?

Habituellement, pour améliorer un chef, on lui donne des notes après chaque plat. S'il fait un bon plat, on le félicite ; s'il fait un mauvais plat, on le gronde. En intelligence artificielle, cela s'appelle le "renforcement".

Mais il y a un gros problème avec les modèles de diffusion (comme ceux qui créent des images) :

• C'est comme si le chef cuisinait un plat en 50 étapes (de l'ébauche floue à l'image finale).
• Pour le féliciter ou le gronder, il faut attendre la fin du plat (l'image finale).
• Le problème, c'est que le chef a déjà oublié ce qu'il a fait aux étapes 1 à 49. Il est très difficile de savoir exactement quelle étape a gâché le plat. C'est comme essayer de corriger un gâteau en train de cuire en regardant seulement la croûte finale.

2. La première solution : P-GRAFT (Le "Dégustateur Intermédiaire")

Les chercheurs ont eu une idée brillante : Pourquoi attendre la fin du plat pour donner la note ?

Imaginez que le chef prépare un gâteau. Au lieu de goûter le gâteau une fois qu'il est sorti du four (l'image finale), vous goûtez la pâte à mi-cuisson (à un niveau de bruit intermédiaire).

• L'analogie : C'est comme si vous arrêtiez le chef à mi-chemin de sa recette. Vous lui dites : "Attends, cette pâte sent déjà très bon, continue comme ça !".
• Comment ça marche ? Le modèle génère beaucoup d'images, mais on s'arrête à mi-chemin. On regarde l'image finale correspondante pour voir si c'est un "bon" résultat. Si oui, on garde la version "mi-cuite" de cette image pour entraîner le chef.
• Le résultat : Le chef apprend à faire les premières étapes de la recette beaucoup plus précisément. C'est comme apprendre à un pianiste à jouer les premières mesures d'une symphonie parfaitement, car c'est là que l'erreur est la plus facile à corriger.

Pourquoi c'est mieux ?
Les auteurs expliquent cela par un équilibre entre précision et bruit.

• Si vous attendez la fin (image finale), le résultat est clair, mais il est difficile de savoir comment y arriver (trop de bruit dans l'apprentissage).
• Si vous arrêtez trop tôt (au début), c'est trop flou pour savoir si c'est bon.
• P-GRAFT trouve le moment parfait (le "juste milieu") où le chef peut encore apprendre facilement, tout en ayant une idée claire de la qualité finale.

3. La deuxième solution : La Correction du "Bruit Inverse" (Inverse Noise Correction)

Cette méthode s'applique à un type de modèle légèrement différent (les modèles "Flow"), mais l'idée est encore plus simple.

Imaginez que le chef utilise une machine à remonter le temps.

• Normalement, il prend du chaos (du bruit blanc, comme de la neige sur une vieille télé) et le transforme en un beau portrait.
• Parfois, la machine est mal calibrée : elle ne prend pas le bon type de "chaos" au départ, donc le portrait final a des défauts.

L'idée géniale :
Au lieu de réapprendre toute la machine, les chercheurs disent : "Et si on apprenait à la machine à générer le bon type de chaos au départ ?"

1. Ils prennent des images parfaites (celles que le chef sait déjà faire).
2. Ils utilisent la machine à l'envers pour transformer ces images parfaites en "bruit".
3. Ils découvrent que ce "bruit" est différent du bruit normal. C'est un "bruit corrigé".
4. Ils entraînent un petit assistant (le "Correcteur de Bruit") qui sait transformer le bruit normal en ce "bruit corrigé".
5. Le résultat : Quand on demande au chef de cuisiner, on lui donne d'abord le "bruit corrigé" préparé par l'assistant. Le chef, qui est déjà bon, produit alors un résultat parfait sans avoir besoin de changer ses compétences de base.

C'est comme si vous donniez au chef des ingrédients pré-mélangés et parfaits, au lieu de lui donner des sacs de farine et de sucre en vrac.

En résumé : Les résultats

Les auteurs ont testé ces idées sur plusieurs tâches :

• Création d'images (Texte vers Image) : Le modèle génère des images qui correspondent beaucoup mieux à ce qu'on lui demande (par exemple, "trois fleurs rouges" au lieu de "une fleur bleue").
• Génération de molécules : Ils créent plus de molécules stables (qui ne s'effondrent pas).
• Efficacité : Ces méthodes sont souvent plus rapides et moins coûteuses en énergie que les méthodes actuelles.

La morale de l'histoire :
Au lieu de forcer l'intelligence artificielle à tout réapprendre de zéro ou de la corriger à la fin (ce qui est difficile et instable), il est souvent plus intelligent de façonner le processus à mi-chemin ou de préparer les ingrédients de départ pour que le résultat soit naturel et parfait. C'est une approche plus douce, plus intelligente et beaucoup plus efficace.

Résumé technique

1. Problématique

Les modèles de diffusion pré-entraînés nécessitent souvent un affinement (fine-tuning) pour corriger des erreurs d'apprentissage ou s'aligner sur des tâches spécifiques (comme la génération d'images à partir de texte). La méthode standard pour l'alignement dans les modèles de langage utilise l'apprentissage par renforcement (RL) avec régularisation KL (ex: PPO). Cependant, l'application de ces méthodes aux modèles de diffusion pose deux problèmes majeurs :

1. Intractabilité de la vraisemblance marginale : Contrairement aux modèles autoregressifs, le calcul de la vraisemblance marginale nécessaire pour la régularisation KL dans PPO est impossible pour les modèles de diffusion.
2. Instabilité et biais : Ignorer la régularisation KL entraîne une instabilité à grande échelle, tandis que l'utilisation de contraintes KL sur les trajectoires (trajectory KL) conduit à des résultats sous-optimaux et à un biais de la fonction de valeur initiale.

Les méthodes existantes basées sur l'échantillonnage par rejet (comme RAFT) sont prometteuses mais n'exploitent pas pleinement la structure temporelle des modèles de diffusion.

2. Méthodologie

Les auteurs proposent une approche unifiée et deux nouvelles stratégies principales :

A. Unification : GRAFT (Generalized Rejection sAmpling Fine-Tuning)

Les auteurs unissent diverses méthodes d'échantillonnage par rejet (comme RAFT, RSO, Best-of-N) sous un cadre théorique appelé GRAFT.

• Principe : Ils démontrent que l'échantillonnage par rejet généralisé (GRS) équivaut implicitement à une maximisation de récompense régularisée par KL, mais avec des récompenses remodelées (reshaped rewards).
• Avantage : Cela permet d'imposer une contrainte de KL marginale pour les modèles de diffusion sans avoir besoin de calculer la vraisemblance marginale intractable.

B. Innovation 1 : P-GRAFT (Partial GRAFT)

Motivés par la structure des modèles de diffusion (qui génèrent des données via un processus de débruitage progressif), les auteurs proposent d'affiner le modèle uniquement jusqu'à un niveau de bruit intermédiaire ($t$), plutôt que jusqu'à la donnée finale ($t=0$).

• Fonctionnement :
  1. Le modèle de référence génère des trajectoires complètes.
  2. Les récompenses sont calculées sur les données finales débruitées.
  3. L'échantillonnage par rejet est appliqué aux états intermédiaires ($X_t$) en fonction de ces récompenses finales.
  4. Le modèle est affiné uniquement sur la phase de débruitage de $T$ à $t$. Pour la phase finale ($t$ à $0$), on utilise le modèle de référence.
• Justification Théorique (Compromis Biais-Variance) :
  • Variance : À un temps $t$ proche du bruit pur, la récompense conditionnelle est très bruitée (variance élevée).
  • Biais : À un temps $t$ proche de la donnée finale, la fonction de score (score function) est complexe et difficile à apprendre (biais élevé).
  • Résultat : P-GRAFT trouve un point d'équilibre optimal en choisissant un temps intermédiaire $t$ où la fonction de score est plus simple à apprendre (proche d'une distribution gaussienne) tout en conservant assez d'information pour guider le modèle.

C. Innovation 2 : Inverse Noise Correction (Correction du Bruit Inverse)

Cette méthode s'applique aux modèles de flux (Flow Models) et vise à corriger les erreurs d'un modèle pré-entraîné sans récompense explicite.

• Principe : Les modèles de flux sont réversibles. Si un modèle pré-entraîné $v_\theta$ génère une distribution imparfaite $p_1$ à partir d'un bruit $p_0$, on peut utiliser l'inversion de l'ODE (via la méthode d'Euler implicite) pour mapper les données réelles vers un "bruit inverse" ($p^{rev}_1$).
• Algorithme :
  1. On génère un jeu de données de "bruit inverse" en faisant passer des échantillons réels à travers le modèle pré-entraîné inversé.
  2. On entraîne un petit modèle "Correcteur de Bruit" (Noise Corrector) pour apprendre à générer ce bruit inverse à partir d'un bruit gaussien standard.
  3. Lors de l'inférence, on utilise le correcteur pour générer le bruit initial, puis le modèle pré-entraîné pour générer l'image finale.
• Théorie : Cela permet de corriger la distribution de sortie en corrigeant la distribution d'entrée, réduisant ainsi l'erreur de génération avec un coût computationnel moindre.

3. Contributions Clés

1. Cadre GRAFT : Une unification théorique prouvant que les méthodes d'échantillonnage par rejet réalisent une maximisation de récompense régularisée par KL marginale pour les modèles de diffusion.
2. P-GRAFT : Une méthode d'affinement efficace exploitant le compromis biais-variance en affinant uniquement la partie initiale du processus de débruitage.
3. Inverse Noise Correction : Une nouvelle méthode pour améliorer les modèles de flux pré-entraînés sans récompense, en apprenant une distribution de bruit corrigée via la réversibilité du modèle.
4. Preuves Mathématiques : Justification rigoureuse via la théorie de Bakry-Émery (convergence exponentielle du score vers le score gaussien) et l'analyse de la variance conditionnelle.

4. Résultats Expérimentaux

Les méthodes ont été évaluées sur plusieurs tâches : génération d'images texte-à-image (T2I), génération de layouts, génération de molécules et génération d'images sans condition.

• Génération Texte-à-Image (Stable Diffusion v2) :

  • P-GRAFT surpasse significativement les méthodes par gradient de politique (comme DDPO) et le modèle de base SDv2.
  • Sur le benchmark GenAI-Bench, P-GRAFT atteint un score VQAScore de 71.94 (contre 66.87 pour SDv2 et 65.70 pour DDPO), soit une amélioration relative de 8,81% par rapport au modèle de base.
  • L'ablation montre que le choix du temps intermédiaire optimal (ex: $0.25N$) est crucial, confirmant la théorie du compromis biais-variance.

• Génération de Molécules et de Layouts :

  • Sur QM9 (molécules), P-GRAFT améliore la stabilité des molécules générées (92.61% vs 90.50% pour la base) avec une seule passe d'échantillonnage, tandis que GRAFT nécessite 9 fois plus d'échantillons pour une amélioration marginale.
  • Sur PubLayNet (layouts), P-GRAFT améliore l'alignement des éléments tout en maintenant une distribution proche de l'entraînement initial (FID plus faible que GRAFT).

• Correction de Bruit Inverse (Flow Models) :

  • Sur CelebA-HQ et LSUN-Church, la méthode améliore le FID (Fréchet Inception Distance) tout en réduisant le coût computationnel (FLOPs).
  • Un correcteur de bruit léger (16M paramètres) combiné au modèle pré-entraîné (65M) surpasse le modèle pré-entraîné seul avec 1000 étapes de débruitage, tout en utilisant beaucoup moins de FLOPs par image.

5. Signification et Impact

Ce travail apporte une avancée théorique et pratique majeure dans l'alignement des modèles génératifs :

• Efficacité Computationnelle : P-GRAFT évite le backpropagation à travers tout le processus de débruitage (contrairement à PPO/DDPO), réduisant considérablement le coût d'entraînement.
• Stabilité : En évitant la régularisation KL sur les trajectoires complètes et en exploitant la simplicité des scores aux temps intermédiaires, la méthode est plus stable et moins sujette à l'effondrement de mode (mode collapse).
• Généralité : La méthode s'applique aussi bien aux modèles de diffusion (SDE) qu'aux modèles de flux (ODE) et aux modèles discrets-continus, offrant un cadre unifié pour l'affinement sans nécessiter de récompenses explicites dans le cas de la correction de bruit.

En résumé, l'article démontre que le façonnage de la distribution à des niveaux de bruit intermédiaires est une stratégie supérieure pour l'affinement des modèles de diffusion, surpassant les approches par gradient de politique tout en étant plus efficace en termes de ressources.