Résumé Technique : Dérive des Caractéristiques de l'Enseignant (TFD)

Énoncé du Problème

Les modèles de diffusion et d'appariement de flux préentraînés génèrent des images de haute fidélité, mais souffrent de coûts d'inférence élevés, nécessitant souvent des dizaines ou des centaines de passes avant pour générer une seule image. Bien que des méthodes de distillation existent pour compresser ces enseignants multi-étapes en générateurs efficaces à une ou quelques étapes, les approches actuelles font face à une complexité significative. Des méthodes comme la Distillation par Appariement de Distribution (DMD) et ses variantes s'appuient sur des réseaux auxiliaires (par exemple, réseaux de score factice, discriminateurs adverses) ou des calendriers d'entraînement multi-étapes. D'autres approches, telles que les Modèles de Cohérence, peuvent souffrir d'un entraînement instable ou nécessiter des produits complexes entre vecteurs et jacobiennes. De plus, les objectifs récents de « Modèle Dérivant », qui offrent une fonction de perte d'attraction-répulsion directe pour la génération en une étape, dépendent généralement d'un extracteur de caractéristiques préentraîné externe (par exemple, modèles auto-supervisés ou latent-MAE) pour définir l'espace de représentation. Cela introduit une surcharge d'entraînement supplémentaire et une complexité architecturale.

Le problème central abordé est de savoir comment réaliser une distillation de diffusion efficace et de haute qualité en une seule étape en utilisant un objectif unique, sans recourir à des réseaux auxiliaires, des pipelines d'optimisation complexes ou des extracteurs de caractéristiques externes.

Méthodologie : Dérive des Caractéristiques de l'Enseignant (TFD)

Les auteurs proposent la Dérive des Caractéristiques de l'Enseignant (TFD), un cadre qui simplifie la distillation en une étape en exploitant les représentations internes du modèle enseignant préentraîné lui-même.

1. Espace de Représentation Induit par l'Enseignant

Contrairement au Modèle Dérivant original, qui nécessite un encodeur de caractéristiques externe, la TFD utilise les états intermédiaires cachés de l'enseignant préentraîné figé comme espace de caractéristiques pour calculer la perte de dérive.

• Mécanisme : Pour un échantillon généré $x$ et un échantillon réel $x^+$, les caractéristiques sont extraites de couches spécifiques $\ell$ du réseau enseignant $T$ : $\phi^\ell_T(x, c) = h^\ell_T(x, t_f, c)$.
• Perte de Dérive : Le champ de dérive est calculé directement dans cet espace induit par l'enseignant. Le générateur $f_\theta$ est optimisé pour minimiser la distance entre sa sortie et une « cible transportée figée » définie par le champ d'attraction-répulsion de l'espace de caractéristiques de l'enseignant. Cela aligne l'objectif de distillation avec le modèle en cours de distillation, éliminant le besoin d'un réseau de représentation séparé.

2. Extraction de Caractéristiques Bruitées

Une découverte empirique clé est que l'extraction de caractéristiques à partir d'entrées modérément bruitées donne de meilleurs résultats que l'utilisation d'échantillons propres.

• Avant de soumettre un échantillon à l'enseignant pour l'extraction de caractéristiques, il est perturbé par le processus de bruitage avant de l'enseignant : $x_{tf} = x + \sigma_{tf} \xi$.
• Les auteurs émettent l'hypothèse que cette perturbation lisse la géométrie des caractéristiques de l'enseignant, supprimant les écarts haute fréquence et rendant l'espace de représentation plus adapté à l'appariement au niveau de la distribution. Les expériences indiquent qu'un niveau de bruit de $\sigma_{tf} = 0.1$ est optimal.

3. Régularisation par Marge d'Ancrage

Pour résoudre le problème de l'effondrement des modes (où le générateur échoue à couvrir les régions diversifiées soutenues par l'enseignant), les auteurs introduisent une perte de marge d'ancrage légère.

• Ancres : Un ensemble de caractéristiques d'ancrage $\{a_i\}$ est extrait d'échantillons réels ou générés par l'enseignant.
• Mécanisme : La perte encourage les échantillons générés à occuper des régions autour de ces ancres. Elle ajuste de manière adaptative le « support » requis pour chaque ancre en fonction de la densité locale de la distribution des ancres elle-même.
• Objectif : Cela fournit un signal de couverture complémentaire à l'objectif local de dérive d'attraction-répulsion, garantissant que le générateur étudiant couvre des régions diversifiées soutenues par l'enseignant sans introduire de nouveaux réseaux.

La fonction objectif finale combine la perte de Dérive des Caractéristiques de l'Enseignant ($L_{TFD}$) et la perte de marge d'ancrage ($L_{anchor}$) :
$$L = L_{TFD} + \lambda_{anchor} L_{anchor}$$

Contributions Clés

1. Cadre Simplifié : La TFD propose un cadre de distillation de diffusion en une étape qui calcule l'objectif de dérive directement dans l'espace de représentation caché de l'enseignant préentraîné, supprimant le besoin d'un réseau de représentation supplémentaire.
2. Extraction de Caractéristiques Bruitées : Le travail démontre que l'extraction des caractéristiques de l'enseignant à partir d'entrées modérément bruitées améliore la géométrie des caractéristiques pour la distillation, conduisant à de meilleurs scores FID.
3. Atténuation de l'Effondrement des Modes : Une perte de marge d'ancrage légère est introduite pour encourager la diversité des échantillons et la couverture des régions de caractéristiques soutenues par l'enseignant, répondant à un mode d'échec courant en distillation.
4. Performance : La méthode atteint une qualité d'image et une diversité compétitives avec une seule passe avant, simplifiant considérablement le pipeline global de distillation par rapport aux bases de référence de l'état de l'art.

Résultats Expérimentaux

Les auteurs ont évalué la TFD sur ImageNet-64×64 (conditionné par la classe) et la génération SDXL (texte vers image).

• ImageNet-64×64 :

  • La TFD a atteint un FID de 1,58 avec une seule passe avant après 200k itérations d'entraînement.
  • Ce résultat est compétitif avec DMD2 (FID 1,51) et surpasse significativement d'autres bases de référence en une étape comme DMD (2,62), EMD (2,20) et CTM (1,92).
  • Convergence : La TFD a atteint un FID $\le$ 10 en 1,0k de mises à jour, tandis que la base de référence DMD2 nécessitait 5,5k de mises à jour (5,5 fois plus rapide). Elle a atteint un FID $\le$ 3 en 12,5k de mises à jour contre 24,0k pour DMD2.

• Texte vers Image SDXL :

  • Sur le jeu de données COCO 2014, la TFD a atteint un FID de 18,39 avec une seule passe avant, surpassant d'autres variantes SDXL distillées (par exemple, DMD2 à 19,01, SDXL-Turbo à 24,57).
  • Elle a maintenu un score CLIP compétitif de 0,332, indiquant une forte alignement texte-image.
  • Les résultats qualitatifs ont montré que la TFD préserve une structure cohérente, des détails visuels et un alignement avec l'invite comparables aux bases de référence distillées en 4 étapes et à l'enseignant en 50 étapes, malgré l'utilisation d'une seule étape.

Importance et Revendications

L'article revendique que la TFD offre une alternative plus simple et plus efficace aux cadres de distillation existants. En utilisant le modèle enseignant lui-même comme extracteur de caractéristiques, la TFD évite la complexité de l'entraînement de réseaux auxiliaires ou de la conception de pipelines multi-étapes. L'introduction de l'extraction de caractéristiques bruitées et de la régularisation par marge d'ancrage répond à des défis spécifiques en matière de géométrie de représentation et de couverture des modes, respectivement.

Les auteurs positionnent la TFD comme une méthode qui réalise une génération efficace en une étape avec une qualité d'image et une diversité compétitives, démontrant qu'un objectif unique dans l'espace de représentation natif de l'enseignant suffit pour une distillation de haute fidélité. Le travail suggère que le concept de « dérive » peut être réalisé efficacement sans encodeurs de caractéristiques externes, à condition que les états internes de l'enseignant soient utilisés correctement.

Limitations : Les auteurs notent que le travail actuel se concentre sur des enseignants basés sur U-Net et n'a pas encore exploré les architectures de Transformateur de Diffusion (DiT). De plus, la sélection des couches de caractéristiques et des niveaux de bruit manque actuellement de guidance théorique, et la méthode n'a pas été testée sur des jeux de données à plus grande échelle ou des paramètres de résolution plus élevés au-delà des expériences actuelles.