DiffuRefill-1B

Status: pretraining, about half way. Step ~100k of 200k, ~46B effective training tokens, merge round 290. The weights here are a live pretraining checkpoint: fluent English, facts still unreliable, not instruction-tuned. Use global_ema.pt; it is the best checkpoint at every round.

A 1.08B masked-diffusion language model decoded with confidence-targeted re-infilling: draft the whole sequence in a handful of parallel denoising passes, then spend a few more passes rewriting only the tokens the model is least sure about. An autoregressive model needs one pass per token; this one needs ~20 passes for a whole sequence, whatever its length.

parameters	1.08B (embedding tied to the output head)
architecture	18 blocks, dim 2048, 16 heads, SwiGLU 5632, RMSNorm, RoPE, bidirectional attention
context	2048 tokens
tokenizer	openbmb/MiniCPM4-0.5B (73,440 tokens; vocab padded to 73,760 for MASK and PAD)
objective	absorbing masked diffusion, loss on masked positions only
training	DiLoCo across 3-5 single-GPU boxes, synchronised through this repository
progress	step ~100k / 200k, ~46B effective tokens

Training progress

Tracked every round on fixed, deterministic batches, so the change between two rounds carries no sampling noise:

measure	what it is	latest
held-out CE (EMA)	masked CE at t = 0.3 / 0.6 / 0.9 on a fixed Ultra-FineWeb slice no worker trains on	2.939 (round 287)
held-out CE (global)	the same for the latest merged global	2.957
fact probe, mc	1,050 PopQA facts written as plain sentences with the answer masked; the gold answer must beat three same-length distractors (chance 0.25)	0.560 (was 0.540 before the data mix)
fact probe, answer NLL	per-token NLL of the masked gold answer	6.08 (was 6.55)
fact probe, exact match (EMA)	the answer decoded token by token and compared with PopQA's accepted answers	0.083 (was 0.075)
decoding accuracy	confidence decoding of 50% masked held-out text in 8 passes: token / adjacent-pair accuracy	0.653 / 0.453 (round 275)

The guard-probe curve at the top is measured on an Ultra-FineWeb slice. When the fact-dense sources were added at step ~96.8k, Ultra-FineWeb fell from 100% to 56% of the documents, and that probe rose by design while the fact probe improved quickly. From here the fact probe and the held-out CE are the numbers to watch. No public benchmark scores are reported yet: at this stage they would sit near chance and say little.

How it is trained

Fleet. Ephemeral single-GPU boxes (RTX PRO 6000 Blackwell, 96 GB) with no interconnect, joined DiLoCo-style through this repository. Each worker trains 150 local AdamW steps, pushes ckpt_wNN.pt, and adopts the merged global_b.pt.

• Merging. The merger averages the round's checkpoints (a plain mean, folded into a running fp32 sum as they arrive) and waits at most 150 s past the first arrival.
• Phase. Workers pick up a new global as soon as it appears and take its step counter, so the fleet starts each interval together.
• EMA. An exponential average of the globals (β 0.9) is published as global_ema.pt.

Objective.

• t ~ U(0.15, 1) per row, with positions masked independently.
• Half of the rows keep a random-length clean prefix, which is the shape of prompted generation.
• Attention stays inside each packed document.
• On a quarter of the rows, names and numbers are masked as whole spans (salient span masking). Masked token by token, a two-token name leaks through its visible half; masked whole, it has to be recalled.

Optimiser. AdamW (0.9, 0.95), weight decay 0.1, clip 1.0, peak LR 1.5e-4, cosine to 10% over 200k steps. Each worker step covers 64 × 2048 = 131,072 tokens. The blocks are torch.compiled.

Data (by document, from step ~96.8k):

source	share
Ultra-FineWeb-L3 Multi-Style + QA synthetic	56%
Wikipedia (en)	15%
FineMath 4+	8%
Cosmopedia v2	8%
OpenCoder annealing: algorithmic corpus + code snippets	8%
Wikipedia leads restated as textbook / lecture / museum-guide prose (own corpus, filtered for copying and for numbers not in the source)	5%

Before step ~96.8k the model saw Ultra-FineWeb only.

What the run taught

Measured on this model or on its 37.9M-parameter stand, with fixed probes:

change	result
peak LR 3e-4 → 1.5e-4	Local drift over one interval fell from 0.076 to 0.013 CE. The merged model improved from 4.336 to 4.288.
keeping local drift across a merge (θ += global − pushed)	Worse in 60 of 63 adoptions, +0.079 CE on average. The run uses plain replacement.
EMA of the globals	Beats the latest global every round. Seeding training from it gave the gain back.
outer Nesterov (μ 0.5)	Raised held-out CE for three rounds. Reverted.
Muon switched on mid-run	Loss rose from 3.16 to 3.32. It was a 2× win from scratch on the stand, but not as a switch on a running model.
span masking, PMI units, stratified t, complementary masks	Each lost to plain independent masking on the stand. Span masking was removed at step ~91.9k.
fact-dense data mix + salient span masking	Fact-probe answer NLL 6.55 → 6.08 within about four hours.
torch.compile per block	17% faster steps, less memory.
FP8 matmuls (torchao, rowwise or tensorwise)	No speed-up on this card. Not used.

The method

1. K denoising steps           -> a full draft + a confidence value per position
2. repeat R times:
     take the p% least confident generated positions
     re-mask them
     one forward pass          -> new tokens, chosen with both sides in view
     refresh their confidence

Cost is K + R forward passes. One refill pass fixes any number of positions at once, because the model processes the whole sequence in parallel. That is what makes targeted repair cheaper than more denoising.

Decoding configuration

The draft and the refill deliberately use different sampling rules. The draft decides what the text is about and needs freedom. The refill touches a fifth of the tokens and needs discipline. Sharing one setting between them trades every metric straight against another.

stage	settings
draft	K=12, temperature 2.2 -> 0.9 annealed, Gumbel noise 2.5 on the unmask order, no truncation
refill	R=8, p=20%, temperature 0.8, min_p=0.10
both	neighbour ban: subtract 6.0 from the logit of the token already standing to the left or right; never commit two positions closer than 2 in one pass

• Gumbel unmask order. Strict top-k confidence reveals the easy positions first, such as articles and punctuation. Content words are left for last, when the context is already frozen. Annealed Gumbel noise on the log confidences breaks that order. It was the largest single effect in the pilot.
• Neighbour ban and commit gap. Two adjacent positions sampled independently both pick the locally likely token. That is how parallel decoding produces repeats, and these two rules target exactly that.
• Split-stage sampling. A hot, free draft followed by a cool, truncated refill.

Pilot results (TinyStories, 37.9M parameters)

Two models of identical architecture were trained from scratch on the same corpus: one autoregressive, one masked-diffusion. An independent judge (roneneldan/TinyStories-33M, related to neither) scored the NLL of each model's generated continuation.

	calls	NLL ↓	doubled tokens ↓	distinct-2 ↑
real text	—	1.15	1.38%	0.478
autoregressive	128	3.08	1.64%	0.466
diffusion + refill, tuned	20	2.87	0.30%	0.459

The diffusion model matched autoregression in quality at a sixth of the passes. It produced five times fewer doubled tokens than the autoregressive baseline, below the rate in the real corpus. Diversity came out at parity, slightly under autoregression.

What did not work in the pilot:

• Autoregressive repair of the draft. It was worse than not repairing at all. An AR model scoring position i cannot see i+1, so it duplicates the neighbour it is not looking at.
• Self-repair training on the model's own drafts. It was harmful at 60–100% corruption and neutral at 10–35%.

Honest limits

• Unfinished. Half of the planned steps remain, and facts are unreliable. Expect confident, fluent mistakes.
• The pilot evidence is small. 37.9M parameters on TinyStories does not show that the decoding results transfer to 1B on natural text.
• The judge is an AR model. A low NLL under it rewards predictable text as well as good text.
• Diffusion costs more to train. It learns only from masked positions, so it reads more tokens than an AR model for the same quality. The saving is entirely at inference.

Files

file	contents
global_ema.pt	EMA of the merged globals; the checkpoint to use
global_b.pt	the latest merged global
ckpt_wNN.pt, log_wNN.txt	per-worker checkpoints and logs, overwritten every interval
config.json	the exact architecture, objective, data mix and training setup
pilots/	held-out and fact-probe logs, stand pilot results

Checkpoints are torch.save dicts with a model state dict for the Net class described in config.json.