README.md

metadata

library_name: pytorch
tags:
  - diffusion
  - language-model
  - discrete-diffusion
  - absorbing-state
  - text-generation
  - tinystories
  - from-scratch
datasets:
  - roneneldan/TinyStories
pipeline_tag: text-generation

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

Architecture	DiT (transformer denoiser), bidirectional attention, adaLN-Zero
Parameters	~142M
Hidden size / depth / heads	768 / 12 / 12
MLP ratio	4.0
Vocab	8,192 (byte-level BPE, trained on TinyStories)
Max sequence length	256
Diffusion	absorbing-state (masked) discrete diffusion
Training data	TinyStories
Eval cross-entropy	2.18

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:

import torch
from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast
from diffusionlm_from_scratch.model import DiT, DiTConfig

repo = "tchauffi/diffusionlm-from-scratch"

ckpt_path = hf_hub_download(repo, "final.pt")
ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)

model = DiT(DiTConfig(**ck["config"]))
model.load_state_dict(ck["model"])   # EMA weights ("raw" also available)
model.eval()

tokenizer = PreTrainedTokenizerFast.from_pretrained(repo)

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