Diffusion Transformers (DiT) — L'Architettura Dietro Sora, Stable Diffusion 3 e Flux

Cos'è: "Scalable Diffusion Models with Transformers" è il paper pubblicato da William Peebles (allora UC Berkeley, poi OpenAI come co-lead di Sora) e Saining Xie (NYU) a dicembre 2022. Propone di sostituire la backbone UNet — l'architettura usata da Stable Diffusion, DALL-E 2, Imagen e tutti i diffusion model fino a quel momento — con un transformer puro che opera sui patch di immagine, analogo concettualmente al Vision Transformer (ViT) di Google. Il paper dimostra che l'architettura scala secondo leggi di potenza pulite: più compute uguale qualità migliore, con FID che migliora prevedibilmente con l'aumentare di GFLOPs. Tre anni dopo è la base architetturale di Sora, Stable Diffusion 3 e Flux.

L'eredità UNet: perché è stata la scelta standard per cinque anni

I diffusion model moderni nascono con DDPM (Denoising Diffusion Probabilistic Models, Ho et al. 2020) e tutti — senza eccezione fino al 2022 — usavano una variante della UNet, un'architettura convolutional introdotta nel 2015 per la segmentazione medica. La UNet ha tre proprietà che la rendevano naturalmente adatta al denoising: è completamente convolutional (gestisce immagini di dimensioni variabili), ha skip connection simmetriche tra encoder e decoder (preservano i dettagli ad alta frequenza durante la ricostruzione), e ha un bias induttivo locale che è ottimo quando il segnale ha forte struttura spaziale.

Stable Diffusion 1.x e 2.x, DALL-E 2, Imagen di Google, eDiff-I di NVIDIA: tutti usavano UNet con cross-attention iniettate per condizionare sul testo. Il problema era lo scaling: la UNet non scala bene oltre una certa dimensione. Aumentare la profondità o la larghezza dei blocchi convolutional dava ritorni decrescenti, e l'architettura non rispondeva bene alle ricette di scaling che funzionavano nei LLM (più parametri, più dati, più compute = qualità migliore in legge di potenza).

L'idea DiT: patch tokenizzati e transformer puro

Peebles e Xie hanno applicato al diffusion model il principio che aveva già funzionato nei Vision Transformer (Dosovitskiy et al., Google, 2020): trattare un'immagine come una sequenza di patch. Nel paper DiT, una latente di dimensione 32×32×4 (l'output dell'autoencoder di Stable Diffusion) viene divisa in patch di 2×2 e linearizzata in una sequenza di 256 token, ognuno una proiezione lineare del patch più un position embedding. Da qui in avanti l'architettura è un transformer standard: N blocchi di self-attention più feed-forward, niente skip connection, niente convoluzioni.

La novità tecnica più interessante è il modo in cui il modello viene condizionato: oltre al rumore e al timestep di diffusion, il modello deve "ricevere" il segnale di classe o di testo. Gli autori sperimentano tre varianti: in-context conditioning (concatenare l'embedding di condizione ai token), cross-attention (come nella UNet di Stable Diffusion), e adaLN-zero — una variante di Adaptive Layer Norm dove i parametri di scala e shift di ogni LayerNorm sono generati dall'embedding di condizione, e sono inizializzati a zero (così il modello inizia come un transformer non condizionato e impara gradualmente a usare il segnale). adaLN-zero risulta sperimentalmente migliore di entrambe le alternative ed è la scelta che resta dominante nei modelli successivi.

Le scaling laws della generazione visiva

Il contributo più importante del paper non è l'architettura in sé, ma la dimostrazione che la qualità di generazione segue leggi di potenza pulite rispetto al compute. Gli autori addestrano quattro varianti — DiT-S, DiT-B, DiT-L, DiT-XL — da 33M a 675M parametri, e mostrano che il FID (Frechet Inception Distance, metrica standard di qualità di generazione su ImageNet) migliora monotonicamente in legge di potenza con i GFLOPs spesi. Questa è la stessa proprietà di scaling osservata da OpenAI nei LLM con Kaplan et al. 2020, applicata per la prima volta in modo rigoroso al dominio visivo.

Questo risultato è quello che ha convinto OpenAI, Stability e gli altri laboratori che il futuro dei modelli generativi visivi era nei transformer. La UNet aveva un soffitto sotto al quale aumentare i parametri smetteva di pagare; i DiT scalano linearmente nel regime testato e — ipotesi confermata dai modelli successivi — continuano a scalare oltre. È il momento "GPT-3" della generazione visiva: la dimostrazione che la ricetta "transformer + più compute" funziona anche qui.

L'eredità: Sora, Stable Diffusion 3, Flux

A febbraio 2024 OpenAI annuncia Sora, il primo modello video text-to-video di qualità cinematografica. Il technical report cita esplicitamente Peebles & Xie come riferimento architetturale: Sora è "un modello di diffusion che opera su patch spazio-temporali di rappresentazioni latenti video", essenzialmente un DiT esteso al video. Bill Peebles, primo autore del paper DiT, è co-lead del progetto Sora in OpenAI. La stessa estensione 3D — patch che coprono blocchi spazio-temporali di N frame — è alla base di Veo (Google DeepMind) e Kling (Kuaishou).

A febbraio 2024 Stability AI rilascia Stable Diffusion 3, basato su un'architettura chiamata MMDiT (Multimodal Diffusion Transformer): una variante di DiT che processa simultaneamente token di testo e token visivi nella stessa sequenza di attenzione, invece di iniettare il testo via cross-attention. È un'evoluzione dell'idea originale che riconosce esplicitamente DiT come base di partenza nel technical report. A luglio 2024 Black Forest Labs — fondata da ex-Stability tra cui Robin Rombach, autore originale di Stable Diffusion — rilascia Flux, anch'esso basato su MMDiT con 12B parametri, considerato il modello text-to-image open weights state-of-the-art fino al 2025.

In meno di tre anni dal paper originale, il DiT è passato da architettura accademica sperimentale a fondazione di tutti i principali modelli generativi visivi commerciali. La UNet, dominante per cinque anni, è oggi un'architettura legacy che sopravvive solo nei modelli più piccoli o nelle distillazioni mirate alla velocità di inferenza locale.

---