Sintesi Tecnica: Set Diffusion

Problematica

I modelli di diffusione discreta sono emersi come un'alternativa competitiva ai modelli linguistici autoregressivi (AR), offrendo velocità di generazione più elevate e una migliore controllabilità. Tuttavia, gli approcci di diffusione esistenti affrontano limitazioni significative:

1. Vincoli di Lunghezza Fissa e Bidirezionalità: La maggior parte dei modelli di diffusione discreta si basa su architetture che richiedono un contesto bidirezionale completo, impedendo l'uso del caching Key-Value (KV) e limitando la generazione a lunghezze fisse.
2. Rigidità della Block Diffusion (BD3LM): Sebbene la Block Diffusion (BD3LM) colmi parzialmente il divario generando blocchi sequenziali di dimensione fissa da sinistra a destra, soffre di una mancanza di flessibilità. Non può supportare inserimenti in posizioni arbitrarie (infilling) o la generazione a lunghezza variabile all'interno di un blocco, e gli aggiornamenti della cache KV sono ritardati fino a quando l'intero blocco non viene decodificato.
3. Limitazioni del Trade-off: Gli attuali metodi interpolano tra AR e diffusione principalmente variando una dimensione di blocco fissa, il che restringe la flessibilità di decodifica e il parallelismo.

La sfida centrale è sviluppare un modello linguistico che interpoli tra autoregressione e diffusione non solo tramite la dimensione del blocco, ma variando l'ordine in cui vengono generati gli insiemi di token, abilitando così la generazione a lunghezza variabile, il caching KV e la decodifica in ordine arbitrario, mantenendo al contempo un'alta qualità del campione.

Metodologia

Gli autori introducono Set Diffusion, un framework che fattorizza il processo di generazione su insiemi di token a posizione e lunghezza flessibile, anziché su blocchi fissi.

1. Modello Probabilistico

La diffusione di set definisce una distribuzione di probabilità autoregressiva su una sequenza di insiemi di token $\sigma = (\sigma_1, \dots, \sigma_N)$, dove ogni $\sigma_n$ è un sottoinsieme di posizioni di token. La verosimiglianza è fattorizzata come:
$$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$$
Qui, $\pi$ è una politica di ordine che seleziona il prossimo insieme di posizioni, e $p_\theta$ modella la verosimiglianza condizionale dei token in quell'insieme utilizzando la diffusione. Il modello ottimizza un Negative Evidence Lower Bound (NELBO) che somma le perdite di diffusione su questi insiemi.

2. Architettura Set-Causal e Caching KV

Una chiave innovazione architettonica è il transformer set-causal. A differenza dei normali modelli di diffusione bidirezionali, questa architettura supporta il caching KV dopo ogni passo di inferenza.

• Meccanismo: Quando un insieme di token $x_{\sigma_n}$ viene generato, i suoi vettori key e value vengono memorizzati in cache e riutilizzati per condizionare le predizioni successive.
• Efficienza di Addestramento: Gli autori propongono un algoritmo di addestramento vettorializzato che elabora simultaneamente il contesto pulito e le query rumorose in un singolo passaggio forward utilizzando una maschera di attenzione specializzata. Questa maschera assicura che i token rumorosi in un insieme prestino attenzione solo ad altri token rumorosi nello stesso insieme e a tutti i token puliti negli insiemi precedenti.

3. Interpolazione degli Ordini di Token

Per controllare il trade-off tra qualità del campione (streitezza del limite) ed efficienza dell'inferenza (parallelismo), gli autori introducono schedules di tempo di rivelazione dipendenti dalla posizione.

• Schedules di Offset di Posizione: Ogni posizione di token $\ell$ riceve un tempo di rivelazione casuale $R_\ell$ estratto da una distribuzione definita da una funzione di ripartizione (CDF) $\alpha^\ell_\tau$.
• Istanziazione Sliding-Window (SW-SetDLM): Gli autori propongono un'istanza specifica in cui i token hanno un intervallo di generazione attivo $[a_\ell, a_\ell + w]$. La dimensione della finestra $w$ controlla il bias "da sinistra a destra":
  • Un $w$ piccolo (che tende a $1/L$) concentra l'ordine sulla generazione da sinistra a destra, recuperando il comportamento AR e rendendo più stretto l'ELBO.
  • Un $w$ grande (che tende a 1) permette una generazione indipendente dall'ordine, recuperando la diffusione a maschera standard.
  • Valori intermedi consentono la decodifica a finestra scorrevole (sliding-window decoding), permettendo al modello di predire insiemi di token a lunghezza flessibile in posizioni arbitrarie, supportando compiti come l'infilling.

4. Inferenza

Durante l'inferenza, il modello campiona le posizioni candidate in base all'ordine temporale corrente e allo schedule di offset di posizione. Successivamente, denoisa queste posizioni in parallelo utilizzando la cache KV. I token accettati vengono confermati e la cache viene aggiornata immediatamente, consentendo la generazione a lunghezza variabile e una decodifica efficiente.

Contributi Chiave

1. Framework Set Diffusion: Una nuova classe di modelli linguistici che fattorizza la generazione su insiemi di token a posizione e lunghezza flessibile, generalizzando l'autoregressione, la diffusione indipendente dall'ordine e la block diffusion.
2. Architettura Set-Causal: Un'architettura transformer che supporta aggiornamenti della cache KV dopo ogni passo di inferenza, abilitando una generazione a lunghezza variabile efficiente.
3. Interpolazione dell'Ordine: Un metodo rigoroso per interpolare tra AR e diffusione sintonizzando gli schedule di ordinamento dipendenti dalla posizione (specificamente gli schedule a finestra scorrevole) piuttosto che dimensioni di blocco fisse.
4. Prestazioni allo Stato dell'Arte: Dimostrazione che Set Diffusion raggiunge un trade-off velocità-qualità superiore in ragionamento matematico, riassunto e generazione incondizionata rispetto ai precedenti modelli di diffusione, e supera significativamente la Block Diffusion nei compiti di infilling.

Risultati Sperimentali

Gli autori hanno valutato Set Diffusion (specificamente SW-SetDLM) su diversi benchmark:

• Ragionamento Matematico (GSM8K): Set Diffusion raggiunge una frontiera Pareto velocità-accuratezza strettamente migliore rispetto alla Block Diffusion. Ad esempio, con una dimensione di finestra di decodifica $S \le 8$, ottiene un'accuratezza pass@1 del 66.41% con un throughput di $\ge 60.42$ token/sec, superando la Block Diffusion ($S=8$) che ottiene il 63.53% di accuratezza a $\ge 55.39$ token/sec.
• Infilling (ROCStories): Set Diffusion supera sostanzialmente la Block Diffusion nei compiti di infilling (riempire 1 o 3 frasi centrali), ottenendo punteggi ROUGE più elevati con una velocità di decodifica fino a circa il 25% superiore. La Block diffusion fatica in questo ambito a causa della sua rigida struttura a blocchi da sinistra a destra.
• Riassunto (CNN/DailyMail): Set Diffusion ottiene punteggi ROUGE competitivi rispetto ai baseline di diffusione e una decodifica fino al 10% più veloce rispetto alla Block Diffusion. Supera anche un baseline AR, che gli autori notano tendere all'overfitting su questo dataset.
• Generazione Incondizionata e Verosimiglianza: Su OpenWebText e LM1B, Set Diffusion migliora i punteggi MAUVE rispetto alla Block Diffusion pur decodificando più velocemente. Su LM1B, raggiunge la migliore perplessità tra i modelli di diffusione (27.83 PPL per $S \le 8$) mantenendo un throughput superiore rispetto ai modelli di diffusione precedenti.

Significato e Rivendicazioni

Il paper sostiene che Set Diffusion rappresenti un avanzamento significativo decuplicando l'interpolazione tra autoregressione e diffusione dalle strutture rigide di blocco. Interpolando tramite l'ordine dei token, il framework offre:

• Maggiore Flessibilità: Supporto per inserimenti in posizioni arbitrarie e generazione a lunghezza variabile, elementi critici per compiti come l'infilling.
• Efficienza: La capacità di aggiornare le cache KV dopo ogni passo, una caratteristica precedentemente non disponibile nei modelli di diffusione, portando a un'inferenza più veloce.
• Prestazioni: Un trade-off superiore tra qualità del campione e velocità di inferenza rispetto agli esistenti modelli di linguaggio di diffusione.

Gli autori posizionano Set Diffusion come un framework unificante che recupera AR, diffusione indipendente dall'ordine e block diffusion come casi speciali, fornendo un meccanismo sintonizzabile per bilanciare parallelismo e qualità del campione in base ai requisiti specifici del compito di inferenza.