| --- |
| license: apache-2.0 |
| tags: |
| - masked-diffusion |
| - discrete-diffusion |
| - llada |
| - educational |
| - experiment |
| language: |
| - en |
| library_name: pytorch |
| --- |
| |
| # KB-Diffusion Model B — word-level masked diffusion |
|
|
| Masked diffusion language models trained on English words. The |
| "generalization companion" experiment from |
| [KB-Diffusion](https://github.com/PastelRuntime/KB-Diffusion-Optimized) |
| (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 |
|
|
| ```python |
| 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. |
|
|
