Résumé technique : Amélioration des références avec les Autoencodeurs de Représentation (RAEv2)

Énoncé du problème

Les Autoencodeurs de Représentation (RAE) sont apparus comme un cadre prometteur pour remplacer les Autoencodeurs Variationnels (VAE) traditionnels dans l'entraînement des transformateurs de diffusion, visant à unifier la tokenisation pour la compréhension visuelle et la génération. Cependant, la formulation originale du RAE fait face à trois goulots d'étranglement critiques entravant son adoption pratique :

1. Reconstruction sous-optimale : Les performances de reconstruction sont en retard par rapport aux VAE spécialisés, en particulier lors de l'entraînement sans données spécialisées (par exemple, texte ou visages).
2. Incompatibilité avec le guidage : Le RAE éprouve des difficultés avec le Guidage Sans Classificateur (CFG) traditionnel, nécessitant l'entraînement d'un modèle de diffusion secondaire et plus faible pour l'AutoGuidage (AG). Cela ajoute une surcharge computationnelle et une complexité significatives.
3. Représentations d'encodeur sous-exploitées : Les travaux antérieurs se contentent d'utiliser uniquement la dernière couche d'un encodeur de vision pré-entraîné comme représentation latente, ignorant les abstractions riches et diversifiées présentes dans les couches intermédiaires.

Méthodologie

L'article propose RAEv2, une référence améliorée qui résout ces problèmes grâce à trois insights clés et changements de conception, tout en maintenant une approche sans entraînement pour l'encodeur.

1. Encodeur de Représentation Généralisé

Au lieu de se fier uniquement à la dernière couche ($K=1$) d'un encodeur de vision pré-entraîné, les auteurs proposent une formulation généralisée où la représentation latente est la somme des $K$ dernières couches de l'encodeur.

• Mécanisme : La sortie est définie comme $\mathbf{x} = \sum_{\ell=L-K+1}^{L} \mathbf{z}_\ell$, où $\mathbf{z}_\ell$ représente la carte de caractéristiques à la couche $\ell$.
• Justification : Différentes couches capturent des caractéristiques complémentaires ; les couches précoces à moyennes conservent des structures spatiales plus fines, tandis que les couches ultérieures mettent l'accent sur la sémantique globale. Les agréger préserve la structure géométrique sans augmenter la dimensionnalité latente ni nécessiter de paramètres appris supplémentaires.
• Implémentation : Les auteurs comparent une addition simple (MLS) à une projection par matrice aléatoire (MLR) et constatent que l'addition simple produit systématiquement de meilleures performances de génération tout en améliorant la reconstruction.

2. Mécanismes complémentaires du RAE et du REPA

L'article remet en question l'hypothèse prévalente selon laquelle le RAE élimine le besoin d'Alignement de Représentation (REPA). Le REPA distille la même représentation pré-entraînée dans les couches intermédiaires de diffusion.

• Découverte : L'analyse empirique sur 27 encodeurs de vision révèle que le RAE et le REPA fonctionnent via des mécanismes complémentaires.
  • RAE fournit un espace latent riche sémantiquement (exploitant la sémantique globale, mesurée par la précision du sondage linéaire).
  • REPA régularise la structure spatiale des caractéristiques de diffusion intermédiaires (mesurée par la Similarité de Distance Locale, LDS).
• Résultat : L'utilisation de la même représentation à la fois comme encodeur (RAE) et comme cible pour les couches intermédiaires (REPA) améliore systématiquement les performances de génération. Cette synergie permet à des encodeurs plus puissants (par exemple, DINOv3-L), excellant à la fois dans les métriques globales et spatiales, d'obtenir des résultats de génération supérieurs.

3. REPA comme AutoGuidage (prédiction de x)

Le RAE original nécessite un modèle faible séparé pour le guidage. Les auteurs reformulent la tête REPA pour permettre un guidage interne.

• Insight : Dans l'espace latent du RAE, la tête de projection REPA (qui mappe les caractéristiques des couches précoces vers le latent propre) effectue efficacement une prédiction de x.
• Mécanisme : En re-paramétrant la sortie complète du modèle pour être également une prédiction de x, la prédiction de la tête REPA ($\hat{\mathbf{x}}_{repa}$) sert de référence « plus faible » par rapport à la prédiction du modèle complet ($\hat{\mathbf{x}}_{full}$).
• Formule de guidage : Le guidage est appliqué internement lors d'un seul passage avant : $\hat{\mathbf{x}}_{guided} = \hat{\mathbf{x}}_{full} + w \cdot (\hat{\mathbf{x}}_{full} - \hat{\mathbf{x}}_{repa})$.
• Avantage : Cela élimine le besoin d'un modèle AutoGuidage séparé et évite le passage avant supplémentaire requis par le CFG, réduisant efficacement de moitié le Nombre d'Évaluations de Fonction (NFE) pour le guidage.

Résultats clés

RAEv2 démontre des améliorations significatives en termes de vitesse de convergence, de qualité de génération et de reconstruction sur ImageNet, le texte-vers-image et les modèles de monde de navigation.

• Efficacité de l'entraînement : RAEv2 atteint une convergence 10 fois plus rapide par rapport au RAE original. Il atteint un gFID guidé de pointe de 1,06 en seulement 80 époques sur ImageNet-256, alors que le RAE original nécessitait 800 époques pour atteindre 3,26.
• Nouvelle métrique d'efficacité (EPFID) : L'article introduit EPFID@k (époques pour atteindre un gFID non guidé $\le k$) comme métrique plus pratique que le gFID absolu. RAEv2 atteint un EPFID@2 de 35 époques, une amélioration massive par rapport au RAE original (177 époques) et aux travaux antérieurs (>480 époques).
• Qualité de génération : Sur la Distance Fréchet de Représentation ($FD_r$) récemment proposée, RAEv2 obtient un score de 2,17 en 80 époques, surpassant le précédent meilleur score de 3,26 (entraîné pendant 800 époques) sans aucun post-entraînement.
• Reconstruction : La formulation généralisée (spécifiquement avec $K=23$) atteint un compromis optimal de Pareto entre reconstruction (rFID 0,18) et génération, surpassant les VAE propriétaires (par exemple, Flux-VAE, SD3.5) malgré un entraînement uniquement sur ImageNet.
• Généralisation :
  • Texte-vers-image : RAEv2 montre des améliorations constantes par rapport à Flux-VAE et au RAE original sur GenEval, DPG-Bench et GenAI-Bench.
  • Modèles de monde de navigation : Dans les tâches de prédiction de futurs cadres (jeu de données RECON), RAEv2-NWM atteint un FVD de 105,61, surpassant significativement NWM (200,97) et RAE (312,01), avec un scintillement réduit dans les déroulements autorégressifs.

Importance et revendications

L'article revendique que RAEv2 fournit une référence simplifiée et hautement efficace pour les autoencodeurs de représentation. Ses contributions principales sont :

1. Simplification : Il élimine le besoin de fine-tuning de l'encodeur, de données d'entraînement spécialisées ou de modèles de guidage séparés.
2. Efficacité : Il réduit drastiquement le temps d'entraînement requis pour atteindre des performances de pointe, déplaçant l'accent des améliorations de métriques absolues vers l'efficacité de l'entraînement (EPFID).
3. Tokenisation unifiée : En préservant les capacités de compréhension sémantique globale de l'encodeur gelé (vérifiées via le sondage linéaire) tout en améliorant la génération et la reconstruction, il renforce la viabilité de l'utilisation d'encodeurs de vision pré-entraînés comme schéma de tokenisation unifié pour les tâches de compréhension et de génération.

Les auteurs espèrent que ces insights faciliteront l'adoption pratique des autoencodeurs de représentation dans divers scénarios de modélisation générative.