Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

1. Il Problema

I modelli di diffusione discreta (utilizzati per la generazione di testo e immagini) richiedono un numero elevato di passaggi di campionamento per produrre dati di alta qualità. A differenza dei modelli autoregressivi (LLM) che generano token uno alla volta, i modelli di diffusione operano su blocchi di token, garantendo un'alta utilizzazione degli acceleratori hardware. Tuttavia, la loro natura iterativa comporta costi computazionali elevati e un numero di FLOP (operazioni in virgola mobile) significativamente più alto rispetto ai modelli autoregressivi.

Le sfide principali sono:

• Inefficienza: La necessità di centinaia o migliaia di passaggi (NFE - Number of Function Evaluations) rende l'inferenza lenta.
• Fallimento delle distillazioni esistenti: I metodi di distillazione sviluppati per la diffusione continua (che riducono i passaggi a pochi) falliscono spesso quando applicati al dominio discreto, portando al collasso della qualità o della diversità dei campioni.
• Limiti delle metriche: Le metriche standard come la "perplexity generativa" sono ingannevoli per i modelli di diffusione discreta, poiché possono premiare campioni privi di senso ma ad alta densità di probabilità (es. ripetizioni di parole).

2. Metodologia: Discrete Moment Matching Distillation (D-MMD)

Gli autori propongono D-MMD, un algoritmo di distillazione che estende il concetto di Moment Matching Distillation (MMD), precedentemente efficace solo nel dominio continuo, al dominio discreto.

Concetti Chiave:

• Generalizzazione dell'MMD: L'algoritmo MMD originale cerca di far sì che l'aspettativa condizionale dei dati puliti sia identica tra la distribuzione dei dati e quella del generatore distillato. D-MMD generalizza questa formulazione per processi di diffusione discreta (sia Masked che Uniform).
• Formulazione Min-Max: Il metodo utilizza un'ottimizzazione alternata tra tre modelli:
  1. Generatore Studente ($\hat{x}_\eta$): Cerca di minimizzare la perdita rispetto al modello insegnante, ma massimizzarla rispetto a un modello ausiliario.
  2. Modello Insegnante ($\hat{x}_\theta$): Fisso, fornisce il target di distribuzione.
  3. Modello Ausiliario ($\hat{x}_\phi$): Viene addestrato per imitare l'aspettativa del generatore studente, ma regolarizzato per rimanere vicino all'insegnante.
• Adattamento al Discreto: Poiché il campionamento discreto (categorico) non è differenziabile, D-MMD non utilizza campioni "hard" per il gradiente. Invece, utilizza il vettore di probabilità "soft" ($\hat{x}_\eta$) e opera direttamente nello spazio dei log-probabilità. La funzione di perdita per lo studente diventa una differenza tra le cross-entropy rispetto all'insegnante e all'ausiliario.
• Gestione del Bias: Per il modello ausiliario, il paper distingue tra diffusione masked (dove è possibile usare target soft) e uniform (dove è necessario usare campioni hard per correggere il bias).
• Condizionamento al Rumore: Per la diffusione mascherata, l'introduzione di una sorgente di rumore di input aggiuntiva è cruciale per permettere al generatore di collassare la sua distribuzione fattorizzata e generare output correlati di alta qualità.

Come funziona l'addestramento (Algorithm 1):

1. Si campiona un dato $x$ e si diffonde fino a un tempo $t$.
2. Lo studente genera una distribuzione di probabilità $\hat{x}_\eta(z_t)$.
3. Si campiona un token $x$ da questa distribuzione e si risale al tempo $s$ ($z_s$) tramite la posterior.
4. Si calcola la perdita alternata: lo studente cerca di ingannare l'ausiliario (massimizzando la perdita dell'ausiliario) mentre cerca di allinearsi all'insegnante.

3. Contributi Chiave

1. Primo metodo di distillazione discreta di successo: D-MMD è in grado di distillare modelli di diffusione discreta in generatori a pochi passaggi mantenendo alta qualità e diversità, risolvendo il problema del collasso osservato in metodi precedenti (come SDTT).
2. Superamento dell'Insegnante: Un risultato controintuitivo ma significativo è che i generatori distillati (studenti) spesso superano le prestazioni dei loro insegnanti originali, specialmente quando si usano un numero ridotto di passaggi. Questo è dovuto alla natura avversariale dell'ottimizzazione che spinge il modello verso le mode della distribuzione (mode-seeking) invece di coprirle tutte (mode-covering), simile a quanto avviene con la riduzione della temperatura.
3. Nuova Metrica di Valutazione (Gradient Moment): Gli autori introducono il GPT-2 Gradient Moment (GPT-2 GM). Poiché la likelihood dei modelli di diffusione non è trattabile, non si può usare la perplexity standard. La nuova metrica misura la norma quadrata del gradiente di un modello autoregressivo di riferimento (GPT-2) rispetto ai campioni generati. Se il gradiente è vicino a zero, i campioni sono indistinguibili dai dati reali per il modello di riferimento.
4. Generalità: Il metodo funziona sia per diffusione mascherata (adatta al testo) che uniforme (adatta alle immagini).

4. Risultati Sperimentali

I risultati sono stati valutati su dataset di immagini (CIFAR-10) e testo (Open Web Text).

• Immagini (CIFAR-10):

  • I modelli D-MMD raggiungono un FID (Fréchet Inception Distance) molto migliore rispetto agli insegnanti con un frazione dei passaggi.
  • Esempio: Un modello Masked D-MMD ottiene un FID di 3.5 in soli 64 passaggi, superando l'insegnante che richiede 1024 passaggi per un FID di 6.4.
  • Confronto con lo stato dell'arte (es. Di4C): D-MMD supera i metodi esistenti con meno passaggi (es. FID 5.0 in 8 passaggi vs FID 8.0 in 40 passaggi per Di4C).

• Testo (Open Web Text):

  • Utilizzando la metrica GPT-2 GM, i generatori D-MMD superano gli insegnanti.
  • Un generatore Masked D-MMD con soli 16 passaggi ottiene un punteggio di 0.236, superando l'insegnante che ne richiede 256 (punteggio 0.275).
  • I modelli distillati mostrano una migliore frontiera di Pareto tra numero di passaggi e qualità.

• Generazione Block-Autoregressive:

  • In uno scenario ibrido (generazione a blocchi condizionata da un encoder autoregressivo), un generatore D-MMD a 16 passaggi eguaglia le prestazioni di un insegnante a 256 passaggi.

5. Significato e Implicazioni

Questo lavoro è fondamentale perché:

• Riduce i costi computazionali: Abilita l'uso pratico dei modelli di diffusione discreta per applicazioni in tempo reale, riducendo drasticamente il numero di passaggi necessari senza sacrificare la qualità.
• Colma il divario di performance: Dimostra che i modelli di diffusione, se correttamente distillati, possono competere o superare i modelli autoregressivi standard in termini di qualità di generazione, pur mantenendo il vantaggio del campionamento parallelo.
• Nuovo paradigma di valutazione: Introduce una metrica robusta (Gradient Moment) che supera i limiti della perplexity generativa, offrendo uno strumento migliore per valutare la qualità dei campioni di modelli generativi non autoregressivi.
• Flessibilità: La capacità di distillare sia modelli masked che uniform rende D-MMD un approccio versatile per la generazione multimodale (testo e immagini).

In sintesi, D-MMD rappresenta un passo avanti cruciale verso l'efficienza e l'adozione pratica dei modelli di diffusione discreta, trasformandoli da modelli lenti e costosi in generatori veloci e di alta qualità.