mdlm-lm1b-gpt2-packed

mdlm-lm1b-gpt2-packed

A Masked Diffusion Language Model (MDLM) trained from scratch on the One Billion Word Benchmark (LM1B) with a GPT-2 BPE tokenizer and packed sequences — no padding anywhere in training. This is an independent retraining using the MDLM codebase, released as a frozen research artifact; it is not an official checkpoint of the MDLM authors and its perplexities are not comparable to the numbers reported in their paper (different tokenizer and sequence protocol; see below).

Protocol: packed GPT-2, no padding

This checkpoint deliberately departs from the paper's LM1B recipe:

Details of the packing: every document is tokenized with GPT-2 BPE, one EOS (<|endoftext|>, id 50256) is inserted per document, the resulting stream is concatenated and split into contiguous blocks of exactly model.length = 128 tokens (6,984,037 training blocks; ~13,641 steps per epoch at global batch 512). Token id 50257 is a [PAD] token appended to the GPT-2 vocabulary (vocab size 50,258), but it is used exclusively as the absorbing MASK state of the diffusion process — it never appears as sequence padding. Any evaluation of this checkpoint must reproduce this packed protocol: the checkpoint bakes in its training distribution, and feeding it padded single sentences measures a different object.

Consequence for numbers: perplexities are on the GPT-2-token basis over packed blocks. The measured validation PPL of ≈ 34.4 at this snapshot corresponds to the same quality class as ≈ 27 PPL on the padded-BERT-token basis used in the literature. Do not compare the raw numbers across bases.

Model

• Architecture: MDLM (SUBS parameterization, time-conditioning off, T=0 continuous-time ELBO), DiT backbone — 12 blocks, hidden 768, 12 heads, rotary position embeddings, context length 128, vocab 50,258. 169.6M parameters.
• Noise schedule: log-linear; antithetic sampling on; loss: ELBO (bf16).
• Training: 4× H100 (DDP), global batch 512, Adam lr 3e-4 (constant after 2,500 warmup steps, β = 0.9/0.999, weight decay 0), gradient clip 1.0, EMA 0.9999, seed 1, PyTorch Lightning 2.2.1.
• Data: billion-word-benchmark/lm1b at the Hub's parquet conversion (refs/convert/parquet), standard train/test splits.

This checkpoint

• Exact 500,000-step milestone (epoch 36) of a planned 1M-step run; the run is currently paused at ~528k steps and may be extended — if a 1M checkpoint is released later it will appear alongside this one, unchanged.
• EMA validation perplexity at this snapshot: ≈ 34.4 (GPT-2 basis, packed protocol; still improving at ~0.1–0.15 PPL per 25k steps when paused; no sign of overfitting after ~37 epochs).
• The file is the full Lightning checkpoint (weights + EMA shadow + optimizer state; resumable), ≈ 2.7 GB, with run-tracking metadata sanitized before release: the embedded config's experiment-tracker fields and local filesystem paths were set to null. All 657 tensors (weights, EMA, and optimizer state) are bit-identical to the internal training milestone — verified tensor-by-tensor.
• sha256(milestone-500000.ckpt) = d5a2f55de42b755a1d72fb7a9abd3a44aa27c9bc4410a3cfce8d2a46096fbfc2

Intended use

Released as a frozen measurement head for information-theoretic analyses of masked diffusion models (per-position conditionals, entropies, coupling measurements) and as a reproducible LM1B diffusion baseline under a modern BPE tokenizer. It is a 128-token research model, not an assistant and not tuned for generation quality.

To load it, use the MDLM codebase (algo=mdlm, model=small, model.length=128) with a GPT-2 tokenizer pipeline that reproduces the packed protocol above (wrap=True, EOS inserted per document, block size 128). The original MDLM repository's LM1B data pipeline uses the padded-BERT recipe and will NOT match this checkpoint.

Attribution

The model architecture and all training code are from MDLM by Sahoo et al. (NeurIPS 2024), code at kuleshov-group/mdlm (Apache-2.0). This checkpoint was trained from scratch by the uploader on a fork of that codebase; the only substantive changes are the packed GPT-2 data pipeline described above and a datasets>=3 compatibility fix for the LM1B loader. All credit for the method and implementation belongs to the MDLM authors; any defects of this particular training run are the uploader's.

@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={Advances in Neural Information Processing Systems},
  year={2024}
}

Dataset: Chelba et al., One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling (2013).

License

Apache-2.0, matching the upstream MDLM code license.