DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Riepilogo Tecnico: DiLaDiff: Diffusione Latente Aumentata Distillata per la Modellazione del Linguaggio

Enunciato del Problema
I modelli di diffusione linguistica (DLM) affrontano un compromesso fondamentale tra qualità del campionamento e throughput. Sebbene i modelli di diffusione discreta (in particolare i Modelli di Linguaggio a Diffusione Mascherata, o MDLM) siano efficaci per dati categoriali, falliscono intrinsecamente nel catturare le correlazioni tra i token decodificati. Questa limitazione deriva dal fatto che il denoiser approssima la posterior come un prodotto di marginali indipendenti per token. Di conseguenza, la generazione parallela di più token porta a frasi incoerenti, costringendo a fare affidamento su una generazione sequenziale o a una riduzione del parallelismo che limita severamente il throughput. Al contrario, i modelli di diffusione continua applicati alle rappresentazioni testuali (ad esempio, embedding) offrono tecniche di campionamento e distillazione principiate, ma spesso faticano con il compromesso espressività-decodificabilità, dove le rappresentazioni contestuali sono difficili da decodificare nuovamente in token discreti, oppure non riescono a eguagliare le prestazioni dei corrispettivi discreti.

Metodologia
Gli autori propongono DiLaDiff (Modello di Linguaggio a Diffusione Latente Aumentata Distillata), un framework ibrido continuo-discreto progettato per risolvere il problema della correlazione tra token mantenendo un throughput elevato. L'architettura è composta da tre componenti principali:

1. Auto-Encoder Latente-Aumentato (LaDiff):

   • Encoder: Un encoder basato su transformer ($E_\phi$) apprende uno spazio latente continuo ($z$) dalle rappresentazioni contestuali (derivate dagli stati nascosti di BERT). Questo spazio latente è progettato per catturare le correlazioni globali tra token e le informazioni semantiche.
   • Decoder: Un decoder discreto ($x_\theta$) è inizializzato da un MDLM pre-addestrato. Genera token testuali condizionati alla variabile latente $z$.
   • Addestramento: Il sistema è addestrato utilizzando un Limite Inferiore della Evidenza (ELBO) modificato che combina l'auto-codifica con la diffusione mascherata. Crucialmente, la distribuzione posterior è fattorizzata sui token condizionati al latente $z$. Poiché $z$ cattura le correlazioni, l'assunzione di indipendenza condizionata vale, permettendo una generazione parallela accurata dei token.
   • Regolarizzazione: Per impedire che lo spazio latente diventi troppo sparso per la modellazione generativa, gli autori impiegano euristiche di rafforzamento della robustezza, tra cui la normalizzazione per coordinata, il mascheramento casuale di feature e latenti, e l'iniezione di rumore gaussiano.

2. Apprendimento del Prior Latente:

   • Poiché la sequenza pulita non è disponibile durante l'inferenza, un modello di diffusione continua (Denoiser Latente $z_\psi$) è addestrato per apprendere la distribuzione a priori $p(z)$ sull'output dell'encoder. Ciò permette al modello di campionare un latente $z$ di alta qualità tramite un'ODE di flusso probabilistico inverso prima della decodifica.

3. Distillazione della Coerenza (DiLaDiff):

   • Per ridurre ulteriormente l'overhead computazionale del passo di diffusione continua, gli autori applicano la distillazione MeanFlow. Un modello studente ($u_\eta$) è addestrato per prevedere la velocità media della traiettoria ODE tra due istanti temporali ($t$ e $r$), piuttosto che la velocità istantanea.
   • Questa distillazione permette di generare il latente $z$ in pochi passi (ad esempio, 5 passi) mantenendo prestazioni vicine a quelle del modello insegnante (LaDiff), che richiede molti passi (ad esempio, 200 passi).

Principali Contributi

• Architettura Ibrida: Il paper introduce una ricetta per l'addestramento di un auto-encoder testuale in cui il decoder è inizializzato da una baseline di diffusione discreta pre-addestrata. Ciò risulta in uno spazio latente controllabile e regolarizzato che cattura informazioni semantiche e correlazioni tra token.
• Parallelismo Guidato dal Latente: Apprendendo un prior latente con diffusione continua, il modello (LaDiff) permette il campionamento di più token in parallelo senza sacrificare la qualità del testo, poiché il latente risolve le dipendenze tra token che i modelli discreti non colgono.
• Prima Distillazione ODE di Successo per Latenti Testuali: Gli autori presentano DiLaDiff, il primo tentativo riuscito di distillare traiettorie ODE per latenti testuali contestuali. Ciò riduce il costo di generazione del latente a livelli trascurabili rispetto alla decodifica discreta.

Risultati
Gli esperimenti sono stati condotti sul dataset OpenWebText con una lunghezza di sequenza di 1024 token.

• Prestazioni vs Baseline: Anche senza distillazione, LaDiff supera la baseline di diffusione mascherata (MDLM). Nella generazione incondizionata, LaDiff ottiene un miglioramento assoluto di 28,5 nella Perplexità Generativa (GenPPL) e un miglioramento assoluto di 0,08 in MAUVE rispetto a MDLM con 1024 passi discreti.
• Throughput: LaDiff accelera l'inferenza di un fattore 7× a una dimensione del batch di 32 rispetto alla configurazione ottimale di MDLM.
• Efficienza della Distillazione: DiLaDiff raggiunge prestazioni vicine al suo insegnante LaDiff utilizzando solo 5 passi di diffusione latente (rispetto a 200 per l'insegnante). L'overhead computazionale della componente di diffusione latente scende a circa il 5% del tempo totale di parete, rendendo la generazione del latente trascurabile rispetto al costo della decodifica discreta.
• Generazione in Pochi Passi: DiLaDiff compete con i metodi all'avanguardia a pochi passi (come MDLM+SDTT e DUO+DCD) senza distillare esplicitamente il decoder discreto, dimostrando l'efficacia dell'aumento latente.
• Diversità Semantica: L'analisi mostra che lo spazio latente preserva le informazioni semantiche (verificate tramite BERTScore su latenti rumorosi) e permette una generazione diversificata, come dimostrato dalla capacità del decoder di mantenere l'entropia anche a temperature più basse.

Significato e Affermazioni
Il paper afferma che DiLaDiff colma con successo il divario tra l'espressività delle rappresentazioni continue e la robustezza della decodifica discreta. Catturando le correlazioni globali in uno spazio latente, il modello supera la limitazione intrinseca dei modelli di diffusione discreta riguardo alla generazione parallela di token. Gli autori sottolineano che il loro approccio permette una generazione di testo di alta qualità con una latenza di inferenza significativamente ridotta.

Il paper rimane modesto riguardo alle sue limitazioni:

• DiLaDiff non eguaglia ancora pienamente le prestazioni del suo modello insegnante (LaDiff) nel regime a pochi passi, suggerendo che è necessario ulteriore lavoro sulla distillazione.
• Il decoder discreto stesso non è distillato; gli autori notano che combinare la loro distillazione latente con la distillazione del decoder discreto è un'estensione futura naturale per raggiungere prestazioni a pochi passi davvero all'avanguardia.
• Le strategie di regolarizzazione e i design degli schedule di diffusione sono attualmente euristici, sebbene gli autori presentino un tentativo più principiato negli allegati.

In sintesi, DiLaDiff dimostra che un approccio ibrido continuo-discreto, aumentato dalla distillazione latente, può accelerare significativamente l'inferenza nella modellazione del linguaggio mantenendo o migliorando la qualità della generazione, risolvendo efficacemente il compromesso throughput-qualità intrinseco nei modelli di linguaggio a diffusione attuali.