Asilarkness / DiffuRefill-1B

DiffuRefill-1B

Status: training in progress. No weights are published yet — this card describes the recipe and the pilot results that motivate it.

A ~1B masked-diffusion language model decoded with confidence-targeted re-infilling : draft the whole sequence in a handful of parallel denoising steps, then spend a few extra passes rewriting only the tokens the model is least sure about.

The point is inference cost. An autoregressive model needs one sequential forward pass per token. This one needs ~20 passes for a whole sequence, regardless of its length.

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

Draft and refill are deliberately given 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.

Three of these were found by search and each does a specific job:

• Gumbel unmask order. Strict top-k confidence systematically reveals the easy positions first — articles, punctuation — leaving content words for last, when the context is already frozen. Adding annealed Gumbel noise to the log confidences breaks that ordering. Largest single effect in the pilot.
• Neighbour ban. One line. It targets exactly the mechanism by which parallel decoding produces repeats: two adjacent positions, sampled independently, both pick the locally likely token.
• Split-stage sampling. Hot free draft, cool truncated refill.

Pilot results (TinyStories, 37.9M parameters)

Two models of identical architecture trained from scratch on the same corpus — one autoregressive, one masked-diffusion — then compared by an independent judge ( roneneldan/TinyStories-33M , related to neither) scoring the NLL of the generated continuation.

Same quality as autoregression at a sixth of the passes, and five times fewer doubled tokens than the AR baseline — below the rate in the real corpus. Diversity lands at parity, slightly under AR.

What did not work

Reported because a negative result saves someone the same detour.

• Autoregressive repair of the draft. The original form of the idea: let an AR model rewrite the shaky tokens. A control sweep over a product-of-experts mixture showed the result degrading monotonically with the weight on the AR expert; pure AR repair was worse than not repairing at all . An AR model scoring position i cannot see position i+1 , so it writes words that duplicate the neighbour it is not looking at. It discards the one advantage the diffusion draft has.
• Span masking during training (contiguous runs instead of isolated tokens). No measurable effect against a control.
• Self-repair training (train on the model's own drafts). Its first form drafted from 60–100% masking and asked for ground-truth tokens amid incoherent context — a mapping fittable only by ignoring context. It made the model four times worse on duplication. Fixing the draft corruption to 10–35% removed the damage but produced no gain either.

In a four-way ablation (control / span / self-repair / both) plain continued training matched or beat every addition.

Planned model

Ultra-FineWeb-L3 is chosen over the larger tiers because a masked-diffusion objective draws its training signal only from the masked positions — on average 57% of them — so quality per token matters more here than raw volume.

Training runs on a fleet of isolated single-GPU boxes with no interconnect, synchronised DiLoCo-style through the Hub: workers train on disjoint shards and periodically average weights through a merged global checkpoint.

Honest limits

• The pilot is small. 37.9M parameters on TinyStories is a toy setting with simple language. Nothing here demonstrates that the results transfer to 1B or to natural text — that is what this run is for.
• The judge is an AR model. Low NLL under it rewards predictable text as well as good text. The diversity number is reported alongside for that reason.
• The tuned configuration is fitted. On the search split it beat the AR baseline on all three metrics; on a held-out split only two of three reproduced, diversity falling 0.007 short. The margins during search were 0.015 nats and 0.001 distinct-2, so some overfitting to 128 samples was inevitable. Quality and duplication are the robust wins.
• Diffusion costs more to train. It reads more tokens than an AR model for the same quality, because it only learns from masked positions. The saving is entirely on the inference side.