Home / Curated / Algorithms / KATE kNN-Augmented In-Context Example Selection

KATE kNN-Augmented In-Context Example Selection

Algorithms AIMech:KATEkNNAugmentedInContextExampleSelection DRAFT

Aliases: KATE, kNN-Augmented In-Context Example Selection, KATE Nearest-Neighbor In-Context Demonstration Retrieval

A non-parametric in-context demonstration selection algorithm for pretrained language models that maps training instances and a test input into a shared vector space using a pretrained or task-fine-tuned sentence encoder, retrieves the k nearest training neighbors to the test input under Euclidean or cosine similarity, and concatenates those retrieved input-output pairs in distance-sorted order as the conditioning context for autoregressive prediction.

Identity Notes & Concept Scope Distinct from datastore-interpolated language modeling such as kNN-LM and kNN-MT, which interpolate a language model's next-token output probability distribution with a non-parametric kNN distribution over cached hidden states, because KATE operates purely at the input prompt level using an external sentence encoder without accessing the target language model's internal hidden states or modifying its decoding distribution. Also distinct from LM-scored contrastive prompt retrieval (AIMech:EPREfficientPromptRetrievalForInContextLearning), which trains a custom bi-encoder using positive and hard-negative labels scored by the target language model itself.

Claims & Grounded Evidence

4 Claims

Mechanism #kate-sentence-encoder-knn-retrieval-mechanism

Given a test query x_test and a training pool D_T = {(x_i, y_i)}, KATE encodes all inputs with a sentence encoder mu_theta, computes similarity scores s_i via negative Euclidean distance or cosine similarity, selects the top-k nearest neighbors ordered by distance, and concatenates their input-target pairs before x_test as the prompt for the frozen language model.

supports Source: What Makes Good In-Context Examples for GPT-3? · Locator: Section 2.3 & Algorithm 1, pp. 3-4

Mechanism #kate-task-related-encoder-finetuning-mechanism

Fine-tuning the retrieval sentence encoder on semantically aligned tasks—such as natural language inference and semantic textual similarity (NLI+STS-B) for open-domain question answering or domain-matched sentiment data—improves the quality of retrieved in-context examples compared to unaligned encoders.

supports Source: What Makes Good In-Context Examples for GPT-3? · Locator: Section 3.1 & Section 4.1-4.3, pp. 4-7, Tables 4-7

Application #kate-nlu-and-nlg-benchmark-application

KATE is applied to few-shot sentiment transfer (SST-2 to IMDB), controlled table-to-text generation (ToTTo), and open-domain question answering (Natural Questions, WebQuestions, and TriviaQA) with GPT-3, substantially outperforming random demonstration sampling and standalone kNN majority voting.

supports Source: What Makes Good In-Context Examples for GPT-3? · Locator: Section 4.1-4.3, pp. 5-7, Tables 4, 5, and 7
supports Source: What Makes Good In-Context Examples for GPT-3? (Crossref metadata) · Locator: DBLP bibliographic record Jiachang Liu et al.

Limitation #kate-dissimilar-encoder-task-mismatch-limitation

When the sentence encoder is fine-tuned on an objective that diverges from the downstream task structure (such as applying an STS-B similarity encoder to table-to-text generation or sentiment transfer), KATE's downstream accuracy can degrade relative to an un-finetuned pretrained RoBERTa encoder.

supports Source: What Makes Good In-Context Examples for GPT-3? · Locator: Section 4.1 & Section 4.2, pp. 5-6, Tables 4-5

⚠️ Curated Limitations & Negative Results

Requires encoding a large labeled training datastore and performing nearest-neighbor lookup for every test query, and its gains depend on the availability of a sufficiently large retrieval pool.
Selects each neighbor independently by surface/semantic input similarity without modeling compositional interactions among the chosen examples.

Artificial Intelligence Ontology (AIO) Grounding

Mapped AIO Classes

AIO Target Class	Mapping Relation	Mapping Rationale
`AIO:ArtificialNeuralNetwork` Artificial Neural Network	related_match	KATE retrieves nearest-neighbor in-context demonstrations in a neural sentence encoder space to condition a frozen neural language model.

Evidence Sources & Primary Literature

3 Sources

Source ID	Publication / Document	Type / Access	Reference & Checksum
`r734-primary-pdf`	What Makes Good In-Context Examples for GPT-3? Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, Weizhu Chen	primary paper full_text	ID: https://dblp.org/search/publ/api?q=What+Makes+Good+In-Context+Examples+for+GPT-3+Liu&h=5&format=xml sha256:2507fe260f51…
`r734-arxiv-abs`	What Makes Good In-Context Examples for GPT-3? (arXiv abstract) Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, Weizhu Chen	primary paper abstract_only	ID: https://arxiv.org/abs/2101.06804 sha256:2e35bb8dc68c…
`r734-metadata-cr`	What Makes Good In-Context Examples for GPT-3? (Crossref metadata) Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, Weizhu Chen	primary paper abstract_only	ID: https://dblp.org/search/publ/api?q=What+Makes+Good+In-Context+Examples+for+GPT-3+Liu&h=5&format=xml sha256:b74d1163241c…

Ontology & Grounding Topology

Open in interactive network graph →

Curation Provenance

Curator: Antigravity

Method: ai_assisted

Model: Gemini

Created: 2026-10-19T13:37:00Z

Updated: 2026-10-19T13:37:00Z

Source File: data/curated/KATEkNNAugmentedInContextExampleSelection.yaml · Schema: curated_concept.yaml

Explore across X-Mech

Related computational & biological Mechs: Embedding Map Network Graph ProteinTraitsMech (ESM-2 pLM) DUFMech (Domain pLM) TraitMech CultureMech All 14 X-Mech Repositories