Diffusion Transformer (DiT)

Un Transformer de diffusion (DiT) est un modèle de diffusion dont le réseau de débruitage est un transformeur plutôt qu’un U-Net convolutionnel. Il découpe le latent d’une image bruitée en patches, les traite comme une séquence de tokens avec une attention automatique, à la manière d’un Vision Transformer (ViT) qui traite une image, et prédit le bruit à supprimer à chaque étape de débruitage. William Peebles et Saining Xie l’ont présenté dans Scalable Diffusion Models with Transformers (ICCV 2023), où DiT-XL/2 a atteint une distance de Fréchet-Inception (FID) de 2,27 sur ImageNet 256×256 conditionné par classe, dépassant le précédent meilleur résultat de 3,60 obtenu par la diffusion latente. L’article a montré que la qualité des images s’améliore régulièrement à mesure que les ressources de calcul d’un DiT augmentent, et les réseaux de débruitage à base de transformeurs ont depuis remplacé les U-Nets dans la plupart des grands générateurs d’images et de vidéos, notamment Stable Diffusion 3, FLUX.1 et Sora d’OpenAI.

Fonctionnement d’un Transformer de diffusion#

Un DiT fonctionne dans l’espace latent de l’autoencodeur variationnel (VAE préentraîné utilisé par Stable Diffusion, dont l’encodeur réduit la résolution par un facteur de 8 ; une image de 256×256×3 devient donc un latent de 32×32×4. Une couche de découpage en patches divise le latent bruité en patches p×p et projette linéairement chacun d’eux sous forme de token ; des tailles de patch de 2, 4 et 8 donnent respectivement 256, 64 et 16 tokens. Des embeddings positionnels sinusoïdaux indiquent l’origine de chaque token, une série de blocs transformeurs traite la séquence, et une couche linéaire finale reconvertit chaque token en un patch de bruit et de variance prédits, qui sont assemblés pour retrouver la forme du latent. Chaque bloc est conditionné par l’étape de diffusion et l’étiquette de classe, ou par le prompt textuel dans un modèle texte-image. La génération répète cette prédiction à partir du bruit pur jusqu’à la dernière étape, puis le décodeur du VAE convertit le latent débruité en pixels :

L’article a comparé quatre façons d’intégrer la condition à chaque bloc : des tokens supplémentaires dans la séquence, l’attention croisée, la normalisation de couche adaptative (adaLN) et adaLN-Zero. adaLN-Zero a obtenu la FID la plus basse et s’est aussi avéré le plus efficace en calcul. Un petit réseau lit la somme des embeddings de l’étape et de la classe, puis estime un décalage et une échelle appliqués après chaque normalisation de couche, ainsi qu’une porte qui pondère chaque branche résiduelle. Les portes commencent à zéro, de sorte que chaque bloc commence l’entraînement comme une fonction identité. Ce patch de code PyTorch découpe un latent en patches et applique la branche d’attention d’un bloc :

import torch
from torch import nn

torch.manual_seed(0)
# A 256x256 image becomes a 4x32x32 latent; 2x2 patches turn it into 256 tokens of width 384 (DiT-S/2)
latent = torch.randn(1, 4, 32, 32)
tokens = nn.Conv2d(4, 384, kernel_size=2, stride=2)(latent).flatten(2).transpose(1, 2)

# adaLN-Zero: the timestep and class embedding regresses a shift, scale, and gate for the attention branch
condition = torch.randn(1, 384)
modulation = nn.Linear(384, 3 * 384)
nn.init.zeros_(modulation.weight)
nn.init.zeros_(modulation.bias)
shift, scale, gate = modulation(condition).unsqueeze(1).chunk(3, dim=-1)

# Self-attention branch of a DiT block, modulated by the condition and scaled by the gate
attention = nn.MultiheadAttention(384, num_heads=6, batch_first=True)
h = nn.functional.layer_norm(tokens, (384,)) * (1 + scale) + shift
out = tokens + gate * attention(h, h, h)[0]
print(f"tokens {tuple(tokens.shape)}, block is identity at initialization: {torch.equal(out, tokens)}")

Il affiche tokens (1, 256, 384), block is identity at initialization: True. Les 256 tokens de largeur 384 correspondent au plus petit modèle, DiT-S/2, et la porte nulle transmet l’entrée sans la modifier jusqu’à ce que l’entraînement la fasse évoluer. Un bloc DiT complet estime un second décalage, une seconde échelle et une seconde porte pour sa branche MLP.

Mise à l’échelle en fonction des tokens et des ressources de calcul#

L’article a entraîné quatre tailles de modèles, de DiT-S (12 couches, largeur 384) à DiT-XL (28 couches, largeur 1 152), chacune avec des tailles de patch de 8, 4 et 2. Pour les 12 modèles, la FID suivait les calculs de propagation avant en Gflops, que ceux-ci proviennent d’un réseau plus profond ou plus large, ou d’un plus grand nombre de tokens. Diviser la taille des patches par deux quadruple le nombre de tokens et les calculs, tout en ajoutant presque aucun paramètre. DiT-XL/2 compte 675 millions de paramètres, sans inclure le VAE de 84 millions de paramètres, et utilise 119 Gflops par propagation avant en 256×256. En 512×512, il traite 1 024 tokens issus d’un latent de 64×64×4 en 525 Gflops et atteint une FID de 3,04. Comme l’attention automatique compare chaque token à tous les autres, les résolutions plus élevées et les vidéos plus longues augmentent son coût de façon quadratique. Ultralytics YOLO26 évite ce coût en plaçant l’attention uniquement à son niveau d’échelle 1/32, qui compte 400 positions pour une image de 640×640, et en conservant des convolutions partout ailleurs pour la détection en temps réel.

DiT, U-Net et ViT : comparaison#

• Réseaux de débruitage U-Net : les U-Nets de DDPM et de Stable Diffusion 1.x et 2.x sont des encodeurs-décodeurs convolutionnels qui travaillent à plusieurs résolutions, avec des connexions de saut entre les niveaux correspondants et de l’attention ajoutée aux résolutions les plus faibles. Un DiT conserve une seule résolution de tokens, applique une attention globale à chaque bloc et comporte moins de choix de conception propres aux images ; il peut donc être mis à l’échelle comme les transformeurs des grands modèles de langage.
• Vision Transformer : ViT encode une image pour des tâches telles que la classification. DiT réutilise son intégration des patches et ses blocs transformeurs, ajoute le conditionnement de l’étape par adaLN-Zero et produit une prédiction pour chaque patch.
• Diffusion latente et correspondance de flux : un modèle de diffusion latente est le cadre général dans lequel le débruitage a lieu dans l’espace latent d’un VAE, et la diffusion ou la correspondance de flux constitue l’objectif d’entraînement. DiT est le réseau de débruitage de ce cadre, et les générateurs actuels associent des architectures DiT à des objectifs de correspondance de flux.

DiT pour la génération d’images et de vidéos à partir de texte#

• DiT multimodal (MMDiT) : Stable Diffusion 3 traite les tokens d’image et de texte avec des poids distincts pour chaque modalité, puis les réunit dans une seule opération d’attention. L’information circule ainsi dans les deux sens entre le prompt et l’image. Cette approche a amélioré la typographie et le respect des prompts, et Stable Diffusion 3.5 Large porte cette architecture à environ 8 milliards de paramètres.
• FLUX.1 : le FLUX.1 de Black Forest Labs est un transformeur de flux rectifié de 12 milliards de paramètres destiné à la génération texte-image.
• Vidéo : le rapport technique de Sora d’OpenAI décrit un transformeur de diffusion qui compresse les vidéos dans un espace latent et les découpe en patches spatio-temporels. Des clips de durées, de résolutions et de formats d’image différents deviennent ainsi des séquences de tokens de longueurs variables. Pour la génération vidéo, le plus grand modèle Movie Gen de Meta est un transformeur de 30 milliards de paramètres dont le contexte de 73 000 tokens vidéo correspond à 16 secondes de vidéo à 16 images par seconde.