Résumé Technique : Distillation de Transformers Dériveurs avec des Autoencodeurs de Représentation

Énoncé du Problème

Bien que les Autoencodeurs de Représentation (RAE) aient considérablement amélioré les performances des modèles génératifs basés sur le flux en utilisant des espaces latents riches sémantiquement et regroupés par étiquettes, dérivés de caractéristiques DINO pré-entraînées, la distillation de ces modèles en générateurs efficaces en une ou quelques étapes reste un défi.

Les auteurs identifient un décalage géométrique fondamental : les espaces latents des RAE présentent une anisotropie sévère et de fortes courbures dans leurs trajectoires d'Équations Différentielles Ordinaires (ODE) sous-jacentes par rapport aux espaces traditionnels des Autoencodeurs Variationnels (VAE). Les méthodes de distillation basées sur les trajectoires existantes reposent sur l'hypothon de trajectoires de l'enseignant lisses et presque rectilignes. Les trajectoires hautement courbes dans les espaces RAE rendent ces méthodes instables ou inefficaces. De plus, bien que le récent Modèle de Dérive (Drifting Model) propose une approche au niveau de la distribution qui évite de correspondre à des trajectoires courbes, sa formulation originale nécessite un Autoencodeur Masqué (MAE) auxiliaire, entraîné séparément, pour fournir des caractéristiques concentrées, ce qui introduit une surcharge de calcul significative.

Méthodologie

Le papier propose Drift-RAE, un cadre de distillation qui adapte le paradigme du Modèle de Dérive spécifiquement pour les espaces latents RAE, éliminant ainsi le besoin d'extracteurs de caractéristiques auxiliaires.

1. Analyse Théorique de la Géométrie et de la Dynamique

Les auteurs analysent d'abord quantitativement les espaces latents, démontrant que les caractéristiques RAE sont substantiellement plus anisotropes (Ratio de Participation et Entropie Spectrale plus faibles) que les caractéristiques SD-VAE. Ils prouvent théoriquement (Théorème 1) que dans des espaces de grande dimension, si les échantillons positifs sont trop dispersés, la composante attractive du champ de dérive tend à disparaître, conduisant à un champ dégénéré.

• Insight Clé : Alors que l'original Modèle de Dérive nécessite un MAE pour concentrer les caractéristiques dispersées, les espaces latents RAE sont naturellement concentrés sémantiquement. Cela permet au Modèle de Dérive d'opérer directement sur les latents RAE sans le module MAE auxiliaire.

2. Le Cadre Drift-RAE

La méthode distille un modèle de flux pré-entraîné $v_\theta$ en un générateur en une étape $G_\theta(z) = z - v_\theta(z, 1)$. Au lieu de faire correspondre les trajectoires d'ODE, le générateur est entraîné pour minimiser un objectif de dérive qui pousse les échantillons générés vers les distributions de données réelles et les éloigne des autres échantillons générés.

• Objectif : La perte est calculée jeton par jeton sur les représentations latentes RAE, en utilisant un champ de dérive normalisé $\tilde{V}$ dérivé d'échantillons positifs (réels) et négatifs (générés).

3. Modifications Pratiques pour la Stabilité

Pour aligner le Modèle de Dérive avec le cadre Diffusion-GAN et améliorer la stabilité de l'entraînement, les auteurs introduisent trois modifications spécifiques :

• Dimension Softmax : Les auteurs suppriment l'astuce du softmax bidirectionnel utilisée dans l'original Modèle de Dérive. Ils conservent un softmax unique sur les indices d'échantillons, ce qui s'aligne mieux avec la direction du gradient du potentiel log-sum-exp sous-jacent.
• Perturbation d'Entrée : Pour prévenir la disparition du gradient causée par des variétés non intersectantes entre les données réelles et générées, les échantillons générés sont perturbés par un bruit échantillonné à partir d'une distribution de Laplace avant de calculer le champ de dérive.
• Détachement Partiel des Échantillons Négatifs : Pour augmenter le nombre d'échantillons négatifs utilisés pour l'approximation de la distribution sans dépasser les budgets de mémoire GPU, les auteurs proposent d'utiliser un ensemble plus large d'échantillons négatifs ($N_{total\_neg}$) où seul un sous-ensemble ($N_{neg}$) porte les gradients, tandis que le reste est détaché. Cela améliore l'approximation empirique de la distribution générée.

Contributions Clés

1. Analyse Géométrique : Une démonstration quantitative que les espaces latents RAE possèdent une courbure et une anisotropie élevées, rendant la distillation basée sur les trajectoires standard instable.
2. Justification Théorique : Une preuve montrant que les Modèles de Dérive échouent sur des espaces trop dispersés (nécessitant un MAE dans la formulation originale) mais prospèrent dans les espaces sémantiques naturellement concentrés des RAE.
3. Drift-RAE : Une nouvelle méthode de distillation qui atteint une génération en une étape de pointe dans les espaces RAE sans nécessiter d'extracteur de caractéristiques MAE auxiliaire.
4. Améliorations de la Stabilité : Des insights théoriques et empiriques concernant la modification de l'objectif de Dérive (dimension softmax, perturbation d'entrée et détachement d'échantillons) pour s'aligner sur les principes de Diffusion-GAN et améliorer la convergence.

Résultats Expérimentaux

La méthode a été évaluée sur le jeu de données ImageNet 256 $\times$ 256 avec des backbones DiTDH-XL et DiTDH-L.

• Performance : Drift-RAE atteint un FID de 1,77 (avec DiTDH-XL) et 2,12 (avec DiTDH-L) en utilisant une seule étape de génération.
• Efficacité : Le modèle atteint ces résultats en seulement 10 000 étapes de distillation (environ 16 époques).
• Comparaison :
  • Il surpasse les méthodes de distillation précédentes dans les espaces RAE (ex: MF-RAE).
  • Il atteint un FID comparable et un score FDDINOv2 amélioré par rapport à l'original Modèle de Dérive, malgré l'absence de l'auxiliaire MAE coûteux en calcul.
  • Il reste compétitif avec les générateurs en une ou quelques étapes entraînés dans d'autres espaces latents (ex: Espace de Pixels, SD-VAE).

Signification et Revendications

Le papier affirme que la Dérive (Drifting) fournit un cadre efficace et prometteur pour distiller les modèles génératifs dans les espaces de représentation. La principale signification réside dans la résolution de l'instabilité de la distillation des modèles basés sur RAE en exploitant leur concentration sémantique inhérente, éliminant ainsi la dépendance aux extracteurs de caractéristiques auxiliaires.

Les auteurs notent modestement les limites, reconnaissant que leur analyse théorique repose sur des modèles de grande dimension simplifiés et que l'entraînement des Modèles de Dérive à partir de zéro (sans distillation) dans les espaces RAE reste un problème ouvert. Ils soulignent également que l'exigence de nombreux échantillons positifs de même classe peut poser des défis de scalabilité pour la génération texte-vers-image, suggérant cela comme une direction pour la recherche future.