Deep Learning · Diffusion & Flow Models · Large Language Diffusion Models· Part 8 of 8

Build a Tiny Diffusion Language Model

September 2026 · Part 8 of 8 · 7 min read

Build something small enough to audit. This lab has two parts: an exact sampler calculation using only Python's standard library, and a 68,294-parameter transformer that learns a tiny synthetic language. No model weights or datasets are downloaded.

Download and run

Save tiny_diffusion.py. The exact oracle needs Python 3.9 or newer. The neural lab additionally needs PyTorch and NumPy. Use a separate environment so the experiment's dependencies stay isolated.

# Run from the directory where you saved the script.
python3 tiny_diffusion.py --oracle-only

python3 -m venv .venv
source .venv/bin/activate
python -m pip install torch numpy
python tiny_diffusion.py --steps 800 --samples 512 --seed 7

The recorded run below used Python 3.12, PyTorch 2.14.0, CPU execution with two threads, and seed 7. Library versions and platforms can change numerical results; the exact oracle is the platform-independent reference. The script prints JSON and optionally saves it with --output result.json.

Experiment 1: remove learning from the problem

The script explicitly propagates probability mass over every two-position state in the red/red–blue/blue example from Part 3. At each reverse interval, it constructs the exact posterior for each masked position, applies independent reveal decisions, and sums the probability of every destination.

This is exact enumeration, not a Monte Carlo simulation. It checks normalization after every interval and verifies that the final distribution contains no masks. For $K\in\{1,2,4,8,16,32\}$, it also checks the mixed-color probability against $1/(2K)$ to numerical tolerance. If this fails, there is a sampler or arithmetic bug; training is not involved.

Experiment 2: a language with three dependencies

Construct a first half by choosing independently from three pairs: red/blue, cat/dog, and runs/rests. Repeat it exactly:

red cat runs red cat runs
blue dog rests blue dog rests
red dog rests red dog rests

There are eight valid sequences, each equally likely. These strings are synthetic symbols arranged to expose copying dependencies, not examples of ordinary English grammar. A valid output must choose the correct category at each position and duplicate the first half exactly. Observing one member of a paired position reveals the other; if both are masked, the clean value remains uncertain.

Training examples are generated from this known distribution. Evaluation uses fresh random corruptions and new generated samples, but the support is the same eight strings. This is a controlled learning-and-sampling demonstration, not a test of generalization to unseen natural language or unseen sequence types.

The network has no special diffusion layer

The model adds learned token and position embeddings, applies two independently initialized bidirectional transformer layers, then projects to six clean-token logits. The input embedding has seven entries because it also accepts the mask symbol. The output head has six entries, so it cannot emit a mask as a clean prediction.

# Shape sketch; the downloadable script contains the complete model.
h = token_embedding(xt) + position_embedding(positions)  # [B, 6, 64]
for layer in layers:
    h = layer(h)  # full attention over the corrupted sequence
logits = head(norm(h))  # [B, 6, 6 clean vocabulary items]

There is no explicit time embedding: the posterior's time-independence argument applies to this homogeneous mask process. The sampler still uses time to decide reveal probabilities. The training loop uses the forced-target estimator, which scores one uniformly selected masked target per sequence. Targets are aligned with the same input position, with no causal next-token shift.

Know the optimal loss before looking at the curve

Condition on a selected target being masked. Its matching partner is also masked with probability $t$. If the partner is visible, the clean target is known and optimal cross-entropy is zero. If the partner is masked, the target is a fair binary choice and optimal cross-entropy is $\log2$. Thus

$$\mathbb E[\mathrm{CE}_{\mathrm{target}}\mid t]=t\log2,\qquad\mathbb E_{t\sim U(0,1)}[\mathrm{CE}_{\mathrm{target}}]=\tfrac12\log2\approx0.3466.$$

A loss near $0.35$ can be excellent for this task. Expecting it to converge to zero would confuse irreducible uncertainty with model error. A single minibatch can fall below the expected optimum because it contains a particular random mix of easy and ambiguous targets.

A recorded CPU run

With 800 optimizer steps, a batch size of 64, learning rate $10^{-3}$, and 512 generated samples per setting, the included script produced:

Intervals $K$	Learned model: valid outputs	Exact posterior: expected validity
1	11.5%	12.5%
2	37.9%	42.2%
4	65.4%	67.0%
8	79.1%	82.4%
16	88.9%	90.9%
32	92.8%	95.4%

Download the full recorded output, including masking-level losses and software version. All eight valid sequence types appeared at every setting; this coverage count alone does not establish that they were generated uniformly. The exact reference in the third column is $(1-1/(2K))^3$, because the three independent duplication pairs must each agree under a perfect posterior and this sampler.

The learned-model percentages are one finite-sample run, not exact probabilities. With 512 samples, a binomial proportion has standard error at most about 2.2 percentage points. The gap from the exact reference combines denoiser error and sampling variation. The much larger change between one and many intervals is already predicted by the oracle analysis, without claiming that extra steps improved the trained network.

Read the implementation in this order

oracle_distribution: verify a normalized transition before touching neural code.
data: identify the exact support and its dependencies.
Denoiser.forward: trace tensor shapes and attention visibility.
The training loop: find the forced target, corruption mask, and unshifted label.
sample: distinguish value proposals from reveal decisions.
metrics: check full-sequence constraints, not only token accuracy.

The script also samples with a fixed first token and asserts that it is preserved. That is a sampler invariant, not a claim that every conditional output will satisfy the full language. The final reveal interval guarantees no masks remain.

Three experiments with specific predictions

Reveal one token at a time with recomputation. Under the oracle, duplication errors disappear. Under the learned model, any residual error exposes imperfect posteriors or out-of-support intermediate states.
Use argmax instead of categorical draws. Ambiguous contexts can collapse toward a small subset of strings. Measure the frequency of all eight sequences, not just validity.
Use one denoising pass after much longer training. Even an exact posterior still yields only $1/8$ valid sequences under independent one-shot sampling. More optimization cannot remove that decoder limitation.

Check your understanding: Why can denoising accuracy look strong while one-shot sequence validity is poor?

The denoiser can correctly express uncertainty at each masked position. Independent sampling does not coordinate the three duplicated pairs. Per-position calibration and joint-sample consistency are different properties.

The next experiment after this lab

Replace the synthetic data generator with a small, licensed text corpus and an explicit train/validation split. Add a tokenizer, padding and attention boundaries, variable lengths, and held-out denoising estimates. Keep a fixed sampler baseline before experimenting with confidence selection. Use the evaluation record from Part 7 to track what changed.

For model-scale implementation details, move to the MDLM implementation or the repositories linked in Part 5. The purpose of this lab is to make their objectives and decoding loops readable, not to reproduce their training scale.