| --- |
| 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. |
|
|
