K-nearest neighbor (KNN) prompting: a complete guide

Few-shot prompting is weirdly fragile. Swap the examples in your prompt and accuracy can swing wildly, even when every example is correct (Liu et al., 2022; Lu et al., 2022). KNN prompting fixes the fragility by stopping the guessing: instead of picking demonstrations by hand or at random, it embeds your test input, finds the most similar labeled examples in a pool, and feeds those as the few-shot demos. The flip is "relevance over quantity," and it pays off — the calibration-free variant (Xu et al., 2023, ICLR) beats state-of-the-art calibration methods by +3.56 on average at 4-shot and +7.07 at 8-shot across 10 classification tasks.

See it work

Say you're classifying restaurant review sentiment and your test input is "The pasta was decent but the wait was too long."

With random few-shot you might get examples about phone battery life, a movie plot, and a hotel — all correctly labeled, none close to the task, so a different draw gives a different answer. With KNN prompting you retrieve the nearest neighbors in embedding space, and they're all mixed-sentiment restaurant reviews — the exact shape of the test case:

Random demos:  "Battery dies fast" → Negative
               "Loved the ending"  → Positive
               "Room was cramped"  → Negative
Test → ??? (unstable, changes per draw)

KNN demos:     "Food was great but service was slow" → Mixed
               "Tasty but overpriced and crowded"    → Mixed
               "Good meal, terrible parking"          → Mixed
Test → Mixed (stable)

Same model, same k, same pool — better examples. That's the whole idea.

The mental model

Think of a junior cook handed an unfamiliar order. Hand them three random recipes from the binder and they flail. Hand them the three recipes closest to the dish in front of them and they nail it. KNN prompting is the act of always pulling the closest recipes.

Relevance beats volume: a few demonstrations that mirror the test input's structure, domain, and reasoning teach more than many random ones.

KNN prompting sits in the example-based, retrieval-augmented family — it's a few-shot optimization technique. Two research lines live under the name:

• KNN exemplar selection (KATE) — Liu et al. (2022), "What Makes Good In-Context Examples for GPT-3?" Uses sentence embeddings (RoBERTa, cosine similarity) to retrieve the most similar training examples as demos, reaching performance nearly comparable to fine-tuning on GPT-3. It chooses which examples go in the prompt.
• kNN prompting for calibration-free inference — Xu et al. (2023), ICLR. Goes further: it uses the LLM's full output probability distribution as the representation and does nearest-neighbor classification directly, never mapping LLM outputs to task labels. Standard deviation across tasks drops from 9.14 (ICL) to 3.83.

Both share the principle of similarity-based retrieval; they just operate at different levels.

How it works

Both variants follow a two-phase, offline-then-per-query structure.

Variant 1 — KATE-style exemplar selection. Offline: collect labeled candidates, encode each with a sentence embedder (Sentence-BERT, RoBERTa, OpenAI embeddings), store the vectors in an index. Per query: encode the test input with the same model, score every candidate by distance (cosine, L2, or dot product), retrieve the k nearest, build a few-shot prompt, and call the LLM once. From the LLM's view this is single-pass — retrieval happens before the call.

Variant 2 — Xu et al. calibration-free inference. A meta-test stage builds a datastore: pick a small set of fixed anchor demos, then for each training example run the anchor prompt plus that example and cache the LLM's full output distribution as the key, paired with the true label as the value. A formal test stage runs the same anchor prompt plus the test input, gets its distribution, computes KL divergence against every cached distribution, takes the k nearest, and majority-votes their labels. This needs multiple LLM calls to build the datastore, but inference then scales beyond the context window.

Why it works

The causal chain: semantic encoding → distance in embedding space → most-relevant demos selected → LLM gets contextually appropriate examples → less task ambiguity → better output. A positive loop reinforces it (better selection → more consistent outputs → easier to tune k and embedder → better selection); a negative one bites if the embedder is poor (superficial matches → worse-than-random performance → false signal that KNN "doesn't work").

It rests on the kNN-LM result (Khandelwal et al., 2020): augmenting a frozen LM with a nearest-neighbor lookup over cached representations hit 15.79 perplexity on Wikitext-103 (a 2.9-point gain, no extra training), and adding kNN retrieval over 3B examples to a model trained on 100M tokens improved perplexity from 19.59 to 13.73 — retrieving from a corpus beat training on it.

The Xu variant's scaling trend holds across roughly 10 orders of magnitude, from 2 shots to 1024 shots, and across model sizes from 0.8B to 30B parameters.

Where it shines

KNN prompting helps anywhere labeled examples exist and example relevance varies by input.

• Text classification — sentiment, topic, intent, spam. Retrieval pulls same-topic or same-pattern demos; Liu et al. (2022) showed retrieval-based ICL on GPT-3 reaching near fine-tuning on multiple classification benchmarks.
• NER and information extraction — examples with matching entity types or domain terminology; strong when entity types shift across domains.
• Question answering — QA pairs matching question structure or reasoning type; multi-hop QA benefits from similar chain-of-reasoning demos.
• Generation and summarization — examples matching input length, style, or content for consistent tone and format.
• Machine translation — pairs sharing vocabulary, structure, or domain; domain-specific translation gains a lot.
• Code generation — examples with similar signatures, libraries, or algorithms; effective for API-specific tasks.

Domain results: clinical NLP (case retrieval for diagnostic reasoning, ICD coding, note summarization — BioSentVec/PubMedBERT embeddings help), legal (precedent and fact-pattern matching for outcome prediction, contract analysis), scientific literature (methodology-similar papers for review and claim verification), financial (similar reports/risk scenarios for earnings and risk analysis), and customer support (similar past tickets and resolutions for routing and suggested replies at scale).

When to use it (and when not)

Reach for it when random few-shot lands around 50-85% with high variance across example sets, you have 100+ labeled candidates, inputs are diverse (topics, domains, structures), you can deploy an embedding model, and you need automated, consistent selection at scale.

Skip it when zero-shot is already above 90%, random few-shot is above 90% with low variance, the pool is under 50 examples, inputs are near-identical (any example works), you can't host an embedding model, or embedding similarity simply doesn't capture task-relevant similarity.

Escalate when KNN-selected few-shot still sits below 60% (consider fine-tuning), you need to leverage thousands of examples (Xu variant or fine-tuning), embedding similarity misses task dimensions (supervised retriever like EPR/UDR), or you need guaranteed format compliance (structured-output APIs or fine-tuning).

Model fit: any few-shot-capable model works as a floor (GPT-3.5, Claude 3 Haiku, Llama 7B+); GPT-4, Claude 3.5 Sonnet, or Llama 70B+ are recommended for best results. Models with tiny context windows (under 2K tokens) or weak in-context learning are poor fits. The Xu variant additionally needs logit access (autoregressive LMs exposing output probabilities). Smaller 7B-13B models often gain more from the distribution-matching variant than from packing examples into context.

Variants and alternatives

Quick read: KATE is the practical default — simple, fast, works with any LLM API and no logit access. Xu's kNN prompting wins on classification when you have logit access and a large training set. Vote-k when you fear redundancy in a unlabeled pool. EPR/UDR when you can train a task-specific retriever (EPR shows 30%+ over random). RAG when you need external knowledge, not just demonstrations.

Components and structure

Required: a candidate pool of labeled examples (50-100 minimum, 500+ recommended), an embedding model (Sentence-BERT, OpenAI, RoBERTa), a distance metric (cosine, L2, dot product), a k value (typically 3-8), and a few-shot prompt template. The Xu variant additionally needs fixed anchor examples, a distribution datastore, and a KL-divergence computation.

Optional: a vector index (FAISS, Annoy, HNSW) for large pools, a fine-tuned or domain-specific embedder, diversity filtering, an ordering strategy, and a reranking model.

The prompt itself is just standard few-shot — the technique is the selection:

Input: {retrieved_input_1}   # most similar to the test input
Output: {retrieved_output_1}

Input: {retrieved_input_2}   # second most similar
Output: {retrieved_output_2}
...
Input: {test_input}
Output:

Design principles: fill every slot with the most relevant available demo; keep diversity within relevance (near-duplicate neighbors waste slots — see MMR below); format all retrieved examples identically regardless of source; and make sure the embedder captures the similarity dimensions that matter for the task. Order usually runs most-similar-first or last — test both, since models differ.

Scenario tweaks: for ambiguous tasks, raise k and add a disambiguating instruction; for complex reasoning, retrieve on reasoning structure (or CoT-annotated examples), not just surface text; for format-critical tasks, pre-filter the pool to correctly formatted examples; for domain tasks, use a domain embedder (PubMedBERT, LegalBERT) and optionally per-domain indices.

The core mechanism

The essence of KATE-style selection is a few lines — embed, score, take the top k:

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer('all-MiniLM-L6-v2')
candidate_embeddings = model.encode([ex['input'] for ex in candidates],
                                     normalize_embeddings=True)

def knn_prompt(test_input, k=5):
    test_emb = model.encode([test_input], normalize_embeddings=True)
    sims = np.dot(candidate_embeddings, test_emb.T).flatten()  # cosine
    top_k = np.argsort(sims)[-k:][::-1]                        # most similar first
    prompt = ""
    for idx in top_k:
        prompt += f"Input: {candidates[idx]['input']}\nOutput: {candidates[idx]['output']}\n\n"
    return prompt + f"Input: {test_input}\nOutput:"

For large pools, swap the brute-force dot product for a FAISS index: IndexFlatIP does exact search under ~10,000 examples; IndexIVFFlat or IndexHNSW handle 10,000-1M; add quantization beyond 1M. Retrieval latency is ~1-10ms — negligible against LLM inference.

The Xu et al. variant replaces embedding similarity with distribution distance. You cache each training example's output distribution, then match the test distribution by KL divergence and vote:

import numpy as np
from scipy.special import rel_entr
from collections import Counter

def predict(test_dist, datastore_keys, datastore_values, k=5):
    # symmetric KL divergence between the test distribution and each cached one
    dists = []
    for stored in datastore_keys:
        f = np.sum(rel_entr(test_dist + 1e-10, stored + 1e-10))
        b = np.sum(rel_entr(stored + 1e-10, test_dist + 1e-10))
        dists.append((f + b) / 2)
    nearest = np.argsort(dists)[:k]
    return Counter(datastore_values[i] for i in nearest).most_common(1)[0][0]

For diversity, Maximal Marginal Relevance (MMR) rescores candidates as lambda * relevance - (1 - lambda) * max_similarity_to_already_selected. Lambda 1.0 is pure relevance (standard KNN), 0.5 is balanced, and 0.7 is a good mildly-diverse default — tune on validation.

A platform shortcut

LangChain ships KNN selection as a built-in, so you rarely write the loop yourself:

from langchain.prompts import FewShotPromptTemplate, PromptTemplate
from langchain.prompts.example_selector import SemanticSimilarityExampleSelector
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

selector = SemanticSimilarityExampleSelector.from_examples(
    examples, OpenAIEmbeddings(), FAISS, k=5)   # KNN under the hood

prompt = FewShotPromptTemplate(
    example_selector=selector,
    example_prompt=PromptTemplate(input_variables=["input", "output"],
                                  template="Input: {input}\nOutput: {output}"),
    suffix="Input: {input}\nOutput:", input_variables=["input"])

LlamaIndex (SimilarityPostprocessor) and DSPy (BootstrapFewShot over KNN-selected examples) offer similar paths.

Configuration

k: 1-2 is too few (high variance), 3-8 is the sweet spot, above 10 brings context pressure and diminishing returns. Start at k=5 and tune. Distance metric: cosine is the default for normalized embeddings; L2 suits unnormalized; dot product is fastest. For normalized vectors all three rank identically. Storage is roughly dimensions × candidates × 4 bytes; typical token use is 4-8 examples × 100-300 tokens = 400-2,400 tokens for demos.

Task tuning: classification k=3-5 with at least one example per class (consider label-balanced retrieval); reasoning/QA k=5-8 prioritizing reasoning-pattern over topic similarity; generation k=3-5 favoring style over topic, optionally embedding input+output; code k=5-8 with code-specific embedders (CodeBERT, UniXcoder) and diverse API patterns.

Implementation workflow

1. Baseline. Measure zero-shot and random few-shot (k=5, run 10+ times). If random few-shot is already above 90% with low variance, you probably don't need KNN.
2. Prepare the pool. Collect and clean labeled examples (500-5,000 is a good range), dedupe near-duplicates, verify label quality — retrieval amplifies bad labels — then split into retrieval pool (80%), validation (10%), test (10%).
3. Embed and index. Encode the pool, build the index sized to the pool (see above), and spot-check retrieval on sample queries.
4. Tune. Sweep k from 3 to 8, compare embedders, test with/without diversity filtering, and test ordering (most-similar first vs last).
5. Evaluate. Full run on validation against the random baseline; analyze failures and retrieval quality; finalize on held-out test.
6. Deploy. Serve the embedder, deploy the index, wire retrieval into the inference pipeline, monitor retrieval similarity and accuracy, and refresh the pool periodically.

Do: validate that embedding similarity tracks task similarity before committing; start with a strong general embedder before fine-tuning; add diversity filtering if neighbors are near-duplicates; normalize embeddings; cache embeddings and retrieval results for repeated inputs; test on diverse inputs, not just typical ones.

Don't: assume any embedder works; push k above 8 without checking context limits; skip the random baseline (you must prove KNN helps); build the index on the test set (leakage); ignore diversity; deploy without monitoring for distribution shift.

Debugging

Testing and optimization

Hold out 10-20% of the pool for validation (never index it), tune on validation, and report on a separate test set. Judge retrieval quality directly — Precision@k, Recall@k, nDCG@k, MRR — not just downstream accuracy, since high similarity scores can hide useless matches. Cover common cases (50%), domain-boundary cases (20%), short/long inputs (15%), out-of-distribution (10%), and adversarial inputs (5%). Track task metrics (accuracy/F1 for classification, BLEU/ROUGE for generation, exact-match/F1 for QA, test pass rate for code) plus improvement over the random baseline, consistency, and retrieval latency.

To prove it beats the baseline, run a paired comparison across trials:

import numpy as np, random
from scipy.stats import ttest_rel

def ab_test(knn, candidate_pool, test_set, k=5, trials=20):
    knn_acc, rand_acc = [], []
    for _ in range(trials):
        knn_acc.append(evaluate(knn, test_set))                 # deterministic
        rand = random.sample(candidate_pool, k)                 # new draw each trial
        rand_acc.append(evaluate_with_fixed_examples(rand, test_set))
    t, p = ttest_rel(knn_acc, rand_acc)
    return np.mean(knn_acc), np.mean(rand_acc), p

Stop tuning when validation gains fall below 1%, Precision@k clears 0.8, the gap over random holds above 5%, more k stops helping, and embedders converge. Keep going if retrieval is visibly poor, the gain over random is under 3%, specific input categories keep failing, or you haven't tried a domain-specific embedder. Optimize cost by caching retrieval for repeated inputs and using approximate search; an open-source embedder removes per-query embedding cost entirely.

Limitations

• Embedding quality is the ceiling. If the embedder misses task-relevant similarity, you retrieve superficially similar but useless examples — worst for cases where surface text doesn't predict usefulness (e.g., look-alike math problems needing different methods).
• Cost for huge pools. Similarity against all candidates is expensive; ANN indices help per-query, but upfront embedding, storage, and index management get nontrivial past millions of examples.
• Context-window pressure. k=5 at ~200 tokens each is ~1,000 tokens before the test input — directly trading against input and response space, and capping k on small-context models.
• No diversity guarantee. Dense clusters yield near-duplicate neighbors; MMR helps but adds a hyperparameter and can lower average relevance.
• Sparse distribution (Xu variant). The kNN distribution only puts mass on neighbors, so it can miss tokens needed for some predictions, especially with a small datastore.
• Static retrieval. It retrieves on the initial input and doesn't adapt to intermediate outputs — initially retrieved demos can go stale in multi-turn or iterative tasks; there's no feedback loop.
• Infrastructure overhead. You now maintain an embedder, an index, and a pool alongside the LLM — qualitatively more than a single API call.

Edge cases worth handling: ambiguous inputs equidistant between clusters (raise k or disambiguate); out-of-distribution inputs where max cosine similarity drops below ~0.5 (fall back to zero-shot below a threshold like 0.3); adversarial inputs crafted to trigger specific demos (monitor retrieval patterns, validate outputs); very short inputs (low-information embeddings — raise k) and very long ones (match on irrelevant detail — use passage-level embeddings); and label imbalance in the retrieved set (stratified retrieval to guarantee per-class representation).

Advanced techniques

What you embed matters. Input-only embedding is simplest but misses output-conditional relevance; input+output embedding suits generation where style matters; task-description-augmented embedding (prepending the task) focuses the embedder on task-relevant features. When examples are long, truncate to key portions, summarize, lower k, or use a tiered approach (full nearest example, abbreviated rest).

Hybrids. Combine KNN with CoT + self-consistency (retrieve relevant CoT-annotated demos, sample several reasoning paths, majority-vote). Pair KNN + RAG — KNN supplies format and reasoning demos while RAG supplies factual grounding. Layer KNN + active prompting — KNN narrows to a top-20 relevant set, active prompting picks the most informative of those. Use graceful fallback: if average retrieval similarity is below a threshold, drop to zero-shot rather than prompting with irrelevant examples.

Model notes. GPT-4 handles k=8-10 and pairs well with text-embedding-3-large; Claude 3.5 Sonnet/Opus often need fewer demos (k=3-5) and have no native embedder (use open-source or OpenAI); Llama 70B/405B benefit strongly but may need k=5-8 and are order-sensitive; 7B-13B models gain less from context-packing and may prefer the distribution-matching variant. The embedder is independent of the LLM, so one index serves many models — but tune k per model.

Domain adaptation, three tiers: general embedder + domain pool (quick, works for most); domain-specific embedder (BioSentVec, LegalBERT, CodeBERT) for better retrieval; fine-tuned embedder on domain or NLI/STS-B data for the best ICL relevance. Bootstrap a new domain with analogous examples from related domains, then replace them with genuine ones as they accumulate.

Risk and ethics

Candidate pools often hold sensitive or proprietary data; once retrieved, it's shipped to LLM APIs and can surface in outputs — anonymize pools, use on-prem models for sensitive data, and gate the index with access controls. Bias amplifies through three channels: a skewed pool, biased embeddings, and proximity bias (near neighbors share the query's biases rather than correcting them). Audit the pool, prefer debiased embedders, and watch output-fairness metrics. For transparency, make retrieval auditable — log which examples were retrieved and why — and document the pool composition and embedder used.

Ecosystem

Embeddings and indexing: Sentence-Transformers (all-MiniLM-L6-v2, all-mpnet-base-v2), OpenAI embeddings, and FAISS (exact + approximate search, GPU support; the same library Khandelwal et al. used for kNN-LM). Frameworks: LangChain (SemanticSimilarityExampleSelector, FewShotPromptTemplate, many vector-store backends), LlamaIndex (SimilarityPostprocessor), and DSPy (BootstrapFewShot). Production vector DBs: Pinecone, Weaviate, Chroma, Milvus, Qdrant. Evaluation: BEIR (retrieval), MTEB (embedding-model comparison), and Ragas (retrieval-augmented systems).

Lineage: kNN-LM (Khandelwal et al., 2020) is the token-level ancestor; KATE (Liu et al., 2022) brought retrieval to example selection with RoBERTa + cosine; EPR (Rubin et al., 2022) added a supervised retriever (BM25 recall → trained reranker, 30%+ over random); UDR (Li et al., 2023) unified retrieval across tasks; Vote-k (Su et al., 2023) added graph-based diversity; CEIL (Ye et al., 2023) modeled the joint probability of the whole example set with conditional DPPs; Xu et al. (2023) moved to distribution matching.

Transitions. From random few-shot: measure the baseline, build the index on your existing pool, compare on validation, deploy if the gain clears ~3%. To a supervised retriever: label which retrieved examples actually helped, train EPR/UDR on that, deploy if it beats unsupervised KNN. To fine-tuning: harvest high-performing example-output pairs from KNN logs and switch if fine-tuning clears ~10%.

Future directions

Recent work pushes several frontiers. Nearest-neighbor speculative decoding (Sun et al., 2024) uses kNN matches to speculatively decode multiple tokens in parallel for faster inference. bias-kNN (IEEE ICSC 2024) treats biased output distributions as features rather than problems, beating ICL in few-shot settings with more stability. kNN-ICL (NAACL 2024, Zhao et al.) layers nearest-neighbor inference over any ICL design and gives access to all demos without context limits. IDEAL (ICLR 2024) does influence-driven selective annotation, and DQ-LoRe (ICLR 2024) uses dual queries plus low-rank approximation for exemplar selection on reasoning tasks.

Open questions: which similarity dimensions matter most for ICL (and can we learn task-specific ones); how to optimize the example set jointly rather than each demo independently; whether k should adapt per query; cross-lingual retrieval; scaling laws for retrieval; when retrieving real examples beats generating synthetic ones (SG-ICL); and privacy-preserving retrieval over sensitive pools. Promising directions include hierarchical retrieval (domain first, then examples), embedding-model co-training for ICL relevance, real-time pool evolution from production traffic, and multimodal KNN prompting.

Summary

• What it is: retrieve the most similar labeled examples as few-shot demos instead of picking them randomly or by hand — relevance over quantity.
• Why it works: relevant demos shrink the example-to-answer gap and cut variance; the Xu variant matches output distributions to sidestep calibration entirely (std 3.83 vs 9.14).
• Two variants: KATE selects which examples enter the prompt (any LLM, no logit access); Xu et al. does calibration-free nearest-neighbor inference over cached distributions (+3.56 at 4-shot, +7.07 at 8-shot, scales to 1024 shots).
• When to use: a pool of 50+ (ideally 500+) labeled examples, diverse inputs, and high-variance random few-shot in the 50-85% range.
• When not: zero-shot or low-variance few-shot already above 90%, tiny or homogeneous pools, or embeddings that don't track task similarity.
• How to run it: embed the pool, index it (FAISS), retrieve k=3-8 nearest, prove the lift against a random baseline, and monitor for distribution shift.
• Watch out for: embedder quality (the ceiling), near-duplicate neighbors (use MMR), context pressure, pool poisoning, and bias amplification.