THE FUSION KIT · three language models, one fair fight, a reliability gate

The fair fight

Three bodies, three comparable numbers

The whole point is comparability: full vocabulary, the canonical corpus and splits, no UNK inflation — so each body’s perplexity can be read against published work instead of an easier private task. The numbers below are field references the kit is built to reach, not measurements taken on this page — the measurements themselves, taken on a laptop, are in the next section.

counts

Modified Kneser-Ney 5-gram

the count baseline · ~48 target

The count body — interpolated modified KN with absolute discounts D=n1/(n1+2·n2) and continuation counts (Chen & Goodman), vectorised in numpy. Full-vocab canonical, so it produces a real, comparable n-gram number (the README aims at ~48). kn_baseline.py

neural

Transformer, from scratch

≈ 20–29 test ppl (by size)

A plain decoder-only GPT (pre-LN blocks, weight-tied head, AdamW + cosine, mixed precision, resumable). Presets 30M / 120M params. Where it lands is where its size puts it on the scaling curve. model.py · train.py

semi-parametric

kNN-LM datastore

−2 to −4 ppl off the net

The transformer’s own hidden states (the input to the final softmax) stored as keys → next-token values; at test, retrieve the K nearest and read their futures. FAISS index or brute-force fallback. build_datastore.py

The piece that won the arc

The shadow gate, made visible

Plain kNN-LM trusts the datastore by a fixed weight. This kit makes that weight adaptive: it scales with how close the nearest neighbour actually is. Retrieval that found a near-identical context is worth leaning on; retrieval that only found something vaguely similar is not. Slide the neighbour distance and watch the gate — and the surprise on the true token — respond.

◧ fusion engine · an illustration of the mechanism, not a trained model — the three distributions are illustrative; the gate math is fusion.py exactly

P_transformer · the net’s guess

P_KN · the counts’ guess

P_kNN · the datastore’s guess

the shadow gate

nearest-neighbour distance d_nn0.60

retrieval reliability e^−d/τ—

w_k = clip( b + s·rel , 0, 0.6 ), w_n = 1 − w_c − w_k
b=0.05 · s=0.55 · w_c=0.25 · τ=1.0 (real kit grid-searches these on val)

▸ P_final · the fused distribution — predicting the word after “…the cat sat on the ___”

Surprise = −log₂ P(true) in bits; lower is better (perplexity = 2^{mean surprise}). Drag d_nn toward 0 to simulate finding a near-identical training context — the gate opens, the datastore’s vote counts, and the true token’s surprise drops. Drag it high and the gate closes to its floor: the fusion falls back to net + counts.

The kit

Five steps, one config

Every script reads config.py (model size, vocab, knobs). Needs a CUDA GPU for the transformer half; the count body is CPU (RAM-heavy).

python prepare_data.py

download WikiText-103, build vocab, tokenize → int32 id arrays

~5 min · CPU

python kn_baseline.py

modified Kneser-Ney 5-gram — the count body (~48)

~10 min · CPU, RAM-heavy

python train.py

train the transformer from scratch (checkpointed, resumable)

hours · GPU

python build_datastore.py

encode train into the kNN datastore (+ FAISS index)

~30 min · GPU

python fusion.py

counts + neural + kNN + shadow gate on test → perplexity

~15 min · GPU

What to expect

The numbers it is built to reach

Word-level test perplexity on WikiText-103. The left column is what this kit’s default config is designed to produce; the right is the published field reference it is meant to stand beside. Nothing here was measured on this page — running the kit takes a GPU and hours.

body / method	this kit (default)	field reference
interpolated Kneser-Ney 5-gram	~45–55 (README target)	varies by setup †
transformer (d512, 8L, ~30M)	~28–35	GPT-2 small ≈ 29*
transformer (d768, 12L, ~120M)	~20–24	GPT-2 large ≈ 18–20*
+ kNN-LM fusion	−2 to −4	kNN-LM ~18.65 → ~16.12
+ shadow gate (adaptive w_k)	measured leak-free: ≈0 (+0.009)	(this kit’s addition)

* GPT-2 uses BPE/subword tokenization, so its perplexity is not strictly comparable to a word-level number — treat the GPT-2 column as an approximate anchor, not an identity.
† Published KN 5-gram figures on WikiText-103 vary with vocabulary and pruning (the ~48 the README targets is one commonly-quoted point; n-gram baselines on the full ~267k vocab often run higher). The value is honest because it is comparable — same algorithm, same setup — not because it hits an exact constant. kNN-LM figures: Khandelwal et al. 2020 report 18.65 → 16.12 (kNN-LM), 15.79 with a continuous cache.

Actually run

Measured on a laptop

Not estimates — the kit was run end-to-end on an RTX 4050 Laptop (6 GB), through the full three-body fusion. Three things surfaced that a teardown alone never could.

The transformer learned

The small preset — 76.5M params by the fingerprint (untied 50k embeddings; the "~30M / 51M" earlier notes both undercounted — see below) — trained from scratch on the GPU: ~2.8 GB VRAM, AMP, micro-batch 8, ~94% util. Validation perplexity, from random init (≈50,000):

the first 3,000 steps shown; training ran on and converged at ~58k steps — val ppl ≈ 51, test 52.45. Short of the ~30 hoped for at this size, but real — and it is the checkpoint the fusion below uses.

The count body was broken — now fixed

At its own default (50k vocab, order 5) the KN 5-gram scored 2,138 ppl — worse than a bigram. The gram keys overflowed int64 and silently hashed while the query stayed exact, so every 4/5-gram lookup missed.

Hashing both sides consistently fixes it — perplexity now falls with order as a real KN must:

order 2→5: 408 · 307 · 286 · 281

Short of ~48 because this is a RAM-capped 12M-token subset (the full 83M run lands lower, and ~48 likely also needs the full 3-discount modified KN). The ~48 was a target — now labelled as one.

The three-body fusion — measured, then re-measured honestly

The first run (v1) reported the fusion at 25.02, credited the datastore with halving the transformer's perplexity, and the shadow gate with a −1.21 win. Two of those three claims did not survive an honest re-run. The v1 mixing weights were grid-searched on the same 200k test slice they were reported on — a leak — so v2 rebuilt the evaluation to tune on a validation slice and report on a held-out one, printing a diagnostic of exactly what the adaptive gate buys. Same checkpoint, same 1M-key datastore (2.5% of ~40M; recall 36.9%), same caps.

body / method	v1 (test-tuned)	v2 (leak-free, held-out)
transformer (76.5M params — see below)	52.45	53.37
Kneser-Ney counts (12M-token cap)	282.61	285.26
best static blend (val-tuned: 0.7 net · 0.3 counts · ~0 kNN)	26.23*	25.25
+ reliability gate	25.02	25.24
what the adaptive gate buys over static	−1.21 (claimed)	+0.009

* v1's static was a hand-set 0.6/0.3/0.1 blend, not a searched one — which is why the gate looked like it won by −1.21. v2 searches the static blend too, and the fight becomes fair.

Correction 1 — the gate's −1.21 was a tuning leak. It compared a hand-set static blend against a gate tuned on the slice it was scored on. Tuned honestly on validation, the gate buys +0.009 ppl — the diagnostic's own verdict: "the gate's win is essentially the static term; rel barely contributes."

Correction 2 — the datastore is not doing the halving. At this 1M-key cap the honest search puts ~0 weight on the kNN channel; transformer(0.7)+KN(0.3) alone already lands 25.25. The 53→25 drop is the count-model mixture, not the datastore. (kNN-LM's published −2 to −4 needs the full ~103M-key store — not a claim 2.5% of it can make.)

What holds: the fusion record itself, ~25.2, now measured on held-out data instead of the slice it was tuned on. The two retracted claims were artifacts of an evaluation grading its own homework — and it was the kit's own new diagnostic that caught them. That is the point of v2: not a lower number, a truer one.

Two instruments the re-run added

The shadow channel (shadow.py) — a witness, not a mixer. It independently recomputes the KN and kNN bodies from their raw inputs and refuses to fold a channel it can't verify, instead of blending poison through. Proven offline (verify.py, no GPU): on a NaN-poisoned channel the old gate returns weights None while the shadow flags it and still fuses; on a silently misaligned KN vector (finite but wrong) the gate passes it through and the shadow flags 85.6% of positions (at the shipped tolerance 0.15; tighter tol catches more). The gate detects neither fault; the witness detects both.

The checkpoint fingerprint (model_id.py) — model identity as a fact on disk (sha256 of arch-signature, of tensor bytes, of training state), deterministic across runs. Its first act was to correct this page: the checkpoint carries 76,552,192 params, not the "51M" v1 asserted (untied embeddings). Identity is now a diff, not a memory.

v3 — the witnesses run on real data, and agree on the weak link

v2 proved the shadow offline; v3 wired it (plus a new uncertainty gate) into the live fusion and ran all three on the real checkpoint. Two independent instruments — one built to save compute, one built to catch corruption — reached the same verdict about the datastore without being told to.

The uncertainty gate (uncertainty_gate.py) retrieves only where the transformer is unsure, and prices the trade. The curve came out inverted — retrieving less is strictly better:

search fraction	skip	ppl (val slice)	vs full search
0% (retrieve nothing)	100%	24.24	−1.14
50%	50%	24.80	−0.58
100% (retrieve all)	0%	25.38	+0.00

A tool built to trade quality for speed measured that there is no quality to trade: the kNN body is net-negative at this cap, so skipping 50% of searches saves 0.23 ppl (the audit calls it "free"). It re-confirms v2's "the datastore isn't the halving" — as a clean monotonic curve, from an instrument that had every incentive to find retrieval useful.

The shadow's first real run then explains why to distrust the datastore. It independently recomputes both aux channels: KN(2) flagged 0.00% (deterministic — agrees to the bit), but kNN(3) flagged 21% — the GPU-fp16 primary and a CPU-fp32 recompute pick different top-16 neighbours on a fifth of positions (precision + tie-breaking, not a crash). Net-negative and precision-unstable: the datastore is the weak link on both quality and reliability, which is exactly why the honest gate weights it ~0.

Assembling and running v3 also caught two integration bugs a reader would miss: the kNN row-min underflow fix had been reverted (re-applied), and the wired-in uncertainty block referenced the KN vector before it was computed (a NameError that only fires when you actually run main()). "Verified on synthetic data" is not "run."

Every number here was computed on the laptop, not taken from a paper. The honest ordering holds — counts 285 > transformer 53 > fusion 25.2 — and each version made the record's shape truer: v1 measured it, v2 removed its leak, v3 named its weak link. None lowered it; all of it is the kit checking its own homework.

Render-not-invent

The actual code

This isn’t a description of a kit — it is the kit. Eight files, ~1,000 lines, browsable and clonable. the repository ↗

config.py

paths, device, vocab, model presets, all knobs

prepare_data.py

download WT-103, word-level vocab, tokenize

kn_baseline.py

interpolated modified Kneser-Ney, vectorised

model.py

decoder-only GPT; final LayerNorm = kNN key

train.py

AdamW, cosine, AMP, grad-accum, resumable

build_datastore.py

hidden→next-token keys/vals + FAISS index

fusion.py

the triple mix + gate; val-tuned, held-out report, diagnostic

shadow.py

the witness channel — independent recompute, refuses to fold poison

verify.py

offline proof the witness catches NaN + silent drift (no GPU)

uncertainty_gate.py

retrieve only where the LM is unsure — prices the speed/quality dial

model_id.py

checkpoint fingerprint — arch / weights / state as sha256

README.md

run order, expectations, the honesty note

The colophon

What is real here, and what is not

Real

The code runs — a complete, coherent pipeline (train/eval/checkpoint, KN recursion, datastore, FAISS), not pseudo-code.
The algorithms are canonical and cited: modified Kneser-Ney (Chen & Goodman); kNN-LM with the pre-softmax hidden state as the key (Khandelwal et al. 2020).
The comparability logic is sound: WT-103, full vocab, canonical splits is the setup that makes a number legitimately comparable. The kNN-LM (Khandelwal 2020) and Chinchilla figures are verified to source; the exact KN5-on-WT103 constant varies across the literature, so it is stated as the kit’s target, not a fixed fact.
The shadow gate — the adaptive, distance-scaled kNN weight — is David’s own addition on top of standard kNN-LM. Measured leak-free, it buys ≈0 over a properly-searched static blend (+0.009 ppl); it is kept as honest plumbing, not sold as a win. The separate shadow channel (shadow.py) is the real addition that earns its place — a witness that catches faults the gate can't see.

Not (and it says so)

Nothing is measured live on this page — the numbers in “Measured on a laptop” were computed by running the kit (RTX 4050, hours), then written in. The “What to expect” table remains targets; the measured section is results, with its caps stated (1M-key datastore, 12M-token counts, test-slice weight search).
The live engine is an illustration of the fusion equation and gate — its three distributions are hand-set to show the mechanism, not read from a trained model.
The earlier CPU-sandbox numbers from the arc (71, 82.7…) were never comparable — 5k-vocab with heavy UNK measured an easier task. This kit is the fix.
The domain-routing / cardinal-hue layer is intentionally absent — WT-103 is unlabelled, so there is nothing to route. That machinery belongs to the labelled 8-domain corpus; this is the neutral, comparable substrate it rides on.

REFERENCES · Merity, Xiong, Bradbury, Socher (2016) Pointer Sentinel Mixture Models — introduced WikiText-103 · Chen & Goodman (1999) An Empirical Study of Smoothing Techniques — modified Kneser-Ney · Kneser & Ney (1995) — the base smoothing · Khandelwal, Levy, Jurafsky, Zettlemoyer, Lewis (2020, ICLR) Generalization through Memorization: Nearest Neighbor Language Models · Hoffmann et al. (2022) Training Compute-Optimal LLMs (Chinchilla). Dataset: Salesforce/wikitext-103-raw-v1.

⧉ the code (GitHub) ⌂ ud0 · the atlas ⌗ the city ↑ top