Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

kNN-LM: A Datastore of Contexts and Next Tokens

4 minute read

Published:

TL;DR: kNN-LM stores hidden representations of observed contexts as keys and their next tokens as values. At prediction time, nearby keys vote for next-token probabilities, which are interpolated with the base model. This external memory operates at the token distribution, not through a passage prompt.

Build a key–value memory

Run a trained language model over a datastore corpus. For each permitted prediction position, save a context representation \(h_i\) and the observed next token \(v_i\). The key describes a prefix; the value is what followed that prefix in the stored example.

At inference, encode the current prefix, retrieve nearby keys and aggregate their values. The original kNN-LM work shows how a fixed neural model can benefit from this nonparametric memory without retraining all of its parameters.

Choosing the representation layer and distance metric is part of the method. A random hidden vector of the same dimension is not an equivalent retrieval key.

Turn distances into token votes

A general temperature-scaled neighbour distribution is

\[ p_{\mathrm{kNN}}(w\mid h)=\frac{\sum_{i\in N_k(h)}\mathbf1[v_i=w]\exp(-d(h,h_i)/T)}{\sum_{i\in N_k(h)}\exp(-d(h,h_i)/T)}. \]

The distance definition and \(T>0\) must be specified. The original formulation can be understood as a particular setting of this distance-weighting scheme. Several neighbours with the same value accumulate probability mass for that token.

Suppose three neighbours have weights 0.5, 0.3 and 0.2 after normalisation, with values charged, charged and empty. The datastore distribution gives charged probability 0.8 and empty probability 0.2. Repeated occurrences can strengthen a vote, but copied duplicates may also dominate it.

Interpolate with the neural model

Combine the memory and base model as

\[ p(w\mid h)=\lambda p_{\mathrm{kNN}}(w\mid h)+(1-\lambda)p_{\mathrm{LM}}(w\mid h),\qquad0\leq\lambda\leq1. \]

If the base model gives charged probability 0.2 and \(\lambda=0.4\), the combined probability is \(0.4(0.8)+0.6(0.2)=0.44\). Other tokens retain some base-model mass even if no retrieved neighbour votes for them.

The interpolation weight is a substantive modelling choice. Too much weight can over-trust irrelevant neighbours; too little can erase the benefit of domain-specific memory. Tune it without repeatedly selecting on the final test set.

A context match is not a factual source

A neighbour can support a likely word continuation without establishing a factual claim in the output. A stored context from an old manual may strongly favour an obsolete product name or date. The model can improve local prediction while producing a statement inappropriate for the current question.

If source-grounded answers are required, preserve the provenance of datastore entries and evaluate claim support. Token votes alone do not supply a clear human-readable explanation of an entire sentence.

Domain adaptation through memory

A target-domain datastore can introduce terminology the base model handles poorly. This makes memory replacement an attractive adaptation mechanism. But if the encoder represents the new domain badly, nearest neighbours may still be unhelpful.

Compare in-domain and out-of-domain memories, inspect neighbour distances and test rare terms separately. Include a no-memory baseline and an exact-search sample to separate representation failures from approximate-index failures.

Storage, duplicates and privacy boundaries

A datastore can contain one entry per prediction position, so storage can grow much faster than a passage index with one vector per chunk. Quantisation, pruning and approximate search change its resource–quality tradeoff.

Copied passages create correlated keys and values. Deduplicate or at least measure this concentration when interpreting improvements. For permissioned data, search must respect the same access boundary as the underlying text; hiding the source text while exposing its token votes is not automatically equivalent to excluding it.

The RETRO chapter retrieves text for neural conditioning, while REPLUG mixes document-conditioned predictions. Comparing the point where memory enters prediction clarifies the family resemblance without collapsing the methods together.

Key Insight: kNN-LM adds a data-dependent distribution at the output layer. The neural model supplies generalisation, while neighbouring stored contexts supply local empirical evidence.
Common pitfall: Do not describe a next-token datastore as ordinary document RAG. It has different storage scale, provenance and inference requirements.

Recap

  • Datastore keys are context representations and values are next tokens.
  • Distance-weighted neighbours define a token distribution.
  • Interpolation controls reliance on memory versus the base model.
  • Evaluate domain transfer, duplication and datastore costs explicitly.

References

  1. Khandelwal et al. Generalization through Memorization: Nearest Neighbor Language Models. ICLR, 2020.
  2. Borgeaud et al. Improving Language Models by Retrieving from Trillions of Tokens. ICML, 2022.
  3. Shi et al. REPLUG: Retrieval-Augmented Black-Box Language Models. NAACL, 2024.

You May Also Enjoy

The Standard Library Modules Worth Knowing Before You pip install

8 minute read

Published:

Python ships with a counter, a queue, a memoiser, a CLI parser, a test framework’s smaller cousin, and a JSON-safe date library already installed, most of what people reach for a package to solve is one import away.

Files and Context Managers: Why with Is Not Optional

9 minute read

Published:

A file object that falls out of scope without with does get closed eventually, but ‘eventually’ means whenever the garbage collector gets around to it, which is not a promise any program handling more than a handful of files can live with.

Errors and Exceptions: EAFP, the Hierarchy, and Reading a Traceback

9 minute read

Published:

Exceptions in Python are not exceptional. They are a normal control-flow mechanism that the language leans on so heavily that the idiomatic style is to try the operation and handle the failure, rather than check first, which is faster, shorter, and free of race conditions.