Consistent Diffusion Language Models

1. Probleemstelling

Diffusion Language Models (DLM's) bieden een veelbelovend alternatief voor autoregressieve (AR) modellen door parallelle token-generatie en inferentie in sublineaire tijd mogelijk te maken. Er blijft echter een kritieke bottleneck bestaan: het genereren van hoogwaardige tekst vereist doorgaans honderden verfijningstappen, wat de efficiëntiewinsten die door parallelisme worden beloofd, tenietdoet.

In continue domeinen (bijvoorbeeld afbeeldingen) hebben Consistency Models de diffusie succesvol versneld door modellen te trainen om invariant te zijn langs een unieke, deterministische trajectorie die wordt gedefinieerd door de Probability Flow Ordinary Differential Equation (PF-ODE).

• De Kloof: In discrete domeinen (zoals tekst) bestaat er geen analoog PF-ODE in de steekproefruimte. Discrete corruptieprocessen (bijvoorbeeld masking of uniforme ruis) volgen geen enkel deterministisch pad; in plaats daarvan omvatten ze stochastische overgangen.
• Het Gevolg: Het direct aanpassen van continue consistentiemethoden aan discrete data is slecht gedefinieerd. Bestaande oplossingen vertrouwen vaak op meerstaps distillatiepijplijnen (met een teacher-model) of surrogaten met continue relaxatie, die complex, rekenkundig duur zijn en vaak lijden onder mode-collapse of verlies aan diversiteit.

2. Methodologie: Multi-Path Discrete Consistency (MPDC)

De auteurs stellen een fundamentele verschuiving voor: in plaats van te zoeken naar een niet-bestaand deterministisch pad, maken ze gebruik van de exacte posterior-bruggen die inherent zijn aan discrete diffusieprocessen.

Kernconcept: De Posterior-Brug

Voor twee ruissniveaus $s < t$ in een discrete diffusieprocess bestaat er een gesloten-vorm exacte posterior-brug $q(x_s | x_t, x_0)$. Deze brug definieert een geldig stochastisch pad dat een ruizige toestand $x_t$ verbindt met een schoner toestand $x_s$, geconditioneerd op de ware data $x_0$.

• In tegenstelling tot de continue PF-ODE is deze brug stochastisch.
• Cruciaal vormen deze bruggen een semigroep (Chapman-Kolmogorov-eigenschap), wat betekent dat paden transitief samengesteld kunnen worden.

Het MPDC-Principe

De auteurs introduceren Multi-Path Discrete Consistency (MPDC), waarbij een denoiser wordt getraind om pad-invariant in verwachting te zijn.

• Definitie: Een model $f_\theta(x_t, t)$ is consistent als zijn voorspelling bij een ruizige toestand $x_t$ overeenkomt (in verwachting) met de voorspelling die wordt gedaan na het nemen van een tussenliggende "sprong" naar een schoner toestand $x_s$ via de exacte posterior-brug.
• Doel: Het model minimaliseert de divergentie tussen zijn directe voorspelling vanuit $x_t$ en de verwachte voorspelling vanuit $x_s$ (waarbij $x_s$ wordt bemonsterd uit de brug).
• Wiskundige Formulering:
  $$\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]$$
  Waarbij $x_s \sim q(x_s | x_t, x_0)$, $\tilde{f}_\theta$ een doelnetwerk is (EMA), en $D$ een divergentiemaat is (typisch Jensen-Shannon Divergentie).

Belangrijke Ontwerpkies voor Stabiliteit

Om mode-collapse (waarbij het model repetitieve tekst produceert om consistentie triviaal te bevredigen) en uniforme drift te voorkomen, introduceren de auteurs drie stabiliserende mechanismen:

1. Stapgrootte als Curriculum: Het bemonsteren van stapgroottes $\delta$ uit een specifiek bereik (bijvoorbeeld $1/8$ tot $3/8$) en het wegen ervan met $1/\delta$ om leer signalen over korte en lange paden in evenwicht te brengen.
2. Max-Step Diffusion Anchor: Een klein deel van de trainingsverlies wordt verankerd aan de ware data ($s=0$) met behulp van een standaard diffusiedoel. Dit voorkomt dat het model convergeert naar degenererende constante oplossingen.
3. Optimalisatie Asymmetrie & Divergentie: Het gebruik van een stop-gradient op het doelnetwerk en het toepassen van Jensen-Shannon Divergentie (JSD) in plaats van Forward/Backward KL. JSD biedt een stabiele afweging tussen mode-zoeken en mode-dekken, en vermijdt de instabiliteit van KL-divergentie bij training vanaf nul.

3. Belangrijkste Bijdragen

1. Een Nieuw Principe voor Discrete Generatieve Modellering: Het paper stelt vast dat exacte posterior-bruggen dienen als een rigoureuze, gesloten-vorm vervanging voor de ontbrekende PF-ODE, waardoor het obstakel "discrete versnelling" wordt omgezet in een analytisch voordeel.
2. Een Gecentraliseerd Trainingskader: CDLM biedt een eenstaps, teacher-vrij doel. De auteurs bewijzen dat standaard masked diffusion, continue consistency-modellen, progressieve distillatie en recente tweestaps distillatiemethoden (zoals SDTT en DUO-DCD) allemaal speciale gevallen, limieten of benaderende koppelingen zijn van het CDLM-kader.
3. State-of-the-Art Prestaties: CDLM behaalt superieure resultaten in zowel onvoorwaardelijke als voorwaardelijke tekstgeneratie, en presteert beter dan sterke basismodellen en komt overeen met of overtreft complexe meerstaps gedistilleerde modellen zonder dat een teacher nodig is.

4. Experimentele Resultaten

De auteurs evalueerden CDLM (specifiek MCDLM voor masked diffusion en UCDLM voor uniforme diffusion) tegenover baselines zoals MDLM, DUO, SDTT en DUO-DCD op het OpenWebText-dataset.

• Onvoorwaardelijke Generatie:

  • Efficiëntie met Weinig Stappen: CDLM vestigt nieuwe state-of-the-art (SOTA) resultaten voor eenstapsmodellen. Het presteert aanzienlijk beter dan basismodellen (MDLM, DUO) over alle bemonsteringsbudgetten heen.
  • Vergelijking met Distillatie: Opmerkelijk is dat het eenstaps CDLM vaak beter presteert dan meerstaps gedistilleerde modellen (SDTT, DUO-DCD) wat betreft Perplexity (PPL), terwijl het veel hogere entropie (diversiteit) behoudt.
  • Mode Collapse: Gedistilleerde modellen die greedy samplers gebruiken (bijvoorbeeld DUO-DCD) vertoonden lage entropie (3,9), wat wijst op ernstige mode-collapse (repetitieve karakters). CDLM behield een gezonde entropie (5,0+), wat bewijst dat het in staat is diversiteit te behouden.
  • Snelheid: CDLM bereikt 32x–64x snelheidswinst ten opzichte van autoregressieve baselines en 64x–128x ten opzichte van standaard MDLM, terwijl het vergelijkbare kwaliteit behoudt.

• Voorwaardelijke Generatie:

  • CDLM toonde robuuste prestaties bij het herstellen van originele zinnen uit 50% gecorrumpeerde invoer over meerdere datasets (OpenWebText, Lambada, Wikitext-103, PTB).
  • Het behaalde betere BLEU-scores en lagere perplexity dan SDTT en DUO, met name bij het behouden van invoercondities waar andere modellen faalden.

5. Betekenis en Impact

• Theoretische Unificatie: Het paper biedt een "unificerend perspectief" voor discrete generatieve modellering, waarbij wordt aangetoond dat diverse versnellingstechnieken fundamenteel met elkaar verbonden zijn door het concept van padconsistentie.
• Schaalbaarheid: Door de noodzaak van meerstaps distillatie en teacher-modellen te elimineren, biedt CDLM een principieel, schaalbaar en teacher-vrij fundament voor next-generation snelle discrete generatieve modellen.
• Generaliseerbaarheid: Hoewel het wordt gedemonstreerd op tekst, is het kader toepasbaar op elk discreet domein met een hanteerbare posterior-brug (bijvoorbeeld biologische sequenties, grafen, programmasynthese).
• Paradigmaverschuiving: Het herformuleert snelle discrete diffusie niet als een benadering van continue ODE's, maar als het leren van een pad-onafhankelijke denoiser over stochastische bruggen, waarmee het probleem van generatie in weinig stappen in discrete ruimten succesvol wordt opgelost.

Concluderend overbrugt CDLM de kloof tussen de efficiëntie van continue consistency-modellen en de realiteit van discrete data, en biedt het een robuuste, eenstaps oplossing die hoogwaardige, snelle tekstgeneratie bereikt zonder de complexiteit van distillatiepijplijnen.