MDLM-TinyShakespeare

A masked diffusion language model (MDLM) trained from scratch on character-level TinyShakespeare. Unlike an autoregressive model, it generates text by iterative denoising: starting from an all-[MASK] sequence and refining it over multiple steps until fully unmasked, rather than left-to-right, one token at a time.

This model was built as a from-scratch learning project — architecture, training loop, and sampler were all implemented directly (no pretrained backbone, no existing diffusion-LM library) — following the MDLM objective from Sahoo et al. (2024) and the LLaDA training recipe (Nie et al., 2025), with an optional ReMDM remasking sampler for higher-quality generation.

Model description

• Architecture: bidirectional Transformer denoiser — RMSNorm, RoPE, SwiGLU feed-forward, sigmoid-gated multi-head attention (non-causal), with sinusoidal timestep conditioning.
• Objective: masked diffusion (absorbing-state forward process + 1/t-weighted ELBO loss), not next-token prediction.
• Tokenization: character-level (vocabulary built directly from the training corpus, plus one reserved [MASK] token).
• Parameters: ~11.4M.

Hyperparameter	Value
dim (embedding / hidden size)	384
n_heads	6
n_layers	6
max_seq_len / BLOCK	256
vocab_size	65 (char-level) + 1 [MASK] token

(Adjust the table above if your actual run used different values — these are the guide's defaults.)

Training details

• Data: TinyShakespeare (~1MB of Shakespeare's plays), 90/10 train/val split, random character-level crops of length 256.
• Optimizer: AdamW, lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1, linear warmup over 200 steps then held constant.
• Batch size: 64.
• Precision: bfloat16 (autocast) where supported, fp16 + GradScaler otherwise.
• Steps: 5,000.

Loss curve (masked-diffusion ELBO, 1/t-weighted — not directly comparable to autoregressive cross-entropy loss numbers):

Step	Train loss	Val loss
0	4.101	4.117
500	2.079	2.148
1000	1.762	1.943
1500	1.633	1.848
2000	1.625	1.827
2500	1.601	1.891
5000 (final)	<fill in>	<fill in>

How to use

Install dependencies and download the weights + config:

pip install huggingface_hub safetensors torch

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import json, torch

REPO_ID = "your-username/mdlm-tinyshakespeare"   # <-- update to your actual repo

cfg = json.load(open(hf_hub_download(REPO_ID, "config.json")))
weights_path = hf_hub_download(REPO_ID, "model.safetensors")

# DiffusionLM class definition -- copy from the training guide (or vendor it alongside this repo)
model = DiffusionLM(cfg["vocab_size"], cfg["dim"], cfg["n_heads"], cfg["n_layers"],
                    max_seq_len=cfg["max_seq_len"]).to("cuda")
model.load_state_dict(load_file(weights_path))
model.mask_id = cfg["mask_id"]
model.eval()

Generating text — base sampler

Confidence-based iterative unmasking. Fast, but once a token is unmasked it's permanent:

print(generate(model, length=256, steps=128, temperature=0.6, top_k=50))

Generating text — ReMDM sampler (recommended for better quality)

Adds a remasking step: low-confidence already-unmasked tokens can be sent back to [MASK] and re-predicted with more context, fixing "committed too early" errors the base sampler can't:

print(generate_remdm(model, length=256, steps=128, temperature=0.6, eta_max=0.15))

Tip: if you see occasional garbled near-miss words in the output (e.g. a word that's almost right), try more steps, a lower temperature, and/or the ReMDM sampler above — all three trade a bit of compute for noticeably cleaner text, with no retraining required.

Sample output

Generated with the base sampler at default settings (steps=64, temperature=0.8):

ine,
And they have cone, they have not no made,
Than in the  name of the king of the  have life:
I have deceived and this father then  aad I am not aou to  thise,
I mame to have talk'd to move be gone,
And live to the people: like to  the caust:
I have see

Limitations

• Small model, small data. 11.4M parameters trained on ~1MB of text — this is a learning/demonstration model, not a general-purpose text generator. Expect Shakespeare-flavored structure (verse lines, character-name capitalization, period-appropriate vocabulary and contractions) rather than coherent long-form text.
• Character-level near-misses. The base sampler occasionally produces words that are one or two characters off from a real word (e.g. "thise" instead of "this") — a known failure mode of confidence-based unmasking with no correction mechanism. The ReMDM sampler mitigates this.
• Not comparable to autoregressive LM benchmarks. Loss values use a masked-diffusion ELBO objective; don't compare them directly to next-token-prediction cross-entropy loss from AR models.
• English, Early Modern spelling/grammar conventions only — trained exclusively on Shakespeare's plays.

Acknowledgements / references

• Sahoo et al., Simple and Effective Masked Diffusion Language Models, arXiv:2406.07524 — the Rao-Blackwellized training objective this model uses.
• Nie et al., Large Language Diffusion Models (LLaDA), arXiv:2502.09992 — training recipe reference.
• Wang et al., Remasking Discrete Diffusion Models with Inference-Time Scaling (ReMDM), arXiv:2503.00307 — the optional remasking sampler included above.

License

This model is released under the MIT License. TinyShakespeare is derived from public-domain text.