Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers

Ce papier présente le Hourglass Diffusion Transformer (HDiT), un modèle de génération d'images qui atteint une mise à l'échelle linéaire avec le nombre de pixels, permettant un entraînement direct en espace pixel à haute résolution sans recourir aux techniques habituelles comme les autoencodeurs latents, et établissant un nouvel état de l'art sur FFHQ-1024².

L'essentiel

🎨 Le Problème : La "Photo Floue" et le "Miroir Magique"

Imaginez que vous voulez créer une image ultra-détaillée (comme un visage humain en 4K) avec un ordinateur. Jusqu'à présent, la méthode la plus populaire (appelée Latent Diffusion, utilisée par des outils comme Stable Diffusion) fonctionnait un peu comme un miroir magique.

1. L'ordinateur prend votre image, la transforme en une version "résumée" et floue (un code secret ou "latent").
2. Il travaille sur ce résumé pour créer l'image.
3. Il doit ensuite retransformer ce résumé en image finale.

Le souci ? Ce processus de résumé et de dé-résumé agit comme un filtre de basse qualité. C'est comme si vous essayiez de peindre un tableau de maître en regardant votre reflet dans une vitre sale : les détails fins (les pores de la peau, les cheveux individuels) se perdent. De plus, pour faire des images très grandes, ces miroirs magiques deviennent extrêmement lents et coûteux en énergie, un peu comme essayer de faire passer un éléphant dans un tuyau d'arrosage.

🏗️ La Solution : HDiT (Le "Sandwich" Intelligent)

Les auteurs de ce papier, Katherine Crowson et son équipe, ont créé un nouveau modèle appelé HDiT (Hourglass Diffusion Transformer). Au lieu d'utiliser un miroir magique, ils ont décidé de travailler directement sur les pixels (les petits points de couleur qui composent l'image), comme un artiste qui peint directement sur la toile sans passer par un brouillon flou.

Pour que cela fonctionne sans exploser le budget informatique, ils ont inventé une architecture en forme de Sablier (Hourglass).

L'Analogie du Sablier et de l'Équipe de Construction

Imaginez que vous devez construire une cathédrale (une image haute résolution) :

1. La partie haute du sablier (L'Encodage) : Au lieu d'envoyer tous les ouvriers (les données de l'image) sur le chantier en même temps, vous les regroupez par équipes.

   • Vous prenez 4 ouvriers et vous en faites 1 chef d'équipe.
   • Vous prenez 4 chefs et vous en faites 1 superviseur.
   • Vous continuez jusqu'à ce qu'il ne reste qu'une toute petite équipe centrale. C'est là que le modèle comprend la "grande idée" de l'image (le ciel est bleu, il y a un visage au centre). Cela coûte très peu d'énergie car il y a peu de gens à gérer.

2. Le fond du sablier (Le Cœur) : C'est là que le modèle travaille sur cette petite idée globale.

3. La partie basse du sablier (Le Décodeur) : Maintenant, on remonte !

   • Le superviseur donne des ordres à ses 4 chefs.
   • Les chefs donnent des ordres à leurs 4 ouvriers.
   • Mais attention : à chaque étape, on réintroduit les détails précis qu'on avait mis de côté au début (comme si on redonnait les plans détaillés aux ouvriers).

La magie de HDiT :
Contrairement aux anciennes méthodes qui devenaient exponentiellement plus lentes quand l'image grossissait (comme si doubler la taille de l'image doublait le travail 4 fois), HDiT reste efficace.

• L'analogie : Si vous doublez la taille de la cathédrale, HDiT n'ajoute qu'une équipe de plus, pas une armée entière. C'est une économie d'énergie colossale.

🚀 Les Résultats : Pourquoi c'est impressionnant ?

Grâce à cette méthode "Sablier", les chercheurs ont pu :

1. Travailler directement sur les pixels : Plus de perte de qualité due au "miroir magique". Les images sont nettes, avec des détails incroyables (comme les reflets dans les yeux ou la texture de la peau).
2. Créer des images géantes : Ils ont entraîné le modèle sur des images de 1024x1024 pixels (une très haute résolution) directement, sans tricher ni utiliser de techniques compliquées.
3. Battre les records : Sur des tests de qualité (comme la création de visages humains), leur modèle a surpassé les meilleurs modèles existants, y compris ceux qui utilisaient des méthodes plus complexes.

🌟 En Résumé

Imaginez que vous vouliez dessiner un portrait ultra-réaliste.

• Les anciennes méthodes vous forçaient à dessiner d'abord un croquis au crayon très flou, puis à essayer de le colorier par-dessus. Le résultat était souvent un peu "mou".
• HDiT, c'est comme avoir un assistant qui vous permet de peindre directement au pinceau fin, mais qui vous aide à ne pas vous perdre dans les détails en vous donnant d'abord une vue d'ensemble, puis en zoomant progressivement sur les détails, le tout sans vous fatiguer.

C'est une avancée majeure qui rend la création d'images haute définition plus rapide, moins coûteuse en énergie, et surtout, beaucoup plus belle et fidèle à la réalité.

Résumé technique

1. Le Problème : Limites de la Synthèse d'Images Haute Résolution

Les modèles de diffusion actuels dominent la génération d'images, mais leur application à des résolutions très élevées (ex: 1024x1024 ou plus) se heurte à plusieurs obstacles majeurs :

• Complexité Quadratique des Transformers : Les architectures basées sur les Transformers (comme DiT - Diffusion Transformers) appliquent une attention globale. Leur complexité computationnelle évolue en $O(n^2)$ par rapport au nombre de tokens (pixels). Cela rend l'entraînement et l'inférence prohibitifs pour des images en espace pixel à haute résolution.
• Limites des Modèles Latents (LDM) : Pour contourner ce problème, la plupart des modèles (ex: Stable Diffusion) utilisent des espaces latents compressés via un VAE (Autoencodeur Variationnel). Cependant, cette compression entraîne une perte de détails fins (hautes fréquences) et une reconstruction imparfaite, ce qui nuit à la qualité des images et limite l'efficacité des tâches d'édition d'image.
• Complexité des Architectures Hybrides : Les approches actuelles pour la haute résolution nécessitent souvent des architectures complexes (cascades de super-résolution, pertes multi-échelles, auto-conditionnement) qui alourdissent l'entraînement et la déduction.

2. Méthodologie : L'Hourglass Diffusion Transformer (HDiT)

Les auteurs proposent HDiT, une architecture purement basée sur les Transformers, conçue pour générer des images directement dans l'espace pixel avec une complexité linéaire.

Architecture Hiérarchique en "Sablier" (Hourglass)

HDiT s'inspire de la structure hiérarchique des U-Net (souvent utilisés dans les CNN) mais l'adapte aux Transformers :

• Structure en Sablier : L'image est traitée à travers plusieurs niveaux hiérarchiques.
  • Encodeur : Les tokens (patchs d'image) sont fusionnés spatialement (2x2 en 1) via des opérations de Pixel-UnShuffle, réduisant la résolution spatiale tout en augmentant la dimension cachée.
  • Cœur (Bottleneck) : Le traitement se fait à une résolution très basse (ex: 16x16 tokens).
  • Décodeur : La résolution est réaugmentée progressivement via des opérations de Pixel-Shuffle.
• Connexions Résiduelles Adaptatives : Contrairement aux U-Net classiques qui concatènent les features, HDiT utilise une interpolation linéaire apprenable (learnable linear interpolation) pour fusionner les informations des connexions résiduelles avec les branches remontantes. Cela permet au modèle d'apprendre l'importance relative de chaque niveau de détail.

Mécanismes d'Attention Efficaces

Pour atteindre une complexité $O(n)$ au lieu de $O(n^2)$ :

• Attention Globale Limitée : L'attention globale (nécessaire pour la cohérence de l'image) n'est appliquée qu'aux niveaux les plus bas de la hiérarchie (résolutions faibles).
• Attention Locale : Aux niveaux supérieurs (haute résolution), le modèle utilise une Attention de Voisinage (Neighborhood Attention). Cela permet de capturer les détails locaux avec une complexité linéaire par rapport au nombre de pixels.
• Encodage de Position : Remplacement des encodages de position additifs classiques par une adaptation 2D des Rotary Positional Embeddings (RoPE), améliorant la généralisation et réduisant les artefacts de patch.

Composants du Bloc

• Normalisation : Utilisation de l'AdaRMSNorm (adaptative) conditionnée par le temps de diffusion et la classe.
• Feed-Forward : Utilisation de l'activation GEGLU (Gated Linear Unit) au lieu de GELU standard, ce qui améliore la qualité de génération.
• Pondération de la Perte : Adoption d'une stratégie de pondération Soft-Min-SNR (Signal-to-Noise Ratio) pour stabiliser l'entraînement et améliorer la convergence.

3. Contributions Clés

1. HDiT : Introduction d'une architecture de diffusion Transformer pure qui atteint une complexité computationnelle linéaire $O(n)$ par rapport au nombre de pixels, contrairement à la complexité quadratique $O(n^2)$ des Transformers standards.
2. Synthèse en Espace Pixel Haute Résolution : Démonstration réussie de la génération d'images de haute qualité à 1024x1024 directement en espace pixel, sans recourir à des espaces latents (VAE) ni à des techniques d'entraînement complexes comme la croissance progressive ou les pertes multi-échelles.
3. Performance Record : Établissement d'un nouvel état de l'art (SOTA) pour les modèles de diffusion sur le dataset FFHQ-1024 (visages) et des performances compétitives sur ImageNet-256, surpassant les modèles basés sur des latents et rivalisant avec les GANs.

4. Résultats Expérimentaux

• FFHQ-1024 (Visages) :

  • Le modèle HDiT (85M de paramètres) atteint un FID de 5.23 (avec guidance) et 8.48 (sans guidance), surpassant nettement les modèles de diffusion précédents comme NCSN++ (FID ~53).
  • Il bat également les GANs récents comme StyleSwin et rivalise avec StyleGAN-XL.
  • Les métriques basées sur DINOv2 (FDD et KDD), qui corrèlent mieux avec la préférence humaine pour les visages, montrent également un nouveau record.
  • Visuellement, les images sont nettes, avec des détails fins et une symétrie faciale correcte, contrairement aux modèles latents qui produisent souvent des images floues.

• ImageNet-256 (Classification) :

  • Un modèle de 557M de paramètres atteint un FID de 6.92 (sans guidance) et 3.21 (avec guidance), se classant parmi les meilleurs modèles de diffusion, bien que légèrement derrière les modèles massifs utilisant l'auto-conditionnement (comme RIN ou VDM++).
  • Il surpasse les modèles U-Net basiques et les modèles latents (LDM, DiT) en termes de rapport qualité/coût computationnel pour une génération en espace pixel.

• Efficacité Computationnelle :

  • À 256x256, HDiT est plus de 10 fois plus efficace qu'un DiT équivalent.
  • À 1024x1024, le gain d'efficacité dépasse 100 fois par rapport à un DiT standard, rendant l'entraînement de modèles Transformers haute résolution en espace pixel viable.

5. Signification et Impact

Ce travail est significatif car il comble le fossé entre l'efficacité des architectures convolutives (U-Net) et la capacité de mise à l'échelle des Transformers.

• Suppression du goulot d'étranglement du VAE : En permettant une génération directe en espace pixel à haute résolution, HDiT élimine les limitations de qualité inhérentes aux VAE, ouvrant la voie à des applications d'édition d'image plus précises et à une génération de détails réalistes.
• Évolutivité : La complexité linéaire permet d'envisager l'entraînement de modèles de diffusion sur des résolutions mégapixel (au-delà de 1024x1024) sans explosion des coûts de calcul.
• Simplicité : L'architecture atteint ces résultats sans nécessiter de "tricks" d'entraînement complexes (comme la cascade de modèles ou l'auto-conditionnement), simplifiant ainsi le pipeline de génération.

En conclusion, HDiT représente une avancée majeure vers des modèles de génération d'images universels, capables de produire des images de très haute fidélité directement en pixels, tout en restant computationnellement efficaces.