Back to Basics: Let Denoising Generative Models Denoise

Di /
stato della ricerca deep learning

JiT Perché “tornare alle basi” del denoising

Il paper propone una critica molto semplice ma potente: i moderni diffusion models non stanno davvero facendo denoising, perché la rete non predice l’immagine pulita ma una forma di rumore o una combinazione rumorosa dell’immagine stessa. Gli autori sostengono che, se si prende sul serio l’idea che le immagini naturali vivono su una manifold di bassa dimensione, allora ha più senso far predire direttamente i dati puliti alla rete invece del rumore, soprattutto quando si lavora in spazi ad altissima dimensionalità come i pixel di immagini ad alta risoluzione.

Per dimostrare questa tesi introducono JiT (Just image Transformers), un diffusion model basato su Transformer che lavora direttamente sui pixel, con patch molto grandi (16×16 e 32×32), senza tokenizer, senza pre-training e senza loss aggiuntive come adversarial o perceptual. Su ImageNet a 256×256 e 512×512, JiT con x-prediction (predizione dell’immagine pulita) ottiene risultati competitivi, mentre varianti che predicono solo il rumore falliscono in modo catastrofico alle stesse dimensionalità. Gli autori rendono disponibile anche un’implementazione open source, tramite un repository GitHub chiamato “LTH14/JiT”.

Paper, Code

Indice

Dentro JiT: approccio e tecniche usate

Denoising diffusion models “classici”

I diffusion models generativi funzionano aggiungendo gradualmente rumore ai dati reali e poi imparando a invertire questo processo, ricostruendo immagini plausibili a partire da stati molto rumorosi. Nella pratica moderna, quasi tutti i diffusion models addestrano la rete a predire il rumore (o una quantità affine al rumore) dato lo stato noised corrente, perché questo si è dimostrato numericamente conveniente e stabile. Il paper parte da questa formulazione “standard” (inclusa la vista come flow/ODE) e la mette in discussione proprio nel punto che di solito si dà per scontato: cosa deve predire la rete.

x-prediction invece di noise-prediction

Gli autori distinguono in modo netto tra il far predire all’architettura l’immagine pulita (x-prediction) e il farle predire il rumore o una combinazione di immagine e rumore. La chiave è la manifold assumption: le immagini naturali occupano una regione di bassa dimensione all’interno dello spazio dei pixel, mentre le quantità rumorose sono distribuite molto più “piene” nello spazio. Se la rete lavora direttamente nello spazio delle immagini pulite, può sfruttare questa struttura, e quindi non ha bisogno di avere una larghezza comparabile alla dimensione del patch per modellare bene la distribuzione.

“Just image Transformers” in pixel space

JiT è essenzialmente un Diffusion Transformer (DiT) applicato non a token latenti ma a patch di pixel grezzi, con patch di dimensione 16×16 per 256×256 e 32×32 per 512×512. Ogni patch ha quindi centinaia o migliaia di canali (3 canali colore moltiplicati per molti pixel), cioè una dimensionalità di ingresso molto alta rispetto alle architetture che lavorano su latent tokens o su patch più piccoli. L’architettura è condizionata sul tempo di diffusione e sull’etichetta di classe tramite meccanismi come adaLN-Zero, in linea con i recenti DiT, ma senza componenti extra come VAE tokenizer, discriminator GAN o reti di feature per la perceptual loss.

Design minimalista: niente tokenizer, niente pre-training

A differenza di molti modelli come Stable Diffusion, che fanno generazione in uno spazio latente compresso e usano un VAE encoder/decoder addestrato separatamente, JiT resta interamente in pixel space. Questo significa niente fase di pre-training di un tokenizer, niente allineamento tra spazi latenti diversi e niente dipendenza da reti esterne per definire le loss. Il messaggio è che, se si sceglie correttamente “cosa” la rete deve predire (i dati puliti) e si rispetta la struttura di manifold dei dati, si può tornare a un design molto semplice e comunque competitivo.

Cosa mostrano davvero i risultati

Confronto sistematico tra tipi di predizione

Il cuore sperimentale del paper è una serie di esperimenti che confrontano diverse scelte su “cosa” predice la rete (immagine pulita, rumore, combinazioni) e “come” viene definita la loss, costruendo una griglia di varianti. Su risoluzioni moderate come 64×64, dove la dimensione per patch è relativamente bassa, tutte le combinazioni funzionano abbastanza bene e le differenze di qualità generativa sono limitate. Ma quando si passa a ImageNet 256×256 e 512×512, con patch ad alta dimensionalità, le varianti che predicono il rumore collassano, mentre le versioni x-prediction restano stabili e performanti.

JiT su ImageNet 256 e 512

JiT viene testato in configurazioni con patch 16×16 (per 256×256) e 32×32 (per 512×512), producendo immagini di qualità alta e competitiva rispetto ai diffusion models recenti. I risultati su metriche standard come FID e Inception Score mostrano che JiT non solo è stabile, ma può raggiungere performance comparabili a modelli più complessi, pur senza utilizzare latenti né pre-training su dataset esterni. In pratica, il paper dimostra che la combinazione “pixel space + patch ad alta dimensione + x-prediction + architettura relativamente compatta” è sufficiente per stare nella stessa fascia dei migliori modelli attuali.

Effetto della dimensionalità e del bottleneck

Gli autori analizzano in dettaglio il rapporto tra dimensione del patch e dimensione dello spazio nascosto della rete, mostrando che non è necessario che la larghezza del Transformer eguagli la dimensionalità del patch. Addirittura, una struttura con bottleneck (dove il numero di feature interne è minore della dimensione di input) può essere benefica, coerentemente con l’idea che i dati vivono su una manifold di dimensione minore. Esperimenti su dati sintetici, ottenuti proiettando dati di bassa dimensione in spazi molto più grandi tramite matrici casuali, rafforzano questa interpretazione.

Concetti chiave da padroneggiare prima di leggere il paper

Diffusion models e flow view

Il paper utilizza la formulazione moderna dei diffusion models che li collega a dinamiche continue e flow-based models, dove il processo di diffusione è visto come un flusso che trasforma gradualmente la distribuzione dei dati in una distribuzione di rumore semplice. La cosa importante per la lettura, al di là delle formule, è avere chiaro che il modello impara un campo di vettori che, passo dopo passo, “riporta” i campioni rumorosi verso il data manifold. Questo permette di capire perché è così centrale il tipo di quantità che si chiede di predire alla rete (immagine pulita o rumore).

Manifold dei dati vs spazio del rumore

L’assunzione di manifold dice che le immagini naturali non riempiono tutto lo spazio dei pixel, ma occupano una superficie molto più sottile e strutturata. Quando si aggiunge rumore a questi dati, ci si sposta fuori dalla manifold e si entra in regioni molto meno strutturate, dove il modello deve gestire variabilità in tutte le direzioni. Se la rete è addestrata a predire direttamente punti sulla manifold (l’immagine pulita), anche con un numero di feature interne più piccolo della dimensione del patch riesce a rappresentare bene la distribuzione, mentre se deve predire vettori di rumore completamente sparsi la capacità richiesta esplode.

x-prediction, noise-prediction, v-prediction

Nel linguaggio del paper, x-prediction indica la scelta di far predire alla rete l’immagine pulita a ogni passo, mentre le alternative sono far predire il rumore o una combinazione lineare di immagine e rumore, spesso chiamata v-prediction. In molti diffusion models di stato dell’arte si usano noise- o v-prediction perché forniscono buone proprietà numeriche e si integrano bene con implementazioni esistenti. Qui invece si mostra che, una volta che i patch diventano molto grandi e si lavora direttamente in pixel space, x-prediction è la scelta cruciale per evitare collassi e ottenere buone prestazioni.

Pixel space vs latent space

Modelli come Stable Diffusion lavorano in latent space: un VAE comprime l’immagine in una rappresentazione più compatta, e il diffusion model opera in quel dominio, spesso con un tokenizer e con loss aggiuntive come adversarial e perceptual per migliorare la qualità visiva. JiT rifiuta esplicitamente questa pipeline e torna al pixel space puro, rinunciando a qualunque tokenizer o loss esterna, per verificare quanto si possa ottenere con un design minimalista ma ben allineato con la struttura dei dati. Capire bene le differenze tra questi due mondi è fondamentale per apprezzare il contributo del paper, perché mostra che non è obbligatorio passare da latenti per scalare a risoluzioni alte.

Transformer patch-based e DiT

JiT è fortemente ispirato ai Diffusion Transformer (DiT), che applicano architetture tipo Vision Transformer alla generazione con diffusion. L’idea di base è dividere l’immagine in patch, linearizzarli come token e processarli con un Transformer, invece di usare una U-Net convoluzionale. Nel paper, questa architettura è resa ancora più “pura”: solo patch di pixel, conditioning sul tempo e sulla classe via adaLN-Zero, e nient’altro, così da isolare l’effetto della scelta di predizione.

Quiz su JIT

  • Domanda 1: Perché gli autori dicono che i diffusion models moderni non fanno davvero denoising “in senso classico”?
  • Risposta: Perché invece di far predire alla rete l’immagine pulita, chiedono di predire il rumore (o una combinazione rumorosa), quindi il modello non “ripulisce” i dati in output ma ricostruisce la quantità rumorosa da cui l’immagine viene poi derivata.
  • Domanda 2: Che cosa rende x-prediction particolarmente importante quando i patch di pixel sono molto grandi?
  • Risposta: Con patch ad altissima dimensionalità, il rumore vive in uno spazio enorme e poco strutturato, quindi predirlo direttamente richiede una rete molto più capace, mentre predire l’immagine pulita sfrutta il fatto che le immagini stanno su una manifold più compatta, permettendo a reti relativamente sotto-dimensionate di funzionare bene.
  • Domanda 3: In che senso JiT è “Back to Basics” rispetto a modelli come Stable Diffusion?
  • Risposta: Perché abbandona tokenizer latenti, pre-training su dataset esterni e loss complesse come adversarial e perceptual, e usa solo un Transformer sui pixel con x-prediction, dimostrando che un design minimalista può comunque raggiungere risultati competitivi su ImageNet ad alta risoluzione.
  • Domanda 4: Cosa mostrano gli esperimenti confrontando diversi tipi di predizione (immagine, rumore, combinazioni) su 64×64 rispetto a 256×256 e 512×512?
  • Risposta: Su 64×64 tutte le varianti funzionano abbastanza bene, mentre sulle risoluzioni alte le versioni che predicono il rumore collassano, e solo le varianti con x-prediction restano stabili e producono immagini di buona qualità.
  • Domanda 5: Perché un design con bottleneck interno può essere addirittura vantaggioso secondo il paper?
  • Risposta: Perché se i dati vivono su una manifold di bassa dimensione, costringere la rete a passare per uno spazio nascosto più piccolo della dimensione del patch la incoraggia a rappresentare proprio quella struttura di bassa dimensione, invece di sprecare capacità su direzioni di variazione non necessarie.

Lavori collegati: altri “ritorni alle basi” nella generazione

Diffusion Transformer (DiT) e modelli Transformer per diffusion

DiT è una famiglia di modelli che sostituiscono le U-Net con Transformer per i diffusion models, trattando patch o token latenti come sequenze su cui applicare attenzione. JiT può essere visto come una variante estremamente minimalista di questa idea, in cui il Transformer lavora direttamente su patch di pixel, senza modulo di compressione o pre-training, e serve a illustrare l’impatto della sola scelta di predizione. Per chi legge il paper, conoscere DiT aiuta a capire che il contributo non è tanto architetturale, quanto concettuale.

Latent diffusion e Stable Diffusion

I latent diffusion models, tra cui Stable Diffusion, hanno mostrato che spostare la generazione in un latent space compresso consente di raggiungere alta qualità su risoluzioni elevate con costi computazionali gestibili. Tuttavia, questi modelli dipendono da un VAE encoder/decoder pre-addestrato, da loss aggiuntive (perceptual, talvolta adversarial) e spesso da grosse pipeline di pre-training, cosa che JiT rifiuta esplicitamente. Il confronto tra latent diffusion e JiT evidenzia due filosofie diverse: compressione + complessità contro pixel space + semplicità.

Representation-Conditioned image Generation (RCG)

RCG, proposto dagli stessi autori, introduce un framework in cui un pixel generator è condizionato da rappresentazioni di alto livello, ottenendo risultati di punta nella generazione non condizionata su ImageNet 256×256. RCG integra bene diversi generatori esistenti, migliorandone di molto le metriche come FID e Inception Score, e mostra quanto sia potente usare buone rappresentazioni come guida per la generazione. Collocato accanto a JiT, suggerisce che ci sono due vie complementari: migliorare la componente di rappresentazione o ripensare il cuore del processo di denoising nei diffusion models.

Autoregressive image generation senza vector quantization

In “Autoregressive image generation without vector quantization”, sempre di Li e colleghi, si esplora un’altra direzione minimalista: fare autoregressive generation in pixel space o su rappresentazioni continue, evitando il passaggio per codici discreti e vector quantization. Questo lavoro elimina la complessità dei tokenizer quantizzati, proponendo architetture che modellano direttamente sequenze di feature continue. Concettualmente è affine a JiT perché mette in discussione lo strato di astrazione (token discreti vs dati continui) e mostra che design più diretti possono essere sufficienti.

xAR e predizione diretta dei contenuti visivi

Il lavoro “Next-X Prediction for Autoregressive Visual Generation” (xAR) propone modelli autoregressivi che predicono direttamente blocchi di contenuto visivo ad alta risoluzione, raggiungendo risultati competitivi su ImageNet a 256×256 e 512×512. xAR dimostra che, con architetture ben scalate, è possibile fare generazione autoregressiva efficiente senza affidarsi a token troppo compressi o quantizzati. Letto insieme a JiT, rafforza l’idea che predire direttamente “cosa si vede” (immagini pulite o patch continui) sia spesso più efficace che lavorare su rappresentazioni indirette come rumore o codici discreti.

Torna in alto