README.md

---
license: apache-2.0
language:
- en
datasets:
- openbmb/Ultra-FineWeb-L3
- wikimedia/wikipedia
- HuggingFaceTB/finemath
- HuggingFaceTB/smollm-corpus
- OpenCoder-LLM/opc-annealing-corpus
tags:
- diffusion-language-model
- masked-diffusion
- parallel-decoding
- diloco
- text-generation
pipeline_tag: text-generation
---

<p align="center"><img src="assets/avatar.svg" width="200" alt="DiffuRefill-1B"></p>

# 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

![training progress](assets/training_progress.png)

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.compile`d.

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

| source | share |
|---|---|
| [Ultra-FineWeb-L3](https://huggingface.co/datasets/openbmb/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`.