Demonstration Selection for In-Context Learning via Reinforcement Learning

TL;DR

This paper tackles the challenge of demonstration selection for in-context learning by introducing Relevance-Diversity Enhanced Selection (RDES), an RL-based framework that jointly optimizes demonstration relevance and diversity. It develops two optimization pipelines, Q-learning and a PPO-based variant, and supports Chain-of-Thought prompting to further boost reasoning performance. Through extensive experiments across four classification benchmarks and a mix of closed- and open-source LLMs, RDES demonstrates significant accuracy gains over ten baselines, with CoT-enhanced RDES/C often achieving the best results. The study highlights the potential of adaptive demonstration selection to improve generalization in data-scarce settings and discusses future work on efficiency, adaptive CoT usage, and extension to broader NLP tasks.

Abstract

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.

Figures (5)

• Figure 1: An example shows how a diversity-based demonstration method works. In this example, the diversity-based method helps the model recognize that the input text expresses a sentiment that is neither strongly positive nor negative, while the no diversity-based method may lead to an inaccurate positive classification due to its lack of varied demonstrations.
• Figure 2: The RDES framework is an adaptive RL approach for few-shot ICL demonstration selection in LLMs. It employs a RL-based agent to dynamically balance the relevance and diversity of selected examples, guided by a reward function that incorporates a label distribution diversity score. This strategy enhances classification accuracy and generalization by mitigating overfitting associated with pure similarity-based methods. The framework involves an Agent interacting with an Environment (including a Knowledge Base and the LLM) to learn an optimal selection policy.
• Figure 3: These figures illustrate the average results across closed-source/open-source models on different datasets, comparing the best results from the prompt engineering (PE) and demonstration selection (DS) methods with our proposed approach.
• Figure 4: These figures illustrate the performance of various closed-source models across different datasets, highlighting the impact of diversity mechanisms.
• Figure 5: These figures illustrate the performance of various open-source models across different datasets, highlighting the impact of diversity mechanisms.

Theorems & Definitions (2)

• Theorem 1.1: Q-Learning Convergence
• Theorem 1.2: Policy Improvement