Title: LASeR: Learning to Adaptively Select Reward Models with Multi-Arm Bandits
Venue: NeurIPS 2025
Accepted: yes
Keywords: reward models

bandit

instruction following

reasoning

LLM alignment
Abstract: Reward Models (RMs) are crucial to aligning large language models (LLMs), but the degree to which an RM specialized to one task (e.g. writing) generalizes to new tasks (e.g. math) is often not known a priori, often making using only one fixed RM to train LLMs suboptimal. However, optimizing LLMs with multiple RMs simultaneously can incur a prohibitively high computational cost and lead to conflicting signals from different RMs that may degrade performance. To address these challenges, we introduce LASeR (Learning to Adaptively Select Rewards), which frames reward model selection as a multi-armed bandit problem, iteratively and efficiently training LLMs using multiple RMs by selecting the most well-suited RM for each instance. On commonsense and math reasoning tasks, we show that LASeR boosts iterative LLM training, improving the absolute average accuracy of Llama-3-8B over three datasets by $2.67$% over an ensemble of RM scores while also showing superior efficiency (e.g., a $2\times$ speedup). Moreover, on WildChat (open-ended instruction-following tasks), LASeR leads to a $72.69$% AlpacaEval win rate over the RM score ensemble baseline. Extending to long-context generation, LASeR improves by $2.96$ F1 points (avg.) on single-document QA tasks and $2.97$ F1 points on few-shot learning over the RM score ensemble baseline with best-of-$n$ sampling. We include our code in the supplementary.


Peer reviews:
comments: SUMMARY:
When finetuning LLMs to align with human preferences (RLHF), it is unclear how to pick the most suitable reward model for a particular domain. Finetuning with multiple reward models at once can be expensive (due to multiple forward passes) and noisy (due to conflicting rewards). 

This paper proposes to incorporate a contextual bandit model (LinUCB) into the RLHF process. The idea is to adaptively pick a _single_, most suitable reward model to score each prompt-answer pair. 

Experiments suggest that this approach is effective in making RLHF both more cost-efficient and effective.

STRENGTHS AND WEAKNESSES:
**Strengths**

1. Very clear and well-written paper
2. Tackling a relevant, real-world problem
3. Extensive experimental validation
4. Straightforward method---a strength since it is easy to implement


**Weaknesses**

1. Limited novelty: The proposed method is essentially just a straightforward application of LinUCB.
2. Lack of theoretical analysis. I believe that the previous point is fine if the authors provide some theoretical analyses for this specific use case. For example, the validity of using the negative log-likelihood as a reward (e.g. the LLM might be misspecified, so this reward might be biased, or easy to "hack", etc.). 
3. It is unclear whether the results are significant. The numbers are very close together, and not even error bars are provided. I suggest the authors provide some non-parametric hypothesis tests on the results.

QUESTIONS:
Could the authors please address the "Weaknesses" list above?

LIMITATIONS:
Not included in the main text. I suggest the authors fix this and be more upfront about the limitations.
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces LASER, an adaptive way to fine-tune large language models. Rather than relying on a single reward model or a static group of them, LASER treats the selection as a dynamic choice. For each new batch of prompts, it picks the single most suitable reward model for the task. This method avoids mixed signals of running multiple models simultaneously, yet still captures the distinct advantages of each one.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths:\n- The authors tackle an interesting problem, essentially addressing the high-level question, \"which model should serve as the teacher/RM at a given iteration,\" if I understand correctly.\n- Their framing of the problem as a contextual multi-armed bandit (MAB) makes sense.\n\n### Weaknesses:\n- My main concern, detailed below (in questions), is that while using a multi-armed bandit approach is intuitive, the current explanation isn't sufficient to fully justify why the proposed method works.\n\nQUESTIONS:\nFrom my understanding, the authors aim to address the meta-level question of selecting the appropriate RM depending on the training stage of the model.\n\n- **Concern 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?\n\n- **Concern 2:** My other concern relates to understanding why the method works effectively. As I understand it, 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?\n\n- **Concern 3:** I believe additional ablation studies could help clarify why the MAB approach works, addressing the concerns mentioned above. For instance, 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?\n\n- **Concern 4:** Additionally, could you perform an ablation study on the parameter $\\alpha$? For example, setting $\\alpha = 0$ represents greedy selection; does introducing exploration ($\\alpha > 0$) significantly impact performance?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to address a key limitation in LLM RL training: using one fixed reward model may not generalize well across different tasks. The authors introduced LASeR, a method that automatically selects the most suitable reward model for each training example using multi-armed bandit algorithms. More specifically, each reward model is treated as an “arm” in the bandit, and the algorithm dynamically selects which reward model to use for each batch of training examples based on contextual information and past performance. The \"reward\" for the bandit is the negative training loss, as an indicator for how well the LLM learned from the preference data generated by that reward model. The author reports consistent improvements for commonsense and math reasoning tasks, as well as open-ended instruction-following tasks.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths:\n* The paper is well motivated and addresses an important issue in RL training, the selection of reward models.\n* Framing reward model selection as a multi-armed bandit problem is elegant and novel.\n* The author conducted comprehensive experiments across multiple models and tasks, and observed meaningful improvement.\n* The paper is well-written , well-structured and easy to follow. \n\n### Weakness: \n* While the paper addresses an important problem—how to dynamically select reward models for different tasks—for math problems where the final answer can be easily verified with a rule-based verifier, 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.\n* 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?\n\nQUESTIONS:\n* Would it be possible to incorporate binary rewards into the LASeR framework? Binary rewards (i.e., whether the final answer is correct or not) tend to work very well in domains like math or code, and have been reported to perform better than learned reward models while being more compute-efficient.\n* 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}