nanoDiffusionLab-TinyStories-106M

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

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.