Research map / Large language models / Retrieval and knowledge
Why do Nearest Neighbor Language Models Work?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's title only (no summary yet). Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use attention-layer output (att) instead of feedforward output (ffn) as kNN query representation (ablation) | kNN-LM · Wikitext-103 | interpolated perplexity: 19.095 vs 21.101 (IP) | Helped |
| Replace kNN datastore with a trainable 1x V x D embedding matrix (ffn input) trained with other params frozen (ablation) | parametric LM · Wikitext-103 | perplexity: 20.920 vs 21.750 base | Helped |
| Add a separate learned V x D embedding output ensembled with the base LM (att input) (ablation) | parametric LM · Wikitext-103 | interpolated perplexity: 20.353 vs 21.750 base; recovers about 55% of kNN gain | Helped |
| Use exact (real) distance scores instead of FAISS approximate scores in kNN (ablation) | kNN-LM · Wikitext-103 | interpolated perplexity: 19.672 vs 19.174 | Hurt |
| Use exact nearest neighbor search (real mask) and exact scores instead of FAISS approximate search (ablation) | kNN-LM · Wikitext-103 | interpolated perplexity: 19.735 vs 19.174 | Hurt |
| Increase softmax temperature/tune temperature for kNN distribution away from 1.0 (ablation) | kNN-LM · Wikitext-103 (full and 5% datastore) | interpolated perplexity: up to 3.7% relative perplexity gap between default and optimal temperature (text) | Helped |
| Increase number of embeddings per word type (multi-embedding softmax, 3x) (ablation) | parametric LM with learned W_ds · Wikitext-103 | interpolated perplexity: 20.395 (3x att) vs 20.353 (1x att); non-monotone in n | No clear effect |
| Adaptive (frequency or loss-weighted) embedding counts per word type (ablation) | learned W_ds · Wikitext-103 | interpolated perplexity: e.g., 20.440 loss-weighted vs 20.395 equal (att, 3x) | No clear effect |
| Mixture of Softmax on top of parametric LM (ablation) | parametric LM · Wikitext-103 | perplexity: MoS-only att R=2 22.552 vs base 21.750 (worse); ffn R=4 21.321 (small gain) | No clear effect |
| Subsample datastore to 5% vs full (ablation) | kNN-LM · Wikitext-103 | interpolated perplexity: perplexity decreases almost linearly with more datastore; 5% still beats base LM | Helped |
| Add kNN-LM to pretrained GPT2 family and DistilGPT2 base models | pretrained GPT2 models · Wikitext-103 | perplexity: e.g., GPT2-large 10.56 to 9.76; Distilled-GPT2 18.25 to 14.84 | Helped |
| Sparsify parametric LM output to top-k tokens (no retrieval) (ablation) | parametric LM · Wikitext-103 | perplexity: does not help at all | No clear effect |
Closest papers by meaning
- Test-Time Training on Nearest Neighbors for Large Language Models (ICLR 2024 · similarity 0.88)
- Nearest Neighbor Speculative Decoding for LLM Generation and Attribution (NeurIPS 2024 · similarity 0.86)
- Language Ranker: A Metric for Quantifying LLM Performance Across High and Low-Resource Languages (AAAI 2025 · similarity 0.85)
- RECOMP: Improving Retrieval-Augmented LMs with Context Compression and Selective Augmentation (ICLR 2024 · similarity 0.85)
- Scaling Retrieval-Based Language Models with a Trillion-Token Datastore (NeurIPS 2024 · similarity 0.85)
- Retrieval meets Long Context Large Language Models (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.