Module 26: Diffusion Language Models

Introduction

Every model in this book so far writes text the same way: one token at a time, left to right. That is the autoregressive factorization \(p(x) = \prod_i p(x_i \mid x_{<i})\) — the transformer (m06) predicts the next token, generation appends it, repeat. Even the alternative architectures — SSMs, linear attention — change how each step is computed but keep the strict left-to-right order.

A diffusion language model throws that order out. It generates by iterative denoising: start from a sequence of all [MASK] and repeatedly fill in real tokens — in any order, several per step — until nothing is masked. The model is a mask predictor: given a half-finished sentence, it predicts what belongs in every hole at once.

Why it matters for LLMs:

  • Parallel, any-order generation. Autoregression is sequential by construction: token 100 waits for token 99. A diffusion LM commits many tokens per step and can decide the middle of a sentence before the ends — a different point on the quality/latency curve.
  • A compute dial, not a fixed cost. You choose the number of denoising steps. Fewer steps is faster and rougher; more steps is slower and sharper. Autoregression has no such knob — it always costs one forward pass per token.
  • It scales. LLaDA-8B (2025) is a masked diffusion model trained from scratch that competes with LLaMA3-8B — evidence that the paradigm is not a toy.

What You’ll Learn

After this module, you can:

  • Describe the absorbing-state forward process and why masking probability equals the diffusion time \(t\).
  • Build a bidirectional mask-predictor transformer from scratch and explain why it must not be causal.
  • Write the reweighted masked cross-entropy loss and connect it to masked language modeling.
  • Sample by confidence-ordered iterative unmasking and trade denoising steps for quality.

Prerequisites

This module requires familiarity with:

Intuition: Corruption and Repair

Think of a photograph dissolving into static and then being restored. Diffusion models are built from exactly two processes:

  • Forward (corruption) — a fixed, untrained process that gradually destroys data. For text, “destroy” means replace a token with [MASK]. At time \(t = 0\) the sentence is intact; as \(t \to 1\) more and more tokens turn into [MASK], until at \(t = 1\) nothing is left.
  • Reverse (repair) — a learned process that undoes one notch of corruption. A neural network looks at the partly-masked sentence and predicts the original tokens. Run it repeatedly and a clean sentence emerges from pure [MASK].

[MASK] is called an absorbing state: once a token is masked, the forward process leaves it masked. All the roads lead to the same fully-masked sink, which is why generation can start from there.

Drag the slider to watch the forward process dissolve a sentence. Each token flips to [MASK] independently once the noise level \(t\) crosses its own random threshold — so the number masked grows with \(t\), but which tokens go first is pure chance.

NoteKey Insight

The forward process has no parameters and needs no training — it is just “mask each token with probability \(t\).” All the learning is in the reverse network that has to undo it. This is the same split as image diffusion (add Gaussian noise forward, learn to denoise backward); only the corruption is different — discrete masking instead of continuous noise.

The Math: The Absorbing-State Process

Forward: masking is a coin flip per token

Fix a noise schedule \(\alpha_t \in [0, 1]\), the probability a token is kept. We use the simplest one, the linear schedule \(\alpha_t = 1 - t\). The forward process corrupts each token independently:

\[ q(x_t \mid x_0) = \alpha_t\,\big[x_t = x_0\big] \;+\; (1 - \alpha_t)\,\big[x_t = \texttt{[MASK]}\big]. \]

So a token survives with probability \(\alpha_t\) and becomes [MASK] with probability \(1 - \alpha_t = t\). The masking probability is the diffusion time. At \(t = 0\) nothing is masked; at \(t = 1\) everything is.

Reverse: a mask predictor

The reverse network \(x_\theta\) takes the corrupted sequence \(x_t\) and outputs, for every position, a distribution over the real vocabulary — its guess at the original token. Crucially it must see the whole sequence, past and future, masked and unmasked, to fill a hole from context. So it is a bidirectional transformer: the same block as m06 with the causal mask removed.

One surprise: the mask predictor takes no time input. It never needs to be told \(t\), because the fraction of [MASK] tokens it can see already reveals how noisy the sequence is. Following MDLM and LLaDA, our model is time-free.

The loss: reweighted masked cross-entropy

Training minimizes the negative evidence lower bound (NELBO). For absorbing-state diffusion this integral collapses into something you already know — a cross-entropy over the masked positions, weighted by the noise level:

\[ \mathcal{L} \;=\; \mathbb{E}_{t}\,\mathbb{E}_{q}\!\left[\;\frac{\alpha_t'}{1 - \alpha_t}\sum_{i \in \text{masked}} \log x_\theta(x_t)_{i, x_0^{(i)}}\;\right]. \]

Under the linear schedule (\(\alpha_t = 1 - t\), so \(\alpha_t' = -1\)), the weight is \(\alpha_t' / (1 - \alpha_t) = -1/t\): each masked position contributes its cross-entropy scaled by \(1/t\). That is the punchline of MDLM (Sahoo et al., 2024): a masked-diffusion language model is a masked language model, trained at every noise level and reweighted by \(1/t\). The \(1/t\) weight down-weights very noisy draws (large \(t\)), where the model has almost no context to work with.

Step by step

Walk one forward-then-reverse cycle:

Code: A Mask-Predictor from Scratch

The full implementation lives in diffusion.py. Start with the forward process — it is just a per-token coin flip:

import torch
from diffusion import forward_mask, linear_alpha

torch.manual_seed(0)
g = torch.Generator().manual_seed(0)

x0 = torch.tensor([[5, 2, 8, 1, 9, 4]])          # a clean sequence
xt, mask = forward_mask(x0, t=0.5, mask_id=99, generator=g)

print("clean x0: ", x0.tolist())
print("noisy xt: ", xt.tolist(), "  (99 = [MASK])")
print("masked?:  ", mask.tolist())
print("alpha_t (keep prob) at t=0.5:", linear_alpha(0.5))
clean x0:  [[5, 2, 8, 1, 9, 4]]
noisy xt:  [[99, 2, 99, 99, 99, 4]]   (99 = [MASK])
masked?:   [[True, False, True, True, True, False]]
alpha_t (keep prob) at t=0.5: 0.5

The marginal masking probability really is \(t\) — average over a big batch:

big = torch.zeros(500, 40, dtype=torch.long)
for t in (0.2, 0.5, 0.8):
    _, m = forward_mask(big, t=t, mask_id=99, generator=g)
    print(f"t = {t}:  fraction masked = {m.float().mean():.3f}")
t = 0.2:  fraction masked = 0.197
t = 0.5:  fraction masked = 0.499
t = 0.8:  fraction masked = 0.804

Now the mask predictor. It is an ordinary pre-norm transformer with one change — the attention is bidirectional (no causal mask), because filling a hole needs the whole sentence. The output head has vocab_size logits, not vocab_size + 1: the model removes [MASK], it never predicts it.

from diffusion import MaskedDiffusionLM

model = MaskedDiffusionLM(
    vocab_size=16, embed_dim=64, num_heads=4, num_layers=2, max_seq_len=8
)

x = torch.full((1, 6), model.mask_id)   # start from all [MASK]
logits = model(x)
print("input (all masked):", x.tolist())
print("logits shape:", tuple(logits.shape), "→ (batch, seq, real vocab)")
print("[MASK] id:", model.mask_id, "(the extra embedding row, never an output)")
input (all masked): [[16, 16, 16, 16, 16, 16]]
logits shape: (1, 6, 16) → (batch, seq, real vocab)
[MASK] id: 16 (the extra embedding row, never an output)

And the loss — cross-entropy on the masked positions, weighted by \(1/t\):

from diffusion import diffusion_loss

# Corrupt with the model's own [MASK] id, predict, and score the holes.
xt_m, mask_m = forward_mask(x0, t=0.5, mask_id=model.mask_id, generator=g)
logits = model(xt_m)
loss = diffusion_loss(logits, x0, mask_m, t=0.5)
print("noisy input:", xt_m.tolist(), f"  ({model.mask_id} = [MASK])")
print("diffusion loss:", round(loss.item(), 4))
print("only the", int(mask_m.sum()), "masked positions contribute")
noisy input: [[16, 2, 16, 1, 16, 4]]   (16 = [MASK])
diffusion loss: 5.4764
only the 3 masked positions contribute

This follows the algorithm in diffusion.py: forward_mask for corruption, MaskedDiffusionLM for the reverse network, diffusion_loss for the objective.

Bidirectional, not Causal

Modules 19 and 22 proved their mixers are causal: editing a future token never changes an earlier output — the property that makes autoregressive decoding valid. The diffusion mask predictor proves the mirror image. It is bidirectional on purpose: a future token must be able to reach backward, or it could not help fill an earlier hole.

model.eval()
x = torch.randint(0, 16, (1, 6))
with torch.no_grad():
    base = model(x)
    x_future_edit = x.clone()
    x_future_edit[0, 5] = (x[0, 5].item() + 1) % 16      # change the LAST token
    edited = model(x_future_edit)

moved = (base[0, 0] - edited[0, 0]).abs().max().item()
print(f"editing the last token moved the FIRST position's logits by {moved:.3f}")
print("bidirectional:", moved > 0)
editing the last token moved the FIRST position's logits by 0.083
bidirectional: True
NoteKey Insight

Causality and any-order generation are opposite design choices. An autoregressive model forbids future→past information flow so it can decode left-to-right; a diffusion model requires it so it can decode in any order. Neither is “better” — they sit at different points on the parallelism/latency frontier.

Sampling: Reveal by Confidence

Generation runs the reverse process. Start from all [MASK], then walk \(t\) from 1 down to 0 in a chosen number of steps. At each step:

  1. The mask predictor scores every masked position.
  2. Keep the highest-confidence predictions — commit them permanently.
  3. Re-mask the rest and continue.

This is LLaDA’s low-confidence remasking: the model finalizes the tokens it is surest about first and leaves the hard ones for later, when more context has been filled in. The number left masked after a step targets \(\text{length} \times t\), so the reveal finishes exactly when \(t\) hits 0. A committed token is never overwritten.

Below is a real trajectory: we train a tiny mask predictor on a handful of memorized sentences (over a 10-word toy vocabulary), then sample from a fully masked sequence. Step through it and watch tokens appear out of order — confidence, not position, decides who goes next.

TipTry This
  1. Watch the order. Note the position numbers as tokens appear — they do not fill left-to-right. The model commits its most confident guess wherever it is.
  2. Spot the bursts. Some steps reveal several tokens, some none — the \(\text{length} \times t\) target reveals in chunks, not one-at-a-time.

Steps ↔︎ Quality: the Compute Dial

The number of denoising steps is a dial you turn. With the same trained model, more steps means each step commits fewer (more confident) tokens, so the final sequence is sharper. Sweep the dial and watch accuracy climb:

With this toy model, 6 steps already recovers the memorized sentence perfectly, while 1–2 steps commits too much too fast and makes mistakes. Real diffusion LMs live on the same curve — you can generate a 100-token answer in far fewer than 100 steps, trading a little quality for a lot of speed.

Common Pitfalls

Pitfall Why it bites Fix
Making the predictor causal A causal mask blocks future→past flow, so a hole can’t be filled from the tokens after it. Use full (bidirectional) attention — no mask.
Scoring unmasked positions The model already sees those tokens; training on them leaks the answer and wastes the signal. Cross-entropy on masked positions only.
Overwriting committed tokens Re-predicting an already-revealed token can flip it and destabilize the reveal. Pin committed tokens; only ever unmask.
Forgetting the \(1/t\) weight Plain masked-CE over-weights very noisy (large-\(t\)) draws where prediction is near-hopeless. Weight each draw by \(1/t\) (the linear-schedule NELBO).
Confusing [MASK] with an output If the head can emit [MASK], sampling can “reveal” a mask and stall. Head has vocab_size logits; [MASK] is input-only.

Exercises

Exercise 1: Confirm the marginal

Show empirically that the fraction of masked tokens equals \(t\) for several noise levels. (You built this above — now add \(t = 0.1\) and \(t = 0.95\).)

from diffusion import forward_mask
gen = torch.Generator().manual_seed(1)
x = torch.zeros(1000, 32, dtype=torch.long)
for t in (0.1, 0.5, 0.95):
    _, m = forward_mask(x, t=t, mask_id=99, generator=gen)
    print(f"t={t}: masked fraction = {m.float().mean():.3f}")
t=0.1: masked fraction = 0.099
t=0.5: masked fraction = 0.500
t=0.95: masked fraction = 0.951

Exercise 2: Any-order proof

Write an assertion that editing the first token changes the last position’s logits (the reverse of the direction shown in the lesson).

# Your implementation here — edit x[0, 0], compare logits at position -1.

Exercise 3: Turn the dial

Sample the trained demo["model"] at steps=1 and steps=10 and print both sequences as words. How many positions differ?

# Your implementation here — use diffusion.sample and the WORDS list.

Summary

Key takeaways:

  1. A different factorization. Diffusion LMs replace left-to-right autoregression with iterative denoising — generate by repairing an all-[MASK] sequence, in any order.
  2. The forward process is free. Corruption is a parameter-free per-token coin flip; under the linear schedule the masking probability is exactly the diffusion time \(t\). All learning is in the reverse network.
  3. The predictor is bidirectional. Filling a hole needs the whole sentence, so the mask predictor drops the causal mask — the deliberate opposite of the causal guarantee that makes autoregression valid.
  4. The loss is reweighted masked-CE. The absorbing-diffusion NELBO reduces to cross-entropy on the masked positions, weighted by \(1/t\) — a masked language model trained at every noise level.
  5. Sampling reveals by confidence. Commit the surest predictions first, re-mask the rest (LLaDA’s low-confidence remasking); the number of denoising steps is a compute-vs-quality dial with no autoregressive analogue.

What’s Next

Diffusion is the book’s second full generation paradigm, sitting beside the autoregressive stack that runs from m06 through m08. From here you can fold it back into what you know: the mask predictor is the same transformer, so quantization and PEFT apply unchanged, and the evaluation harness scores its samples the same way. The frontier open questions — closing the last quality gap to autoregression, and block (semi-autoregressive) diffusion that interleaves the two paradigms — are where LLaDA and its successors are pushing now.

Going Deeper

Core Papers:

Practical Resources: