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.
| 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 |
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.
|
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 |
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
|
|
| parameters |
~1.0B |
| dim / layers / heads |
2048 / 18 / 16 |
| FFN |
SwiGLU, 5632 |
| context |
2048 |
| vocab |
32768 BPE |
| objective |
masked diffusion, mask rate t ~ U(0.15, 1), loss on masked positions |
| corpus |
Ultra-FineWeb-L3 |
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.
