nanoDiffusionLab TinyStories 106M: AR + Masked Diffusion
This repository contains the paired inference weights used by nanoDiffusionLab: a GPT-style autoregressive (AR) model and a masked-diffusion language model (MDLM) trained from scratch for a controlled comparison.
Both models use the same 12-layer, 576-dimensional Transformer backbone, GPT-2 BPE tokenizer, 1,024-token context window, pinned TinyStories dataset revision, and 2.000B-token training budget. The AR model has 106.28M parameters; the MDLM has 106.61M parameters because it includes a mask embedding and time-conditioning MLP.
Files
| File | Purpose |
|---|---|
ar.safetensors |
BF16 inference weights for the causal AR baseline |
ar_config.json |
AR architecture, tokenizer, dataset, and training metadata |
mdlm.safetensors |
BF16 inference weights for masked-diffusion generation |
mdlm_config.json |
MDLM architecture, tokenizer, dataset, and training metadata |
The compact exports omit optimizer, scaler, RNG, and distributed-resume state from the original training checkpoints. They are intended for inference, visualization, and reproducible comparisons.
Results and limitations
On one NVIDIA A40 using BF16 and batch size 1, the MDLM was 13.63x faster than cached AR at 8 denoising steps and 1.76x faster at 64 steps. This speed does not imply equal quality: across 8,000 blinded pairwise comparisons, AR won the large majority. MDLM utility increased from 0.0028 at 8 steps to 0.0745 at 64 steps. The current MDLM commonly produces repetition and incoherent stories.
These figures are implementation- and hardware-specific. They describe this compact baseline, not diffusion language models in general. See the full benchmark for prompts, sampling settings, uncertainty, memory, and judge-audit details.
Interactive demo
Try the paired models in the nanoDiffusionLab AR vs MDLM Studio hosted on ModelScope.
Reproducibility
The architecture and sampler are implemented in the source repository. The export utility is
scripts/export_hf_checkpoint.py. Generation uses GPT-2 BPE, temperature 1.0, top-k 40, cached AR
decoding, and confidence-based parallel unmasking for MDLM by default.
License
MIT. The TinyStories dataset has its own terms and provenance; consult its dataset card before use.