Résumé Technique : MaskAttn-SDXL

Énoncé du Problème

Bien que les modèles de diffusion, en particulier Stable Diffusion XL (SDXL), aient atteint un succès significatif dans la génération d'images à partir de texte, ils font face à des défis persistants lors du traitement de prompts structurés et multi-objets. Les limitations principales sont doubles :

1. Compromis Architecturaux : Les architectures U-Net dominantes reposent sur des couches de convolution à forte localité spatiale, ce qui assure une stabilité et une efficacité de l'entraînement, mais entrave la modélisation des dépendances à long terme et de la composition globale. À l'inverse, les modèles de diffusion basés sur des Transformers (par exemple, DiT) offrent une auto-attention globale mais entraînent des coûts computationnels et mémoire substantiellement plus élevés, souvent au détriment de la structure locale fine.
2. Fiabilité Compositionnelle : Même les modèles les plus avancés échouent fréquemment à maintenir une cohérence au niveau des tokens dans des prompts complexes. Les modes d'échec courants incluent :
   • Fuite d'Attributs : Attribuer incorrectement des couleurs ou des propriétés entre les objets.
   • Violations Spatiales : Inverser les relations spatiales (par exemple, la position gauche/droite).
   • Omissions d'Objets : Échouer à générer les entités demandées.
     Ces erreurs ne sont souvent pas capturées par des métriques globales comme le FID ou les scores CLIP, qui mesurent l'alignement de la distribution globale plutôt que la liaison spécifique token-région. La cause racine est identifiée comme étant l'interférence inter-token, où les concepts textuels entrent en compétition pour une attention spatiale partagée, conduisant à des sorties visuelles entremêlées.

Méthodologie : MaskAttn-SDXL

Les auteurs proposent MaskAttn-SDXL, un module léger et plug-in conçu pour améliorer SDXL sans modifier son pipeline central, ses poids préentraînés ou sa procédure d'inférence.

Mécanisme Central

La méthode injecte un gating spatial conditionné par le token dans les logits de l'attention croisée du U-Net SDXL avant l'opération softmax.

1. Intégration Architecturale :

   • La modification est appliquée spécifiquement au sein des blocs d'attention croisée des couches de résolution moyenne sélectionnées du U-Net SDXL. Cet emplacement a été choisi pour équilibrer l'abstraction sémantique et la précision spatiale.
   • Tous les composants préentraînés (squelette U-Net, encodeurs de texte, VAE) restent figés. Seules des "têtes de masquage" légères sont entraînées.

2. Génération de Masque :

   • Une petite tête de porte partagée $f(\cdot)$ (par exemple, une couche de convolution peu profonde ou un MLP) prend en entrée la carte de caractéristiques latentes actuelle $X$ et un embedding de token spécifique $e_t$.
   • Elle produit une carte de probabilité spatiale continue $\hat{G}_t \in (0, 1)^{H \times W}$, représentant la croyance du modèle quant à la force avec laquelle le token $t$ devrait contribuer à la génération de l'image à chaque coordonnée spatiale.
   • Cette carte est binarisée en utilisant un seuil de 0,5 et un estimateur à passage direct (STE) pour permettre la propagation du gradient pendant l'entraînement, résultant en une porte binaire dure $G_t$.

3. Modulation des Logits :

   • La porte binaire est convertie en une matrice de masque additive $M$. Pour une position spatiale $i$ et un token $t$ :
     • Si la porte est active ($G_t(i)=1$), la valeur du masque est 0.
     • Si la porte est inactive ($G_t(i)=0$), la valeur du masque est $-\infty$ (approximée par une grande constante négative).
   • Ce masque est ajouté directement aux logits de l'attention croisée :
     $$\text{MaskAttn}(Q, K, V; M) = \text{Softmax}\left(\frac{QK^\top}{\sqrt{d}} + M\right)V$$
   • Définir un logit à $-\infty$ force le poids softmax correspondant à zéro, supprimant efficacement l'interaction entre un token spécifique et une région spatiale spécifique.

4. Connexion Résiduelle :

   • La sortie de l'attention masquée est passée à travers un réseau feed-forward standard (FFN) et combinée à l'entrée originale via une connexion résiduelle, assurant que le flux d'information original est préservé tout en affinant l'alignement token-région.

Contributions Clés

1. Attention Masquée dans l'Espace des Logits : Les auteurs proposent un mécanisme qui sparse les connexions token-à-emplacement dans SDXL en injectant des masques appris et conditionnés par le token directement dans les logits d'attention.
2. Efficacité Plug-and-Play : La méthode ne nécessite aucune entrée auxiliaire (comme des boîtes englobantes), ajoute des paramètres négligeables (environ 3,6 M) et laisse le squelette préentraîné et le processus d'échantillonnage standard inchangés.
3. Ancrage Compositionnel Amélioré : L'approche améliore la fidélité spatiale et des attributs sous des prompts complexes multi-entités sans sacrifier l'efficacité ni la flexibilité de déploiement.

Résultats Expérimentaux

Les auteurs ont évalué MaskAttn-SDXL par rapport à diverses bases de référence, notamment SD-1.5, SD-2.1, SDXL, PixArt-$\alpha$/$\Sigma$, et d'autres méthodes d'édition compositionnelle (par exemple, Attend-and-Excite, Structured Diffusion).

1. Fidélité Globale et Alignement

Évalué sur les ensembles de données MS-COCO et Flickr30k :

• FID : MaskAttn-SDXL a réduit le score FID sur MS-COCO de 25,77 (SDXL) à 24,57, et sur Flickr30k de 209,80 à 206,98.
• Score CLIP : Amélioration de 31,53 à 31,75 sur MS-COCO et de 33,03 à 33,54 sur Flickr30k.
• Précision-Rappel : A montré une couverture améliorée (Rappel) tout en maintenant une fidélité comparable.

2. Benchmarks Compositionnels

Évalué sur T2I-CompBench++ (Spatial) et GenEval :

• T2I-CompBench++ (Spatial) : A obtenu un score de 0,2248, surpassant SDXL (0,2090) et toutes les autres bases de référence.
• GenEval : A obtenu un score de 0,66, surpassant SDXL (0,62) et d'autres méthodes spécialisées.
• Ces résultats indiquent une réduction significative des fuites d'attributs et des violations spatiales.

3. Efficacité et Surcharge

• Paramètres : Ajout de seulement ~3,6 M de paramètres au modèle SDXL de 2,6 Md.
• Mémoire et Latence : Augmentation légère de l'utilisation mémoire (de 12,4 Go à 12,7 Go) et de la latence de 0,03 seconde par image (de 1,00 s à 1,03 s), confirmant que la méthode est légère.

4. Études d'Ablation

• Sélection d'Étape : Les blocs de résolution moyenne ont produit les meilleurs scores compositionnels, validant la stratégie de placement par défaut.
• Placement : L'insertion du module dans les blocs encodeur et décodeur a fourni le meilleur alignement global, bien que les gains principaux soient attribués à la sélection de résolution moyenne.

Importance et Revendications

L'article revendique que MaskAttn-SDXL adresse l'enjeu critique de l'interférence inter-token dans la génération d'images à partir de texte multi-objets. En supprimant explicitement les liaisons token-région non pertinentes dans l'espace des logits, la méthode atteint une correction compositionnelle plus forte sans la surcharge computationnelle du réentraînement de grands modèles ni la complexité au moment de l'inférence des outils d'édition basés sur des heuristiques.

Les auteurs soulignent que leur approche offre une solution pratique pour améliorer la fiabilité des modèles de diffusion dans des contextes purement textuels, maintenant l'efficacité et la flexibilité du pipeline SDXL standard tout en réduisant significativement les modes d'échec courants tels que l'omission d'objets et l'inversion spatiale. Les résultats suggèrent qu'une modulation ciblée et sparse de l'attention est une voie viable pour améliorer la cohérence globale dans les squelettes de diffusion convolutifs.