Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning

1. Problématique

Les Modèles de Diffusion (DM) ont établi un état de l'art dans la génération d'images, notamment grâce à la technique de Guidage sans Classificateur (Classifier-Free Guidance - CFG). Le CFG permet de contrôler le compromis entre la fidélité de l'image générée et sa diversité en interpolant entre les prédictions conditionnelles et non conditionnelles lors de l'inférence. Cependant, l'application du CFG aux DMs est coûteuse en temps de calcul car elle nécessite de multiples étapes de débruitage itératif.

Les Modèles de Cohérence (Consistency Models - CMs) ont été développés pour résoudre ce problème de coût en permettant une génération d'images de haute qualité en une seule étape ou en quelques étapes. Néanmoins, les méthodes de guidage existantes pour les CMs souffrent de limitations majeures :

• Elles reposent généralement sur la Distillation de Cohérence (Consistency Distillation - CD), qui nécessite un modèle DM pré-entraîné (enseignant) pour distiller le guidage.
• Pour les modèles entraînés par Entraînement de Cohérence (Consistency Training - CT), qui sont entièrement basés sur les données sans nécessiter de DM enseignant, il n'existe pas de méthode de guidage efficace a posteriori (post-hoc).
• L'absence de guidage pour les modèles CT limite leur contrôle et leur qualité par rapport aux modèles distillés.

Objectif du papier : Proposer une méthode de guidage a posteriori pour les CMs pré-entraînés (notamment ceux issus de l'entraînement CT) qui fonctionne sans dépendre d'un modèle DM enseignant, en permettant un contrôle ajustable similaire au CFG.

---

2. Méthodologie : Joint Flow Distribution Learning (JFDL)

Les auteurs proposent JFDL, une méthode légère d'alignement qui permet d'ajouter des capacités de guidage à un CM pré-entraîné en exploitant la connexion théorique entre les Modèles de Diffusion, les Modèles basés sur le Flux (Flow-based Models - FMs) et les CMs.

A. Fondements Théoriques

• Correspondance CM/FM : Les auteurs établissent que, dans un schéma d'explosion de variance (VE), l'objectif d'entraînement d'un CM (ECT) recouvre une approximation précise du champ de vecteurs de l'ODE de flux de probabilité (PF-ODE), similaire à l'objectif de Flow Matching (CFM) des FMs.
• Hypothèse de bruit pseudo-Gaussien : Le cœur de la méthode repose sur l'idée que la trajectoire hybride combinant un échantillon conditionnel et un ancrage non conditionnel (ou "null") génère un bruit résiduel qui suit une distribution gaussienne.
  • Ils définissent une solution hybride d'ODE : $y_{0}^{\emptyset, t} = \text{Solver}_{t \to 0}^{\emptyset}(\text{Solver}_{0 \to t}^{c}(x_0^c))$.
  • Ils démontrent théoriquement (Proposition 1) que la loi marginale de cette solution hybride correspond à la distribution de données non conditionnelle $p_0(x)$.
  • Ils prouvent (Proposition 2) et vérifient expérimentalement que le "bruit pseudo" résultant de cette trajectoire est effectivement gaussien, validant ainsi l'utilisation de ce mécanisme pour le guidage.

B. Algorithme JFDL

La méthode fonctionne en deux variantes principales :

1. JFDL "Naïf" (Naive JFDL) :

   • Nécessite un CM pré-entraîné capable de résoudre l'ODE pour la classe "null" ($\emptyset$) (c'est-à-dire un modèle capable de génération non conditionnelle).
   • Le processus consiste à interpoler les directions de débruitage conditionnelle ($u_{cls}$) et non conditionnelle ($u_{\emptyset}$) via un facteur d'échelle $\omega$ (similaire au CFG).
   • Une perte de cohérence ($L_{JFDL}$) est calculée pour forcer le modèle à apprendre cette trajectoire guidée.

2. JFDL "Aléatoire" (Random JFDL) :

   • Conçu pour les cas où le CM pré-entraîné n'a pas appris explicitement la classe $\emptyset$ (cas courant des modèles CT).
   • Au lieu d'utiliser une classe "null", la méthode remplace l'ancrage non conditionnel par un échantillon d'une classe aléatoire tirée de la distribution des données.
   • Cette variante s'avère surprenante : elle génère des images à plus fort contraste et s'aligne même mieux avec les effets du CFG que la version naïve, élargissant ainsi la généralité de la méthode.

C. Optimisation

• Le modèle est affiné (fine-tuned) en optimisant une perte multi-tâches combinant la perte d'entraînement original (ECT) et la perte JFDL.
• Un mécanisme de GradNorm est utilisé pour équilibrer dynamiquement les poids entre les deux pertes, assurant que le modèle conserve ses capacités de génération tout en acquérant le guidage.

---

3. Contributions Clés

1. Proposition de JFDL : Une méthode de guidage a posteriori pour les CMs qui ne dépend pas de modèles DM enseignants, comblant ainsi un vide majeur pour les modèles entraînés par CT.
2. Validation Théorique et Empirique : Fourniture d'analyses théoriques (basées sur les modèles de flux) et de vérifications expérimentales (tests de normalité de Shapiro-Wilk, Anderson-Darling, Kolmogorov-Smirnov) prouvant que les trajectoires hybrides respectent les hypothèses gaussiennes nécessaires au guidage.
3. Efficacité sur Modèles CT : Démonstration qu'un CM pré-entraîné sans conception explicite pour l'échantillonnage non conditionnel peut être efficacement adapté au guidage via la variante "Random JFDL".
4. Amélioration des performances : Réduction significative du FID (Fréchet Inception Distance) sur CIFAR-10 et ImageNet 64x64, prouvant que le guidage améliore la qualité perceptive même pour les modèles à une seule étape.

---

4. Résultats Expérimentaux

Les expériences ont été menées sur CIFAR-10 et ImageNet 64x64, en comparant le JFDL à une baseline ECT (Consistency Training) non guidée.

• Qualité de Génération (FID) :

  • CIFAR-10 : Le JFDL (Naïf et Random) réduit le FID en 1 étape par rapport à la baseline. Par exemple, pour le Random JFDL, le FID passe de 3.40 (baseline) à 3.24 (avec $\omega \approx 1.15$).
  • ImageNet 64x64 : Des améliorations notables sont observées, passant de 5.84 à 4.38 (Naïf) et 4.68 (Random) en 1 étape.
  • La courbe FID en fonction de la force de guidage $\omega$ suit la même tendance que le CFG dans les DMs : amélioration pour de faibles valeurs de $\omega$ (généralement < 2) avant une dégradation due à une perte de diversité.

• Efficacité du Calcul :

  • L'affinage (fine-tuning) avec JFDL ne consomme que ~7,5 % des données d'entraînement de la baseline ECT et environ 30 % des heures GPU, rendant la méthode très économe.

• Qualité Visuelle :

  • Les images générées montrent une fidélité accrue (contours plus nets, détails plus précis, comme observé sur les images "hotdog" et "jay" d'ImageNet) au prix d'une légère réduction de la diversité, reproduisant le compromis classique fidélité-diversité du CFG.

---

5. Signification et Impact

Ce travail est une avancée significative pour le domaine de la génération d'images rapide :

• Indépendance vis-à-vis des DMs : Il brise la dépendance des modèles de cohérence vis-à-vis des modèles de diffusion lourds pour obtenir un guidage, permettant aux modèles CT (plus rapides et moins coûteux à entraîner) d'atteindre un niveau de contrôle équivalent aux modèles distillés.
• Accessibilité : En rendant le guidage accessible a posteriori avec un coût de calcul minime, JFDL démocratise l'utilisation de modèles de génération rapide et de haute qualité pour des applications nécessitant un contrôle précis (visualisation scientifique, création de contenu).
• Nouvelles Perspectives : La méthode ouvre la voie à l'exploration de schémas de guidage adaptatifs pour l'échantillonnage multi-étapes et suggère que les concepts de flux de probabilité peuvent être exploités de manière plus flexible que prévu pour le contrôle des modèles génératifs.

En résumé, JFDL comble un fossé technologique crucial en permettant aux modèles de cohérence les plus rapides et les plus économes d'offrir un contrôle de génération de niveau industriel sans sacrifier leur efficacité.