9. Discrete diffusion

Discrete states and masked diffusion Draft

Markov chains on finite sets and their reversal

A token belongs to a finite vocabulary, so adding Gaussian noise does not make sense in the token's own state space. We can still use the diffusion strategy: corrupt the data with a Markov process, learn what is needed to reverse it, and sample backward. Probability ratios between neighbouring states take the place of the continuous score.

We will build this construction for finite state spaces, then specialize it to masking. With exact conditional predictions, revealing one token at a time can recover the data law. Revealing several independently in one step introduces an approximation when those tokens depend on one another. That tradeoff is central to the speed claims of the diffusion language models discussed at the end of this part.

Definition 9.1.1 (continuous-time Markov chain, rate matrix). On a finite set \(\mathcal X\), a continuous-time Markov chain is specified by rate matrices \(Q_t\) with \(Q_t(x, y) \ge 0\) for \(y \ne x\) and \(\sum_y Q_t(x, y) = 0\): in a short time \(h\) the chain at \(x\) jumps to \(y \ne x\) with probability \(Q_t(x, y)\, h + o(h)\). Its generator is \((\mathcal L_t \phi)(x) = \sum_y Q_t(x, y)\, (\phi(y) - \phi(x))\), and its marginals \(p_t\), as row vectors, satisfy the forward equation \(\dot p_t = p_t\, Q_t\), that is, \(\dot p_t(y) = \sum_x p_t(x)\, Q_t(x, y)\).

The forward equation is the Fokker–Planck equation of Section 1.2 with the differential operator replaced by a matrix; it is again linear in \(p_t\), so the whole construction of Section 6.2 will go through.

A rate is probability per unit time. The off-diagonal entries describe possible jumps; the negative diagonal records the total rate of leaving a state. That diagonal is bookkeeping, not a negative jump probability. In the forward equation, probability entering one state is probability leaving another, just as fluxes balanced across cell boundaries in Section 1.

Theorem 9.1.2 (time reversal of a chain). Let \(p_t > 0\) on \(\mathcal X\). The process \(y_s = x_{1 - s}\) is a Markov chain whose rate from \(y\) to \(x \ne y\) at time \(s\) is

\[\bar Q_s(y, x) \;=\; Q_{1 - s}(x, y)\, \frac{p_{1 - s}(x)}{p_{1 - s}(y)} . \tag{9.1}\]

Proof of the marginal statement. Let \(q_s = p_{1 - s}\), and compute the forward equation of the chain with rates (9.1) at a state \(x\), everything at time \(1 - s\):

\[\begin{aligned} \text{inflow} &= \sum_{y \ne x} q_s(y)\, \bar Q_s(y, x) = \sum_{y \ne x} p(y)\, Q(x, y)\, \frac{p(x)}{p(y)} \\ &= p(x) \sum_{y \ne x} Q(x, y) = -p(x)\, Q(x, x), \\ \text{outflow} &= q_s(x) \sum_{y \ne x} \bar Q_s(x, y) = \sum_{y \ne x} p(y)\, Q(y, x). \end{aligned}\]

Their difference is \(-\sum_y p(y)\, Q(y, x) = -\dot p_{1-s}(x) = \frac{d}{ds} q_s(x)\). So \(q_s\) solves the forward equation of the reversed chain from \(q_0 = p_1\), and the reversed chain has marginals \(p_{1 - s}\). ∎The hypothesis \(p_t > 0\) is not idle for masking. A partially masked sequence whose revealed tokens cannot occur together under the data has \(p_t = 0\), and the ratio in (9.1) is undefined there; the forward chain never visits such a state, but the sampler of Section 9.3 with \(k \ge 2\) can produce one, and what a network answers there is outside the theory. In Example 9.3.3 every sequence has positive probability, so the question does not arise.

The reverse rate from \(y\) to \(x\) is the forward rate from \(x\) to \(y\) multiplied by \(p_t(x)/p_t(y)\).For masking the ratio is explicit. If \(y\) is \(x\) with one more position \(\ell\) masked, then \(p_t(x)/p_t(y) = \frac{\alpha_t}{1 - \alpha_t}\, p_0(x^\ell \mid \text{the other revealed tokens of } x)\), and with the forward rate \(-\dot\alpha_t/\alpha_t\) of masking a position, (9.1) gives \(\frac{-\dot\alpha_t}{1 - \alpha_t}\, p_0(x^\ell \mid \text{revealed})\) as the rate of revealing the token \(x^\ell\): the unmasking density of Theorem 9.2.3 times the conditional of Proposition 9.2.4. The factor \(\alpha_t/(1 - \alpha_t)\) is the time-dependent scalar of Ou et al. in the sidenote there. For sequences of \(L\) tokens with single-token transitions, we need ratios only between sequences that differ at one position. Lou, Meng and Ermon call this collection the concrete score. It plays the role that local density information, encoded by \(\nabla \log p_t\), plays in continuous diffusion.

The ratio has the same inferential meaning as the continuous score correction. With equal forward jump rates, a state containing twice as much probability sends twice as many particles along that edge. On observing the arrival state and looking backward, we must account for that difference in possible origins. Reversing the arrows alone would lose it.

left   three states A, B, m at t = 0.5 under the masking chain with alpha_t = 1 − t and p_0 = (0.667, 0.333):
       occupancies 0.333, 0.167, 0.500; the forward arrows A → m and B → m at the rate 2.000 carry the
       currents 0.667 and 0.333; the reversed arrows m → A and m → B at the rates 1.333 and 0.667 of (9.1)
       carry the same currents back (arrows only: the rate 2.000 on both, currents 1.000 and 1.000)
right  the marginals p_t(A), p_t(B), p_t(m) against t, with the reversed chain's marginals as rings on the
       curves, ending at p_0 = (0.667, 0.333) by (9.1) and at (0.500, 0.500) under arrows only
(the live figure needs JavaScript)
Reversing a discrete chain requires matching probability currents. Left: state probabilities and forward and reverse arrows, labelled with currents and rates. Right: the forward marginals and the distributions reached by the reversed chain. Correct reversal multiplies each rate by the ratio of source to destination probability. Reversing only the arrows loses that weighting. Change the initial odds or switch between masking and uniform replacement. The masking calculation is exact; the uniform reverse chain is integrated numerically.

Their score entropy loss learns these ratios, much as denoising score matching learns a continuous score. Campbell et al. develop the continuous-time framework and its evidence lower bound. Austin et al. study the discrete-time version, where transition matrices \(Q_i\) can replace tokens uniformly, follow a structured corruption rule, or send them to an absorbing mask.J. Austin, D. D. Johnson, J. Ho, D. Tarlow and R. van den Berg, "Structured denoising diffusion models in discrete state-spaces", NeurIPS 2021 (arXiv 2107.03006), for the matrices "that mimic Gaussian kernels", "based on nearest neighbors in embedding space", and "that introduce absorbing states", the last drawing "a connection between diffusion models and autoregressive and mask-based generative models"; A. Campbell, J. Benton, V. De Bortoli et al., "A continuous time framework for discrete denoising models", NeurIPS 2022 (arXiv 2205.14987); A. Lou, C. Meng and S. Ermon, "Discrete diffusion modeling by estimating the ratios of the data distribution", ICML 2024, best paper (arXiv 2310.16834): "score entropy, a novel loss that naturally extends score matching to discrete spaces", with models that reduce "perplexity by 25-75%" against earlier diffusion language models and are "competitive with autoregressive models, in particular outperforming GPT-2" on most of its zero-shot tasks.

Masked diffusion

Masking is a particularly simple corruption process. Each position either retains its original token or becomes a mask. Once masked, it stays masked during the forward process. At the noise end, every position is a mask; reversing the process fills the tokens back in.

Definition 9.2.1 (masked, or absorbing, diffusion). Let the vocabulary \(\mathcal V\) be extended by a mask token \(\mathbf m\). Each of the \(L\) positions of a sequence \(x_0 \in \mathcal V^L\) is masked independently, at a rate chosen so that the probability of a position being still unmasked at time \(t\) is \(\alpha_t\), decreasing from \(\alpha_0 = 1\) to \(\alpha_1 = 0\):

\[q(x_t \mid x_0) = \prod_{\ell = 1}^{L} \Bigl(\alpha_t\, \mathbf 1[x_t^\ell = x_0^\ell] + (1 - \alpha_t)\, \mathbf 1[x_t^\ell = \mathbf m]\Bigr).\]

A masked position never returns, and \(x_1\) is the all-mask sequence with probability one.

Proposition 9.2.2 (the posterior of a step). For \(s < t\) and a position \(\ell\): if \(x_t^\ell \ne \mathbf m\) then \(x_s^\ell = x_t^\ell\); if \(x_t^\ell = \mathbf m\) then, given \(x_0\), \(x_s^\ell = \mathbf m\) with probability \((1 - \alpha_s)/(1 - \alpha_t)\) and \(x_s^\ell = x_0^\ell\) otherwise, independently across positions.

Proof. Bayes' rule position by position: a position unmasked at \(t\) was unmasked at \(s\); a position masked at \(t\) was masked at \(s\) with probability \((1 - \alpha_s)/(1 - \alpha_t)\), the ratio of the two masking probabilities, and otherwise carried its token. ∎

The reverse step therefore needs one thing from a model: the law of the clean token at a masked position given the current partially masked sequence, \(p_\theta(x_0^\ell \mid x_t)\). Sahoo et al. and Shi et al. showed that the continuous-time evidence lower bound of this model is a weighted cross-entropy on the masked positions, the loss of a masked language model with a particular weighting.The bound also does not depend on the schedule. For a network that does not look at \(t\) the integrand of (9.2) depends on \(t\) only through \(\alpha_t\), and the substitution \(u = \alpha_t\) turns (9.2) into \(\int_0^1 \frac{\mathbb E_u[\cdots]}{1 - u}\, du\), with \(\mathbb E_u\) the expectation of the same sum when each position is masked with probability \(1 - u\), for every decreasing schedule: the counterpart of Theorem 3.4.1; with \(\alpha_t = 1 - t\) the weight is \(1/t\). A schedule chooses only how the training time of Algorithm 9.2.5 is spread over masking rates, which is the lever Hong et al. use in Section 9.5.

The training labels are available for the same reason as in Gaussian denoising: we kept the original sequence before masking it. The target is now a categorical distribution. Cross-entropy rewards the probability assigned to the original token, and its ideal predictor returns the full conditional probabilities. Always choosing the most likely token would discard the uncertainty that a sampler needs.

Theorem 9.2.3 (the bound is a weighted masked cross-entropy). For a model that predicts each masked token independently, \(p_\theta(x_0 \mid x_t) = \prod_{\ell: x_t^\ell = \mathbf m} p_\theta(x_0^\ell \mid x_t)\), and a differentiable \(\alpha_t\),

\[-\log p_\theta(x_0) \;\le\; \int_0^1 \frac{-\dot\alpha_t}{1 - \alpha_t}\; \mathbb E_{x_t \sim q(\cdot \mid x_0)}\Bigl[\sum_{\ell:\, x_t^\ell = \mathbf m} -\log p_\theta(x_0^\ell \mid x_t)\Bigr] dt . \tag{9.2}\]

Proof sketch. Write the bound of Theorem 3.3.1 for a grid of times. Each term is a relative entropy between the true posterior of a step, Proposition 9.2.2, and the model's, which unmasks the same positions with the same probabilities but draws the token from \(p_\theta\); the relative entropy of a point mass at \(x_0^\ell\) from the categorical \(p_\theta(\cdot \mid x_t)\) is \(-\log p_\theta(x_0^\ell \mid x_t)\), and it is incurred only when the position is unmasked in the step, which happens with probability \((\alpha_s - \alpha_t)/(1 - \alpha_t)\) per masked position. Summing over the grid and letting it refine gives \(\int \frac{-\dot\alpha_t}{1 - \alpha_t}\, dt\) as the density of unmasking per masked position, and (9.2). ∎S. S. Sahoo, M. Arriola, Y. Schiff et al., "Simple and effective masked diffusion language models", NeurIPS 2024 (arXiv 2406.07524): "a simplified, Rao-Blackwellized objective … a mixture of classical masked language modeling losses"; J. Shi, K. Han, Z. Wang, A. Doucet and M. K. Titsias, "Simplified and generalized masked diffusion for discrete data", NeurIPS 2024 (arXiv 2406.04329): "the continuous-time variational objective of masked diffusion models is a simple weighted integral of cross-entropy losses", with \(2.75\) and \(3.40\) bits per dimension on CIFAR-10 and ImageNet \(64 \times 64\) as image models. The perplexity bounds reached at the scale of GPT-2 on OpenWebText were \(\le 23.21\) for MDLM against \(17.54\) for an autoregressive transformer trained alike, and \(\le 27.04\) against \(22.32\) on LM1B.

Here a visible token is exact evidence: the corruption process either preserves it or hides it completely. The time does not change how trustworthy that visible token is. This differs from a Gaussian observation, whose reliability depends on its noise level. Independent masking also means that the pattern of missing positions carries no extra information about the hidden token values.

Proposition 9.2.4 (the conditional does not depend on the time). Under Definition 9.2.1, the posterior of the clean sequence given the masked one depends only on the revealed tokens: \(q(x_0 \mid x_t) = p_0(x_0 \mid x_0^\ell = x_t^\ell \text{ for all unmasked } \ell)\), without reference to \(t\). The optimal predictor in (9.2) is therefore a time-independent function of the partially masked sequence, and it is the conditional law of the data at the masked positions given the revealed ones.

Proof. \(q(x_t \mid x_0) = \alpha_t^{\#\text{unmasked}}(1 - \alpha_t)^{\#\text{masked}}\, \mathbf 1[x_t \text{ agrees with } x_0 \text{ where unmasked}]\); the first factor depends on \(x_t\) only through its number of masks, so by Bayes' rule \(q(x_0 \mid x_t) \propto p_0(x_0)\, \mathbf 1[\text{agreement}]\), in which \(t\) does not appear. The optimality claim is Lemma 2.3.6 for the cross-entropy, whose minimizer is the conditional law. ∎J. Ou, S. Nie, K. Xue et al., "Your absorbing discrete diffusion secretly models the conditional distributions of clean data", ICLR 2025 (arXiv 2406.03736): "the concrete score in absorbing diffusion can be expressed as conditional probabilities of clean data, multiplied by a time-dependent scalar in an analytic form", whence a model "without time-condition" that can cache its output while the sequence does not change, and the unification of "absorbing discrete diffusion and any-order autoregressive models".

The optimal model predicts a token from whichever other tokens are visible. It can handle any reveal order and any pattern of gaps. Under this masking process, the visible context contains all the relevant information; time only controls the training weight and the sampling schedule. Ou et al. therefore omit time from the network input.This is special to the absorbing process. Under the uniform transitions of Austin et al. a token of \(x_t\) may have been replaced, so \(q(x_0 \mid x_t)\) depends on how long the chain has run and the network needs \(t\) to know how far to trust what it sees; masking is legible in the state itself, which is also why its reverse process can only add tokens and never revise one, the limitation that the editing systems of Sections 9.3 and 9.5 address.

the integrand of (9.2) on the clock for the six-token sticky chain, P(stay) = 0.8, with the exact
conditionals: w(t) E_t under the linear schedule alpha_t = 1 − t runs from 2.50 at t = 0 to 4.16 at t = 1
through 3.13 at t = 0.5, where 3.00 positions are masked on average, E_t = 1.565 nats and w = 2; the
cosine schedule's integrand runs from 0 to 6.53 and the quadratic's from 4.99 to 0; the area under each
is 3.195 nats = log 2 + 5 H_b(0.8), the entropy of the data law
(the live figure needs JavaScript)
Masking schedules redistribute the same ideal training loss. On six correlated binary tokens, the dashed curve is expected cross-entropy over masked positions; multiplying by the unmasking weight gives the loss integrand. Linear, cosine and quadratic schedules place it differently in time but have the same area: the data entropy, reached here using exact conditionals. Switch to the masking-rate axis to make the curves coincide. The ticks show how uniformly sampled training times cover masking rates.
Algorithm 9.2.5  MASKED-DIFFUSION-TRAIN
                 (the bound (9.2) with t uniform)

Picture: one example. A fraction 1 - alpha_t of the
positions is masked, the network predicts the clean
token at every masked position, and the cross-entropy
of those predictions is weighted by -alpha'_t/(1-alpha_t).

   x_0:  the cat sat on the mat
   x_t:  the [m] sat [m] the [m]        t drawn, 3 masked
   loss: w(t) [ -log p(cat|x_t) - log p(on|x_t) - log p(mat|x_t) ]

Input   sequences of tokens; a masking schedule alpha_t;
        a network p_theta(. | x_t) over the vocabulary at
        every position.
Output  theta.

 1. repeat until converged:
 2.    draw x_0 and t ~ U(0, 1)
 3.    mask each position independently with probability
       1 - alpha_t, giving x_t
 4.    L <- (-alpha'_t / (1 - alpha_t)) sum over masked
       positions of -log p_theta(x_0^l | x_t)
 5.    take a gradient step on L

One token at a time is exact, many at a time is not

Sampling starts with an all-mask sequence. At each step, choose positions to reveal and draw their tokens from the model. We can change both the reveal order and the number of tokens drawn together. Drawing several at once reduces the number of network calls, but we must understand what dependence is lost when those draws are made independently.

Algorithm 9.3.1  MASKED-SAMPLER
                 (reveal k positions per step)

Picture: six positions, two revealed per step; the
two tokens of one step are drawn independently from
the model's conditionals given the state before it.

   step 1   [m] [m] [m] [m] [m] [m]  ->  [m]  A  [m] [m]  B  [m]
   step 2    .   A   .   .   B   .   ->   A   A  [m]  B   B  [m]
   step 3                                 A   A   B   B   B   A

Input   a model p_theta(. | x) of the clean token at each
        masked position; a length L; a number k per step;
        a rule for choosing positions (uniform at random, or
        the k whose predicted law is most concentrated).
Output  a sequence.

 1. x <- the all-mask sequence
 2. while x has masked positions:
 3.    choose a set U of min(k, #masked) masked positions
 4.    for each l in U, independently:
          x^l <- a draw from p_theta(. | x)   (x as before step 3)
 5. return x

Consider a two-token dataset containing only AA and BB, equally often. With both positions masked, each position individually predicts A or B with equal probability. Independent parallel draws then produce AA, AB, BA and BB equally often, including two sequences that never occur in the data. Revealing one token first resolves the ambiguity: after A the other must be A, and after B it must be B. Both predictors were exact; the difference lies in when the second prediction receives the first token as context.

Proposition 9.3.2 (exactness and its failure). Suppose the model's conditionals are exact, \(p_\theta(\cdot \mid x) = p_0(x_0^\ell \mid \text{the revealed tokens of } x)\). (a) With \(k = 1\), Algorithm 9.3.1 produces an exact sample of \(p_0\), for any rule choosing the positions, including rules that look at the revealed tokens. (b) With \(k \ge 2\), the \(k\) tokens of a step are drawn from the product of their marginal conditionals, which equals their joint conditional only when the \(k\) positions are conditionally independent given the revealed tokens; otherwise the output law differs from \(p_0\).Part (a) needs the exact conditionals. With a learned model the \(k = 1\) sampler is an any-order autoregressive model whose output law does depend on the order, because the model's errors differ from context to context; a rule by confidence then keeps the sampler in the contexts in which the model is surest. Ou et al.'s identification with any-order autoregressive models, in the sidenote of Proposition 9.2.4, is the same object seen from the other side.

Proof. (a) By the chain rule, the probability of producing \(x\) along an order \(\ell_1, \dots, \ell_L\) is \(\prod_j p_0(x^{\ell_j} \mid x^{\ell_1}, \dots, x^{\ell_{j-1}}) = p_0(x)\), whatever the order; if the order is chosen by a rule that depends on the tokens revealed so far, the same identity holds conditionally on each choice, and averaging over the rule's choices gives \(p_0(x)\) again. (b) The joint law of two tokens given a context factorizes into the product of their marginals given the context exactly when they are conditionally independent given it; the sampler draws from the product. ∎

Parallel draws use an independence approximation: each token sees the revealed context, but not the other tokens being drawn in that step. How much this matters depends on the dependencies in the data. The six-token example below is small enough to compute the resulting output law exactly.

left   one run of Algorithm 9.3.1 on two to eight binary tokens from a sticky two-state chain,
       revealing k per step from the exact conditionals, the last step's tokens framed in green;
       bars: P(A | revealed) at each position, hatched once revealed
right  the exact law of the sampler's output against the data law over the 64 sequences, with the
       total variation and KL between them; or the budget of that KL, one term per step; or TV and
       KL against k
(the live figure needs JavaScript)
Parallel unmasking loses dependencies between tokens revealed together. Left: reveal a correlated binary sequence in batches, using exact single-token conditionals. Right: compare the sampler’s full output law with the data law, inspect its stepwise error budget, or plot error against batch size. Change the reveal order, batch size and stickiness. One token at a time is exact; choosing the most likely token changes the distribution even then. Output laws are computed by dynamic programming over all partial sequences.

Example 9.3.3 (six tokens). With one token per step the output law is the data law to the last digit, by Proposition 9.3.2(a), in six steps and for either rule of choosing positions. With all six revealed at once, one step, each token is drawn from its marginal, which is uniform, so the output is uniform over the 64 sequences: the total variation from the data law is \(0.550\), the relative entropy \(0.964\) nats, and the alternating sequences ABABAB and BABABA, whose data probability is \(0.5 \times 0.2^5 = 1.6 \times 10^{-4}\), are produced \(97.7\) times too often. Two per step, three steps, leaves a total variation of \(0.117\); three per step, two steps, \(0.240\). Which positions are revealed together matters as much as how many.

For a rule that fixes its sets in advance, the relative entropy of the sampler's output is exactly the sum, over its steps, of the expected total correlation of the tokens drawn together given the tokens already known, \(\mathrm{KL}(p \| q) = \sum_s \mathbb E_p\, \mathrm{TC}(x^{U_s} \mid \text{context}_s)\), by the chain rule of relative entropy applied to the sampler's own factorization; for a random rule the average of the right side over the rule's choices is an upper bound.The bound is loose. For the random rule with \(k = 2\) the averaged budget is \(0.172\) nats against the \(0.049\) the output incurs: the random rule's output is the mixture, over its choices, of the outputs of fixed rules, and relative entropy is convex in its second argument, so mixing the orders helps. The figure's budget view draws the two as the rules Σ and KL. Revealing the neighbouring pairs \(\{1, 2\}, \{3, 4\}, \{5, 6\}\) costs \(0.193 + 0.126 + 0.126 = 0.446\) nats, the random rule \(0.049\), and the pairs \(\{1, 4\}, \{2, 5\}, \{3, 6\}\) only \(0.0235\): once positions 1 and 4 are known the rest are conditionally independent, because the chain is Markov. A rule that reveals the most confident positions first picks neighbours, since the neighbour of a known token is the most predictable one, and on this law it is the worst of the three; the figure's second view draws the budget term by term. The program below computes these laws. ◁

# Masked diffusion over six binary tokens from a sticky two-state chain,
# P(stay) = 0.8, uniform start: the exact law of a sampler that reveals k
# masked positions per step, each from its exact conditional, by dynamic
# programming over the 3^6 partial sequences (-1 masked, 0 = A, 1 = B).
# With a partition, the positions revealed at each step are fixed in advance,
# and the relative entropy of the output is the sum, over the steps, of the
# expected total correlation of the tokens drawn together given the context.
#
# Picture: one step of the sampler with k = 2 on a state
# with four masked positions; the two new tokens are drawn
# independently, from conditionals given the state before.
#
#   [m] A [m] [m] B [m]   -- choose {1, 3}, draw each --> [m] A  B  [m] B  A
#                                                              ^        ^
import itertools, math

L, Q = 6, 0.8
seqs = list(itertools.product((0, 1), repeat=L))
p = {}
for s in seqs:
    pr = 0.5
    for a, b in zip(s, s[1:]): pr *= Q if a == b else 1 - Q
    p[s] = pr

def agrees(rev, s): return all(r < 0 or r == v for r, v in zip(rev, s))

def conditional(rev, pos):                  # P(x_pos = 1 | the revealed tokens)
    num = den = 0.0
    for s in seqs:
        if agrees(rev, s):
            den += p[s]; num += p[s] * s[pos]
    return num / den

def output_law(k, partition=None):          # the sets of a step: every k-subset at random, or the partition's
    layer = {tuple([-1] * L): 1.0}
    out = {s: 0.0 for s in seqs}
    step = 0
    while layer:
        nxt = {}
        for rev, pr in layer.items():
            masked = [i for i in range(L) if rev[i] < 0]
            if not masked:
                out[rev] += pr; continue
            if partition is None:
                sets = list(itertools.combinations(masked, min(k, len(masked))))
            else:
                sets = [tuple(partition[step])]
            for st in sets:
                probs = [conditional(rev, pos) for pos in st]    # all from the state before the step
                for vals in itertools.product((0, 1), repeat=len(st)):
                    r2, q = list(rev), pr / len(sets)
                    for pos, v, pb in zip(st, vals, probs):
                        r2[pos] = v; q *= pb if v else 1 - pb
                    nxt[tuple(r2)] = nxt.get(tuple(r2), 0.0) + q
        layer = nxt
        step += 1
    return out

def budget(partition):                      # per step, E_p[ TC(x^U | the tokens known before the step) ]
    terms, known = [], []
    for U in partition:
        term = 0.0
        for ctx in itertools.product((0, 1), repeat=len(known)):
            rev = [-1] * L
            for pos, v in zip(known, ctx): rev[pos] = v
            pc = sum(p[s] for s in seqs if agrees(rev, s))       # P(context) under the data law
            margs = [conditional(rev, pos) for pos in U]
            for vals in itertools.product((0, 1), repeat=len(U)):
                r2 = list(rev)
                for pos, v in zip(U, vals): r2[pos] = v
                joint = sum(p[s] for s in seqs if agrees(r2, s)) / pc
                prod = 1.0
                for v, m in zip(vals, margs): prod *= m if v else 1 - m
                if joint > 0: term += pc * joint * math.log(joint / prod)
        terms.append(term)
        known += list(U)
    return terms

for k in (1, 2, 3, 6):
    q = output_law(k)
    tv = 0.5 * sum(abs(p[s] - q[s]) for s in seqs)
    kl = sum(p[s] * math.log(p[s] / q[s]) for s in seqs if q[s] > 0)
    worst = max(seqs, key=lambda s: q[s] / p[s])
    name = "".join("AB"[v] for v in worst)
    print(f"k = {k}: {math.ceil(L / k)} steps, total variation {tv:.4f}, KL {kl:.4f}, "
          f"most over-produced {name} with q/p = {q[worst] / p[worst]:.1f}")

for name, parts in (("blocks {1,2} {3,4} {5,6}", ((0, 1), (2, 3), (4, 5))),
                    ("pairs {1,4} {2,5} {3,6}", ((0, 3), (1, 4), (2, 5)))):
    q = output_law(2, parts)
    tv = 0.5 * sum(abs(p[s] - q[s]) for s in seqs)
    kl = sum(p[s] * math.log(p[s] / q[s]) for s in seqs if q[s] > 0)
    terms = budget(parts)
    print(f"{name}: 3 steps, total variation {tv:.4f}, KL {kl:.4f} = budget "
          + " + ".join(f"{t:.4f}" for t in terms) + f" = {sum(terms):.4f}")
k = 1: 6 steps, total variation 0.0000, KL 0.0000, most over-produced AAABAB with q/p = 1.0
k = 2: 3 steps, total variation 0.1174, KL 0.0492, most over-produced BABABA with q/p = 6.9
k = 3: 2 steps, total variation 0.2402, KL 0.1794, most over-produced BABABA with q/p = 13.3
k = 6: 1 steps, total variation 0.5498, KL 0.9637, most over-produced BABABA with q/p = 97.7
blocks {1,2} {3,4} {5,6}: 3 steps, total variation 0.3871, KL 0.4457 = budget 0.1927 + 0.1265 + 0.1265 = 0.4457
pairs {1,4} {2,5} {3,6}: 3 steps, total variation 0.1080, KL 0.0235 = budget 0.0235 + 0.0000 + 0.0000 = 0.0235

Token accuracy and sequence correctness ask different questions. Even in a simple model with independent errors of probability \(\varepsilon\) per position, the chance of a completely correct length-\(L\) sequence is \((1-\varepsilon)^L\). Small local errors can therefore accumulate into frequent global failures. Dependencies can make this harder still. The theorem below studies this distinction for parallel decoding.

left   the total variation from the data law against L = 2, 4, 6, 8, 10 under the random rule with
       P(stay) = 0.8: two steps 0.000, 0.176, 0.240, 0.297, 0.352; one step 0.300, 0.396, 0.550, 0.625,
       0.660; two tokens per step (L/2 steps) 0.300, 0.176, 0.117, 0.105, 0.089; one token per step 0
right  the relative entropy per token, KL(p || q)/L in nats: two steps 0.000, 0.020, 0.030, 0.035, 0.038;
       one step 0.096, 0.145, 0.161, 0.169, 0.173; two tokens per step 0.096, 0.020, 0.008, 0.004, 0.003
(the live figure needs JavaScript)
A small per-token error can hide a large sequence error. For correlated binary sequences of increasing length, compare one step, two steps, two tokens per step, and one token per step. Left: total variation of the full output law. Right: relative entropy per token. Holding the step budget fixed increasingly exposes dependencies that parallel draws ignore. Change the reveal order or stickiness; each plotted law is computed exactly. The single-token sampler remains exact at every length.

Theorem 9.3.4 (Feng et al.: what parallel decoding can and cannot do). Under mild conditions on the data, a masked diffusion model with a learned conditional of error \(\varepsilon_{\mathrm{learning}}\) can achieve a near-optimal token error rate, a perplexity-type metric, with a number of sampling steps that is independent of the sequence length \(L\); for an \(n\)-gram language the steps needed are \(O((n - 1)/\varepsilon^n)\) for an excess of order \(\varepsilon\). But there is a hidden Markov model over a vocabulary of sixteen symbols such that, for any masking schedule and any learned conditional with error below \(1/128\), if the number of steps is \(CL\) for a certain constant \(C\) then the sequence error rate, the probability that the output is not a correct sequence, exceeds one half: for correct sequences the steps must grow linearly with \(L\).G. Feng, Y. Geng, J. Guan et al., "Theoretical benefit and limitation of diffusion language model", NeurIPS 2025 (arXiv 2502.09622), Theorems 4.2 to 4.4: "when using perplexity as the metric, MDMs can achieve near-optimal perplexity in sampling steps regardless of sequence length … when using the sequence error rate … the required sampling steps must scale linearly with sequence length to obtain 'correct' sequences, thereby eliminating MDM's efficiency advantage over autoregressive models".

The theorem makes precise a distinction already visible in the six-token example. A metric that averages over tokens can tolerate errors that a requirement for whole-sequence correctness cannot.The extreme case is in that example. With \(k = 6\) every token is drawn from its exact marginal, so a statistic of one position at a time is exactly right, while the law of the sequence is uniform, at total variation \(0.550\) from the data law; a per-token metric and a per-sequence metric can disagree completely about one sampler. The systems in Section 9.5 address this gap in several ways.

One approach limits which tokens are generated together. On a Markov chain, the dependence cost of a group is governed by its neighbouring members; this explains the factor-of-twenty difference between the block and interleaved orders in Example 9.3.3. Arriola et al.'s block diffusion generates blocks of \(L'\) tokens autoregressively and uses masked diffusion within each block. This confines the independence approximation to a block and allows attention over earlier blocks to be cached. With \(L' = 4\), their OpenWebText perplexity bound was \(\le 20.73\) against MDLM's \(\le 22.98\) and the autoregressive \(17.54\).

Another approach revises earlier decisions. Remasking and token editing let later steps correct a revealed token, which the absorbing process alone cannot do. Sahoo et al.'s diffusion-augmented models instead retain an autoregressive target law and use diffusion to propose several tokens. An exact acceptance test keeps or rejects the proposals, preserving the target law while gaining whatever speed the proposals allow.M. Arriola, A. Gokaslan, J. T. Chiu et al., "Block diffusion: interpolating between autoregressive and diffusion language models", ICLR 2025 (arXiv 2503.09573); the LLaDA2.1 report of Section 9.5 for token-to-token editing; S. S. Sahoo, L. Chen, K. Pham et al., "Unlocking lossless speedups in LLMs via discrete diffusion", arXiv 2609.04010 (September 2026): "a new class of models that defines an AR model distribution while using diffusion to draw multiple tokens in parallel from that distribution … up to 3× speedups over the base AR model".

The all-mask start also changes where randomness enters. In continuous flows, different initial Gaussian draws can produce different samples under a deterministic map. An all-mask sequence is the same input every time. Diversity must come from random token draws, a random reveal rule, or another source of randomness introduced by the sampler. A fully deterministic procedure from that single state would produce one sequence.

Discrete flow matching and generator matching

Flow matching also extends to discrete states: learn transition rates in place of a velocity. Gat et al. choose a conditional path that switches each position from a source token to a data token at a random time. In this series's convention, it is \(p_t(x^\ell \mid x_0, x_1) = (1 - \kappa_t)\, \delta_{x_1^\ell} + \kappa_t\, \delta_{x_0^\ell}\), where \(\kappa_t\) is the share of positions at their data value.

The rates for each conditional path are explicit. Averaging them under the posterior gives the marginal rates, by the discrete counterpart of Theorem 6.2.2. A network predicts the data token, the \(x\)-prediction, or the source token; the required rates follow from that prediction. Masked diffusion is the special case with an all-mask source.

At 1.7 billion parameters, Gat et al.'s model reaches \(6.7\%\) pass@1 and \(13.4\%\) pass@10 on HumanEval. Campbell et al.'s multiflow applies the same construction jointly to a protein's discrete sequence and continuous structure. Generator matching, introduced in Section 6.5, provides a common framework for these cases.I. Gat, T. Remez, N. Shaul et al., "Discrete flow matching", NeurIPS 2024 (arXiv 2407.15595): "a general family of probability paths interpolating between source and target distributions", "a generic formula for sampling from these probability paths using learned posteriors such as the probability denoiser (x-prediction) and noise-prediction", and "by scaling Discrete Flow Matching models up to 1.7B parameters, we reach 6.7% Pass@1 and 13.4% Pass@10 on HumanEval and 6.7% Pass@1 and 20.6% Pass@10 on 1-shot MBPP"; A. Campbell, J. Yim, R. Barzilay, T. Rainforth and T. Jaakkola, "Generative flows on discrete state-spaces: enabling multimodal flows with applications to protein co-design", ICML 2024 (arXiv 2402.04997); Holderrieth et al. (2025) and Lipman et al. (2024), Chapters 6 to 9, cited in Section 6.5.

The systems, with dates

Diffusion language models aim for high throughput: many tokens per second per accelerator, at quality comparable to autoregressive models of similar size. The table records systems from 2025 and 2026 with their dates and sources. Speeds are those reported by the vendor or paper on the named hardware; quality claims summarize the makers' benchmark results. Those details matter when comparing the rows.

date system size and recipe speed, as reported source
Feb 2025 Mercury Coder, Inception Labs commercial masked diffusion for code Mini 1,109 and Small 737 tokens/s on an H100, measured by Artificial Analysis on a February 2025 endpoint arXiv 2506.17298 (June 2025)
Feb 2025 LLaDA 8B, Nie et al. trained from scratch on 2.3 trillion tokens, 0.13 million H800 GPU-hours not the claim; "competitive with strong LLMs like LLaMA3 8B in in-context learning" arXiv 2502.09992
Apr 2025 Dream 7B, HKU and Huawei initialized from Qwen2.5 7B, 580 billion tokens not the claim; matches "top-tier autoregressive language models of similar size" on the makers' suite hkunlp.github.io/blog/2025/dream, arXiv 2508.15487
May 2025 Gemini Diffusion, Google DeepMind research model, announced at I/O on 20 May 2025 1,479 tokens/s average sampling speed across the reported evaluations, overhead 0.84 s; LiveCodeBench 30.9% against 28.5% for Gemini 2.0 Flash-Lite, GPQA Diamond 40.4% against 56.5% deepmind.google/models/gemini-diffusion, read 24 Oct 2026
Aug 2025 Seed Diffusion Preview, ByteDance code model 2,146 tokens/s on H20 GPUs, "significantly faster than contemporary Mercury and Gemini Diffusion" arXiv 2508.02193
Dec 2025 LLaDA2.0, Bie et al. 16B and 100B mixture-of-experts, converted from autoregressive models by a three-phase block-diffusion schedule   arXiv 2512.15745
Dec 2025 Efficient-DLM, Fu et al. (NVIDIA) 8B, converted with block-wise attention and position-dependent masking 4.5× and 2.7× the throughput of Dream 7B and Qwen3 4B at higher accuracy, as reported arXiv 2512.14067
Feb 2026 LLaDA2.1, Bie et al. 16B and 100B, token-to-token editing added to mask-to-token decoding, with a speedy and a quality mode 892, 801 and 663 tokens/s on HumanEval+, BigCodeBench and LiveCodeBench for the 100B model arXiv 2602.08676
Feb to Oct 2026 Mercury 2, 2.5 and 3, Inception Labs commercial; 3 is "more than 40× larger than Mercury 1" with a context of one million tokens 1,009 tokens/s on Blackwell (24 Feb), 1,107 (8 Sep), 1,068 (19 Oct) inceptionlabs.ai/blog, read 24 Oct 2026
Jul 2026 Nemotron-Labs-Diffusion, Fu et al. (NVIDIA) 3B, 8B and 14B; one model with autoregressive, diffusion and self-speculation decoding modes the 8B model "decodes 6x more tokens per forward than Qwen3-8B with comparable accuracy, translating to 4x higher throughput" on a GB200 arXiv 2607.05722
Jul 2026 DiffusionGemma, Google DeepMind 25.2B total, 3.8B active, fine-tuned from Gemma 4 with under a tenth of its token budget; blocks of 256 tokens "around 20 tokens per forward pass and … roughly 1,500 output tokens per second on a single NVIDIA H100" arXiv 2608.00146
Sep 2026 Uno, Sahoo et al. 8B diffusion-augmented autoregressive model; lossless by construction "up to 3× speedups over the base AR model" arXiv 2609.04010
Diffusion language models of 2025 and 2026 as read on 24 October 2026, with the numbers each source reports; the hardware differs between rows (H100, H20, Blackwell, GB200) and the speeds are not comparable across them. Sources: the arXiv abstract pages named, the DeepMind model page, the HKU blog of 2 April 2025, and the Inception Labs blog index.

Three results from 2026 connect the systems back to the theory. Hong et al. trace slow masked-diffusion training to language's locality: nearby context carries much of the information needed to predict a token. A bell-shaped distribution for \(t\) in Algorithm 9.2.5 speeds training by about a factor of four, paralleling the time-weighting choices in Section 3.4.

Chen et al. study a different obstacle. Every forward masking trajectory ends at the same all-mask state, leaving no terminal randomness for a consistency-style few-step sampler. They introduce multiple masks to restore it. Takahashi, Takahashi and Kabashima study discrete diffusion on Ising data and find the speciation and collapse transitions discussed in Section 10.3, with collapse described by the random energy model.C. Hong, S. Lee, C.-H. Lai et al., "Understanding and accelerating the training of masked diffusion language models", arXiv 2605.13026 (May 2026): "the main factor is the locality bias of language", with "bell-shaped time sampling" reaching the same validation loss "up to ∼4× faster" on LM1B; S. Chen, Y. Ren, H. Zhao et al., "Multi-mask diffusion language models for few-step generation", COLM 2026 (arXiv 2607.19686): "all forward trajectories collapse to a single fully masked state, leaving no terminal entropy for consistency-style few-step generation"; T. Takahashi, T. Takahashi and Y. Kabashima, "Dynamical regimes of discrete diffusion models", Journal of Statistical Mechanics (2026) 083402 (arXiv 2604.10961).

Remark 9.5.1 (takeaway). With exact conditionals, revealing one token at a time recovers the data law. Revealing several independently can be faster, but loses their conditional dependence. Block structure, reveal order, editing, and acceptance tests offer different ways to manage that tradeoff. Theorem 9.3.4 shows why it cannot be eliminated for every data law.

The learning problem remains familiar. Replace squared-error regression with cross-entropy, and the ideal prediction is the conditional law of a token given the revealed context, as Proposition 9.2.4 states.

← Back to all posts