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.
  • Add remasking (ReMDM) so the sampler can un-write a committed token and self-correct — turning extra steps into real inference-time scaling.
  • Choose among the ReMDM schedule family — cap, rescale, and loop — three ways to spend the same marginal-preserving band [0, \sigma_t^{\max}].
  • Distinguish the stochastic (coin-per-token) and deterministic (\lfloor \sigma_t n\rfloor) remask realizations, and prove they agree in expectation — the count is \mathrm{Binomial}(n,\sigma_t) with mean n\sigma_t.

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 — a limitation we lift in Remasking below.

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.

But look closely at that curve: past a point it stops climbing. On the real OpenWebText benchmark, MDLM’s sample quality (MAUVE) is stuck near 0.035 and barely moves as you pour in more steps. Something is capping it — and it is the one rule we have taken for granted since the first sample.

Remasking: Sampling That Can Take a Token Back

Re-read the reveal loop’s invariant: a committed token is never overwritten. That is not a shortcut — it is the exact MDLM reverse posterior. For an already-unmasked position, q(z_s \mid z_t, x) = \mathrm{Cat}(z_s; z_t): all the probability mass sits on the token already there, so it is copied forward, forever. Unmasking is absorbing.

That invariant has a dark side. Suppose the model reveals a token early with high confidence, and it is wrong — a plausible mistake that a little more context would have corrected. Too late: it is frozen. Every later step builds on the error, and adding steps cannot fix it, because no step is allowed to touch it. This is exactly why quality saturates. The model has no eraser.

NoteKey Insight

Autoregression can’t revise the past either — but it never sees the future, so it never has reason to. A diffusion model fills tokens out of order, so an early guess is made with less context than the tokens revealed after it. Freezing that early guess is the worst of both worlds: it was the least-informed decision, and it is the one you can never revisit.

ReMDM (Wang, Schiff, Sahoo & Kuleshov, NeurIPS 2025) hands the sampler an eraser. It generalizes the reverse step with a single new knob, a remask probability \sigma_t: at each step, an already-committed token is sent back to [MASK] with probability \sigma_t, where the bidirectional predictor gets to re-decide it next step — now with everything else filled in. The masked-position rule is rebalanced by the same \sigma_t so the process still lands on a clean sequence:

q_\sigma(z_s \mid z_t, x) = \begin{cases} \mathrm{Cat}\big(z_s;\ (1-\sigma_t)\,x + \sigma_t\, m\big), & z_t \neq m \ \ (\text{committed}),\\[8pt] \mathrm{Cat}\!\left(z_s;\ \dfrac{\alpha_s-(1-\sigma_t)\alpha_t}{1-\alpha_t}\,x + \dfrac{1-\alpha_s-\sigma_t\alpha_t}{1-\alpha_t}\,m\right), & z_t = m \ \ (\text{masked}). \end{cases}

Set \sigma_t = 0 and the top case collapses to \mathrm{Cat}(z_s; z_t) (the frozen carry-over) and the bottom to the ordinary reveal mixture — plain MDLM, recovered exactly. ReMDM is a strict generalization of the sampler you already built, not a different one.

\sigma_t is not free, though. To keep the noisy-sequence marginals identical to standard masked diffusion — the property that lets us reuse the same trained model with no retraining — it must stay inside a band:

0 \;\le\; \sigma_t \;\le\; \min\!\left(1,\ \frac{1-\alpha_s}{\alpha_t}\right) \;=:\; \sigma_t^{\max}.

Under this module’s linear schedule \alpha_t = 1-t, stepping from time t down to s, that is \sigma_t^{\max} = \min(1,\ s/(1-t)). Two edges matter: at the final step s = 0, so \sigma_t^{\max} = 0 — the last step can only reveal, guaranteeing a clean finish; and early on, when little is committed, there is little to remask anyway. The simplest schedule in the band is ReMDM-cap — a constant rate \eta_\text{cap}, clipped to the bound:

\sigma_t = \min\!\left(\eta_\text{cap},\ \sigma_t^{\max}\right).

(Wang et al. also propose rescale, conf, and loop schedules; cap is the one we build.) Drive \eta_\text{cap} and watch the schedule live inside its valid band:

Notice the ceiling collapse to zero as t \to 0: the schedule cannot remask on the run-in to a clean sequence, so generation always terminates.

Three schedules for one band

cap clips a flat line against the ceiling — but the ceiling has a shape, and that shape is information. \sigma_t^{\max} is near zero at both ends (little is committed early; nothing may be remasked on the clean run-out) and bulges through the middle, exactly where the sampler has decoded enough context to judge an early token but has not yet finalized it. Wang et al. give two more schedules that read that shape, and both still land clean for free:

\underbrace{\sigma_t = \eta\,\sigma_t^{\max}}_{\textbf{rescale}} \qquad\qquad \underbrace{\sigma_t = \begin{cases}\min(\eta,\ \sigma_t^{\max}), & t_\text{off} \le t \le t_\text{on},\\[4pt] 0, & \text{otherwise.}\end{cases}}_{\textbf{loop}}

rescale takes a constant fraction of the band instead of a constant rate: it rides the bulge, remasking hardest when there is most to fix and tapering to zero on its own. loop switches remasking on only inside a middle time window [t_\text{off},\, t_\text{on}] (recall t\!:\!1\to0, so it turns on once t drops past t_\text{on} and off again below t_\text{off}). That splits sampling into three phases — reveal like plain MDLM → a burst of correction → a clean pure-reveal finish — concentrating the extra compute where it pays.

NoteKey Insight

All three schedules are the same generalized reverse step; they differ only in how they spend the band. cap spends a fixed amount every step, rescale spends a fixed fraction (so it follows the band’s natural rise and fall), and loop spends nothing until a chosen window and then a lot. None can exceed \sigma_t^{\max}, so all three reuse the pretrained model and all three finish clean — the choice is purely where in generation to place the eraser.

Pick a schedule and watch how differently it fills the same band. (Our loop keeps the module’s linear-\alpha march and simply gates remasking to the window; the paper’s loop additionally holds \alpha constant inside it — re-predicting at one fixed noise level — which we leave as a follow-up. The phase structure is the same.)

Now measure what the shape buys. Seed every single-token error as before, but sweep all three schedules at the same rate against the reveal-only floor. On this toy model the differences are modest and budget-dependent — the point is that each schedule is a distinct, valid way to spend the band, not that one dominates.

TipTry This

Set the schedule to rescale and drag \eta up: the whole curve swells to fill the band, peaking in the middle of generation — no clipping corner. Switch to loop and the line goes flat-zero outside the dashed window and jumps inside it. Then compare the recovery chart: all three schedules climb off the reveal-only floor, but they place their gains at different step budgets — the band is one resource, and the schedule is your policy for spending it.

Code: the remasking sampler

remask_sample in diffusion.py is the reveal loop with one extra pass. After committing the surest masked positions exactly as before, it re-scores the committed tokens and sends the \lfloor \sigma_t \cdot n_\text{committed}\rfloor least confident ones back to [MASK] — the deterministic realization of ReMDM-conf’s “remask inversely to confidence.” A committed token the model now doubts is precisely the one worth reopening. The rate \sigma_t comes from remdm_sigma(...), the family dispatcher, so schedule="cap" (the default), "rescale", or "loop" all flow through the same sampler unchanged.

import inspect
from diffusion import remask_sample, remdm_cap_sigma, sigma_max

# σ_t = 0 makes the remask sampler bit-for-bit the reveal-only sampler.
print("ReMDM-cap σ_t at t=0.5→s=0.4, η_cap=0.1 :", round(remdm_cap_sigma(0.5, 0.4, 0.1), 3))
print("valid ceiling σ_max there              :", round(sigma_max(0.5, 0.6), 3))
print("last step (s=0) can never remask       :", remdm_cap_sigma(0.1, 0.0, 0.9))
print()
print("remask pass (source):")
src = inspect.getsource(remask_sample).split("Remask phase")[1].split("trajectory.append")[0]
print("        # --- Remask phase" + src.rstrip())
ReMDM-cap σ_t at t=0.5→s=0.4, η_cap=0.1 : 0.1
valid ceiling σ_max there              : 0.8
last step (s=0) can never remask       : 0.0

remask pass (source):
        # --- Remask phase (ReMDM): reopen committed slots to self-correct ---
        sigma = remdm_sigma(t, s, eta_cap, schedule=schedule, t_on=t_on, t_off=t_off)
        if sigma > 0.0 and i < steps - 1:
            committed = x != mask_id
            if remask == "stochastic":
                # The honest per-token process: each committed slot flips w.p. σ_t.
                x[stochastic_remask(committed, sigma, generator=generator)] = mask_id
            else:
                # The ⌊σ_t·n⌋ least-confident committed slots (ReMDM-conf, fixed count).
                # Confidence the model *now* assigns to the token sitting in each slot.
                committed_conf = probs.gather(-1, x.clamp(max=model.vocab_size - 1).unsqueeze(-1)).squeeze(-1)
                # Non-committed slots can't be remasked → sort last.
                committed_conf = committed_conf.masked_fill(~committed, float("inf"))
                for b in range(batch_size):
                    n_committed = int(committed[b].sum())
                    n_remask = int(sigma * n_committed)
                    if n_remask <= 0:
                        continue
                    order = torch.argsort(committed_conf[b], descending=False)
                    remask_idx = order[:n_remask]
                    remask_idx = remask_idx[committed[b][remask_idx]]
                    x[b, remask_idx] = mask_id

Watch it self-correct

Here is the payoff on a real trained model. We take a memorized sentence, plant one wrong token in it — a fully-committed sequence with a single planted error, exactly what a greedy early reveal produces — and denoise it two ways from the same seed. Reveal-only (\eta_\text{cap} = 0) is inert: nothing is masked, so nothing changes, and the error survives. ReMDM spots the model’s low confidence in the planted token, remasks it, and re-predicts it correctly.

TipTry This
  1. Find the ↩︎. Step forward until position 3 flips from its wrong red word to the cyan ↩︎ — that is the model erasing its own mistake — then watch it come back green and correct.
  2. Compare the accuracies. Reveal-only (\eta_\text{cap}=0) ends at token accuracy (the error is permanent); ReMDM ends at .

Inference-time scaling: an eraser turns steps into quality

The saturating “steps → quality” curve from the last section had a hidden ceiling: without an eraser, extra steps can only reveal faster, never fix. With remasking, extra steps buy correction passes. To see it cleanly, seed every single-token error into every memorized sentence and measure how much each sampler recovers as the step budget grows. Reveal-only stays pinned at the one-token-wrong floor; ReMDM climbs.

This is the headline of the ReMDM paper, in miniature. On real benchmarks the same shape holds at scale: OpenWebText sample quality (MAUVE) jumps from MDLM’s 0.035 to 0.656 at 4096 steps, and wrapping the pretrained LLaDA-8B model in the ReMDM sampler lifts Countdown reasoning accuracy from 45.2 to 46.1 — all with no retraining, just a better reverse process.

NoteKey Insight

Autoregression’s quality knob is model size; you cannot spend more compute at inference to fix a bad sample, only resample it. Reveal-only diffusion added a steps knob but capped it with the freeze. Remasking is what makes that knob actually buy quality — a genuine inference-time-compute axis, the same idea as reasoning’s test-time compute, living in the sampler instead of the chain of thought.

The Count or the Coin: Stochastic Remasking

Look again at the committed-token case of the ReMDM posterior:

q_\sigma(z_s \mid z_t, x) = \mathrm{Cat}\big(z_s;\ (1-\sigma_t)\,x + \sigma_t\, m\big), \qquad z_t \neq m.

Read literally, this is a coin flip, one per token: every already-committed position is sent back to [MASK] independently with probability \sigma_t, and kept otherwise. Nothing in the equation counts tokens, ranks them by confidence, or coordinates them — the reverse process treats each committed slot on its own.

Yet the sampler you built does none of that. It reopens exactly \lfloor \sigma_t \cdot n_\text{committed}\rfloor slots, and picks them deliberately — the least confident ones. So which is ReMDM: the coin, or the count?

Both. They are two realizations of the same per-step budget, and the bridge is a single line of probability. If each of n committed tokens is remasked independently with probability \sigma_t, the number that actually flip is a binomial random variable,

K \sim \mathrm{Binomial}(n,\ \sigma_t), \qquad \mathbb{E}[K] = n\,\sigma_t, \qquad \mathrm{Var}[K] = n\,\sigma_t(1-\sigma_t).

The deterministic sampler spends \lfloor n\sigma_t\rfloor every step; the stochastic one spends n\sigma_t on average. They agree in expectation — the ⌊·⌋ count is just the stochastic process’s mean, made exact. What the fixed count throws away is the variance: the coin sometimes reopens far more or far fewer than n\sigma_t, while the count never wavers.

NoteKey Insight

This is exactly the systematic-vs-multinomial resampling choice from particle filters. Multinomial resampling draws each particle independently (the coin); systematic resampling fixes the count and stratifies the draw (our \lfloor n\sigma_t\rfloor). Same expectation, lower variance for the stratified one — and here we get a bonus variance reduction, because instead of picking the \lfloor n\sigma_t\rfloor slots uniformly we pick the least-confident ones, spending the budget where the model itself says correction is most needed (ReMDM-conf).

stochastic_remask in diffusion.py is the coin — one vectorized Bernoulli draw over the committed slots — and remask_sample(..., remask="stochastic") threads it through the sampler in place of the fixed-count pass. The default stays "deterministic", so everything above is unchanged; flip the switch and you get the honest per-token process. Because \sigma_t^{\max} \to 0 as s \to 0 either way, the stochastic path still finishes clean.

from diffusion import demonstrate_remask_count

# n committed tokens, each remasked independently with prob σ. Average the count.
demonstrate_remask_count(n=40, sigma=0.25, trials=4000, seed=0)
n = 40, σ = 0.25, trials = 4000
stochastic count : mean 9.936 (want n·σ = 10.000), std 2.698 (want √(nσ(1-σ)) = 2.739)
deterministic    : ⌊σ·n⌋ = 10 every step (std 0)
{'n': 40,
 'sigma': 0.25,
 'mean': 9.93575,
 'std': 2.697706792351597,
 'expected': 10.0,
 'std_analytic': 2.7386127875258306,
 'floor': 10,
 'counts': [11,
  12,
  12,
  10,
  11,
  10,
  11,
  9,
  7,
  13,
  10,
  13,
  7,
  12,
  11,
  8,
  8,
  9,
  12,
  12,
  4,
  5,
  9,
  15,
  10,
  12,
  10,
  11,
  16,
  12,
  13,
  11,
  10,
  7,
  7,
  13,
  8,
  9,
  9,
  11,
  8,
  12,
  11,
  7,
  7,
  11,
  6,
  12,
  9,
  13,
  7,
  7,
  10,
  11,
  6,
  11,
  8,
  10,
  7,
  9,
  12,
  10,
  13,
  14,
  8,
  13,
  3,
  14,
  12,
  7,
  5,
  13,
  7,
  6,
  9,
  7,
  5,
  7,
  13,
  8,
  7,
  15,
  11,
  10,
  9,
  11,
  12,
  13,
  13,
  15,
  12,
  9,
  9,
  5,
  12,
  14,
  6,
  6,
  13,
  12,
  10,
  10,
  12,
  8,
  12,
  12,
  10,
  9,
  12,
  12,
  11,
  6,
  6,
  12,
  11,
  8,
  8,
  13,
  10,
  7,
  9,
  10,
  17,
  10,
  12,
  12,
  7,
  9,
  7,
  9,
  9,
  11,
  9,
  11,
  8,
  9,
  10,
  13,
  8,
  8,
  15,
  11,
  11,
  8,
  8,
  7,
  6,
  16,
  6,
  9,
  11,
  8,
  8,
  14,
  14,
  13,
  8,
  10,
  10,
  7,
  8,
  12,
  13,
  10,
  16,
  7,
  7,
  12,
  7,
  6,
  8,
  13,
  13,
  8,
  8,
  10,
  11,
  9,
  5,
  6,
  10,
  9,
  9,
  11,
  9,
  10,
  11,
  8,
  9,
  10,
  13,
  13,
  11,
  5,
  11,
  11,
  9,
  11,
  6,
  8,
  11,
  12,
  18,
  9,
  6,
  12,
  11,
  7,
  6,
  10,
  14,
  7,
  10,
  5,
  11,
  4,
  6,
  9,
  11,
  12,
  10,
  10,
  10,
  9,
  10,
  3,
  9,
  12,
  10,
  12,
  11,
  6,
  8,
  7,
  9,
  9,
  9,
  8,
  8,
  11,
  12,
  10,
  12,
  6,
  11,
  12,
  10,
  10,
  12,
  7,
  6,
  14,
  16,
  13,
  7,
  13,
  9,
  11,
  10,
  8,
  9,
  11,
  11,
  12,
  13,
  9,
  7,
  8,
  10,
  11,
  13,
  11,
  8,
  12,
  10,
  7,
  4,
  10,
  8,
  13,
  4,
  7,
  7,
  10,
  9,
  12,
  6,
  13,
  9,
  9,
  12,
  9,
  5,
  13,
  7,
  9,
  12,
  16,
  9,
  10,
  10,
  8,
  9,
  9,
  13,
  7,
  8,
  10,
  8,
  8,
  13,
  8,
  13,
  13,
  13,
  12,
  9,
  13,
  10,
  10,
  9,
  15,
  9,
  10,
  9,
  10,
  13,
  10,
  9,
  5,
  13,
  5,
  7,
  12,
  12,
  14,
  6,
  10,
  8,
  9,
  7,
  12,
  10,
  12,
  11,
  11,
  11,
  9,
  9,
  8,
  12,
  9,
  8,
  12,
  6,
  13,
  7,
  7,
  12,
  10,
  9,
  14,
  9,
  9,
  10,
  15,
  8,
  8,
  8,
  14,
  8,
  11,
  5,
  18,
  11,
  9,
  10,
  12,
  11,
  8,
  9,
  7,
  11,
  10,
  12,
  9,
  6,
  13,
  10,
  7,
  11,
  12,
  12,
  11,
  11,
  7,
  1,
  12,
  9,
  12,
  10,
  10,
  7,
  11,
  12,
  7,
  9,
  12,
  11,
  10,
  10,
  13,
  11,
  8,
  11,
  6,
  7,
  10,
  8,
  10,
  9,
  9,
  11,
  10,
  4,
  10,
  13,
  10,
  12,
  5,
  9,
  13,
  10,
  7,
  8,
  6,
  13,
  12,
  8,
  10,
  13,
  8,
  11,
  7,
  9,
  4,
  8,
  12,
  16,
  6,
  13,
  6,
  17,
  11,
  8,
  11,
  11,
  8,
  8,
  12,
  11,
  13,
  12,
  6,
  6,
  15,
  12,
  7,
  8,
  9,
  11,
  9,
  11,
  17,
  14,
  8,
  10,
  12,
  5,
  5,
  10,
  7,
  10,
  7,
  10,
  6,
  12,
  5,
  14,
  9,
  14,
  9,
  10,
  13,
  10,
  9,
  12,
  13,
  7,
  9,
  11,
  9,
  10,
  9,
  12,
  10,
  11,
  8,
  10,
  5,
  9,
  10,
  16,
  10,
  10,
  15,
  11,
  10,
  13,
  6,
  12,
  12,
  12,
  4,
  12,
  7,
  12,
  6,
  5,
  9,
  13,
  13,
  7,
  15,
  10,
  13,
  11,
  12,
  9,
  12,
  16,
  14,
  8,
  9,
  11,
  12,
  8,
  12,
  12,
  15,
  12,
  9,
  9,
  12,
  9,
  7,
  13,
  11,
  11,
  9,
  9,
  11,
  7,
  10,
  16,
  8,
  11,
  6,
  11,
  9,
  11,
  13,
  8,
  9,
  11,
  10,
  11,
  9,
  9,
  12,
  9,
  8,
  8,
  9,
  11,
  11,
  12,
  13,
  9,
  7,
  6,
  9,
  8,
  15,
  8,
  10,
  11,
  13,
  6,
  11,
  9,
  15,
  10,
  7,
  7,
  11,
  10,
  9,
  9,
  14,
  6,
  4,
  8,
  14,
  7,
  10,
  13,
  13,
  13,
  8,
  10,
  10,
  5,
  13,
  6,
  10,
  9,
  9,
  15,
  7,
  12,
  11,
  6,
  9,
  10,
  7,
  8,
  9,
  7,
  10,
  9,
  11,
  9,
  10,
  9,
  9,
  8,
  7,
  8,
  6,
  14,
  15,
  13,
  10,
  8,
  9,
  11,
  10,
  7,
  8,
  9,
  11,
  7,
  14,
  10,
  7,
  12,
  13,
  10,
  7,
  9,
  11,
  8,
  11,
  11,
  8,
  9,
  8,
  11,
  10,
  11,
  14,
  9,
  10,
  10,
  7,
  10,
  8,
  9,
  6,
  10,
  10,
  8,
  9,
  9,
  8,
  9,
  10,
  10,
  7,
  13,
  6,
  5,
  5,
  10,
  11,
  10,
  9,
  12,
  9,
  7,
  11,
  9,
  8,
  12,
  8,
  9,
  8,
  8,
  12,
  10,
  8,
  9,
  5,
  9,
  10,
  9,
  4,
  8,
  9,
  11,
  9,
  10,
  10,
  10,
  6,
  12,
  9,
  10,
  16,
  14,
  14,
  9,
  7,
  10,
  8,
  6,
  8,
  10,
  13,
  8,
  13,
  9,
  10,
  7,
  10,
  7,
  11,
  14,
  10,
  7,
  9,
  9,
  11,
  14,
  8,
  9,
  7,
  16,
  13,
  8,
  12,
  10,
  6,
  11,
  10,
  10,
  11,
  15,
  4,
  13,
  10,
  13,
  9,
  7,
  9,
  9,
  7,
  9,
  12,
  7,
  12,
  15,
  7,
  8,
  17,
  9,
  15,
  13,
  12,
  13,
  10,
  9,
  9,
  7,
  13,
  14,
  11,
  7,
  10,
  12,
  7,
  15,
  9,
  6,
  10,
  9,
  6,
  9,
  8,
  13,
  6,
  12,
  8,
  6,
  5,
  9,
  8,
  11,
  10,
  12,
  10,
  8,
  7,
  11,
  10,
  5,
  9,
  12,
  11,
  10,
  13,
  12,
  11,
  10,
  7,
  12,
  15,
  11,
  10,
  8,
  7,
  12,
  8,
  12,
  9,
  9,
  7,
  14,
  9,
  13,
  12,
  9,
  15,
  15,
  10,
  11,
  9,
  11,
  9,
  12,
  12,
  9,
  12,
  12,
  10,
  9,
  14,
  6,
  9,
  12,
  7,
  10,
  15,
  9,
  6,
  12,
  7,
  12,
  10,
  10,
  12,
  10,
  5,
  12,
  14,
  10,
  6,
  7,
  11,
  8,
  13,
  7,
  5,
  12,
  14,
  11,
  10,
  10,
  11,
  9,
  6,
  10,
  11,
  16,
  16,
  10,
  12,
  13,
  14,
  12,
  8,
  13,
  8,
  8,
  11,
  8,
  9,
  14,
  11,
  11,
  10,
  10,
  14,
  7,
  9,
  7,
  10,
  11,
  9,
  7,
  5,
  8,
  7,
  7,
  12,
  12,
  11,
  10,
  8,
  10,
  10,
  8,
  10,
  12,
  11,
  14,
  10,
  11,
  11,
  8,
  9,
  11,
  9,
  12,
  9,
  10,
  10,
  7,
  8,
  7,
  16,
  9,
  16,
  5,
  11,
  11,
  7,
  7,
  10,
  11,
  11,
  5,
  10,
  14,
  12,
  11,
  10,
  6,
  9,
  11,
  10,
  7,
  10,
  7,
  7,
  12,
  14,
  10,
  9,
  ...]}

The empirical mean lands on n\sigma = 10 — the equivalence in expectation, measured — while the spread matches \sqrt{n\sigma(1-\sigma)} \approx 2.74. Drive n and \sigma below and watch the whole distribution of “how many got reopened”: the deterministic spike sits at \lfloor n\sigma\rfloor, the stochastic coin smears a binomial around the same mean, and as n grows the fraction remasked concentrates on \sigma (its spread shrinks like 1/\sqrt{n}) — the law of large numbers quietly turning the coin back into the count.

TipTry This
  1. Watch the count reappear. Hold \sigma fixed and slide n up: the binomial stays centered on n\sigma but its fractional spread \mathrm{std}/n shrinks toward zero — at large n the coin and the count are nearly the same thing.
  2. Widen the coin. Push \sigma toward 0.5: variance n\sigma(1-\sigma) is largest there, so the deterministic count (the orange bar at \lfloor n\sigma\rfloor) is at its most unlike any single stochastic draw.

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.
Remasking everything, or with no schedule Re-predicting committed tokens ad hoc destabilizes the reveal and may never terminate. Remask on purpose, but bound it: keep \sigma_t \le \sigma_t^{\max} (ReMDM) so the marginals hold and s\!=\!0 forces a clean finish.
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.
A loop window that never closes If t_\text{off} = 0 the remask window runs to the final step, but \sigma_t^{\max}\to0 there anyway — worse, a window that opens too late leaves no steps to correct in. Set t_\text{on} < 1 and t_\text{off} > 0 so all three phases (reveal, correct, finish) actually get steps.
Reading \sigma_t as a count The posterior sets a per-token probability, not a fixed number reopened; treating \lfloor \sigma_t n\rfloor as the definition hides that the honest process is \mathrm{Binomial}(n,\sigma_t) and can reopen more or fewer. The count is the expectation n\sigma_t; use remask="stochastic" for the true per-token coin, "deterministic" for its variance-reduced realization.

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.

Exercise 4: Confirm ReMDM generalizes MDLM

Show that remask_sample(..., eta_cap=0.0) reproduces reveal-only sample(...) exactly (same seed → identical tokens), then rerun with eta_cap=0.3 and count how many trajectory snapshots contain a token that was reopened (real → [MASK]).

# Your implementation here — compare sample vs remask_sample at eta_cap=0.0,
# then scan a remask_sample trajectory for real→[MASK] transitions.

Exercise 5: Spend the band three ways

Using remdm_sigma, tabulate \sigma_t at a mid-generation step (say t=0.5, s=0.4) for cap, rescale, and loop at \eta=0.5. Which one is largest, and why? Then sample the same seeded error with remask_sample(..., schedule="loop", t_on=0.7, t_off=0.1) and confirm from the trajectory that no slot is reopened before t drops past t_on.

# Your implementation here — compare remdm_sigma(0.5, 0.4, 0.5, schedule=...) for
# the three schedules, then check the loop window on a remask_sample trajectory.

Exercise 6: The count vs. the coin

Draw stochastic_remask many times over n=30 committed slots at sigma=0.3 and confirm the mean number reopened is \approx n\sigma = 9 (the equivalence in expectation), then compare it to the deterministic \lfloor n\sigma\rfloor. Bonus: show the sample fraction’s standard deviation shrinks when you double n.

# Your implementation here — use diffusion.stochastic_remask (or
# demonstrate_remask_count) and compare the empirical mean to n*sigma.

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.
  6. Remasking makes that dial pay off. Reveal-only decoding freezes early mistakes, so quality saturates. ReMDM adds a bounded remask probability \sigma_t \in [0, \sigma_t^{\max}] that reopens the least-confident committed tokens for self-correction — \sigma_t = 0 is plain MDLM, and more steps now fix rather than merely reveal. Real inference-time scaling, no retraining.
  7. A schedule is a spending policy for the band. The bound \sigma_t^{\max} has a shape — near zero at both ends, bulging through the middle. cap spends a flat rate, rescale = \eta\,\sigma_t^{\max} rides the shape, and loop confines remasking to a middle window (denoise → correct → finish). All three stay inside the band, so all three reuse the pretrained model and finish clean.
  8. The count and the coin agree in expectation. The posterior remasks each committed token independently with probability \sigma_t (the stochastic realization), so the number reopened is \mathrm{Binomial}(n,\sigma_t) with mean n\sigma_t. The sampler’s fixed \lfloor \sigma_t n\rfloor count is that mean made exact — a variance-reduced, confidence-informed draw of the same budget (systematic vs. multinomial resampling).

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.

Next, Module 27: Model Merging steps back from making models to combining them — folding several finished checkpoints into one by arithmetic in weight space, no training required.

Going Deeper

Core Papers:

Practical Resources: