Reward Mixology: Crafting Hybrid Signals for Reinforcement Learning Driven In-Context Learning (2025)

Abstract: In-context learning (ICL) performance heavily relies on the quality and ordering of demonstrations. Iterative selection (IS) is a promising approach to address this issue, but existing IS methods face two key challenges: the oversimplification of process reward signals that guide intermediate steps (often using single-dimensional metrics) and the lack of outcome reward signals that directly optimize final-task accuracy (relying solely on binary terminal feedback like correct/incorrect predictions). To address these issues, we propose a reinforcement learning method R-Mix which models iterative demonstration selection as a Markov Decision Process (MDP), crafting hybrid reward signals — combining outcome-based accuracy signals (i.e., outcome rewards) with process-oriented signals (i.e, process rewards) like stepwise influence and label entropy improvement. Our analysis reveals a positive but trade-off relationship between outcome rewards and process rewards, underscoring the importance of both components for effective policy optimization. We further introduce a dual-head policy architecture that explicitly decouples input-semantic relevance and label-content compatibility. Experiments across NLP benchmarks demonstrate superior performance over state-of-the-art methods, with ablation studies validating the necessity of both reward components and architectural disentanglement. Our work has deeply explored the effective potential of ICL through demonstration selection.

## introduction in-context learning (icl), an emergent ability of large language models @xcite , enables these models to achieve impressive performance on diverse tasks-such as text summarization, dialogue management, and semantic analysis-by merely incorporating demonstration input-label pairs as prompts, all without any parameter updates @xcite @xcite @xcite @xcite . demonstration selection critically impacts icl efficacy, where demonstrations can be manually designed or retrieved from datasets. research has demonstrated that both the quality and order of demonstrations significantly influence the icl capabilities of large language models @xcite . while point-wise similarity matching @xcite efficiently retrieves test-relevant demonstrations, it neglects inter-demonstration synergies. sample-then-select methods @xcite @xcite address this via permutation sampling and list-wise metrics (e.g., permutation entropy, mutual information), yet suffer from computational overhead and sampling bias. recent iterative approaches @xcite @xcite , while circumventing these issues, adopt a myopic optimization strategy-exclusively relying on oversimplified intermediate benefits (referred to as process rewards)-that fails to align with the improvement of final-task performance. this creates a critical gap that hinders unleashing the effective potential of llms for icl through demonstration selection. in response to these challenges, we model the demonstration selection process as a markov decision process (mdp) and optimize it by training a small reinforcement learning model (named r-mix) utilizes crafted hybrid reward signals. these include an outcome reward (or) aligned with finaltask accuracy, which however suffers from sparsity since it is only obtainable after the entire selection process is completed. to mitigate optimization difficulties associated with sparse rewards ( @xcite )), we introduce several rich process rewards (pr), such as improvements in influence and label entropy. we analyze the re-lationship between process rewards and outcome rewards, finding that while they exhibit a positive correlation, it's also a trade-off, highlighting the need for all types of rewards to guide the policy effectively. furthermore, since both the input and output in the demonstrations affect the final selection, we adopt a dual-head policy model to decouple their respective contributions. we validate the effectiveness of our model on three datasets and provide extensive experimental analysis to demonstrate its robustness and generalizability. ## related work demonstration selection can critically affect the performance of icl. several studies have proposed methods that leverage point-wise matching to retrieve demonstrations that are most relevant to the task at hand @xcite @xcite . uncertainty @xcite , influences @xcite , entropy @xcite , diversity @xcite , sensitivity @xcite and llm scores @xcite have also been utilized as criteria for demostration selection. while these methods are efficient, they fail to account for the interactions and synergies between different demonstrations. order sensitivity plays a crucial role in maximizing icl performance. research has shown that the sequence in which demonstrations are presented significantly influences model behavior @xcite . ordering strategies based on list-wise metrics, such as permutation entropy and mutual information, help capture the relationships between demonstrations, optimizing the presentation for the model @xcite . however, these sample-thenselect methods often suffer from computational overhead and potential biases introduced by the sampling process, raising concerns about their scalability and efficiency in practical applications. iterative methods have been proposed to optimize demonstration selection and ordering in a more dynamic and adaptive manner. these methods focus on refining the demonstration sequence iteratively @xcite @xcite . @xcite trains a ranker to iteratively approximate the optimal demonstration list. reinforcement learning (rl) has also been explored as a key technique for optimizing demonstration selection through iterative refinement with process influence improvement @xcite . while such iterative approaches have shown promise in improving performance, they are optimized using intermediate benefits (referred to as process rewards) and lack outcome-oriented reward signals to directly optimize final-task performance. additionally, their process rewards remain oversimplified. our work builds on this gap by introducing outcome rewards with process rewards that capture intermediate activation and guide the model to refine demonstration selection. ## reward mixolog (r-mix): our method to address the challenge of example selection in icl, our k-shot iterative demonstration retrieval method models the demonstration selection process as a markov decision process and optimizes the policy based on the hybrid reward, as shown in figure 1. rmix operates through three components: ## discussions and analysis 5.1 ablation studies to further evaluate the key components of r-mix, we conducted a series of ablation studies. ## conclusion we propose a reinforcement learning method r-mix for icl, modeling iterative demonstration selection as an mdp. by integrating hybrid rewards combining outcome and process signals, we ad-dress oversimplified and misguided rewards in existing methods. analysis shows a positive yet tradeoff between process and outcome rewards, highlighting their joint necessity. a dual-head policy architecture enhances performance by decoupling input relevance and label compatibility. experiments across nlp benchmarks demonstrate superior performance over state-of-the-art methods.