Improving Diffusion Generalization with Weak-to-Strong Segmented Guidance

1. Problématique

Les modèles de diffusion (et les modèles de flux) sont devenus la norme pour la synthèse d'images grâce à leur capacité à générer des images réalistes via un processus de raffinement itératif. Cependant, ces modèles souffrent d'un biais d'exposition (exposure bias) : il existe un décalage entre l'objectif d'entraînement (simulation sans erreur, conditionnel à l'état initial $x_0$) et le processus d'échantillonnage itératif (où le modèle doit prédire l'état suivant en se basant sur une prédiction précédente bruitée).

Ce décalage entraîne une accumulation d'erreurs de gradient le long de la trajectoire d'échantillonnage, ce qui conduit à :

• Une généralisation médiocre, en particulier pour les tâches conditionnelles complexes (texte-à-image).
• La production d'échantillons hors distribution ou perceptuellement inacceptables.

Pour contrer ce phénomène, des techniques de guidage comme le Classifier-Free Guidance (CFG) et l'AutoGuidance (AG) sont utilisées.

• CFG (Condition-Dependent Guidance - CDG) : Utilise la différence entre la prédiction conditionnelle et non conditionnelle du même modèle. Il est robuste mais peut réduire la diversité des échantillons (effondrement de mode).
• AG (Condition-Agnostic Guidance - CAG) : Utilise un modèle "faible" (moins performant mais aligné sur la condition) pour guider un modèle "fort". Il améliore la diversité mais peut échouer à corriger les erreurs de structure globale si le modèle est sous-entraîné.

Le problème central : Les régimes opérationnels efficaces de ces méthodes ne sont pas bien compris. Le choix entre CDG et CAG est souvent arbitraire, et aucune méthode unique ne semble optimale pour tous les scénarios (bruit élevé vs bruit faible, conditions floues vs conditions précises).

2. Méthodologie

Les auteurs proposent une approche fondée sur le principe Weak-to-Strong (W2S) (de faible à fort), analysant d'abord les régimes de validité des méthodes existantes, puis proposant une solution hybride et une intégration en phase d'entraînement.

A. Analyse des régimes opérationnels (CDG vs CAG)

À travers des expériences synthétiques et une analyse sur ImageNet, les auteurs démontrent que l'efficacité du guidage dépend de deux facteurs clés : la granularité de la condition et la capacité d'ajustement (fitting capacity) du modèle.

• Bruit élevé (étapes précoces) : Le modèle doit identifier la bonne variété (manifold) conditionnelle. Le CDG (CFG) excelle ici car il force une séparation inter-classe forte et corrige les erreurs de structure globale.
• Bruit faible (étapes tardives) : Le modèle doit affiner les détails intra-classe. Le CAG (AG) est plus efficace ici car il évite l'effondrement de mode et préserve la diversité des détails sans dévier de la condition.

B. Guidage Segmenté (SGG - Segmented Guidance)

Pour combiner les avantages des deux approches, les auteurs proposent SGG, un mécanisme hybride qui segmente le processus d'échantillonnage dans le temps :

1. Phase à haut bruit ($t > \tau$) : Utilisation du CDG (CFG) pour localiser la bonne variété conditionnelle et corriger les erreurs structurelles.
2. Phase à faible bruit ($t \le \tau$) : Basculement vers le CAG (utilisant un modèle faible ou une perturbation du modèle) pour affiner les détails intra-classe et maintenir la diversité.

La formule du guidage $g(x_t, t, c)$ est définie comme :
$$g(x_t, t, c) = \begin{cases} v(x_t, t, c) - v(x_t, t, \emptyset) & \text{si } t > \tau \quad (\text{CDG}) \\ v(x_t, t, c) - \tilde{v}(x_t, t, c) & \text{si } t \le \tau \quad (\text{CAG}) \end{cases}$$
où $v$ est le modèle fort et $\tilde{v}$ le modèle faible.

C. Intégration à l'entraînement (Training Integration)

Au lieu d'appliquer SGG uniquement lors de l'inférence (ce qui coûte un calcul supplémentaire), les auteurs migrent ce principe dans la fonction de perte d'entraînement.

• L'objectif de régression standard est modifié pour inclure explicitement le terme de guidage W2S.
• Cela permet au modèle "fort" d'apprendre directement à extrapoler au-delà de l'ajustement conservateur du MSE, améliorant ainsi sa capacité de généralisation intrinsèque.
• Pour l'entraînement, ils proposent des variantes de modèles faibles efficaces, notamment une branche auxiliaire (BR - Branch) qui extrait un signal faible depuis une couche intermédiaire du transformateur, évitant ainsi d'entraîner un modèle séparé coûteux.

3. Contributions Clés

1. Analyse théorique et empirique : Une catégorisation claire des méthodes de guidage en CDG et CAG sous l'angle du principe Weak-to-Strong, démontrant leurs régimes de succès et d'échec respectifs via des expériences synthétiques et réelles.
2. Proposition de SGG : Introduction d'une méthode de guidage hybride simple mais efficace qui combine dynamiquement CDG et CAG selon le niveau de bruit, surpassant les méthodes existantes en inférence.
3. Migration vers l'entraînement : Première intégration systématique du principe W2S et de SGG dans l'objectif d'entraînement, permettant d'améliorer la généralisation des modèles de diffusion non guidés et de réduire la dépendance au guidage coûteux lors de l'inférence.
4. Validation complète : Des expériences exhaustives sur des modèles de pointe (SD3, SD3.5, SiT) et des tâches variées (ImageNet, texte-à-image, vidéo).

4. Résultats

Inférence (Texte-à-Image)

Sur les modèles SD3 et SD3.5 (benchmarks MS-COCO et LAION) :

• SGG surpasse toutes les variantes de guidage sans entraînement (CFG, SLG, AG, etc.).
• Il atteint un équilibre optimal entre l'adhérence au prompt (mesurée par HPSv2.1) et la qualité esthétique (Aesthetic Score).
  • Exemple (SD3.5 / MS-COCO) : SGG obtient un score HPSv2.1 de 29.736 et un score esthétique de 5.717, surpassant CFG (29.199 / 5.279) et SLG (27.295 / 5.714).

Entraînement (Généralisation)

Sur le modèle SiT-B/2 entraîné sur ImageNet :

• L'intégration de SGG dans l'entraînement améliore significativement la performance des modèles non guidés (unguided).
• Le modèle entraîné avec SGG atteint un FID de 4.58 sans guidage à l'inférence, surpassant le modèle de base guidé (FID 6.02 avec CFG).
• La variante BR (Branche) offre un gain de performance notable avec un surcoût computationnel minimal (seulement 2% de temps d'itération supplémentaire).

5. Signification et Impact

Ce travail apporte une compréhension fondamentale sur la dynamique des erreurs dans les modèles de diffusion et comment le guidage peut être optimisé temporellement.

• Efficacité : En déplaçant l'intelligence du guidage vers la phase d'entraînement, SGG permet de générer des images de haute qualité avec moins de calculs à l'inférence (moins d'appels au modèle faible).
• Généralité : La méthode est applicable à divers architectures (Transformers, modèles de flux) et modalités (images, vidéos).
• Principe Unificateur : Le cadre "Weak-to-Strong" offre une nouvelle perspective pour concevoir des algorithmes de génération, suggérant que l'utilisation stratégique de signaux "faibles" (imparfaits) peut renforcer la robustesse des modèles "forts".

En résumé, SGG résout le compromis traditionnel entre adhérence au texte et diversité esthétique en adaptant dynamiquement la stratégie de guidage au stade de la débruitage, tout en améliorant la robustesse intrinsèque du modèle grâce à une intégration intelligente dans l'entraînement.