💻 computer science

Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis

Nemotron-Labs-Diffusion-Image è un modello di diffusione discreta mascherata allo stato dell'arte che potenzia la sintesi testo-immagine ad alta risoluzione introducendo un meccanismo di editing dei token per la correzione dinamica dell'inferenza e un obiettivo di Cross-Entropy raggruppata con un operatore fuso personalizzato per superare la scarsità di segnale in contesti a vocabolario ampio, raggiungendo prestazioni superiori nei benchmark GenEval, DPG e HPSv3.

Autori originali: Shufan Li, Greg Heinrich, Hanrong Ye, Yonggan Fu, Aditya Grover, Jan Kautz, Pavlo Molchanov

Pubblicato 2026-06-30

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Shufan Li, Greg Heinrich, Hanrong Ye, Yonggan Fu, Aditya Grover, Jan Kautz, Pavlo Molchanov

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di cercare di insegnare a un robot come dipingere un capolavoro basandosi su una descrizione che gli dai, come "un gatto con un cappello in una giornata di sole". Per molto tempo, i migliori robot hanno utilizzato un metodo simile alla scultura da un blocco di argilla: partivano da un ammasso sfocato e lo perfezionavano lentamente, passo dopo passo, correggendo gli errori man mano che procedevano. Questo è chiamato Diffusione Continua.

Recentemente, un nuovo approccio chiamato Diffusione Discreta è diventato popolare. Invece dell'argilla, questo metodo utilizza un enorme sacco di mattoncini Lego (token). Il robot inizia con una parete di mattoncini nascosti e li rivela uno alla volta. Il problema? Una volta che un mattoncino viene rivelato, è bloccato. Se il robot sceglie il mattoncino sbagliato per l'orecchio del gatto, non può tornare indietro a sostituirlo. È come uno scultore che, una volta scolpito un pezzo di pietra, non può mai più correggere un errore.

Il documento presenta Nemotron-Labs-Diffusion-Image, un nuovo "scultore" che risolve due problemi principali legati a questo approccio Lego.

1. Il tasto "Annulla" (Editing dei Token)

Il Problema: Nella pittura Lego standard, se posizioni un mattoncino rosso dove dovrebbe essercene uno blu, sei costretto a conviverci. Il robot non ha modo di dire: "Aspetta, ho cambiato idea", e correggerlo in seguito. Questo porta a piccoli errori che si accumulano finché l'immagine non appare strana.

La Soluzione: Gli autori hanno dato al robot un tasto "Annulla". Lo hanno addestrato non solo a rivelare i mattoncini nascosti, ma anche a guardare i mattoncini che ha già posizionato e dire: "In realtà, quello sembra sbagliato; cambiamolo".

L'Analogia: Immagina uno scultore che lavora su una statua. Nel vecchio metodo, una volta rimosso un pezzo di pietra, era perso per sempre. In questo nuovo modo, lo scultore può rimuovere delicatamente un pezzo appena rimosso, levigarlo e rimodellarlo se non sembra corretto. Questo permette al robot di perfezionare iterativamente l'immagine, correggendo gli errori mentre procede, proprio come potevano fare i vecchi modelli di "argilla".

2. Il problema del "Grande Dizionario" (Cross-Entropy Raggruppata)

Il Problema: Per creare immagini di alta qualità e dettagliate, il robot ha bisogno di un vocabolario massiccio di mattoncini Lego (un "codebook"). Più mattoncini ha, più dettagliata sarà l'immagine. Tuttavia, con un dizionario di oltre 100.000 mattoncini unici, il robot vede raramente lo stesso specifico mattoncino due volte durante l'addestramento. È come cercare di imparare una lingua in cui vedi la parola "elefante" solo una volta ogni pochi anni. Il robot si confonde perché il segnale è troppo rado.

Inoltre, l'addestramento standard tratta ogni mattoncino sbagliato come ugualmente cattivo. Se il robot ha bisogno di un mattoncino "blu chiaro" ma ne sceglie uno "blu scuro", il vecchio metodo urla: "SBAGLIATO!" con la stessa intensità con cui urlerebbe se avesse scelto un mattoncino "rosso". Non riesce a capire che "blu chiaro" e "blu scuro" sono in realtà vicini nel significato.

La Soluzione: Gli autori hanno introdotto una nuova regola di addestramento chiamata Cross-Entropy Raggruppata (GCE).

L'Analogia: Invece di punire il robot per aver scelto il mattoncino esatto sbagliato, l'insegnante dice: "Hai scelto un mattoncino blu scuro quando volevo un blu chiaro. È vicino! Ti do un credito parziale".
Hanno organizzato i oltre 100.000 mattoncini in piccoli gruppi (come "Blu", "Rossi", "Verdi"). Durante l'addestramento, se il robot sceglie un mattoncino dal gruppo corretto (anche se non è quello perfettamente esatto), riceve un segnale positivo. Questo aiuta il robot a imparare le relazioni tra i mattoncini, rendendo molto più facile imparare un vocabolario enorme senza perdersi.

3. La Spinta di Velocità (Operatore Personalizzato)

Il Problema: Fare questo calcolo "raggruppato" è computazionalmente pesante. Di solito richiede molta memoria del computer (VRAM) e rallenta le cose, rendendo difficile l'addestramento su modelli enormi.

La Soluzione: Il team ha costruito uno strumento specializzato e personalizzato (un "operatore fuso") per eseguire questo calcolo.

L'Analogia: Immagina di fare una contabilità complessa. Il vecchio modo era usare una calcolatrice, scrivere il numero, prendere una penna, scriverlo in un registro e poi usare di nuovo la calcolatrice. Il nuovo modo è una macchina specializzata che esege il calcolo e scrive la voce nel registro in un unico movimento, fulmineo. Questo ha risparmiato una quantità enorme di memoria e tempo.

I Risultati

Il documento dimostra che questo nuovo robot, Nemotron-Labs-Diffusion-Image, è un campione:

Qualità Migliore: Crea immagini più nitide e accurate rispetto ai precedenti metodi basati sui Lego.
Più Veloce: Può generare immagini in meno passaggi perché può correggere i propri errori lungo il percorso.
Efficiente: Può gestire un vocabolario enorme di dettagli senza mandare in crash la memoria del computer.

In breve, hanno preso un metodo rigido e incline agli errori, gli hanno dato la capacità di autocorreggersi, lo hanno istruito a comprendere il "quartiere" del suo vocabolario e hanno costruito un motore più veloce per far girare il tutto. Il risultato è un generatore di immagini ad alta risoluzione che è allo stesso tempo più intelligente ed efficiente.

Riepilogo Tecnico: Nemotron-Labs-Diffusion-Image

Definizione del Problema

Sebbene i Modelli di Diffusione Latente (LDM) abbiano raggiunto prestazioni allo stato dell'arte nella sintesi text-to-image, la generazione di immagini discrete offre vantaggi distinti, tra cui la compatibilità nativa con i Large Language Models (LLM) e la capacità di sfruttare le tecniche di ottimizzazione consolidate degli LLM (ad es. sequence packing, KV caching). Tuttavia, gli esistenti Masked Discrete Diffusion Models (MDM) affrontano due sfide fondamentali che ne ostacolano la scalabilità e la fedeltà rispetto agli LDM continui:

Mancanza di Autocorrezione: A differenza degli LDM continui, che raffinano iterativamente le rappresentazioni latenti, gli standard MDM si impegnano su un token una volta che questo viene de-mascherato. Poiché i token discreti non possono essere modificati dopo la de-mascheratura, gli errori si accumulano durante il processo di inferenza. L'ipotesi standard di campionamento di indipendenza inter-token fallisce nel catturare le forti dipendenze presenti nei token delle immagini, portando a distribuzioni congiunte subottimali.
Sparsità del Codebook in Vocabolari Grandi: La generazione di immagini ad alta fedeltà richiede codebook (vocabolari) ampi per ridurre l'errore di discretizzazione. Tuttavia, l'aumento della dimensione del vocabolario porta a una grave sparsità dei segnali di addestramento; i singoli token appaiono infrequenti nel dataset. La standard Cross-Entropy (CE) one-hot tratta tutti i token non corrispondenti alla verità fondamentale come ugualmente negativi, fallendo nel catturare le somiglianze semantiche tra i token che sono vicini nello spazio di embedding ma distinti nel codebook. Ciò crea difficoltà di ottimizzazione e richiede una capacità di modello e dati massicci per la convergenza.

Metodologia

Gli autori propongono Nemotron-Labs-Diffusion-Image (NLD-Image), un modello fondazionale da 8 miliardi di parametri progettato per affrontare queste sfide attraverso innovazioni architettoniche e di obiettivo inedite.

1. Architettura: Transformer Decoder-Only

A differenza dei precedenti generatori di immagini discreti (ad es. Meissonic) che utilizzano un framework encoder-decoder con un encoder di testo congelato, NLD-Image impiega un singolo transformer decoder-only.

Inizializzazione: Il modello è inizializzato da un modello di diffusione linguistica pre-addestrato, sfruttando le esistenti capacità di comprensione del linguaggio.
Vantaggi: Questo design rimuove i vincoli di lunghezza del contesto degli encoder congelati (ad es. il limite di 77 token di CLIP), supporta prompt più lunghi e consente il sequence packing per l'efficienza dell'addestramento. L'head di language modeling finale viene sostituito con un nuovo MLP inizializzato per predire i token delle immagini.

2. Meccanismo di Token Editing (Autocorrezione)

Per consentire il raffinamento iterativo simile alla diffusione continua, gli autori introducono un meccanismo di token editing.

Addestramento: Il processo di diffusione forward è aumentato includendo la corruzione dei token. Oltre a mascherare un sottoinsieme di token, il modello corrompe una porzione dei token visibili (puliti). Il modello è addestrato a predire le probabilità per tutte le posizioni: de-mascherando i token mascherati e correggendo i token visibili corrotti.
Inferenza: Durante il campionamento, il modello non solo de-maschera i token, ma valuta anche i token già de-mascherati. Se il modello predice un token diverso con una confidenza superiore a una soglia $\tau$ , il token esistente viene sostituito. Ciò consente al modello di autocorregger l'errore dinamicamente durante il processo di generazione.

3. Obiettivo Grouped Cross-Entropy (GCE)

Per mitigare la sparsità del codebook, gli autori propongono la Grouped Cross-Entropy (GCE), una funzione di perdita gerarchica.

Clustering: Il grande codebook (ad es. 131k token) è raggruppato in gruppi più piccoli (ad es. 16k, 8k) utilizzando il K-means basato sulla prossimità dello spazio di embedding.
Funzione di Perdita: L'obiettivo combina la standard CE one-hot a livello di token con perdite CE ausiliarie a livello di cluster. Il modello predice le probabilità per tutti i token, e queste vengono sommate per ottenere le probabilità del cluster. La perdita è applicata sia al token specifico della verità fondamentale che al cluster della verità fondamentale a molteplici granularità.
Beneficio: Questo fornisce supervisione positiva ai token che sono semanticamente simili (nello stesso cluster) alla verità fondamentale, anche se non sono l'esatto match, alleviando così la scarsità del segnale.
Implementazione: Un operatore fuso personalizzato è implementato per calcolare la GCE in modo efficiente, riducendo significativamente l'uso di VRAM e la latenza rispetto alle implementazioni naive di PyTorch.

Contributi Chiave

Nemotron-Labs-Diffusion-Image: Un modello di diffusione discreta mascherata da 8B all'avanguardia per la sintesi text-to-image ad alta risoluzione (1024px).
Token Editing: Un meccanismo innovativo che consente ai MDM di raffinare iterativamente e autocorregere gli output durante l'inferenza, affrontando la mancanza di auto-raffinamento nei modelli discreti standard.
Grouped Cross-Entropy (GCE): Un obiettivo di addestramento che allevia la sparsità del codebook assegnando segnali di apprendimento positivi ai token semanticamente vicini, migliorando la stabilità e l'efficienza dell'ottimizzazione per i generatori a vocabolario ampio.
Implementazione Efficiente: Un operatore fuso personalizzato per la GCE che riduce l'impronta di memoria e l'overhead computazionale, facilitando l'addestramento di modelli a vocabolario ampio.

Risultati Sperimentali

Il modello è stato addestrato per 300k step su 64 GPU H100 utilizzando una strategia di upscaling progressivo (da 256px a 1024px).

Prestazioni nei Benchmark:
- GenEval: Ha raggiunto un punteggio di 0.90, superando modelli specialistici come Meissonic (0.54) e modelli multimodali unificati come LaViDa-O (0.77), e competendo con modelli di frontiera come Qwen-Image-2507 (0.87).
- DPG: Ha raggiunto 86.9, superando SD3-Medium (83.5) e Flux-Dev.
- MJHQ-30k: Ha raggiunto un punteggio HPSv3 di 10.76 (dopo il fine-tuning), indicando un alto allineamento con la preferenza umana.
Studi di Ablazione:
- Token Editing: I risultati qualitativi e quantitativi (HPSv3) mostrano che il token editing migliora costantemente la fedeltà dell'immagine e riduce l'accumulo di errori. Consente al modello di raggiungere una qualità comparabile con 32 step di campionamento (NFE) che precedentemente richiedevano 64 step senza editing, raddoppiando efficacementamente la velocità di inferenza per un target di qualità fissato.
- GCE: Nella generazione condizionale per classe su ImageNet, la GCE ha costantemente superato sia la vanilla Cross-Entropy che la SNCE (Soft Nearest Neighbor Cross-Entropy) nei punteggi FID attraverso diversi modelli e epoche di addestramento.
Efficienza: L'operatore GCE ottimizzato ha ridotto la latenza da 44ms a 20ms e l'uso massimo di VRAM da 25GB a 16GB rispetto alle implementazioni eager. NLD-Image ha dimostrato una velocità di inferenza 42.4x più veloce rispetto a Emu3.5 (AR) pur ottenendo punteggi GenEval più elevati.
Generazione a pochi Step: Il modello genera immagini di ragionevole qualità in soli 4 step senza distillazione, convergendo più velocemente dei baselines di flow-matching continuo che spesso rimangono sfocati a bassi conteggi di step.

Significato e Rivendicazioni

Il paper sostiene che NLD-Image rappresenti un avanzamento significativo nella generazione di immagini discrete, riuscendo a colmare il divario di prestazioni tra gli LDM continui e gli MDM discreti. Introducendo il token editing, gli autori dimostrano che i modelli discreti possono superare la loro intrinseca mancanza di autocorrezione, consentendo il raffinamento iterativo. Inoltre, l'obiettivo GCE fornisce una soluzione scalabile al problema della sparsità del codebook, permettendo l'uso di vocabolari ampi ed espressivi senza costi di addestramento proibitivi.

Gli autori posizionano questo lavoro come un passo verso i modelli multimodali unificati, sfruttando l'architettura decoder-only per elaborare testo e immagini all'interno di un unico framework. I risultati suggeriscono che i modelli di diffusione discreta, quando dotati di queste specifiche innovazioni, possono raggiungere una fedeltà ed efficienza allo stato dell'arte, aprendo la strada a sistemi di sintesi di immagini più potenti e scalabili.