It's Never Too Late: Noise Optimization for Collapse Recovery in Trained Diffusion Models

1. Énoncé du problème

Les modèles de diffusion texte-vers-image contemporains (par exemple, SDXL-Turbo, Flux.1) souffrent d'effondrement de mode. Lors de l'échantillonnage de plusieurs images à partir d'une même invite textuelle en utilisant différentes graines aléatoires (bruit i.i.d.), les modèles produisent souvent des sorties presque identiques avec une variation visuelle minimale.

• Limitations actuelles : Les solutions existantes reposent sur :
  • Mécanismes de guidage : Pilotage du modèle via un conditionnement modifié (par exemple, Particle Guidance, DPP), qui sacrifie souvent la qualité au profit de la diversité.
  • Méthodes basées sur la recherche : Génération d'un grand pool de candidats et élagage de ceux-ci (par exemple, "Best-of-N"), ce qui est coûteux en calcul et gourmand en mémoire.
  • Ingénierie de l'invite : Modification de l'entrée textuelle, qui ne peut pas facilement capturer toutes les variations dans l'espace de sortie.
• Question centrale : Peut-on optimiser directement le bruit initial ($x_0$) pour se remettre de l'effondrement de mode et générer des images diversifiées et de haute qualité sans modifier le modèle pré-entraîné ni l'invite utilisateur ?

2. Méthodologie

Les auteurs proposent un cadre d'optimisation de bruit de bout en bout qui traite le vecteur de bruit initial comme un paramètre apprenable. Le processus implique :

A. Objectif d'optimisation du bruit

Étant donné une invite fixe $c$ et un modèle de diffusion figé $g_\theta$, la méthode optimise un lot d'échantillons de bruit initial $\{x_0^{(i)}\}_{i=1}^B$. La fonction objectif minimise une perte $\mathcal{L}$ composée de trois termes :

1. Récompense de qualité ($r_s$) : Garantit que les images individuelles respectent l'invite et maintiennent une haute fidélité (par exemple, en utilisant CLIPScore ou HPSv2).
2. Métrique de diversité ($v_\mathcal{B}$) : Maximise la dissimilarité entre les images du lot. Le papier explore diverses métriques :
   • Par paires : Patchs de DINOv2, DreamSim, LPIPS, histogrammes de couleurs.
   • Niveau d'ensemble : Processus ponctuels déterminantaux (DPP) et Score Vendi, qui mesurent le nombre effectif d'éléments uniques dans un ensemble.
3. Régularisation ($\text{reg}$) : Empêche le bruit optimisé de s'écarter trop de la distribution a priori $N(0, I)$. Cela est réalisé en pénalisant les écarts par rapport à la distribution $\chi_d$ de la norme du bruit.

L'optimisation est effectuée par rétropropagation à travers l'échantillonneur figé $g_\theta$.

B. Stratégie d'initialisation du bruit (bruit rose)

Une idée clé du papier est l'analyse de l'évolution du bruit. Au cours de l'optimisation, les auteurs ont observé que les changements les plus significatifs se produisent dans les composantes de basse fréquence du spectre du bruit, tandis que les hautes fréquences restent relativement stables.

• Hypothèse : Les images naturelles suivent un spectre de puissance en $1/f$ (plus de puissance dans les basses fréquences). Le bruit blanc standard est spectralement plat.
• Solution : Les auteurs introduisent une initialisation par bruit rose. Ils transforment le bruit blanc dans le domaine fréquentiel via une FFT, appliquent un facteur de mise à l'échelle de puissance $f^{-\alpha/2}$ (où $\alpha \approx 0,2$), puis effectuent la transformation inverse.
• Avantage : Cette initialisation aligne le bruit de départ avec les propriétés statistiques des images naturelles et la dynamique d'optimisation, conduisant à une convergence plus rapide et à une plus grande diversité.

C. Modes d'optimisation

• Optimisation par lot : Optimise conjointement un ensemble d'images (par exemple, 4) pour maximiser la diversité au niveau de l'ensemble.
• Optimisation séquentielle : Génère les images une par une, en optimisant chacune pour qu'elle diffère des sorties précédemment générées. Cela réduit la surcharge mémoire et permet de générer de grands ensembles diversifiés.

3. Contributions clés

1. Optimisation du bruit pour la diversité : Démontre que l'optimisation directe du bruit initial est une stratégie supérieure pour se remettre de l'effondrement de mode par rapport aux méthodes basées sur le guidage ou la recherche.
2. Objectifs de diversité au niveau de l'ensemble : Identifie que les DPP et le Score Vendi sont plus efficaces que les métriques par paires simples pour optimiser des ensembles diversifiés, car ils empêchent la métrique "moyenne par paires" d'être contournée en rendant une image très différente tout en gardant les autres similaires.
3. Initialisation par bruit rose : Propose une stratégie simple et sans paramètre (utilisant le bruit rose) qui améliore constamment la diversité à travers différents modèles (SDXL-Turbo, Flux.1) et lignes de base, même sans optimisation supplémentaire.
4. Compromis qualité-diversité : Montre que leur méthode peut augmenter la diversité sans sacrifier la qualité de l'image ni l'adhésion à l'invite, surmontant le compromis typique observé dans les méthodes basées sur le guidage.

4. Résultats expérimentaux

La méthode a été évaluée sur GenEval, DPG-Bench (invites complexes) et T2I-CompBench en utilisant des modèles tels que SDXL-Turbo et Flux.1 [schnell].

• Performance quantitative :
  • Diversité : Surpasse l'échantillonnage i.i.d. et la méthode de recherche de l'état de l'art [Parmar et al., 2025] sur toutes les métriques (DINO, DreamSim, LPIPS, Score Vendi).
  • Qualité : Maintient ou améliore légèrement les scores de qualité d'image (HPSv2, CLIPScore) par rapport aux lignes de base.
  • Efficacité : Avec l'initialisation par bruit rose, la méthode atteint des scores de diversité plus élevés que la méthode de recherche de base en beaucoup moins d'itérations (par exemple, 6–8 itérations contre 9–12 pour le bruit blanc).
• Résultats qualitatifs :
  • Génère des variations diversifiées dans la pose des objets, la couleur, le style et l'arrière-plan tout en adhérant strictement à des invites complexes.
  • Gère avec succès des invites complexes (par exemple, "Une scène pittoresque d'automne...") où l'échantillonnage i.i.d. échoue à produire de la variation.
• Étude utilisateur : Les évaluateurs humains ont préféré les sorties de la méthode proposée (en particulier celles optimisées avec DPP/Vendi) par rapport à l'échantillonnage i.i.d. et à la ligne de base de recherche dans 85–91 % des cas.

5. Importance

Ce travail déplace le paradigme de la génération de diversité dans les modèles de diffusion des stratégies d'échantillonnage (attendre une bonne graine) ou du guidage (modifier le chemin de génération) vers l'optimisation de l'entrée.

• Applicabilité générale : L'approche est agnostique au modèle et fonctionne à la fois sur les modèles de diffusion distillés par étapes (rapides) et sur les modèles de diffusion standard.
• Efficacité : En optimisant directement le bruit, elle évite la surcharge mémoire massive de la génération et du stockage de centaines de candidats pour l'élagage.
• Insight sur l'espace latent : La découverte que les composantes de bruit de basse fréquence pilotent la diversité offre une nouvelle compréhension de la façon dont les modèles de diffusion interprètent le bruit initial, suggérant que la "structure" est établie tôt dans le domaine fréquentiel.

En conclusion, le papier présente une solution robuste, efficace et de haute qualité au problème de l'effondrement de mode, prouvant qu'il n'est "jamais trop tard" pour optimiser le bruit d'entrée afin de débloquer le plein potentiel de diversité des modèles de diffusion entraînés.