Diffusion and Non-Autoregressive Decoding
iterative unmasking; confidence-ranked commits; the speed-vs-coherence dial.
iterative unmasking; confidence-ranked commits; the speed-vs-coherence dial.
The Concept
Picture a crossword grid, empty. Every cell is a blank.
An autoregressive model solves it like someone forced to fill the cells in reading order, top-left to bottom-right. They write a final letter in each cell before moving to the next. No erasing, ever. That works, but it is a strange rule to impose on a crossword. The natural way to solve one is different. You scan all the clues. You fill in the three or four answers you are certain about. Those committed letters become constraints, and the crossings you were unsure about turn solvable. You do several passes. Each pass, you write down what you are now confident of and leave the doubtful cells blank for the next pass.
That second solver is a diffusion language model, or dLLM. The rest of this chapter is a careful account of how it turns a grid of blanks into a sentence. And why that buys speed, at a cost you can measure.
Start with the words. An autoregressive model, abbreviated AR, generates one token at a time in a fixed left-to-right order, each token conditioned only on the tokens before it. That is every GPT in this book. Non-autoregressive means the model does not commit to that order: it can produce or revise tokens at many positions in the same step. A dLLM is non-autoregressive.
Now the mechanism, in the exact form the research uses. The dLLMs this chapter is about (LLaDA, Mercury, Gemini Diffusion) are masked, or discrete, diffusion models. Here is the one sentence to hold onto before any code. In a text diffusion model the "noise" is a fully-masked sequence, not continuous random values. Generation runs by iteratively unmasking that sequence. Say it that way and you will not fall into the common wrong picture. Image diffusion models start from a field of random Gaussian pixels and denoise them. Text dLLMs do not. Their starting state is a row of blanks. A blank is a special vocabulary symbol, the mask token, that means "no token decided here yet." The LLaDA paper (arXiv:2502.09992, from Nie and colleagues, ICLR 2025) builds exactly this: an all-masked sequence and a reverse process that fills it in.
The thing doing the filling is a mask predictor. That is a Transformer with no causal mask, so every position can read every other position in both directions. Call that bidirectional attention: unlike the causal attention of a GPT, where a token sees only its past, here a masked cell in the middle of the sequence sees the whole sequence at once, left and right. You built the bidirectional version already, without the name, in the vision encoder of Project 29: Multimodal: A Tiny Vision-Language Model. Same idea, different data.
A denoising step is one forward pass of the mask predictor followed by a decision about which blanks to fill. The model reads the current partly-filled sequence. It predicts a probability distribution over the vocabulary for every masked position at once. And it reports a confidence for each: how sure it is of its best guess there, read straight off the top probability. Then comes the rule that makes the whole thing work. Keep the high-confidence guesses. Commit them as real tokens. Leave the low-confidence positions blank, or blank them back out if an earlier step filled them badly. That last move has a name, remasking: a position that was written can be returned to the mask state on a later step, which is how the model revises. Run this for T steps, from an all-masked sequence toward a fully-filled one, and the sequence resolves coarse-to-fine: the easy, high-confidence words first, the words that depend on them last.
Researchers have a formal name for this setup: absorbing-state diffusion. The mask token is the absorbing state. A forward process corrupts real text by sending positions to the mask, and once a position is masked it stays masked until the reverse process fills it. So the entire model reduces to one learned skill. Given a partly-masked sequence, guess the masked tokens. That is all it ever does. This is why an ordinary masked language model can stand in for the decode loop at all: it already has that skill, learned at a low masking rate. A dLLM has the same skill at every masking rate, all the way up to the all-mask start.
Map the analogy onto the tensors before the build, so the crossword becomes concrete. The grid is the token sequence. A blank cell is a position holding the mask id. A pencilled-in answer is a position holding a real vocabulary id. The solver's certainty is the top softmax probability at that position. A pass over the grid is one forward call of the mask predictor. Filling the sure clues is committing the highest-confidence positions. Rubbing out a doubtful cell is remasking. Every noun in the analogy is a value you can print in the build below.

Contrast the two orders plainly. The AR model asks, at every step, one question: given everything so far, what is the single next token? The dLLM asks a different question at every step: given everything decided anywhere in the sequence, which blanks are now sure enough to fill? The AR model's answer is one token. The dLLM's answer is a batch of tokens, chosen by confidence rather than by position. That is the entire architectural difference, and everything else in this chapter follows from it.
One honest word before the payoff, because it is the part marketing tends to drop. Filling many positions at once is fast, but it is not free of trouble. When the model commits several blanks in a single step, it decides each of them from the same context, without seeing the others it is committing alongside. If two of those positions needed to agree with each other, nothing forced them to. Fewer steps means more positions decided in parallel per step, which means more of these missed agreements. That tension is the whole game. It is what BREAK IT will make concrete.
Why It Matters
The reason anyone builds these is throughput. An autoregressive model's speed has a hard floor. It runs one forward pass per token, in order. So a 200-token answer is 200 sequential passes, no matter how much hardware you own. A dLLM breaks that floor by deciding many tokens per pass, and the reported numbers are large.
State them with their conditions attached, because the conditions are the whole story. Inception Labs reports its Mercury Coder models running at over 1000 tokens per second on NVIDIA H100 GPUs: Mercury Coder Mini at 1109 tok/s and Mercury Coder Small at 737 tok/s, per the Mercury paper (arXiv:2506.17298) and the company's launch blog. Those are code-optimized models, single-request throughput, on commodity H100s rather than custom silicon, and they are the vendor's own benchmark. Inception Labs positions this as up to roughly 10x faster than speed-optimized frontier autoregressive models, with the blog phrasing it 5 to 10x. That is measured against AR baselines in the 200-tok/s class. It is not a quality-matched comparison. Google DeepMind reports a second system, Gemini Diffusion, at about 1479 tok/s, framed as the quality of Gemini 2.0 Flash-Lite at roughly 5x the speed. Read the label on that one too. DeepMind calls it an experimental demo, not a shipped product. The speed figure excludes a fixed overhead in their own chart. When a 200-token answer arrives in a fifth of the time, an interactive coding assistant stops feeling like a typewriter and starts feeling like pasting the whole answer at once.
Now the other half, stated just as plainly, because a chapter that only sold the speed would be lying by omission. Diffusion LLMs do not beat autoregressive models on reasoning. On DeepMind's own evaluation table, Gemini Diffusion scores 15.0% on BIG-Bench Extra Hard against 21.0% for Gemini 2.0 Flash-Lite. A clear reasoning gap, not parity. The Efficient-DLM paper (arXiv:2512.14067) notes that diffusion models' learning efficiency lags autoregressive models when trained from scratch. Where dLLMs are competitive is code and short-form generation. On the same DeepMind table, Gemini Diffusion posts 89.6% on HumanEval against Flash-Lite's 90.2%. On LiveCodeBench it is 30.9% against 28.5%. Even there, close, not a rout. The open research anchor is LLaDA, an 8B diffusion model trained from scratch. The paper shows it competitive with a comparable autoregressive baseline (LLaMA3 8B) on in-context learning and instruction following. It also mitigates the "reversal curse," even beating GPT-4o on a reverse-poem-completion task. Competitive with. Not uniformly better. The picture is a speed-versus-capability frontier, strongest on code and short answers, still behind on hard reasoning and long-form coherence.
The production status matters when you decide whether to reach for one. As of mid-2026 the three anchors sit at three different stages. Mercury is shipped. Inception Labs offers a playground, an API, and on-prem deployment. Gemini Diffusion is an experimental demo from Google DeepMind, not a generally available product. LLaDA is open research with public weights on Hugging Face. So the one you can pull down and run today, weights and all, is LLaDA. The one you can call in production is Mercury. The build below imitates the decode loop both of them use, on hardware you already own.
There is a connection back into the book that is easy to miss, and worth pinning down because it prevents a real confusion. The phrase "block diffusion" already appeared in this book, in a completely different job. Project 14: Speculative Decoding used a lightweight block-diffusion model as a drafter. The DFlash method (arXiv:2602.06036) drafts a whole block of guess tokens in one non-causal forward pass. Then a large autoregressive model verifies every one of them and throws out the wrong guesses. The DFlash paper reports up to about 6x lossless acceleration across models (around 4.9x average on Qwen3-8B, greedy, temperature 0). That is diffusion used as a fast intern whose work a left-to-right model checks. This chapter is the other use entirely: here the diffusion model is the decoder, and nothing downstream re-does its work left to right. Same word, "diffusion," two roles. In Project 14 it is a draft mechanism inside AR serving. Here it is the full non-autoregressive model. Both are masked/discrete diffusion, and in neither case is the "noise" continuous Gaussian noise.
The serving side connects too. Deciding many tokens per pass raises GPU utilization, but a long context makes the key-value cache from Project 13: Fast Inference: The KV Cache grow until it dominates memory. The HERALD system (arXiv:2606.21633) targets exactly that for block-diffusion serving. It reports up to 1.59x lower per-block latency and 2.47x higher throughput than a GPU-only baseline, at near-lossless accuracy with a 5 to 10% KV budget. Those figures were measured across three block-diffusion LLMs on five long-context tasks, built on SGLang. That is the same production-serving concern as Project 17: Production Serving: Continuous Batching and PagedAttention, aimed at a non-autoregressive decoder.
The Build
build.py$python projects/31_diffusion-and-non-autoregressive-decoding/build.py --tiny
"""
Project 31: Diffusion and Non-Autoregressive Decoding
Canonical runnable build for this project. Generated by tools/extract_code.py
from the book's chapter markdown — re-run that tool to regenerate.
"""
# === Step 1: Load the mask predictor and locate the blank ===
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForMaskedLM
MODEL = "bert-base-uncased"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForMaskedLM.from_pretrained(MODEL, dtype=torch.float32).eval()
MASK = tok.mask_token_id # the integer id of the [MASK] blank
print("mask token:", tok.mask_token, "id:", MASK)
# === Step 2: Build the all-masked starting sequence ===
def build_sequence(prefix_text, n_blanks):
prefix = tok(prefix_text, add_special_tokens=False)["input_ids"]
ids = ([tok.cls_token_id] + prefix
+ [MASK] * n_blanks + [tok.sep_token_id])
x = torch.tensor([ids]) # shape (1, seq_len)
start = 1 + len(prefix) # first blank position
fill = torch.arange(start, start + n_blanks) # positions we may fill
return x, fill
x, fill = build_sequence("the weather today is", n_blanks=6)
print(tok.decode(x[0]))
# [CLS] the weather today is [MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [SEP]
# === Step 3: One denoising step, and the confidence read ===
@torch.no_grad()
def predict(x, fill):
logits = model(x).logits[0] # (seq_len, vocab)
probs = F.softmax(logits[fill], dim=-1) # (n_blanks, vocab)
conf, pred = probs.max(dim=-1) # best prob + best token id
return conf, pred # each shape (n_blanks,)
# === Step 4: The unmask schedule and the loop ===
@torch.no_grad()
def diffusion_decode(prefix_text, n_blanks, steps, verbose=True):
x, fill = build_sequence(prefix_text, n_blanks)
committed = torch.zeros(n_blanks, dtype=torch.bool)
for s in range(steps):
conf, pred = predict(x, fill)
conf = conf.masked_fill(committed, -1.0) # ignore filled blanks
target = round(n_blanks * (s + 1) / steps) # how many filled by now
reveal = target - int(committed.sum())
if reveal <= 0:
continue
order = torch.argsort(conf, descending=True) # most confident first
for j in order[:reveal]:
x[0, fill[j]] = pred[j] # write the token into the blank
committed[j] = True # mark it committed
if verbose:
filled = int(committed.sum())
print(f"step {s+1:>2}/{steps} filled {filled}/{n_blanks} "
f"| {tok.decode(x[0, fill])}")
return tok.decode(x[0, fill])
out = diffusion_decode("the weather today is", n_blanks=6, steps=6)
The Steps
4 reference filesEach step is the code for one piece, in isolation — read them as the book introduces them. They assemble into build.py above.
Step 01Load The Mask Predictor And Locate The Blankstep_01_load-the-mask-predictor-and-locate-the-blank.py
"""
Project 31: Step 1 — Load the mask predictor and locate the blank
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForMaskedLM
MODEL = "bert-base-uncased"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForMaskedLM.from_pretrained(MODEL, dtype=torch.float32).eval()
MASK = tok.mask_token_id # the integer id of the [MASK] blank
print("mask token:", tok.mask_token, "id:", MASK)
Step 02Build The All Masked Starting Sequencestep_02_build-the-all-masked-starting-sequence.py
"""
Project 31: Step 2 — Build the all-masked starting sequence
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
def build_sequence(prefix_text, n_blanks):
prefix = tok(prefix_text, add_special_tokens=False)["input_ids"]
ids = ([tok.cls_token_id] + prefix
+ [MASK] * n_blanks + [tok.sep_token_id])
x = torch.tensor([ids]) # shape (1, seq_len)
start = 1 + len(prefix) # first blank position
fill = torch.arange(start, start + n_blanks) # positions we may fill
return x, fill
x, fill = build_sequence("the weather today is", n_blanks=6)
print(tok.decode(x[0]))
# [CLS] the weather today is [MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [SEP]
Step 03One Denoising Step And The Confidence Readstep_03_one-denoising-step-and-the-confidence-read.py
"""
Project 31: Step 3 — One denoising step, and the confidence read
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
@torch.no_grad()
def predict(x, fill):
logits = model(x).logits[0] # (seq_len, vocab)
probs = F.softmax(logits[fill], dim=-1) # (n_blanks, vocab)
conf, pred = probs.max(dim=-1) # best prob + best token id
return conf, pred # each shape (n_blanks,)
Step 04The Unmask Schedule And The Loopstep_04_the-unmask-schedule-and-the-loop.py
"""
Project 31: Step 4 — The unmask schedule and the loop
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
@torch.no_grad()
def diffusion_decode(prefix_text, n_blanks, steps, verbose=True):
x, fill = build_sequence(prefix_text, n_blanks)
committed = torch.zeros(n_blanks, dtype=torch.bool)
for s in range(steps):
conf, pred = predict(x, fill)
conf = conf.masked_fill(committed, -1.0) # ignore filled blanks
target = round(n_blanks * (s + 1) / steps) # how many filled by now
reveal = target - int(committed.sum())
if reveal <= 0:
continue
order = torch.argsort(conf, descending=True) # most confident first
for j in order[:reveal]:
x[0, fill[j]] = pred[j] # write the token into the blank
committed[j] = True # mark it committed
if verbose:
filled = int(committed.sum())
print(f"step {s+1:>2}/{steps} filled {filled}/{n_blanks} "
f"| {tok.decode(x[0, fill])}")
return tok.decode(x[0, fill])
out = diffusion_decode("the weather today is", n_blanks=6, steps=6)
Break It
break_it.py$python projects/31_diffusion-and-non-autoregressive-decoding/break_it.py --tiny
"""
Project 31: BREAK IT experiment.
Deliberately sabotages one mechanism from build.py to show what happens
when it's removed. Compare outputs to the working version.
"""
print("baseline, one blank per step:")
diffusion_decode("the weather today is", n_blanks=6, steps=6)
print("\nbroken, all blanks in one step:")
diffusion_decode("the weather today is", n_blanks=6, steps=1)
for T in (1, 2, 3, 6):
out = diffusion_decode("the weather today is", n_blanks=6,
steps=T, verbose=False)
print(f"steps={T}: {out}")
Read it in the book
This page ships the code. The full derivation — why the chain rule is mechanical not magical, what the computational graph buys you, the long-form reasoning — lives in Chapter 31 of Under the Hood.
Get the book →