---
title: "5.4 diffusion LM micro-implementation"
sidebar:
  order: 504
---

<!-- PRD metadata (moved from frontmatter for Blume schema compatibility) -->
**status:** not-started  
**owner:** unassigned  
**created:** 2026-06-13  
**parent_plan:** docs/PLAN.md §3 Tier 5 (5.4)  
**related_prds:** 5.1-reasoning-on-22M.md, 5.2-testtime-compute-scaling.md (sibling Tier-5 research arcs)

# PRD — Diffusion language model toy via masked denoising

## Goal

Train a small (22M-class) diffusion language model from scratch on
the masked-denoising objective ([Austin et al. 2021](https://arxiv.org/abs/2107.03006);
[DiffusionBERT, Sahoo et al. 2024](https://arxiv.org/abs/2406.07524)).
Different generation paradigm than autoregressive: model learns to
denoise a fully-masked sequence over T diffusion steps. Publish as
"smallest from-scratch text diffusion model on consumer hardware."

Tier-5 research arc. ~1–2 weeks. Publishable shape: same
"smallest-X-on-consumer-hardware" frame as 5.3 / 5.6.

## Why now

- The platform has no point on the discrete-diffusion curve.
  Adding one is a clean side-experiment that doesn't interfere
  with the AR mainline.
- Discrete diffusion has been quietly gaining ground (Sahoo's
  DiffusionBERT, Google's recent Discrete Diffusion LM work).
  Worth a research-grade artifact.

## Scope — in

- **Model:** same transformer decoder we have, but training
  objective changes to masked-token denoising over T (default 20)
  diffusion steps. The model takes `(x_t, t)` as input.
- **Time embedding:** sinusoidal step embedding added to the
  position embedding.
- **Loss:** cross-entropy on masked positions only, weighted by
  the diffusion noise schedule.
- **Sampling:** T-step iterative refinement. Start from
  fully-masked; at each step, unmask the model's top-confidence
  predictions for some fraction of positions; stop when fully
  unmasked.
- **Eval:** generation quality on the standard from-scratch tinybench
  (a few hundred shakespeare prompts; manual quality + length-
  match rubric), plus a held-out PPL-via-denoising-likelihood
  (the standard discrete-diffusion likelihood lower bound).
- **Artifact:** `docs/research/diffusion-lm-toy.md` —
  paper-shaped report.

## Scope — out

- **Continuous-token diffusion** (operating in a latent embedding
  space). V1 is discrete-masked.
- **Distillation to a 1-step or few-step sampler.** Defer.
- **Conditional generation** (text prompts → continuation). V1 is
  unconditional + length-conditioned only.

## Files to touch

| File | Change |
|---|---|
| `Sources/TinyGPTModel/DiffusionLM.swift` | new — masked-denoising training + sampling |
| `Sources/TinyGPT/TrainDiffusion.swift` | new — recipe |
| `Sources/TinyGPT/SampleDiffusion.swift` | new — T-step iterative refinement |
| `Sources/TinyGPT/TinyGPT.swift` | two new `case`s |
| `docs/research/diffusion-lm-toy.md` | new — results |
| `evals/diffusion-smoke.sh` | new — 100-step train + 20-sample generate; assert non-degenerate |
| `docs/PLAN.md` | 5.4 ⬜ → ✅ on ship |

## Don't touch

- The autoregressive training paths — diffusion is a sibling, not
  a refactor. Both objectives live in `TinyGPTModel/` side by side.

## Acceptance criteria

- [ ] `posttrainllm train-diffusion --depth 12 --corpus shakespeare.txt
  --steps 5000 --t-steps 20 --out d.tinygpt` runs end-to-end on
  M5 Pro.
- [ ] `posttrainllm sample-diffusion d.tinygpt --length 64 --t-steps 20`
  produces strings that resemble shakespeare-like text (manual
  rubric: most samples grammatical English; not random tokens).
- [ ] `docs/research/diffusion-lm-toy.md` reports the recipe + a
  qualitative comparison to a same-budget AR baseline.
- [ ] Smoke passes in CI.

## Reference patterns

- `Sources/TinyGPT/Train.swift` — the loop shape; we replace the
  objective.
- [Austin et al. 2021](https://arxiv.org/abs/2107.03006) — the
  D3PM paper introducing discrete diffusion. Cite for the math.
- [Sahoo et al. 2024](https://arxiv.org/abs/2406.07524) — modern
  practical recipe at small scale.

## Open questions

- Whether to share embeddings + transformer weights with the AR
  path (multi-objective training, like UL2) or keep separate.
  **Recommendation:** keep separate for V1; multi-objective is V2
  research.