knn-lm (k-nearest neighbor language model)
**kNN-LM (k-Nearest Neighbor Language Model)** is a retrieval-augmented language modeling approach that enhances any pre-trained neural language model by interpolating its output distribution with a non-parametric distribution derived from k-nearest neighbor search over a datastore of cached (context, target) pairs. At inference time, the model's hidden representation retrieves similar contexts from the datastore and uses their associated target tokens to construct an alternative prediction distribution, which is then combined with the model's own softmax output.
**Why kNN-LM Matters in AI/ML:**
kNN-LM provides **significant perplexity improvements without any additional training** by leveraging a datastore of examples, enabling domain adaptation, knowledge updating, and improved rare-word prediction through pure retrieval augmentation.
• **Datastore construction** — A single forward pass over the training data stores each token's (key, value) pair where key = the transformer's hidden representation at that position and value = the next token; this creates a non-parametric memory of all training contexts
• **kNN retrieval at inference** — For each generated token, the model's current hidden state queries the datastore for the k nearest neighbors (typically k=1024) using L2 distance, retrieving similar contexts and their associated next tokens
• **Distribution interpolation** — The kNN distribution p_kNN (softmax over negative distances to retrieved neighbors, grouped by target token) is interpolated with the model's parametric distribution p_LM: p_final = λ · p_kNN + (1-λ) · p_LM, where λ controls the retrieval weight
• **No additional training** — kNN-LM improves a pre-trained model's perplexity by 2-7 points without any gradient updates, weight modifications, or fine-tuning—only requiring a forward pass to build the datastore
• **Domain adaptation** — Swapping the datastore to domain-specific text instantly adapts the model to new domains (medical, legal, scientific) without retraining, providing a practical mechanism for rapid specialization
| Component | Specification | Notes |
|-----------|--------------|-------|
| Datastore | (h_i, w_{i+1}) pairs | Hidden state → next token |
| Index | FAISS (IVF + PQ) | Approximate nearest neighbor |
| k | 1024 (typical) | Number of retrieved neighbors |
| Distance | L2 norm | On hidden representations |
| Temperature | 10-100 | Sharpens kNN distribution |
| Interpolation λ | 0.2-0.5 | Tuned on validation set |
| Perplexity Gain | -2 to -7 points | Without any training |
**kNN-LM demonstrates that augmenting any pre-trained language model with non-parametric nearest-neighbor retrieval over cached representations provides substantial quality improvements without additional training, establishing a powerful paradigm for domain adaptation, knowledge updating, and retrieval-augmented generation that separates memorization from generalization.**