Diffusion As Self-Distillation: End-to-End Latent Diffusion In One Model

Di /
stato della ricerca deep learning

Perché “Diffusion As Self-Distillation” è una buona notizia per i generative model

Il paper propone Diffusion Self-Distillation (DSD), un framework che fonde encoder, decoder e diffusion model in un unico network end-to-end, evitando l’architettura a tre pezzi tipica dei Latent Diffusion Model standard.
Gli autori mostrano che un tentativo ingenuo di addestrare tutto insieme porta a “latent collapse”, cioè a una rappresentazione latente che perde dimensionalità ed espressività, e propongono una lettura del problema attraverso la lente della self-distillation non supervisionata.

La parte interessante, anche per chi studia foundation model di visione, è che DSD collega in modo esplicito diffusion model e self-distillation (alla DINO/BYOL) e dimostra che, con piccole modifiche all’obiettivo di training, è possibile avere un unico Vision Transformer che funge da tokenizer, autoencoder e modello di diffusion.
Su ImageNet 256×256, DSD ottiene FID pari a 13.44, 6.38 e 4.25 con modelli da 42M, 118M e 205M parametri, senza classifier-free guidance e con solo 50 epoche di training, competendo con Latent Diffusion Model molto più grandi (fino a ~700M di parametri).

Paper

Indice

Dentro DSD: un solo modello per encodare, decodere e diffondere

Il problema: la modularità dei Latent Diffusion Model

I Latent Diffusion Model classici usano un Variational AutoEncoder (encoder + decoder) separato dal diffusion network, spesso con tre reti distinte addestrate in più fasi, ognuna con il proprio obiettivo.
Questa struttura è costosa da addestrare, poco efficiente in inferenza e soprattutto non si integra bene con il paradigma “single backbone” dei vision foundation model (per esempio un singolo Vision Transformer più pochi head).

In più, il VAE viene normalmente addestrato solo con una reconstruction loss, senza tener conto delle esigenze del diffusion model che userà quel latent space, portando a un’ottimizzazione sub-ottimale a livello di sistema.
Molti lavori recenti migliorano o il tokenizer (per esempio VAVAE o detok) o il diffusion model (per esempio DiT, SiT, FasterDiT, REPA), ma raramente affrontano il problema come modello unico end-to-end.

L’idea di base: diffusion come self-distillation

Gli autori osservano che un Latent Diffusion Model, visto dal punto di vista dell’encoder, assomiglia moltissimo a un metodo di self-distillation tipo DINO o BYOL: un encoder produce una rappresentazione latente usata in due rami, uno passato direttamente alla loss e uno passato attraverso un predictor più complesso (qui il diffusion head).
Nei metodi di self-distillation ben progettati esiste un meccanismo (rank differentiation) che impedisce il collasso dimensionale, mentre nel caso naive di end-to-end Latent Diffusion questo meccanismo viene violato.

Partendo dall’analogia, il paper rilegge la loss di diffusion come una particolare forma di self-distillation non supervisionata, in cui il diffusion head gioca il ruolo del predictor e il rumore gioca il ruolo della “vista perturbata” dei dati.
Questa rilettura permette di capire perché il training ingenuo collassa e quali modifiche introdurre per riportare il sistema nel regime stabile tipico della self-distillation moderna.

Le due radici del “latent collapse”

Dall’analisi teorico-sperimentale, gli autori identificano due cause principali del collasso quando si prova a fare un end-to-end training ingenuo:

  • La diffusion loss standard tende a comprimere la varianza nel latent space, spingendo le rappresentazioni a raggrupparsi attorno a un unico punto.
  • Il diffusion head produce un’uscita ad alta dimensionalità effettiva, influenzata dal rumore, che viola la condizione di rank differentiation richiesta per la stabilità della self-distillation.

Per misurare questo fenomeno gli autori usano la “effective rank” della matrice di embedding, che cala rapidamente verso valori molto bassi nei setup che collassano e cresce in quelli stabilizzati da DSD.

Come funziona DSD: dal decoupling alla loss “pulita”

Step 1: decoupling con stop-gradient

Il primo passo è eliminare il canale di gradiente che spinge direttamente a ridurre la varianza del latent, usando un operatore di stop-gradient sulla rappresentazione “target” nella diffusion loss.
In pratica, il diffusion head continua a vedere la rappresentazione latente, ma la parte della loss che dipende solo dal target non ha più un gradiente che schiaccia il latent space, mantenendo la sua ricchezza informativa.

Questa idea richiama la struttura teacher-student tipica della self-distillation, dove il teacher è aggiornato con Exponential Moving Average e di solito non riceve gradiente diretto dalla loss principale.
Applicata al contesto diffusion, questa modifica rallenta il collasso ma ancora non basta a risolverlo completamente, perché il problema di rank differentiation rimane.

Step 2: trasformare la loss in clean latent prediction

La seconda mossa chiave è reinterpretare l’obiettivo della diffusion non come previsione della “velocity” tra rumore e immagine pulita, ma come previsione diretta del latent pulito a partire da un latent rumoroso.
Gli autori mostrano che le due formulazioni sono matematicamente equivalenti dal punto di vista dell’ottimo di training, ma la versione “predict clean latent” ha proprietà di stabilità molto migliori.

Intuitivamente, obbligare il diffusion head a produrre direttamente il latent pulito significa trasformarlo in un denoiser, cioè in un filtro a bassa frequenza che riduce naturalmente la dimensionalità effettiva del suo output.
Questo rende valida la condizione di rank differentiation (il target ha rank effettivo più alto del predictor), che nei metodi di self-distillation è ciò che impedisce il collasso anche in assenza di explicit regularizer sulla distribuzione dei feature.

Step 3: abbracciare fino in fondo la self-distillation

Una volta introdotti decoupling e clean latent prediction, gli autori proseguono fino in fondo nell’analogia con la self-distillation moderna:

  • Inseriscono un target encoder aggiornato con EMA, che fornisce target più stabili e smussati nel tempo.
  • Usano forti data augmentation (masking, blur, color jitter, solarization) sul ramo “online”, in stile DINOv3, così che il predictor impari a riportare molte viste rumorose allo stesso latent pulito.
  • Aggiungono anche una self-distillation loss sugli stati intermedi del Vision Transformer per allineare in profondità il backbone.

In più, per la fase di generazione viene aggiunto un secondo diffusion head “detached”, addestrato con una velocity loss standard ma senza retropropagare gradiente verso il backbone, così da non disturbare la rappresentazione appresa tramite DSD.

Architettura DSD: un ViT, tanti head

Un unico Vision Transformer come backbone

DSD usa un solo Vision Transformer con patch 16×16 come backbone condiviso, che produce una sequenza di latent token compatti (256 token con dimensione molto bassa).
Su questo backbone vengono innestati diversi head: uno per l’encoding/decoding delle immagini, uno per la clean latent prediction, uno per la velocity e uno per la classificazione ImageNet.

Il modello agisce contemporaneamente da tokenizer (come un VAE), da autoencoder (ricostruzione immagine) e da diffusion model (generazione a partire dal rumore nello spazio latente).
Grazie al design unificato, non c’è più bisogno di preaddestrare un VAE separato: tutto viene ottimizzato insieme rispetto a obiettivi condivisi, coerenti con il compito di generazione finale.

Le loss usate durante il training

Nel training entrano tre gruppi principali di loss, tutte calcolate sulla stessa passata avanti nel backbone:

  • Una DSD loss che fa self-distillation tra online e target encoder, con il predictor che mappa un latent rumoroso verso il latent “pulito” del teacher.
  • Una detached velocity loss su un head separato, usata solo per avere una velocity ben stimata durante il sampling, ma senza toccare il backbone.
  • Una reconstruction loss che fa funzionare il backbone anche come autoencoder, usando un decoder conv per ricostruire l’immagine a partire da un latent parzialmente rumoroso.

A queste si aggiungono loss ausiliarie di allineamento con un teacher esterno (per esempio DINOv3) su alcuni layer del ViT, e una classification loss su ImageNet, che aiutano a strutturare un latent space semanticamente ricco e stabile.

Risultati: cosa guadagniamo davvero con DSD

Setup sperimentale in breve

DSD viene valutato su ImageNet 256×256 con generazione condizionata alla classe, seguendo i protocolli standard di valutazione (FID, sFID, Inception Score) su 50k immagini generate.
Il training riprende quasi interamente il setup di LightningDiT (optimizer Muon, patch size, numero di step di sampling, ecc.), così da rendere il confronto equo con la famiglia DiT/SiT/LightningDiT.

Gli autori addestrano tre dimensioni di modello: DSD-S (42M parametri), DSD-M (118M) e DSD-B (205M), tutte su 8 GPU NVIDIA A800 con tempi per epoca dell’ordine di qualche decina di minuti.
Il backbone ViT usa pochi layer ma una struttura di token molto compatta, con un numero ridotto di canali nel latent space, a conferma dell’enfasi sulla rappresentazione efficiente piuttosto che sulla pura scala di parametri.

Qualità di generazione e scaling

Il risultato forse più impressionante è che DSD-B ottiene FID 4.25 senza classifier-free guidance, con soli 205M di parametri e 50 epoche, raggiungendo o superando modelli Latent Diffusion che usano combinazioni VAE + diffusion da oltre 700M di parametri e training più lunghi.
Anche DSD-M e DSD-S mostrano un chiaro miglioramento scalando la dimensione del modello, indicando che il framework segue una buona legge di scaling pur rimanendo param-efficient.

Quando si attiva la guidance, le metriche migliorano ulteriormente e rimangono competitive con metodi state-of-the-art, pur mantenendo l’architettura unificata a singolo backbone.
Nel confronto di Tabella 1 con modelli autoregressivi (per esempio MaskGIT, LlamaGen, VAR, MAGViT-v2) e con la famiglia DiT/SiT/REPA/LightningDiT, DSD ottiene un ottimo compromesso tra qualità di generazione, numero di parametri e numero di epoche richieste.

Efficienza e limiti dichiarati

DSD elimina del tutto il costo di un VAE preaddestrato separatamente, sia in termini di parametri che in termini di complessità ingegneristica della pipeline, riducendo la latenza complessiva del modello generativo.
Il design a ViT unico si sposa bene con la visione di foundation model multimodali in cui lo stesso backbone può essere riusato per classificazione, representation learning e generazione.

Gli autori però sottolineano che, per limiti di risorse, non hanno ancora scalato DSD a dimensioni paragonabili ai più grandi baseline, e che non hanno esplorato in profondità il potenziale di DSD come puro metodo di representation learning auto-supervisionato al di là della generazione.
Questi aspetti rappresentano linee di ricerca aperte interessanti, soprattutto per chi lavora su pretraining unificato per visione e generazione.

Cose da capire bene prima di leggere il paper

Latent Diffusion Model e VAE tokenizer

Un Latent Diffusion Model classico comprime l’immagine in uno spazio latente tramite un VAE e applica il diffusion process in questo spazio compresso, per ridurre il costo rispetto al lavorare direttamente sui pixel.
Questo significa che la qualità del latent space (quanto è ricco, strutturato e regolare) influenza fortemente sia la ricostruzione che la qualità della generazione finale.

Capire cosa fanno encoder e decoder del VAE e come il diffusion model li usa aiuta a motivare perché ha senso voler unificare tutto in un unico backbone.
Nel contesto DSD, il ViT prende il ruolo sia di tokenizer sia di backbone generativo, quindi è utile avere chiaro il pipeline standard dei Latent Diffusion Model prima di vedere la variante unificata.

Self-distillation e rank differentiation

Nei metodi di self-distillation (DINO, DINOv2, DINOv3, BYOL, I-JEPA) c’è un encoder online e un encoder target, di solito aggiornato con EMA, e una loss che spinge l’output dell’online a imitare quello del target sotto forti augmentation.
Nonostante non ci siano espliciti vincoli di dispersione, questi metodi non collassano perché il predictor apprende a “filtrare” le informazioni, e il target mantiene un rank effettivo più alto, innescando un loop di arricchimento della rappresentazione.

Il lavoro RankDiff ha formalizzato questa intuizione introducendo la nozione di rank differentiation come condizione per la stabilità della self-distillation: la rappresentazione target deve avere dimensionalità effettiva maggiore di quella proiettata.
DSD si appoggia pesantemente a questa teoria per diagnosticare perché il diffusion head tradizionale rompe la condizione e propone modifiche per ripristinarla.

Dimensional collapse ed effective rank

Per “dimensional collapse” si intende il caso in cui le rappresentazioni latenti, che dovrebbero occupare una regione ad alta dimensionalità dello spazio, si concentrano intorno a pochi assi o addirittura a un singolo punto, perdendo informazione.
Nel paper questo fenomeno viene misurato tramite l’effective rank, una metrica continua che riassume quante direzioni significative sono presenti in una matrice di feature, osservando il suo andamento durante il training.

Vedere come l’effective rank del latent e del predictor evolve nei vari ablation (naive joint training, solo decoupling, decoupling + loss transformation, aggiunta dell’EMA, aggiunta delle augmentation) è fondamentale per capire perché certe modifiche funzionano.
Per leggere bene queste parti è utile familiarizzarsi concettualmente con il fatto che “più alto effective rank” corrisponde a una rappresentazione più espressiva e meno collassata.

Vision Transformer come backbone generativo

DSD usa un Vision Transformer standard come backbone unico, con patch embedding e token aggiuntivi per migliorare la capacità di rappresentazione, come in vari tokenizer moderni.
Capire come un ViT processa una sequenza di token e come gli head specifici (encoder/decoder, diffusion, classifier) si innestano sopra è utile per seguire la descrizione dell’architettura unificata.

Molti modelli generativi moderni (DiT, LightningDiT) già usano Transformer nello spazio latente, quindi DSD si inserisce naturalmente in questa linea di ricerca, ma con la novità dell’autoencoder integrato.
Avere chiaro il ruolo degli head e il flusso di dati tra token rumorosi, token puliti e immagini aiuta a leggere il paper senza perdersi nei dettagli implementativi.

Quiz per fissare le idee (con soluzioni)

Domande

  1. Qual è il problema principale che gli autori osservano quando si prova a addestrare end-to-end encoder, decoder e diffusion model senza alcuna modifica all’objective?
  2. Perché la loss di diffusion standard tende a comprimere la varianza nel latent space?
  3. In che senso l’output del diffusion head tradizionale viola la condizione di rank differentiation tipica della self-distillation?
  4. Qual è il vantaggio di riformulare l’obiettivo di training come clean latent prediction invece che come velocity prediction, dal punto di vista della stabilità?
  5. Che ruolo giocano EMA target encoder e forti data augmentation nel framework DSD?
  6. Perché DSD può essere considerato un “single-backbone generative model”, e quali vantaggi offre rispetto ai Latent Diffusion Model modulari?
  7. In che cosa DSD si differenzia dal lavoro REPA-E riguardo al modo in cui la diffusion loss interagisce con il latent space?

Risposte commentate

  1. Il training congiunto ingenuo porta a un forte latent collapse, cioè la reconstruction loss non scende e l’effective rank del latent cade rapidamente verso valori molto bassi, segnalando che la rappresentazione perde dimensionalità.
  2. La loss di diffusion contiene implicitamente un termine che dipende solo sulla varianza del latent e che, quando ottimizzato con gradiente, spinge l’encoder a ridurre la dispersione delle rappresentazioni, facendo convergere tutti i latent verso un punto medio.
  3. Il diffusion head tradizionale somma al suo output un rumore ad alta dimensionalità effettiva, producendo rappresentazioni con rank molto alto, spesso superiore a quello del target, rompendo la condizione “target rank > predictor rank” che la teoria di RankDiff indica come necessaria per evitare il collasso.
  4. Predire direttamente il latent pulito costringe il predictor a comportarsi come un denoiser, cioè un filtro a bassa frequenza che sopprime le componenti casuali, riducendo la dimensionalità effettiva dell’output e ripristinando la condizione di rank differentiation favorevole alla stabilità.
  5. L’EMA target encoder fornisce target più stabili e meno rumorosi, facilitando l’apprendimento, mentre le forti data augmentation obbligano il predictor a riportare molte viste deformate allo stesso latent pulito, rafforzando il comportamento da filtro a bassa frequenza e la robustezza della rappresentazione.
  6. DSD usa un singolo Vision Transformer come backbone per encoding, decoding, diffusion e classificazione, evitando VAE e diffusion network separati, il che semplifica la pipeline, riduce i parametri totali e rende più naturale riutilizzare il modello come foundation model di visione.
  7. REPA-E isola il VAE dalla diffusion loss staccando del tutto il gradiente sul latent, così che la diffusion non possa guidare l’apprendimento della rappresentazione, mentre DSD usa stop-gradient solo sul target e una loss riformulata in modo che la diffusion continui a plasmare il latent senza causare collasso.

Studi collegati: DSD nella costellazione di diffusion e self-distillation

End-to-end latent generative model e REPA-E

Latent Score-Based Generative Models (LSGM) avevano già esplorato l’idea di addestrare insieme encoder, decoder e modello generativo, usando però un termine di regolarizzazione esplicito sulla varianza del latent per evitare il collasso.
DSD si differenzia perché non aggiunge un nuovo regularizer ma rilegge l’intero objective come self-distillation, modificando la struttura della loss affinché la stabilità emerga dalle proprietà del predictor.

REPA-E, invece, propone un “partial autoencoder” che riusa feature intermedie del diffusion U-Net per la ricostruzione e risolve il problema del collasso staccando i gradienti della diffusion loss rispetto al VAE, mantenendo però encoder, decoder e diffusion come moduli distinti.
DSD si pone come passo successivo, riuscendo a unificare l’architettura in un singolo ViT e offrendo anche una diagnosi teorica del collasso basata su rank differentiation che REPA-E non fornisce.

Self-distillation per representation learning (DINO, DINOv3, RankDiff)

I metodi di self-distillation come DINO, DINOv2, DINOv3 e BYOL hanno mostrato che è possibile apprendere rappresentazioni visive di alta qualità senza supervisioni esplicite, usando teacher EMA, augmentation forti e predictor leggeri.
Il lavoro RankDiff analizza perché questi metodi non collassano, introducendo la condizione di rank differentiation e mostrando che, se il predictor riduce la dimensionalità effettiva rispetto al target, il training rafforza progressivamente la ricchezza della rappresentazione.

DSD importa direttamente queste idee nel mondo delle diffusion, reinterpretando il diffusion head come predictor di self-distillation e progettando loss e architettura in modo che la condizione di RankDiff sia soddisfatta.
Per chi conosce già DINOv3, leggere DSD è quindi naturale: molte scelte (augmentation, EMA, allineamento dei layer intermedi) sono parallele, ma applicate ora a un modello generativo.

Distillation per accelerare i diffusion model

Parallelamente, esiste una linea di lavori che usa distillation per accelerare modelli di diffusion già addestrati, riducendo il numero di step necessari in sampling tramite teacher-student training.
Per esempio, lavori come “SELF-DISTILLATION FOR DIFFUSION” e “Few-Step Diffusion via Score Identity Distillation” distillano un diffusion model in versioni più veloci, concentrandosi sul tempo di generazione piuttosto che sulla struttura del latent space o sull’architettura unificata.

Rispetto a questi approcci, DSD non parte da un modello preaddestrato da distillare ma ripensa la fase di training principale, integrando già in origine concetti di self-distillation per ottenere stabilità ed efficienza parametrica.
In altre parole, mentre la distillation “classica” per diffusion è un post-processing di compressione, DSD è un disegno di training di base che mira a unire representation learning e generazione in un solo schema.

Self-distillation per task specifici (zero-shot customization)

Un’altra direzione collegata è usare self-distillation per adattare diffusion model a compiti specifici senza dover fare per-instance fine-tuning costoso.
Per esempio, “Diffusion Self-Distillation for Zero-Shot Customized Image Generation” usa un meccanismo di self-distillation per personalizzare il comportamento di un modello text-to-image senza supervisioni pesanti, mostrando che il paradigma funziona anche su task di customization.

Questi lavori confermano che il concetto di self-distillation è estremamente flessibile nel contesto dei diffusion model, dalla compressione alla personalizzazione.
DSD aggiunge a questo quadro la prospettiva “di sistema”, mostrando come la self-distillation possa diventare il principio organizzatore dell’intero design di un generative model, non solo un trucco accessorio.

Torna in alto