MDLM-TinyStories

A small Masked Diffusion Language Model (MDLM) trained on TinyStories.

Model Details

How It Works

Unlike autoregressive LMs, MDLM generates text through iterative denoising:

1. Start with all [MASK] tokens
2. At each step, the bidirectional model predicts clean tokens for all masked positions
3. Gradually unmask tokens over ~100 steps
4. Uses bidirectional attention — every position attends to every other position

Based on Simple and Effective Masked Diffusion Language Models (Sahoo et al., NeurIPS 2024).

Usage

import torch
from model import MDLMConfig, MDLM, sample
from transformers import AutoTokenizer

# Load model
model = MDLM.from_pretrained("youraveragedev/mdlm-tiny-stories", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("youraveragedev/mdlm-tiny-stories")
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Generate text (unconditional)
generated_ids = sample(model, seq_len=128, batch_size=1, num_steps=100, temperature=0.7, device=device)
text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
print(text)

Training Recipe

From the MDLM paper (NeurIPS 2024):

• Noise process: Mask each token with probability t ~ U(0,1)
• Loss: Cross-entropy on masked positions, weighted by 1/t (ELBO)
• Optimizer: AdamW, lr=3e-4, linear warmup
• Schedule: Constant after warmup

Citation

@inproceedings{sahoo2024simple,
  title={Simple and Effective Masked Diffusion Language Models},
  author={Sahoo, Subham Sekhar and Arriola, Marianne and Schiff, Yair and Gokaslan, Aaron and Marroquin, Edgar and Chiu, Justin T and Rush, Alexander and Kuleshov, Volodymyr},
  booktitle={NeurIPS},
  year={2024}
}