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.

Training progress

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

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):

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:

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.

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

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

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