---
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.