Minimal Masked Diffusion Language Model

A clean, minimal implementation of Masked Diffusion Models for language modeling.

Files

model.py              # Bidirectional transformer (RoPE, RMSNorm, SwiGLU)
diffusion.py          # Forward masking + reverse sampling
data.py               # Data loading (nanoGPT-style binary format)
train.py              # Training script (Hydra + wandb)
sample.py             # Text generation
eval.py               # Evaluation (loss, perplexity)
configs/
  config.yaml         # Default configuration
  experiment/         # Experiment configs (small, medium)

Quick Start

# Install dependencies
pip install torch numpy hydra-core omegaconf wandb

# Train (auto-downloads Shakespeare)
python train.py

# Generate text (checkpoint saved in outputs/{date}/{time}/checkpoints/)
python sample.py checkpoint=outputs/2024-01-01/12-00-00/checkpoints/ckpt.pt prompt="ROMEO:"

# Evaluate
python eval.py checkpoint=outputs/2024-01-01/12-00-00/checkpoints/ckpt.pt

Dataset

Shakespeare Character-Level

The default dataset is Tiny Shakespeare (~1MB, 1.1M characters), automatically downloaded on first run.

Data format (nanoGPT-style):

• train.bin, val.bin: Binary files of uint16 token IDs
• meta.pkl: Vocabulary metadata (stoi, itos, vocab_size)

Preparation (automatic or manual):

# Auto-downloads and prepares data on first train.py run
# Or manually:
python -c "from data import prepare_shakespeare_char; prepare_shakespeare_char('data/shakespeare_char')"

Adding Custom Datasets

To use your own data, create train.bin and val.bin with uint16 token IDs:

import numpy as np
import pickle

# Your tokenized data
train_ids = [...]  # list of token IDs
val_ids = [...]

# Save binary files
np.array(train_ids, dtype=np.uint16).tofile('data/mydata/train.bin')
np.array(val_ids, dtype=np.uint16).tofile('data/mydata/val.bin')

# Save metadata
meta = {'vocab_size': YOUR_VOCAB_SIZE, 'stoi': {...}, 'itos': {...}}
with open('data/mydata/meta.pkl', 'wb') as f:
    pickle.dump(meta, f)

Then train with: python train.py data.dir=data/mydata

Evaluation

Metrics

Monte Carlo Estimation

Since MDMs use random masking, evaluation uses Monte Carlo sampling:

• Average loss over mc_samples different random maskings per batch
• Default: 16 MC samples × 20 eval batches = 320 forward passes

# Standard evaluation
python eval.py checkpoint=out/ckpt.pt

# More accurate (more MC samples)
python eval.py checkpoint=out/ckpt.pt mc_samples=64 eval_iters=50

# Evaluate on train split
python eval.py checkpoint=out/ckpt.pt split=train

Expected Results (Shakespeare char-level)

Results vary with training iterations and hyperparameters.

Configuration (Hydra)

All scripts use Hydra for configuration management.

# Override any config parameter
python train.py training.max_iters=10000 model.n_layer=6

# Use experiment preset
python train.py +experiment=medium

# Enable wandb logging
python train.py wandb.enabled=true wandb.project=my-project

# Multiple overrides
python train.py model.n_layer=6 model.n_head=6 model.n_embd=384 \
                training.learning_rate=6e-4 wandb.enabled=true

Weights & Biases Integration

Organized wandb logging with metric groups:

• loss/: train, train_avg
• perf/: iter_per_sec, tokens_per_sec, time_ms
• optim/: lr, grad_norm
• eval/: train_loss, val_loss, train_ppl, val_ppl, train_bpc, val_bpc
• samples: Generated text samples (logged periodically)
• Config: Full Hydra config logged automatically
• Model: Parameter count, architecture details

# Enable wandb
python train.py wandb.enabled=true

# Custom project/run name
python train.py wandb.enabled=true wandb.project=mdm-experiments wandb.name=my-run

Core Concepts

Forward Process (Masking)

p_mask(t) = (1 - ε) * t + ε    # ε=0.001

• t=0: Almost no masking
• t=1: Complete masking

Training Loss

L = E[CE(model(x_masked), x_original) / p_mask]

Importance weighting by 1/p_mask ensures proper likelihood estimation.

Generation (Reverse Process)

1. Start fully masked
2. Model predicts all positions
3. Unmask fraction 1 - s/t each step
4. Repeat until clean

Strategies: stochastic (random) or confidence (highest confidence first)

Architecture

• Bidirectional attention (no causal mask)
• RMSNorm + RoPE + SwiGLU
• Mask token = vocab_size (appended to vocabulary)

Training Examples

# Default small model
python train.py

# Medium model with wandb
python train.py +experiment=medium wandb.enabled=true

# Custom model
python train.py model.n_layer=8 model.n_head=8 model.n_embd=512 \
                training.max_iters=20000

# CPU training
python train.py system.device=cpu system.dtype=float32

Sampling Examples

# Basic sampling (use checkpoint path from training output)
python sample.py checkpoint=outputs/2024-01-01/12-00-00/checkpoints/ckpt.pt

# With prompt
python sample.py checkpoint=outputs/.../checkpoints/ckpt.pt prompt="To be or not"

# Adjust quality/diversity
python sample.py checkpoint=outputs/.../checkpoints/ckpt.pt \
                sampling.steps=128 sampling.temperature=0.8

# Confidence-based sampling
python sample.py checkpoint=outputs/.../checkpoints/ckpt.pt sampling.strategy=confidence

# Multiple samples
python sample.py checkpoint=outputs/.../checkpoints/ckpt.pt num_samples=5

References

• SMDM: Scaling up Masked Diffusion Models
• D3PM: Discrete Denoising Diffusion
• nanoGPT: Minimal GPT