Échelonnage de l'inférence des modèles de diffusion un nouveau paradigme

Introduction

Les avancées récentes dans les Grands Modèles de Langage (LLM) ont démontré l'efficacité de l'échelonnage pendant l'inférence. Des modèles tels que o1, o3, DeepSeek R1, QwQ et Step Reasoner mini ont montré qu'une puissance de calcul accrue pendant l'inférence peut améliorer considérablement les performances. Cela soulève la question : ce principe peut-il également être appliqué aux modèles de diffusion ?

Une équipe dirigée par Xie Saining à l'Université de New York a exploré cette question. Grâce à une étude systématique utilisant un cadre de recherche général, ils ont constaté que l'échelonnage au moment de l'inférence est effectivement efficace pour les modèles de diffusion. L'augmentation du calcul pendant l'inférence conduit à une amélioration significative de la qualité des échantillons générés. De plus, la complexité des images permet différentes combinaisons de composants au sein du cadre, adaptées à divers scénarios d'application.

Principales conclusions

• L'échelonnage au moment de l'inférence est efficace pour les modèles de diffusion : Allouer plus de ressources de calcul pendant l'inférence entraîne des échantillons de meilleure qualité.
• Flexibilité des combinaisons de composants : Le cadre permet différentes configurations de composants, répondant à diverses applications.
• Au-delà des étapes de débruitage : La recherche suggère que la recherche d'un meilleur bruit pendant l'échantillonnage est une autre dimension pour l'échelonnage NFE, au-delà de l'augmentation simple des étapes de débruitage.
• Deux axes de conception : Le cadre se concentre sur deux axes de conception clés :
  • Vérificateurs : Fournir un retour d'information pendant le processus de recherche.
  • Algorithmes : Trouver de meilleurs candidats de bruit.

Méthodologie de recherche

L'équipe a exploré trois scénarios différents pour les vérificateurs, simulant divers cas d'utilisation :

1. Scénarios où des informations privilégiées sur l'évaluation finale sont disponibles.
2. Scénarios où des informations conditionnelles sont disponibles pour guider la génération.
3. Scénarios sans informations supplémentaires disponibles.

Pour les algorithmes, ils ont étudié :

1. Recherche aléatoire : Sélectionner le meilleur parmi un ensemble fixe de candidats.
2. Recherche d'ordre zéro : Améliorer de manière itérative les candidats de bruit à l'aide du retour d'information du vérificateur.
3. Recherche de chemin : Améliorer de manière itérative les trajectoires d'échantillonnage de diffusion à l'aide du retour d'information du vérificateur.

L'étude a initialement exploré ces conceptions dans une configuration de génération conditionnelle de classe ImageNet relativement simple. Par la suite, ils ont appliqué ces conceptions à une génération conditionnelle de texte à plus grande échelle et ont évalué leur cadre proposé.

Échelonnement du temps d'inférence

L'article propose un cadre pour échelonner le temps d'inférence dans les modèles de diffusion, en encadrant le défi comme une recherche du bruit d'échantillonnage optimal. Le processus implique deux composants essentiels :

• Vérificateurs : Ce sont des modèles pré-entraînés qui évaluent la qualité des échantillons générés. Ils prennent en entrée les échantillons générés, ainsi que des conditions facultatives, et produisent un score scalaire.
• Algorithmes : Ces algorithmes utilisent les scores des vérificateurs pour trouver de meilleurs échantillons candidats. La fonction prend en entrée un vérificateur (V), un modèle de diffusion pré-entraîné (_θ) et un ensemble d'échantillons et de conditions générés, produisant le meilleur bruit initial.

Le budget total d'inférence est mesuré par le nombre total d'évaluations de fonctions (NFE), comprenant à la fois les étapes de débruitage et les coûts de recherche.

Vérificateurs de recherche

Les chercheurs ont commencé avec un vérificateur Oracle, qui possède des informations complètes sur l'évaluation finale des échantillons sélectionnés. Pour ImageNet, cela comprenait des mesures telles que FID et IS. Ils ont ensuite exploré des modèles pré-entraînés plus accessibles comme vérificateurs supervisés, tels que CLIP et DINO. Ces modèles ont été utilisés pour classifier les échantillons, en sélectionnant l'échantillon avec le logit le plus élevé correspondant à l'étiquette de classe.

Cependant, ils ont observé que ces classificateurs, fonctionnant point par point, ne s'alignent que partiellement sur les objectifs du score FID. Cela a conduit à une réduction de la variance des échantillons et à un effondrement de mode à mesure que le calcul augmentait. Ce phénomène, appelé « piratage du vérificateur », a été accéléré par l'espace de recherche non contraint de l'algorithme de recherche aléatoire.

Il est intéressant de noter que l'étude a révélé que les vérificateurs n'ont pas nécessairement besoin d'informations conditionnelles pour guider efficacement la recherche. Ils ont observé une forte corrélation entre les logits des classificateurs DINO/CLIP et la similarité cosinus de l'espace de caractéristiques entre la prédiction x à un faible niveau de bruit et l'échantillon propre final. Cela a conduit à l'utilisation de vérificateurs auto-supervisés, qui ne nécessitent pas d'informations conditionnelles supplémentaires et ont toujours présenté un comportement de mise à l'échelle efficace.

Algorithmes de recherche

Pour atténuer le piratage des vérificateurs, les chercheurs ont exploré des algorithmes de recherche plus affinés qui optimisent progressivement les échantillons candidats. Cela comprenait une méthode de recherche d'ordre zéro :

1. Commencer avec un bruit gaussien aléatoire comme point pivot.
2. Trouver N candidats dans le voisinage du point pivot.
3. Exécuter les candidats via le solveur ODE pour obtenir des échantillons et des scores de vérificateur.
4. Mettre à jour le point pivot avec le meilleur candidat et répéter les étapes 1 à 3.

Ils ont également étudié un algorithme de recherche de chemin, qui explore la possibilité de rechercher le long de la trajectoire d'échantillonnage :

1. Échantillonner N échantillons de bruit initial et exécuter le solveur ODE jusqu'à un niveau de bruit σ.
2. Ajouter du bruit à chaque échantillon et simuler un processus de bruitage avant.
3. Exécuter un solveur ODE sur chaque échantillon bruité et conserver les N meilleurs candidats en fonction des scores du vérificateur, en répétant jusqu'à ce que le solveur ODE atteigne σ=0.
4. Rechercher aléatoirement les N échantillons restants et conserver le meilleur.

Les algorithmes de recherche d'ordre zéro et de chemin maintiennent une forte localité par rapport à la recherche aléatoire.

Échelonnement dans les scénarios texte-image

L'équipe a examiné les capacités d'échelonnage du cadre de recherche dans des tâches de texte-image à plus grande échelle. Ils ont utilisé les ensembles de données DrawBench et T2I-CompBench pour l'évaluation, avec le modèle FLUX.1-dev comme base. Ils ont également élargi la sélection de vérificateurs supervisés, notamment Aesthetic Score Predictor, CLIPScore et ImageReward. De plus, ils ont créé un ensemble de vérificateurs en combinant ces trois vérificateurs.

Analyse : alignement vérificateur-tâche

L'étude a comparé les résultats de diverses combinaisons vérificateur-algorithme sur différents ensembles de données. Sur DrawBench, ils ont constaté que l'utilisation de tous les vérificateurs améliorait généralement la qualité des échantillons. Cependant, ils ont observé que l'utilisation isolée des vérificateurs Aesthetic et CLIP pouvait entraîner une suradaptation de leurs biais, ce qui avait des impacts négatifs l'un sur l'autre. Cela découle d'une inadéquation dans leur objectif d'évaluation :

• Aesthetic Score se concentre sur la qualité visuelle, favorisant souvent les images très stylisées,
• tandis que CLIP privilégie l'alignement visuel-texte, sacrifiant parfois la qualité visuelle.

Ils ont noté que certains vérificateurs sont plus adaptés à des tâches spécifiques et que l'efficacité d'un vérificateur dépend de son alignement avec les exigences de la tâche.

Performances de l'algorithme

Les trois algorithmes de recherche (aléatoire, d'ordre zéro et de chemin) ont tous amélioré efficacement la qualité de l'échantillonnage sur DrawBench. Cependant, la recherche aléatoire a surpassé certains aspects en raison de la nature locale des deux autres méthodes. La recherche aléatoire a convergé plus rapidement vers le biais du vérificateur, tandis que les deux autres algorithmes nécessitent une amélioration sur des candidats moins qu'optimaux.

Compatibilité avec le réglage fin

L'équipe a étudié la compatibilité de leur méthode de recherche avec les modèles réglés avec précision. Ils ont utilisé un modèle Stable Diffusion XL réglé avec précision DPO et ont constaté que la méthode de recherche pouvait être généralisée à différents modèles et améliorer les performances des modèles déjà alignés.

Effets des différentes dimensions du calcul d'inférence

L'étude a exploré comment différents aspects du calcul d'inférence affectent les résultats :

• Nombre d'itérations de recherche : L'augmentation des itérations rapproche le bruit de l'optimum.
• Calcul par itération de recherche : L'ajustement du nombre d'étapes de débruitage par itération révèle différentes régions optimales en termes de calcul.
• Calcul de génération finale : L'équipe a utilisé des paramètres optimaux pour les étapes de débruitage finales afin d'assurer la meilleure qualité d'échantillon finale.

Efficacité de l'investissement dans le calcul

Les chercheurs ont exploré l'efficacité de l'échelonnage au moment de l'inférence sur des modèles de diffusion plus petits. Ils ont constaté que, pour ImageNet, l'échelonnage de modèles plus petits peut être très efficace. Dans certains cas, la recherche sur un modèle plus petit peut surpasser des modèles plus grands sans recherche. Cependant, l'efficacité dépend des performances de base du modèle plus petit.

Dans les paramètres basés sur le texte, PixArt-Σ, utilisant seulement une fraction du calcul, a surpassé FLUX-1.dev. Ces résultats démontrent que d'importantes ressources de calcul dépensées pendant l'entraînement peuvent être compensées par de plus petites quantités de calcul pendant la génération, ce qui donne des échantillons de meilleure qualité plus efficacement.