DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Technische Samenvatting: DiLaDiff: Gedestilleerde Latent-Verrijkte Diffusie voor Taalmodellering

Probleemstelling
Diffusietalenmodellen (DLM's) staan voor een fundamentele afweging tussen steekproefkwaliteit en doorvoer. Hoewel discrete diffusiemodellen (specifiek Masked Diffusion Language Models, of MDLM's) effectief zijn voor categorische data, slagen ze er inherent niet in om correlaties tussen gedecodeerde tokens vast te leggen. Deze beperking vloeit voort uit het feit dat de denoiser de posterior benadert als een product van onafhankelijke token-voor-token marginaalverdelingen. Bijgevolg leidt het parallel genereren van meerdere tokens tot oncoherente zinnen, wat dwingt tot een afhankelijkheid van sequentiële generatie of een vermindering van parallelisme dat de doorvoer ernstig beperkt. Omgekeerd bieden continue diffusiemodellen toegepast op tekstreprentaties (bijv. embeddings) gevestigde steekproef- en destillatietechnieken, maar worstelen ze vaak met de afweging tussen expressiviteit en decodeerbaarheid, waarbij contextuele representaties moeilijk terug te decoderen zijn naar discrete tokens, of ze slagen er niet in om de prestaties van discrete tegenhangers te evenaren.

Methodologie
De auteurs stellen DiLaDiff (Distilled Latent-Augmented Diffusion Language Model) voor, een hybride continu-discreet raamwerk ontworpen om het probleem van tokencorrelatie op te lossen terwijl een hoge doorvoer behouden blijft. De architectuur bestaat uit drie primaire componenten:

1. Latent-Verrijkte Auto-Encoder (LaDiff):

   • Encoder: Een op transformers gebaseerde encoder ($E_\phi$) leert een continue latente ruimte ($z$) van contextuele representaties (afgeleid van BERT-verborgen staten). Deze latente ruimte is ontworpen om globale tokencorrelaties en semantische informatie vast te leggen.
   • Decoder: Een discrete decoder ($x_\theta$) wordt geïnitieerd vanuit een voorgetraind MDLM. Het genereert teksttokens geconditioneerd op de latente variabele $z$.
   • Training: Het systeem wordt getraind met een gewijzigde Evidence Lower Bound (ELBO) die auto-encoding combineert met gemaskerde diffusie. Cruciaal wordt de posteriorverdeling gefactoriseerd over tokens geconditioneerd op de latent $z$. Aangezien $z$ de correlaties vastlegt, geldt de aanname van conditionele onafhankelijkheid, wat nauwkeurige parallelle tokengeneratie mogelijk maakt.
   • Regularisatie: Om te voorkomen dat de latente ruimte te schaars wordt voor generatieve modellering, passen de auteurs robuustheidsbevorderende heuristieken toe, waaronder coördinaatgewijze normalisatie, willekeurige masking van features en latents, en injectie van Gaussisch ruis.

2. Leren van de Latente Prior:

   • Aangezien de schone sequentie tijdens inferentie niet beschikbaar is, wordt een continue diffusiemodel (Latent Denoiser $z_\psi$) getraind om de priorverdeling $p(z)$ over de output van de encoder te leren. Dit stelt het model in staat een hoogwaardige latent $z$ te steekproeven via een omgekeerde probability-flow ODE voordat er wordt gedecodeerd.

3. Consistentiedestillatie (DiLaDiff):

   • Om de rekenkundige overhead van de continue diffusiestap verder te verminderen, passen de auteurs MeanFlow-destillatie toe. Een studentmodel ($u_\eta$) wordt getraind om de gemiddelde snelheid van de ODE-trajectorie tussen twee tijdstappen ($t$ en $r$) te voorspellen, in plaats van de momentane snelheid.
   • Deze destillatie stelt de latent $z$ in staat om in enkele stappen (bijv. 5 stappen) gegenereerd te worden, terwijl de prestaties dicht bij het docentmodel (LaDiff) blijven, dat veel stappen vereist (bijv. 200 stappen).

Belangrijkste Bijdragen

• Hybride Architectuur: Het artikel introduceert een recept voor het trainen van een tekstauto-encoder waarbij de decoder wordt geïnitieerd vanuit een voorgetrainde discrete diffusie-baseline. Dit resulteert in een controleerbare, geregulariseerde latente ruimte die semantische informatie en tokencorrelaties vastlegt.
• Latent-Geleide Parallelisme: Door een latente prior te leren met continue diffusie, stelt het model (LaDiff) het steekproeven van meerdere tokens parallel in staat zonder tekstkwaliteit te offeren, aangezien de latent de tokenafhankelijkheden oplost die discrete modellen missen.
• Eerste Succesvolle ODE-Destillatie voor Tekstlatents: De auteurs presenteren DiLaDiff, de eerste succesvolle poging om ODE-trajectoires voor contextuele tekstlatents te destilleren. Dit verlaagt de kosten voor latentgeneratie tot verwaarloosbare niveaus in vergelijking met discrete decoding.

Resultaten
Experimenten werden uitgevoerd op het OpenWebText-dataset met een sequentielengte van 1024 tokens.

• Prestaties versus Baseline: Zelfs zonder destillatie presteert LaDiff beter dan de gemaskerde diffusie-baseline (MDLM). Bij onvoorwaardelijke generatie bereikt LaDiff een absolute verbetering van 28,5 in Generative Perplexity (GenPPL) en een absolute verbetering van 0,08 in MAUVE ten opzichte van MDLM met 1024 discrete stappen.
• Doorvoer: LaDiff versnelt inferentie met een factor van 7× bij een batchgrootte van 32 in vergelijking met de optimale MDLM-instelling.
• Destillatie-efficiëntie: DiLaDiff bereikt prestaties die dicht bij zijn LaDiff-docent liggen met slechts 5 latent-diffusiestappen (tegenover 200 voor de docent). De rekenkundige overhead van de latent-diffusiecomponent daalt tot ongeveer 5% van de totale wandtijd, waardoor de latentgeneratie verwaarloosbaar wordt in vergelijking met de kosten voor discrete decoding.
• Generatie in Weinige Stappen: DiLaDiff concurreert met state-of-the-art methoden voor weinig stappen (zoals MDLM+SDTT en DUO+DCD) zonder de discrete decoder expliciet te destilleren, wat de effectiviteit van latent-verrijking aantoont.
• Semantische Diversiteit: Analyse toont aan dat de latente ruimte semantische informatie behoudt (geverifieerd via BERTScore op ruisige latents) en diverse generatie toelaat, zoals blijkt uit het vermogen van de decoder om entropie te behouden, zelfs bij lagere temperaturen.

Betekenis en Claims
Het artikel beweert dat DiLaDiff succesvol de kloof overbrugt tussen de expressiviteit van continue representaties en de robuustheid van discrete decoding. Door globale correlaties in een latente ruimte vast te leggen, overwint het model de inherente beperking van discrete diffusiemodellen met betrekking tot parallelle tokengeneratie. De auteurs benadrukken dat hun aanpak hoogwaardige tekstgeneratie mogelijk maakt met aanzienlijk verminderde inferentielatentie.

Het artikel blijft bescheiden over zijn beperkingen:

• DiLaDiff komt de prestaties van zijn docentmodel (LaDiff) in het regime van weinig stappen nog niet volledig overeen, wat suggereert dat verder werk aan destillatie nodig is.
• De discrete decoder zelf wordt niet gedestilleerd; de auteurs merken op dat het combineren van hun latentdestillatie met destillatie van de discrete decoder een natuurlijke toekomstige uitbreiding is om echte state-of-the-art prestaties in weinig stappen te bereiken.
• De regularisatiestrategieën en diffusie-schema-ontwerpen zijn momenteel heuristisch, hoewel de auteurs in de appendices een meer gevestigde poging presenteren.

Kortom, DiLaDiff toont aan dat een hybride continu-discrete aanpak, verrijkt met latentdestillatie, de inferentie van taalmodellering aanzienlijk kan versnellen terwijl de generatiekwaliteit behouden of verbeterd blijft, waardoor de in huidige diffusietalenmodellen inherente afweging tussen doorvoer en kwaliteit effectief wordt opgelost.