KB-Diffusion Model B — word-level masked diffusion
A 4.75M-parameter masked diffusion language model trained on 8,000 five-letter English words. This is the "generalization companion" experiment suggested in 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, with a much harder hypothesis space. Trained in ~11 minutes on a free Kaggle T4.
The recipe (identical in shape to the keyboard model)
- Bidirectional transformer, no causal mask: vocab 27 (a-z + [MASK]), seq len 5, dim 256, 6 layers, 4 heads, ff 1024
- LLaDA-style training: masking ratio t ~ U(0.05, 1), cross-entropy on masked positions only, weighted by 1/t
- AdamW, lr 3e-4, 8000 steps, batch 1024
Results
256 samples per sampler, checked against the full ~16k list of English 5-letter words:
| Sampler | Valid English | Unique |
|---|---|---|
| One-shot parallel (all 5 letters at once) | 2.0% | 256/256 |
| k=2 commits per step | 39.5% | 256/256 |
| Ancestral (commit 1, re-condition, repeat) | 68.4% | 254/256 |
| Greedy ancestral (no sampling) | 100% | 1/64 ("bales") |
The model's from-scratch letter predictions also track exact analytic unigram statistics to mean total variation 0.0374 — the "transformer learns Bayes' rule from corrupted examples" effect, at 8000 classes instead of 4.
Two takeaways:
- The sampler is half the model. Identical weights produce 2% or 68% valid output depending only on how commitments are made. Coherence is bought with iteration, not parameters.
- 68.4% of outputs were valid English but only 66.4% came from the training list — it generalizes past its vocabulary, producing words it never saw.
Usage
import torch
from model_b_word_diffusion import Net, CH, MASK, N # from the GitHub repo
model = Net()
sd = torch.load("modelb_v2.pt", map_location="cpu", weights_only=True)
model.load_state_dict(sd)
model.eval()
# then use the ancestral confidence-commit sampler from the repo's script
See the GitHub repo for the full training/eval script and writeup.
Intended use & limitations
An educational experiment, not a production model. 27-token vocabulary, 5-position sequences, memorization-adjacent scale on purpose. It exists to make the masked-diffusion mechanism (parallel prediction, confidence commits, re-masking, posterior sharpening) measurable and visible.