2.12 Diffusion & Non-Autoregressive Language Models¶

Every model we have built so far in this book generates text the same way: one token at a time, strictly left to right, each token conditioned on all the tokens before it. This is the autoregressive (AR) factorization of the joint distribution over a sequence \(x = (x_1, \dots, x_L)\):

It is a beautiful, exact factorization — and it is also the source of a stubborn latency wall. Generating a 1000-token response requires 1000 sequential forward passes, each of which must finish before the next can begin, because \(x_t\) is an input to the computation of \(x_{t+1}\). No amount of GPU parallelism removes that serial dependency. The hardware is busy but starved: at decode time you are matrix-vector limited, reading the entire model’s weights from memory to produce a single token (see The Anatomy of LLM Inference: Prefill, Decode & The KV Cache and Inference Economics: Latency, Throughput & Cost).

This chapter is about a different bet: non-autoregressive (NAR) language models, and in particular the discrete diffusion / masked diffusion family that has, since roughly 2024–2025, produced the first genuinely competitive non-AR LLMs — LLaDA, Dream, the commercial Mercury system from Inception Labs, and, in a signal that the idea has reached the frontier labs, Google DeepMind’s experimental Gemini Diffusion. Instead of emitting tokens left-to-right, these models start from a fully masked sequence and iteratively denoise all positions in parallel, refining the whole sequence over a small number of steps. The promise is a fundamentally different point on the latency–quality curve: tokens-per-second numbers that can be several times higher than AR models of comparable size, plus the ability to reason bidirectionally and fill text in any order.

We will develop the absorbing-state diffusion objective from first principles, contrast it sharply with next-token prediction (covered in The Pretraining Objective & Loss), implement a tiny masked-diffusion sampler from scratch, and then engineer our way toward the production tricks — semi-autoregressive block diffusion, KV-cache reuse, remasking schedules, and flexible-length generation — that make these systems practical. The mathematical machinery overlaps with continuous image diffusion (Diffusion Models & Generative Modeling (Breadth)), but the discrete-token setting has its own elegant simplifications that are worth seeing carefully.

Why Non-Autoregressive? The Latency Argument¶

Let us quantify the wall we are trying to climb over. In the decode phase, an AR transformer produces one token per forward pass. The wall-clock time to generate \(L\) tokens is

where \(t_{\text{step}}\) is the per-step latency, dominated on modern GPUs by the time to stream the model weights from HBM into the compute units. For a dense model with \(P\) parameters at precision \(b\) bytes/param and memory bandwidth \(\beta\) bytes/s, the floor is roughly \(t_{\text{step}} \gtrsim P b / \beta\) — and crucially this floor is independent of how clever your kernel is, because decode is memory-bandwidth bound, not compute bound.

A non-autoregressive diffusion model instead runs \(N\) denoising steps, where each step is one forward pass over the entire length-\(L\) sequence in parallel, and ideally \(N \ll L\):

Here \(t_{\text{step}}'\) is somewhat larger than \(t_{\text{step}}\) because each step processes all \(L\) positions at once (a compute-bound batched operation that uses the hardware far more efficiently), but the win is that \(N\) can be 16, 32, or 64 instead of \(L = 1000\). If a diffusion model produces acceptable text in \(N = 64\) steps for a 512-token output, it has done \(64\) sequential forward passes instead of \(512\) — an \(8\times\) reduction in serial depth, even before accounting for the fact that each diffusion forward pass keeps the matrix multiply units saturated.

There is a second, qualitatively different motivation. AR models are causally blind to the future: when predicting \(x_t\) they cannot see \(x_{>t}\). This makes some tasks awkward — infilling, editing, satisfying global constraints, and the famous reversal curse (a model trained on “A is B” often fails to answer “what is B? → A”). A bidirectional denoiser conditions every position on every other position at every step, which structurally sidesteps these issues. Non-AR models are not just “AR but faster”; they have a different inductive bias.

Discrete & Masked Diffusion From First Principles¶

Continuous diffusion (DDPM, score-based models) gradually adds Gaussian noise to a real-valued vector and learns to reverse it. Text is discrete — there is no meaningful “add a little Gaussian noise to the token cat.” We need a corruption process that lives on a finite vocabulary. The family that has proven both simple and powerful is absorbing-state (masked) diffusion.

The forward (corruption) process¶

Fix a special absorbing token [MASK], with index \(m\), that is not a normal vocabulary token. The forward process is defined on a continuous time variable \(t \in [0, 1]\). At \(t = 0\) the sequence is clean data; at \(t = 1\) it is entirely [MASK]. For each token independently, define a monotone masking schedule \(\alpha_t\) with \(\alpha_0 = 1\) and \(\alpha_1 = 0\). Conditioned on the clean token \(x_0^i\) at position \(i\), the corrupted token at time \(t\) is

So \(\alpha_t\) is literally the probability that a given token survives unmasked at time \(t\). The absorbing property is what makes this clean: once a token becomes [MASK] it stays [MASK] for all later times — the only transition is data \(\to\) mask, never mask \(\to\) data and never data \(\to\) different-data. This is far simpler than the general discrete-diffusion transition matrices of Austin et al.’s D3PM, where any token can flip to any other.

The reverse (denoising) process and the objective¶

The model is a bidirectional transformer (an encoder-style stack with no causal mask — see Architecture Variants: Encoder-Decoder, Decoder-Only & Prefix-LM) that takes a partially-masked sequence \(x_t\) and predicts a distribution over the clean token at every masked position simultaneously. Call this \(p_\theta(x_0^i \mid x_t)\).

The remarkable simplification of absorbing diffusion is that, after the full variational derivation, the training loss collapses to a weighted cross-entropy on the masked positions only — essentially a continuous-time generalization of BERT’s masked language modeling. For a clean sequence \(x_0\), sample a time \(t \sim \mathcal{U}(0,1)\), mask each token independently with probability \(1 - \alpha_t\) to get \(x_t\), and minimize

Read this carefully, because three ideas are packed into it:

1. Only masked positions contribute. The sum runs over \(i\) where \(x_t^i\) is [MASK]. Unmasked positions provide context but no loss — the model is rewarded only for recovering what was hidden.
2. The weight \(1/(1-\alpha_t)\) corrects for the masking rate. When \(t\) is near 1, almost everything is masked, \(1-\alpha_t \approx 1\), and the model must reconstruct from almost nothing — a hard denoising problem. When \(t\) is near 0, only a few tokens are masked and the weight is large, so each rare masked position counts heavily. This weighting is exactly what makes the masked-LM loss a valid (upper bound on the) negative log-likelihood, i.e. a true generative objective, not just a representation-learning trick. The MDLM and RADD analyses (Sahoo et al.; Ou et al., 2024) show this weighted form is a tight Evidence Lower Bound. (Strictly, the continuous-time ELBO carries the weight \(-\alpha_t'/(1-\alpha_t)\), where \(\alpha_t' = \mathrm{d}\alpha_t/\mathrm{d}t\). With the standard linear schedule \(\alpha_t = 1 - t\) we get \(\alpha_t' = -1\) and the weight collapses to \(1/(1-\alpha_t) = 1/t\), the form we use throughout this chapter. A pleasant result of the MDLM/Shi et al. analyses is that the ELBO’s value is invariant to the choice of monotone schedule — unlike continuous diffusion, where the noise schedule is a hyperparameter you tune — so the schedule only affects gradient variance, not the objective being optimized.)
3. There is no left-to-right ordering. Unlike next-token prediction, where the chain rule dictates the factorization order, here the model learns to fill any subset of positions given any other subset. This is the source of the bidirectional advantage.

Contrast with next-token prediction, where the loss is \(-\sum_t \log p_\theta(x_t \mid x_{<t})\) over every position with a strict causal mask. AR predicts the next token from a left context; masked diffusion predicts a random subset from a bidirectional context. AR gives you an exact autoregressive likelihood and trivially correct sampling order; diffusion gives you parallelism and bidirectionality at the cost of an approximate, order-agnostic factorization.

Time conditioning, or the lack of it¶

Continuous diffusion models almost always feed the timestep \(t\) into the network (via sinusoidal embeddings added to the input). A pleasant surprise in the discrete-masked setting is that \(t\) is largely redundant: the number of [MASK] tokens in the input already tells the model roughly where it is in the denoising process. Several strong masked-diffusion LLMs (including LLaDA) drop explicit time conditioning entirely and rely on the mask count as an implicit clock. This is one fewer thing to get right and lets the architecture stay a vanilla bidirectional transformer.

The training step, end to end¶

The objective above is short enough to implement in a dozen lines, and — this is the practically important point — it bolts onto the GPT you already built. Everything from Building a GPT From Scratch (nanoGPT-style) and the Stack-100M block in The Stack-100M Architecture carries over unchanged: RoPE, RMSNorm, SwiGLU, GQA, the tokenizer, the packed uint16 shards, the optimizer, the WSD schedule. Exactly three things change.

1. Delete the causal mask. F.scaled_dot_product_attention(q, k, v, is_causal=False) with no additive mask. That single flag converts the decoder into a bidirectional denoiser.
2. Add one vocabulary entry for [MASK] — one extra row in the token embedding and in lm_head (a tied-embedding model gets both for free). Real checkpoints do this: LLaDA reserves a dedicated mask id rather than reusing an existing special token.
3. Change the loss, and drop the shift. This is the most common bug in a first implementation. An AR model’s logits at position \(i\) predict token \(i+1\), so training shifts targets by one. A diffusion denoiser predicts the token at position \(i\) from the corrupted sequence — no shift. If you copy an AR training loop verbatim you get a model that trains to a plausible-looking loss and generates garbage.

import torch
import torch.nn.functional as F

def masked_diffusion_loss_batch(model, x0, mask_id, eps=1e-3):
    """
    Batched absorbing-diffusion training loss (the LLaDA/MDLM form).

    model: bidirectional transformer, (B, L) ids -> (B, L, V) logits. NO causal mask.
    x0:    (B, L) clean token ids from your ordinary packed pretraining shards.
    Linear schedule alpha_t = 1 - t, so P(mask) = t and the ELBO weight is 1/t.
    Returns a scalar ready for .backward().
    """
    B, L = x0.shape
    # One t PER SEQUENCE, not per token: the forward process draws a single noise
    # level for the sequence and then masks its tokens i.i.d. at that level.
    # Clamp t away from 0 so the 1/t weight cannot explode a single microbatch.
    t = torch.rand(B, 1, device=x0.device) * (1.0 - eps) + eps      # (B, 1)
    masked = torch.rand(B, L, device=x0.device) < t                 # broadcast -> (B, L)

    # A row with zero masked positions contributes no gradient; force one.
    # (Bias is negligible: P(no mask) = (1-t)^L is tiny for L in the thousands.)
    masked[:, 0] |= ~masked.any(dim=1)

    xt = torch.where(masked, torch.full_like(x0, mask_id), x0)
    logits = model(xt)                                              # (B, L, V)

    # Cross-entropy on MASKED positions only. NOTE: no shift -- logits[b, i]
    # predicts x0[b, i], unlike the next-token objective.
    ce = F.cross_entropy(logits[masked], x0[masked], reduction="none")   # (n_masked,)

    # Weight each masked position by 1/t of ITS OWN sequence, then normalize by
    # B*L (per-token normalization) so the loss is comparable across batches.
    w = t.expand(B, L)[masked]                                      # (n_masked,)
    return (ce / w).sum() / (B * L)

Two consequences worth internalizing. First, the number reported by this loss is not a perplexity you can compare to an AR run — it is a Monte-Carlo estimate of an ELBO, and its variance is high early in training because a single \(t\) per sequence is a coarse estimator (antithetic or low-discrepancy sampling of \(t\) across the batch is the standard variance reduction). Second, the diffusion model sees strictly less supervision per forward pass: an AR pass gets a gradient at all \(L\) positions, a diffusion pass only at the \(\approx tL\) masked ones (in expectation, half of them). Empirically this is a real part of why from-scratch masked-diffusion LLMs need more tokens or more epochs to match AR quality at equal parameters, and it is the single strongest argument for the AR-to-diffusion adaptation route that Dream took.

For SFT the recipe is a small variant: never mask the prompt. Concatenate prompt and response, corrupt only the response positions, and compute the loss only there — the prompt is permanently-clean context, exactly as is_prompt is clamped in the sampler below. That is the diffusion analogue of the prompt-masked SFT loss in Supervised Fine-Tuning & Instruction Tuning.

Concretely, this is a weekend project on top of the capstone: keep the corpus and packing of Data: Sourcing, Filtering, Dedup, Tokenize & Pack ~20B Tokens and the training loop of The Pretraining Run: A Complete Single-GPU Training Loop exactly as they are, set is_causal=False, grow the vocab by one, and swap the shifted cross-entropy for masked_diffusion_loss_batch. You get a ~100M-parameter masked-diffusion LM you can sample with the code in the next section — an ablation against the AR Stack-100M that costs one extra run and teaches more about the two paradigms than any benchmark table. Budget more tokens than the AR baseline for the reason just given, and compare on generation quality rather than on loss, since the two numbers measure different things.

Iterative Parallel Denoising: A Tiny Sampler¶

Time to make this concrete. The inference loop for masked diffusion is conceptually simple: start with everything masked, and repeatedly (a) predict all masked positions in parallel, (b) commit some of those predictions by unmasking them, and © leave the rest masked for the next round. The art is entirely in which positions to commit each step — the remasking schedule.

Here is a complete, from-scratch sampler. The “model” is abstracted behind a denoiser callable so the loop is crystal clear; in practice it is a bidirectional transformer returning logits of shape (L, V).

import torch
import torch.nn.functional as F

MASK_ID = 0  # reserve index 0 in the vocab for the [MASK] absorbing token

@torch.no_grad()
def masked_diffusion_sample(
    denoiser,            # callable: (LongTensor[L]) -> FloatTensor[L, V] logits
    L: int,              # sequence length to generate
    num_steps: int,      # number of denoising iterations N (the serial depth)
    vocab_size: int,
    temperature: float = 0.0,   # 0.0 = greedy/argmax per position
    prompt: torch.Tensor = None,   # optional LongTensor of clamped prefix tokens
    clamp_mask: torch.Tensor = None,    # optional (L,) bool mask of arbitrary clamped positions
    clamp_values: torch.Tensor = None,  # (L,) values at those positions; used with clamp_mask
    device: str = "cpu",
):
    """
    Absorbing-state masked diffusion sampler.

    Strategy ("confidence-based remasking", a la LLaDA's low-confidence remasking):
    at each step we predict ALL masked positions, but only *keep* (unmask) the
    most confident predictions, sized so that the fraction of masked tokens
    follows a linear schedule from 1.0 down to 0.0 over num_steps. Less confident
    positions are returned to [MASK] and revisited in later steps.
    """
    # 1. Start fully masked.
    x = torch.full((L,), MASK_ID, dtype=torch.long, device=device)

    # 2. Clamp fixed context (these positions are never masked and never predicted).
    #    `clamp_mask`/`clamp_values` clamp an ARBITRARY subset of positions (e.g. a
    #    prefix AND a suffix for infilling); `prompt` is the common special case of
    #    clamping only a contiguous prefix.
    is_prompt = torch.zeros(L, dtype=torch.bool, device=device)
    if clamp_mask is not None:
        is_prompt = clamp_mask.to(device)
        x[is_prompt] = clamp_values.to(device)[is_prompt]
    elif prompt is not None:
        x[: prompt.numel()] = prompt.to(device)
        is_prompt[: prompt.numel()] = True

    # 3. Denoising schedule: target number of STILL-masked tokens after step k.
    #    Linear from "all generatable positions masked" down to 0.
    n_gen = int((~is_prompt).sum().item())   # positions we actually generate
    # masked_target[k] = how many gen-positions remain masked AFTER step k+1
    masked_target = [
        round(n_gen * (1.0 - (k + 1) / num_steps)) for k in range(num_steps)
    ]

    for step in range(num_steps):
        masked = (x == MASK_ID) & (~is_prompt)   # positions still to fill
        if masked.sum() == 0:
            break

        logits = denoiser(x)                     # (L, V) — full bidirectional pass
        logits[:, MASK_ID] = -float("inf")       # never predict [MASK] itself

        if temperature > 0.0:
            probs = F.softmax(logits / temperature, dim=-1)
            pred = torch.multinomial(probs, num_samples=1).squeeze(-1)  # (L,)
            conf = probs.gather(-1, pred.unsqueeze(-1)).squeeze(-1)     # (L,)
        else:
            probs = F.softmax(logits, dim=-1)
            conf, pred = probs.max(dim=-1)       # greedy + its confidence

        # Candidate fill: tentatively set every masked position to its prediction.
        x_candidate = x.clone()
        x_candidate[masked] = pred[masked]

        # Decide how many of the currently-masked positions to KEEP unmasked.
        n_keep = int(masked.sum().item()) - masked_target[step]
        n_keep = max(0, n_keep)

        # Rank masked positions by confidence; keep the top-n_keep, remask the rest.
        conf_masked = conf.clone()
        conf_masked[~masked] = -float("inf")     # only compete among masked positions
        if n_keep > 0:
            keep_idx = torch.topk(conf_masked, n_keep).indices
            new_x = x.clone()
            new_x[masked] = MASK_ID              # provisionally remask all
            new_x[keep_idx] = x_candidate[keep_idx]  # commit the confident ones
            x = new_x
        # if n_keep == 0 we commit nothing this step (rare; only at the very start)

    # Final cleanup: fill any leftover masks greedily (last step should handle this).
    leftover = (x == MASK_ID) & (~is_prompt)
    if leftover.any():
        logits = denoiser(x)
        logits[:, MASK_ID] = -float("inf")
        x[leftover] = logits[leftover].argmax(dim=-1)

    return x

The whole behavior of the model lives in how n_keep is chosen each step. A few canonical schedules:

• Random remasking. Pick the positions to keep uniformly at random. Simple, unbiased, but wastes steps committing to low-confidence guesses early.
• Confidence-based (low-confidence remasking). Keep the highest-confidence predictions, return the rest to [MASK] (what the code above does). This is LLaDA’s default and works well because the model “locks in” the tokens it is sure about and keeps reconsidering the uncertain ones with more context each round.
• Top-\(k\) / greedy decoding orders. Variants that commit a fixed number of tokens per step, or use entropy rather than max-prob as the confidence measure.

The crucial subtlety is that already-committed tokens become context for the next step. When the model unmasks “The capital of France is” it makes the later positions far easier to predict — the bidirectional pass now sees both left and right context flowing into each remaining mask. This is the iterative-refinement engine: each round of commitments sharpens the conditional distribution for everything still masked.

How many steps do you actually need?¶

This is the central quality–latency knob. With \(N = L\) steps and a one-token-per-step schedule, masked diffusion essentially reduces to an (order-flexible) autoregressive model and matches AR quality — but you have thrown away the speed advantage. With \(N\) very small (say 8), you get blazing speed but the conditional-independence error bites and quality drops. Production systems live in the interesting middle: enough steps that each commits a small block of tokens, few enough that serial depth stays well below \(L\). The empirical finding across LLaDA and Dream is that you can often reach AR-comparable quality at \(N\) on the order of \(L/4\) to \(L/2\), and trade further quality for speed below that.

Semi-Autoregressive Block Diffusion¶

Pure diffusion over a fixed-length block has two weaknesses that show up immediately in production: (1) it commits to a fixed sequence length up front, which is unnatural for open-ended generation, and (2) it cannot reuse a KV cache across steps, because every position is recomputed every step. Block diffusion (Arriola et al., 2025) is the hybrid that fixes both by interpolating between autoregressive and diffusion modeling.

The idea: split the sequence into contiguous blocks of \(B\) tokens each. Generate the blocks autoregressively (left to right, one block after another), but generate the tokens within each block by diffusion (parallel denoising). Formally, with blocks \(x^{(1)}, x^{(2)}, \dots\),

where each block factor is itself a small masked-diffusion model conditioned on all previously finalized blocks. Set \(B = 1\) and you recover a pure autoregressive model; set \(B = L\) (one block) and you recover pure diffusion. Block diffusion is the dial between them.

   Block 1            Block 2            Block 3
 [denoise B tok] -> [denoise B tok] -> [denoise B tok] -> ...
   (parallel)         (parallel)         (parallel)
   |__________________|__________________|
        previous blocks are FINALIZED and cached in the KV cache;
        the current block attends to them causally (block-causal mask)

This buys three things at once:

• KV-cache reuse. Because earlier blocks are finalized before the current block starts, their keys and values are fixed. We compute them once and cache them, exactly like AR decoding (Prefix Caching & KV-Cache Reuse, PagedAttention & KV-Cache Memory Management). Within-block denoising recomputes only the \(B\) positions of the current block — a huge saving versus pure diffusion, which recomputes all \(L\) positions every step.
• Flexible / arbitrary length. Generation proceeds block by block until an end-of-sequence condition fires, so you are no longer locked into a length chosen before you start. You can keep emitting blocks until the model produces [EOS], just like AR.
• Parallelism where it pays. Inside a block you get the diffusion speedup (parallel denoising of \(B\) tokens in a few steps); across blocks you get the AR structure that captures long-range left-to-right dependencies and gives you the cache.

The attention mask is the key implementation detail. Within a block, attention is bidirectional (every position sees every other position in the block — that is what gives diffusion its power). Across blocks, attention is causal at block granularity: tokens in block \(b\) can attend to all tokens of blocks \(1..b\) but not \(b+1..\). This block-causal mask is what makes the KV cache valid.

def block_causal_mask(L: int, block_size: int) -> torch.Tensor:
    """
    Build an additive attention mask for block diffusion.
      - bidirectional WITHIN a block (no causal restriction)
      - causal ACROSS blocks (block b sees blocks 1..b only)
    Returns (L, L) with 0.0 where attention is allowed, -inf where blocked.
    """
    idx = torch.arange(L)
    block_of = idx // block_size                       # (L,) block index per position
    # allowed if key's block <= query's block
    allowed = block_of[None, :] <= block_of[:, None]   # (L_query, L_key) boolean
    mask = torch.zeros(L, L)
    mask.masked_fill_(~allowed, float("-inf"))
    return mask

Training block diffusion has one wrinkle the sampler hides. Each block must be denoised while attending to the clean (finalized) versions of all earlier blocks, but the earlier blocks in your training batch are themselves noised. BD3-LM’s solution is a vectorized two-forward-pass scheme: one pass over the clean sequence under the block-causal mask to materialize keys and values for every block, then a second pass over the noised sequence in which each block’s queries attend to those cached clean keys/values for blocks \(<b\) and to its own noised block bidirectionally. That keeps training \(O(1)\) passes per batch instead of one pass per block. The paper also reports that the naive \(t \sim \mathcal{U}(0,1)\) draw has high gradient variance at block granularity and uses clipped noise schedules — sampling \(t\) from a narrower interval \([\beta, \omega] \subset (0,1)\) tuned per block size — to reduce it, which is the main reason BD3-LM reaches better perplexities than earlier discrete-diffusion models.

Block diffusion is the architecture that most directly maps onto the existing high-performance AR serving stack: you keep continuous batching, paged KV cache, and prefix caching, and you simply swap the inner decode kernel from “one token per step” to “a block of tokens per few diffusion steps.” This is a large part of why the first commercially viable diffusion LLM, Mercury, is built on this semi-autoregressive structure rather than pure parallel diffusion.

Production Systems & Inference Economics¶

We now have the pieces — masked-diffusion objective, iterative denoising, block structure, KV reuse — to understand the real systems and their economics.

The landscape¶

• LLaDA (Nie et al., 2025) is the proof of concept at scale: an 8B-parameter masked-diffusion LLM trained from scratch with the absorbing objective, with an instruction-tuned chat variant. Its headline result is that a pure (non-block) masked-diffusion model can be competitive with similarly-sized autoregressive LLaMA-class models on standard benchmarks — the first time a from-scratch diffusion LLM closed most of that gap. It also concretely demonstrated the bidirectional/reversal advantage (below).
• Dream (2025) is a 7B masked-diffusion LLM that, rather than training from scratch, initializes from an existing autoregressive checkpoint (a Qwen2.5-7B base model) and adapts it to the diffusion objective — a clever way to reuse the enormous compute already spent on AR pretraining, at a small fraction of from-scratch cost. It uses context-adaptive token-level noise and confidence-based decoding orders to push quality. This adaptation route is the one to copy if you want a diffusion LM without a pretraining budget: take any open base model, drop the causal mask, add a mask token, and continue training under the loss above.
• Mercury (Inception Labs, 2025) is the commercial diffusion-LLM system, marketed around coding (Mercury Coder) and built on a block-diffusion-style architecture for flexible length and KV reuse. Its technical report puts Mercury Coder Mini and Small at roughly 1109 and 737 tokens/sec on an NVIDIA H100 — up to about 10× the throughput of comparable speed-optimized autoregressive code models at similar quality — by exploiting parallel decode. (The mechanism — parallel denoising of blocks with cache reuse — is the durable takeaway; treat exact numbers as hardware- and configuration-dependent.)
• Gemini Diffusion (Google DeepMind, 2025) and Seed Diffusion (ByteDance, 2025) are the clearest signal that the approach has reached the frontier and the speed frontier respectively: the former is an experimental text-diffusion model reported at on the order of 1,400 tokens/sec while matching Google’s fast autoregressive models on coding and math benchmarks; the latter is a code-focused discrete-diffusion model reporting above 2,000 tokens/sec on datacenter GPUs. Their existence — more than any single vendor figure — is the evidence that parallel-decode diffusion is now a mainstream research direction rather than an academic curiosity.

The open-source stack: what you actually run¶

The from-scratch code above is the mechanism; here is the tooling that implements it. Open-weight diffusion LLMs ship on the HuggingFace Hub with custom modeling code, so they load through transformers with trust_remote_code=True — but note that model.generate() in its usual AR form does not apply, because the generation loop is a denoiser, not a next-token loop. The repos therefore ship their own generate function implementing the remasking schedule.

# Loading a released masked-diffusion LLM through HuggingFace transformers.
# Repo ids move; confirm the exact one on the Hub before depending on it.
import torch
from transformers import AutoModel, AutoTokenizer

repo = "GSAI-ML/LLaDA-8B-Instruct"          # LLaDA's instruction-tuned release
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModel.from_pretrained(
    repo, trust_remote_code=True, torch_dtype=torch.bfloat16
).to("cuda").eval()

# The absorbing token is a REAL vocabulary entry in released checkpoints -- read
# it from the config/tokenizer instead of hardcoding index 0 as our toy sampler did.
mask_id = getattr(model.config, "mask_token_id", None) or tok.mask_token_id

# From here the loop is the one we wrote by hand: start from a fully-[MASK]ed
# answer block after the prompt, forward the whole sequence, commit the most
# confident positions, repeat. The reference `generate.py` in the model's repo
# implements exactly this, plus semi-autoregressive block decoding.

The pieces of the ecosystem worth knowing by name:

• Reference training code. The Kuleshov group’s mdlm and bd3lms repositories are the canonical implementations of the simplified masked-diffusion objective and of block diffusion respectively — Hydra-configured PyTorch Lightning training on OpenWebText/LM1B, and the cleanest place to read a correct ELBO implementation. The LLaDA authors’ repo (ML-GSAI/LLaDA) is the reference for the 8B-scale sampler, including low-confidence remasking and semi-autoregressive block decoding.
• Inference acceleration. Fast-dLLM (NVIDIA, 2025) is the notable training-free accelerator: it approximates a KV cache for block-wise bidirectional attention (caching prefix and suffix blocks) and adds confidence-thresholded parallel decoding — commit every position whose probability exceeds a threshold in one shot, rather than a fixed count — reporting large end-to-end speedups on LLaDA and Dream with minimal quality loss. The mechanism is the durable part: adaptive commit counts beat the fixed linear schedule our toy sampler uses.
• Serving stacks. The mature engines (vLLM: Architecture, PagedAttention & Internals, SGLang: RadixAttention & Structured Programs, TensorRT-LLM, TGI & Other Serving Stacks) are built around the AR decode loop — a per-request single-token step, paged KV cache, continuous batching. First-class diffusion-LM support is emerging rather than settled; check the current release notes rather than assuming parity. In practice today you serve a diffusion LLM from the model repo’s own loop wrapped in torch.compile and CUDA graphs (Kernel Fusion, torch.compile, CUDA Graphs & Compilers), and you get continuous batching only if you write it. Quantization (GPTQ/AWQ/bitsandbytes) applies unchanged, since it operates on the weights and is indifferent to the decode loop.
• Evaluation. lm-evaluation-harness (Building Eval Harnesses) mostly works, with one caveat that matters: its loglikelihood request type assumes an exact AR log-probability. A diffusion model can only supply a Monte-Carlo ELBO estimate, so multiple-choice tasks scored by log-likelihood are noisy and not strictly comparable to AR numbers. Generative tasks (generate_until) are the honest comparison; wire the model’s denoising generate into a custom LM subclass.

Where the speed actually comes from — and where it doesn’t¶

The diffusion speedup is a serial-depth win, not a FLOPs win. A diffusion model often does more total floating-point work than an AR model (it recomputes positions across steps), but it organizes that work into fewer sequential, more parallel steps. This has two important consequences for the economics (Inference Economics: Latency, Throughput & Cost):

1. Single-stream latency: diffusion wins big. When you have one request and want the answer fast — interactive coding, low-batch agentic loops — cutting serial depth from \(L\) to \(N\) is a direct latency reduction. This is diffusion’s home turf.
2. High-throughput, large-batch serving: the gap narrows. AR decoding is memory-bandwidth bound at small batch but becomes compute-efficient at large batch, because continuous batching packs many requests’ single-token steps into one fat matrix multiply that saturates the GPU. Diffusion already saturates the GPU per step, so it has less headroom to gain from batching. At very high concurrency, an AR system’s aggregate tokens/sec can rival or exceed a diffusion system’s, because the diffusion model’s extra FLOPs per token now compete for the same saturated compute. Diffusion’s advantage is largest at low-to-moderate batch and long single outputs.

The honest summary: diffusion LLMs are not a free lunch that beats AR everywhere. They occupy a different region of the latency–throughput–quality surface — superb single-stream latency and bidirectional capability, in exchange for an approximate factorization, more total compute per token, and (today) a smaller ecosystem of training data, tooling, and post-training recipes than the mature AR stack.

Sampling controls carry over — mostly¶

Temperature, top-\(k\), and top-\(p\) (Sampling Strategies & Decoding Algorithms) all apply per position inside a denoising step, exactly as in AR. What is new is the remasking schedule as a first-class decoding hyperparameter: how many tokens to commit per step, and by what confidence criterion. Constrained/structured generation (Structured & Constrained Generation) is actually more natural in some respects — because the model sees the whole sequence, you can pin known tokens (a closing brace, a required JSON key) as permanently-unmasked context and let denoising fill around them, which is exactly the infilling capability AR models lack.

The Bidirectional Advantage: Infilling, Constraints & the Reversal Curse¶

The most intellectually interesting reason to care about non-AR LLMs is not speed — it is bidirectionality. Because a masked-diffusion model conditions every position on every other position, it can do things that are structurally hard for a left-to-right model.

Infilling and editing. Given a document with a hole in the middle, an AR model must either be specially trained with fill-in-the-middle objectives or awkwardly re-prompted. A diffusion model treats infilling as its native operation: clamp the known prefix and suffix as unmasked context, mask the hole, and denoise. The hole is filled conditioned on both sides, which is exactly what coherent editing requires.

Global constraints. Tasks like “write a sentence that ends with this exact word” or “produce code with this signature and this return statement” require coordinating the beginning and end of the output. Left-to-right generation can paint itself into a corner; bidirectional denoising can place the constrained tokens first and grow the rest around them.

The reversal curse. AR models trained on “A is B” famously struggle to answer the reversed query “who/what is B?” because the training gradient only ever flowed in the A→B direction. A bidirectional diffusion model is trained to reconstruct any masked subset from any context, so it sees both directions during training and is markedly more robust to reversal. LLaDA’s authors highlighted exactly this: on reversal-style tasks the masked-diffusion model can outperform a same-scale autoregressive baseline, sometimes by a wide margin, precisely because its objective is not directionally biased.

# Infilling with the same sampler: clamp prefix AND suffix, denoise the middle.
def infill_example(denoiser, prefix_ids, suffix_ids, hole_len, vocab_size):
    L = len(prefix_ids) + hole_len + len(suffix_ids)
    x = torch.full((L,), MASK_ID, dtype=torch.long)
    is_clamped = torch.zeros(L, dtype=torch.bool)

    # left context
    x[: len(prefix_ids)] = torch.tensor(prefix_ids)
    is_clamped[: len(prefix_ids)] = True
    # right context (note: this is the future, which AR cannot use!)
    x[len(prefix_ids) + hole_len :] = torch.tensor(suffix_ids)
    is_clamped[len(prefix_ids) + hole_len :] = True

    # Reuse the diffusion loop, but treat ALL clamped positions (prefix AND
    # suffix) like a "prompt" so they are never masked and never overwritten.
    # The hole is denoised conditioned on BOTH sides simultaneously.
    return masked_diffusion_sample(
        denoiser, L=L, num_steps=hole_len, vocab_size=vocab_size,
        clamp_mask=is_clamped, clamp_values=x,
    ), is_clamped

This is the capability that AR models have to bolt on and diffusion models get for free, and it is the strongest argument that non-AR LMs are a genuinely different tool rather than a faster clone.

Practicalities, Limits & When to Reach for This¶

A few engineering realities to keep the picture honest:

• Likelihood is a bound, not exact. AR models give you an exact log-likelihood (useful for perplexity, ranking, watermarking). Masked diffusion gives a variational bound; reported perplexities use the ELBO and are not directly comparable to AR perplexity. Be careful when comparing benchmark tables across the two paradigms.
• Post-training is younger. The instruction-tuning, RLHF, and preference-optimization stack (Part V) was built around AR generation. SFT ports cleanly (mask only the response, as above). The hard part is RL and preference optimization, because PPO/GRPO/DPO all need a per-sequence log-probability \(\log \pi_\theta(y \mid x)\) and a well-defined generation order (Policy Gradients & PPO for Language Models, GRPO, RLOO & Critic-Free RL, Direct Preference Optimization & Its Variants) — and a diffusion LM has neither. The workaround in the literature is to substitute the ELBO estimate for the exact log-prob: d1 (2025) introduces diffu-GRPO, a GRPO variant that uses a one-forward-pass, per-token ELBO surrogate for the policy log-prob so that RLVR (RL with Verifiable Rewards (RLVR) & The Reasoning Recipe) can run on a masked-diffusion LLM; LLaDA 1.5 attacks the same problem on the preference side with variance-reduced preference optimization (VRPO), whose whole content is reducing the variance of that Monte-Carlo ELBO so DPO-style gradients are not swamped by noise. Both are the same lesson: the missing exact likelihood is the central obstacle to diffusion post-training, and every method here is a variance-reduction story around estimating it.
• Length handling. Pure diffusion needs a length up front; block diffusion fixes this but adds the block-size hyperparameter and a block-causal mask. Either way, length control is more involved than AR’s natural “generate until [EOS].”
• Tooling maturity. vLLM, SGLang, speculative decoding (Speculative Decoding: Draft Models, Medusa, EAGLE & Lookahead), and the entire kernel ecosystem are tuned for AR decode. Diffusion-specific serving is catching up but is not yet at parity.

So when should you reach for a diffusion LLM today? When single-stream latency dominates your objective (interactive coding assistants, low-concurrency agents), when the task is natively bidirectional (infilling, editing, constraint satisfaction), or when you want robustness to reversal-style queries. When you need exact likelihoods, the mature post-training/RL stack, maximal throughput at very high batch, or simply the lowest-risk path, autoregressive transformers remain the default — for now.

Diffusion and non-AR language models are best understood as a third major branch of the sequence-modeling tree, alongside the attention transformers of this Part and the recurrent/SSM alternatives of Beyond Attention: SSMs, Mamba, RWKV & Linear Attention. All three are attempts to escape the \(O(L)\)-serial, \(O(L^2)\)-attention cost of the vanilla transformer; diffusion attacks the serial-depth axis specifically, and pays in approximate factorization and extra compute for the privilege.

Key Takeaways¶

Further Reading¶

• Austin, J., et al. (2021). Structured Denoising Diffusion Models in Discrete State-Spaces (D3PM). NeurIPS 2021.
• Lou, A., Meng, C., & Ermon, S. (2023). Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution (SEDD). ICML 2024.
• Sahoo, S., et al. (2024). Simple and Effective Masked Diffusion Language Models (MDLM). NeurIPS 2024.
• Shi, J., et al. (2024). Simplified and Generalized Masked Diffusion for Discrete Data. NeurIPS 2024.
• Ou, J., et al. (2024). Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data (RADD).
• Nie, S., et al. (2025). Large Language Diffusion Models (LLaDA).
• Arriola, M., et al. (2025). Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models (BD3-LM). ICLR 2025.
• Inception Labs (2025). Mercury: Ultra-Fast Language Models Based on Diffusion. arXiv:2506.17298 — commercial diffusion LLM (incl. Mercury Coder).
• Google DeepMind (2025). Gemini Diffusion — experimental frontier-lab text-diffusion model.
• Song, Y., et al. (2025). Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference. arXiv:2508.02193.
• Gong, S., et al. (2025). Dream 7B — autoregressive-to-diffusion adaptation from a Qwen2.5-7B base checkpoint.
• Zhao, S., et al. (2025). d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning — diffu-GRPO.
• Zhu, F., et al. (2025). LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Models.
• Wu, C., et al. (2025). Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding. NVIDIA.
• Reference code: kuleshov-group/mdlm, kuleshov-group/bd3lms, ML-GSAI/LLaDA.

Exercises¶

1. Continuous image-diffusion models almost always inject the timestep \(t\) into the network (via sinusoidal embeddings). Several strong masked-diffusion LLMs, including LLaDA, drop explicit time conditioning entirely and still generate coherent text. What implicit signal lets the network infer “where it is” in the denoising process without being told \(t\), and why does that signal exist specifically in the absorbing-state formulation?

2. A dense \(3\text{B}\)-parameter model runs in bf16 (\(b = 2\) bytes/param) on a GPU with HBM bandwidth \(\beta = 2.0\) TB/s. Using the decode-step memory floor \(t_{\text{step}} \gtrsim P b / \beta\), compute, for a single sequence of \(L = 256\) output tokens:

(a) the autoregressive wall-clock time and tokens/sec;

(b) the diffusion wall-clock time and tokens/sec for \(N = 32\) denoising steps, taking \(t_{\text{step}}' = 9\) ms per (compute-heavier) parallel step;

© the single-stream latency speedup of diffusion over AR.

3. Consider a block-diffusion model generating \(L = 256\) tokens with block size \(B = 16\) and \(N_B = 4\) denoising steps per block. Compute (a) the number of blocks, (b) the total number of block-denoising forward passes and the resulting serial depth, © how the serial depth compares to the equivalent AR model, and (d) the total position-evaluations by the denoiser versus AR. What does the comparison of © and (d) tell you about the nature of the diffusion speedup?

4. Within a single denoising step the model predicts every masked position independently: \(p_\theta(x_0 \mid x_t) = \prod_i p_\theta(x_0^i \mid x_t)\). Explain why this factorization is wrong for natural language, what failure you would observe if you filled all masked positions in a single step, and how the iterative remasking loop recovers the missing structure. What choice of \(N\) (number of steps) makes the error vanish, and what is the cost?

5. The sampler in the chapter uses confidence-based remasking: each step it keeps the n_keep highest-confidence predictions. Modify the sampler to implement random remasking instead — keep a uniformly random subset of the currently-masked positions of size n_keep. Give the replacement code block and explain, in terms of the chapter’s argument, why confidence-based remasking generally produces better text at the same step budget.

    # Random remasking: keep a uniformly random subset of the masked
    # positions of size n_keep, remask the rest.
    masked_idx = masked.nonzero(as_tuple=True)[0]   # indices still masked
    if n_keep > 0:
        perm = masked_idx[torch.randperm(masked_idx.numel(), device=device)]
        keep_idx = perm[:n_keep]                     # random subset to commit
        new_x = x.clone()
        new_x[masked] = MASK_ID                      # provisionally remask all
        new_x[keep_idx] = x_candidate[keep_idx]      # commit the random ones
        x = new_x

6. Implement the training-side counterpart of the sampler: a function masked_diffusion_loss(denoiser, x0, alpha_fn) that returns one Monte-Carlo estimate of the absorbing-diffusion loss $$ \mathcal{L}(\theta) = \mathbb{E}{t \sim \mathcal{U}(0,1)}\, \mathbb{E} -\log p_\theta(x_0^i \mid x_t) \right] $$ for a single clean sequence }\left[ \frac{1}{1-\alpha_t} \sum_{i:\, x_t^i = \texttt{[MASK]}x0. Follow the chapter’s conventions (MASK_ID, a denoiser returning (L, V) logits). Then answer: which positions contribute to the loss, and what would go wrong if you dropped the \(1/(1-\alpha_t)\) weight?

import torch
import torch.nn.functional as F

MASK_ID = 0  # same absorbing-token convention as the sampler

def masked_diffusion_loss(denoiser, x0, alpha_fn, device="cpu"):
    """
    One Monte-Carlo estimate of the absorbing-diffusion loss for a clean
    sequence x0 (LongTensor[L]). `alpha_fn(t)` maps t in [0,1] to the
    survival probability alpha_t (alpha_0 = 1, alpha_1 = 0).
    `denoiser(xt)` returns (L, V) logits over clean tokens.
    """
    x0 = x0.to(device)
    L = x0.numel()

    # 1. Sample a time t ~ U(0,1) and its survival probability.
    t = torch.rand(1, device=device).item()
    alpha_t = alpha_fn(t)                          # scalar in [0, 1]

    # 2. Mask each token independently with prob (1 - alpha_t): data -> [MASK].
    keep = torch.rand(L, device=device) < alpha_t  # True = survives unmasked
    xt = x0.clone()
    xt[~keep] = MASK_ID
    masked = ~keep                                 # positions that got masked

    # Degenerate draw: nothing masked -> no loss signal this sample.
    if masked.sum() == 0:
        return torch.zeros((), device=device)

    # 3. One bidirectional pass predicting the clean token at every position.
    logits = denoiser(xt)                          # (L, V), no causal mask

    # 4. Cross-entropy on MASKED positions only (unmasked = context, no loss).
    ce = F.cross_entropy(logits[masked], x0[masked], reduction="sum")

    # 5. Weight by 1 / (1 - alpha_t) to make the masked-LM loss a valid
    #    negative-log-likelihood bound across all masking rates.
    return ce / (1.0 - alpha_t)