{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Module 26: Diffusion Language Models" ], "id": "9fea4d47-1174-4f30-98bc-ece9d1fa9209" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [], "id": "86ba5530-fe18-4244-b7b3-b499135df74f" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "f4cc4326-19eb-4b11-a382-144bedfab74c" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction\n", "\n", "Every model in this book so far writes text the same way: **one token at\n", "a time, left to right**. That is the *autoregressive* factorization\n", "$p(x) = \\prod_i p(x_i \\mid x_{ **Key Insight**\n", ">\n", "> The forward process has **no parameters** and needs no training — it\n", "> is just “mask each token with probability $t$.” All the learning is in\n", "> the *reverse* network that has to undo it. This is the same split as\n", "> image diffusion (add Gaussian noise forward, learn to denoise\n", "> backward); only the corruption is different — discrete masking instead\n", "> of continuous noise.\n", "\n", "## The Math: The Absorbing-State Process\n", "\n", "### Forward: masking is a coin flip per token\n", "\n", "Fix a noise schedule $\\alpha_t \\in [0, 1]$, the probability a token is\n", "**kept**. We use the simplest one, the **linear schedule**\n", "$\\alpha_t = 1 - t$. The forward process corrupts each token\n", "independently:\n", "\n", "$$q(x_t \\mid x_0) = \\alpha_t\\,\\big[x_t = x_0\\big] \\;+\\; (1 - \\alpha_t)\\,\\big[x_t = \\texttt{[MASK]}\\big].$$\n", "\n", "So a token survives with probability $\\alpha_t$ and becomes `[MASK]`\n", "with probability $1 - \\alpha_t = t$. **The masking probability *is* the\n", "diffusion time.** At $t = 0$ nothing is masked; at $t = 1$ everything\n", "is.\n", "\n", "### Reverse: a mask predictor\n", "\n", "The reverse network $x_\\theta$ takes the corrupted sequence $x_t$ and\n", "outputs, for every position, a distribution over the *real* vocabulary —\n", "its guess at the original token. Crucially it must see the **whole**\n", "sequence, past and future, masked and unmasked, to fill a hole from\n", "context. So it is a **bidirectional** transformer: the same block as\n", "[m06](../m06_transformer/lesson.qmd) with the causal mask **removed**.\n", "\n", "One surprise: the mask predictor takes **no time input**. It never needs\n", "to be told $t$, because the *fraction of `[MASK]` tokens it can see*\n", "already reveals how noisy the sequence is. Following MDLM and LLaDA, our\n", "model is time-free.\n", "\n", "### The loss: reweighted masked cross-entropy\n", "\n", "Training minimizes the negative evidence lower bound (NELBO). For\n", "absorbing-state diffusion this integral collapses into something you\n", "already know — a **cross-entropy over the masked positions**, weighted\n", "by the noise level:\n", "\n", "$$\\mathcal{L} \\;=\\; \\mathbb{E}_{t}\\,\\mathbb{E}_{q}\\!\\left[\\;\\frac{\\alpha_t'}{1 - \\alpha_t}\\sum_{i \\in \\text{masked}} \\log x_\\theta(x_t)_{i, x_0^{(i)}}\\;\\right].$$\n", "\n", "Under the linear schedule ($\\alpha_t = 1 - t$, so $\\alpha_t' = -1$), the\n", "weight is $\\alpha_t' / (1 - \\alpha_t) = -1/t$: each masked position\n", "contributes its cross-entropy scaled by $1/t$. That is the punchline of\n", "**MDLM** (Sahoo et al., 2024): *a masked-diffusion language model is a\n", "masked language model, trained at every noise level and reweighted by\n", "$1/t$.* The $1/t$ weight down-weights very noisy draws (large $t$),\n", "where the model has almost no context to work with.\n", "\n", "### Step by step\n", "\n", "Walk one forward-then-reverse cycle:" ], "id": "78ff600a-9c50-4d8d-8019-c57c6f440ec5" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "60adc30d-0eae-4d8e-90e8-805aeba5bdbf" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "c89c1da6-8124-4ffd-9dc9-52d668fd93a7" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "7583ae76-aeff-4039-b6e2-0b88e716672c" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Code: A Mask-Predictor from Scratch\n", "\n", "The full implementation lives in `diffusion.py`. Start with the forward\n", "process — it is just a per-token coin flip:" ], "id": "1a7ec2b9-9d82-411f-ab81-b5c6beb8b082" }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "clean x0: [[5, 2, 8, 1, 9, 4]]\n", "noisy xt: [[99, 2, 99, 99, 99, 4]] (99 = [MASK])\n", "masked?: [[True, False, True, True, True, False]]\n", "alpha_t (keep prob) at t=0.5: 0.5" ] } ], "source": [ "import torch\n", "from diffusion import forward_mask, linear_alpha\n", "\n", "torch.manual_seed(0)\n", "g = torch.Generator().manual_seed(0)\n", "\n", "x0 = torch.tensor([[5, 2, 8, 1, 9, 4]]) # a clean sequence\n", "xt, mask = forward_mask(x0, t=0.5, mask_id=99, generator=g)\n", "\n", "print(\"clean x0: \", x0.tolist())\n", "print(\"noisy xt: \", xt.tolist(), \" (99 = [MASK])\")\n", "print(\"masked?: \", mask.tolist())\n", "print(\"alpha_t (keep prob) at t=0.5:\", linear_alpha(0.5))" ], "id": "cc4eabb5" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The marginal masking probability really is $t$ — average over a big\n", "batch:" ], "id": "9b012879-8c43-48c0-9fae-ecc83895492f" }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "t = 0.2: fraction masked = 0.197\n", "t = 0.5: fraction masked = 0.499\n", "t = 0.8: fraction masked = 0.804" ] } ], "source": [ "big = torch.zeros(500, 40, dtype=torch.long)\n", "for t in (0.2, 0.5, 0.8):\n", " _, m = forward_mask(big, t=t, mask_id=99, generator=g)\n", " print(f\"t = {t}: fraction masked = {m.float().mean():.3f}\")" ], "id": "bb32d3cd" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now the mask predictor. It is an ordinary pre-norm transformer with\n", "**one change** — the attention is bidirectional (no causal mask),\n", "because filling a hole needs the whole sentence. The output head has\n", "`vocab_size` logits, not `vocab_size + 1`: the model *removes* `[MASK]`,\n", "it never predicts it." ], "id": "9e0df950-34a6-43cc-a204-32b78dfc3eea" }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "input (all masked): [[16, 16, 16, 16, 16, 16]]\n", "logits shape: (1, 6, 16) → (batch, seq, real vocab)\n", "[MASK] id: 16 (the extra embedding row, never an output)" ] } ], "source": [ "from diffusion import MaskedDiffusionLM\n", "\n", "model = MaskedDiffusionLM(\n", " vocab_size=16, embed_dim=64, num_heads=4, num_layers=2, max_seq_len=8\n", ")\n", "\n", "x = torch.full((1, 6), model.mask_id) # start from all [MASK]\n", "logits = model(x)\n", "print(\"input (all masked):\", x.tolist())\n", "print(\"logits shape:\", tuple(logits.shape), \"→ (batch, seq, real vocab)\")\n", "print(\"[MASK] id:\", model.mask_id, \"(the extra embedding row, never an output)\")" ], "id": "17cf2950" }, { "cell_type": "markdown", "metadata": {}, "source": [ "And the loss — cross-entropy on the masked positions, weighted by $1/t$:" ], "id": "e5116243-321c-4d4b-a90f-611d4c20afa3" }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "noisy input: [[16, 2, 16, 1, 16, 4]] (16 = [MASK])\n", "diffusion loss: 5.4764\n", "only the 3 masked positions contribute" ] } ], "source": [ "from diffusion import diffusion_loss\n", "\n", "# Corrupt with the model's own [MASK] id, predict, and score the holes.\n", "xt_m, mask_m = forward_mask(x0, t=0.5, mask_id=model.mask_id, generator=g)\n", "logits = model(xt_m)\n", "loss = diffusion_loss(logits, x0, mask_m, t=0.5)\n", "print(\"noisy input:\", xt_m.tolist(), f\" ({model.mask_id} = [MASK])\")\n", "print(\"diffusion loss:\", round(loss.item(), 4))\n", "print(\"only the\", int(mask_m.sum()), \"masked positions contribute\")" ], "id": "e044c9ea" }, { "cell_type": "markdown", "metadata": {}, "source": [ "This follows the algorithm in `diffusion.py`: `forward_mask` for\n", "corruption, `MaskedDiffusionLM` for the reverse network,\n", "`diffusion_loss` for the objective.\n", "\n", "## Bidirectional, not Causal\n", "\n", "Modules 19 and 22 proved their mixers are **causal**: editing a *future*\n", "token never changes an *earlier* output — the property that makes\n", "autoregressive decoding valid. The diffusion mask predictor proves the\n", "**mirror image**. It is bidirectional *on purpose*: a future token must\n", "be able to reach backward, or it could not help fill an earlier hole." ], "id": "ea5e705b-e162-41a4-9aa9-23b230c4511a" }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "editing the last token moved the FIRST position's logits by 0.083\n", "bidirectional: True" ] } ], "source": [ "model.eval()\n", "x = torch.randint(0, 16, (1, 6))\n", "with torch.no_grad():\n", " base = model(x)\n", " x_future_edit = x.clone()\n", " x_future_edit[0, 5] = (x[0, 5].item() + 1) % 16 # change the LAST token\n", " edited = model(x_future_edit)\n", "\n", "moved = (base[0, 0] - edited[0, 0]).abs().max().item()\n", "print(f\"editing the last token moved the FIRST position's logits by {moved:.3f}\")\n", "print(\"bidirectional:\", moved > 0)" ], "id": "c201260b" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> Causality and any-order generation are opposite design choices. An\n", "> autoregressive model *forbids* future→past information flow so it can\n", "> decode left-to-right; a diffusion model *requires* it so it can decode\n", "> in any order. Neither is “better” — they sit at different points on\n", "> the parallelism/latency frontier.\n", "\n", "## Sampling: Reveal by Confidence\n", "\n", "Generation runs the reverse process. Start from all `[MASK]`, then walk\n", "$t$ from 1 down to 0 in a chosen number of **steps**. At each step:\n", "\n", "1. The mask predictor scores every masked position.\n", "2. Keep the **highest-confidence** predictions — commit them\n", " permanently.\n", "3. Re-mask the rest and continue.\n", "\n", "This is **LLaDA’s low-confidence remasking**: the model finalizes the\n", "tokens it is surest about first and leaves the hard ones for later, when\n", "more context has been filled in. The number left masked after a step\n", "targets $\\text{length} \\times t$, so the reveal finishes exactly when\n", "$t$ hits 0. A committed token is never overwritten.\n", "\n", "Below is a **real** trajectory: we train a tiny mask predictor on a\n", "handful of memorized sentences (over a 10-word toy vocabulary), then\n", "sample from a fully masked sequence. Step through it and watch tokens\n", "appear **out of order** — confidence, not position, decides who goes\n", "next." ], "id": "3f2c8369-8e79-4853-a79a-ac9e91c83558" }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [], "id": "b62c8342" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "b2e68a29-26db-49b3-a0a0-f833792fb701" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "aa76c704-3940-4ca2-8536-2a1fbd9d8dfb" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. **Watch the order.** Note the position numbers as tokens appear —\n", "> they do *not* fill left-to-right. The model commits its most\n", "> confident guess wherever it is.\n", "> 2. **Spot the bursts.** Some steps reveal several tokens, some none —\n", "> the $\\text{length} \\times t$ target reveals in chunks, not\n", "> one-at-a-time.\n", "\n", "## Steps ↔ Quality: the Compute Dial\n", "\n", "The number of denoising steps is a **dial you turn**. With the *same*\n", "trained model, more steps means each step commits fewer (more confident)\n", "tokens, so the final sequence is sharper. Sweep the dial and watch\n", "accuracy climb:" ], "id": "a9c3bc24-47ad-4b16-82c7-bfa87e5e5e70" }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [], "id": "f3ec7c82" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "78123354-9dba-4781-ac38-32c73f923aea" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "8065645c-eb7d-4a03-bbea-4f10d58ecb84" }, { "cell_type": "markdown", "metadata": {}, "source": [ "With this toy model, **6 steps** already recovers the memorized sentence\n", "perfectly, while 1–2 steps commits too much too fast and makes mistakes.\n", "Real diffusion LMs live on the same curve — you can generate a 100-token\n", "answer in far fewer than 100 steps, trading a little quality for a lot\n", "of speed.\n", "\n", "## Common Pitfalls\n", "\n", "| Pitfall | Why it bites | Fix |\n", "|-----------------------|------------------------------------|--------------|\n", "| **Making the predictor causal** | A causal mask blocks future→past flow, so a hole can’t be filled from the tokens after it. | Use full (bidirectional) attention — no mask. |\n", "| **Scoring unmasked positions** | The model already sees those tokens; training on them leaks the answer and wastes the signal. | Cross-entropy on masked positions *only*. |\n", "| **Overwriting committed tokens** | Re-predicting an already-revealed token can flip it and destabilize the reveal. | Pin committed tokens; only ever unmask. |\n", "| **Forgetting the $1/t$ weight** | Plain masked-CE over-weights very noisy (large-$t$) draws where prediction is near-hopeless. | Weight each draw by $1/t$ (the linear-schedule NELBO). |\n", "| **Confusing `[MASK]` with an output** | If the head can emit `[MASK]`, sampling can “reveal” a mask and stall. | Head has `vocab_size` logits; `[MASK]` is input-only. |\n", "\n", "## Exercises\n", "\n", "### Exercise 1: Confirm the marginal\n", "\n", "Show empirically that the fraction of masked tokens equals $t$ for\n", "several noise levels. (You built this above — now add $t = 0.1$ and\n", "$t = 0.95$.)" ], "id": "69d0913e-f625-4744-a898-ca8f95723d5b" }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "t=0.1: masked fraction = 0.099\n", "t=0.5: masked fraction = 0.500\n", "t=0.95: masked fraction = 0.951" ] } ], "source": [ "from diffusion import forward_mask\n", "gen = torch.Generator().manual_seed(1)\n", "x = torch.zeros(1000, 32, dtype=torch.long)\n", "for t in (0.1, 0.5, 0.95):\n", " _, m = forward_mask(x, t=t, mask_id=99, generator=gen)\n", " print(f\"t={t}: masked fraction = {m.float().mean():.3f}\")" ], "id": "3a5d3a45" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2: Any-order proof\n", "\n", "Write an assertion that editing the *first* token changes the *last*\n", "position’s logits (the reverse of the direction shown in the lesson)." ], "id": "5a30759e-a96a-4a4c-a23a-3fe81dbdeda9" }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "# Your implementation here — edit x[0, 0], compare logits at position -1." ], "id": "19450327" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 3: Turn the dial\n", "\n", "Sample the trained `demo[\"model\"]` at `steps=1` and `steps=10` and print\n", "both sequences as words. How many positions differ?" ], "id": "267be8b9-8c62-468c-a34f-f9c69904492f" }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "# Your implementation here — use diffusion.sample and the WORDS list." ], "id": "4677167d" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\n", "\n", "Key takeaways:\n", "\n", "1. **A different factorization.** Diffusion LMs replace left-to-right\n", " autoregression with iterative denoising — generate by *repairing* an\n", " all-`[MASK]` sequence, in any order.\n", "2. **The forward process is free.** Corruption is a parameter-free\n", " per-token coin flip; under the linear schedule the masking\n", " probability is exactly the diffusion time $t$. All learning is in\n", " the reverse network.\n", "3. **The predictor is bidirectional.** Filling a hole needs the whole\n", " sentence, so the mask predictor drops the causal mask — the\n", " deliberate opposite of the causal guarantee that makes\n", " autoregression valid.\n", "4. **The loss is reweighted masked-CE.** The absorbing-diffusion NELBO\n", " reduces to cross-entropy on the masked positions, weighted by $1/t$\n", " — a masked language model trained at every noise level.\n", "5. **Sampling reveals by confidence.** Commit the surest predictions\n", " first, re-mask the rest (LLaDA’s low-confidence remasking); the\n", " number of denoising **steps** is a compute-vs-quality dial with no\n", " autoregressive analogue.\n", "\n", "## What’s Next\n", "\n", "Diffusion is the book’s second full generation paradigm, sitting beside\n", "the autoregressive stack that runs from\n", "[m06](../m06_transformer/lesson.qmd) through\n", "[m08](../m08_generation/lesson.qmd). From here you can fold it back into\n", "what you know: the mask predictor is the same transformer, so\n", "[quantization](../m14_quantization/lesson.qmd) and\n", "[PEFT](../m25_peft/lesson.qmd) apply unchanged, and the\n", "[evaluation](../m17_evaluation/lesson.qmd) harness scores its samples\n", "the same way. The frontier open questions — closing the last quality gap\n", "to autoregression, and *block* (semi-autoregressive) diffusion that\n", "interleaves the two paradigms — are where LLaDA and its successors are\n", "pushing now.\n", "\n", "### Going Deeper\n", "\n", "**Core Papers:**\n", "\n", "- [Structured Denoising Diffusion Models in Discrete\n", " State-Spaces](https://arxiv.org/abs/2107.03006) — Austin et al., 2021\n", " (D3PM). Introduced discrete diffusion with an **absorbing-state**\n", " (`[MASK]`) variant and drew the link to masked and autoregressive\n", " models.\n", "- [Discrete Diffusion Modeling by Estimating the Ratios of the Data\n", " Distribution](https://arxiv.org/abs/2310.16834) — Lou, Meng & Ermon,\n", " 2023 (SEDD). The **score-entropy** objective; the first discrete\n", " diffusion competitive with GPT-2 on perplexity.\n", "- [Simple and Effective Masked Diffusion Language\n", " Models](https://arxiv.org/abs/2406.07524) — Sahoo et al., 2024 (MDLM).\n", " Showed the absorbing-diffusion NELBO is a **$1/t$-weighted masked\n", " cross-entropy** — the loss this lesson builds.\n", "- [Large Language Diffusion Models](https://arxiv.org/abs/2502.09992) —\n", " Nie et al., 2025 (LLaDA). An **8B** masked diffusion model trained\n", " from scratch, competitive with LLaMA3-8B; the source of\n", " **low-confidence remasking**.\n", "\n", "**Practical Resources:**\n", "\n", "- [The Annotated D3PM / discrete diffusion\n", " notes](https://arxiv.org/abs/2107.03006) — work the absorbing-state\n", " transition matrix by hand to see why `[MASK]` is a fixed point." ], "id": "56db4f08-40b4-4195-9bd5-dfe9094793db" } ], "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "name": "python3", "display_name": "Python 3 (ipykernel)", "language": "python", "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3" }, "language_info": { "name": "python", "codemirror_mode": { "name": "ipython", "version": "3" }, "file_extension": ".py", "mimetype": "text/x-python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.14" } } }