Diffusion Language Models

Chapter 1: Embedding Words into Diffusion

Chapter 0 named diffusion as the alternative. Now build it, starting from the version that already works on images, and find out exactly what breaks the moment you point it at text.

The vanilla diffusion model, in one page

A diffusion model treats data x₀ — for an image, a grid of real-valued pixel intensities — as the endpoint of a chain of increasingly noisy latent variables x_T, x_T-1, …, x₁, x₀, where x_T is (almost) pure Gaussian noise. The forward process defines how to get from clean data to noise by adding a little Gaussian noise at each step:

where β_t is a hyperparameter controlling how much noise gets added at step t. This process has no trainable parameters — it's a fixed recipe. The interesting part is the reverse process: a neural network p_θ(x_t-1 | x_t) trained to undo one noising step at a time, so that running it T times, starting from pure noise, produces something that looks like it came from the real data distribution.

Training uses a variational lower bound on the data's log-likelihood, and in practice a simplified surrogate loss (Ho et al. 2020) that turns each denoising step into an ordinary mean-squared-error regression problem: predict the mean of the true posterior q(x_t-1|x₀,x_t) — a closed-form Gaussian, since every step is Gaussian — using a network μ_θ(x_t,t). None of this machinery says anything about discreteness. It was built for continuous data.

Words are not continuous

A pixel intensity is already a real number; Gaussian noise can be added to it directly. A word is a discrete symbol from a fixed vocabulary — there is no native sense in which you can add “0.3 units of Gaussian noise” to the word “Cambridge.” Before any of the machinery above can run on text at all, you need a function that maps discrete words into the continuous space the diffusion process actually operates on:

Diffusion-LM's key architectural move is a new Markov transition, added in front of the standard diffusion chain, from the discrete sentence w to the continuous starting point x₀:

What makes this “end-to-end”

Crucially, Emb is not fixed in advance — it is trained jointly with the diffusion model itself, using the reparametrization trick to backpropagate through the sampling of x₀. The paper's preliminary experiments tried fixed embeddings (random Gaussian vectors, or off-the-shelf pretrained word vectors) and found them consistently worse than letting the embedding table be learned end-to-end alongside everything else. The resulting objective, extending the standard variational bound from Chapter 0's recap, picks up two extra terms for the new embedding step:

Read the three pieces in order: ℒ_vlb(x₀) is the ordinary diffusion denoising loss you already had. log q_φ(x₀|w) scores how close the embedding step landed to the real word embeddings. And −log p_θ(w|x₀) is a brand-new term — a rounding loss, scoring how well the model's clean latent x₀ can be mapped back to the actual discrete words that produced it. That third term is going to be the entire subject of Chapter 2: nothing here guarantees that a denoised x₀ lands exactly on a word's embedding, and it turns out that, unmodified, it usually doesn't.

Worked example: the noise schedule, by hand

Diffusion-LM uses a BERT-base Transformer of about 80 million parameters, sequence length n = 64, and T = 2,000 diffusion steps. The embedding dimension d is a hyperparameter: d = 16 for the E2E dataset (50K restaurant reviews, 8 labeled fields like food type and price), d = 128 for the harder ROCStories dataset (98K five-sentence stories, an 11K-word vocabulary — the same story-infilling task from Chapter 0).

Standard image-diffusion noise schedules turned out to be poorly suited to text. The paper's diagnosis: adding a small amount of Gaussian noise to a word embedding rarely changes which word is its nearest neighbor, so the early denoising steps (t near 0) are almost trivially easy under a schedule built for continuous pixels — wasted compute. Their fix is a custom square-root noise schedule:

Check it by hand at both ends of the chain, the way you should check any formula before trusting it. At t = 0 (the cleanest point):

The corresponding noise standard deviation is √(1−ᾱ₀) = √0.01 = 0.1 — matching exactly the “initial standard deviation of 0.1” the paper reports for this schedule. Now the other end, t = T = 2,000:

ᾱ near zero means the noise standard deviation √(1−ᾱ) is near 1 — x_T is (almost exactly) standard Gaussian noise, exactly the boundary condition the forward process is supposed to hit. Both ends check out from a two-line hand calculation, which is worth doing yourself before trusting any paper's noise schedule.

Qualitatively, the square-root shape front-loads noise injection — increasing rapidly over roughly the first 50 steps — then decelerates for the remaining ~1,950 steps, spending fewer of the schedule's fixed step budget on the too-easy near-clean regime this dataset turned out to have, and more on the harder, higher-noise problems in between.

Checking the schedule choice against the alternatives

Chapter 1's square-root schedule wasn't chosen by intuition alone — the paper's appendix ablation compares it directly against the two standard image-diffusion schedules, linear (Ho et al. 2020, the original DDPM schedule) and cosine (Nichol & Dhariwal 2021, a later refinement), across every embedding dimension and parametrization choice tested. The reported finding: the sqrt schedule attains consistently good and stable performance everywhere, while the paper notes it matters comparatively less once you've already applied the x₀-parametrization from Chapter 2 — but becomes substantially more important, a genuinely more robust choice, under the alternative ε-parametrization that Chapter 2 shows collapsing at higher dimensions. In other words: the noise schedule and the loss parametrization aren't independent design choices you can tune one at a time and expect the same answer either way — a schedule built for text can partially compensate for a parametrization built for images, and vice versa, which is exactly the kind of interaction effect that makes ablating “does the sqrt schedule matter” in isolation misleading without also fixing which parametrization you're testing it against.

Worked example: what training actually costs

Numbers worth having, both to ground “80M parameters” in something concrete and to set up the scale contrast Chapter 5 draws later. Training uses AdamW with a linearly-decaying learning rate starting at 1×10⁻⁴, dropout 0.1, and batch size 64, run for 200,000 iterations on E2E and 800,000 iterations on the harder, larger ROCStories dataset — a 4× difference in iteration count directly tracking the roughly 2× larger dataset (98K stories vs. 50K reviews) and its greater per-example difficulty. On a single NVIDIA A100, the paper reports roughly 5 hours to complete 200,000 iterations:

Hold that single-GPU, single-afternoon training budget in mind. Chapter 5's LLaDA run used 0.13 million GPU-hours — roughly 26,000× this chapter's 5-GPU-hour estimate for one full E2E training run. That ratio is a rough proxy for how much the entire project scaled up between 2022 and 2025, and it's worth returning to when Chapter 8 asks whether Diffusion-LM's honest costs (slower decoding, worse likelihood) survive at LLaDA's much larger scale, or get renegotiated by scale itself.

One more engineering detail worth surfacing: training under the full variational bound ℒ_vlb^e2e from Chapter 1 (rather than the simplified surrogate) needed two explicit stabilization tricks to converge at all — gradient clipping at 1.0, and importance sampling to reweight each term of ℒ_vlb during training (borrowed from Nichol & Dhariwal's image-diffusion work). Neither trick was necessary once the paper switched to the simplified surrogate objectives this chapter and Chapter 2 actually use. This is a small but telling data point about how much of the “engineering” in diffusion language modeling is really about taming training instability that the underlying math permits but doesn't prevent.

In code, so the shapes are undeniable

python
import torch, math

def sqrt_schedule(t, T=2000, s=1e-4):
    # alpha_bar_t = 1 - sqrt(t/T + s)
    return 1 - math.sqrt(t / T + s)

print(sqrt_schedule(0))      # 0.99   -> noise std 0.1, matches the paper
print(sqrt_schedule(2000))   # ~0.0   -> noise std ~1.0, pure Gaussian

# the new Markov step this chapter adds, on top of ordinary diffusion:
class WordEmbedding(torch.nn.Module):
    def __init__(self, vocab_size, d):
        super().__init__()
        self.emb = torch.nn.Embedding(vocab_size, d)   # TRAINED end-to-end, not fixed

    def forward(self, word_ids, sigma_0=0.1):
        mu = self.emb(word_ids)                       # Emb(w), shape (n, d)
        return mu + sigma_0 * torch.randn_like(mu)    # q_phi(x0 | w)

Worked example: the embedding table's own parameter budget

The embedding dimension d is not a free lunch — every dimension you add multiplies the size of the table that has to be learned end-to-end, jointly with a much larger diffusion network trying to converge on a shared objective. Work out the actual parameter counts for both datasets this chapter's architecture targets.

E2E's restaurant-review vocabulary is small (the dataset is 50K reviews across 8 structured fields, so the word vocabulary stays compact); at d = 16, a vocabulary of roughly 800 words costs:

ROCStories is a different story: 98K five-sentence stories draw on an 11,000-word vocabulary, and the paper uses a much wider d = 128 for exactly this dataset:

Notice the pattern before Chapter 5 revisits it at a completely different scale: embedding-table size scales with vocabulary × dimension, and a harder, more diverse dataset (ROCStories' larger vocabulary) earns itself a wider embedding space, at a real but manageable parameter cost. The same arithmetic, run again for LLaDA's 126,464-token vocabulary at Transformer-scale dimension 4,096, produces the roughly 1-GiB embedding table Chapter 5 derives by hand.

The forward process, once more, with the embedding step attached

It's worth writing the complete forward chain in one place, since every later chapter refers back to pieces of it. Starting from discrete words and ending at pure noise, four transitions in sequence:

Every quantity in this chain except step 1 is continuous, which is exactly what lets you borrow the entire machinery Chapter 1's recap opened with — the reverse process, the mean-squared-error surrogate loss, the whole apparatus — unmodified, once you've paid the one-time cost of defining Emb.

Worked example: checking the schedule at its midpoint too

The two boundary checks earlier in this chapter (t=0 and t=T) confirm the schedule starts and ends correctly. It's worth checking one interior point as well, since a formula can pass both endpoints and still misbehave in the middle. Take t = 1,000, exactly halfway through the T = 2,000-step schedule:

The corresponding noise standard deviation is √(1−0.2928) = √0.7072 ≈ 0.841 — already most of the way to 1.0 (pure noise), despite being only halfway through the schedule by step count. This confirms the qualitative claim from earlier in the chapter directly: the sqrt schedule front-loads noise injection heavily, reaching 84% of maximum noise standard deviation by the schedule's midpoint, rather than reaching it linearly (which would put the midpoint noise level at roughly 0.5, not 0.841). Half the schedule's steps remain after this point, but the sequence is already nearly fully noised — exactly the front-loaded behavior the schedule was designed to produce, confirmed by plugging in a number rather than just trusting the qualitative description.

The embedding dimension search, and why 16 beat 256 for E2E

Chapter 1 stated d=16 for E2E and d=128 for ROCStories as settled choices; they weren't chosen by intuition alone. The paper's hyperparameter search tested d ∈ {16, 64, 128, 256} on both datasets and selected per dataset based on downstream performance. That E2E's optimal dimension (16) is nearly an order of magnitude smaller than ROCStories' (128) is itself informative: E2E's vocabulary is small and its 8 structured fields constrain the space of plausible sentences considerably, so a low-dimensional embedding space is enough to separate the words that actually matter for this task; ROCStories, with its 11K-word vocabulary and much more open-ended five-sentence narratives, needs more dimensions to keep semantically distinct words from crowding together. This is a useful general instinct to carry forward: embedding dimension is not a “bigger is always safer” hyperparameter for Diffusion-LM the way it often is for a plain classifier — too many dimensions on an easy task can hurt, not just cost more compute, which is consistent with Chapter 2's finding that the ε-parametrization specifically collapses at higher dimensions rather than merely plateauing.

What a sequence length of 64 actually buys and costs

One more architectural number worth sitting with before Chapter 2 takes over: sequence length n=64. This is a genuinely short context by the standards even of 2022-era language models, and it's not an accident — every position in the sequence is present, fully, in the diffusion latent at every single one of the T diffusion steps, unlike an autoregressive model where only the tokens generated so far occupy any compute at a given point. A longer sequence multiplies the size of x₀ ∈ ℝ^nd directly (n × d numbers, all of them denoised together, at every step), which multiplies the compute cost of every single one of the 200 or 2,000 diffusion steps, not just the final output length the way autoregressive generation cost scales. This is the concrete, architecture-level reason Diffusion-LM's own scope stayed at short sequences (E2E reviews, individual ROCStories sentences) rather than the much longer documents autoregressive models routinely handle — and it's a cost Chapter 5 will show LLaDA inheriting in a different form: even though LLaDA drops the continuous embedding machinery, it still has to run full bidirectional attention over the entire current sequence at every sampling step, for exactly the same structural reason. Short-sequence Diffusion-LM and long-sequence-but-few-steps LLaDA are both symptoms of the same underlying fact: diffusion pays for the whole sequence, every step, no matter which corruption process it uses.

Chapter 2: The Rounding Problem

Chapter 1 got you a clean latent x₀ at the end of denoising. It lives in ℝ^nd — a continuous vector. Somewhere, that vector has to become an actual sentence again. This chapter is about the gap between “a vector near some word embeddings” and “a vector that unambiguously is one particular word,” and why that gap turned out to be the hardest part of the whole system.

Rounding, defined

The reverse process ends with a trainable rounding step, a softmax over the vocabulary at each position:

Ideally, the denoising process ends with x₀ sitting exactly on the embedding of some real word, making argmax rounding trivial and unambiguous. Empirically, the paper reports, it doesn't: the model fails to generate x₀ vectors that commit to a single word. They hover in the space between two or three plausible embeddings instead of landing decisively on one.

Diagnosing why

Recall the simplified surrogate loss from Chapter 1's recap — the network μ_θ(x_t,t) predicts the mean of the reverse transition p_θ(x_t-1|x_t) at every one of the T diffusion steps:

Here's the problem, and it's subtle: the constraint that x₀ has to commit to a single word embedding is only implicitly encoded, and only shows up strongly in the loss terms where t is very close to 0 — a small handful out of 2,000 total terms. The other ~1,995 terms are busy predicting an intermediate noisy mean, several steps away from any word-commitment decision at all. The paper found this parametrization required careful, fragile hyperparameter tuning to force the objective to emphasize the few terms that actually mattered for rounding.

Fix #1: predict x₀ directly, at every step

The fix re-parametrizes the network to predict x₀ itself — not the intermediate posterior mean — at every single one of the T diffusion steps:

Now every one of the 2,000 loss terms directly forces the network's output to be something that ought to sit on a real word embedding, not just the handful near t = 0. This is a genuinely different training signal, not a cosmetic relabeling — predicting x₀ and predicting x_t-1 are mathematically interchangeable up to a closed-form scaling (since the forward process gives x_t-1 = √ᾱ x₀ + √(1−ᾱ) ε directly), but which one you make the network's literal output changes what the loss function emphasizes at every step. The paper's own ablation (their Figure 4) shows this is not a small effect: parametrizing by the noise term ε, the standard choice for image diffusion, works fine at small embedding dimensions but collapses badly as the dimension grows, while the x₀-parametrization stays reliable across the dimensions they tested.

Fix #2: the clamping trick, at decoding time

The x₀-parametrization fixes training. A second, complementary fix improves decoding. Standard sampling from an x₀-parametrized model computes an estimate f_θ(x_t,t) and then samples the next latent conditioned on it:

The clamping trick inserts one extra operation: before using f_θ(x_t,t) in that formula, snap it to its nearest actual word embedding:

This forces the model to commit to a word identity at every intermediate step, not just at the very end — each subsequent denoising step then works around an actually-valid word vector, instead of an ambiguous blend of two or three candidates.

Worked example: clamping on a number line

Strip this down to a toy 1D vocabulary to see the mechanics without the vector-space complexity. Suppose three word embeddings sit at scalar positions {−2, 0, 3}, and at some intermediate diffusion step the network's raw prediction is f_θ(x_t,t) = 1.4.

Without clamping: the sampling formula uses 1.4 directly — a value sitting between the embeddings for “0” and “3,” genuinely ambiguous about which word it represents. Later steps inherit that ambiguity.

With clamping: compute the distance to each embedding — |1.4−(−2)| = 3.4, |1.4−0| = 1.4, |1.4−3| = 1.6 — and snap to the nearest, which is 0:

This position is now locked to “word 0” going forward. Every subsequent step for this position denoises around that committed value, rather than re-deciding between “0” and “3” from scratch at every step. (The paper notes clamping every step, including the earliest, noisiest ones, doesn't hurt performance much in practice, even though intuitively you'd expect early commitments — made before the model has much signal — to sometimes be wrong; they treat when to start clamping as an optional hyperparameter.)

Why ε-parametrization specifically collapses at high dimension

It's worth understanding, not just citing, why predicting the noise term ε — the standard choice for image diffusion, and a mathematically valid alternative to predicting x₀ directly — degrades as the embedding dimension d grows, since the paper's ablation reports exactly this pattern (Figure 4) without spelling out the mechanism. The two parametrizations are related by the forward process's closed form, x_t = √ᾱ x₀ + √(1−ᾱ) ε, which rearranges to:

Whatever error the network makes predicting ε gets divided by √ᾱ when converted back to an x₀ estimate — and ᾱ shrinks toward 0 as noise increases, meaning that division amplifies small ε-prediction errors into large x₀-estimate errors, especially at high-noise steps. In a wide embedding space (large d), the network has more coordinates in which to make a small mistake, and the amplification compounds across all d of them at once. Predicting x₀ directly skips that division entirely — whatever error the network makes IS the x₀ error, with no noise-schedule-dependent amplification factor multiplying it. That's the mechanical reason the paper's ablation shows ε-parametrization “working fine for small dimensions, but quickly collapsing for larger ones,” while x₀-parametrization stays reliable across the dimensions tested.

The rounding loss term, revisited

Chapter 1 flagged a third term in the full training objective, −log p_θ(w|x₀), and deferred explaining it to this chapter. Now you have the full picture: that term is exactly the softmax rounding step p_θ(w_i|x_i) discussed above, scored against the actual training words — it's what teaches the softmax classifier itself to map embedding-space vectors back to vocabulary indices accurately, a separate (if related) concern from teaching the diffusion network to land its x₀ predictions near valid embeddings in the first place. Two failure modes, two fixes: the x₀-parametrization (this chapter's Fix #1) makes the diffusion network commit to something embedding-shaped; the rounding loss term (Chapter 1's third term) makes the softmax head interpret that something correctly once it gets there.

A second architecture ablation: Transformer vs. U-Net

One more design choice worth naming, since it's easy to assume the Transformer backbone this session has used throughout was the only option ever considered. Image diffusion models overwhelmingly use a U-Net — a convolutional encoder-decoder with skip connections between matching resolutions. The paper explicitly tests this alternative for text, building a U-Net variant that mirrors Ho et al.'s original 2020 architecture exactly, except swapping every 2D convolution (suited to image height/width) for a 1D convolution (suited to a single sequence axis). The result, reported in their extended ablations: the Transformer architecture outperforms the U-Net variant. This is a useful data point precisely because it isn't obvious in advance — U-Nets earned their dominance in image diffusion through years of refinement specific to that domain, and there was no guarantee that advantage would transfer to a domain (sequences of discrete symbols embedded into vectors) with a very different notion of locality than a 2D pixel grid. It's a small confirmation that borrowing diffusion's process (the forward/reverse Markov chain) doesn't obligate you to borrow diffusion's most common backbone — the two choices are separable, and for text, a Transformer suited to sequences beat a convolutional network adapted from images.

Worked example: clamping in two dimensions

The number-line example above is honest but artificially simple — real embeddings live in d ≥ 16 dimensions, where “nearest” means Euclidean distance across all of them at once, not a single scalar comparison. Take a 2D toy case to see the generalization: three word embeddings at (0,0), (3,0), and (1,4), and a raw prediction f_θ(x_t,t) = (1.5, 1.2). Compute the squared distance to each:

This particular prediction happens to sit exactly equidistant from “word A” at the origin and “word B” at (3,0) — a genuinely ambiguous point, the kind of case where clamping's argmin has to break a tie (in practice, by whichever embedding numerically compares first, or by a small amount of noise in the actual floating-point distances breaking the exact tie). This is worth sitting with: clamping doesn't make ambiguity impossible, it just forces a hard decision at every step instead of letting the ambiguity persist and compound across all T steps. One bad clamp, corrected by later steps denoising around it, is far less damaging than an unresolved blend carried all the way to the final output.

What happens if you clamp too early — and why the paper says it's fine anyway

This chapter flagged, without fully resolving, an intuitive worry: clamping at very early steps (t close to T, almost pure noise) forces a word-identity decision before the model has had much chance to denoise anything at all. Isn't that just committing to noise?

Work through what actually happens next if an early clamp is wrong. Say at t = T−5 (still very noisy), the model clamps position 3 to the wrong word. That wrong commitment doesn't get treated as ground truth for the rest of the process — it becomes the input to the next denoising step, exactly like any other x_t-1. If later steps' evidence (the growing context from neighboring positions, increasingly resolved as t continues to fall) strongly disagrees with that early guess, the x₀-parametrized network at the next step is free to predict a completely different word for that position, and clamping will snap to that new nearest embedding instead — the wrong early commitment gets overwritten, not locked in permanently. This is precisely the same self-correction mechanism Chapter 4 documents in full with a real generated example (the “Family friendly Indian food” span getting fixed via a later change elsewhere in the sentence): commitment at any single step, including an early, uncertain one, is always revisable by subsequent steps, because clamping only fixes what happens within one denoising transition, never across the whole trajectory at once. That's the concrete reason the paper reports applying clamping at every step, even the noisiest ones, “doesn't hurt performance much in practice” — early mistakes are cheap precisely because there's a long remaining trajectory that gets a chance to fix them.

Worked example: what the softmax rounding head actually computes

Chapter 1 named p_θ(w_i|x_i) as a softmax and moved on; make it concrete with an actual small computation. Suppose, at the final denoising step, x₀ for one position is a vector close to (but not exactly on) three candidate word embeddings, and the rounding head produces three raw scores (logits) by comparing x₀ against each candidate — say −1.2 for “the,” 2.5 for “a,” and 0.3 for “an.” Softmax turns these into probabilities:

Argmax rounding picks “a,” the highest-probability candidate, with 88.1% confidence — a decisive commitment, the kind Chapter 2's fixes were designed to produce reliably. If instead the three logits had come out close together (say 0.4, 0.5, 0.3 — roughly the ambiguous, uncommitted case this chapter opened with), the resulting softmax probabilities would sit close to a uniform 1/3 each, and argmax rounding would be picking essentially at random among three nearly-tied options. That's rounding failure, made numerically concrete: not a mysterious phenomenon, just a softmax over logits that never separated cleanly, which is exactly what the x₀-parametrization and clamping trick are designed to prevent from happening in the first place.

This chapter's fix, restated as a one-sentence engineering principle

Strip away the notation and Chapter 2 reduces to a single transferable engineering principle, worth carrying into any generative model that has to commit to a discrete choice at the end of a continuous process: if a loss function only weakly penalizes an undesirable outcome (here, an ambiguous, uncommitted latent) in a small minority of its terms, don't expect the trained network to avoid that outcome reliably — reparametrize the objective so that every term penalizes it, or the network will happily leave that failure mode under-optimized. Diffusion-LM's specific instance of this principle was the switch from predicting an intermediate posterior mean (where word-commitment pressure appeared only near t=0) to predicting x₀ directly (where it appears in all T terms). The same principle resurfaces, in a different costume, essentially everywhere a paper reports needing to “reweight” or “reparametrize” a loss to fix a qualitative failure — including, as it happens, LLaDA's own 1/t reweighting in Chapter 5, which exists for a related but distinct reason (compensating for how many terms a given draw happens to contain, rather than concentrating commitment pressure across all terms) worth keeping straight rather than conflating.

Chapter 2, in one paragraph

Diffusion-LM's denoising ends with a continuous vector that has to become a discrete word, and nothing about a plain mean-squared-error training objective forces that vector to land decisively on any single word's embedding. Two fixes close the gap: retrain the network to predict x₀ at every one of the T steps rather than only near the end, so word-commitment pressure applies throughout training rather than in a small fraction of it; and at decoding time, snap every intermediate estimate to its nearest real word embedding before continuing, so ambiguity never gets a chance to accumulate across steps. Neither fix is optional, both are cheap, and together they turn a model that hovers between words into one that reads as ordinary text — the exact prerequisite Chapters 3 and 4 need before any of Diffusion-LM's control-generation results can mean anything at all. Without this chapter's two fixes, every downstream result in this session's first half would be built on top of a generator that couldn't reliably produce a single coherent sentence, let alone one satisfying a syntax tree or an exact word count.

Chapter 3: Steering the Latent

Chapters 1 and 2 built a working text diffusion model. Now put it to work on the problem Chapter 0 opened with: controllable generation, on constraints too global for left-to-right steering to handle well.

Control on the latent, not the words

Recall the Bayesian framing from Chapter 0: p(w|c) ∝ p_lm(w) · p(c|w). Diffusion-LM's generation process runs through a whole sequence of continuous latents x_0:T, so instead of steering the discrete words directly, the control target gets applied to the latent trajectory: p(x_0:T|c) = ∏_t=1^T p(x_t-1|x_t,c). Each per-step conditional factors, via Bayes' rule and a conditional-independence assumption borrowed from image diffusion guidance work, into:

Both terms on the right are differentiable: the first comes from Diffusion-LM itself, the second from a small classifier trained directly on the diffusion latents. That gives a gradient you can actually compute and climb:

Run gradient ascent on this at every diffusion step, and the whole latent sequence — every position, simultaneously — gets nudged toward satisfying c, while the first term keeps it anchored to what a fluent sentence's latent should look like.

Two modifications beyond straight image-style guidance

Image diffusion guidance takes one gradient step toward ∇log p(c|x_t-1) per diffusion step and calls it done. Text needs two adjustments:

Fluency regularization. Optimize a weighted objective, λ log p(x_t-1|x_t) + log p(c|x_t-1), where λ trades off fluency against control satisfaction. Without the λ-weighted fluency term, the paper found gradient ascent on the classifier term alone tends to produce text that satisfies the control but stops reading as grammatical language — the classifier gradient, left unchecked, will happily push the latent somewhere that scores well on the control and terribly on everything else.

Multiple gradient steps per diffusion step. Instead of one gradient update per step, Diffusion-LM runs 3 steps of Adagrad optimization at every diffusion step. To afford the extra compute this costs, they downsample the diffusion process from 2,000 steps to 200 for controllable generation, which the paper reports speeds up decoding without hurting sample quality much.

Worked example: why this beats PPLM by ~60×, mechanically

The paper reports Diffusion-LM's controllable generation is about 1.5× slower than FUDGE but roughly 60× faster than PPLM. Trace where that factor actually comes from, step by step, rather than taking it on faith.

Diffusion-LM's total gradient-step count, for a 64-token sequence: 200 diffusion steps × 3 Adagrad steps each = 600 gradient updates, total, for the entire generation.

PPLM's total gradient-step count, for the same 64 tokens: the paper's own speed appendix reports PPLM needs 30 gradient-ascent steps per token, with no batching across positions:

By raw count, PPLM does about 3.2× more gradient steps than Diffusion-LM (1,920 vs. 600) — nowhere near the full 60× gap. The rest of the gap comes from structure, not count: Diffusion-LM's 600 updates each touch all 64 positions at once, in parallel, inside a single forward/backward pass. PPLM's 1,920 updates are locked into strict per-token sequential order — token 2's activations can't be steered until token 1 has already been finalized, because that's how an autoregressive forward pass works. Fewer total steps, applied in parallel, beats more total steps, applied one position at a time, compounding into the reported ~60× wall-clock difference (PPLM: about 80 minutes for 50 samples, unbatched; Diffusion-LM: roughly 1.5× FUDGE's 50 seconds, so on the order of 75 seconds for the same 50 samples).

Worked example: steering toward a syntax tree

Table 1 of the paper gives a concrete target: the syntax tree (TOP (S (NP (*)(*)(*)) (VP (*) (NP (NP (*)(*)))))), and Diffusion-LM's output for it: “The Twenty Two has great food.” Here is what “steering denoising” concretely means for a target like this. At each of the 200 diffusion steps, a syntax-parsing classifier — trained directly on Diffusion-LM's own latents, so it understands the same representation the denoiser does — computes ∇_xt-1 log p(c|x_t-1): a direction, for every position in the sequence simultaneously, that would make the resulting parse closer to the target tree. That gradient gets added (scaled by λ) to the ordinary denoising direction, and 3 Adagrad steps are taken along the combined direction before moving to the next diffusion step. Over 200 iterations of this tug-of-war between “stay fluent” and “match the tree,” the sentence converges on something that does both — and because every position updates together, an early phrase can still shift in response to a later constraint, which Chapter 0 established a left-to-right steering method structurally cannot do.

Worked example: length control needs no classifier at all

One control target gets special treatment. For Diffusion-LM specifically, length is classifier-free — you never train a separate p(c|x) for it. The reason is structural: sequence length is fixed the moment you decide how many latent positions to start denoising, before the process even begins. Set n = 14 positions, and every valid output has exactly 14 words, by construction, with no classifier gradient required. Other methods (PPLM, FUDGE, the fine-tuning baselines) don't have this shortcut — an autoregressive model decides word by word without knowing in advance where “the end” should land, so it needs an explicit length signal, usually a classifier, to hit a target.

The reported numbers make the structural advantage concrete:

Near-perfect length control, with no classifier, because the constraint is already baked into how many latent positions exist — there was never anything to steer toward.

Two more control tasks, same pattern

On both structured tasks, Diffusion-LM doesn't just beat the plug-and-play baseline (FUDGE) — it beats the fine-tuning oracle, a model directly trained on (control, text) pairs, which isn't even plug-and-play. Chapter 4 explains the mechanism behind that surprising result in full.

Why Adagrad, specifically

The paper's own footnote on this choice is worth surfacing: they tried replacing Adagrad with plain SGD for the 3-steps-per-diffusion-step control optimization, and found Adagrad “substantially less sensitive to hyperparameter tuning.” Adagrad adapts its effective step size per-coordinate based on the accumulated history of past gradients — coordinates that have seen large gradients get smaller future steps, coordinates that have seen small gradients get larger ones. For a control objective combining two very differently-scaled gradient sources (a smooth fluency term from a large pretrained diffusion network, and a sharper, more volatile control term from a much smaller classifier trained on far less data), that per-coordinate adaptivity matters more than it would for training the base model itself, where a single well-tuned learning rate already works. This is a small implementation detail with a real consequence: get the optimizer choice wrong here and the λ balance from the widget above would need re-tuning for every new control task.

Worked example: FUDGE's call count, for comparison

Chapter 0 introduced FUDGE as a plug-and-play baseline; here is exactly what makes it slower than Diffusion-LM despite being lighter-weight per call. FUDGE reweights each next-token prediction using a lightweight classifier that scores every candidate continuation — the paper's speed appendix reports FUDGE requires 200 classifier calls per token, which for a 64-token sequence is:

These calls can be batched (unlike PPLM's sequential gradient steps), which is why FUDGE lands at a middle speed — about 50 seconds for 50 samples, faster than PPLM's 80 minutes but still slower than Diffusion-LM's roughly 75 seconds despite Diffusion-LM doing genuine gradient optimization rather than FUDGE's cheaper reweighting. The lesson: raw operation count is not the whole story once you factor in what can run in parallel versus what's forced sequential — Chapter 8 returns to exactly this distinction when comparing LLaDA's sampling-step count against autoregressive decoding.

Worked example: steering toward a part-of-speech sequence

A second, more mechanical worked example, using the paper's own POS control task. Given a target tag sequence — say PRONOUN VERB DETERMINER NOUN — the classifier p(c|x_t-1) is a POS tagger trained directly on Diffusion-LM's continuous latents, one that has learned to read off, from a noisy intermediate x_t-1, a distribution over which grammatical category each position is drifting toward. Early in denoising (t close to T), the classifier's gradient is necessarily coarse — it can only say “position 2 needs to look more verb-like in general,” because the noisy latent doesn't yet distinguish between, say, “ate” and “ran.” Late in denoising (t close to 0), the gradient sharpens into something much more specific, because by then the latent is close enough to a real embedding that the classifier can discriminate between individual verb choices, not just the coarse category. This coarse-to-fine structure is exactly why the paper attributes Diffusion-LM's strength on POS and Syntax Tree control specifically to its hierarchy of intermediate representations — a control signal that's naturally forced to operate at the right granularity for how confident the model actually is at each point in denoising, something a single-pass classifier reweighting an autoregressive model's output never gets.

What the four classifiers actually are

Every control task in Chapter 3's tables leans on a purpose-built classifier trained on Diffusion-LM's latents. It's worth knowing, concretely, what each one is, rather than treating “the classifier” as an unexamined black box:

Notice the pattern: none of these four are exotic. Every one is an ordinary, independently-trainable classifier or tagger, of a kind that existed and worked well before this paper — the innovation isn't in the classifiers, it's in where they're applied (Diffusion-LM's continuous latents, at every diffusion step) and how their gradients get combined with the fluency term (Chapter 3's core equation). This matters for reproducibility and generality: adding a new control task to Diffusion-LM doesn't require a new kind of machinery, just a new classifier trained on the same shared latent space the existing four already use.

The hyperparameters actually tuned, for a fair comparison

Every method in Chapter 3's tables gets its own tunable knob, searched over a grid, so the comparisons aren't an accident of one method being tuned harder than another. Diffusion-LM tunes its Adagrad learning rate over lr ∈ {0.05, 0.1, 0.15, 0.2} and its fluency tradeoff over λ ∈ {0.1, 0.01, 0.001, 0.0005} per task. PPLM tunes its per-token gradient-update count k ∈ {10, 30} (recall Chapter 3's PPLM worked example used k=30, the upper end of this exact range). FUDGE tunes its own tradeoff parameter λ_FUDGE ∈ {16, 8, 4, 2}. Every method's other hyperparameters — PPLM's learning rate (0.04) and KL-scale (0.01); FUDGE's precondition top-K (200) and post top-K (10) — are set by following each original paper's own recommended defaults, not re-tuned in Diffusion-LM's favor. This is worth knowing before trusting any of Chapter 3's control-success numbers: the comparison was built to be fair on its own terms, not stacked by under-tuning the competition.

Worked example: reading the λ grid against what actually gets composed

It's worth noticing something in the reported λ grid that connects back to Chapter 3's tables directly. Diffusion-LM's per-task λ values, drawn from {0.1, 0.01, 0.001, 0.0005}, span more than two orders of magnitude — meaning the optimal fluency/control balance is not remotely the same number across tasks. Length control (nearly classifier-free, structurally almost guaranteed to succeed) can afford a smaller λ and let the control term dominate more, since there's little risk of the constraint fighting fluency too hard. Syntax Tree control, a much harder global constraint actively competing with grammaticality at every position, needs careful tuning of exactly how much fluency pressure to apply so the sentence doesn't degenerate while still bending toward the target structure. This variation is a quiet piece of evidence for something Chapter 3's widget already demonstrated visually: the λ tug-of-war isn't a fixed dial you set once for “controllable generation” in general — it's a per-task tradeoff that has to be re-discovered whenever the nature of the constraint changes, because how hard a constraint fights fluency depends entirely on what the constraint is asking for.

What “training the classifier on the latents” actually requires

One subtlety worth making explicit, since it's easy to read past: every classifier in Chapter 3's table — the GPT-2-small semantic classifier, the BERT-base POS tagger, the Transformer parser — is trained to operate on Diffusion-LM's own noisy intermediate latents x_t, not on clean text. That means training these classifiers requires running the forward diffusion process on labeled training data first (add noise at some sampled t to a real sentence's embedding, exactly as Chapter 1 defined), and training the classifier to still recover the correct label from that noised representation. A classifier trained only on clean, unnoised text would have no idea what to do with a half-denoised x_t-1 at diffusion step 1,500 — it would be evaluating on data unlike anything it saw during its own training. This is a real extra cost the plug-and-play framing sometimes obscures: “keep the language model frozen and add a classifier” sounds cheap, but the classifier itself has to be trained specifically against the noise distribution the base generative model produces, which ties the classifier's training data generation process to the base model in a way a classifier trained purely on clean AR-generated text wouldn't need.

Composing more than two controls, and where the mechanism would strain

Chapter 4's composition results only test pairs of controls (semantic+syntax, semantic+POS). It's worth reasoning through what the gradient-summing mechanism predicts for three or more simultaneous constraints, since the paper itself doesn't test beyond two. The mechanism — summing ∇ log p(c_i|x_t-1) across every active constraint i, applied to one shared latent — scales additively in principle: nothing in the gradient-sum formula caps the number of terms. But the fluency term λ log p(x_t-1|x_t) does not scale up alongside them unless λ is explicitly retuned, so a naive three- or four-way composition risks the same failure mode Chapter 3's widget demonstrated for a single control at too-low λ: the combined control gradient overwhelming the fluency anchor, producing text that satisfies several narrow constraints while reading as disfluent overall. This is a reasoned extrapolation from the mechanism this session derived, not a result either paper reports directly — exactly the kind of distinction Chapter 9's reading checklist asks you to keep clear.

Chapter 3, in one paragraph

Controllable generation on Diffusion-LM works by treating the whole latent trajectory, not the final discrete text, as the thing being controlled: at every diffusion step, a classifier's gradient (measuring how well the current latent satisfies a target constraint) gets added to the ordinary denoising gradient, weighted by a fluency term λ that keeps the result readable rather than merely constraint-satisfying. Because that steering happens on a representation the whole sequence shares simultaneously, an early word can still respond to a constraint that only fully resolves many positions later — the exact capability Chapter 0 argued left-to-right steering methods structurally cannot have. The next chapter asks how well this actually pays off, against real baselines, on real numbers.

Chapter 4: Diffusion-LM's Report Card

Chapter 3 showed Diffusion-LM winning on control. This chapter asks the harder, more honest question: what did it cost, and does the win survive scrutiny? Read the paper's own numbers straight, including the ones that don't flatter it.

The likelihood gap, converted to something you can feel

In negative log-likelihood (NLL, lower is better — and because exact likelihood is intractable for a diffusion model, what's actually reported is the variational upper bound ℒ_vlb^e2e from Chapter 1, which could overstate the true gap somewhat), Diffusion-LM loses to an equivalent autoregressive Transformer on both datasets:

Nats-per-token is not an intuitive unit. Convert to perplexity (PPL = e^loss, recall from earlier in this course) to feel the gap the way you'd feel it reading the model's actual output:

Diffusion-LM is a noticeably worse pure language model by this measure — roughly 1.7× the perplexity on E2E, over 2× on the harder ROCStories task. Scaling up model and dataset size partially closes the gap (the paper reports ROCStories NLL improving from 3.88 to 3.10 with a larger model), which converts to:

Hold both facts at once: this session's opening didn't promise diffusion language models would be better raw next-word predictors. It's worth being precise about which claim actually held up.

Why worse likelihood didn't stop it from winning on control

Chapter 3's numbers weren't a fluke of one lucky task. Diffusion-LM “almost doubles the control success rate of prior plug-and-play methods” across all five classifier-guided tasks the paper tests, and on two of them — Syntax Tree and Syntax Spans — it beats even the fine-tuned oracle, a model that isn't plug-and-play at all. This looks contradictory until you separate what each number is actually measuring: likelihood asks “how well does the model predict ordinary held-out text,” while control success asks “can the model be steered, after the fact, toward a specific structural target.” Those are different questions, and a model can be a slightly worse general-purpose predictor while being structurally much better suited to steering — because it sees and can revise the entire sequence at every denoising step, instead of committing to early tokens (autoregressive) or a single unrevisable pass (fine-tuning, in the standard sampling or beam-search setting) before later constraints are even visible.

The paper's qualitative example makes this concrete. Steering toward a target syntax tree, Diffusion-LM first generates the span “Family friendly Indian food” — one word too long for its target slot. Rather than propagating that error forward the way a left-to-right model would, Diffusion-LM's later denoising steps correct for it: the model drops a conjunction later in the sentence to compensate, and the suffix ends up matching the target tree anyway. The fine-tuned baseline, hitting an analogous one-word-short error in a different example (“The Mill”), never recovers — the mismatch propagates through the rest of the generated span. Self-correction, made possible by seeing the whole sequence at every step, is the mechanism behind the surprising win over the oracle.

Composing controls: where the mechanism shows up again

Given classifiers for two independent tasks, plug-and-play control composes them just by summing gradients on the same shared latent — no retraining required:

FT-PoE multiplies two separately fine-tuned models' probabilities at inference — and it composes badly, because neither model was ever trained expecting the other's constraint; their preferences fight each other. Diffusion-LM's approach sums gradients onto one shared latent within a single coherent denoising trajectory, so both constraints get reconciled together rather than negotiated between two separate models after the fact — the same “see everything, revise together” mechanism from the last section, applied to two controls instead of one.

Infilling: the classifier-free case, with Minimum Bayes Risk decoding

Recall Chapter 0's aNLG task — generate a middle sentence connecting a known left context (O₁) and right context (O₂). Like length, infilling gets no classifier at all. Instead, Diffusion-LM samples a set 𝒮 of candidate completions and applies Minimum Bayes Risk (MBR) decoding: pick the candidate that minimizes expected loss against every other candidate in the set, ŵ = argmin_w∈𝒮 ∑_w′∈𝒮 (1/|𝒮|) ℒ(w,w′), using negative BLEU as ℒ. The intuition: a bad, weird sample looks dissimilar from the rest of the batch and gets penalized by the loss; a good, central sample resembles most of the others and wins.

Diffusion-LM roughly matches the AR-infilling baseline — the one Chapter 0 flagged as needing its training data physically reordered to work around the left-to-right constraint — while significantly outperforming both energy-based baselines (COLD, DELOREAN), and does it without any specialized retraining for the infilling task at all.

Worked example: MBR decoding, by hand, on a toy candidate set

The MBR formula from earlier in this chapter — ŵ = argmin_w∈𝒮 ∑_w′∈𝒮 (1/|𝒮|) ℒ(w,w′) — is easiest to trust once you've run it on numbers small enough to check by hand. Suppose Diffusion-LM samples four candidate completions for the tennis-balls infilling task from Chapter 0, and a simplified loss ℒ(w,w′) = 1 − BLEU(w,w′) ∈ [0,1] (0 = identical, 1 = no overlap at all) scores every pair:

Candidate 2 wins — lowest average loss against the rest of the set, at 0.375. Candidate 4, the one weird outlier, is penalized hard: precisely because it doesn't resemble anything else the model sampled, its average distance to the rest of the batch is large. This is the mechanism behind the paper's own intuition, made concrete with actual numbers instead of just asserted: MBR doesn't need to know why candidate 4 is bad — degenerate, off-topic, ungrammatical, or just an unlucky sample — it only needs candidate 4 to look unlike what the model otherwise tends to produce for this prompt, which a genuinely bad sample usually does.

Worked example: the correction, narrated in full

Chapter 3's control walkthrough mentioned Diffusion-LM's self-correcting behavior in passing; it's worth narrating the paper's own qualitative example completely, because it's the clearest illustration in this whole session of what “seeing the entire sequence at every step” buys you in practice. The target parse tree calls for a specific span structure. Partway through denoising, the model's current best guess for one span is “Family friendly Indian food” — four words, one more than the target span allows. Under strict left-to-right generation, this is already a fixed mistake: whatever comes next has to live with a span that's structurally wrong, because there's no mechanism to reach back and shorten something already committed. Diffusion-LM has no such constraint — later denoising steps are free to revise any position, including ones the model output would suggest are “already decided.” The actual fix the model finds: it drops a conjunction later in the sentence, which brings the total word count for the constrained span back in line with the target, even though the fix happens at a completely different position than the original error. The fine-tuned baseline, hitting an analogous problem in a different example — a span one word short, “The Mill” instead of the required longer phrase — has no equivalent recovery mechanism (standard sampling and beam search both commit to tokens as they're generated), and the mismatch propagates through every downstream span for the rest of the sentence, degrading the output further with each subsequent word.

Worked example: what the full 2,000-step decode costs, and why 200 is enough for E2E but not ROCStories

Chapter 3 established the 200-step downsampled schedule used for controllable generation specifically. It's worth being precise about where that number came from and where it stops working. The paper's decoding-speed analysis reports about 1 minute to decode 50 sequences of length 64 using the full, undownsampled T=2,000 schedule — the number this chapter's “7× slower than AR” comparison is built on. Downsampling to 200 steps (by setting ᾱ_t = ᾱ_10t, effectively treating every 10th step of the original schedule as one new step) speeds this up roughly 10×, and the paper reports this naive downsampling doesn't hurt sample quality much for the simpler E2E task — but it does measurably hurt quality on the harder ROCStories task, which is why the paper keeps the full 2,000-step schedule for ROCStories generation and control experiments specifically, accepting the extra decoding cost in exchange for not degrading the harder dataset's output. The lesson generalizes past this one paper: how aggressively you can compress a diffusion process's step count is itself a property of how hard the underlying distribution is to model, not a universal constant you can tune once and reuse everywhere.

Worked example: verifying “almost doubles the success rate” by hand

The paper's abstract claims Diffusion-LM “almost doubles the success rate of previous plug-and-play methods.” Rather than taking that on faith, average the ctrl% numbers from Chapter 3's tables across all five classifier-guided tasks, for FUDGE (the stronger of the two plug-and-play baselines) versus Diffusion-LM:

The claim checks out almost exactly — a genuinely rare thing for a paper abstract's rounded-off summary phrase to survive this literally when you go back and compute it from the paper's own reported numbers table by table.

The honest costs, straight from the paper's own conclusion

The paper doesn't hide its limitations, and neither should this lesson. Three, stated plainly: (1) higher perplexity — the gap this chapter opened with. (2) decoding is substantially slower — even after downsampling from 2,000 to 200 steps, it takes about 1 minute to decode 50 sequences of length 64, roughly 7× slower than decoding the equivalent autoregressive model. (3) training converges more slowly than an equivalent AR model.

The third cost is worth a mechanistic explanation, since it's easy to state and easy to skip past. An autoregressive model's training signal at each gradient step is “predict this one next token, given everything before it” — a single, comparatively narrow prediction problem, repeated across positions. A diffusion model's training signal at each gradient step is “denoise this latent, at this particular noise level, back toward the clean data” — and because t is sampled randomly at every step (Chapter 1), the network has to become competent across the entire range of noise levels simultaneously, from nearly-clean (t near 0, where the job is subtle refinement) to nearly-pure-noise (t near T, where the job is closer to unconditional generation from scratch). That's a strictly harder, higher-variance learning problem for the same number of gradient steps, and it's a large part of why the paper reports slower convergence — not a training bug, but a direct consequence of what the objective is asking the network to learn across every diffusion timestep at once.

Three independent costs, not one

It's tempting to treat “Diffusion-LM is worse than AR” as a single verdict, but this chapter has actually surfaced three genuinely separate costs, and it's worth being precise that fixing one doesn't automatically fix the others. Scaling up model and dataset size (this chapter's ROCStories NLL improving from 3.88 to 3.10) addresses the likelihood cost — but says nothing directly about the decoding-speed cost, which is a property of how many diffusion steps a sample needs, not of how good the underlying network is; a bigger, better-trained network still has to take roughly the same number of denoising steps to produce a sample of comparable quality. Similarly, the 200-step downsampling trick that addresses decoding speed for E2E specifically doesn't transfer cleanly to ROCStories (this chapter's earlier worked example), so a fix for one dataset's speed problem isn't automatically a fix for a harder dataset's speed problem. And neither of those two fixes touches the third cost, slower training convergence, which is a property of the objective's per-gradient-step difficulty (the previous section's derivation), not of scale or step count. Three costs, three largely independent levers — a useful discipline for reading any paper's reported “X is worse than Y” result: ask whether the fix for one weakness is expected to help the others, or whether they need to be addressed one at a time.

A road not taken: fixed embeddings for better perplexity

One honest footnote in the paper is worth surfacing, because it complicates the clean “learned embeddings win” story from Chapter 1 in an interesting way. The paper reports that while trainable end-to-end embeddings perform best on control and generation quality (everything Chapters 3 and 4 have shown), a different choice — fixed embeddings placed onto the vocabulary simplex, not random Gaussian vectors, not learned, but a specific fixed geometric construction — were found to help specifically with held-out perplexity, the exact likelihood metric this chapter has been comparing unfavorably to autoregressive models throughout. The paper explicitly set this direction aside, stating plainly that “the focus of this work is generation quality and not perplexity.” That's a deliberate scope decision, clearly labeled as such, not an oversight — but it's worth knowing it exists, because it suggests the likelihood gap this chapter has treated as a cost of the diffusion approach in general might, at least partially, be a cost of the specific embedding choice this paper optimized for control quality, rather than an unavoidable property of continuous text diffusion itself. A different paper, optimizing for the opposite tradeoff, might report a smaller likelihood gap and a weaker control result — the same underlying machinery, tuned toward a different target.

A short scorecard, before moving to LLaDA

Four chapters, six datasets and metrics, a lot of numbers. Reduce it to the essential shape before Chapter 5 starts a new architecture from scratch: Diffusion-LM loses on raw likelihood (Chapter 4, this chapter), wins decisively on structured classifier-guided control (Chapter 3: 2.09× FUDGE's average success rate, and outright beating a fine-tuned oracle on the two hardest structural tasks), and roughly ties on classifier-free tasks it was never expected to dominate outright (Chapter 4's infilling result, matching rather than beating specialized AR-infilling). Three different verdicts for three different kinds of task, not one uniform “ diffusion is better” or “diffusion is worse” — which is exactly the nuance a reader skimming only the abstract's “almost doubles the success rate” headline would miss entirely.

One question this scorecard leaves for Chapter 5 to answer

Every number in this chapter comes from an 80-million-parameter model trained on two datasets small enough to fit comfortably on a single GPU for a single afternoon (this chapter's own 5-hour, one-A100 training estimate). That scale is small enough to raise an honest question this chapter cannot answer on its own: does control-task superiority survive when the underlying model gets dramatically larger, is trained on dramatically more data, and stops being a narrow proof of concept aimed at two specific benchmarks? Nothing about the mechanism this chapter identified — seeing and revising the whole sequence at every step — is obviously scale-dependent, but nothing guarantees it either. That's precisely the gap Chapter 5 exists to close, at a scale one hundred times larger and with an entirely different, simpler corruption process.

Chapter 5: LLaDA at Language-Model Scale

Diffusion-LM proved the concept at 80 million parameters on two small, narrow datasets. Three years later, LLaDA (Large Language Diffusion with mAsking, Nie et al., 2025) asks a much blunter question at a much bigger scale: train an 8-billion-parameter diffusion language model completely from scratch — no continuous embeddings, no rounding step, none of Chapters 1–2's machinery — using a radically simpler discrete formulation, and see if it can genuinely compete with a real, deployed-scale autoregressive model.

Masked diffusion, derived from scratch

Instead of embedding words into a continuous space and adding Gaussian noise, LLaDA corrupts a sequence directly and discretely: replace some tokens with a literal [MASK] token. For a clean sequence x₀ of length L, define the forward process per token, independently, for time t ∈ [0,1]:

At t = 0, every token is unmasked with certainty — the clean sequence, no corruption. At t = 1, every token is masked with certainty — a Dirac distribution on the fully-masked sequence. This is dramatically simpler than Diffusion-LM's pipeline: no embedding space, no Gaussian noise, no argmax-rounding ambiguity from Chapter 2 at all — a masked position is just, literally, the [MASK] token, unambiguous by construction.

How this differs from BERT, precisely

If “randomly mask some tokens and predict them” sounds like BERT, that's not a coincidence — but one design choice separates them completely. BERT masks a fixed ratio, about 15%, on every training example. LLaDA samples t uniformly from [0,1] for every example — sometimes almost nothing is masked, sometimes almost everything is. That single difference is what makes LLaDA a genuine generative model rather than a representation-learning objective: training across the full range, including t very close to 1 (nearly everything masked), teaches the network to reconstruct a sequence from almost no visible context at all — exactly the situation at the very start of generation, when the response begins as a fully-masked blank. BERT, trained only ever at a light 15% masking ratio, was never asked to operate in that near-total-masking regime, and was never designed to sample text from scratch.

The training loss, and why it's reweighted by 1/t

The mask predictor p_θ(·|x_t) is a bidirectional Transformer — crucially, no causal mask, so it can see the entire corrupted sequence, tokens both before and after any given position, when predicting what a masked position should be. It's trained with cross-entropy, computed only on the actually-masked positions:

The indicator 1[x_tⁱ=M] restricts the sum to masked positions only — the model is never penalized for tokens it can already see. The 1/t factor out front is the part worth deriving intuition for.

Worked example: what 1/t actually corrects for

Take a 100-token training sequence, sampled twice at two different masking ratios. At t = 0.1, in expectation about 10 tokens are masked — an “easy” reconstruction problem, most of the context is visible. At t = 0.9, in expectation about 90 tokens are masked — a “hard” problem, almost nothing is visible.

The raw, un-reweighted sum ∑_i 1[masked] log p_θ(·) would naturally be smaller at t = 0.1 just because there are fewer masked terms being summed — an artifact of how many terms happened to get included, not a meaningful statement about how confident or accurate the model is per token. The 1/t reweighting compensates directly:

The few masked tokens in the easy (t=0.1) case get their loss amplified 10×; the many masked tokens in the hard (t=0.9) case get amplified only about 1.11×. This isn't a cosmetic balancing trick — it's exactly the factor that makes ℒ(θ) a mathematically valid upper bound on the true negative log-likelihood of the data distribution, rather than an arbitrary masked-reconstruction heuristic with no principled connection to log p(x₀). Chapter 6 derives why the 1/t factor specifically is what's required.

Worked example: the scale, in numbers you can check

LLaDA 8B has 8.02 billion total parameters, of which 6.98 billion are non-embedding. The roughly 1-billion-parameter gap is mostly the vocabulary embedding table: 126,464 vocabulary entries × 4,096-dimensional hidden size × 2 bytes (fp16) per number:

— roughly matching the ~1.04B-parameter gap between total and non-embedding counts once you account for both the input embedding table and a separate, untied output projection of the same size.

Pretraining ran on 2.3 trillion tokens, using 0.13 million H800 GPU-hours. Divide one by the other for a throughput sanity check:

Supervised fine-tuning followed on 4.5 million (prompt, response) pairs, for 3 epochs, with the SFT version of the same loss — the prompt tokens are left unmasked, and only the response tokens are ever masked during training, so the model learns to fill in an answer conditioned on a fixed, always-visible prompt.

Architecture: a close cousin of LLaMA3, with one deliberate difference

LLaDA borrows LLaMA3's recipe almost wholesale — RMSNorm, SwiGLU activations, RoPE positional encoding — but makes one pointed departure: vanilla multi-head attention (32 query heads, 32 key/value heads) instead of LLaMA3's grouped-query attention (32 query heads, only 8 key/value heads, which shrinks the KV cache). The reason isn't an oversight — it's explained fully in Chapter 8, but the short version: LLaDA is architecturally incompatible with KV caching in the first place (its bidirectional attention means every position's representation can change at every sampling step), so there's no cache to shrink, and GQA's whole purpose disappears.

Nearly identical parameter budgets everywhere else — the KV-head count is the single number that gives away the whole architectural bet this model is making.

Worked example: the training FLOPs, and where Chapter 9's compute cap comes from

The paper reports using the standard 6ND formulation (Kaplan et al. 2020; Hoffmann et al. 2022) to estimate training compute, where N is the count of non-embedding parameters and D is the total number of training tokens. Plug in LLaDA 8B's own numbers from this chapter:

That number lands just under 10²³ FLOPs — and Chapter 9 previews the paper's own admission that its direct, matched comparisons against an autoregressive baseline were capped at exactly that threshold, under 10²³ FLOPs, for resource reasons. This isn't a coincidence worth glossing over: LLaDA 8B's own full pretraining run sits right at the edge of the compute budget the paper says it could afford to compare cleanly. Scaling LLaDA any further — more parameters, more tokens, or both — would push past the range where the paper's own reported comparisons remain apples-to-apples, which is exactly the limitation Chapter 9 names directly: the 8B headline results are not a compute-matched contest against an equivalently-scaled autoregressive model, because affording that second model at the same budget wasn't possible alongside training LLaDA 8B itself.

Worked example: the learning-rate schedule, stage by stage

LLaDA's pretraining uses a Warmup-Stable-Decay (WSD) schedule, and the paper reports its exact milestones — worth laying out as a concrete timeline over the full 2.3T-token run, since “we used a WSD schedule” on its own tells you nothing about what actually happened during training.

Two things worth noticing in this timeline. First, the “stable” middle stage — roughly 1.2 trillion tokens, over half the entire pretraining run — sits at a single fixed learning rate with no decay at all, which is the whole point of WSD: it lets training be extended or shortened without needing to know the final token budget in advance, unlike a cosine schedule that has to commit to a decay endpoint up front. Second, both decay stages together (0.8T + 0.3T = 1.1T tokens) are nearly as long as the stable stage itself — the decay isn't a brief cooldown tacked onto the end, it's a substantial fraction of the whole budget, deliberately.

The optimizer is AdamW, weight decay 0.1, global batch size 1,280, with a local batch size of just 4 sequences per GPU — meaning the 1,280 global batch is assembled by gradient-accumulating or all-reducing across a large number of GPUs simultaneously (1,280 ÷ 4 = 320 GPU-equivalent slots, at minimum, before accounting for any additional accumulation steps). The paper is explicit that the full 8B run was executed once, with no hyperparameter tuning — a detail worth taking seriously when reading any of this session's LLaDA numbers: nothing here benefited from a second attempt.

Worked example: variable-length training, and why 1% is enough

Pretraining uses a fixed sequence length of 4,096 tokens for the overwhelming majority of training, but the paper deliberately sets aside 1% of the pretraining data to a random length uniformly sampled from [1, 4096], following prior work on variable-length handling. The reasoning: a model trained exclusively at one fixed length has no direct training signal for what a much shorter (or oddly-sized) sequence should look like at inference time, and the paper reports this small fraction is enough to teach the model to handle variable generation lengths robustly — consistent with Section 2.4's claim that final results are largely insensitive to the generation-length hyperparameter (Chapter 7 verifies this with the Appendix B.4 ablation directly).

SFT follows a parallel but separately-tuned schedule: linear warmup from 0 to 2.5×10⁻⁵ over the first 50 iterations, held constant, then linearly decayed to 2.5×10⁻⁶ over the final 10% of iterations — the same overall shape as pretraining's WSD schedule, but compressed to fit 3 epochs over 4.5 million pairs instead of a single pass over 2.3 trillion tokens. Global batch size drops to 256, local batch size to 2 per GPU, and weight decay stays at 0.1.

Worked example: turning a multi-turn dialogue into training pairs

One SFT detail worth understanding concretely, since it explains how LLaDA learns to hold a multi-turn conversation using an objective that only ever conditions a response on a single fixed prompt. Given an n-turn dialogue (p₀⁰, r₀⁰, p₀¹, r₀¹, …, p₀^n-1, r₀^n-1), the paper partitions it into n separate single-turn (prompt, response) training pairs, where each successive pair's “prompt” is the entire conversation history up to that point:

and then randomly samples one of these n pairs per training example, rather than using every pair from every dialogue every time. This means “single-turn SFT” and “multi-turn capability” are not actually in tension — every training example genuinely is a single (prompt, response) pair in the exact same shape Section 2.3's loss expects, it's just that the “prompt” for a later turn happens to contain the entire visible conversation history as ordinary unmasked context. Nothing about the training objective itself had to change to support multi-turn dialogue; only how the training data gets sliced up did.

Worked example: the SFT loss, on a toy (prompt, response) pair

Ground the SFT loss the same way Chapter 6 will later ground the pretraining loss — with actual numbers. Take a toy pair: prompt p₀ = (“translate”, “hello”), response r₀ = (“bonjour”, “|EOS|”), so L′ = 2 response tokens. The prompt is never masked — only r₀ is subject to the forward process, at some sampled t. Suppose t = 0.5 masks just the first response token, giving r_t = (M, “|EOS|”), fed to the network together with the fully-visible prompt.

Suppose the mask predictor, seeing the prompt “translate hello” and the still-visible “|EOS|”, assigns probability 0.7 to the correct completion “bonjour” at the one masked position. The loss contribution for this single training example is:

Notice what this toy example makes concrete about “controlling response length by learning to generate |EOS|”, a claim Chapter 5's earlier discussion of SFT padding made in passing: the second response position, “|EOS|”, is treated as an ordinary token subject to the exact same masking and loss machinery as “bonjour” — it just happened not to be masked in this particular draw. On a different training draw, at a higher t, |EOS| itself might be the masked token the network has to predict, which is exactly how the model learns to place it at the right position: by being asked to reconstruct it from context, the same way it reconstructs any other word.

Why no causal mask, stated as a one-line architectural diff

It's worth being precise about exactly what code changes when you strip the causal mask out of an otherwise ordinary Transformer, since “bidirectional attention” can sound like a bigger architectural leap than it actually is. An autoregressive Transformer's self-attention applies a triangular mask before the softmax — position i's attention scores against positions j>i get set to −∞ before normalizing, so those positions receive exactly zero attention weight, structurally guaranteeing position i never sees anything after it. LLaDA's mask predictor removes that triangular mask entirely: every position attends to every other position, mask token or not, with no directional restriction at all. The rest of the Transformer block — the Q/K/V projections, the softmax, the FFN, RMSNorm, RoPE — is unchanged. This is worth holding next to Chapter 0's chain-rule argument: the causal mask is the autoregressive factorization, implemented as a single line of attention-score masking, not an emergent property of anything else in the architecture. Removing that one line is the entire mechanical difference between an autoregressive Transformer and LLaDA's bidirectional one — everything else this chapter has described (the masking process, the loss, the reweighting) is about what you train that bidirectional Transformer to do, not about the Transformer block itself.

What positional encoding has to mean differently without a fixed order

One more consequence of dropping the causal mask, easy to overlook: RoPE (rotary position embeddings) still tells the network where each token sits in the sequence, but “where a token sits” no longer implies anything about generation order the way it does for an autoregressive model. In an AR Transformer, position and generation order are the same fact viewed two ways — token 5 is both the fifth position and the fifth token generated. In LLaDA, position is purely a spatial coordinate: token 5's position tells the network where it sits relative to its neighbors for the purpose of computing attention, but says nothing about whether it was committed at sampling step 1 or sampling step 200 (Chapter 7's remasking schedule decides that, independently of position). This is a subtle but real conceptual split that Chapter 9's “no diffusion-native architecture” limitation gestures at directly: RoPE was designed for, and tuned around, a setting where position and order coincide, and nothing in LLaDA's borrowed positional encoding was redesigned to reflect that those two concepts have come apart.

Chapter 6: The Likelihood Bound

Chapter 5 asserted that ℒ(θ) upper-bounds the true negative log-likelihood of the data. This chapter derives the key steps behind that claim, and shows why it isn't a technicality — it's the reason LLaDA gets to call itself a principled generative model at all, rather than a masked-reconstruction heuristic that happens to also generate text.

The reverse process, worked out per token

For 0 ≤ s < t ≤ 1, the true reverse transition q_s|t(x_s|x_t) factorizes across positions, and each position follows exactly one of four cases:

Read each case as a plain sentence. Case 1: a token that's already unmasked stays exactly as it is — once revealed, under the true reverse process, a token is never re-masked or changed. Case 2: a masked token stays masked, with probability s/t. Case 3: a masked token “resolves” to an actual word, with probability (t−s)/t, and conditional on resolving, the value it resolves to follows the distribution q_0|t. Case 4: everything else — an unmasked token spontaneously becoming masked, or resolving to the wrong value — has probability exactly zero. Impossible, by construction of the forward process.

Sanity-checking the formula at both boundaries

Before trusting a four-case formula, do what Chapter 1 did with the noise schedule: check both ends by hand.

At s = t (no time has elapsed — the “transition” should do nothing): Case 2's probability is s/t = t/t = 1 — a masked token stays masked with certainty. Case 3's probability is (t−s)/t = (t−t)/t = 0 — a masked token never resolves. Exactly what you'd want: zero elapsed time means zero chance of anything changing.

At s = 0 (the reverse process has fully finished — every mask should be resolved): Case 2's probability is s/t = 0/t = 0 — a masked token never stays masked. Case 3's probability is (t−s)/t = (t−0)/t = 1 — a masked token resolves with certainty. Again exactly right: by the time you reach s=0, nothing should remain uncertain, and the formula enforces that on its own, without needing to special-case the endpoint.

Both checks pass with the formula exactly as stated, no adjustment needed — a genuinely reassuring sign before building an entire sampling algorithm (Chapter 7) on top of it.

The time-invariance shortcut

The one piece left to estimate is q_0|t(x_sⁱ|x_t) — what does a mask resolve to? A key theoretical result the paper leans on (Ou et al. 2024) is that this distribution has an equivalent, time-free form:

where x_t^UM is the collection of currently-unmasked tokens. Notice: t does not appear on the right-hand side at all. The distribution over what a mask resolves to depends only on which tokens are currently visible, not on the specific value of t at which you're asking. This is a genuine simplification diffusion for continuous data doesn't get for free — it's why LLaDA's mask predictor takes no time-embedding input whatsoever, unlike Diffusion-LM's network, which explicitly conditions on t at every step (Chapter 1). The masking formulation buys this simplicity as a direct consequence of its structure, not as a separate engineering choice.

Worked example: two equivalent loss forms, wildly different sample efficiency

The main text's ℒ(θ) samples t uniformly, then masks each token independently with probability t. An equivalent form in the paper's appendix instead samples an integer l uniformly from {1,…,L} and masks exactly l tokens, chosen uniformly without replacement:

These two forms have the same expectation — but the paper reports the first needs over 1,000 Monte Carlo samples to estimate reliably, while the second reaches stability with only 128. Here's the mechanism, worked by hand: under independent per-token masking at ratio t, the actual number of masked tokens in any single draw is itself a random variable — a Binomial(L,t) draw — with standard deviation √(Lt(1−t)).

Even at the “same” t=0.5, one draw might mask 45 tokens and another 55 — a 10-token spread purely from sampling noise, layered on top of whatever difficulty the actual prediction problem has. Fixing the exact count l removes that entire extra source of variance, which is why the deterministic-count form (Eq. 14) needs roughly 8× fewer samples for equally reliable loss estimates.

Where the extra variance bites hardest: small t

The 8× figure is an average across the whole [0,1] range of t. It's worth seeing where the problem is actually worst, because it connects directly back to Chapter 5's 1/t reweighting. Take t=0.05 on the same L=100 sequence — an expected 5 masked tokens:

A standard deviation of about 2.2 tokens around an expected count of only 5 is a huge relative swing — one draw might genuinely mask 3 tokens, another 7, more than doubling or nearly halving how many terms appear in the loss sum. Now remember Chapter 5's reweighting factor at this same t: 1/0.05 = 20 — every one of those few masked tokens' prediction losses gets amplified twenty-fold. Put the two effects together: at low t, you have both the fewest masked tokens (so each one that IS included dominates the sum more) and the largest reweighting multiplier applied to each of them. A single unlucky or lucky prediction on one of only 3–7 masked positions, amplified 20×, can swing the whole loss estimate for that training example substantially. This is precisely the regime where the deterministic-count form (Eq. 14) earns its keep most — by fixing exactly how many terms enter the sum, it removes the one source of variance that would otherwise be worst exactly where the 1/t weighting is already largest.

Why the bound matters beyond bookkeeping

Because ℒ(θ) is a principled upper bound on true negative log-likelihood — not a heuristic reconstruction loss with no guaranteed relationship to log p(x₀) — LLaDA inherits properties that follow directly from being a genuine maximum-likelihood-style generative model. The paper's own central claim, stated in its introduction, is that capabilities like in-context learning are a consequence of the generative modeling principle itself — minimizing KL(p_data ‖ p_θ) — and not a special trick unique to autoregression. A model trained on a heuristic loss with no likelihood interpretation would have no theoretical claim to that property at all; a model trained to minimize a genuine (if imperfect) bound on log-likelihood does.

Fisher consistency, in plain terms

The paper's introduction leans on a specific piece of statistical theory to justify why a genuine likelihood bound matters for scalability, not just for correctness on paper: Fisher consistency is the property that, given infinite data, a sufficiently large network, and optimal training, minimizing the objective recovers the true data distribution exactly. This is a property of the generative-modeling principle from Chapter 5's opening equation — max_θ ℰ_pdata log p_θ(x) — not a property that depends on autoregression specifically. Any sufficiently expressive model family trained to minimize a genuine bound on KL(p_data ‖ p_θ) inherits it, diffusion transformers included. The paper points to diffusion transformers' success on visual data as independent evidence for the same underlying claim: the scalability large language models exhibit is plausibly a consequence of Fisher consistency plus Transformer architecture plus data and model scale, none of which is unique to next-token prediction. LLaDA's likelihood bound (this chapter) is what lets that argument apply to LLaDA specifically, rather than remaining a claim about diffusion models on images that text has no principled connection to.

Worked example: the loss, computed by hand on a 3-token toy sentence

Every quantity in this chapter has been symbolic so far. Ground it once, completely, in numbers. Take the toy sentence w = (“the”, “cat”, “sat”), L = 3, and suppose a training draw samples t = 0.6, and the actual masking draw happens to mask positions 1 and 3 (“the” and “sat”) but not position 2 (“cat” stays visible) — a perfectly plausible outcome at t=0.6, even though the Chapter 6 variance discussion above shows the exact count is itself random.

The corrupted sequence x_t the network actually sees is (M, “cat”, M). Suppose the mask predictor, conditioned on seeing “cat” in the middle, assigns probability 0.4 to the correct word “the” at position 1, and probability 0.25 to the correct word “sat” at position 3 (lower, because “the ___ ___” with a known middle word “cat” still leaves several plausible verbs). Position 2 contributes nothing to the sum — it's unmasked, so the indicator 1[x_tⁱ=M] zeroes it out. The inner sum is:

Apply the 1/t reweighting, 1/0.6 ≈ 1.667, and negate (since the loss is defined as the negative of this expectation):

That single number, 3.84, is this one training example's contribution to the Monte Carlo estimate of ℒ(θ) — one (t, x₀, x_t) draw, run through the exact formula from earlier in this chapter, entirely by hand. Training accumulates millions of draws like this one, at every value of t, across every sentence in the corpus, and gradient descent pushes the network's predicted probabilities (0.4 and 0.25 here) upward over time — which is visibly the only way to make this specific number smaller.

A concrete comparison: masked diffusion's ELBO vs. Diffusion-LM's

It's worth explicitly contrasting the derivation this chapter just walked through with Chapter 1's continuous version, now that you've seen both in full. Diffusion-LM's bound (Chapter 1) required an entire extra Markov transition — the embedding step q_φ(x₀|w) — bolted onto the front of a standard continuous-diffusion ELBO, plus a rounding term to get back to discrete words, plus careful re-parametrization (Chapter 2) just to make the bound behave well in practice. LLaDA's bound needed none of that: because the forward process corrupts discrete tokens directly into a literal mask state, there's no embedding step to define, no rounding term to add, and (Chapter 6's time-invariance result) not even a time-conditioning input for the network. The masked formulation isn't simpler because the underlying math is easier in some abstract sense — it's simpler because discreteness, treated head-on instead of routed through a continuous embedding space, removes entire categories of approximation error that Diffusion-LM had to engineer around one at a time.

Worked example: five draws, made visceral

Every argument this chapter has made about Binomial variance is easiest to actually feel by looking at what five independent draws might genuinely look like, rather than only their standard deviation formula. Take L=10, t=0.5 (expected count: 5). Five independent simulated draws of “how many of the 10 positions happen to get masked”:

Five draws at the exact same t, and the number of terms actually appearing in the loss sum ranges from 3 to 7 — more than doubling between the smallest and largest. Every one of those five loss estimates is an unbiased estimate of the same underlying quantity in expectation, but any single one of them, taken alone, could be built from anywhere between 3 and 7 masked-token prediction terms. Eq. 14's fix — always masking exactly 5 (or whatever the drawn l happens to be, chosen once and applied exactly) — removes this particular axis of draw-to-draw variation entirely, leaving only the genuine variance from the network's prediction quality itself. That's the whole mechanism behind the 1,000-vs-128 sample-count gap from earlier in this chapter, made as concrete as five numbers in a row can make it.

Case 4 deserves one more look: what “impossible” is protecting you from

It's easy to skim past Case 4 — “everything else has probability zero” — as a throwaway catch-all. It's worth being explicit about what it actually rules out, because each ruled-out event corresponds to a real bug you could otherwise accidentally implement. It rules out an unmasked token becoming masked again (a bug: overwriting already-committed context, which would make generation non-monotonic and could cause the same position to flicker between different words across steps). It rules out a masked position resolving to a word inconsistent with q_0|t's distribution (a bug: the sampled value disagreeing with what the network actually predicted, silently breaking the connection between the network's output and what gets generated). Formalizing these as probability-zero events, rather than leaving them as implicit assumptions in a sampling loop's code, is exactly the kind of rigor that makes Chapter 7's algorithm implementable without hidden edge cases — every branch the code needs to handle is named and given a probability by the theory first, rather than discovered later as a bug during implementation.

What this chapter did NOT need to derive, and why that's notable

It's worth explicitly noticing an absence. Chapter 1's derivation for Diffusion-LM needed a full extra Markov transition, a rounding loss term, a re-parametrization trick, and a custom noise schedule — four separate pieces of machinery, each solving a specific problem continuous diffusion's adaptation to text introduced. This chapter's derivation, covering the theoretically harder-sounding topic of proving a genuine likelihood bound for a discrete stochastic process, needed none of that extra scaffolding: one forward-process definition, one four-case reverse formula, one time-invariance shortcut, and one variance argument about sampling l tokens versus masking independently. The absence of extra machinery isn't an accident of this chapter being written more concisely — it's the direct payoff of Chapter 5's architectural choice to corrupt tokens directly rather than routing through a continuous embedding space at all. Every piece of scaffolding Chapter 1 needed existed specifically to bridge discrete words and continuous math; a formulation that never crosses that bridge never needs the scaffolding built to cross it.

A quick self-check before Chapter 7

Two questions worth answering out loud before continuing, since Chapter 7 builds directly on both. First: why does an unmasked token, once revealed under the true reverse process, never change again (Case 1)? Because the forward process only ever moves probability mass toward the masked state as t increases — it never defines a transition that un-reveals a token, so the reverse process, which exactly inverts the forward one, has no mechanism to do so either. Second: why does the mask predictor need no time-conditioning input at all? Because Eq. 11 showed that what a mask resolves to depends only on the currently-unmasked tokens, not on the specific value of t — a fact special to the masking formulation, absent from Diffusion-LM's continuous version, where the network explicitly takes t as an input at every step (Chapter 1). If both answers came readily, Chapters 5 and 6's machinery has landed; if not, both are worth a second pass before Chapter 7's algorithm asks you to trust them without re-deriving them each time.

Bridge to Chapter 7

One loose end remains from this chapter's four-case reverse-transition formula: Case 1 stated that an unmasked token, once revealed, never changes again under the true reverse process. Combined with the time-invariant resolution rule from Eq. 11, this has a striking consequence the next chapter derives in full: it means LLaDA's training objective is provably equivalent to training an autoregressive model over every possible generation order at once, which is the mechanistic explanation for the bidirectional reasoning and reversal-curse results the rest of this session builds toward.

Chapter 7: Any Order, One Model (showcase)

This chapter has two jobs: derive why LLaDA is, in a precise sense, equivalent to a model trained over every possible generation order simultaneously; and turn the abstract reverse-process equations from Chapter 6 into the actual step-by-step decoding algorithm, worked by hand on a small concrete example.

The any-order autoregressive equivalence

An any-order autoregressive model (AO-ARM) factorizes the same joint distribution p(x₀) autoregressively, but over an arbitrary permutation π of the L positions, not just left-to-right, trained by averaging the negative log-likelihood over all possible orderings, sampled uniformly:

The paper proves this objective is equivalent to the masked-diffusion loss from Chapters 5–6. The translation, in words: x₀^π(<i) — the tokens already “generated” before position i, under order π — corresponds exactly to the currently-unmasked tokens in some masked sequence x_t. The not-yet-generated positions {π(≥i)} correspond exactly to the currently-masked ones. A single masked pattern x_t is therefore simultaneously a valid intermediate state for every ordering π consistent with “the currently-unmasked positions came before the currently-masked ones.” Training on random masking patterns — which is all Chapters 5–6 ever did — is therefore, in expectation, training on the union of all consistent generation orders at once, via one weight-sharing network.

The sampling algorithm, from equations to procedure

Turn Chapter 6's four-case reverse-transition formula into an actual, runnable algorithm. Start from r₁, a fully-masked sequence of the target response length L (generation length is itself a hyperparameter, and the paper reports results are insensitive to it, since both pretraining and SFT trained on variable lengths). Divide [0,1] into N equal sampling steps.

The single most structurally different step from autoregressive decoding: one forward pass proposes values for every currently-masked position simultaneously, regardless of how many remain. An autoregressive model needs one forward pass per token, always; LLaDA needs one forward pass per sampling step, which can be — and typically is — far fewer than the number of tokens being generated.

Worked example: the masking schedule, by hand, L=8, N=4

Walk a concrete instance all the way through, tracking exactly how many of 8 toy positions commit at each of 4 sampling steps.

Read the shape, not just the arithmetic. Step 1 is cautious: with almost no context visible yet (t=1, everything masked), only 25% of proposed guesses get kept, and the rest are thrown back into the pool for another try. Step 4 is fully confident: by construction, s/t = 0/(1/4) = 0 at the final step, so every still-masked position commits with certainty — nothing is allowed to remain uncertain by the time the schedule ends. This front-loaded caution and back-loaded confidence isn't hand-tuned per step; it falls directly out of dividing the fixed interval [0,1] into N equal-sized steps of t.

One more precision worth having: Algorithm 4 uses greedy (argmax) prediction for r₀ at every single step, but which positions get remasked — the coin flip at probability s/t — is genuinely random. So the overall trajectory is stochastic even though each individual per-step prediction is deterministic. This baseline is called random remasking.

Two smarter remasking strategies

Random remasking (above) is the baseline the equations directly imply. Two deterministic alternatives, both inspired by prior work on masked image generation (Chang et al. 2022), do better in practice:

Low-confidence remasking. Instead of randomly choosing which s/t fraction of masked positions to remask, deterministically remask whichever proposed guesses the model is least confident about (lowest predicted probability), keeping only its most confident proposals each step.

Semi-autoregressive remasking. Divide the full generation length into fixed-size blocks (32 tokens, in the paper's ablations) and generate blocks left to right; within each block, apply the reverse process (with low-confidence remasking) as usual. This deliberately reintroduces left-to-right order at the block level, while staying fully parallel and bidirectional within each block.

Low-confidence remasking, as an actual step-by-step change to Algorithm 4

It's worth pinning down exactly which line of the earlier algorithm changes, since “remask the least confident ones instead of randomly” is easy to nod along with but less obvious to actually implement. Algorithm 4's inner loop, for each masked position, flips a coin with probability s/t to decide whether that position stays masked. Algorithm 5 (the paper's low-confidence variant) replaces that per-position coin flip with a single ranking step applied once per sampling step: after the mask predictor proposes a value for every masked position, compute each proposal's own predicted probability (how confident the model was in its own guess), sort all currently-masked positions by that confidence, and deterministically commit exactly the top ⌊(1−s/t)·(number currently masked)⌋ most-confident proposals — the same expected count Chapter 7's worked table computed, but now chosen by rank instead of by chance. Every other position, regardless of how confident its own proposal was, gets remasked and tried again next step. This is a genuinely different algorithm from Algorithm 4, not a minor variation: it makes the trajectory deterministic given the model's own confidence estimates (removing the coin-flip randomness entirely), and it concentrates commitments on exactly the positions the model itself is most sure about at each point in the schedule — which is precisely why it measurably outperforms random remasking in the ablation below.

Worked ablation, GSM8K accuracy, answer length 512, N=256 steps, block length 32:

For the Base model, confidence-based commit ordering is a clear win (52.3 → 64.7). For the Instruct model, pure low-confidence remasking collapses, catastrophically, to 12.9. The explanation is specific to how SFT trains the model: short training responses get padded with many |EOS| tokens to equalize lengths within a batch, and |EOS| is an extremely easy, extremely confident prediction almost everywhere. Pure confidence-based remasking greedily commits to a flood of premature |EOS| tokens near the start of generation, truncating the real answer before it's written — unless the semi-autoregressive block structure forces the model to fill in genuine content within each block before it can reach the ones padded with |EOS|. Combined, the two give the Instruct model's best result of the three: 73.8.

How much does the number of sampling steps actually matter?

The worked L=8, N=4 example fixed N=4 for arithmetic convenience. The paper runs a direct ablation on this exact question, fixing response length at 1,024 tokens and sweeping the number of sampling steps N, measured on GSM8K and HumanEval with the LLaDA 8B Base model. The finding is unambiguous: accuracy keeps improving as N increases, all the way up toward N = L (one sampling step per token, the point at which the schedule commits at most one new token per step on average, closest in spirit to sequential decoding). There is no sweet spot where adding more steps stops helping and plateaus early — the paper reports the trend continuing smoothly across the entire range tested. This has a direct, sobering consequence for Chapter 8's latency discussion: choosing N << L to save wall-clock time is not a free discount, it is a real point on an accuracy-vs-compute curve, and the curve keeps climbing right up to the expensive end.

Worked example: does generation length itself matter?

A related but distinct ablation: holding the ratio of steps to length fixed (one sampling step per two tokens, so N = length/2 throughout, matching Chapter 5's variable-length training setup), how sensitive is accuracy to the length hyperparameter itself — the size of the fully-masked starting sequence, chosen before generation even begins? This matters practically in a way autoregressive decoding never has to worry about: an AR model's generation length is discovered dynamically, one |EOS| decision at a time; LLaDA has to commit to a maximum length up front, before denoising even starts. The paper reports the exact GSM8K accuracy at three different length settings, for both models:

Read the spread, not just the individual numbers: across a 4× range in length (256 to 1024), Base moves by only 3.4 points (62.5 → 65.9) and Instruct moves by less than 1.5 points and isn't even monotonic (75.3 → 73.8 → 75.3). Compare that to the reversal-curse gaps from Chapter 8 — tens of points — and the length hyperparameter looks almost inert by comparison. That near-inertness is precisely the payoff of Chapter 5's 1%-random-length training trick: it's what makes an otherwise-awkward hyperparameter safe to set generously (as the paper does — 1,024 for most conditional-generation benchmarks, per its own evaluation appendix) and mostly forget about.

A second scale check: the iGSM synthetic benchmark

Beyond the standard GSM8K numbers Chapter 8 compares in full, the paper runs a second, harder mathematical check specifically designed to rule out data contamination: iGSM, an infinite, procedurally-generated GSM8K-style dataset where every problem is synthesized fresh, so no version of it could possibly appear in either model's training data. Problems are parametrized by how many solution steps they require:

LLaDA leads at every difficulty level, and by the widest margin at the shortest problems (26 points at 4 steps). This particular result can't be explained away by GSM8K-specific data leakage, since iGSM problems are synthesized on demand — it's independent evidence that LLaDA's mathematical reasoning strength (also visible in Table 1's ordinary GSM8K and Math benchmarks, Chapter 8) generalizes past the exact benchmark it was measured on.

Where Chapter 6's Monte Carlo discussion shows up in actual benchmark evaluation

Chapter 6 derived, abstractly, that the deterministic-count loss form (Eq. 14) reaches stability with 128 Monte Carlo samples. It's worth seeing where that exact number gets used in practice, not just in a training loss. Several of LLaDA's standard benchmarks (Chapter 8's Table 1) are evaluated via conditional likelihood estimation: given a prompt and several candidate answers, compute each candidate's likelihood under the model and pick the highest-scoring one. For benchmarks like MMLU, CMMLU, and C-Eval, where each candidate answer is a single token (the letter A, B, C, or D), a single Monte Carlo estimate is enough — there's no masking-count variance to average over when there's only one token to ever mask in the first place. For every other likelihood-based benchmark, where candidate answers span multiple tokens, the paper reports needing the same 128 samples Chapter 6 derived the reasoning for. This is a genuinely satisfying full circle: an abstract variance argument from Chapter 6, and a concrete number quoted without derivation there, turns out to be the exact number the paper actually used to evaluate several of Chapter 8's headline benchmark comparisons.

Greedy vs. sampled: one more decoding choice worth naming

Algorithm 4's line 4 specifies argmax — greedy — prediction for r₀ at every step, and this session has treated that as the default throughout. It's worth being explicit about what this choice costs and buys, since it's a genuine design decision, not a forced one. Greedy decoding is deterministic given the remasking randomness: run the same schedule with the same random remasking draws twice, and you'd get the same output both times (modulo Chapter 7's low-confidence variant, which removes even that source of randomness). This mirrors an analogous choice in autoregressive decoding — greedy next-token selection versus temperature-scaled sampling — and the tradeoffs run in the same direction: greedy tends toward safer, more predictable, occasionally repetitive output; sampling with some temperature trades a bit of that reliability for more diverse generations. Nothing in Algorithm 4's structure prevents replacing the argmax at line 4 with a temperature-scaled sample from p_θ(r₀|p₀,r_t) instead — the remasking machinery (deciding which positions commit) is entirely separate from the question of how a committing position picks its value, which is exactly why the paper can treat remasking strategy (random vs. low-confidence vs. semi-autoregressive) and value-selection strategy (greedy vs. sampled) as two independent knobs rather than one entangled choice.

Recap: everything this chapter's showcase actually demonstrated

This chapter carried the most machinery of any in the session, so it's worth compressing it once before moving on. The AO-ARM equivalence proved LLaDA's training is, in expectation, equivalent to training over every generation order simultaneously — a theoretical result. Algorithm 4 turned Chapter 6's abstract reverse-transition formula into a concrete, runnable procedure — an implementation result. The worked L=8, N=4 table showed that procedure's cautious-then-confident shape falls directly out of dividing [0,1] into N equal steps, not from any hand-tuning — an arithmetic result. And the remasking ablation showed that the theoretically simplest strategy (random remasking) is measurably beaten by two smarter, still-principled alternatives, with the right choice depending on which model (Base or Instruct) and why — an empirical result. Four different kinds of claim, each verified a different way, is exactly the standard the “never trust a headline number you haven't traced back to its arithmetic” discipline from this session's closing chapter asks for.

Chapter 8: AR vs Diffusion, Honestly

Time to put everything from this session side by side and be precise about what actually won, what didn't, and where the paper's own headline claims need a careful second read.

The tradeoff table

Why LLaDA gave up KV caching, precisely

An autoregressive model's KV cache works because a token's key and value, once computed, never change — future tokens only ever attend to them, never rewrite them. LLaDA's bidirectional attention breaks that invariant completely: at every sampling step, any position's representation can change, because the mask predictor recomputes a full forward pass over the entire current sequence, masked and unmasked positions alike, every single time. There is nothing to safely cache and reuse, because there's no guarantee a previously-computed key or value is still valid one step later. That's the real reason, from Chapter 5, LLaDA uses 32 KV heads instead of LLaMA3's 8 — grouped-query attention exists specifically to shrink a KV cache LLaDA never gets to build in the first place.

Worked example: sizing the cache LLaDA would have needed, if it could cache at all

Put an actual number on what Chapter 5's KV-head difference would have cost, using the same per-token cache formula this course derived in Session 08 (Attention Sinks & Streaming): 2 × n_{kv heads} × head_dim × dtype_bytes × n_layers. LLaMA3 8B, with its GQA-shrunk 8 KV heads and head dimension 4,096/32 = 128, in fp16:

LLaDA, with 32 KV heads (same head dimension, same layer count, same precision) — the configuration Chapter 5's architecture table showed — would cost, hypothetically, if it could safely cache at all:

Four times LLaMA3's per-token cache cost, purely from the KV-head count — a direct, quantified illustration of why grouped-query attention exists at all, and exactly the number Chapter 5 gestured at without computing. But this comparison is entirely hypothetical, and that's the whole point of this chapter's argument: LLaDA never actually builds this 512-KiB-per-token structure, because there's nothing safe to cache in the first place under bidirectional attention. The 4× multiplier above isn't a cost LLaDA pays; it's a number that would only matter if you tried to bolt caching onto an architecture that structurally can't support it — which is exactly why the paper didn't bother with GQA's savings at all.

Worked example: sampling steps vs. sequence length, the actual latency lever

Chapter 7's remasking ablation ran with response length fixed at 512 and N = 256 sampling steps — on average, 512/256 = 2 tokens newly committed per sampling step (unevenly, front-loaded cautious and back-loaded confident, as the L=8 worked table showed). Compare directly to autoregressive decoding of the same 512 tokens: exactly 512 sequential forward passes, one per token, no way around it — each individual pass is cheap thanks to KV caching, but the 512 passes cannot be parallelized across positions within one sample, because token 2 structurally depends on token 1 already being finalized.

LLaDA's 256 forward passes are fewer in count — a 2× reduction in pass count — but each individual pass is asymptotically more expensive than an AR model's cached step, since it processes the whole sequence with full bidirectional attention every single time, with nothing cached from the previous pass. Whether LLaDA ends up faster wall-clock depends entirely on hardware batching and parallelism assumptions, and the paper is candid that its inference behavior remains sensitive to these hyperparameters. Section B.6's finding sharpens the tradeoff further: accuracy on GSM8K and HumanEval keeps improving as the number of sampling steps increases, all the way up toward matching the response length itself — meaning the “cheap” N = L/2 regime used in Chapter 7's ablation is already trading away some accuracy for speed, not a free lunch.

The reversal curse, read honestly

The reversal curse (Berglund et al. 2023): a model trained on “A is B” often fails to answer questions requiring “B is A”—evidence that autoregressive training bakes in a direction-dependent asymmetry that shouldn't, in principle, exist in the underlying facts. LLaDA tests this with a poem-completion task: given one line of a famous poem, generate the next line (forward) or the previous line (reversal), with no additional fine-tuning for either direction.

Read this table slowly, because the honest version of the story is more interesting than the headline. LLaDA does not win on either individual number — its forward score (51.8) is dramatically lower than GPT-4o's (82.7) and Qwen2.5's (75.9), by 30.9 and 24.1 points respectively. Some of that gap is simply resourcing (2.3T training tokens vs. presumably much larger, proprietary training runs for the other two); some of it is structural — forward generation is exactly the task an AR model's entire training procedure optimizes for.

What LLaDA demonstrates is balance, not superiority. Its reversal score (45.6) is actually the highest of the three models compared — and the gap between its own forward and reversal performance (6.2 points) is roughly 6–8× smaller than GPT-4o's or Qwen2.5's internal gap. This is exactly what Chapter 7's AO-ARM equivalence predicts structurally: LLaDA was never trained to prefer one direction over the other, marginalizing over all orderings during training, so it has no structural reason to be lopsided the way a strictly left-to-right model is. But “less lopsided” is a claim about internal consistency, not a claim of beating GPT-4o outright — a reader who only remembers the headline “LLaDA breaks the reversal curse” without this breakdown would badly overestimate what's actually been shown. The paper's own stated claim is careful on exactly this point: “consistent zero-shot performance across both forward and reversal tasks,” which is a claim about the model relative to itself.

Zooming out: how competitive is LLaDA overall?

The reversal-curse result is one narrow, carefully-designed test. It's worth checking it against the broader picture from Table 1 of the paper — LLaDA 8B Base against LLaMA2 7B Base and LLaMA3 8B Base, all evaluated under the same protocol, across the general, mathematical, code, and Chinese-language categories:

LLaDA beats LLaMA2 7B decisively across every category, and leads the larger, much-more-heavily-trained LLaMA3 8B (15T tokens, versus LLaDA's 2.3T) on every row shown here — general knowledge included, not just math and Chinese. That includes code: LLaDA edges LLaMA3 on standard HumanEval too (35.4 vs. 34.8), not just the fill-in-the-middle variant below.

Read that the way this chapter has insisted on reading every other table, though: these seven rows are a curated subset of the paper's full fifteen-row Table 1, and this lesson — not the paper — chose which seven to show. On the six general-purpose reasoning rows this table leaves out, LLaMA3 8B leads every single time: BBH (62.1 vs. 49.7, a 12.4-point LLaDA deficit), Hellaswag (79.1 vs. 70.5, 8.6 points), PIQA (80.6 vs. 73.6, 7.0 points), ARC-C (53.1 vs. 45.9, 7.2 points), WinoGrande (77.3 vs. 74.8, 2.5 points), and GPQA (25.9 vs. 25.2, an essential tie). MBPP, the code benchmark this table also omits, goes to LLaMA3 as well (48.8 vs. 40.0, 8.8 points). The paper's own summary is more measured than either headline version of this table would suggest: LLaDA 8B Base is described as “surpassing LLaMA2 7B Base on nearly all tasks, and overall competitive with LLaMA3 8B Base,” with a specific, stated edge in math and Chinese — not a clean sweep in either direction, on either side of the ledger. For the strengths it does show, the paper's own conjecture is that they “stem from the same factors as its relatively weaker performance in some tasks — differences in data quality and distribution,” largely unavoidable given that the exact training corpora behind closed competitor models aren't public, so a perfectly matched comparison isn't fully possible for anyone in this space.

The HumanEval-FIM row still deserves its own callout, since it's the one place this session's infilling story (Chapters 0 and 4, argued through Diffusion-LM's small-scale aNLG task) shows up as a real, standard, widely-used benchmark number at full 8B scale, purpose-built for exactly this comparison: LLaDA 8B (73.8) edges out LLaMA3 8B (73.3) on fill-in-the-middle code completion — a genuine, if narrow, win on exactly the task type Chapter 0 argued autoregressive models are structurally worst-suited for. Standard HumanEval, above, tells the same story in miniature once read correctly — the fill-in-the-middle variant is just the cleaner, purpose-built test of the same underlying claim.

A serving-side footnote: batching and the KV-cache gap

One more practical wrinkle worth naming, since it compounds with Chapter 8's core latency point rather than replacing it. An autoregressive server's KV cache is what makes serving many concurrent users efficient: each user's cache grows incrementally and independently, and the expensive part (attention over history) gets reused, not recomputed, at every step. LLaDA's inability to cache means every sampling step, for every active generation, re-processes the full current sequence from scratch with full bidirectional attention. At the single-request level this is the tradeoff Chapter 8 already worked through in detail (fewer, more expensive passes vs. more, cheaper ones). At the fleet level, it also removes one of the standard levers production LLM serving systems use to pack many concurrent conversations onto limited GPU memory — a genuinely different memory-and-batching profile that a team deploying a diffusion language model has to design around, not just a slower version of the same autoregressive serving stack.

What “sensitive to inference hyperparameters” looks like in practice

Chapter 9 previews the paper's own admission that LLaDA is sensitive to inference hyperparameters; here is what that sensitivity actually looks like as a table, not just a sentence. Evaluating LLaDA 8B Instruct across its benchmark suite required a per-benchmark choice of both answer length and block length (the semi-autoregressive chunk size from Chapter 7):

Notice the block length varies wildly relative to answer length — GSM8K uses tiny 8-token blocks inside a 256-token answer (many small semi-autoregressive chunks), while Math uses a block length equal to the full answer length (equivalent to no block structure at all, pure parallel decoding within one block). This isn't a single robust default applied everywhere; it's a per-task configuration choice the paper made deliberately, benchmark by benchmark. An autoregressive model's inference-time knobs (temperature, top-p, max tokens) are comparatively benchmark-agnostic by contrast — you don't typically retune the sampling strategy itself per task, only the stopping conditions. This table is the concrete evidence behind Chapter 9's “more tuning than an AR model's forgiving knobs” claim, not just an assertion to take on faith.

Guidance was only affordable for some benchmarks, not others

One more honest wrinkle worth surfacing, since it qualifies several of this chapter's headline comparisons. The classifier-free guidance formula Chapter 9 introduces (a hyperparameter w controlling how strongly a response follows the prompt versus a masked-prompt baseline) was only searched over for the likelihood-estimation benchmarks — ARC-C, HellaSwag, TruthfulQA, WinoGrande, PIQA, and GPQA — using a grid w ∈ {0, 0.5, 1, 1.5, 2}, with the paper reporting whichever value in that grid scored best. For every benchmark evaluated by conditional generation instead (GSM8K, HumanEval, MBPP, and the rest of Table 1's generation-based rows), the paper explicitly states guidance was not applied, due to computational resource constraints. This means the Table 1 and Table 2 numbers this chapter has leaned on throughout are not uniformly “LLaDA's best possible result” — some rows benefited from a guidance-scale search, others didn't, purely because of compute budget rather than because guidance wouldn't have helped the second group. It's a small but real asterisk on any benchmark-by-benchmark comparison drawn from this session: the playing field, even within one paper's own reported numbers, wasn't perfectly level row to row.

Putting a number on “structurally hard” from Chapter 0

Chapter 0 opened with a claim stated in purely structural terms: a left-to-right factorization has no conditioning term linking an early token to a later one. This chapter's evidence lets you attach real numbers to that claim, gathered in one place: the syntax-tree control gap (86.0% vs. 17.9%, Chapter 3), the composition-of-controls gap (74.8% vs. 15.4% syntax success under composition, Chapter 4), and the near-zero forward/reversal gap (6.2 points vs. 48.4, this chapter). Three completely different experimental setups — one about steering with an external classifier, one about combining two constraints at once, one about which direction a poem gets read — and all three trace back to the exact same structural fact Chapter 0 stated in one sentence: nothing in a left-to-right factorization lets an early decision see a later one. That's the kind of convergent evidence worth trusting more than any single number in isolation — three independently designed experiments, run by two different research groups three years apart, all landing on the same underlying mechanism.

What this chapter deliberately did not resolve

In the spirit of the honesty this chapter has insisted on throughout, it's worth naming what remains genuinely open rather than implying this session settled it. Whether the wall-clock latency tradeoff (fewer, more expensive forward passes) favors diffusion or autoregressive decoding in a real production serving system depends on hardware, batch size, and workload in ways neither paper this session studied was built to answer — Chapter 8's arithmetic shows the shape of the tradeoff, not its resolution on any particular GPU fleet. Whether LLaDA's reversal-curse balance would survive at a scale genuinely matched to GPT-4o or Qwen2.5, rather than at LLaDA's comparatively modest 2.3T training tokens, is untested by this session's sources. And whether the specific numbers in this chapter's tables would hold up under a compute-matched, RL-aligned rerun is precisely the open question Chapter 9 spends its own opening section naming directly. Read this chapter's verdicts as the honest state of the evidence as of two specific published papers, not as a final ruling on autoregression versus diffusion as approaches.

The tradeoff table, revisited with everything now grounded

Look back at this chapter's opening table with all eight rows now individually derived rather than merely asserted: KV-caching incompatibility traced to the 512-KiB-vs-128-KiB-per-token comparison above, the latency knob traced to the 256-vs-512 pass-count arithmetic, controllability traced to Chapter 0's chain-rule enumeration and Chapter 3's syntax-tree numbers, infilling traced to both Chapter 4's MBR result and this chapter's HumanEval-FIM comparison, likelihood traced to Chapter 4's nats-to-perplexity conversion, and the reversal-reasoning row traced to the honest read this section just walked through in full. Every cell in that table now has a chapter and a derivation behind it — which was always the point of building this session the long way rather than starting with the summary table and working backward.

Chapter 9: Where This Goes

Nine chapters, two papers, one underlying idea: replace the fixed left-to-right factorization with a process that starts vague everywhere and refines everywhere together. Close by reading the limitations honestly, and by pointing at where the idea is heading next.

LLaDA's own stated limitations

The paper's own conclusion is unusually direct about what it hasn't shown yet. Four items, worth taking as seriously as the headline results:

Compute-matched comparisons are incomplete. Due to resource constraints, LLaDA's direct comparisons against a matched autoregressive baseline were capped under 10²³ FLOPs; the 8B-scale headline comparisons to LLaMA3 8B reuse an existing 7B AR baseline trained separately, not a perfectly matched control.

No diffusion-native architecture. LLaDA borrows LLaMA3's recipe wholesale — RoPE, RMSNorm, SwiGLU — rather than anything designed specifically for bidirectional masked prediction. The one architectural change (Chapter 5's KV-head count) was a necessity, not an optimization.

Inference guidance remains preliminary. The classifier-free guidance formula LLaDA borrows for inference,

balances how strongly the response follows the prompt (p₀) against a masked-prompt baseline (m), controlled by hyperparameter w. The paper explicitly calls this exploration preliminary and describes LLaDA as sensitive to inference hyperparameters more broadly — getting good samples out currently takes more tuning than an AR model's comparatively forgiving temperature/top-p knobs.

No RL-based alignment. LLaDA Instruct went through SFT only — the paper explicitly defers RLHF/DPO-style alignment to future work. Benchmark gaps against RL-aligned competitors (Table 2's HumanEval: 49.4 for LLaDA vs. 59.8 for LLaMA3 8B Instruct) conflate what SFT-only training achieves with what masked diffusion's architecture is capable of — a real confound this session shouldn't let you forget when reading those gaps.

Worked example: where the SFT-only gap shows up, and where it doesn't

The fourth limitation is worth quantifying rather than just naming, because the gap isn't uniform across task types — a detail that supports reading it as an alignment-stage confound rather than an architectural ceiling. Table 2 compares LLaDA 8B Instruct (SFT only) against LLaMA3 8B Instruct (SFT plus RL alignment) directly:

The pattern is not quite as clean as “code only, everything else fine.” The two largest gaps are still the code benchmarks (HumanEval at 10.4 points, MBPP at 16.6 points), and ARC-C is the one place LLaDA is unambiguously ahead (6.1 points), with MMLU the closest call (2.9 points). But GSM8K breaks the tidy story: it shows a real, non-trivial 8.9-point deficit — not the near-parity a quick skim of “general knowledge, reasoning, and math hold up fine” might lead you to expect. Grouping HumanEval and MBPP against “everything else” overstates how contained the gap actually is.

GSM8K is worth a second, more careful look, because it's the row where the SFT-only-confound explanation is on shakiest ground — and where the base-model comparison from earlier in this chapter makes the shakiness concrete. Chapter 8's Table 1 showed LLaDA 8B Base beating LLaMA3 8B Base on GSM8K by a wide margin: 70.3 vs. 48.7, a 21.6-point LLaDA lead, before either model saw a single alignment-stage gradient update. After LLaMA3's SFT+RL pipeline, its GSM8K score jumps to 78.3 — a +29.6-point gain from alignment alone. After LLaDA's SFT-only pipeline, its GSM8K score moves to 69.4 — a −0.9-point net change, essentially flat, arguably a slight regression. The ordering flips entirely: the model that led by 21.6 points pre-alignment trails by 8.9 points post-alignment, purely because one competitor's alignment stage added nearly 30 points on this specific benchmark and the other's added none. That's consistent with reading the Instruct-level gap as an alignment-stage confound rather than an architectural ceiling — the underlying masked-diffusion base model clearly can do GSM8K-style math — but it's also evidence that whatever RL contributes on GSM8K specifically is doing real, substantial work, not just output-formatting cleanup, a claim this session's sources don't test directly for LLaDA and so can't fully settle either way.

Where the idea is heading

Shortly after LLaDA, commercial-scale diffusion language models began shipping: Inception Labs' Mercury and Google DeepMind's Gemini Diffusion both apply the same masked, parallel-decoding idea LLaDA demonstrates at 8B, at production scale, explicitly marketing the same latency lever Chapter 8 worked through by hand — fewer sequential forward passes, more tokens committed per step — as a genuine serving-speed advantage over autoregressive decoding. The bet these systems are making is exactly the one Chapters 7 and 8 walked through with real numbers: if you can push sampling steps N well below response length L without losing too much of the accuracy Section B.6's steps-vs-length curve showed you'd trade away, the parallel-commit structure becomes a genuine wall-clock win, not just a parameter-count curiosity. Whether that holds up broadly, beyond the specific tasks this session's two papers evaluated, is still an open, actively-worked question, and it is worth being explicit that neither Mercury nor Gemini Diffusion is a paper this session grounded the way it grounded Diffusion-LM and LLaDA — treat this paragraph as context for where the field is headed, not as a third set of verified numbers to memorize alongside Chapters 1–8.

What would actually settle the open question

Chapter 0 opened with LLaDA's own framing: is autoregression the only viable path to the capabilities large language models exhibit, or one implementation of something more general? This session's evidence is real but partial, and it's worth being precise about what's still missing before treating the question as settled in either direction. A genuinely compute-matched comparison — same data, same FLOPs budget, same amount of hyperparameter search on both sides — between a diffusion language model and an autoregressive one at a scale beyond 10²³ FLOPs would say more than any of this chapter's benchmark tables, precisely because LLaDA's own paper flags that comparison as something resource constraints prevented it from running cleanly. An RL-aligned LLaDA Instruct, closing the SFT-only gap this chapter already named, would separate “masked diffusion has a lower ceiling” from “masked diffusion hasn't finished the same training pipeline yet.” And a diffusion-native architecture — attention and positional encoding designed for bidirectional masked prediction from scratch, rather than borrowed wholesale from an autoregressive recipe that had years of head start to mature — would test whether LLaDA's current results are close to what the idea can do, or a first working draft of it.

A practical checklist: when would you actually reach for one of these

Put the whole session to work as a decision, not just a history lesson. Reach for something in the diffusion family, based on what's actually been shown here, when: the task has a genuine global constraint that spans the whole output — a target parse structure, an exact length, a fixed left-and-right context to infill between — where Chapter 3's mechanism (seeing and revising every position at every step) has a real structural advantage over left-to-right steering, not just a marginal one. Reach for an ordinary autoregressive model, without a second thought, when the task is open-ended left-to-right generation with no unusual bidirectional constraint, decoding latency at long sequence lengths genuinely matters and your serving stack already leans hard on KV caching, or you need the alignment maturity (RL-tuned instruction-following, extensively red-teamed safety behavior) that Chapter 9's limitations section shows autoregressive models currently have and LLaDA does not yet have. Between those two clear cases sits a genuine judgment call — infilling-adjacent tasks like code completion, where Chapter 8's HumanEval-FIM result (73.8 vs. 73.3) suggests the gap has narrowed to roughly a coin flip rather than a clear win either way.

One last piece of intellectual honesty, echoing both papers' own closing statements: like any language model, diffusion language models inherit the same broader concerns as their autoregressive counterparts — compute and energy cost at training scale, potential misuse for generating misleading or harmful content at volume, and the risk of amplifying biases already present in training data. Nothing about generating text via denoising instead of next-token prediction makes those concerns disappear; the generation mechanism changed, the surrounding responsibilities didn't.

The whole session, in the numbers that matter

Nine chapters produce a lot of tables. If only a handful of numbers survive in memory, these are the ones worth keeping, each tied to the single idea it demonstrates:

Every one of these numbers has a derivation behind it somewhere in the nine chapters above — that's the whole discipline this session tried to model: never trust a headline claim you haven't personally traced back to the arithmetic underneath it.

Bridges to the rest of the course

If Chapters 1 and 2's continuous-diffusion machinery — the forward noise process, the reverse denoising network, the noise schedule — felt like it assumed intuition you hadn't built yet, start with Diffusion Models, which derives image diffusion from zero before this session ever borrowed it for text.

Flow Matching reformulates the same denoising idea around learning a velocity field rather than predicting noise or x₀ directly — worth comparing side by side with Chapter 2's x₀-parametrization fix, since both are really answering the same question: what should the network's output actually represent at each step of the process?

And if Chapter 8's KV-cache-incompatibility argument raised more questions than it answered, Session 08 of this course, Attention Sinks & Streaming, derives exactly what a KV cache stores and why it exists in the first place, from zero — useful background for feeling the full weight of what LLaDA gave up by choosing bidirectional attention over a causal mask.

Comparison, one more time

A reading checklist for the next diffusion-LM paper you meet

This field moves fast, and by the time you read this, new discrete- and continuous-diffusion language models will exist that this session never saw. Rather than re-deriving everything from scratch each time, carry forward the specific questions this session's nine chapters taught you to ask of any new result:

• Bound or exact? Is the reported likelihood/NLL an exact quantity, or a variational upper bound (Chapters 1, 4, and 6 all leaned on this distinction)? A bound being worse than an autoregressive baseline is a much weaker claim than an exact likelihood being worse.
• Compute-matched? Are the headline comparisons run at the same FLOPs budget, same data, same amount of hyperparameter search on both sides (Chapter 9's own admission that LLaDA's weren't, above roughly 9.6×10²² FLOPs)?
• Which training stage? SFT only, or SFT plus RL alignment (Chapter 9's HumanEval/MBPP gap, concentrated almost entirely in code tasks — exactly where RL alignment tends to matter most)?
• What's the control task, really? A control success rate on a narrow, easy attribute (sentiment, topic) is a different, much weaker claim than success on a genuinely global structural constraint (Chapter 3's Syntax Tree and Syntax Spans results, which is where Diffusion-LM's advantage was largest and most convincing).
• What's the actual latency lever? Fewer sampling steps than response length is a real speed advantage only if the accuracy-vs-steps curve (Chapter 7's Section B.6 finding) hasn't already been quietly traded away to get there.

None of these questions requires re-deriving this session's math from scratch every time — they're the compressed, reusable version of everything Chapters 0 through 9 built up the long way.

One closing derivation, stated as a challenge

Before moving on, it's worth testing whether Chapter 7's central result actually stuck, using a version you haven't seen worked out yet. Chapter 7 proved the AO-ARM equivalence using L=8, N=4 as its running example throughout Chapter 7's worked table. Try the same masking-schedule arithmetic for L=6, N=3 by hand, before reading on: three steps, t starting at 1, step size 1/3. Step 1: t=1, s=2/3, s/t=0.667, commit probability 0.333, expected commits from 6 masked ≈ 2. Step 2: t=2/3, s=1/3, s/t=0.5, commit probability 0.5, expected commits from the remaining 4 ≈ 2. Step 3: t=1/3, s=0, s/t=0, commit probability 1.0, the remaining 2 all commit. Running totals: 2, 4, 6 — the same qualitative shape as the L=8, N=4 table (cautious first, confident last, everything resolved by the final step), just compressed into three steps instead of four. If your own hand computation landed on those same three numbers, Chapter 7's mechanism has genuinely transferred, not just been read.

The single sentence each paper would want you to keep

If this entire session had to compress to two sentences, one per paper, they would be these. Diffusion-LM: a non-autoregressive language model, built from continuous Gaussian diffusion over learned word embeddings, can outperform autoregressive plug-and-play methods — and even a fine-tuned oracle — on control tasks that require reasoning about a sequence's global structure, precisely because it never commits to any token before the whole sequence has had a chance to inform every other token. LLaDA: the same core idea, radically simplified into discrete token masking and scaled to 8 billion parameters with a genuine likelihood bound behind it, competes directly with a leading autoregressive model of the same size while inheriting a structurally different, more balanced relationship to generation order — evidence, not proof, that autoregression is a choice about how to factorize a distribution, not the only way to build a language model capable of the things large language models are known for.

“What is now proved was once only imagined.” — William Blake, quoted at the opening of LLaDA's introduction