tchauffi/diffusionlm-from-scratch

diffusionlm-from-scratch — masked diffusion LM (DiT, 142M)

A masked (absorbing-state) diffusion language model, built and trained from scratch on TinyStories. Instead of generating left-to-right one token at a time, it starts from a sequence of pure [MASK] tokens and denoises the whole sequence in parallel — committing the tokens it is most confident about first, in whatever order the meaning falls into place.

• —Code, training & sampling: https://github.com/tchauffi/diffusionlm-from-scratch
• —Course / write-up: `RESEARCH.md` — a from-scratch course on discrete/text diffusion (D3PM → absorbing-state → sampling).
• —Demo site: animated real generations live in `docs/`.

Model

Key finding: uniform loss weighting (w(t) = 1), not the textbook ELBO weight 1/σ(t), is what turned word-salad into coherent stories.

Files

• —final.pt — checkpoint with two state dicts, model (EMA, preferred) and raw, plus the config used to build the model.
• —tokenizer.json, tokenizer_config.json — the byte-level BPE tokenizer (PreTrainedTokenizerFast; special tokens [PAD] [UNK] [MASK] <|endoftext|>).

Usage

Install the model code from the GitHub repo, then generate stories in two lines — DiffusionLM bundles the model, tokenizer, and absorbing-state scheduler:

from diffusionlm_from_scratch import DiffusionLM

lm = DiffusionLM.from_pretrained("tchauffi/diffusionlm-from-scratch")
for story in lm.generate(n=4, seq_len=80, temperature=0.9):
    print(story)

generate exposes the sampler knobs (order, steps, corrector_frac, confidence_threshold, …). For lower-level access, load just the model:

from diffusionlm_from_scratch.model import DiT

model = DiT.from_pretrained("tchauffi/diffusionlm-from-scratch")  # downloads final.pt
# the raw checkpoint carries ck["config"], ck["model"] (EMA), and ck["raw"].

See `scripts/capture_trajectories.py` in the repo for the full parallel-denoising sampling loop.