KITE: Kernelized and Information Theoretic Exemplars for In-Context Learning

2025/09/19 by Vaibhav Singh, Singh, Vaibhav, Soumya Suvra Ghosal +7
Computer Science · #Artificial Intelligence (cs.AI) #Computation and Language (cs.CL) #Data Stream Mining Techniques #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning and Algorithms #Machine Learning and Data Classification

paper · pdf · doi:10.48550/arxiv.2509.15676

openalex publication_date 2025/09/19 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/01

Abstract

In-context learning (ICL) has emerged as a powerful paradigm for adapting large language models (LLMs) to new and data-scarce tasks using only a few carefully selected task-specific examples presented in the prompt. However, given the limited context size of LLMs, a fundamental question arises: Which examples should be selected to maximize performance on a given user query? While nearest-neighbor-based methods like KATE have been widely adopted for this purpose, they suffer from well-known drawbacks in high-dimensional embedding spaces, including poor generalization and a lack of diversity. In this work, we study this problem of example selection in ICL from a principled, information theory-driven perspective. We first model an LLM as a linear function over input embeddings and frame the example selection task as a query-specific optimization problem: selecting a subset of exemplars from a larger example bank that minimizes the prediction error on a specific query. This formulation departs from traditional generalization-focused learning theoretic approaches by targeting accurate prediction for a specific query instance. We derive a principled surrogate objective that is approximately submodular, enabling the use of a greedy algorithm with an approximation guarantee. We further enhance our method by (i) incorporating the kernel trick to operate in high-dimensional feature spaces without explicit mappings, and (ii) introducing an optimal design-based regularizer to encourage diversity in the selected examples. Empirically, we demonstrate significant improvements over standard retrieval methods across a suite of classification tasks, highlighting the benefits of structure-aware, diverse example selection for ICL in real-world, label-scarce scenarios.

Citations

• Relic: Enhancing Reward Model Generalization for Low-Resource Indic Languages with Few-Shot Examples
• PromptRefine: Enhancing Few-Shot Performance on Low-Resource Indic Languages with Example Selection from Related Example Banks
• Parameter-Efficient Fine-Tuning in Large Models: A Survey of Methodologies
• The Llama 3 Herd of Models
• NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models
• Prompt Optimization with EASE? Efficient Ordering-aware Automated Selection of Exemplars
• In-Context Learning Demonstration Selection via Influence Analysis
• In-context Learning with Retrieved Demonstrations for Language Models: A Survey
• Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena
• Dr.ICL: Demonstration-Retrieved In-context Learning
• What In-Context Learning "Learns" In-Context: Disentangling Task Recognition and Task Learning
• Symbol tuning improves in-context learning in language models
• Unified Demonstration Retriever for In-Context Learning
• Larger language models do in-context learning differently
• In-context Example Selection with Influences
• Compositional Exemplars for In-context Learning
• Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning
• Transformers as Algorithms: Generalization and Stability in In-context Learning
• A Survey on In-context Learning
• Self-Adaptive In-Context Learning: An Information Compression Perspective for In-Context Example Selection and Ordering
• Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers
• Text Embeddings by Weakly-Supervised Contrastive Pre-training
• In-context Learning and Induction Heads
• What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
• Ground-Truth Labels Matter: A Deeper Look into Input-Label Demonstrations
• Improving In-Context Few-Shot Learning via Self-Supervised Training
• On the Effect of Pretraining Corpora on In-context Learning by a Large-scale Language Model
• Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?
• Learning To Retrieve Prompts for In-Context Learning
• An Explanation of In-context Learning as Implicit Bayesian Inference
• MetaICL: Learning to Learn In Context
• Calibrate Before Use: Improving Few-Shot Performance of Language Models
• What Makes Good In-Context Examples for GPT-3?
• A Mathematical Exploration of Why Language Models Help Solve Downstream Tasks
• RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models
• Bandit Algorithms
• Language Models are Few-Shot Learners
• MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers
• Towards a Human-like Open-Domain Chatbot
• DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
• An Evaluation Dataset for Intent Classification and Out-of-Scope Prediction
• RoBERTa: A Robustly Optimized BERT Pretraining Approach
• SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems
• HellaSwag: Can a Machine Really Finish Your Sentence?
• HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering
• SWAG: A Large-Scale Adversarial Dataset for Grounded Commonsense Inference
• GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding
• TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for\n Reading Comprehension
• A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference
• Submodular meets Spectral: Greedy Algorithms for Subset Selection, Sparse Approximation and Dictionary Selection
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
• CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge

Related

• Large Language Models Know What Makes Exemplary Contexts
• Revisiting In-Context Learning with Long Context Language Models
• Finding Support Examples for In-Context Learning
• Does In-Context Learning Really Learn? Rethinking How Large Language Models Respond and Solve Tasks via In-Context Learning
• Large Language Models are Demonstration Pre-Selectors for Themselves