diffusion-51M-base

A small English masked diffusion base model, trained from random initialization as part of Tiny llm ablation. 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. These weights are not fine-tuned from that model.

The absorbing-mask objective follows the LLaDA family and its official guidelines: 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, 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. Equal source-token budgets across architectures do not imply equal supervision or FLOPs.

Evaluation

Full splits, zero-shot, lm-eval 0.4.12, GPU BF16, context 2048, 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; brackets are 95% Wilson intervals.

Dataset	Split	Examples	Metric	Score (%)	95% CI (%)
HellaSwag	validation	10,042	acc_norm	27.95	[27.08, 28.84]
ARC-Easy	test	2,376	acc_norm	33.63	[31.76, 35.55]
ARC-Challenge	test	1,172	acc_norm	22.27	[19.98, 24.74]
PIQA	validation	1,838	acc_norm	53.86	[51.58, 56.13]
WinoGrande	validation	1,267	acc	50.12	[47.37, 52.87]
OpenBookQA	test	500	acc_norm	26.00	[22.35, 30.02]
BoolQ	validation	3,270	acc	52.97	[51.25, 54.67]
LAMBADA OpenAI reconstruction	test	5,153	acc	42.21	[40.87, 43.56]

WikiText-2 raw test: 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. 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.

Usage

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

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:

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 retain 750 training-loss points, development probes every 100 steps, and 82 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.