LASeR: Learning to Adaptively Select Reward Models with Multi-Arm Bandits

We thank you for your review and suggestions. We appreciate your recognition of our paper as “well-written”, our efforts to address “a relevant, real-world” problem, and conduct “extensive” experiments. Below, we respond to your comments and suggestions in detail.

Novelty of LASeR (beyond using LinUCB)

We would like to clarify that the primary contribution of our work is the formulation of preference-based LLM training as a MAB problem over multiple RMs, a novel and previously unexplored setting. This problem is challenging because RM quality can vary across tasks and examples, yet most existing methods assume a fixed RM throughout training. Our key insight is to treat RM selection as an instance-level, contextual MAB problem, allowing the model to adaptively choose the most informative RM for each input. This framing enables us to integrate RM feedback directly into the bandit loop during preference fine-tuning, rather than using MABs only for data selection (Dwaracherla et al., 2024) or model selection at inference time without RMs (Nguyen et al., 2024; Li, 2025). To our knowledge, LASeR is the first method to do that in the context of aligning LLMs with multiple RMs.

As discussed in lines 207–210, our focus is not on proposing a new MAB algorithm. Moreover, the effectiveness of LASeR does not rely on LinUCB specifically, but rather on the exploration-exploitation trade-off enabled by the MAB framework. LASeR is agnostic to the specific choice of bandit algorithm (see Section 3, lines 207-210), and we demonstrate this empirically in Appendix C (lines 1225-1234, Table 9) by comparing LinUCB with alternative MAB methods (Exp3 and NeuralUCB). The consistent performance across algorithms supports our claim that the contribution lies in the integration of MABs into the LLM training pipeline.

Theoretical analysis

We acknowledge that theoretical analysis is a valuable direction for future work. However, developing such analysis in our setting falls outside the primary focus of this work. Theoretical modeling is complicated for this setting by several factors: the non-stationary dynamics of LLM training, the non-convex nature of the optimization landscape, and the use of reward signals from imperfect and potentially biased RMs. These challenges are shared across recent work in preference-based LLM alignment (e.g., Pang et al., 2024; Dong et al., 2025; Ye et al., 2025), which also focus on empirical performance rather than theoretical guarantees. While we recognize that RMs may be biased or misspecified, LASeR mitigates the risk of reward hacking by explicitly avoiding optimization against any single fixed RM.

Moreover, we believe our use of the negative DPO loss as the MAB reward is well-justified for the following reasons:

• Preference modeling: The DPO loss is aligned with the preference model in common RLHF setups and reflects how clearly the LLM distinguishes preferred from rejected responses. A lower DPO loss corresponds to higher preference margin, making it a reasonable proxy for reward quality. In our framework, an RM that provides more informative and consistent rankings is given a higher reward compared to a less-suited RM providing uninformative rankings (as discussed in Section 3, lines 180-187).
• No more susceptible to reward hacking than optimizing RMs directly: Unlike settings where a single RM is directly optimized (e.g., via RL or DPO with a fixed RM), LASeR avoids such issues by never optimizing against any one RM. Instead, it adaptively evaluates and switches between multiple RMs, reducing the risk of overfitting to a misspecified or hackable reward (as discussed in Section 2, lines 102-117).
• Empirical results: To support the validity of our MAB reward reward, we empirically compare our approach (using negative training loss as the MAB reward) with a variant that uses accuracy difference ($\text{Acc}(y_w) - \text{Acc}(y_l)$), which relies on ground-truth answers. As shown in Appendix C (lines 1280–1283, Table 13), using negative training loss results in stronger performance gains, suggesting it serves as a reliable proxy for RM signal quality.

I suggest the authors provide some non-parametric hypothesis tests on the results.

We note that we included an analysis in Appendix B (lines 1146–1156), where we report results over 5 random seeds to demonstrate the robustness of LASeR across datasets. To further address the reviewer’s suggestion, we extended our experiments to 10 random seeds and conducted non-parametric hypothesis testing on the reasoning benchmarks. Specifically, we test the following:

• Null hypothesis ($\mathcal H_0$): LASeR’s accuracy is equal to the competing method’s accuracy.

• Alternative hypothesis ($\mathcal H_a$): LASeR’s accuracy is greater than the competing method’s accuracy.

We use the non-parametric one-sided Wilcoxon signed-rank test to paired accuracy values across seeds. The p-value of the test between LASeR and the two best baselines is reported in the following table. At a significant level of 0.05, the table indicates that we should reject the null hypothesis, which shows that LASeR significantly outperforms both the strongest single RM and the most competitive ensemble baseline across all three tasks.

We hope our response has addressed all of your questions and will allow you to revisit your score.

References

Dwaracherla, Vikranth, et al. "Efficient exploration for llms." ICML 2024.

Nguyen, Quang H., et al. "Metallm: A high-performant and cost-efficient dynamic framework for wrapping llms." arXiv preprint arXiv:2407.10834 (2024).

Li, Yang. “LLM Bandit: Cost-Efficient LLM Generation via Preference-Conditioned Dynamic Routing.” arXiv preprint arXiv:2502.02743 (2025).

Pang, Richard Yuanzhe, et al. "Iterative reasoning preference optimization." NeurIPS 2024.

Dong, Qingxiu, et al. "Self-boosting large language models with synthetic preference data." ICLR 2025.

Ye, Ziyi, et al. "Learning LLM-as-a-judge for preference alignment." ICLR 2025.

We thank you for your detailed review and for acknowledging the “interesting” problem we solve and that the problem framing “makes sense”. Please find our responses to your comments below.

From my understanding, the authors aim to address the meta-level question of selecting the appropriate RM depending on the training stage of the model.

To clarify, LASeR does not solely select RMs based on the training stage of the LLM. Rather, our method performs instance-level RM selection using a contextual MAB that takes into account both the input prompt (via its embedding) and the current training dynamics (via the model’s evolving response quality and loss) (as discussed in Section 3, lines 195-206). We also showed that LASeR can be applied to training-free settings such as long-context understanding tasks for reranking responses (Section 4, lines 384-370).

1. It's unclear to me how the last-token embedding alone sufficiently encodes the necessary information to predict which RM to use. Could you empirically elaborate on why this is the case?

We use the last-token embedding as the context representation because in Transformer-based LLMs, this position typically aggregates information from the entire sequence. Moreover, this method is a standard and commonly used approach for extracting LLM embeddings (Wang et al., 2024, BehnamGhader et al., 2024). To empirically validate this choice, we conducted an ablation comparing three context embedding methods: first-token, average-token, and last-token embeddings. These results show that last-token embeddings outperform the alternatives across all three benchmarks, showing the effectiveness of last-token embeddings method.

2. …only one RM is selected per batch, but batches typically contain diverse samples with varying difficulties and type of problems. Consequently, averaging embeddings from multiple samples could produce a noisy estimate of context. Could the authors clarify or provide insights into why their method succeeds despite this issue?

As noted in footnote 2 of the submission, we experimented with different batch sizes to evaluate their impact. Using a batch size of 1 yielded comparable performance to a batch size of 16 but was significantly less efficient in training the LLM due to the increased computational overhead. Based on this, we opted to use a batch size of 16 for a better trade-off between performance and efficiency. We will include this discussion in the main text in our final version.

For datasets like Wildchat, which contain clearly defined and diverse categories, we structure the batches such that each batch consists of data belonging to a single category. This setup minimizes the risk of mismatches between the RM and the batch data, as each RM is evaluated on its most relevant data category. For reasoning datasets where such predefined categories do not exist, shuffling the data during training ensures diverse data and good training signal for LLM within each batch. Even if individual batches vary, the repeated exposure to diverse inputs allows the bandit to learn which RMs are more helpful overall. Over training iterations, this appears to provide a robust signal for the bandit to learn effective RM selection, as demonstrated by LASeR’s consistent gains across datasets and the convergence behavior of MAB rewards shown in Appendix B, Fig. 9.

3. How does the accumulated selection frequency for each RM behave—does it change abruptly, or is there a gradual transition between RMs as certain abilities are learned?

We would like to point to an ablation study in which we demonstrated that LASeR selected RM adjusts to the query in Appendix B (lines 1157-1176, Fig. 6 and Fig. 7). These figures show that LASeR selects different RMs depending on the task type. For example, on WildChat, creative prompts favor Olmo and Eurus, while math prompts are dominated by Qwen. Similarly, on LongBench, Zephyr is strongly preferred for summarization and few-shot learning, while Qwen is favored for multi-document QA. This task-aware behavior indicates that LASeR automatically learns to associate RM effectiveness with query type, validating the contextual adaptation capability of our approach.

In addition to these per-task dynamics, we showed in Appendix B (lines 1183-1189, Fig. 9) that the normalized MAB rewards stabilize after 6-8 training iterations. RMs like Qwen and Zephyr, which are more informative for reasoning tasks, retain higher average rewards, indicating that the bandit consistently identifies and prioritizes stronger RMs over time.

4. Additionally, could you perform an ablation study on the parameter \alpha. For example, setting \alpha represents greedy selection; does introducing exploration (\alpha > 0) significantly impact performance?

Based on this suggestion, we conduct an ablation study to assess the impact of the exploration parameter $\alpha$ in the LinUCB algorithm. The results show that for reasoning datasets, disabling exploration $\alpha$ leads to suboptimal performance, particularly in early training stages where the bandit may overfit to a suboptimal RM. In contrast, enabling exploration ($\alpha > 0$) improves performance by allowing the model to discover and leverage more informative RMs over time. We find that setting $\alpha$ in the range of 0.4 to 0.6 consistently yields the best results across tasks, offering a good balance between exploration and exploitation. This is consistent with our observations on the performance on dev set for reasoning tasks.

We hope our response has addressed all of your questions and will allow you to revisit your score.

References

Wang, Liang, et al. "Improving text embeddings with large language models." ACL 2024.

BehnamGhader, Parishad, et al. "Llm2vec: Large language models are secretly powerful text encoders." COLM 2024.

We thank you for your comments and suggestions, and we are glad that you found our method to be “well motivated” and “novel”, and our results to represent a “meaningful improvement.” Below, we provide specific responses to each of your points.

While the paper addresses an important problem…a common approach in recent literature is to use binary rewards in RL. This method has been reported to outperform reward models and should be included as an important baseline for the math tasks.

We thank the reviewer for this suggestion. We implement a Correctness Reward baseline for the reasoning datasets. Specifically, we assign a reward of 1 if the generated answer matches the correct final answer and 0 otherwise. We maintain the same iterative training setup as in our main experiments, sampling $n= 30$ responses and $P = 10$ pairs of (correct, incorrect) per query. We report the results in the table below. These results show that while the correctness-based reward performs competitively, especially on GSM8K (74.56% vs. 74.75% accuracy), LASeR outperforms it across all tasks, including commonsense reasoning tasks (StrategyQA and MMLU).

This is likely because correctness-based rewards are well-suited for tasks with clearly verifiable outputs, such as math or code, but they struggle to capture the multi-step reasoning and contextual understanding needed in factual or open-domain QA, where answers are often ambiguous or depend on background knowledge. In these cases, other forms of supervision such as process-based rewards (Ma et al., 2025) or human preference signals modeled by reward models (Peng et al., 2025) are typically more effective. LASeR is designed to leverage this by selecting among multiple RMs, allowing it to perform well across a variety of task types. As we show below, LASeR can also integrate correctness-based rewards, combining their strengths with those of preference-based models.

Would it be possible to incorporate binary rewards into the LASeR framework?

Yes, binary rewards can be incorporated into the LASeR framework. In our setup, each RM is separated, so we can simply define a new RM that assigns a reward of 1 to correct final answers and 0 to incorrect ones, and include it in the MAB’s candidate set.

Setup. To test this, we run an additional experiment on GSM8K, where we add the correctness-based RM to the existing pool of learned RMs (Olmo, Eurus, Qwen, Zephyr). To further test LASeR’s ability to adaptively select appropriate RMs, we create a mixed dataset combining GSM8K with WildChat (subsampled to 1.2K examples per category to ensure balance). Since correctness-based signals are not applicable to WildChat’s open-ended prompts, we assign a random binary reward (0 or 1) for those examples to simulate a misleading RM.

Results. LASeR achieves an accuracy of 75.06% accuracy on the GSM8K test set, improving over LASeR without correctness rewards by 0.31%. We also report the utilization rate (%) of each RM in the pool (Olmo, Eurus, Qwen, Zephyr, correctness/random) for two types of queries: GSM8K and WildChat in the table below. On GSM8K, LASeR heavily leverages the correctness-based RM (40.26%) alongside strong learned RMs like Qwen (43.58%), indicating that it effectively identifies and uses the most reliable supervision. In contrast, for WildChat, LASeR downweights the misleading correctness-based RM (2.84%) and shifts toward more suitable RMs like Zephyr and Qwen. These results demonstrate LASeR’s ability to adapt RM usage based on task domain, even in the presence of noisy or misleading supervision. This is consistent with our results on the robustness of LASeR to noisy rewards (Section 6, line 426 and Appendix C, lines 1295-1306).

LASeR introduces additional compute overhead during LLM training, and the learned multi-armed bandits are limited to a pre-selected set of reward models. What if I want to add or update a reward model? Do I have to retrain the MAB?

Although LASeR introduces additional compute overhead during LLM training, this is minimal compared to iterative DPO training. We showed in lines 1042-1047 that LASeR only adds 1.27 seconds to the training pipeline per batch. In comparison, a forward and backward pass of the LLM takes 36.39 seconds on 4 RTX A6000-48G GPUs.

Adding or updating a reward model in LASeR does not require retraining the entire MAB from scratch. LinUCB maintains independent statistics ($A_k$, $B_k$) for each RM arm. To add a new RM, we can initialize its parameters and begin exploration, allowing the bandit to learn its utility online. When adding a new RM to LASeR, only the new RM's bandit parameters need to be initialized (e.g., $A_k=I$, $b_k=0$). The existing RMs' statistics remain unchanged, meaning the MAB retains all prior learning about their utility. This allows the bandit to continue exploiting well-performing existing RMs while gradually exploring the new RM. Similarly, if an existing RM is updated (e.g., improved or retrained), the MAB can re-learn its statistics gradually through continued training.

To empirically demonstrate this, we conduct an experiment where we resume training Llama-3-8B from a saved checkpoint and add a new, stronger RM nicolinho/QRM-Llama3-8B based on the latest RewardBench leaderboard to the current pool of RMs. We observe that the MAB quickly adapts to the new RM, with training converging after just 2 additional iterations and achieving 75.58% accuracy on the GSM8K test set, improving over LASeR without QRM-Llama3-8B by 0.83%. We also run an additional experiment where we reset the MAB parameters for an existing RM (Qwen) and then resume training Llama-3-8B from a saved checkpoint. Training converges after 3 additional iterations and achieves comparable accuracy to the original checkpoint before the reset (74.72% vs. 74.75% accuracy).

Besides accuracy, we also show the utilization rates (%) of each RM before (Checkpoint), after (Transfer) adding the new RM, and restart (Restart) in the table below. These results show that LASeR can effectively integrate new RMs and adapt its selection strategy, shifting toward using stronger RMs when they offer more informative supervision, while still retaining and utilizing previously effective RMs like Qwen for reasoning tasks. Additionally, in our restart experiment, where we reset the MAB parameters of Qwen and resume training, demonstrates that LASeR can quickly relearn the utility of a strong RM, returning to similar usage levels with just a few additional iterations.

While using the negative training loss as the "reward" for the bandit makes intuitive sense, have you explored other learning signals? How do they compare?

We empirically compared our approach (using negative training loss as the MAB reward) with a variant that uses accuracy difference ($\text{Acc}(y_w) - \text{Acc}(y_l)$), which relies on ground-truth answers. As noted in Section 3, lines 148-149 and shown in Appendix C (lines 1280-1283, Table 13), using the negative training loss results in stronger performance gains, suggesting it serves as a reliable proxy for RM signal quality. The intuition is that DPO loss directly captures how clearly the model distinguishes between preferred and rejected responses, helping the model to sharpen preference rankings and increase confidence in distinguishing between chosen and rejected responses (see lines 174-187 in our submission).

We hope our response has addressed all of your questions and will allow you to revisit your score.

References

Ma, Ruotian, et al. "S $^ 2$ R: Teaching LLMs to Self-verify and Self-correct via Reinforcement Learning." ACL 2025.

Peng, Hao, et al. "Agentic reward modeling: Integrating human preferences with verifiable correctness signals for reliable reward systems." ACL 2025.