KB-Diffusion Model B — word-level masked diffusion
Masked diffusion language models trained on English words. The "generalization companion" experiment from KB-Diffusion (an educational masked-diffusion project by Bijan Bowen / OminousIndustries): swap the repo's four keyboard layouts for thousands of words and see if the same recipe still works. It does — and iterating on decoding strategy turned out to matter as much as architecture, with sequence length flipping which sampler wins.
Three checkpoints, same LLaDA-style recipe (t ~ U(0.05, 1) masking, 1/t-weighted CE, bidirectional transformer, no causal mask):
| v2 (N=5) | v3 (N=5) | N=10 | |
|---|---|---|---|
| Params | 4.75M (6 layers) | 6.33M (8 layers) | 6.34M (8 layers) |
| Steps | 8,000 | 12,000 + cosine LR | 12,000 + cosine LR |
| Best valid English | 95.5% (T=0.5) | 98.4% (T=0.5) | 74.2% (revision + T=0.5) |
| Unique words / 512 | 409 | 428 | 200 |
| Unigram TV vs exact Bayes | 0.0374 | 0.0135 | 0.019 |
The headline findings
1. Temperature is the free win (v2, frozen weights): ancestral sampling at T=1.0 gives 68.8% valid English; T=0.5 gives 95.5%. Same weights, same 5 forward passes, +27 points.
2. Decoding strategy > extra parameters: v2 read well (95.5%) beats v3 read poorly (82.6% at T=1.0).
3. Sequence length flips the sampler winner. At N=5, revision-capable sampling (un-commit weak letters, re-mask, retry) loses to plain low-temperature sampling (77.1% vs 95.5%). At N=10 it wins (74.2% vs 63.3%) — early mistakes poison enough downstream positions that un-committing them pays for its 4x compute. The "diffusion can revise" capability has a measured regime where it wins.
4. The parallel/iterative gap explodes with length: one-shot sampling falls 2.0% (N=5) → 0.0% of 512 samples (N=10). This is why real diffusion LMs commit few tokens at a time.
Full methodology, negative results, and per-sampler tables: docs/model-b.md
in the GitHub repo.
Usage
import torch
from model_b_word_diffusion_v3 import Net, CH, MASK, N # from the GitHub repo
model = Net(layers=8) # v3; use model_b_word_diffusion_n10.py for N=10
sd = torch.load("modelb_v3.pt", map_location="cpu", weights_only=True)
model.load_state_dict(sd)
model.eval()
# ancestral confidence-commit sampler, temperature 0.5 — see repo scripts
Intended use & limitations
Educational artifact, not a production model: 27-token vocab, 5/10-position sequences. It exists to make the masked-diffusion mechanism (parallel prediction, confidence commits, re-masking, revision, posterior sharpening) measurable — and to map how decoding strategy and sequence length interact on frozen weights.