kNN-LM's adaptive variant decides how much to trust its retrieval memory by looking at
the language model's own uncertainty: the blend weight is λ = λ_max · σ(a·H + b),
where H is the next-token entropy. Lookaside implements that gate and fits a, b
to held-out negative log-likelihood.
Fitting it kept producing a gate that barely moved. This note works out why, and the answer turned out to be worth changing the shipped configuration over: entropy predicts which tokens the memory will win and almost nothing about how much those wins are worth, and a gate fitted to likelihood needs the second thing.
A 28M-parameter model trained on thirty books. Two further Doyle novels, The Valley of Fear and His Last Bow, were held out of training entirely. The memory holds 600k entries (83% coverage) built from the six Doyle books the model did train on plus the first half of each never-trained one. Two evaluation sets come out of that same memory:
| eval set | scored positions | in the weights? | in the memory? |
|---|---|---|---|
| never-trained books (second halves) | 89,541 | no | no |
| trained books (tail slices) | 33,483 | yes | no |
Nothing evaluated is in the datastore. That is guaranteed by how build_split carves
the corpus, and separately by a test pinning the self-retrieval invariant, rather than
by inspecting outputs. No-memory perplexity differs sharply between the two — 42.5 on
the never-trained report split against 21.9 on the trained one — which is consistent
with the split working, though the two sets are also different text and so differ in
difficulty for reasons beyond training exposure.
Gates are fitted on one book and reported on a different one. Splitting positions at random would put neighbouring tokens of the same passage on both sides, which tests very little; holding out a whole book is a real test.
The signals on offer at each step are the model's entropy H and its top-1
probability, plus, from the search: d1 (squared L2 distance to the nearest key),
d_mean (mean distance over the 64 neighbours), d_spread = d_mean − d1 (how much
better the best match is than a typical one), and the peak and entropy of the retrieval
distribution.
λ_max = 0.5 throughout. Perplexity on the complete held-out halves:
| the gate conditions on | The Valley of Fear | His Last Bow |
|---|---|---|
| nothing — no memory at all | 42.53 | 38.94 |
| a constant λ, fitted | 38.59 | 37.54 |
entropy H — kNN-LM's signal |
38.59 | 37.50 |
entropy + p_LM top-1 |
38.60 | 37.52 |
d1 alone |
38.45 | 37.52 |
d_spread alone |
37.94 | 37.38 |
d1 + d_spread |
37.86 | 37.34 |
| all five neighbour signals | 37.86 | 37.34 |
| all seven signals | 37.90 | 37.34 |
| per-token oracle (cheating) | 31.31 | 31.14 |
Start with what is not the gate. A single fitted λ applied to every token gets 9.26% of the 10.97% total on the first book — about 84% of everything the memory is worth here is one number. Any gating scheme is competing over what is left, so the comparison that matters is what each adds on top of that constant. Intervals are 95%, from 2,000 bootstrap resamples over contiguous 256-token blocks; a per-token bootstrap would badly understate them, since tokens within a passage are strongly dependent.
| improvement over a fitted constant λ | The Valley of Fear | His Last Bow |
|---|---|---|
entropy H |
+0.009% [−0.061, +0.080] |
+0.104% [+0.078, +0.130] |
entropy + p_LM top-1 |
−0.007% [−0.076, +0.065] |
+0.047% [+0.009, +0.086] |
d1 alone |
+0.362% [+0.282, +0.444] |
+0.055% [−0.035, +0.141] |
d_spread alone |
+1.703% [+1.496, +1.905] |
+0.426% [+0.250, +0.593] |
d1 + d_spread |
+1.889% [+1.659, +2.115] |
+0.519% [+0.337, +0.696] |
| all five neighbour signals | +1.897% [+1.661, +2.130] |
+0.522% [+0.336, +0.699] |
Entropy's interval spans zero on one book and tops out at +0.13% on the other. Distance and margin together are worth 0.52% and 1.89%, significant in both directions, and two signals capture essentially everything five do. Adding all seven signals is worth nothing more either way — slightly worse on The Valley of Fear (+1.798%) and indistinguishably better on His Last Bow (+0.521% against +0.519%).
Note how much the whole effect varies between two books of one author — 9.3% against 3.6% for the memory itself, 1.89% against 0.52% for the gate. The ordering is robust; the magnitude is not established.
A gate exists to predict the blend weight a token wants, so measure that. Split all 89,541 held-out tokens into fifths by each signal and find the single λ minimising NLL inside each fifth:
| fifth | by neighbour distance | by model entropy |
|---|---|---|
| 1 (nearest / most confident) | 0.163 | 0.157 |
| 2 | 0.107 | 0.080 |
| 3 | 0.071 | 0.077 |
| 4 | 0.062 | 0.084 |
| 5 (farthest / least confident) | 0.049 | 0.066 |
| spread | 3.3×, monotone | 2.4×, not monotone |
By distance the optimal weight falls monotonically across the whole range, and
σ(w·d + b) is monotone too, so the form fits the fact. By entropy, four of the five
fifths want the same weight, 0.066–0.084. Only the most confident fifth departs from
it, at 0.157, and those tokens sit at perplexity 3.7, already so well predicted that
hardly any total likelihood rides on their λ.
The underlying confusion is between two different questions:
| signal | correlation with the memory winning the token | correlation with nats actually saved |
|---|---|---|
entropy H |
−0.358 | +0.031 |
p_LM top-1 |
+0.375 | −0.019 |
d1 |
−0.186 | −0.025 |
d_mean |
−0.094 | +0.093 |
d_spread |
+0.069 | +0.176 |
p_kNN top-1 |
+0.169 | +0.075 |
H of p_kNN |
−0.170 | −0.081 |
Entropy is the strongest available predictor of which tokens the memory will win, and
essentially useless for how many nats those wins are worth. The margin is the reverse.
A gate fitted to likelihood optimises the second column, so entropy is answering a
question it was not asked. Put plainly: entropy measures how bad p_LM is; λ depends
on whether p_kNN is better, and knowing the model is lost tells you nothing about
whether the memory is found.
That also explains a result that would otherwise be confusing. The five fifths above are not monotone in entropy, so an obvious guess is that the sigmoid's monotonicity is the binding constraint. It is not.
| control | perplexity recovered |
|---|---|
| a constant λ | 9.259% |
| entropy gate, fitted by Adam | 9.267% |
entropy gate, best of a 19,481-point grid over (a, b) |
9.267% |
entropy gate, allowed a quadratic in H (non-monotone) |
9.263% |
An exhaustive grid search over the whole parameter space finds what Adam found to three
decimals, so this is not an optimisation artefact. Giving λ a quadratic in H, which
can express the non-monotone shape, does not help either — so the sigmoid's form was
never the problem.
Both controls point the same way on the other two eval sets, though not as tidily. On
the reverse book split the grid does beat Adam, by 0.015 points (3.711% against
3.696%) — so the entropy gate there is mildly underfitted, and its honest value is the
grid's, which is still only 0.115 points above a constant. On the trained-text set the
grid's own optimum comes out at a = +0.00: the best entropy gate available is exactly
a constant.
The single best signal is d_spread, the margin between the best match and a typical
one, ahead of the absolute distance d1. That fits: absolute distances drift with the
scale of the hidden states, while a margin is a comparison within one search and is
largely free of that.
Worth keeping in view: the best gate here still closes only about a tenth of the gap between a constant and the per-token oracle. Most of what a perfect per-token decision would buy is unclaimed.
A gate that opens further can pay for it on text the weights already handle. So the gate fitted on a never-trained book was applied, unchanged, to the trained-book slices (30,200 tokens, from five of the six books; the sixth was the fit set):
| gate | perplexity | vs no memory |
|---|---|---|
| no memory | 21.93 | — |
| the gate this project used to ship | 21.88 | 0.24% |
entropy H, refitted here |
21.71 | 1.04% |
| a constant λ | 21.71 | 1.04% |
d1 + d_spread, fitted on a never-trained book |
21.68 | 1.15% |
No measurable penalty. Its +0.109% over the constant has interval [−0.017, +0.241],
so the honest reading is "no worse", not "better". The refitted entropy gate ties the
constant to four decimals, which is the same finding again on different text.
Where the two gates genuinely differ is in how selective they are. Generating against
the shipped 750k-entry memory, the confidence gate's λ ranges over 0.028–0.279 while
the entropy gate spans 0.082–0.114 (runs/demo_receipt.log) — it shuts when the search
returns nothing close. That is what lets it open further where it counts without paying
for it elsewhere, and it is why checkpoints/lookaside-demo/ now defaults to
gate_mode="confidence".
Bucketing the same tokens by how often the model saw them in the six Doyle books it trained on:
| times seen in training | tokens | gain from the memory | best λ |
|---|---|---|---|
| never | 116 | −1.5% | 0.011 |
| 1–10 | 2,752 | 8.5% | 0.054 |
| 11–100 | 17,389 | 15.0% | 0.086 |
| 101–1,000 | 24,647 | 9.2% | 0.092 |
| more than 1,000 | 44,637 | 1.3% | 0.069 |
The payoff concentrates in the middle of the frequency range: words the model has met but not often enough to pin down. Common words gain almost nothing. Tokens it never saw are the honest failure — 116 of them, the memory makes them slightly worse, and the λ they want is 0.011, essentially closed, because a token missing from the training corpus is usually missing from the datastore too.
The "of headroom" framing used elsewhere is scale-dependent. Expressing a gate's
contribution as a share of the constant-to-oracle gap is convenient, but the oracle is
capped at λ_max = 0.5 while every fitted gate averages λ ≈ 0.07. Halving λ_max
would roughly double those percentages without changing a single gate, so the
relative-perplexity deltas with intervals are the numbers to trust.
The distance features are functions of the search whose output they gate. That is fine at inference — they are free, the search has already happened — but it does make "features of retrieval predict retrieval quality" less surprising than it first sounds. The interesting part is not that it works; it is that entropy does not.
The standardisation constants are datastore-specific. gate_mu and gate_sd are
the mean and spread of distances in the datastore fitted against (600k entries, 83%
coverage). The shipped checkpoint applies them to a slightly denser memory (750k
entries over all eight books, 88%). The measured λ range above confirms they transfer
here, but a substantially different datastore warrants refitting rather than reuse.
d_spread should travel better than d1 for the reason given above, which is an
argument, not a measurement.
Unswept hyperparameters. top_k = 64 and retrieval_temp = 20.0 set how peaked
p_kNN is, which bounds how much d_spread can possibly matter. Neither was varied.
Scale. One model at 27.8M parameters, one author's prose, a 725k-token memory, two held-out books. The mechanism — entropy predicting the wrong quantity — is general in principle; nothing here establishes the effect size beyond this setup.
python3 scripts/gate_study.py --eval-set holdout # the main tables and controls
python3 scripts/gate_study.py --eval-set holdout --fit-book the_valley_of_fear
python3 scripts/gate_study.py --eval-set trained \
--gate-from checkpoints/lookaside-demo/retrieval_config.json
python3 scripts/demo_receipt.py # lambda ranges and throughputThe first run caches the datastore and the per-token features, so the rest are quick.
runs/gate_study.log holds all three gate runs and runs/demo_receipt.log the
generation-time measurements.