Technische Samenvatting: Set Diffusion

Probleemstelling

Discrete diffusiemodellen zijn opgekomen als een competitief alternatief voor autoregressieve (AR) taalmodellen, omdat ze snellere generatiesnelheden en verbeterde controleerbaarheid bieden. Bestaande diffusiebenaderingen kampen echter met significante beperkingen:

1. Vaste lengte en bidirectionele beperkingen: De meeste discrete diffusiemodellen vertrouwen op architecturen die volledige bidirectionele context vereisen, wat het gebruik van Key-Value (KV) caching verhindert en de generatie beperkt tot vaste lengtes.
2. Rigiditeit van Block Diffusion: Hoewel Block Diffusion (BD3LM) de kloof gedeeltelijk overbrugt door sequentiële blokken van vaste grootte van links naar rechts te genereren, lijdt het aan inflexibiliteit. Het kan geen willekeurige invoegingen op willekeurige posities (infilling) ondersteunen of variabele lengtes genereren binnen een blok, en KV-cache updates worden vertraagd tot een heel blok is gedecodeerd.
3. Trade-off beperkingen: Huidige methoden interpoleren tussen AR en diffusie primair door een vaste blokgrootte te variëren, wat de decodeervrijheid en parallellisme beperkt.

De kernuitdaging is het ontwikkelen van een taalmodel dat interpoleert tussen autoregressie en diffusie, niet alleen door de blokgrootte te variëren, maar door de volgorde waarin tokensets worden gegenereerd te variëren, waardoor flexibele generatie met variabele lengte, KV-caching en decodering in willekeurige volgorde mogelijk wordt, terwijl een hoge steekproefkwaliteit behouden blijft.

Methodologie

De auteurs introduceren Set Diffusion, een framework dat het generatieproces factoriseert over tokensets met flexibele posities en flexibele lengtes in plaats van vaste blokken.

1. Probabilistisch Model

Set diffusion definieert een autoregressieve waarschijnlijkheidsverdeling over een sequentie van tokensets $\sigma = (\sigma_1, \dots, \sigma_N)$, waarbij elke $\sigma_n$ een deelverzameling is van tokenposities. De waarschijnlijkheid is gefactoriseerd als:
$$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$$
Hierbij is $\pi$ een ordebeleid dat de volgende set posities selecteert, en $p_\theta$ modelleert de conditionele waarschijnlijkheid van tokens in die set met behulp van diffusie. Het model optimaliseert een Negative Evidence Lower Bound (NELBO) die de diffusieverliezen over deze sets optelt.

2. Set-Causal Architectuur en KV Caching

Een belangrijke architecturale innovatie is de set-causal transformer. In tegen tegenstelling tot standaard bidirectionele diffusiemodellen, ondersteunt deze architectuur KV-caching na elke inferentiestap.

• Mechanisme: Wanneer een tokenset $x_{\sigma_n}$ wordt gegenereerd, worden de key- en value-vectoren hiervan gecachet en hergebruikt om daaropvolgende voorspellingen te conditioneren.
• Efficiëntie van training: De auteurs stellen een vectoriseerbaar trainingsalgoritme voor dat schone context en ruisige queries gelijktijdig verwerkt in een enkele forward pass met behulp van een gespecialiseerde attention-mask. Deze mask zorgt ervoor dat ruisige tokens in een set alleen aandacht besteden aan andere ruisige tokens in dezelfde set en aan alle schone tokens in voorgaande sets.

3. Interpolatie van Token-orderingen

Om de trade-off tussen steekproefkwaliteit (strakheid van de grens) en efficiëntie tijdens inferentie (parallellisme) te controleren, introduceren de auteurs positieafhankelijke reveal-time schema's.

• Position-Offset Schema's: Elke tokenpositie $\ell$ krijgt een willekeurige reveal-tijd $R_\ell$ getrokken uit een distributie gedefinieerd door een cumulatieve distributiefunctie (CDF) $\alpha^\ell_\tau$.
• Sliding-Window Instantiatie (SW-SetDLM): De auteurs stellen een specifieke instantiatie voor waarbij tokens een actief generatie-interval $[a_\ell, a_\ell + w]$ hebben. De venstergrootte $w$ controleert de "links-naar-rechts" bias:
  • Kleine $w$ (naderend aan $1/L$) concentreert de volgorde op links-naar-rechts generatie, wat het AR-gedrag herstelt en de ELBO strakker maakt.
  • Grote $w$ (naderend aan 1) staat volgorde-onafhankelijke generatie toe, wat de standaard gemaskeerde diffusie herstelt.
  • Intermediaire waarden maken sliding-window decoding mogelijk, waardoor het model flexibele lengtes aan tokens in willekeurige posities kan voorspellen, wat taken zoals infilling ondersteunt.

4. Inferentie

Tijdens de inferentie worden kandidaat-posities gesampled op basis van de huidige ordeningstijd en het position-offset schema. Vervolgens worden deze posities parallel gedenoised met behulp van de KV-cache. Geaccepteerde tokens worden vastgelegd (committed) en de cache wordt onmiddellijk bijgewerkt, wat generatie met variabele lengte en efficiënte decodering mogelijk maakt.

Belangrijkste Bijdragen

1. Set Diffusion Framework: Een nieuwe klasse taalmodellen die generatie factoriseert over flexibele-positie, flexibele-lengte tokensets, waarmee autoregressie, volgorde-onafhankelijke diffusie en block diffusion worden gegeneraliseerd.
2. Set-Causal Architectuur: Een transformer-architectuur die KV-cache updates ondersteunt na elke inferentiestap, wat efficiënte generatie met variabele lengte mogelijk maakt.
3. Order Interpolation: Een fundamentele methode om te interpoleren tussen AR en diffusie door positionele ordeningsschema's (specifiek sliding-window schema's) af te stemmen, in plaats van vaste blokgroottes.
4. State-of-the-Art Prestaties: Demonstratie dat Set Diffusion een superieure snelheid-kwaliteit trade-off bereikt bij wiskundige redenering, samenvatting en onvoorwaardelijke generatie vergeleken met eerdere diffusiemodellen, en Block Diffusion aanzienlijk verslaat bij infilling-taken.

Experimentele Resultaten

De auteurs evalueerden Set Diffusion (specifiek SW-SetDLM) op verschillende benchmarks:

• Wiskundig Redeneren (GSM8K): Set diffusion bereikt een strikt betere snelheid-nauwkeurigheid Pareto-frontier dan Block Diffusion. Bijvoorbeeld, met een decoding venstergrootte $S \le 8$, bereikt het een 66,41% pass@1 nauwkeurigheid met een throughput van $\ge 60,42$ tokens/sec, waarmee het Block Diffusion ($S=8$) verslaat dat 63,53% nauwkeurigheid behaalt bij $\ge 55,39$ tokens/sec.
• Infilling (ROCStories): Set diffusion presteert aanzienlijk beter dan Block Diffusion bij infilling-taken (het invullen van 1 of 3 middelste zinnen), waarbij hogere ROUGE-scores worden behaald met tot wel ~25% snellere decodering. Block diffusion heeft hier moeite met vanwege de rigide links-naar-rechts blokstructuur.
• Samenvatting (CNN/DailyMail): Set diffusion behaalt concurrerende ROUGE-scores vergeleken met diffusie-baselines en een tot 10% snellere decodering dan Block Diffusion. Het presteert ook beter dan een AR-baseline, die volgens de auteurs de neiging heeft tot overfitting op deze dataset.
• Onvoorwaardelijke Generatie & Likelihood: Op OpenWebText en LM1B verbetert Set diffusion de MAUVE-scores ten opzichte van Block Diffusion terwijl het sneller decodeert. Op LM1B bereikt het de beste perplexiteit onder de diffusiemodellen (27.83 PPL voor $S \le 8$) terwijl het een hogere throughput behoudt dan eerdere diffusiemodellen.

Betekenis en Claims

De auteurs stellen dat Set Diffusion een significante vooruitgang vertegenwoordigt door de interpolatie tussen autoregressie en diffusie los te koppelen van rigide blokstructuren. Door te interpoleren via token-ordening, biedt het framework:

• Grotere Flexibiliteit: Ondersteuning voor invoegingen op willekeurige posities en generatie met variabele lengte, wat cruciaal is voor taken zoals infilling.
• Efficiëntie: Het vermogen om KV-caches na elke stap bij te werken, een functie die voorheen niet beschikbaar was in diffusiemodellen, wat leidt tot snellere inferentie.
• Prestaties: Een superieure trade-off tussen steekproefkwaliteit en inferentiesnelheid vergeleken met bestaande diffusie taalmodellen.

De auteurs positioneren Set Diffusion als een verenigend framework dat AR, volgorde-onafhankelijke diffusie en block diffusion als speciale gevallen herstelt, en een afstembare mechanisme biedt om de balans tussen parallellisme en steekproefkwaliteit te vinden op basis van de specifieke eisen van de inferentietaken.