Diffusion language models

Beyond left to right. Instead of writing one token after another, a masked diffusion model starts from a page of blanks and fills them in, several at a time, in any order. The probability behind it, and what parallel unmasking costs.

Key idea
Masked (discrete) diffusion: a forward masking process, a learned reverse process, the ELBO training loss, and the cost of unmasking in parallel
Real world
A crossword is solved in no fixed order. You fill in the clues you're surest of first, and each answer constrains the ones that cross it. Filling in two crossing answers at once, without checking each against the other, is how you end up with a clash.
In an LLM
LLaDA, Dream and Gemini Diffusion start from a fully masked answer and reveal tokens over a fixed number of steps, many per step. Fewer steps means faster generation, but tokens revealed in the same step are chosen independently, which can produce text that doesn't fit together.

Two ways to factor a sentence

An ordinary language model uses the chain rule left to right:

p(x1,…,xn)=p(x1) p(x2∣x1)⋯p(xn∣x1,…,xn−1)p(x_1, \dots, x_n) = p(x_1)\, p(x_2 \mid x_1) \cdots p(x_n \mid x_1, \dots, x_{n-1})

But the chain rule works in any order. You could predict the last word first, then the third, then the rest. A model that can predict any subset of positions given any other subset can generate in any order. That's the idea behind masked diffusion language models.

The forward and reverse process

  • Forward process (fixed, no learning): take real text and replace each token with [MASK] independently with probability tt. At t=0t = 0 the text is intact; at t=1t = 1 it's all masks.
  • Reverse process (learned): a network looks at partially masked text and predicts, for each masked position, a distribution over what token belongs there.

To generate, start at all masks and run the reverse process in steps: predict every masked position, reveal some of them by sampling, repeat until nothing is masked.

example · restoring a damaged manuscript

An archivist restores a page with half the words smudged. They fill in the obvious ones first ("the", "of"), then use those to pin down the harder ones. No left-to-right order is needed; each word is guessed from all the legible words around it, before and after.

example · training on masked text

Take "the capital of France is Paris". Pick t=0.5t = 0.5 and mask about half: "the [MASK] of France [MASK] Paris". Train the model to put high probability on "capital" and "is" in those slots. Repeat with every tt between 0 and 1. It's BERT's masked-language-model objective, with a random masking ratio and a weighting that turns it into a proper likelihood bound.

The catch: parallel tokens are independent

Within one step, the model gives each masked position its own distribution, and samples them independently. But the true tokens are not independent of each other. If two positions depend on each other strongly, sampling both at once can produce a combination the data never contains.

example · France and Tokyo

After "the capital of [MASK] is [MASK]", the model's slot predictions might be: country 40% France, 35% Japan, 25% Italy; city 30% Paris, 25% Tokyo, and so on. Each is right on its own. Sampled independently in the same step, "France ... Tokyo" comes out about 10% of the time. Reveal one first, and the other is predicted given it, and the clash disappears.

That's the core speed–quality trade-off of diffusion LMs:

  • One token per step: an exact sample from the model, but as slow as a normal language model.
  • Many tokens per step: fast, but with independence errors between tokens revealed together.
  • Which tokens to reveal: revealing in each step every position the model is confident about, and only those, avoids most clashes. Confident positions barely depend on the others, so sampling them together is safe; uncertain, dependent ones go one at a time. LLaDA's low-confidence remasking and Fast-dLLM's threshold decoding both work this way.

Try it

[mask]the100%
[mask]capital70%food20%weather10%
[mask]of100%
[mask]France40%Japan35%Italy25%
[mask]is100%
[mask]Paris30%Tokyo25%great20%
steps so far0
result–
over 2,000 runs: average steps3.0
incoherent18.4%
Several words per step are sampled independently of each other. The model knows “France” and “Tokyo” never go together, but if both are revealed in the same step, nothing stops it. Fewer steps means faster decoding and more of these errors.
A model trained on six sentences starts from all masks. Each step it predicts every masked word given the revealed ones, then reveals some, each sampled on its own.

Why it's a likelihood model

Masked diffusion isn't a heuristic. Its training loss is an evidence lower bound (latent variables) on the log-likelihood of the text, with the masked versions as latent variables. So diffusion LMs can be compared with autoregressive ones on perplexity, and trained with RL using the same tools, with one difficulty: the exact log-probability of a whole answer, needed for importance ratios and policy gradients, isn't available, only a bound. Methods like d1 estimate it.

Where this shows up

  • LLaDA, Dream, Mercury, Gemini Diffusion: masked diffusion at scale, competitive with autoregressive models of similar size on many tasks.
  • RL for diffusion LMs (d1 and successors): policy gradients with estimated sequence log-probabilities.
  • The any-order view connects diffusion to autoregressive models over random orders.
  • It's a Markov chain from all-masked to complete text.
going deeper

Forward process

Each token is independently masked with probability t∈[0,1]t \in [0, 1]:

q(xt∣x0)=∏i=1n[(1−t) 1[xti=x0i]+t 1[xti=M]]q(x_t \mid x_0) = \prod_{i=1}^{n} \Big[(1 - t)\, \mathbb{1}[x_t^i = x_0^i] + t\, \mathbb{1}[x_t^i = \texttt{M}]\Big]

Training loss

The network pθ(x0i∣xt)p_\theta(x_0^i \mid x_t) predicts the clean token at each masked position. LLaDA and MDLM minimise

L(θ)=− Et∼U(0,1), x0, xt∼q(xt∣x0)[1t∑i=1n1[xti=M] log⁡pθ(x0i∣xt)]\mathcal L(\theta) = -\,\mathbb{E}_{t \sim U(0,1),\ x_0,\ x_t \sim q(x_t \mid x_0)}\Big[\frac{1}{t} \sum_{i=1}^{n} \mathbb{1}[x_t^i = \texttt{M}]\, \log p_\theta(x_0^i \mid x_t)\Big]

which upper-bounds the negative log-likelihood: −log⁡pθ(x0)≤L-\log p_\theta(x_0) \le \mathcal L. The 1/t1/t weight balances lightly and heavily masked examples; with it, the expected number of masked tokens per example is matched to their information content.

Sampling

From s>ts \gt t to tt (fewer masks), each masked position is revealed with probability (s−t)/s(s - t)/s, its token drawn from pθ(x0i∣xs)p_\theta(x_0^i \mid x_s), independently across positions. The joint distribution of tokens revealed in one step is therefore the product of marginals, not the true conditional joint, which is the source of the errors the widget counts.

The any-order view

Averaging over masking ratios, the loss equals an expected negative log-likelihood under a random generation order σ\sigma:

L=Eσ[−∑k=1nlog⁡pθ(xσ(k)∣xσ(<k))]\mathcal L = \mathbb{E}_{\sigma}\Big[-\sum_{k=1}^{n} \log p_\theta\big(x^{\sigma(k)} \mid x^{\sigma(\lt k)}\big)\Big]

an autoregressive model trained over every order at once.

Papers

Related

  • latent-variables - Treat the chain of thought as an unobserved variable between question and answer. Majority voting, rejection-sampling fine-tuning and several RL methods turn out to be standard tools for latent-variable models.
  • mcmc - Random walks that forget their past, and how to design one that spends its time in proportion to any distribution you want. Behind recent methods that sample better answers from a base model without any training.
  • decoding - How a probability distribution becomes an actual word. Inverse-CDF sampling, greedy and beam search, temperature, top-k, top-p and min-p, and the Gumbel-max trick, with what each one does to the distribution you're really sampling from.