---
license: mit
language:
- en
library_name: pytorch
tags:
- diffusion-language-model
- masked-diffusion
- autoregressive
- text-generation
- tinystories
datasets:
- roneneldan/TinyStories
---

# nanoDiffusionLab TinyStories 106M: AR + Masked Diffusion

This repository contains the paired inference weights used by
[nanoDiffusionLab](https://github.com/tutao0123/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](https://github.com/tutao0123/nanoDiffusionLab/blob/main/reports/tinystories_106m_generation.md)
for prompts, sampling settings, uncertainty, memory, and judge-audit details.

## Interactive demo

Try the paired models in the
[nanoDiffusionLab AR vs MDLM Studio](https://modelscope.cn/studios/taoshuai/nanoDiffusionLab)
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.