Hoe bouw je een Diffusion Language Model?

Introductie: Autoregressieve en Diffusion Language Models

Tegenwoordig worden twee families van generatieve AI-algoritmen veel gebruikt. Voor continue data, zoals afbeeldingen of video, is de state-of-the-art aanpak gebaseerd op diffusion-modellen. Voor discrete data, zoals tekst of code, is de standaardaanpak daarentegen gebaseerd op autoregressieve modellen. Dit artikel verkent een alternatief voor discrete data, gebouwd op het moderne paradigma van diffusie.

Mainstream taalmodellen zijn autoregressief: ze genereren tokens van links naar rechts, één voor één, waarbij elk token geconditioneerd is op de tokens die eraan voorafgingen. Deze aanpak is krachtig, maar kent inherente beperkingen:

• Geen foutcorrectie: zodra een token is uitgezonden, kan het niet meer worden herzien, waardoor vroege fouten zich opstapelen.
• Trage generatie: het produceren van een sequentie vereist evenveel stappen als er tokens zijn, en leent zich niet natuurlijk voor snelle, parallelle generatie.
• Causale aandacht: bij generatie wordt alleen naar het verleden gekeken, nooit naar de toekomstige context.

Diffusion-modellen hanteren een andere aanpak. In plaats van tekst token voor token te produceren, genereren ze de hele sequentie in één keer, beginnend bij een eerste schatting die over een aantal stappen iteratief wordt verfijnd. Dit biedt verschillende voordelen: de generatiesnelheid en kwaliteit kunnen worden afgewogen door meer of minder stappen te gebruiken, fouten kunnen onderweg worden gecorrigeerd, en elke stap maakt gebruik van bidirectionele context.

Autoregressieve LLM's genereren één token per keer, van links naar rechts, en nemen evenveel stappen als er tokens zijn. Diffusion LLM's — zoals Gemma Diffusion — beginnen daarentegen met een ruwe conceptversie van volledige lengte en verfijnen elke positie parallel over een paar rondes, waarbij de hele sequentie bij elke stap wordt herschreven in plaats van een enkel token uit te zenden.

Het toepassen van diffusie op taal was lange tijd een open probleem. In 2024 bereikte het veld een keerpunt, toen diffusion-modellen qua kwaliteit concurrerend werden met autoregressieve modellen. Tegen 2026 zijn diffusion LLM's een realiteit, met releases van toonaangevende industriële labs — Mercury 2 (Inception Labs), Gemma Diffusion (Google) en Nemotron Diffusion (NVIDIA). Dit artikel volgt de ideeën en publicaties die ten grondslag liggen aan deze moderne modellen.

Achtergrond: Gaussische Diffusie

Voordat we diffusie voor taal introduceren, beginnen we met een kort overzicht van Gaussische diffusie voor beeldgeneratie. We zullen vervolgens discrete diffusie door analogie opbouwen.

Genereren door iteratieve ruisonderdrukking (denoising)

Het centrale concept achter diffusion-modellen is denoising. In plaats van een afbeelding in één keer te schilderen, produceert een diffusion-model afbeeldingen stap voor stap, beginnend bij pure willekeurige ruis en deze bij elke stap een beetje verwijderen totdat er een coherente afbeelding ontstaat. Het genereren van een afbeelding via vele kleine stappen blijkt veel eenvoudiger dan het in één keer produceren, en dit is wat diffusion-modellen zo effectief maakt.

Hoe leert een model om ruis te verwijderen? De truc is om het te leren door voorbeelden te tonen van ruis die geleidelijk wordt getransformeerd in een afbeelding. Diffusie bereikt dit via twee complementaire processen. Ten eerste transformeert een forward process een schone bronafbeelding stap voor stap in pure ruis. Ten tweede leert een reverse process deze transformatie om te keren, waardoor pure ruis weer in een afbeelding wordt veranderd; het wordt getraind op de trajecten van afbeelding naar ruis die door het forward process zijn geproduceerd.

Forward process

Het forward process neemt een schone trainingsafbeelding en produceert een reeks steeds ruisiger afbeeldingen die een pad volgen van schone data naar pure ruis. Dit gebeurt door bij elke stap een groeiende hoeveelheid willekeurige Gaussische ruis toe te voegen, totdat de afbeelding oplost in pure statische ruis. Deze stap vereist geen enkel leerproces — we voegen simpelweg ruis toe — maar is enorm nuttig, omdat het een eindeloze voorraad trainingsdata genereert: voorbeelden van afbeeldingen die in ruis worden getransformeerd, en vice versa.

Reverse process

Het reverse process is waar het eigenlijke leren plaatsvindt. We trainen een model om ruis in afbeeldingen om te zetten door de stappen van het forward process in omgekeerde volgorde te volgen.

Concreet trainen we, gegeven een ruisige afbeelding, een machine learning-model om de ruis van de onderliggende afbeelding te scheiden of, equivalent daaraan, om ofwel de toegevoegde ruis of de schone afbeelding zelf te voorspellen. Aangezien de ruisige invoer bekend is, bepaalt het weten van de één de ander. Zodra het model dit kan, is generatie eenvoudig: begin bij pure ruis, vraag het model om een klein deel ervan te schatten en te verwijderen, en herhaal dit. Elke passage brengt het monster iets dichter bij iets dat lijkt op echte data, totdat er een schone afbeelding overblijft.

Dit recept van forward/reverse — data corrupt maken met ruis, en vervolgens leren die corruptie stap voor stap ongedaan te maken — is het blauwdruk voor elk diffusion-model.

Eenvoudige Masked Diffusion Modellen

Het belangrijkste obstakel bij het brengen van diffusie naar taal is het bepalen van wat "ruis" moet betekenen voor discrete tokens. De ruis in klassieke diffusie is bijvoorbeeld Gaussisch, en het toevoegen van continue Gaussische ruis aan categorische variabelen is niet goed gedefinieerd. Hieronder introduceren we een eenvoudige maar effectieve aanpak die ruis definieert via masking.

Masked Diffusion in een notendop

De eenvoudigste manier om masked diffusion te begrijpen is als een unmasking transformer. We trainen het model door schone sequenties te nemen, een willekeurig deel van hun tokens te maskeren en een bidirectionele transformer te vragen de gaten in te vullen. Als je BERT kent, is dit in essentie BERT met een gerandomiseerde masking-rate — maar in tegenstelling tot BERT is het resulterende model generatief. Je kunt masked diffusion zien als een generatieve BERT.

Zodra we de unmasking transformer hebben getraind, kunnen we tekst genereren door te beginnen met een volledig gemaskeerde sequentie en twee stappen herhaaldelijk uit te voeren:

1. Infilling: vraag het model om elk gat in de huidige sequentie in te vullen, wat resulteert in een ruwe schatting van de schone tokens.
2. Remasking: voeg opnieuw willekeurig ruis toe aan de ingevulde sequentie door tokens te vervangen door masks, maar houd een paar meer tokens ongemaskeerd dan in de vorige ronde.

Elke ronde laat minder posities gemaskeerd achter, totdat de sequentie convergeert naar een schoon monster van het model. Generatie komt dus neer op het starten met een sequentie vol gaten en het geleidelijk invullen van woorden in een willekeurige volgorde.

Masked Diffusion begrijpen: Een probabilistisch perspectief

We kunnen beter begrijpen waarom dit proces werkt door het te kaderen als een analoog van het Gaussische diffusiemodel. Net als Gaussische diffusie kan masked diffusion worden beschreven als een model bestaande uit een forward- en een reverse-proces.

Forward process

Het doel van het forward process is het genereren van trainingsdata voor het reverse process. De output is een traject dat begint bij een datapunt en eindigt bij een sequentie van pure ruis.

De belangrijkste uitdaging is het bepalen van wat "ruisig" betekent. In masked diffusion maskeren we in plaats van witte ruis willekeurig een fractie van de tokens in een discrete sequentie. De hoeveelheid masking wordt beheerst door een schema $\alpha_t$ — de waarschijnlijkheid dat een gegeven token ongemaskeerd blijft — wat de rol van de signaal-ruisverhouding speelt. Het begint bij $1$ wanneer $t = 0$ (een schone sequentie) en neemt af naar $0$ wanneer $t = 1$ (een volledig gemaskeerde sequentie). De tijdvariabele $t$ indexeert een pad van schone naar ruisige data.

We implementeren dit proces als een Markov-keten over een reeks variabelen $zt$ geïndexeerd door $t$, waarbij $z0$ de schone, ongemaskeerde sequentie is. Voor $s < t$ definieert de keten $q(zt \mid zs)$ door elk nog ongemaskeerd token van $zs$ met waarschijnlijkheid $(\alphas - \alphat)/\alphas$ te maskeren.

Reverse process

Vervolgens trainen we het reverse process om de reeks steeds sterker gemaskeerde latents in omgekeerde volgorde te doorlopen — beginnend bij een volledig gemaskeerde sequentie en uiteindelijk outputs te genereren die lijken op schone data.

Met de regel van Bayes kunnen we het mathematisch optimale reverse process $q(zs \mid zt, x)$ afleiden wanneer de schone sequentie $x$ bekend is. Dit optimale proces heeft twee stappen: (1) gegeven een gedeeltelijk gemaskeerde $zt$, kijken we naar $x$ om de werkelijke schone tokens te vinden; (2) vorm $zs$ door elke gemaskeerde positie van $zt$ met waarschijnlijkheid $(\alphas - \alphat) / (1 - \alphat)$ te vervangen door de waarde in $x$, en deze anders gemaskeerd te laten.

In de praktijk is de uiteindelijke output $x$ onbekend tijdens generatie. We trainen daarom een model $x\theta(zt)$ om de uiteindelijke schone sequentie te voorspellen gegeven de huidige staat $zt$ en passen het ideale reverse process $q(zs \mid zt, x)$ toe met de schatting $x\theta(zt)$ in plaats van de echte $x$. Formeler definiëren we het reverse process als een waarschijnlijkheid $p(zs \mid zt) = q\big(zs \mid zt, x\theta(z_t)\big)$.

Een probabilistisch model voor masked diffusion

Samen vormt dit de mathematische definitie van een Masked Diffusion Language Model (MDLM). Het forward process $q(zt \mid zs)$ produceert een traject van schone naar volledig gemaskeerde data, en het reverse process $p(zs \mid zt)$ leert dit ongedaan te maken.

We kunnen de likelihood $\log p(x)$ van het model gebruiken om de kwaliteit te beoordelen. In latent-variabele modellen is dit onbehandelbaar, dus maken we gebruik van benaderingen via variationele inferentie. Voor een MDLM heeft de evidence lower bound (ELBO) een verrassend eenvoudige vorm (ervan uitgaande dat $\alpha_t = 1-t$):

\[\mathcal{L}\text{ELBO} = \mathbb{E}{t \sim \mathcal{U}[0,1]}\; \mathbb{E}{q(zt \mid x)} \left[ \frac{1}{t}\, \log p\theta\!\left(x \mid zt\right) \right].\]

De term $\log p\theta(x \mid zt)$ is de likelihood van een schone sequentie $x$ gegeven een gedeeltelijk gemaskeerde sequentie $z_t$. In andere woorden, dit is de cross-entropy loss tussen de voorspellingen van onze unmasking transformer en de werkelijke tokens — dit is exact de BERT-loss!

In tegenstelling tot BERT wordt deze loss gemiddeld over alle $t$ (alle mogelijke masking-rates) en genormaliseerd door $t$, de verwachte fractie gemaskeerde tokens.

Samenvatting en Evaluatie van Masked Diffusion

Samenvattend is MDLM zeer vergelijkbaar met BERT, met twee belangrijke verschillen:

1. Het staat principiële sampling-algoritmen toe, overeenkomend met ancestral sampling in een latent-variabel model.
2. Training maakt gebruik van een gerandomiseerde masking-rate, wat correspondeert met een principiële variationele ondergrens van de log-likelihood.

De ELBO maakt een principiële vergelijking mogelijk tussen autoregressieve en diffusion-taalmodellen met behulp van log-likelihood (of perplexiteit). Hoewel er lange tijd een aanzienlijk gat in perplexiteit bestond, waren vereenvoudigde masked diffusion-modellen een van de eersten die dit gat grotendeels wisten te dichten.

Een Real-World Diffusion Language Model bouwen

Hoewel masked diffusion-modellen nuttig zijn voor intuïtie, zijn ze niet direct productierijp: ze genereren alleen sequenties van vaste lengte, ondersteunen geen iteratieve verfijning (foutcorrectie) out-of-the-box, en zijn niet bijzonder snel zonder extra post-training.

Block Diffusion voor generatie van variabele lengte

Standaard MDLMs zijn beperkt tot vaste lengtes. Block diffusion lost dit op door diffusie uit te voeren over blokken, geconditioneerd op eerder gegenereerde tokens. Deze blokken kunnen een willekeurige grootte hebben, van een dozijn tot duizenden tokens, afhankelijk van het toepassingsgebied.

In biologische toepassingen kunnen we de blokgrootte instellen op de minimale lengte die nodig is om specifieke interacties vast te leggen. Bij algemene taalmodellering kiezen we de blokgrootte vaak om de GPU-benutting te maximaliseren. Daarnaast ondersteunt block diffusion natuurlijk KV-caching, een techniek die generatie versnelt door keys en values van eerdere blokken te hergebruiken.

Andere benaderingen voor variabele lengte zijn onder meer Set Diffusion (die opereert over willekeurige sets van posities) en methoden zoals Edit Flows of FlexMDM, die generatie modelleren als een reeks insertie-, deletie- en substitie-operaties.

Architecturen: Encoder, Decoder en Encoder-Decoder

Standaard masked diffusion-modellen zijn effectief encoder-only (zoals BERT), in tegenstelling tot de decoder-only autoregressieve modellen (zoals GPT). Een encoder-only architectuur vereist sampling-algoritmen die bij elke denoising-stap het volledige netwerk aanroepen, wat rekentechnisch duur is.

Een belangrijk inzicht is dat diffusie twee soorten berekeningen uitvoert: (1) het berekenen van een representatie van de tokens die tot nu toe zijn gegenereerd, en (2) het ontdoen van de ruis in de corrupted tokens. Dit suggereert het gebruik van een encoder-decoder architectuur. Hierbij representeert een encoder de schone tokens en verfijnt een lichtgewicht decoder iteratief de ruisige sequentie. Deze architectuur vormt de kern van state-of-the-art open-source diffusion LLM's, zoals Gemma Diffusion en Nemotron Diffusion.

Iteratieve Verfijning en Ingebouwde Foutcorrectie

Standaard MDLMs missen iteratieve verfijning: zodra een token is ongemaskeerd, kan het nooit meer worden bijgewerkt. Moderne diffusion LM's passen het forward- en reverse-proces aan om dit te herstellen.

Remasking Diffusion

De eenvoudigste oplossing is remasking: bij elke stap houden we sommige nieuw ongemaskeerde tokens, maar we maskeren ook een klein deel van de eerder ongemaskeerde tokens opnieuw, zodat ze opnieuw kunnen worden gegenereerd. Dit stelt het model in staat om grammaticale fouten die vroeg in het proces zijn gemaakt, later te corrigeren wanneer er meer context beschikbaar is.

Remasking kan als plug-in sampler worden toegevoegd aan bestaande pretrained MDLMs. Het introduceert een vorm van inference-time compute scaling: door het aantal sampling-stappen te verhogen, nadert de kwaliteit die van autoregressieve generatie.

Uniform State Diffusion

Een alternatief is uniform state diffusion. In plaats van masking vervangt het forward process tokens door willekeurige tokens uit de vocabulaire. Het reverse process begint met een willekeurige sequentie en wijzigt tokens totdat het resultaat op data lijkt.

Omdat er geen masks zijn, kan het model bij elke stap beslissen om elk token te vervangen. Dit ondersteunt foutcorrectie op een natuurlijke manier. Hoewel masked diffusion-modellen meestal sneller trainen, faciliteren Uniform Diffusion Language Models (UDLMs) snellere sampling en betere controleerbaarheid. Gemma Diffusion is bijvoorbeeld een UDLM.

Versnellen van Diffusion Sampling via Distillatie

Omdat diffusie meerdere tokens per stap kan genereren of verfijnen, kan het 5–10× sneller zijn dan autoregressieve generatie. Echter, het tegelijkertijd genereren van veel tokens kan inconsistenties introduceren. Daarom vertrouwen snelle modellen op foutcorrigerende processen zoals remasking of UDLM.

Versnelling wordt vaak bereikt via progressive distillation: een model wordt getraind op zijn eigen generaties om een stap over te slaan, wat recursief wordt herhaald om het aantal stappen te halveren. In diffusion LM's vertaalt dit zich naar technieken als self-distillation through time en discrete consistency distillation.

Dit kan worden gezien als een vorm van on-policy training. Standaard training is off-policy (getraind op samples van het forward process). Door in een post-processing stap op de eigen samples van het model te trainen, kan de sampling-snelheid aanzienlijk worden verhoogd.

Diffusie maakt Controleerbare Generatie mogelijk

Diffusion-modellen blinken uit in controleerbare generatie: het produceren van een sample $x$ die voldoet aan een doeleigenschap $y$. Omdat ze globaal verfijnen in plaats van onomkeerbare lokale wijzigingen aan te brengen, navigeren ze effectiever door de sample-ruimte.

In de praktijk is er een Pareto-afweging tussen natuurlijkheid (lijkt het op echte data?) en eigenschap-voldoening (heeft het de gewenste eigenschap?). Diffusie presteert hier beter dan autoregressie.

Algoritmen voor controleerbare generatie

Technieken zoals classifier-based guidance (CBG) en classifier-free guidance (CFG) worden hiervoor gebruikt. Bij CBG stuurt een predictor-model $p(y \mid x)$ de generatie bij elke stap.

Discrete classifier-based guidance (D-CBG): de wiskunde

Via de regel van Bayes kan elk conditioneel reverse process worden ontbonden: $$ \log p(zs \mid zt, y) = \log p(y \mid zt, zs) + \log p(zs \mid zt) + c $$ Met een classifier $p(y \mid zt)$ en een onconditioneel diffusiemodel kunnen we een conditioneel model construeren: $$ \log p^{(\gamma)}(zs \mid zt, y) = \gamma \log p(y \mid zs, zt) + \log p(zs \mid z_t) $$ Hierbij is $\gamma$ de sturingssterkte die de balans bepaalt tussen de eigenschap en de voorkeuren van het model.

Discrete classifier-free guidance (D-CFG): de wiskunde

Bij CFG wordt geen aparte classifier gebruikt, maar een combinatie van een conditioneel model $p(zs \mid zt, y)$ en een onconditioneel model $p(zs \mid zt)$: $$ \log p^{(\gamma)}(zs \mid zt, y) = \gamma \log p(zs \mid zt, y) + (1 - \gamma) \log p(zs \mid zt) $$ Dit is praktisch omdat het een aparte classifier vermijdt; beide distributies worden geparametriseerd door hetzelfde model, dat is getraind door het conditioneringssignaal $y$ willekeurig weg te laten.

Post-Training van Diffusion Language Models

Diffusion-taalmodellen kunnen via reinforcement learning (RL) worden gepost-trained om redeneervermogen te verbeteren. De complicatie is dat RL-algoritmen de likelihood van een traject nodig hebben, wat duur is voor masked diffusion-modellen.

Hier introduceert diffu-GRPO een critic-free policy-gradient algoritme dat traject-log-probabilities schat met een mean-field benadering. Andere verbeteringen richten zich op nauwkeuriger likelihood-estimators, wat leidt tot nieuwe state-of-the-art resultaten op logische en mathematische benchmarks zonder afhankelijk te zijn van supervised fine-tuning.

RL is ook cruciaal in biologische toepassingen, waar de beloning vaak een experimenteel gemeten eigenschap is (bijv. bindingsaffiniteit). Methoden zoals DRAKES propageren deze beloningen terug door het sampling-traject om het model direct te finetunen voor DNA- en eiwitontwerp.

Diffusion Large Language Models van Vandaag

De afgelopen jaren zijn diffusion-taalmodellen geschaald naar miljarden parameters.

Biologische en Wetenschappelijke Domeinen

Diffusie is hier zeer succesvol omdat biologische sequenties minder last hebben van de links-naar-rechts bias van autoregressieve modellen en omdat controleerbare generatie essentieel is.

• Eiwitsequentie-modellering: Het ESM3-model implementeert MDLM met tot 100B parameters. Het model kan genereren en conditioneren over drie tracks: sequentie, structuur en functie, en presteert beter dan autoregressieve baselines.
• Nucleotide-sequentie-modellering: De Nucleotide Transformer v3 (NT-v3) familie schaalt MDLM naar miljarden parameters over biljoenen DNA-tokens. Door gebruik te maken van discrete CFG met remasking kan NT-v3 DNA-sequenties genereren die genexpressie effectiever moduleren dan eerdere baselines.

Diffusion-Based Large Language Models

De laatste 18 maanden is er een snelle opkomst van algemene diffusion LLM's:

• LLaDA: Schaalt MDLM naar 8B parameters, compatibel met remasking en post-training. Het is concurrerend met modellen als LLaMA2-7B en LLaMA3-8B.
• Mercury: De eerste commerciële diffusion LLM (2025), gericht op snelheid. Het haalt meer dan 1.000 tokens/sec/user op standaard GPU's, wat 5–10× sneller is dan vergelijkbare frontier-modellen.
• Gemma Diffusion: Een open-weights model van Google dat een UDLM-backbone combineert met block diffusion, een encoder-decoder architectuur en versnelde sampling.
• Nemotron Diffusion: Een familie van NVIDIA-modellen getraind met een joint autoregressive-diffusion objective. Met tot 35B parameters bieden ze 2-8× de throughput van AR-modellen met behoud van bijna volledige kwaliteit.

Conclusie en Slotoverwegingen

Het veld van diffusion-taalmodellen is geëxplodeerd. Diffusie biedt voordelen in snelheid (via parallelle generatie), controleerbaarheid, multimodaliteit en inference-time scaling. Hoewel ze nog niet tot dezelfde schaal (data/compute) zijn gebracht als autoregressieve modellen, tonen experimenten tot 100B parameters veel potentieel.

Is Diffusie een Pad naar Intelligentere Modellen? Een Perspectief op Scaling Laws

Kunnen diffusion-modellen fundamentele verbeteringen in pure intelligentie opleveren t.o.v. autoregressieve modellen?

Vanuit het perspectief van scaling laws was de belangrijkste vooruitgang tussen 2019 en 2024 het schalen van pre-training. Dit werd mogelijk door de Transformer, die een volledig parallel trainingsalgoritme introduceerde, in tegenstelling tot de sequentiële RNN's.

Sinds 2024 vlakt de winst uit pre-training af; winst komt nu vooral uit post-training en inference-time compute. Echter, zowel post-training als inferentie worden nu beperkt door de sequentiële aard van generatie. Als inferentie volledig parallel gemaakt kan worden, zou dit de intelligentie-winst kunnen versnellen op een manier die vergelijkbaar is met de transitie van RNN naar Transformer tijdens de pre-training fase.

Diffusie is de aanpak die inferentie volledig parallel kan maken. In essentie zou diffusie voor inferentie- en post-training scaling laws kunnen zijn wat de Transformer was voor pre-training scaling laws. Door meer FLOPs per seconde te kunnen gebruiken, ontsluit diffusie een betere hardware-benutting, wat de weg opent naar verdere schaling van intelligentie.