Consistent Diffusion Language Models

1. Enunciato del Problema

I Modelli di Diffusione per il Linguaggio (DLM) offrono un'alternativa promettente ai modelli autoregressivi (AR) abilitando la generazione parallela di token e l'inferenza in tempo sublineare. Tuttavia, rimane un collo di bottiglia critico: generare testo di alta qualità richiede tipicamente centinaia di passi di raffinamento, annullando i guadagni di efficienza promessi dal parallelismo.

Nei domini continui (ad esempio, immagini), i Modelli di Coerenza hanno accelerato con successo la diffusione addestrando i modelli a essere invarianti lungo una traiettoria unica e deterministica definita dall'Equazione Differenziale Ordinaria del Flusso di Probabilità (PF-ODE).

• Il Divario: Nei domini discreti (come il testo), non esiste un PF-ODE analogo nello spazio dei campioni. I processi di corruzione discreta (ad esempio, mascheramento o rumore uniforme) non seguono un singolo percorso deterministico; invece, coinvolgono transizioni stocastiche.
• La Conseguenza: Adattare direttamente i metodi di coerenza continui ai dati discreti è mal definito. Le soluzioni esistenti spesso si affidano a pipeline di distillazione multi-stadio (utilizzando un modello insegnante) o a surrogati di rilassamento continuo, che sono complesse, computazionalmente costose e spesso soffrono di collasso modale o perdita di diversità.

2. Metodologia: Coerenza Discreta Multi-Percorso (MPDC)

Gli autori propongono un cambiamento fondamentale: invece di cercare un percorso deterministico inesistente, sfruttano i ponti posteriori esatti intrinseci ai processi di diffusione discreta.

Concetto Chiave: Il Ponte Posteriore

Per qualsiasi due livelli di rumore $s < t$ in un processo di diffusione discreto, esiste un ponte posteriore esatto in forma chiusa $q(x_s | x_t, x_0)$. Questo ponte definisce un percorso stocastico valido che collega uno stato rumoroso $x_t$ a uno stato più pulito $x_s$, condizionato ai dati veri $x_0$.

• A differenza del PF-ODE continuo, questo ponte è stocastico.
• Crucialmente, questi ponti formano un semigruppo (proprietà di Chapman-Kolmogorov), il che significa che i percorsi possono essere composti transitivamente.

Il Principio MPDC

Gli autori introducono la Coerenza Discreta Multi-Percorso (MPDC), che addestra un denoiser a essere invariante per percorso in media.

• Definizione: Un modello $f_\theta(x_t, t)$ è coerente se la sua previsione in uno stato rumoroso $x_t$ concorda (in media) con la previsione effettuata dopo aver compiuto un "salto" intermedio verso uno stato più pulito $x_s$ tramite il ponte posteriore esatto.
• Obiettivo: Il modello minimizza la divergenza tra la sua previsione diretta da $x_t$ e la previsione attesa da $x_s$ (dove $x_s$ è campionata dal ponte).
• Formulazione Matematica:
  $$\mathcal{L}_{CDLM} = \mathbb{E}_{t, \delta, x_t} \left[ \sum_{i \in \mathcal{M}(x_t)} w(t, \delta) \cdot D\left( f_\theta(x_t, t)_i \parallel \tilde{f}_\theta(x_s, s)_i \right) \right]$$
  Dove $x_s \sim q(x_s | x_t, x_0)$, $\tilde{f}_\theta$ è una rete target (EMA) e $D$ è una metrica di divergenza (tipicamente Divergenza Jensen-Shannon).

Scelte Chiave di Progettazione per la Stabilità

Per prevenire il collasso modale (dove il modello genera testo ripetitivo per soddisfare trivialmente la coerenza) e la deriva uniforme, gli autori introducono tre meccanismi stabilizzanti:

1. Dimensione del Passo come Curriculum: Campionare dimensioni del passo $\delta$ da un intervallo specifico (ad esempio, da $1/8$ a $3/8$) e ponderarle con $1/\delta$ per bilanciare i segnali di apprendimento tra percorsi brevi e lunghi.
2. Ancoraggio alla Diffusione a Passo Massimo: Una piccola frazione della funzione di perdita di addestramento è ancorata ai dati veri ($s=0$) utilizzando un obiettivo di diffusione standard. Questo impedisce al modello di convergere verso soluzioni costanti degeneri.
3. Asimmetria di Ottimizzazione e Divergenza: Utilizzare un stop-gradient sulla rete target e impiegare la Divergenza Jensen-Shannon (JSD) invece della KL Forward/Backward. La JSD offre un compromesso stabile tra la ricerca di modalità e la copertura delle modalità, evitando l'instabilità della divergenza KL quando si addestra da zero.

3. Contributi Chiave

1. Un Nuovo Principio per la Modellazione Generativa Discreta: Il documento stabilisce che i ponti posteriori esatti fungono da sostituto rigoroso e in forma chiusa per il PF-ODE mancante, trasformando l'ostacolo della "accelerazione discreta" in una risorsa analitica.
2. Un Framework di Addestramento Unificato: CDLM fornisce un obiettivo a stadio singolo e senza insegnante. Gli autori dimostrano che la diffusione mascherata standard, i modelli di coerenza continui, la distillazione progressiva e i recenti metodi di distillazione a due stadi (come SDTT e DUO-DCD) sono tutti casi speciali, limiti o accoppiamenti approssimati del framework CDLM.
3. Prestazioni all'Avanguardia: CDLM ottiene risultati superiori sia nella generazione di testo incondizionata che condizionata, superando modelli base forti e uguagliando o superando modelli complessi distillati in più stadi senza richiedere un insegnante.

4. Risultati Sperimentali

Gli autori hanno valutato CDLM (nello specifico MCDLM per la diffusione mascherata e UCDLM per la diffusione uniforme) rispetto a baseline come MDLM, DUO, SDTT e DUO-DCD sul dataset OpenWebText.

• Generazione Incondizionata:

  • Efficienza a Pochi Passi: CDLM stabilisce nuovi risultati all'avanguardia (SOTA) per modelli a stadio singolo. Supera significativamente i modelli base (MDLM, DUO) in tutti i budget di campionamento.
  • Confronto con la Distillazione: In modo notevole, il CDLM a stadio singolo spesso supera i modelli distillati in più stadi (SDTT, DUO-DCD) in termini di Perplexity (PPL) mantenendo un'entropia molto più alta (diversità).
  • Collasso Modale: I modelli distillati che utilizzano campionatori greedy (ad esempio, DUO-DCD) hanno mostrato bassa entropia (3.9), indicando un grave collasso modale (caratteri ripetitivi). CDLM ha mantenuto un'entropia sana (5.0+), dimostrando la sua capacità di preservare la diversità.
  • Velocità: CDLM ottiene accelerazioni da 32x a 64x rispetto alle baseline autoregressive e da 64x a 128x rispetto al MDLM standard, mantenendo una qualità comparabile.

• Generazione Condizionata:

  • CDLM ha dimostrato prestazioni robuste nel recuperare frasi originali da input corrotti al 50% su più dataset (OpenWebText, Lambada, Wikitext-103, PTB).
  • Ha ottenuto punteggi BLEU migliori e una perplexità inferiore rispetto a SDTT e DUO, in particolare nel preservare le condizioni di input dove altri modelli fallivano.

5. Significato e Impatto

• Unificazione Teorica: Il documento fornisce una "lente unificante" per la modellazione generativa discreta, mostrando che tecniche di accelerazione diverse sono fondamentalmente connesse attraverso il concetto di coerenza del percorso.
• Scalabilità: Eliminando la necessità di distillazione multi-stadio e modelli insegnante, CDLM offre una base principiata, scalabile e senza insegnante per la prossima generazione di modelli generativi discreti veloci.
• Generalizzabilità: Sebbene dimostrato sul testo, il framework si applica a qualsiasi dominio discreto con un ponte posteriore trattabile (ad esempio, sequenze biologiche, grafi, sintesi di programmi).
• Cambiamento di Paradigma: Riformula la diffusione discreta veloce non come un'approssimazione di ODE continui, ma come l'apprendimento di un denoiser indipendente dal percorso su ponti stocastici, risolvendo con successo il problema della generazione a pochi passi negli spazi discreti.

In conclusione, CDLM colma il divario tra l'efficienza dei modelli di coerenza continui e le realtà dei dati discreti, offrendo una soluzione robusta a stadio singolo che raggiunge una generazione di testo veloce e ad alta fedeltà senza la complessità delle pipeline di distillazione.