Résumé Technique : SEGA – Attention Guidée par l'Énergie Spectrale pour l'Extrapolation de Résolution dans les Transformers de Diffusion

Énoncé du Problème

Les Transformers de Diffusion (DiT) se sont imposés comme l'architecture dominante pour la génération d'images à partir de texte (T2I). Cependant, leurs performances se dégradent considérablement lors de la génération d'images à des résolutions dépassant leur plage d'entraînement (généralement de 1024² à 2048²). L'extrapolation vers des résolutions plus élevées entraîne souvent des textures floues, des motifs répétitifs et des effondrements structurels.

Les solutions existantes sans entraînement tentent d'atténuer ce problème en modifiant le comportement de l'attention au moment de l'inférence, principalement par l'extrapolation des Embeddings de Position Rotatifs (RoPE) combinée à une mise à l'échelle de l'attention. Les méthodes actuelles, telles que YaRN, appliquent une mise à l'échelle uniforme et agnostique au contenu sur tous les composants RoPE. Cette approche traite les composantes de courte longueur d'onde (gouvernant les textures fines) de manière identique aux composantes de longue longueur d'onde (façonnant la structure globale). L'article soutient que cette uniformité statique induit un compromis fondamental : les stratégies qui préservent la structure globale échouent souvent à récupérer les détails fins, et vice versa. De plus, la distribution spectrale de la représentation latente évolue dynamiquement au cours du processus de débruitage et varie selon le contenu de l'image, rendant la mise à l'échelle statique sous-optimale.

Méthodologie : SEGA

Les auteurs introduisent SEGA (Spectral-Energy Guided Attention), une méthode sans entraînement qui met à l'échelle dynamiquement l'attention sur les composants RoPE en fonction de la structure en fréquence spatiale du latent à chaque étape de débruitage.

Insight Central

SEGA est motivé par l'observation que les composants RoPE sont couplés à des fréquences spatiales spécifiques. Au lieu d'appliquer un facteur de mise à l'échelle fixe, SEGA dérive des magnitudes de mise à l'échelle par dimension en analysant la distribution d'énergie du latent intermédiaire courant.

Flux Algorithmique

1. Analyse Spectrale : À chaque étape de débruitage, le latent intermédiaire $Z$ est transformé du domaine spatial au domaine fréquentiel à l'aide d'une Transformée de Fourier Rapide (FFT) 2D.

   • Profils par Axe : Le spectre de puissance 2D est marginalisé sur les axes orthogonaux pour obtenir des profils d'énergie 1D ($E_H, E_W$) pour les fréquences horizontales et verticales.
   • Profil Radial : Le spectre de puissance est moyenné dans des anneaux concentriques pour obtenir un profil invariant par rotation ($E_{iso}$) représentant la distribution d'énergie à travers les échelles spatiales.

2. Correction par Dimension ($s_d$) :

   • Pour chaque dimension RoPE $d$, la longueur d'onde correspondante est identifiée pour la mapper à une bande de fréquence spécifique dans le profil par axe.
   • L'énergie logarithmique de cette bande est standardisée.
   • Une non-linéarité (tanh) est appliquée, et le résultat est centré sur zéro pour assurer une redistribution stricte de la mise à l'échelle. Les dimensions associées à des bandes de faible énergie reçoivent une mise à l'échelle plus forte pour préserver la discrimination positionnelle, tandis que les bandes de haute énergie reçoivent une mise à l'échelle plus faible pour éviter la sur-amplification.

3. Facteur d'Amplitude Global ($\sigma$) :

   • Le profil radial est réduit à un scalaire unique en utilisant la platitude spectrale (entropie de Wiener).
   • Ce facteur contrôle la force de la mise à l'échelle. Si le spectre du latent est plat (dominé par le bruit), $\sigma \to 0$, supprimant la correction. À mesure que le contenu structurel émerge et que le spectre devient moins plat, $\sigma \to 1$, appliquant la correction à pleine force.

4. Formule de Mise à l'Échelle Finale :
   Le facteur de mise à l'échelle final $m_d$ pour une dimension $d$ est calculé comme suit :
   $$m_d = m_{ref} \cdot (1 - \sigma \cdot s_d)$$
   Où $m_{ref}$ est une échelle de référence déterminée par le rapport de résolution (en utilisant une formulation en loi de puissance). Cela permet à SEGA de s'adapter continuellement au contenu spectral du latent, en affinant l'attention sur les fréquences sous-résolues et en l'adoucissant sur celles qui sont sur-émphasisées.

Contributions Clés

• Mise à l'Échelle Dynamique et Consciente du Contenu : Contrairement aux méthodes de mise à l'échelle uniforme précédentes, SEGA adapte la mise à l'échelle de l'attention à la structure spectrale spécifique du latent à chaque étape de débruitage.
• Résolution du Compromis Structure-Détail : En différenciant les composants RoPE en fonction de leurs bandes de fréquence associées et de leurs niveaux d'énergie, SEGA améliore simultanément la cohérence structurelle globale et la fidélité des détails fins.
• Sans Entraînement et Agnostique de l'Architecture : La méthode ne nécessite aucun réglage fin, aucun paramètre apprenable et aucun changement d'architecture. Elle s'intègre directement dans les pipelines standards basés sur RoPE (démontré sur Flux et Qwen).
• Performances Supérieures aux Résolutions Ultra-Élevées : SEGA permet une génération stable à des résolutions dépassant 36 millions de pixels (par exemple 6144²), un régime où d'autres méthodes échouent souvent complètement.

Résultats Expérimentaux

Les auteurs ont évalué SEGA sur les architectures Flux et Qwen dans plusieurs configurations haute résolution (jusqu'à 6144²) en utilisant le jeu de données Aesthetic-4K et un benchmark personnalisé en zéro-shot.

• Performance Quantitative : SEGA a systématiquement surpassé les références de l'état de l'art (incluant NTK, YaRN, DyPE, UltraImage et les méthodes de guidage multi-étapes comme HiFlow et I-Max) sur toutes les métriques : ImageReward (IR), PickScore (PS), CLIP Score (CS), MUSIQ (MSQ) et FID. Notamment, à 4096², SEGA a obtenu les meilleurs scores sur chaque métrique évaluée.
• Améliorations Qualitatives : Les comparaisons visuelles ont démontré que SEGA prévient efficacement les modes de défaillance courants tels que les artefacts de tuilage répétitifs, l'effondrement structurel et la perte de clarté sémantique. Il préserve les géométries complexes et les textures fines (par exemple, le pelage, les détails des tissus) que d'autres méthodes floutent ou déforment.
• Robustesse : La méthode a maintenu sa stabilité sur des rapports d'aspect non carrés et des facteurs d'extrapolation extrêmes (jusqu'à 36 fois la zone d'entraînement), alors que les méthodes concurrentes ont montré une dégradation significative des performances à mesure que la résolution augmentait.

Importance et Revendications

L'article revendique que SEGA représente une solution minimale mais efficace pour débloquer les capacités de génération haute résolution des DiT pré-entraînés. En passant d'une mise à l'échelle statique basée sur des heuristiques à une approche dynamique guidée par l'énergie spectrale, les auteurs abordent une limitation fondamentale des stratégies d'extrapolation actuelles.

Les auteurs soulignent que leur travail fournit une « perspective spectrale » sur la mise à l'échelle de l'attention, révélant un couplage entre les dimensions RoPE et les fréquences spatiales qui peut être exploité pour guider l'adaptation au moment de l'inférence. Ils positionnent SEGA non pas comme un remplacement des approches basées sur l'entraînement, mais comme une amélioration puissante, plug-and-play, permettant aux modèles existants de généraliser bien au-delà de leurs résolutions d'entraînement natives, sans le coût computationnel du réentraînement ni la complexité des pipelines multi-étapes. La méthode est présentée comme une étape vers une meilleure compréhension et un meilleur contrôle du comportement des mécanismes d'attention dans les tâches d'extrapolation de résolution.