# Demonstration Selection for In-Context Learning via Reinforcement Learning

arXiv: https://arxiv.org/abs/2412.03966 · HTML: https://mlsanity.com/paper/2412.03966/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 36 of its 49 references that are in the map.
- First posted on arXiv 5 Dec 2024.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)

## Builds on (most cited first, 30 of 36)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Least-to-Most Prompting Enables Complex Reasoning in Large Language Models](https://mlsanity.com/paper/2205.10625/) (2022)
- [Calibrate Before Use: Improving Few-Shot Performance of Language Models](https://mlsanity.com/paper/2102.09690/) (2021)
- [Automatic Chain of Thought Prompting in Large Language Models](https://mlsanity.com/paper/2210.03493/) (2022)
- [Text Summarization with Pretrained Encoders](https://mlsanity.com/paper/1908.08345/) (2019)
- [Efficient Intent Detection with Dual Sentence Encoders](https://mlsanity.com/paper/2003.04807/) (2020)
- [An Evaluation Dataset for Intent Classification and Out-of-Scope Prediction](https://mlsanity.com/paper/1909.02027/) (2019)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Generated Knowledge Prompting for Commonsense Reasoning](https://mlsanity.com/paper/2110.08387/) (2021)
- [Active Example Selection for In-Context Learning](https://mlsanity.com/paper/2211.04486/) (2022)
- [The Next Decade in AI: Four Steps Towards Robust Artificial Intelligence](https://mlsanity.com/paper/2002.06177/) (2020)
- [Benchmarking Natural Language Understanding Services for building Conversational Agents](https://mlsanity.com/paper/1903.05566/) (2019)
- [A Comprehensive Survey of Few-shot Learning: Evolution, Applications, Challenges, and Opportunities](https://mlsanity.com/paper/2205.06743/) (2022)
- [Compositional Exemplars for In-context Learning](https://mlsanity.com/paper/2302.05698/) (2023)
- [Diverse Demonstrations Improve In-context Compositional Generalization](https://mlsanity.com/paper/2212.06800/) (2022)
- [In-Context Learning with Iterative Demonstration Selection](https://mlsanity.com/paper/2310.09881/) (2023)
- [Coverage-based Example Selection for In-Context Learning](https://mlsanity.com/paper/2305.14907/) (2023)
- [Skill-Based Few-Shot Selection for In-Context Learning](https://mlsanity.com/paper/2305.14210/) (2023)
- [Which Examples to Annotate for In-Context Learning? Towards Effective and Efficient Selection](https://mlsanity.com/paper/2310.20046/) (2023)

## Built on by (most cited first, 8 of 8)

- [DICE: Dynamic In-Context Example Selection in LLM Agents via Efficient Knowledge Transfer](https://mlsanity.com/paper/2507.23554/) (2025)
- [Thinking with Nothinking Calibration: A New In-Context Learning Paradigm in Reasoning Large Language Models](https://mlsanity.com/paper/2508.03363/) (2025)
- [Learning to Select Visual In-Context Demonstrations](https://mlsanity.com/paper/2603.26775/) (2026)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [GRIP: Feedback-Guided Prompt Retrieval for Large Multimodal Models](https://mlsanity.com/paper/2606.12744/) (2026)
- [Downstream-Aware Context Selection for Online In-Context Reinforcement Learning](https://mlsanity.com/paper/2609.33166/) (2026)
- [You Only Edit Once: Incentivizing In-Context Capability of LLMs via Local Demonstration Refinement](https://mlsanity.com/paper/2609.33609/) (2026)
- [Query-Adaptive Neuron Steering for LLMs](https://mlsanity.com/paper/2610.14519/) (2026)

## Abstract (arXiv)

Diversity in demonstration selection is critical for enhancing model generalization by enabling broader coverage of structures and concepts. Constructing appropriate demonstration sets remains a key research challenge. This paper introduces the Relevance-Diversity Enhanced Selection (RDES), an innovative approach that leverages reinforcement learning (RL) frameworks to optimize the selection of diverse reference demonstrations for tasks amenable to in-context learning (ICL), particularly text classification and reasoning, in few-shot prompting scenarios. RDES employs frameworks like Q-learning and a PPO-based variant to dynamically identify demonstrations that maximize both diversity (quantified by label distribution) and relevance to the task objective. This strategy ensures a balanced representation of reference data, leading to improved accuracy and generalization. Through extensive experiments on multiple benchmark datasets, including diverse reasoning tasks, and involving 14 closed-source and open-source LLMs, we demonstrate that RDES significantly enhances performance compared to ten established baselines. Our evaluation includes analysis of performance across varying numbers of demonstrations on selected datasets. Furthermore, we investigate incorporating Chain-of-Thought (CoT) reasoning, which further boosts predictive performance. The results highlight the potential of RL for adaptive demonstration selection and addressing challenges in ICL.