5.4 diffusion LM micro-implementation

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; DiffusionBERT, Sahoo et al. 2024). 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

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 — the D3PM paper introducing discrete diffusion. Cite for the math.
• Sahoo et al. 2024 — 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.