Résumé Technique : Maîtrise des Tokens Anormaux dans les Transformers de Diffusion

Énoncé du Problème

L'article examine le phénomène des tokens anormaux au sein des Transformers de Diffusion (DiT) utilisés pour la génération d'images. Alors que des travaux antérieurs (par exemple, Darcet et al.) ont établi que les Transformers de Vision (ViT) dans des tâches de reconnaissance produisent souvent un petit sous-ensemble de tokens aux normes de caractéristiques anormalement élevées, absorbant une attention disproportionnée tout en véhiculant une information locale limitée, le rôle de ces tokens dans les modèles génératifs restait inexploré.

Les auteurs identifient deux lacunes critiques :

1. Omniprésence dans les Pipelines Génératifs : Les tokens anormaux apparaissent non seulement dans les encodeurs ViT préentraînés (tokeniseurs) des pipelines modernes d'Autoencodeurs de Représentation (RAE), mais également en interne au sein des Transformers de Diffusion (débruiteurs) eux-mêmes.
2. Nature du Défaut : La dégradation causée par ces tokens n'est pas simplement le résultat de quelques valeurs de perte extrêmes. Au contraire, les auteurs émettent l'hypothèse que les tokens anormaux sont symptomatiques d'une sémantique de patchs locaux corrompue. Le simple masquage des tokens à forte norme durant l'entraînement échoue à améliorer les performances, suggérant que le problème réside dans la perte de structure spatiale et d'information locale plutôt que dans une simple instabilité numérique.

Méthodologie : Registres à Double Étape (DSR)

Pour remédier à ces problèmes, les auteurs proposent les Registres à Double Étape (DSR), une stratégie d'intervention unifiée appliquant des tokens de registre aux deux composantes du pipeline RAE-DiT : l'encodeur de représentation et le générateur de diffusion.

1. Registres dans les Encodeurs de Vision

Les auteurs traitent les tokens anormaux à l'étape de l'encodeur (le tokeniseur) en utilisant deux approches selon la disponibilité de registres préentraînés :

• Registres Entraînés : Pour des encodeurs comme DINOv2 qui prennent en charge les registres, les auteurs utilisent des tokens de registre préentraînés pour supprimer les artefacts à forte norme dans les cartes de caractéristiques de l'encodeur.
• Registres Récursifs au Moment du Test (TTR) : Pour les encodeurs sans registres intégrés (par exemple, SigLIP2), les auteurs ajoutent des tokens supplémentaires non entraînés à la séquence d'entrée lors de l'inférence. Dans les cas où une deuxième vague d'anomalies est détectée dans la représentation résultante, ils appliquent les TTR de manière récursive pour stabiliser la sortie.

2. Registres dans les Transformers de Diffusion

Les auteurs observent que même avec un encodeur stabilisé, le Transformer de diffusion lui-même développe des tokens anormaux internes, principalement dans les couches intermédiaires (contrairement aux ViT standards où les anomalies apparaissent souvent dans les couches finales).

• Registres de Diffusion : Ils introduisent un petit ensemble de tokens de registre apprenables directement dans le Transformer de diffusion. Ces tokens sont entraînés conjointement avec le générateur pour absorber les puits d'attention et stabiliser l'espace de représentation.
• Protocole d'Inférence : Lors de l'inférence, les sorties de tous les tokens de registre (provenant à la fois de l'encodeur et du modèle de diffusion) sont rejetées, ne conservant que les sorties des tokens d'image.

Contributions Clés

• Découverte d'Anomalies Internes : L'article démontre que les tokens anormaux sont omniprésents dans les DiT, apparaissant à la fois dans l'encodeur et le débruiteur, avec un motif distinct de concentration dans les couches intermédiaires du modèle de diffusion.
• Rejet de l'Hypothèse des "Valeurs Extrêmes" : Grâce à des expériences de masquage de perte, les auteurs montrent que l'exclusion simple des tokens à forte norme de la perte d'entraînement n'améliore pas la qualité de génération. Cela conduit à la conclusion que le problème central est la corruption de la sémantique des patchs locaux, et non la simple présence de valeurs extrêmes.
• Intervention à Double Étape : La proposition des DSR offre un mécanisme simple et léger pour stabiliser l'ensemble du pipeline de génération en traitant les anomalies à la source (encodeur) et durant le processus de génération (débruiteur).
• Généralisabilité : La méthode s'avère efficace à travers différents espaces d'entrée (espace pixel, latents VAE et espaces de représentation) et diverses architectures (SiT, JiT, conceptions basées sur RAE).

Résultats Expérimentaux

Les auteurs évaluent les DSR sur la génération conditionnelle par classe ImageNet-1K (256×256) et la génération texte-à-image à grande échelle.

• Performance ImageNet-1K :
  • En utilisant SigLIP2-B comme encodeur, les DSR réduisent le FID de 5,89 à 4,58 (avec 800 époques) et améliorent les scores GenEval sur les tâches texte-à-image de 0,426 à 0,466.
  • Combinés avec des encodeurs DINOv2-B, les DSR améliorent le FID de 4,16 à 3,95.
  • Les DSR accélèrent la convergence, atteignant une qualité comparable avec 4 fois moins d'époques d'entraînement par rapport à la base.
• Évolutivité : Les améliorations sont cohérentes à travers les tailles de modèles (DiT-S, B, XL), avec une augmentation marginale du coût computationnel (environ 10 % d'augmentation des GFLOPs).
• Études d'Ablation :
  • Profondeur d'Insertion : Les registres sont plus efficaces lorsqu'ils sont introduits dans les couches débutantes à intermédiaires (par exemple, à partir du bloc 8). Leur introduction trop tôt ou trop tard donne des rendements décroissants.
  • Nombre de Registres : Un nombre modéré de registres (par exemple, 36) produit des performances optimales ; un nombre trop faible (1) ou trop élevé (100) dégrade les performances.
• Comparaison avec les Alternatives : Les DSR surpassent le "conditionnement in-context" (qui utilise des tokens dépendants de l'entrée) à mesure que la représentation d'entrée devient plus sémantique, suggérant que les registres offrent un mécanisme plus robuste pour la capacité globale.

Importance et Revendications

L'article revendique que le contrôle des tokens anormaux est un composant critique, mais jusque-là négligé, dans la construction de pipelines de diffusion robustes basés sur les Transformers. En démontrant que ces artefacts dégradent la sémantique des patchs locaux et que les interventions basées sur les registres peuvent efficacement les atténuer, les auteurs positionnent les DSR comme une amélioration fondamentale pour la stabilité de l'encodeur et du générateur.

Le travail souligne que la flexibilité des Transformers, bien que bénéfique pour les interactions globales, introduit des artefacts spécifiques (tokens anormaux) qui entravent la fidélité spatiale dans la génération. La solution proposée est présentée comme une intervention simple, prête à l'emploi, qui améliore constamment la qualité de génération dans des contextes divers sans nécessiter de changements architecturaux complexes ni de réentraînement de l'ensemble du pipeline à partir de zéro.