Résumé technique : Les registres sont déterminants pour les Transformers de diffusion dans l'espace des pixels

Énoncé du problème

Les Transformers de vision (ViT) sont connus pour présenter des tokens de patch à norme élevée, des valeurs aberrantes qui dégradent la qualité des cartes de caractéristiques et créent des artefacts dans les cartes d'attention. Ce problème est efficacement atténué dans les ViT par l'introduction de « tokens de registre », qui agissent comme des puits pour ces valeurs aberrantes. Alors que les Transformers de diffusion (DiT) adoptent de plus en plus des architectures de type Transformer et se tournent vers un entraînement directement dans l'espace des pixels (s'éloignant de l'espace latent), ils ressemblent structurellement davantage aux ViT. Cela soulève deux questions critiques :

1. Les DiT héritent-ils des tokens de patch à norme élevée observés dans les ViT ?
2. Les tokens de registre sont-ils bénéfiques pour les DiT, en particulier dans les régimes d'entraînement dans l'espace des pixels ?

Bien que les puits d'attention et les comportements des tokens spéciaux aient été étudiés dans les grands modèles de langage (LLM) et les modèles de diffusion linguistique, leur présence et leur rôle fonctionnel dans les transformers de diffusion d'images restent peu explorés.

Méthodologie

Les auteurs mènent une analyse systématique des tokens de registre à travers différents espaces de modèles de diffusion (espace des pixels, espace VAE et espace RAE) et les comparent aux ViT standards (spécifiquement DINOv2).

1. Analyse empirique des valeurs aberrantes : L'étude examine les normes de caractéristiques au niveau des tokens et les cartes d'attention dans les DiT (spécifiquement des pDiT entraînés avec du matching de flux sur ImageNet) par rapport aux ViT. Ils analysent si les tokens de patch présentent des valeurs aberrantes à norme élevée en l'absence de registres.
2. Intégration des tokens de registre : Les tokens de registre sont implémentés comme des tokens appris supplémentaires ajoutés à la séquence de tokens de patch, suivant l'approche utilisée dans les ViT. Ces tokens sont exclus de l'objectif d'entraînement (fonction de perte).
3. Analyse des représentations internes : Les auteurs analysent les cartes de caractéristiques intermédiaires en utilisant la variation totale (TV) pour mesurer la régularité spatiale et le sondage linéaire pour évaluer le contenu sémantique des tokens de registre. Ils examinent comment les registres affectent les normes de caractéristiques et la cohérence des représentations à différents pas de diffusion (niveaux de bruit).
4. Exploration architecturale : Motivés par les rôles distincts des registres et des tokens de patch, les auteurs proposent une architecture à double flux économe en paramètres. Cette conception traite les tokens de registre et les tokens de patch via des voies spécialisées (spécifiquement en découplant les couches RMSNorm, MLP et adaLN) tout en partageant les mécanismes d'attention et la plupart des paramètres du transformer.
5. Études d'ablation : L'étude ablate le nombre de registres, la profondeur à laquelle ils sont introduits et les composants spécifiques à dualiser pour déterminer les configurations optimales.

Résultats clés et contributions

1. Les DiT diffèrent fondamentalement des ViT concernant les valeurs aberrantes

Contrairement aux ViT, les DiT standards (tant dans l'espace latent que dans l'espace des pixels) ne présentent pas de tokens de patch à norme élevée. Dans les ViT, les valeurs aberrantes sont localisées sur des tokens spécifiques (souvent dans des régions à faible information), tandis que dans les DiT, les normes des tokens de patch restent relativement uniformes et les cartes d'attention manquent des artefacts à faible information observés dans les ViT.

2. Les registres améliorent les DiT dans l'espace des pixels malgré l'absence de valeurs aberrantes préexistantes

De manière surprenante, l'ajout de tokens de registre aux DiT améliore considérablement la convergence et la qualité de génération, en particulier dans les modèles dans l'espace des pixels.

• Mécanisme : Contrairement aux ViT où les registres absorbent des valeurs aberrantes préexistantes, les registres dans les DiT créent des tokens à norme élevée en leur sein.
• Effet : La présence de registres réduit systématiquement les normes de caractéristiques de tous les tokens de patch. Cela conduit à des cartes de caractéristiques intermédiaires plus propres et plus lisses, en particulier à hauts niveaux de bruit (premiers pas de diffusion).
• Double rôle : Les tokens de registre se spécialisent en deux rôles distincts :
  • Puits de norme : Tokens avec des normes extrêmement élevées qui absorbent l'amplitude des tokens de patch.
  • Porteurs d'informations globales : Tokens avec des normes modérées qui encodent des informations sémantiques diversifiées sur l'image (par exemple, objets au premier plan, contexte d'arrière-plan), vérifiés par sondage linéaire et visualisation de cartes d'attention.

3. Efficacité spécifique à l'espace

Les avantages des tokens de registre dépendent fortement de l'espace d'entraînement :

• Espace des pixels : Les registres produisent les améliorations les plus significatives de la qualité de génération (FID). Les modèles dans l'espace des pixels produisent intrinsèquement des représentations intermédiaires plus bruyantes avec des normes de caractéristiques plus élevées, rendant l'effet de régularisation des registres crucial.
• Espace latent (VAE/RAE) : Les registres apportent des gains modérés ou, dans le cas de modèles basés sur RAE, dégradent légèrement les performances. Les espaces latents sont plus structurés et de dimension inférieure, réduisant le besoin de la forte régularisation fournie par les registres.

4. Registres implicites dans les architectures existantes

Les auteurs observent que les architectures récentes de DiT dans l'espace des pixels (par exemple JiT) qui utilisent un conditionnement de classe in-context (ajout d'embeddings de classe à la séquence de tokens) intègrent implicitement des mécanismes de type registre. Ces tokens in-context présentent un comportement similaire aux registres explicites : certains agissent comme des puits de norme, tandis que d'autres encodent des informations sémantiques globales. Cela suggère que les gains de performance de tels modèles peuvent découler largement de ce comportement de type registre plutôt que de l'information de classe elle-même.

5. Architecture à double flux économe en paramètres

Pour tirer parti de ces insights, les auteurs proposent une architecture compacte à double flux qui spécialise le traitement des tokens de registre.

• Conception : Les registres sont introduits uniquement dans les couches plus profondes (à partir du bloc 4), où la structure sémantique est établie. L'architecture utilise une attention partagée mais découple RMSNorm, MLP et adaLN pour les registres et les patchs.
• Efficacité : Cette approche n'augmente la taille du modèle que de ~14 % (contre ~77 % pour une duplication complète naïve) tout en maintenant des coûts d'inférence similaires (GFLOPs).
• Performance : Cette architecture améliore systématiquement la qualité de génération à travers différentes échelles de modèles (B, L, H) et résolutions (256x256, 512x512) par rapport aux bases à flux unique.

Résultats

• Qualité de génération : Sur ImageNet 256x256, les pDiT avec registres surpassent systématiquement ceux sans registres. Par exemple, un modèle pDiT-B/16 a amélioré son FID de 7,39 à 5,30 à 200 époques.
• Régularité des caractéristiques : Les modèles avec registres affichent des ratios de variation totale (TV) plus faibles dans les caractéristiques intermédiaires à hauts niveaux de bruit, indiquant des représentations plus lisses et plus cohérentes.
• Ablation : Les registres sont les plus efficaces lorsqu'ils sont introduits après les quatre premières couches et lorsqu'un plus grand nombre de tokens est utilisé (par exemple, 32 registres pour pDiT-B, contre 4 pour les ViT).

Importance et affirmations

L'article affirme que les tokens de registre sont un composant critique, mais jusque-là négligé, pour une génération de diffusion de haute qualité dans l'espace des pixels. L'importance réside dans :

1. Redéfinition du rôle des registres : Démontrer que les registres ne sont pas simplement une solution aux valeurs aberrantes spécifiques aux ViT, mais qu'ils remplissent une fonction de régularisation plus large dans les modèles de diffusion en gérant les normes de caractéristiques et en améliorant la cohérence spatiale.
2. Explication des succès récents : Fournir une justification pour la forte performance empirique des DiT récents dans l'espace des pixels qui utilisent un conditionnement in-context, attribuant une partie de leur succès à un comportement de registre implicite.
3. Conception architecturale efficace : Offrir une conception à double flux économe en paramètres qui capitalise sur les rôles distincts des registres et des patchs, atteignant une qualité de génération de pointe avec une surcharge d'exécution négligeable.

Les auteurs concluent que, bien que leurs conceptions capturent la promesse de la spécialisation des tokens, le potentiel complet de l'utilisation des registres dans les transformers de diffusion dans l'espace des pixels reste une question ouverte pour la recherche future.