diffusion-51M-base / README.md
d0rj's picture
Add eight full comparison benchmarks and TensorBoard metrics at step 15000
d532205 verified
|
Raw History Blame Contribute Delete
19 kB
---
language:
- en
library_name: transformers
pipeline_tag: other
tags:
- custom_code
- tensorboard
- tiny-llm-ablation
- from-scratch
- diffusion
- masked-language-modeling
datasets:
- HuggingFaceFW/fineweb-edu
model-index:
- name: diffusion-51M-base
results:
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: Rowan/hellaswag
name: HellaSwag
config: default
split: validation
metrics:
- type: pll_acc_norm
name: PLL acc_norm
value: 0.2795259908384784
args:
num_few_shot: 0
ci95_low: 0.2708343265409323
ci95_high: 0.2883862705438426
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: allenai/ai2_arc
name: ARC-Easy
config: ARC-Easy
split: test
metrics:
- type: pll_acc_norm
name: PLL acc_norm
value: 0.3362794612794613
args:
num_few_shot: 0
ci95_low: 0.3175609903739864
ci95_high: 0.3555264763913292
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: allenai/ai2_arc
name: ARC-Challenge
config: ARC-Challenge
split: test
metrics:
- type: pll_acc_norm
name: PLL acc_norm
value: 0.2226962457337884
args:
num_few_shot: 0
ci95_low: 0.19980419390295
ci95_high: 0.2474001931137188
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: baber/piqa
name: PIQA
config: default
split: validation
metrics:
- type: pll_acc_norm
name: PLL acc_norm
value: 0.5386289445048966
args:
num_few_shot: 0
ci95_low: 0.5157819596226829
ci95_high: 0.5613147955404465
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: allenai/winogrande
name: WinoGrande
config: winogrande_xl
split: validation
metrics:
- type: pll_acc
name: PLL acc
value: 0.5011838989739542
args:
num_few_shot: 0
ci95_low: 0.47369054474157507
ci95_high: 0.5286700959026042
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: allenai/openbookqa
name: OpenBookQA
config: main
split: test
metrics:
- type: pll_acc_norm
name: PLL acc_norm
value: 0.26
args:
num_few_shot: 0
ci95_low: 0.22348572318538285
ci95_high: 0.3001739602361623
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: aps/super_glue
name: BoolQ
config: boolq
split: validation
metrics:
- type: pll_acc
name: PLL acc
value: 0.5296636085626911
args:
num_few_shot: 0
ci95_low: 0.5125316201633632
ci95_high: 0.5467259836150076
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL / reconstruction
dataset:
type: EleutherAI/lambada_openai
name: LAMBADA OpenAI reconstruction
config: default
split: test
metrics:
- type: pll_acc
name: PLL acc
value: 0.42208422278284496
args:
num_few_shot: 0
ci95_low: 0.4086621968305708
ci95_high: 0.4356223315164309
ci_method: Wilson
protocol: single-mask continuation PLL
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: AxiomicLabs/Arithmark-3.0
name: ArithMark-3
config: default
split: train
metrics:
- type: pll_acc_norm
name: pll_acc_norm (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.332
args:
dtype: bfloat16
num_few_shot: 0
max_length: 1024
standard_error: 0.014899597242811565
evaluation_date: '2026-10-01'
ci95_low: 0.30350363489239063
ci95_high: 0.36178215595875596
ci_method: Wilson 95%; item independence approximation
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: pkavumba/balanced-copa
name: Balanced COPA
config: default
split: train
metrics:
- type: pll_acc
name: pll_acc (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.522
args:
dtype: bfloat16
num_few_shot: 0
max_length: 2048
standard_error: 0.01580397942816194
evaluation_date: '2026-10-01'
ci95_low: 0.4910152521271681
ci95_high: 0.5528163704994674
ci_method: Wilson 95%; item independence approximation
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: tau/commonsense_qa
name: CommonsenseQA
config: default
split: validation
metrics:
- type: pll_acc
name: pll_acc (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.21457821457821458
args:
dtype: bfloat16
num_few_shot: 0
max_length: 2048
standard_error: 0.011753423094216953
evaluation_date: '2026-10-01'
ci95_low: 0.19246524693116476
ci95_high: 0.2384815135817527
ci_method: Wilson 95%; item independence approximation
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: allenai/sciq
name: SciQ (with support)
config: default
split: test
metrics:
- type: pll_acc_norm
name: pll_acc_norm (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.688
args:
dtype: bfloat16
num_few_shot: 0
max_length: 2048
standard_error: 0.014658474370509057
evaluation_date: '2026-10-01'
ci95_low: 0.6586108249018996
ci95_high: 0.7159503139075316
ci_method: Wilson 95%; item independence approximation
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: truthfulqa/truthful_qa
name: TruthfulQA MC2
config: multiple_choice
split: validation
metrics:
- type: pll_acc
name: pll_acc (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.45836334441298165
args:
dtype: bfloat16
num_few_shot: 0
max_length: 2048
standard_error: 0.0158524239677764
evaluation_date: '2026-10-01'
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: BananaMind/BananaMind-Base-Bench-1.1
name: BananaMind Base 1.1
config: default
split: test
metrics:
- type: pll_raw_accuracy
name: pll_raw_accuracy (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.39714285714285713
args:
dtype: bfloat16
num_few_shot: 0
max_length: 2048
standard_error: 0.02619195222772307
evaluation_date: '2026-10-01'
ci95_low: 0.34726441786546686
ci95_high: 0.4492546213676228
ci_method: Wilson 95%; item independence approximation
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: cais/mmlu
name: MMLU continuation
config: 57 subjects
split: test
metrics:
- type: pll_acc
name: pll_acc (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.24483691781797465
args:
dtype: bfloat16
num_few_shot: 0
max_length: 2048
standard_error: 0.0036208715458871405
evaluation_date: '2026-10-01'
- task:
type: fill-mask
name: Experimental continuation PLL
dataset:
type: nyu-mll/blimp
name: BLiMP
config: 67 minimal-pair subsets
split: train
metrics:
- type: pll_acc
name: pll_acc (fraction; lm-eval 0.4.12 comparison protocol)
value: 0.6420746268656716
args:
dtype: bfloat16
num_few_shot: 0
max_length: 2048
standard_error: 0.0016890890018278603
evaluation_date: '2026-10-01'
---
# diffusion-51M-base
A small English **masked diffusion base model**, trained from random initialization as part of [Tiny llm ablation](https://huggingface.co/collections/d0rj/tiny-llm-ablation-6aafca336122dd2c2868f923). This is the successfully trained **v2** checkpoint: **51,392,512 stored parameters**, **50,867,200 optimized parameters**, exactly **3,932,160,000 processed source tokens**.
## Architecture and references
10 bidirectional transformer layers, width 512, 8 attention heads (head dimension 64), SwiGLU intermediate size 1536, RoPE, RMSNorm, tied input/output embeddings, context 2048. The unchanged 32,768-entry tokenizer and small-model comparison reference are [Q-50M-Base](https://huggingface.co/q-project/Q-50M-Base). These weights are not fine-tuned from that model.
The absorbing-mask objective follows the [LLaDA](https://arxiv.org/abs/2502.09992) family and its [official guidelines](https://github.com/ML-GSAI/LLaDA/blob/main/GUIDELINES.md): sample t uniformly, mask tokens independently with probability t, predict original tokens at masked positions, and minimize `sum(masked CE / t) / source_token_count`. Local t is clamped at 1e-5. This is an independent small-scale implementation, not an exact LLaDA reproduction.
**Time conditioning is disabled.** Removing the legacy additive time branch restored context learning after an optimization collapse. Its 525,312 unused, frozen parameters remain in the checkpoint for state-dictionary compatibility. The mask uses a separate learned input vector with ID 32768, outside the output vocabulary; it is not an added tokenizer token. Generic config validation may warn about this intentional custom mask ID.
## Training
[FineWeb-Edu sample-10BT](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu), local parquet shards, seed 2026, shuffle buffer 100,000. **15,000 steps × 8 sequences × 16 accumulation × 2048 tokens = 3,932,160,000 source-token exposures**. This is not a count of unique or masked target tokens. The v2 model started fresh; short debugging runs are not included in these weights or this token count.
One RTX 5070 Ti 16 GB, BF16 compute / FP32 weights, compiled forward, fused AdamW: peak LR 0.001, 150-step warmup then cosine decay to 0.0001, betas (0.9, 0.95), weight decay 0.1 excluding 1D parameters, clipping 1.0. Training took **439.5 minutes**; final logged denoising loss **5.1272**. [Exact training settings](training_config.json). Equal source-token budgets across architectures do not imply equal supervision or FLOPs.
## Evaluation
Full splits, no added few-shot examples, lm-eval 0.4.12, GPU BF16, context 2048 (ArithMark: 1024), no chat template. **Experimental continuation pseudo-log-likelihood (PLL):** mask one answer token at a time while all other answer tokens remain visible. `acc_norm` is the harness length-normalized option score. LAMBADA is full-token reconstruction accuracy, **not autoregressive final-word generation**. Scores are percentages ± one standard error; brackets are **95% Wilson intervals**.
| Dataset | Split | Examples | Metric | Score ± SE (%) | 95% CI (%) |
|---|---|---:|---|---:|---:|
| [HellaSwag](https://huggingface.co/datasets/Rowan/hellaswag) | validation | 10,042 | acc_norm | 27.95 ± 0.45 | [27.08, 28.84] |
| [ARC-Easy](https://huggingface.co/datasets/allenai/ai2_arc) | test | 2,376 | acc_norm | 33.63 ± 0.97 | [31.76, 35.55] |
| [ARC-Challenge](https://huggingface.co/datasets/allenai/ai2_arc) | test | 1,172 | acc_norm | 22.27 ± 1.22 | [19.98, 24.74] |
| [PIQA](https://huggingface.co/datasets/baber/piqa) | validation | 1,838 | acc_norm | 53.86 ± 1.16 | [51.58, 56.13] |
| [WinoGrande](https://huggingface.co/datasets/allenai/winogrande) | validation | 1,267 | acc | 50.12 ± 1.41 | [47.37, 52.87] |
| [OpenBookQA](https://huggingface.co/datasets/allenai/openbookqa) | test | 500 | acc_norm | 26.00 ± 1.96 | [22.35, 30.02] |
| [BoolQ](https://huggingface.co/datasets/aps/super_glue) | validation | 3,270 | acc | 52.97 ± 0.87 | [51.25, 54.67] |
| [LAMBADA OpenAI reconstruction](https://huggingface.co/datasets/EleutherAI/lambada_openai) | test | 5,153 | acc | 42.21 ± 0.69 | [40.87, 43.56] |
| [ArithMark-3](https://huggingface.co/datasets/AxiomicLabs/Arithmark-3.0) | train | 1,000 | `acc_norm` | 33.20 ± 1.49 | [30.35, 36.18] |
| [Balanced COPA](https://huggingface.co/datasets/pkavumba/balanced-copa) | train | 1,000 | `acc` | 52.20 ± 1.58 | [49.10, 55.28] |
| [CommonsenseQA](https://huggingface.co/datasets/tau/commonsense_qa) | validation | 1,221 | `acc` | 21.46 ± 1.18 | [19.25, 23.85] |
| [SciQ (with support)](https://huggingface.co/datasets/allenai/sciq) | test | 1,000 | `acc_norm` | 68.80 ± 1.47 | [65.86, 71.60] |
| [TruthfulQA MC2](https://huggingface.co/datasets/truthfulqa/truthful_qa) | validation | 817 | `acc` | 45.84 ± 1.59 | — |
| [BananaMind Base 1.1](https://huggingface.co/datasets/BananaMind/BananaMind-Base-Bench-1.1) | test | 350 | `raw_accuracy` | 39.71 ± 2.62 | [34.73, 44.93] |
| [MMLU continuation](https://huggingface.co/datasets/cais/mmlu) | test | 14,042 | `acc` | 24.48 ± 0.36 | — |
| [BLiMP](https://huggingface.co/datasets/nyu-mll/blimp) | train | 67,000 | `acc` | 64.21 ± 0.17 | — |
[WikiText-2 raw test](https://huggingface.co/datasets/Salesforce/wikitext): 291 nonoverlapping 1024-token blocks, prefix 512 + scored suffix 512, 148,992 scored tokens, 335 tail tokens omitted. Single-mask continuation PLL, GPU BF16, TF32 disabled. **Pseudo-perplexity 15.865 [15.198, 16.559]**; NLL **2.764107 [2.721156, 2.806927]**. 95% percentile block bootstrap, 10,000 resamples, seed 2026; exponentiate NLL endpoints for pseudo-perplexity CI. **This is not AR PPL**, and is not directly comparable with Q50M/T5/PrefixLM/Looped continuation PPL.
[Detailed metrics, standard errors and provenance](evaluation/results.json). Model-index entries use explicit `pll_*` metric labels and are author-reported results; no official leaderboard submission is claimed. Intervals do not include training-seed variation or all within-document dependence. Earlier test diagnostics informed the collapse investigation; benchmark contamination was not audited.
Full selected splits; lm-eval 0.4.12; seed 1234; BF16 on RTX 5070 Ti;
context cap 2048 (ArithMark 1024), TF32 disabled, no chat template and no added
few-shot examples. TruthfulQA retains the harness's fixed six-QA preamble.
ArithMark and BananaMind normalize by continuation token count; ordinary harness
acc_norm uses its own length normalization. BananaMind is raw accuracy, not Elo.
SciQ includes the support passage. Balanced COPA uses the mirrored 1000-item
train-named evaluation split; cRia's split was inferred, not confirmed.
MMLU scores full answer continuations across 57 subjects, weighted by item count;
BLiMP averages 67 equal-sized minimal-pair subsets. Standard errors are retained
from each evaluator. Wilson intervals are reported only where the runner logged
binary item accuracy; MC2 is probability mass, not binary accuracy. These
intervals do not model dependence between paired/templated examples or training
seed variation. UL2, PrefixLM and experimental diffusion PLL use their documented
conditional scoring protocols; PLL exposes the other answer tokens and is not
autoregressive likelihood. cRia's published scores used a different precision
and benchmark-adapted checkpoint; this completes our comparison coverage, not
an independent reproduction of cRia or an official leaderboard submission.
[Full results, provenance and group scores](evaluation/comparison-20261001/results.json). [Updated machine-readable results](evaluation/results.json). [TensorBoard events](tensorboard/) contain these new scores at step 15,000.
## Usage
Install [requirements.txt](requirements.txt). Load with **AutoModelForMaskedLM** and `trust_remote_code=True`; the custom API requires a `timesteps` tensor even though v2 does not condition on it.
```python
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
repo = "d0rj/diffusion-51M-base"
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(repo, trust_remote_code=True).eval()
ids = tokenizer("The purpose of science is to understand the world.",
return_tensors="pt", add_special_tokens=False)["input_ids"]
ids[:, 2] = model.config.mask_token_id
with torch.no_grad():
logits = model(ids, timesteps=torch.tensor([0.5])).logits
print(tokenizer.decode([logits[0, 2].argmax().item()]))
# Optional unconditional confidence-based iterative denoising.
tokens = model.generate_masked(batch_size=1, seq_len=32, steps=16, temperature=0.)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))
```
`generate_masked` starts from all masks and commits confident tokens on a linear schedule. It is an experimental unconditional sampler, not standard causal `.generate()` or an instruction-following interface. The reported core benchmarks use the PLL adapter, not this sampler. The generic HF inference widget is not configured for this custom interface.
Reproduce from a downloaded repository after installing `evaluation/requirements.txt`:
```bash
python evaluation/run_core.py --device cuda:0 --dtype bfloat16 --batch-size 1 --output core-results
python evaluation/run_continuation.py --output continuation-results.json
```
`--limit 2` on the core runner is a smoke test only. [TensorBoard files](tensorboard/) retain 750 training-loss points, development probes every 100 steps, and test evaluation scalars at step 15,000, including CI bounds. The failed v1 weights are not part of this release. Small English research ablation; not instruction-tuned, and performance improvements are not uniform across tasks.
To reproduce after downloading this model repository, accept the BananaMind dataset terms, authenticate with `hf auth login`, then run in a suitable CUDA environment:
```bash
pip install -r evaluation/comparison-20261001/repro/requirements.txt
python evaluation/comparison-20261001/repro/run.py --device cuda:0 --dtype bfloat16 --batch-size 1 --output comparison-rerun
```
The bundled runner uses the published model classes with the exact evaluation adapters and tokenizer. `--limit` produces smoke results only. Raw dataset examples are not included in this release.