The Mathematics of Discrete Diffusion

A companion lesson introduced diffusion language models at the level of mechanism: mask tokens, predict them in parallel, unmask iteratively. This path takes the machinery apart, and the right place to start is what a diffusion process must supply before any of that works.

Continuous diffusion leans on three properties of Gaussian noise. You can jump to any corruption level in closed form without simulating the intermediate steps. The endpoint is a distribution you can sample from trivially. And the reverse of a small noising step is itself approximately Gaussian, so a network predicting a mean is enough.

Discrete data has none of this for free. Adding noise to token 17,432 is undefined, because token IDs are labels, not quantities. Any discrete process has to be built so that the three properties hold.

The construction that works treats corruption as a Markov chain over the vocabulary, and once you write it that way, the entire apparatus of diffusion follows from linear algebra on categorical distributions.

Jacob Austin and colleagues formalised this in "Structured Denoising Diffusion Models in Discrete State-Spaces" (NeurIPS 2021), the paper that named D3PM and set the framework everything since has used.

Represent a token as a one-hot row vector $x$ over a vocabulary of size $V$. Define the corruption at step $t$ by a $V \times V$ transition matrix $Q_t$, where entry $[Q_t]_{ij}$ is the probability that token $i$ becomes token $j$:

$q(x_t \mid x_{t-1}) = \text{Cat}(x_t; \, p = x_{t-1} Q_t)$

The crucial consequence is that composing steps is just matrix multiplication. Define the cumulative matrix $\bar{Q}_t = Q_1 Q_2 \cdots Q_t$, and you get the marginal at any level in one shot:

$q(x_t \mid x_0) = \text{Cat}(x_t; \, p = x_0 \bar{Q}_t)$

That recovers the first required property. During training you sample a random $t$, corrupt directly to that level, and never simulate the chain.

Everything now hinges on one design decision: what to put in $Q_t$.

D3PM catalogued several structured choices, and the differences between them are substantive rather than cosmetic.

Uniform. With some probability, replace the token with one drawn uniformly from the vocabulary. This is the discrete analogue of isotropic noise and was the earlier multinomial diffusion approach. Its weakness is that the model cannot tell a corrupted token from a genuine one, so it must second-guess every position at every step.

Absorbing. With some probability, replace the token with a special [MASK] state that has no outgoing transitions. Once masked, always masked.

Ordinal or Gaussian-like. Concentrate transition probability on nearby token indices, mimicking a Gaussian kernel. Sensible for quantised images where index order means something, meaningless for text where token 500 and token 501 are unrelated.

Embedding-nearest-neighbour. Transition preferentially to tokens that are close in embedding space, which injects semantic structure into the corruption.

Text converged on the absorbing kernel, and it is worth being precise about why, because the reason is structural rather than empirical taste.

The absorbing kernel has a property no other choice offers: corruption is self-identifying. A [MASK] symbol announces that this position is noise; every other position is known-clean. The model never wastes capacity deciding what is signal.

Its cumulative matrix is also trivially simple. Under a schedule where the mask probability at time $t$ is $\alpha_t$, a token has either survived untouched or been absorbed:

$q(x_t \mid x_0) = (1 - \alpha_t)\,\delta_{x_0} + \alpha_t\,\delta_{\texttt{[MASK]}}$

No matrix powers, no vocabulary-sized products. Sampling a corrupted sequence is one Bernoulli draw per position.

The reverse process inherits the same simplicity. Because masking is irreversible in the forward direction, the reverse only ever needs to answer one question: given the visible context, what was this masked token? Unmasked positions require no decision at all.

The uniform kernel, by contrast, forces the reverse model to consider revising every position at every step, a strictly harder problem with no compensating benefit for text. The empirical gap follows from the structural one.

Every discrete diffusion model is assembled from the same four decisions. Fixing the kernel to absorbing collapses most of the downstream complexity, which is why the masked branch of this tree became the mainstream path.

Like every diffusion model, a discrete one is trained by maximising a variational bound on the log-likelihood. The generic D3PM bound is a sum over timesteps of KL divergences between the true reverse posterior $q(x_{t-1} \mid x_t, x_0)$ and the model's $p_\theta(x_{t-1} \mid x_t)$, plus a reconstruction term.

For the absorbing kernel this collapses dramatically. Because an unmasked token is known with certainty and a masked one is predicted directly from context, each KL term reduces to a cross-entropy on masked positions, weighted by the noise schedule. In continuous time the bound becomes

$\mathcal{L} = \mathbb{E}_{t,\,x_t}\left[\frac{\alpha_t'}{\alpha_t} \sum_{i \,:\, x_t^i = \texttt{[MASK]}} \log p_\theta(x_0^i \mid x_t)\right]$

Read that carefully, because it is the punchline of the whole framework: the loss is masked language modelling with a schedule-dependent weight. BERT's objective, generalised over all masking ratios instead of a fixed 15 percent, and reweighted so it forms a genuine bound on sequence likelihood.

That single observation is why diffusion language models could be trained with existing infrastructure almost immediately.

Stripped to its essentials, a training step is short enough to read in full:

def mdlm_loss(model, x0, mask_id):
    B, L = x0.shape

    # 1. sample a corruption level per sequence, uniform in (0, 1]
    t = torch.rand(B, 1, device=x0.device).clamp_min(1e-3)

    # 2. corrupt: mask each token independently with probability t
    masked = torch.rand(B, L, device=x0.device) < t
    xt = torch.where(masked, mask_id, x0)

    # 3. predict the originals at every position
    logits = model(xt)                      # bidirectional, no causal mask

    # 4. cross-entropy on masked positions only, weighted by 1/t
    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")
    return ((ce * masked) / t).sum() / masked.sum().clamp_min(1)

Three details carry the weight. The masking ratio is resampled every batch rather than fixed, so the model learns to denoise from any corruption level. The loss is computed only on masked positions, since unmasked ones are given. And the $1/t$ factor is the schedule weighting that turns a heuristic into a likelihood bound: lightly masked examples are easy and get downweighted.

The bound above is correct but was not immediately competitive. Two papers closed most of the remaining gap by changing how the reverse model is parameterized, not by changing the underlying process.

MDLM, from Subham Sahoo and colleagues (NeurIPS 2024), introduced the SUBS parameterization, built on two constraints that are true by construction and so should be hard-coded rather than learned. Zero masking probability: the model should never predict [MASK] as an original token. Carry-over unmasking: an already-revealed token stays fixed and is never re-predicted. Enforcing both simplifies the absorbing-state bound into a clean weighted mixture of masked-language-modelling losses, and a Rao-Blackwellised continuous-time version cuts gradient variance further. The result set state-of-the-art perplexity among diffusion models on standard benchmarks.

SEDD, from Aaron Lou, Chenlin Meng, and Stefano Ermon (ICML 2024), took a different route: rather than predicting tokens, estimate ratios between probabilities of neighbouring sequences, the discrete analogue of a score. Its score-entropy objective is the closest thing to score matching that discrete spaces admit, and it substantially narrowed the perplexity gap to autoregressive models.

The three formulations that matter, judged on what they actually give you:

One gotcha deserves emphasis before the next lesson, because it recurs in every published comparison. These models report bounds on log-likelihood, not exact values. An autoregressive model factorises the sequence probability exactly and reports true perplexity.

So a diffusion model reporting perplexity 20 against an autoregressive model reporting 18 is not straightforwardly worse: the diffusion number is an upper bound on its true perplexity, and the size of the gap between bound and truth is unknown. Comparisons in this literature are apples to oranges unless the evaluation protocol is stated, and this is the single most common way results here get over-read.