Training a Masked Diffusion Language Model on a MacBook

Autoregressive models generate text one token at a time, left to right. Diffusion models do something stranger: start from pure noise and iteratively denoise the whole sequence at once.

For text, "noise" means masks. I trained a masked discrete diffusion LM from scratch on TinyStories — on a MacBook (M4 Pro, 24 GB unified RAM).

Repo: github.com/Adityaadpandey/diffusion_llm

How masked discrete diffusion works

Forward process (corruption): take clean text, randomly mask tokens. More timesteps = more masks. At t=T the sequence is 100% `[MASK]`.

Reverse process (generation): a model predicts the original tokens at every masked position. At inference you start fully masked and progressively unmask over N steps.

The model is a BERT-style bidirectional transformer — and it has to be. The reverse process must see all currently-revealed tokens to predict any masked position. Causal attention would be throwing away exactly the information you need.

Key design choices:

• Timestep conditioning — sinusoidal timestep embeddings added to token embeddings, so the model knows how corrupted the input is
• Cosine corruption schedule — spends more steps near the clean end, where fine-grained predictions matter most
• ELBO loss — cross-entropy only on masked positions, averaged per masked token

The sampling gotcha

There are two ways to decide which tokens to unmask each step:

• Confidence-first: unmask the tokens the model is most sure about. Sounds smart — and on an undertrained model it collapses into repetitive junk, because the model is most confident about boring high-frequency tokens.
• Random posterior: use the analytic masked-diffusion posterior to pick positions randomly. More diverse, more stable, and much better early in training.

Random posterior is the default for a reason. I learned this the annoying way.

Training on Apple Silicon

The whole thing runs on MPS (Metal). Practical notes that cost me time:

• Python 3.11/3.12 only — 3.14 is too new for stable ML wheels
• Default config: ~30M params (d_model=384, 6 layers, seq_len 128) — fits comfortably in 24 GB unified memory
• A ~22M param proof-of-concept run (50k steps) already produces grammatical English
• If the MPS dataloader gets flaky, `num_workers: 0` is your friend
• Gradient accumulation over bigger batches when memory gets tight

Targets: validation perplexity < 30 after the first small run, < 15 after 100k–200k steps at default size — at which point the generated TinyStories are coherent little stories.

Why bother?

Because bidirectional generation is a genuinely different capability: infilling, editing, and constrained generation come naturally when the model denoises the whole sequence instead of committing left to right. And because training one from scratch on a laptop teaches you more about diffusion LMs than reading ten papers.