Diffusion Transformer (DiT)

Un Diffusion Transformer (DiT) è un modello di diffusione la cui rete di denoising è un transformer invece di una U-Net convoluzionale. Suddivide la rappresentazione latente di un'immagine rumorosa in patch, le elabora come una sequenza di token con self-attention, come fa una Vision Transformer (ViT) con un'immagine, e prevede il rumore da rimuovere a ogni passaggio di denoising. William Peebles e Saining Xie lo hanno introdotto in Scalable Diffusion Models with Transformers (ICCV 2023), dove DiT-XL/2 ha raggiunto una distanza di Fréchet-Inception (FID) di 2,27 su ImageNet 256×256 condizionato per classe, superando il precedente record di 3,60 della diffusione latente. L'articolo ha mostrato che la qualità delle immagini migliora costantemente all'aumentare del calcolo di un DiT; da allora, i denoiser transformer hanno sostituito le U-Net nella maggior parte dei grandi generatori di immagini e video, tra cui Stable Diffusion 3, FLUX.1 e Sora di OpenAI.

Come funziona un Diffusion Transformer#

Un DiT opera nello spazio latente dell'autoencoder variazionale (VAE) preaddestrato usato da Stable Diffusion, il cui encoder esegue un downsampling di 8 volte: un'immagine 256×256×3 diventa quindi una rappresentazione latente 32×32×4. Un livello patchify divide la rappresentazione latente rumorosa in patch p×p e incorpora linearmente ciascuna come token; patch di dimensioni 2, 4 e 8 producono rispettivamente 256, 64 e 16 token. Gli embedding posizionali seno-coseno indicano la posizione originaria di ciascun token, una serie di blocchi transformer elabora la sequenza e un livello lineare finale riconverte ogni token in una patch di rumore e varianza previste, che vengono ricomposte nella forma della rappresentazione latente. Il timestep di diffusione e l'etichetta della classe, oppure il prompt testuale in un modello text-to-image, condizionano ogni blocco. La generazione ripete questa previsione partendo da rumore puro fino all'ultimo passaggio, poi il decoder del VAE converte la rappresentazione latente pulita in pixel:

L'articolo ha confrontato quattro modi di fornire la condizione a ciascun blocco: token aggiuntivi nella sequenza, cross-attention, normalizzazione adattiva del livello (adaLN) e adaLN-Zero. adaLN-Zero ha ottenuto il FID più basso ed è risultato anche il più efficiente dal punto di vista computazionale. Una piccola rete legge la somma degli embedding del timestep e della classe e stima una traslazione e una scala da applicare dopo ogni normalizzazione del livello, oltre a un gate che scala ciascun ramo residuo. I gate partono da zero, perciò ogni blocco inizia l'addestramento come funzione identità. Questa patch di codice PyTorch divide in patch una rappresentazione latente e applica il ramo di attenzione di un blocco:

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)}")

Stampa tokens (1, 256, 384), block is identity at initialization: True. I 256 token con larghezza 384 corrispondono al modello più piccolo, DiT-S/2, e il gate a zero lascia invariato l'input finché l'addestramento non lo modifica. Un blocco DiT completo stima una seconda traslazione, scala e gate per il proprio ramo MLP.

Scalabilità con token e calcolo#

L'articolo ha addestrato quattro dimensioni di modello, da DiT-S (12 livelli, larghezza 384) a DiT-XL (28 livelli, larghezza 1.152), ciascuna con patch di dimensione 8, 4 e 2. In tutti e 12 i modelli, il FID seguiva il calcolo del passaggio in avanti in Gflops, indipendentemente dal fatto che derivasse da una rete più profonda o più larga oppure da un maggior numero di token. Dimezzare la dimensione delle patch quadruplica il numero di token e il calcolo, aggiungendo pochissimi parametri. DiT-XL/2 ha 675 milioni di parametri, esclusi gli 84 milioni del VAE, e richiede 119 Gflops per passaggio in avanti a 256×256. A 512×512 elabora 1.024 token da una rappresentazione latente 64×64×4 con 525 Gflops e raggiunge un FID di 3,04. Poiché la self-attention confronta ogni token con tutti gli altri, risoluzioni più alte e video più lunghi ne aumentano il costo in modo quadratico. Ultralytics YOLO26 evita questo costo applicando l'attenzione solo allo stadio con scala 1/32, che ha 400 posizioni per un'immagine 640×640, e mantenendo le convoluzioni ovunque altrove per il rilevamento in tempo reale.

DiT a confronto con U-Net e ViT#

• Denoiser U-Net: le U-Net di DDPM e Stable Diffusion 1.x e 2.x sono encoder-decoder convoluzionali che operano a diverse risoluzioni, con connessioni skip tra livelli corrispondenti e attenzione aggiunta alle risoluzioni più basse. Un DiT mantiene un'unica risoluzione dei token, applica l'attenzione globale in ogni blocco e presenta meno scelte progettuali specifiche per le immagini, quindi scala come i transformer usati nei modelli linguistici di grandi dimensioni.
• Vision Transformer: ViT codifica un'immagine per attività come la classificazione. DiT riutilizza il suo embedding delle patch e i blocchi transformer, aggiunge il condizionamento sul timestep tramite adaLN-Zero e produce una previsione per ogni patch.
• Diffusione latente e flow matching: un modello di diffusione latente è il framework complessivo che esegue il denoising nello spazio latente di un VAE, mentre la diffusione o il flow matching costituisce l'obiettivo di addestramento. DiT è la rete di denoising all'interno di quel framework e i generatori attuali abbinano backbone DiT a obiettivi di flow matching.

DiT per immagini e video da testo#

• DiT multimodale (MMDiT): Stable Diffusion 3 elabora i token delle immagini e quelli testuali con pesi separati per ciascuna modalità e li riunisce in un'unica operazione di attenzione, così le informazioni fluiscono in entrambe le direzioni tra il prompt e l'immagine. Questo ha migliorato la tipografia e il rispetto dei prompt; Stable Diffusion 3.5 Large porta il progetto a circa 8 miliardi di parametri.
• FLUX.1: FLUX.1 di Black Forest Labs è un transformer a flusso rettificato con 12 miliardi di parametri per la generazione text-to-image.
• Video: il rapporto tecnico di Sora di OpenAI descrive un transformer di diffusione che comprime i video in uno spazio latente e li suddivide in patch spazio-temporali, così clip di durata, risoluzione e rapporto d'aspetto diversi diventano sequenze di token di lunghezza variabile. Per la generazione video, il più grande modello Movie Gen di Meta è un transformer da 30 miliardi di parametri il cui contesto di 73.000 token video corrisponde a 16 secondi di video a 16 fotogrammi al secondo.