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

## Abstract

Abstract Reward Models (RMs) play a crucial role in aligning large language models (LLMs) with human preferences, enhancing their performance by ranking outputs during inference or iterative training. However, the degree to which an RM generalizes to new tasks is often not known a priori . For instance, some RMs may excel at scoring creative writing, while others specialize in evaluating math reasoning. Therefore, using only one fixed RM while training LLMs can be suboptimal . Moreover, optimizing LLMs with multiple RMs simultaneously can be prohibitively computationally-intensive and challenging due to conflicting signals from different RMs, potentially degrading performance. To address these challenges, we introduce LASeR ( L earning to A daptively Se lect R ewards), which iteratively trains LLMs using multiple RMs, selecting and utilizing the most well-suited RM for each instance to rank outputs and generate preference data, framed as a multi-armed bandit problem. Our empirical results on commonsense and math reasoning tasks demonstrate that LASeR can boost iterative LLM optimization by optimizing for multiple RMs, improving the absolute average accuracy of Llama-3-8B over three datasets by 2.67 % percent 2.67 2.67\% 2.67 % over training with ensemble RM scores while also showing superior training efficiency (e.g., a 2 × 2\times 2 × speedup). Moreover, on WildChat, a benchmark of instruction-following prompts in open-form generation, we find that using Llama-3-8B LASeR leads to a 71.45 % percent 71.45 71.45\% 71.45 % AlpacaEval win rate over sequentially optimizing multiple RMs. Extending to long-context generation tasks, we find that on Llama-3-8B, LASeR achieves an average improvement of 2.64 2.64 2.64 2.64 F1 points on single-document QA tasks and 2.42 2.42 2.42 2.42 F1 points on multi-document QA over random RM selection when used with best-of- n 𝑛 n italic_n sampling. Our analysis shows that LASeR is robust to noisy rewards and generalizes to multiple settings. Finally, we demonstrate that LASeR ’s RM selection changes depending on the underlying task or instance, and we verify the presence of conflicting preferences from multiple RMs, which can be mitigated using LASeR . 1 1 1 Code: https://github.com/duykhuongnguyen/LASeR-MAB

## 1 Introduction

When comparing two responses, human preferences often differ depending on factors like the underlying task, who the annotators are (Santurkar et al., ; Ahmadian et al., ), and how preferences are elicited (Bansal et al., ).
Therefore, models of preference data are also likely to differ, and might include noise as well as any biases contained in the preference data used to train them.
This can pose a problem when using such models as “reward models” (RMs) to align large language models (LLMs) to human preferences using reinforcement learning with human feedback (Christiano et al., ; Ziegler et al., ; Ouyang et al., ).
Recent work has focused on aligning LLMs through iterative training, using reward models as proxies for human judgment (Gulcehre et al., ), leveraging the LLM to act as an implicit RM or judge (Yuan et al., ; Chen et al., ), or using the gold answer to compute a reward (Pang et al., ).
Under this paradigm, there are three stages to training LLMs:
(i) *generating multiple responses* to a query from an LLM;
(ii) *scoring the responses with an RM* to create preference data with better and worse responses; and (iii) using model-generated preference data to further train the LLM.
Note that for most domains, the gold reward is not readily available, making the quality of the RM or the degree to which it reflects human preferences (i.e., the gold reward) crucial to improving LLM performance. Indeed, several prior efforts aim to train new RMs that better reflect human preferences (Lambert et al., ).

However, *selecting one reward model* to guide LLM training can be *suboptimal* for three main reasons:
(1) A single RM may not be generalized to heterogeneous sets of examples.
RMs are typically designed to reflect specific objectives and may be trained on offline preference datasets.
Thus, an RM that performs well on one dataset or domain may not generalize effectively to others, leading to misaligned outputs across different tasks or domains (Kirk et al., ; Chen et al., ; Casper et al., ; Gao et al., ). For instance, creativity plays a key role in evaluating the quality of a story, whereas correctness is more important in scoring math solutions.
(2) RM performance leaderboards (e.g., Lambert et al. ()) that rely on human-annotated preferences can have unreliable rankings due to the presence of incorrect and ambiguous preferences (Yu et al., ; Hejna et al., ).
(3) Lastly, over-optimization on a particular RM can lead to *reward hacking* (Skalse et al., ; Rafailov et al., ), resulting in minimal gain or even drops in downstream performance.

To mitigate these issues, a prevalent approach is to ensemble multiple reward models (Coste et al., ; Eisenstein et al., ; Zhang et al., ; Ramé et al., ).
However, these methods also come with significant challenges: as RMs are typically based on LLMs, training with multiple RMs often requires loading and managing several large models simultaneously, which can be computationally expensive, becoming infeasible as models increase in size.
Moreover, aggregating multiple RM scores together is susceptible to
noisy rewards or conflicting preferences from RMs, especially RMs that are not well-suited for the specific task (Rita et al., ).
This in turn, can degrade the quality of the preference data, leading to low-quality updates during training (Wang et al., ).
Finally, manually selecting a subset of RMs to combine is a labor-intensive process that involves training many different variants on a combinatorially large set of RM groupings.
This underscores the need for more efficient methods that efficiently and robustly optimize LLMs using multiple RMs.

In this work, we introduce Learning to Adaptively Select Rewards (LASeR), that, given a set of RMs, adaptively and efficiently *selects* a suitable RM for each instance by casting selection as a *multi-armed bandit* problem (Vermorel & Mohri, ; Audibert et al., ).
Specifically, during training, the RM (arm) is chosen dynamically based on contextual information about the model’s performance and past interactions.
The LLM is then fine-tuned based on the RM-annotated data, and the bandit’s parameters are updated accordingly to reflect the performance of the LLM after training on preference data annotated using selected RM (see Fig. 1).
By design, LASeR’s adaptive instance-level RM selection (c.f. Sec. 3) addresses the three shortcomings of choosing one reward model (lack of generalization, unreliable rankings, and over-optimization) and outperforms using the same RM across all instances, yielding higher downstream performance and better generalization.
Moreover, unlike previous multi-RM methods that require simultaneously loading and running multiple RMs (Ramé et al., ; Coste et al., ), our method selects one reward model at a time (Sec. 4). This makes the training more efficient and improves performance by allowing the model to adaptively focus on the most suitable RM for each specific instance or phase of training.

Empirically, we demonstrate the effectiveness of LASeR for iteratively training LLMs using multiple RMs on three broad domains: reasoning, instruction-following in text generation, and long-context understanding (Sec. 4.2).
We show that on reasoning benchmarks such as StrategyQA (Geva et al., ) (testing commonsense reasoning) and GSM8K (Cobbe et al., ) (testing math reasoning), LASeR with Llama-3-8B outperforms training with the single best RM by $1.45\%$, and jointly optimizing with ensemble (average) RM scores by $2.67\%$ in absolute accuracy averaged over datasets.
LASeR is also effective on general instruction-following: we show that using LASeR with four strong 7B RMs from RewardBench to finetune Llama-3-8B on a subset of WildChat (Zhao et al., ) beats LLMs trained with the best RM in the ensemble and a sequential baseline, with $56.34\%$ and $71.45\%$ win rates (respectively) on length-controlled AlpacaEval (Dubois et al., ). Moreover, our results show the effectiveness of LASeR’s RM selection strategy at inference-time on long-form generation tasks; on LongBench (Bai et al., ), we find LASeR beats random RM selection baseline by $2.64$ F1 points on single-document QA tasks and $2.42$ F1 points on multi-document QA
when using best-of-$n$ sampling for Llama-3-8B.
Our analysis reveals that LASeR is more efficient than sequential multi-RM and RM ensemble baselines in terms of training time (wall-clock hours) by a factor of $3\times$, and $2\times$, respectively while being more robust to noisy rewards and conflicting preferences from multiple RMs (Sec. 5). Finally, we demonstrate that LASeR effectively selects RMs based on the underlying instance and generalizes to multiple settings, including out-of-distribution datasets, different training loss functions, etc.

## 2 Related work

#### Multiple Reward Ensembles.

Training large language models (LLMs) with multiple reward functions is an emerging research area focused on aligning model outputs with complex objectives that require diverse evaluation metrics. One prior approach involves using ensembles of multiple rewards (Ramé et al., ; Wu et al., ; Coste et al., ; Zhang et al., ; Wang et al., ; Jang et al., ; Eisenstein et al., ).
Unlike these methods, which train multiple RMs or use multiple RMs during LLM training, LASeR selects only a single pretrained RM at each LLM training step.
Not only is optimizing for one RM at a time more efficient, but it also avoids the problem of having conflicting rewards (Rita et al., ).
In Sec. 4.2, we show that LASeR outperforms an offline variant of optimizing LLMs with an ensemble of RMs (consistent with the mean optimization setting from Coste et al. ()).
Another line of work uses mixture-of-experts (MoE) techniques for training an RM from multiple interpretable objectives (Wang et al., ) or jointly training task-specific RMs and a sparse MoE router (Quan, ).
Instead of relying on static datasets with human-annotated attributes for RM training (as Wang et al. () does), we employ existing RMs to train LLMs on its own generations, a strategy that has been shown to be more effective (Ivison et al., ).
Unlike Quan (), who jointly train multiple RMs and sparse MoE router (requiring loading all models at once), LASeR implicitly and efficiently learns a sparse router and does not need to train RMs.
LASeR performs RM selection via multi-armed bandits over existing off-the-shelf RMs with strong performance on leaderboards (Lambert et al., ).

#### Multi-Armed Bandits (MABs).

There is a long-standing history of using multi-armed bandit algorithms for diverse applications in machine learning spanning online advertising (Chen et al., ), recommendation systems (Li et al., ), hyperparameter optimization (Li et al., ), curriculum learning (Graves et al., ), with some recent work at the intersection of MABs and language models (Pasunuru et al., ; Krishnamurthy et al., ; Dwaracherla et al., , inter alia).
In the realm of RLHF specifically, Dwaracherla et al. () use double Thomson Sampling (Wu & Liu, ) to select which of the sampled responses should be annotated and paired using a single fixed RM, improving LLM performance.
In contrast, LASeR first selects which RM (i.e., model of preferences) should be used to annotate LLM’s responses to a query and then creates multiple preference pairs from these responses.
Pasunuru et al. () optimize text generation models for different evaluation metrics such as ROUGE-L and BLEU via policy gradients (Williams, ) over existing (static) question and data-to-text generation datasets.
In contrast, LASeR adopts an iterative training recipe and improves downstream generation performance across a wide range of tasks from instruction-following to math and commonsense reasoning by selecting relevant RMs and scoring the LLM’s own generated responses without access to true rewards, i.e., gold labels, and without optimizing for the downstream evaluation metric.

#### Iterative LLM Training.

Recent works on training LLMs incorporate reinforcement learning with human feedback (RLHF) due to its effectiveness in improving instruction following the ability of LLMs over their pretrained counterparts (Ouyang et al., ; Bai et al., ; Touvron et al., ). The standard RLHF training pipeline, which comprises finetuning the LLM on static datasets with human-annotated preferences, is bottlenecked by the size and quality of annotated preference data as well as the limited effectiveness of off-policy optimization (Xu et al., ; Xiong et al., ; Yuan et al., ; Guo et al., ).
To remedy this, recent work focuses on training LLMs iteratively,
scoring the LLM’s own generations to create feedback data for RLHF.
This line of work obtains scores either from gold labels (Singh et al., ; Pang et al., ), from a single RM (Gulcehre et al., ), or from the generating LLM (Yuan et al., ; Chen et al., ).
In contrast, we take advantage of the abundance of publicly-available RMs and growing interest in developing RMs (Lambert et al., ) by using *multiple* RMs.
This has several advantages over past work: it works in cases where gold labels are not available (e.g., generation tasks), deals with the three issues associated with using a single RM (unreliable rankings, lack of generalizability, and over-optimization), reduces the burden on the user to pick the right RM, and avoids problems stemming from the inability of LLMs to judge their own responses for certain domains (Huang et al., ).

## 3 LASeR : L earning to A daptively Se lect R ewards

In this section, we describe LASeR in detail.
First, we expand on the training pipeline (in one iteration) with a general reward function (Sec. 3.1).
Then, in Sec. 3.2, we describe how LASeR dynamically selects an RM from a set of multiple RMs using MAB algorithms, i.e., how we dynamically assign the reward function for a given instance or batch of instances.
Finally, in Sec. 3.3, we describe the overall training setup across iterations and specifically how we update the parameters of the MAB at the end of each iteration. A detailed illustration of LASeR is shown in Fig. 1.

Figure: Figure 1: Overview of LASeR. Given the query, the multi-armed bandit selects an RM depending on the underlying query and the bandit’s parameters (based on the usage of each RM and the expected MAB reward). At iteration $m$, the LLM generates multiple responses that are scored based on the selected RM for that query, scoring each response. These responses are ranked into preference pairs, which are then used to fine-tune the model. The same train loss $\mathcal{L}^{m}$ is used to update the parameters of the LLM as well as the MAB for the next iteration, making the entire pipeline iterative.
Refer to caption: x1.png

### 3.1 Training LLMs using a Reward Function

LASeR involves training with multiple RMs using a multi-arm bandit (MAB), which selects one model at a time.
Therefore, we first describe how we train LLMs with generated data assuming a single RM; this corresponds to the top-right in Fig. 1 (in blue).

#### Notation.

Following Yuan et al. () and Pang et al. (), we adopt an iterative training pipeline to finetune the LLM for $M$ iterations.
Let $\pi_{m}$ be the LLM at iteration $m$; we assume that we start from an initial pretrained model $\pi_{0}$.
Let $\mathcal{D}=\{x_{1},x_{2},\dots,x_{N}\}$ represent the training inputs, where $x_{i}$ is an input query or prompt. Corresponding to each input query $x_{i}$, we sample a set of $n$ responses from the LLM at the current $m^{\text{th}}$ iteration as $\mathbf{y}_{i}=\{y^{1}_{i},y^{2}_{i},\dots,y^{n}_{i}\}\sim\pi_{m}(y|x_{i})$. Let $R^{\star}\!:\!(y_{i}^{j}|x_{i})\rightarrow\mathbb{R}$ be a reward function that can score an LLM generated response $y_{i}^{j}$ to a query $x_{i}$ based on how well it aligns with specific task objectives or instructions. Note that $R^{\star}(.)$ can be any reward function and may correspond to a single RM, one of the multiple RMs selected by the MAB (as in our case), or even the true reward.

#### Generating Preference Pairs.

We evaluate each response $y_{i}^{j}$ using the reward function $R^{\star}(y_{i}^{j}|x_{i})$.
By comparing the rewards assigned to different responses, we can form preference pairs $(y_{i}^{w},y_{i}^{l})$, where $y_{i}^{w}$ is preferred over $y_{i}^{l}$ if $R^{\star}(y_{i}^{w}|x_{i})>R^{\star}(y_{i}^{l}|x_{i})$, thereby building a preference dataset:(^2^22Following Pang et al. (), in our experiments, we randomly sample $P\!=\!10$ pairs corresponding to each prompt $x_{i}$. For brevity, we omit this in the notation of $\mathcal{D}_{\mathrm{pref}}$; but in our setting $|\mathcal{D}_{\mathrm{pref}}|=P\times|\mathcal{D}|$.)

$$ $\mathcal{D}_{\mathrm{pref}}=\{(x_{i},y_{i}^{w},y_{i}^{l})\mid x_{i}\in\mathcal {D},R^{\star}(y_{i}^{w})>R^{\star}(y_{i}^{l})\}.$ $$

#### Training Loss Function ( ℒ m superscript ℒ 𝑚 \mathcal{L}^{m} caligraphic_L start_POSTSUPERSCRIPT italic_m end_POSTSUPERSCRIPT ).

In each iteration, we fine-tune the model using the generated preference dataset $\mathcal{D}_{\mathrm{pref}}$, resulting in $M$ models $\pi_{1},\pi_{2},\dots,\pi_{M}$.
Specifically, we update the model using the DPO loss (Rafailov et al., ) for learning from the preference pairs. In this work, we use the following loss functions for training the LLM at iteration $m$:

$$ $\displaystyle\begin{split}\mathcal{L}^{m}_{\textsc{dpo}}(\pi_{m})\!&=\!- \mathbb{E}_{(x_{i},y_{i}^{w},y_{i}^{l})\sim\mathcal{D}_{\mathrm{pref}}}\left[ \log\sigma\left(\beta\log\frac{\pi_{m}(y_{i}^{w}\!\mid\!x)}{\pi_{m-1}(y_{i}^{w }\mid x)}-\beta\log\frac{\pi_{m}(y_{i}^{l}\!\mid\!x)}{\pi_{\scriptstyle m-1}(y _{i}^{l}\!\mid x)}\right)\right]\\ \mathcal{L}^{m}_{\textsc{nll}}(\pi_{m})\!&=\!-\mathbb{E}_{(x_{i},y_{i}^{w}) \sim\mathcal{D}_{\mathrm{pref}}}\left[\frac{\log\pi_{m}(y_{i}^{w}\!\mid\!x_{i} )}{|y_{i}^{w}|}\right],\end{split}$ (1) $$

where $\pi_{m}$, and $\pi_{m-1}$ denotes the LLM in the current iteration $m$ and the previous iteration $m-1$ (used as the reference model in DPO loss).
Following Yuan et al. (), we use the standard DPO loss for instruction-finetuning.
Following Pang et al. (), we use the NLL loss on the preferred responses as an additional regularizer for reasoning tasks, i.e., $\mathcal{L}^{m}\!=\!\mathcal{L}^{m}_{\textsc{dpo}}+\mathcal{L}^{m}_{\textsc{
nll}}$. In Appendix C, we show that LASeR outperforms baselines irrespective of the choice of the loss function $\mathcal{L}^{m}$.

### 3.2 Bandit Algorithms for Adaptive RM Selection

Sec. 3.1 described the data creation and LLM training procedure for our method when using a general RM (Fig. 1; top-right), which trains the LLM for a single mini-batch.
Here, we describe the process by which we adaptively *select an RM for each batch of queries* using bandit algorithms (shown in Fig. 1-left, in yellow) and update the parameters of the bandit (more details in Sec. A.2).

#### Background: Multi-Armed Bandits.

The multi-armed bandit (MAB) problem addresses the challenge of balancing exploration and exploitation in sequential decision-making (Vermorel & Mohri, ; Audibert et al., ).
The goal is to maximize cumulative MAB rewards over time by selecting arms that yield the highest MAB rewards.(^3^33In order to distinguish between rewards or scores generated by RMs and the rewards used in MAB literature, we refer to the latter as “MAB rewards”.)
A decision-making agent faces a trade-off: whether to exploit the arm with the highest known MAB reward based on past observations or explore other, less familiar arms to gather more information that might lead to even better rewards in the future.
In a contextual MAB setting, the agent is also provided with additional information in the form of a context, such as current state and input, to help inform arm selection accordingly.

#### RM selection in LASeR .

LASeR uses mini-batch training for each iteration, i.e., we use MABs to select a single RM for a batch of prompts $\mathbf{x}_{t}$ for $t^{\text{th}}$ batch or training step (total of $T$ steps/batches in each iteration).(^4^44Note that LASeR can switch between RMs at the instance level if the batch size is set to 1; however, for the sake of efficiency, we batch instances together both for LASeR and the baselines, as this reduces the computational overhead associated with loading RMs onto the GPU.)
Let the set of $K$ reward models (or arms) denoted by $\mathcal{R}=\{R_{1},R_{2},\dots,R_{K}\}$, where each $R_{k}$ corresponds to a different RM.
We employ LinUCB (Li et al., ), a contextual bandit algorithm for the arm or RM selection. Li et al. () assume that the MAB reward – in our case, the cumulative train loss function on the batch ($\mathcal{L}^{m}$) at the given iteration $m$ – can be modeled linearly as a function of context features and computes the expected MAB reward of each arm with an upper confidence bound to ensure exploration (Garivier & Moulines, ; ).
In each step $t$, we have a batch of input prompts $\mathbf{x}_{t}$ for which we compute sentence embeddings, using the policy model $\pi_{m}$, and use the mean sentence embedding as the context $c(t)$ to the MAB, i.e., $c(t)=\sum_{x\in\mathbf{x}_{t}}e_{m}(x)/|\mathbf{x}_{t}|$, where $e_{m}(.)\in\mathbb{R}^{d}$ yields the sentence embedding from the model $\pi_{m}$. We calculate the embedding for a prompt as the last token embedding from model $\pi_{m}$.
The learned parameters of the LinUCB bandit include $\hat{\theta}_{k}\in\mathbb{R}^{d}$ which represents the learned weights for the features of each reward model and $A_{k}\in\mathbb{R}^{d\times d}$ (a covariance matrix) and a bias vector $b_{k}\in\mathbb{R}^{d}$ corresponding to each arm or RM $R_{k}$. We initialize the parameters for LinUCB by randomly initializing $b_{k}$ and setting parameter $A_{k}$ to the identity matrix.
Based on the LinUCB algorithm, for each batch, the selected RM $R^{\star}_{t}$ is determined by:

$$ $R^{\star}_{t}=\arg\max_{k\in[1,K]}\left(c(t)^{\top}\hat{\theta}_{k}+\alpha \sqrt{c(t)^{\top}A_{k}^{-1}c(t)}\right),$ (2) $$

where $\hat{\theta}_{k}=A_{k}^{-1}b_{k}$. $A_{k}$ and $b_{k}$ are updated based on the MAB reward for each RM, which corresponds to the normalized negative train loss $-\hat{\mathcal{L}}^{m}$ (described in detail in Sec. A.2):

$$ $A_{k}\leftarrow A_{k}+c(t)c(t)^{\top};b_{k}\leftarrow b_{k}-\hat{\mathcal{L}}^ {m}(t)c(t).$ (3) $$

### 3.3 LLM and Bandit Training in LASeR

A key aspect of our approach is the generation of new preference training data in each iteration using the generations of the LLM itself and the RM selected by the MAB.
Fig. 1 presents our training procedure, broken down into three stages: (i) the MAB selects an RM $R^{\star}_{t}$ (see Sec. 3.2; Fig. 1 left), generating preference pairs by scoring the LLM’s outputs using the RM (Fig. 1 (top-right)), and parameter updates to the LLM and MAB.
In this way, the model continuously learns from its own outputs, guided by the selected reward model.
After each LLM train step (i.e., one mini-batch), the MAB’s parameters are updated based on the observed MAB reward, i.e., how much the LLM’s loss decreased from using the selected RM.
In the case of LinUCB, this involves updating the parameter estimates $b_{k},A_{k}$ (see Fig. 1; bottom in green). This entire process – selection of reward models, generation of new supervision data, fine-tuning, and bandit updates – repeats for a total of $M$ iterations (summarized in Algorithm 1).

#### LASeR with Best-of- n 𝑛 n italic_n Sampling.

For settings where finetuning the LLM is not desirable or feasible, LASeR can also be applied to learn the MAB parameters without training the LLM.
Rather than fine-tuning the model with preference data, we employ best-of-$n$ sampling (Lightman et al., ; Sun et al., ), where multiple responses are generated, and the best one is selected based on the RM. The bandit parameters are then updated using equation (3), with the MAB reward calculated as the negative normalized NLL loss on the train data.
This updated bandit can subsequently be used for inference on the test set.
This approach is particularly useful for tasks such as long-context understanding, where training would be too computationally intensive (example setting in Sec. 4.2).

## 4 Experiments and Results

### 4.1 Experimental Setup

#### Models.

We conduct our experiments on the Llama-3-8B base (AI@Meta, ) and the Mistral-7b-v3-instruct (Jiang et al., ) models. For training, all models are fine-tuned
using Low-Rank Adaptation (LoRA) (Hu et al., ) for efficiency. For both training and inference, we do 0-shot prompting and sampling $n\!=\!30$ responses per prompt with temperature $0.8$ (see Sec. A.1).

#### Reward models.

We select $K=4$ strong 7B RMs from RewardBench (Lambert et al., ), which include Zephyr-7B-Alpha, Qwen1.5-7B-Chat, Eurus-7B-KTO, and OLMo-7B-Instruct. Following the pipeline outlined in Lambert et al. (), for these models, we compute the reward for each response as the log likelihood of the RM for that response (details in Sec. A.1).

#### Datasets and Metrics.

Our experiments cover a range of tasks and datasets (see Sec. A.1):

- •
Reasoning: Evaluating reasoning abilities is crucial for testing the model’s capacity to handle complex, multi-step tasks and has presented a challenge to iterative preference optimization methods (Yuan et al., ; Chen et al., ).
We train and evaluate on
StrategyQA (Geva et al., ), MMLU (Hendrycks et al., ; ), and GSM8K (Cobbe et al., ).
- •
Instruction-Following: We further evaluate our method on heterogeneous tasks without gold labels. We use user prompts from WildChat dataset (Zhao et al., ), which contains a collection of natural user-chatbot interactions.
This dataset has five primary categories of instruction-following prompts: creative writing, analysis, coding, factual information, and math reasoning.
Due to computational constraints, we randomly subsample 5K prompts from each category for model training.
We compare models trained with LASeR against baselines (described below) using length-controlled AlpacaEval (Dubois et al., ) that pairs responses from two different LLMs and uses GPT-4 as a judge to pick the winner, accounting for the length of both responses.
- •
Long-Context Understanding: As finetuning LLMs on long-context inputs is computationally intensive, we demonstrate the effectiveness of LASeR using Best-of-$n$ sampling on
LongBench (Bai et al., ) which consists of multiple tasks, such as single-document QA, multi-document QA, summarization, and few-shot learning.
For the QA and few-shot learning tasks, we measure performance with F1 score, while for summarization we use Rouge-L (Lin, ).

#### Baselines.

We compare our models against the following baselines:

- •
Best RM: From our collection of RMs, we pick the RM that corresponds to the best overall score on RewardBench (Lambert et al., ): Zephyr-7B-Alpha.
We use this *single* RM during training (c.f. Sec. 3.1).
This baseline reflects the performance gain a user could expect when selecting the best RM from a leaderboard without knowing *a priori* how it generalizes to a particular domain.
- •
Avg. RM: Here, we perform single RM training over all the RMs in our collection and report the average performance. A comparison with this baseline represents an expected gain from a randomly-picked RM
from a leaderboard.
- •
Random RM Selection: In this baseline, we randomly sample a single RM from the set of RMs (from a uniform distribution) for each training batch in every iteration. This comparison demonstrates whether, without prior knowledge of which RM is best suited for a downstream task, LASeR can outperform random sampling of RMs during training.
- •
Sequential RM Selection: In training, this method explores different RM sequentially and based on a set order in each iteration to examine their impact on model training, demonstrating that, instead of optimizing with all RMs, LASeR can adaptively select the best RM for each batch.
- •
Offline RM Ensemble: We generate multiple responses for each query, which are scored (offline) using each RM, and the preference dataset is created by averaging the scores across all RMs (following Coste et al. ()); thus, comparing LASeR with using all RMs simultaneously.

Conceptually, the best RM baseline serves as an “exploit-only” setting that exploits the best available RM based on RewardBench.
On the other hand, the random and sequential selection baselines are “explore-only” in that they pick a new RM either randomly or via a predefined sequence irrespective of the performance of each arm (RM).
In our experiments, we train each model to convergence. In particular, LASeR, “Best RM”, “Avg. RM”, and RM ensemble baselines were trained for 10 iterations.
For both the sequential and random RM selection, we found LLM training took longer to converge, and consequently, the model was trained for 25 iterations.
The number of iterations for each approach was chosen based on stabilization in performance on the dev set (see Sec. A.1).

**Table 1: Performance on reasoning benchmarks. The baselines also include supervised fine-tuning on human-written responses (SFT) as a reference for performance without preference optimization. The highest accuracy is shown in bold, and the second-highest accuracy is underlined. Across both Llama-3-8B and Mistral-7B models, LASeR yields the highest accuracy for each task.**
| Method | Llama-3-8B | Mistral-7B |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| StrategyQA | GSM8K | MMLU | Avg. | StrategyQA | GSM8K | MMLU | Avg. |  |
| SFT | 80.41 | 69.43 | 65.66 | 71.83 | 68.57 | 43.62 | 56.48 | 56.22 |
| Best RM | 84.29 | 73.16 | 67.15 | 74.87 | 70.06 | 45.81 | 62.04 | 59.30 |
| Avg. RM | 82.62 | 71.57 | 66.67 | 73.62 | 69.62 | 45.47 | 59.58 | 58.22 |
| Offline RM Ensemble | 82.96 | 70.94 | 67.04 | 73.65 | 68.89 | 44.53 | 58.23 | 57.22 |
| Random RM Selection | 84.37 | 71.99 | 67.85 | 74.74 | 69.97 | 45.12 | 59.88 | 58.32 |
| Sequential RM Selection | 83.90 | 72.94 | 68.02 | 74.95 | 70.59 | 46.11 | 59.66 | 58.79 |
| LASeR (Ours) | 85.96 | 74.75 | 68.24 | 76.32 | 73.06 | 46.89 | 62.27 | 60.74 |

### 4.2 Main Results

#### LASeR achieves the best average accuracy on reasoning tasks.

Table 1 shows our method consistently outperforms the baselines across multiple benchmarks, particularly in the StrategyQA and GSM8K tasks.
For example, using the Llama-3-8B model, LASeR yields the highest reasoning accuracy average across all tasks, with improvements of approximately $2\%$ absolute accuracy over the sequential baseline on both GSM8K and StrategyQA.
On the Mistral-7B model, LASeR also improves average accuracy by roughly $2\%$ over the sequential baseline. Additionally, our method outperforms the best RM (based on RewardBench) baseline in average accuracy by $1.45\%$ and $1.44\%$ absolute accuracy with Llama-3-8B and Mistral-7B, respectively.
In cases where the best RM is not known beforehand, LASeR surpasses the performance of the average RM by $2.7\%$ on Llama-3-8B and using RM ensemble for each instance by $2.67\%$ and $3.52\%$ on Llama-3-8B and Mistral-7B, respectively (in accuracy averaged over the three tasks).
Overall, LASeR provides *consistent* results while the underlined second-place models show *inconsistent* performance across datasets and models.
These results also emphasize the benefit of LASeR, as it eliminates the need to choose a different RM in advance, such as random, sequential selection, or the best RM.

#### LASeR beats baselines at instruction-following.

Often, LLMs are used by large numbers of people with a diverse set of queries, goals, and intentions, and their preferences vary based on the underlying query. To demonstrate the effectiveness of LASeR in such settings, we compare the instruction-following performance, i.e., AlpacaEval win rates, of LLMs trained using LASeR with the baselines using WildChat (see Sec. 4.1).
Specifically, with Llama-3-8B, LASeR achieves substantial win rates compared to the random and sequential baselines, with $78.33\%$, and $71.45\%$, respectively.
We also outperform training with the single best RM (per RewardBench) by a $56.34\%$ win rate.
We hypothesize the lower win rate of the baselines stems from the inability of these baselines to deal with conflicting signals from multiple RMs (see Fig. 5 for further analysis).
The results are also applicable to the Mistral-7B model, which achieved a win rate of $58.72\%$ against the best-RM baseline and a win rate of $63.72\%$ against the sequential selection baseline. Overall, these results highlight that LASeR excels in tasks without gold labels and performs consistently well at following instructions across various user queries, showcasing its adaptability to diverse tasks.

Figure: Figure 2: Length-controlled AlpacaEval win rates comparing LASeR against baselines on the instruction-following tasks on WildChat using Llama-3-8B.
Refer to caption: x2.png

#### LASeR ’s adaptive RM selection helps long-context understanding.

Given the cost of training long-context systems,
for LongBench (Bai et al., ), rather than finetuning a model using RMs, we employ the selected RM to rerank generation in Best-of-$n$ sampling (see Sec. 3.3).
In Table 2, we observe that LASeR consistently outperforms the baselines across tasks on Llama-3-8B and Mistral-7B except on summarization, where we achieve comparable performance. LASeR improves single-doc QA by $3.58$ F1 points over the base Llama-3-8B model and $2.64$ F1 points over random RM selection.
On multi-doc QA, our approach improves performance over the Llama-3-8B and Mistral-7B models by $\approx\!4$ F1 points each, beating out random RM selection by $2.42$ F1 points on Mistral-7B.
Lastly, on few-shot learning tasks, LASeR provides over $3$ points gain in F1 compared to the base model for both Llama-3-8B and Mistral-7B, surpassing the average RM performance by up to $2.4$ F1 points (on Llama-3-8B) and demonstrating its effectiveness across tasks.

**Table 2: LASeR outperforms baselines in long-context understanding tasks with Llama-3-8B and Mistral-7B. Sequential RM selection is not applicable in this setting as only inference is conducted. For QA and few-shot learning tasks, we report F1 scores, and for summarization we report Rouge-L.**
| Method | Single-Doc QA | Multi-Doc QA | Summarization | Few-shot Learning |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3 | Mistral | Llama-3 | Mistral | Llama-3 | Mistral | Llama-3 | Mistral |  |
| Base model | 33.89 | 26.01 | 32.96 | 24.06 | 29.54 | 26.47 | 70.23 | 64.93 |
| Best RM | 35.12 | 28.93 | 35.83 | 27.93 | 34.26 | 30.42 | 71.79 | 68.34 |
| Avg. RM | 34.96 | 27.47 | 35.39 | 26.38 | 31.71 | 27.34 | 71.28 | 66.94 |
| Random RM Selection | 34.83 | 27.44 | 35.19 | 25.38 | 31.57 | 27.19 | 70.91 | 66.72 |
| LASeR (Ours) | 37.47 | 29.14 | 36.94 | 27.80 | 34.13 | 30.08 | 73.31 | 68.47 |

## 5 Additional Analysis of LASeR

#### Robustness to Noisy Rewards.

To examine the robustness of our method in the presence of noisy or irrelevant rewards, we conduct the following analysis using Llama-3-8B on GSM8K. We add varying amounts of Gaussian noise $\sigma$ to the rewards generated by RMs sampled from the distribution $\mathcal{N}(0,\sigma I)$, where $I$ is the identity matrix, to simulate noisy rewards when using RMs in out-of-distribution settings.
In addition to LASeR using the LinUCB algorithm (c.f. Sec. 3.2),

Figure: Figure 3: Impact of the magnitude of Gaussian noise on the accuracy of LASeR and sequential baseline on RewardBench.
Refer to caption: x3.png

we also use Exp3 (Auer et al., ) designed for adversarial bandit settings.
In Fig. 3, we find that even as the degree of noise in RM scores increases (from $\sigma\!=\!$ 0.1 to 0.4), LASeR’s selection strategy continues to perform robustly, mitigating the effects of noise compared to the sequential baseline.
Specifically, LASeR has an average performance drop of only $0.55\%$ accuracy at a noise level of $\sigma\!=\!0.3$, whereas the sequential baseline suffers a $1.6\%$ accuracy drop (3 times as much) under the same conditions.
Furthermore, LASeR using Exp3, the most noise-robust method, maintains consistent performance, with only a $0.26\%$ accuracy drop.

Figure: Figure 4: Training efficiency of LASeR (LinUCB, and Exp3) vs. sequential selection and RM ensemble on GSM8K.
Refer to caption: x4.png

#### Training Efficiency of LASeR .

As we noted in Sec. 4.1, standard multi-reward baselines
such as sequential and random RM selection are slow to converge. We now concretely
show the accuracy-training time tradeoff in Fig. 4 by comparing the GSM8K performance of training with LASeR (with LinUCB and Exp3) with sequential RM selection and the offline RM ensemble, along with the corresponding wall clock training time.(^5^55Wall clock time is measured as the training time of a model (hours), keeping compute resource consistent.) We find that sequentially optimizing over each RM performs the worst in terms of training time ($3\times$ of LASeR) while offline RM ensemble has the worst accuracy (and takes $2\times$ the training time of LASeR). Among the MAB algorithms, LinUCB achieves the best balance. Exp3, while slower than LinUCB, reaches the highest accuracy, demonstrating its strength in effectively managing exploration across RMs.

Figure: Figure 5: Agreement in preference rankings between RMs on output pairs generated by Llama-3-8B on MMLU (left) and WildChat (right).
Refer to caption: x5.png

#### Presence of Conflicting Signals among RMs.

In Sec. 4.2, we find that LASeR consistently outperforms other multi-reward baselines across a wide variety of tasks. We attribute some of these performance gains to the inability of the multi-reward baseline to handle conflicting signals, resulting in subpar training data from multiple RMs. To study this, we
sample pairs of outputs generated by Llama-3-8B on MMLU as well as WildChat and evaluate the consistency of response preferences measured by multiple RMs.
Since pair-wise preferences are binary, we compute F1 to measure consistency with one RM’s preferences serving as the reference. Fig. 5 (left on MMLU) reveals that Qwen and Zephyr have the highest agreement rate at $0.77$, while Qwen’s agreement with Eurus and Olmo is lower at $0.58$ and $0.43$, respectively. Zephyr also shows low agreement with Olmo at $0.57$. This is expected as Qwen and Zephyr are the top-performing models in reasoning according to RewardBench, while Olmo ranks the lowest in reasoning ability among the four models. We observe similar trends in agreement across RMs on WildChat (albeit with different agreement scores), which contains user queries asked LLMs in the wild; see Fig. 5 (right). It appears that for more heterogeneous datasets with more categories, the level of disagreement among RMs (or conflict) increases. This also highlights the benefit of LASeR over multi-RM baselines that do not adequately address conflicts in RMs.

Figure: Figure 6: Utilization (%) of each RM on instruction-following queries from WildChat. The bars are arranged based on their overall scores on RewardBench, from lowest to highest. LASeR dynamically selects from different RMs depending on the nature of the underlying instance.
Refer to caption: x6.png

#### LASeR ’s selected RM adjusts to the Query.

Fig. 6 shows the relative utilization rates of each arm (i.e., RM) of the bandit on WildChat.
We observed vastly different RM utilization rates depending on the underlying query within the *same dataset*. On queries requiring creativity in LLM responses, we find that Olmo and Eurus RMs are utilized about $20\%$ more often than Qwen RM, despite Qwen RM being ranked higher on RewardBench. This can be explained by the fact that the Qwen RM largely underperforms on the “chat” subsplit of RewardBench (behind Olmo and Eurus by nearly 40 points in chat score). On the other hand, Qwen RM is used roughly half the time for user prompts involving math, while Olmo and Eurus are used sparingly. This is consistent with Qwen RM’s ranking on the “reasoning” split of RewardBench, outperforming Eurus and Olmo RMs by 15-20 points. Note that LASeR *automatically* deduces these relative rankings of RMs and uses them depending on the underlying query *without having access to the RewardBench leaderboard*.
Therefore, RM utilization of LASeR can serve as an analysis tool for future work when assessing performance on untested domains.

#### Generalization ability of LASeR .

While recent work focuses on building RMs that reflect preferences across domains, a large body of prior work developed a suite of evaluation metrics catered to specific domains such as reasoning (Golovneva et al., ; Prasad et al., ). In Table 5 (Appendix C), we show that Llama-3-8B trained using LASeR, to adaptively select relevant evaluation metrics, outperforms baselines by $1.62\%$ on average and can effectively filter underperforming metrics without degrading performance (c.f. Fig. 8). Furthermore, from Table 6, we observe that on reasoning datasets LASeR outperforms sequential RM optimization under four different choices of loss functions: NLL, DPO (Rafailov et al., ), DPO + NLL (Pang et al., ), and KTO (Ethayarajh et al., ). Finally, in Table 7, we find that models trained with LASeR also exhibit the highest generalization to out-of-distribution settings such as on CommonsenseQA (Talmor et al., ) and MATH (Hendrycks et al., ), reiterating the broad generalizability of LASeR.

## 6 Discussion and Conclusion

#### LASeR with different “kinds” of RMs.

In Sec. 4.2, we show that LASeR can choose from a set of candidate RMs, and our analysis in Fig. 3 highlights the fact that LASeR is robust to noisy RMs.
Furthermore, in Appendix C, we show that LASeR can also be used with metric-based rewards (Golovneva et al., ).
These experiments reflect a conceptual split between the generator (the LLM) and the scorer (the RM or metric).
Thus, LASeR is applicable to other settings that follow this paradigm, e.g., using an LLM-as-a-judge (Zheng et al., ), where LASeR could be used to choose between different judge models or prompts or different combinations of RMs and metrics. We leave adding such models into LASeR to future work.

#### Quality of RMs used with LASeR .

Methods that rely on RMs for scoring generally assume that these RMs have a strong correlation with human judgments.
LASeR tempers this assumption in a number of ways:
first, by ensembling multiple RMs, LASeR weakens the effect of noisy RMs; this can be seen in Fig. 3, where LASeR mitigates the negative impact of a noisy RM even as the level of noise is increased.
Moreover, the fact that LASeR can select RMs at an instance-level means that there need not be a *single* RM that always correlates well across all instances.
However, LASeR does require at least one RM to be positively correlated with human judgments on each instance.
If this assumption is not met (i.e., *all* RMs are poorly-correlated across *all* instances), then optimizing for the RMs will yield poor results. Note that this holds true for any method optimizing for RMs.
Because LASeR selects from multiple RMs, its contributions are complementary to developments in RMs, which can easily be integrated into LASeR, as well as improvements to preference optimization loss functions (see Appendix C).
Such RM improvements are likely to be necessary as LLMs are deployed in domains that are out of scope for existing systems and domains with heterogeneous requirements (e.g., our generation domains in Sec. 4.2).
In these cases, there will be no single “perfect” existing RM, and successful solutions will likely involve mixing multiple RMs.
A core benefit of LASeR is its the ability to automatically filter RMs; in Fig. 6 we see that utilization differs across domains.
This allows users to avoid expensive experimentation with subsets of RMs: they can simply offload this task to LASeR, which will automatically select the more useful RM(s).

#### Conclusion.

We present LASeR, an adaptive method for selecting RMs and iteratively training LLMs using multiple RMs.
We formulate RM selection as a contextual multi-armed bandit problem, learning to select the RM that most improves the LLM conditioned on the given input or query.
We test LASeR across diverse settings, demonstrating its utility on reasoning tasks, open-ended instruction-following tasks, as well as long-context generation, where we use the selected RM to rerank outputs without training the LLM.
Across domains, we show that LASeR *consistently* results in superior performance, whereas multi-RM baselines that select RMs using random or fixed strategies, or ensemble multiple RMs uniformly, have lower and more variable performance.
In our analysis, we show that LASeR is robust to noisy RMs, and flexibly uses different RMs depending on the domain, and generalizes to multiple settings.
Moreover, by selecting one RM at a time, LASeR provides the best of both worlds: consistently outperforming all baselines while still maintaining efficiency by only optimizing for one model at a time.
Taken together, our results indicate that LASeR is a powerful tool for improving LLMs, providing us with the robustness and performance improvements associated with training against multiple RMs without incurring the downsides of having to select only one fixed RM or having to handle conflicting reward signals.

## Acknowledgements

We thank Justin Chih-Yao Chen for his feedback on the figures. This work was supported by DARPA ECOLE Program No. HR00112390060, NSF-AI Engage Institute DRL2112635, DARPA Machine Commonsense (MCS) Grant N66001-19-2-4031, the Accelerating Foundation Models Research Program, and NSF CAREER Award 1846185. The views contained in this article are those of the authors and not of the funding agency.

## Appendix A Experiments

### A.1 Experimental Settings

Training setup. For training with LoRA, we set the rank to 16 and alpha to 32. We fine-tune the model for 10 iterations using a learning rate of $5e\!-\!6$ and a batch size of 16. Following Pang et al. (), we generate $P=10$ pairs per problem for training with our loss in Sec. 3.1. For all experiments, we trained each method to converge. The number of iterations for each method was selected based on the observed convergence, with a performance metric threshold of $0.1$ across training batches on the dev set. Our experiments are
run on 4 RTX A6000 with 48G memory each.

RewardBench. Following Lambert et al. (), rewards are computed with no reference model and only used the log-likelihood of the reward model. For instance, given a reward model $\pi_{R^{\star}}$, the reward for an input $x_{i}$ and response $y_{i}$ is calculated as: $\log\pi_{R^{\star}}(y_{i}\mid x_{i})$. There is no need for normalization since we use this log-likelihood to rank the responses. Specifically, on GSM8K, we computed the average and standard deviation of the rewards for the chosen responses as $-4.2965$ and $1.4156$, respectively, and for the rejected responses as $-5.9861$ and $1.6546$, respectively.

**Table 3: Number of examples in train, val, and test sets.**
| Task | Dataset/Category | Train | Dev | Test | Total |
| --- | --- | --- | --- | --- | --- |
| Reasoning | StrategyQA | 1946 | 278 | 556 | 2780 |
|  | GSM8K | 6750 | 750 | 1000 | 8500 |
|  | MMLU | 11135 | 1591 | 3182 | 15908 |
| WildChat | Creative | 3500 | 500 | 1000 | 5000 |
|  | Analysis | 3500 | 500 | 1000 | 5000 |
|  | Coding | 3500 | 500 | 1000 | 5000 |
|  | Factual | 3500 | 500 | 1000 | 5000 |
|  | Math | 3500 | 500 | 1000 | 5000 |
| LongBench | Single-doc QA | 3534 | 505 | 1010 | 5049 |
|  | Multi-doc QA | 3500 | 500 | 1000 | 5000 |
|  | Summarization | 3500 | 500 | 1000 | 5000 |
|  | Few-shot learning | 3500 | 500 | 1000 | 5000 |

Datasets.
For StrategyQA, GSM8K, and MMLU, we divided each dataset into training and test sets. The model is fine-tuned on the training set and dev set, then evaluate on the test set. For WildChat, the dataset was split into a 70/10/20 ratio for training, development, and testing. Following Zhao et al. (), prompt categorization is done using a small off-the-shelf classifier.(^6^66Link: [https://huggingface.co/valpy/prompt-classification](https://huggingface.co/valpy/prompt-classification)) For LongBench, we subsample 5K examples for three tasks: multi-document QA, summarization, and few-shot learning. Each category was split into a 70/10/20 ratio, and the bandit model was trained and validated on the training and development sets and then tested on the test set. We report the detailed number of instances for train, development, and test sets in Table 3.

Prompts. We present a collection of prompt templates designed
for each dataset as follows:

### A.2 Details of Bandit Algorithms

Algorithm for Sec. 3.3. We provide the detailed algorithm for Sec. 3.3 in Algorithm 1.

Figure: Algorithm 1 Bandit-based Reward Model Selection for LLM Training

Exp3. Exp3 is a non-contextual bandit algorithm designed for adversarial settings. It maintains a probability distribution over the arms and selects arms based on the exponential weighting of past rewards. The probability for choosing arm $a_{k}$ at round $t$ is calculated as follows:

$$ $p_{k}(t)=(1-\gamma)\frac{\exp(S_{k}(t))}{\sum_{a_{k}\in\mathcal{A}}\exp(S_{k}( t))}+\frac{\gamma}{K},$ $$

where $S_{k}(t)$ is the cumulative score for arm $a$ up to time $t$ and $\gamma$ is a parameter controlling the exploration rate.

The arm $a_{k}$ is selected by sampling the following categorical distribution

$$ $a_{t}\sim\mathrm{Categorical}(p_{1}(t),\dots,p_{K}(t))$ (4) $$

The score for arm $a_{t}$ is updated based on the observed normalized reward $-\hat{\mathcal{L}}^{m}(t)$ and the probability $p_{k}(t)$ of selecting that arm:

$$ $S_{k}(t+1)=S_{k}(t)-\frac{\hat{\mathcal{L}}^{m}(t)}{p_{k}(t)}\cdot\mathbb{1}(a _{t}=a_{k}),$ (5) $$

where $\mathbb{1}(a_{t}\!=\!a_{k})$ is an indicator function that equals 1 if arm $a_{k}$ was selected at time $t$, and 0 otherwise.

MAB reward normalization. To maintain a consistent scale and magnitude MAB rewards across training, we apply scaled rewards based on the quantiles of the reward history, following the method outlined by Graves et al. (). Let $L=\{-\mathcal{L}^{m}(1),\dots,-\mathcal{L}^{m}(t-1)\}$ represent the unscaled reward history up to time step $t$. This history’s lower and upper quantiles are denoted as $q_{t}^{lo}$ and $q_{t}^{hi}$, respectively. We set $q_{t}^{\text{lo}}$ and $q_{t}^{\text{hi}}$ to be $20^{th}$ and $80^{th}$ quantiles. The scaled reward, $-\hat{\mathcal{L}}^{m}(t)$, is then defined as:

$$ $-\hat{\mathcal{L}}^{m}(t)=\begin{cases}0&\text{if }-\mathcal{L}^{m}(t)<q_{t}^{ \text{lo}}\\ 1&\text{if }-\mathcal{L}^{m}(t)>q_{t}^{\text{hi}}\\ \frac{-{\mathcal{L}}^{m}(t)-q_{t}^{\text{lo}}}{q_{t}^{\text{hi}}-q_{t}^{\text{ lo}}}&\text{otherwise}.\end{cases}$ $$

## Appendix B Additional empirical results

LASeR adaptively selects from multiple RMs on LongBench. On LongBench (Fig. 7), we observe distinct utilization patterns for the QA tasks vs. summarization and few-shot learning. QA tasks exhibit nearly equal utilization of the top-2 RMs on RewardBench (Zephyr-7B-Alpha and Qwen1.5-7B-Chat in decreasing order), with the utilization of the Qwen RM even *exceeds* that of Zephyr RM for multi-document QA. In contrast, on summarization and few-shot learning the top RM (Zephyr) is far more preferred by LASeR with margins of $59\%$ and $31\%$ over the second-best RM and the least performant RM being utilized less that $3\%$ of the times.

Detailed results for each RM. Here, we provide detailed reasoning results for each chosen RM where we use a single RM during training (c.f. Sec. 3.1) in Table 4. These results demonstrate that Zephyr-7b-alpha outperforms other RMs on StrategyQA and MMLU, except for GSM8K, where Qwen1.5-7B-Chat has the best performance.

Figure: Figure 7: Utilization (%) of each RM on long-context understanding tasks. The bars are arranged based on their overall scores on LongBench, from lowest to highest. LASeR dynamically selects from different RMs depending on the nature of the underlying instance.
Refer to caption: x7.png

**Table 4: Performance of 4 RMs including OLMo, Eurus, Zephyr and Qwen1.5.**
| Method | Llama-3-8B | Mistral-7B |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| StrategyQA | GSM8K | MMLU | StrategyQA | GSM8K | MMLU |  |
| OLMo-7B-Instruct | 80.23 | 68.91 | 65.74 | 68.73 | 44.96 | 56.94 |
| Eurus-7b-kto | 81.15 | 71.13 | 66.26 | 68.64 | 45.37 | 56.96 |
| Zephyr-7b-alpha | 84.29 | 73.16 | 67.15 | 70.06 | 45.81 | 62.04 |
| Qwen1.5-7B-Chat | 84.79 | 73.07 | 67.53 | 71.05 | 45.74 | 62.38 |

**Table 5: Comparison of LASeR and baselines on ROSCOE. The baselines include supervised fine-tuning (SFT), sequential optimization, uniform rewards selection, and base model optimized with one specific evaluation metric (Perplexity, Informativeness).**
| Method | Llama-3-8B | Mistral-7B |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| StrategyQA | GSM8K | MMLU | Avg. | StrategyQA | GSM8K | MMLU | Avg. |  |
| SFT | 80.41 | 69.43 | 65.66 | 71.83 | 68.57 | 43.62 | 56.48 | 56.22 |
| Perplexity | 80.55 | 69.21 | 65.62 | 71.79 | 68.83 | 43.47 | 57.14 | 56.48 |
| Informativeness | 82.87 | 73.55 | 66.69 | 74.37 | 70.29 | 44.98 | 59.29 | 58.19 |
| Random RM Selection | 82.72 | 70.93 | 66.10 | 73.25 | 69.24 | 44.05 | 57.68 | 56.99 |
| Sequential RM Selection | 83.15 | 73.38 | 66.17 | 74.23 | 70.40 | 44.79 | 59.07 | 58.09 |
| LASeR (Ours) | 83.54 | 73.80 | 66.79 | 74.71 | 70.91 | 44.93 | 59.63 | 58.49 |

## Appendix C Generalization Capabilites of LASeR

LASeR Training with Domain-specific Evaluation Metrics. While recent works focus on building RMs that reflect preferences across domains, an extensive body of prior work develop a suite of evaluation metrics catered to specific domains such as reasoning (Golovneva et al., ; Prasad et al., ). To show that LASeR can be used to select any kind of evaluation metric from a collection of metrics during training, in Table 5, we present results with training LLMs on model-based metrics from ROSCOE (Golovneva et al., ) by replacing RMs with informativeness, faithfulness, reasoning alignment, hallucination, common sense error, semantic, coherence and perplexity in Sec. 3. Llama-3-8B models trained using LASeR yields 1.62% accuracy improvement over baselines on average across three datasets. These results are also generalized to Mistral-7B, except for GSM8K, where we achieve comparable performance to the Base + Informativeness baselines. Note that the perplexity of most responses is nearly identical, making it difficult to distinguish between them, explaining why perplexity shows little to no improvement compared to the base model.

Figure: Figure 8: Impact of irrelevant metrics from ROSCOE on the GSM8K accuracy of LASeR and sequential baseline.
Refer to caption: x8.png

#### Robustness to Underperforming Evaluation Metrics.

Similar to our analysis on noise in rewards in Sec. 5, we investigate how adding ROSCOE metrics with poor correlation to human-annotated labels in meta-evaluation by Golovneva et al. () impacts the performance of Llama-3-8B on GSM8K. Once again, even with ROSCOE metrics, demonstrates LASeR can maintain consistent performance by adaptively prioritizing the most relevant reward signals, outperforming the sequential baseline, which fails to filter out irrelevant information effectively. Fig. 8 shows that as the number of irrelevant metrics increases, LASeR’s selection strategy continues to perform robustly.
Specifically, LASeR has an average performance drop of only $0.13\%$, whereas the sequential baseline suffers a $2.15\%$ accuracy drop under the same conditions.
Lastly, LASeR using Exp3 maintains a consistent performance level with a $0.4\%$ accuracy drop.

**Table 6: Across different training loss functions (Sec. 3.1), optimizing with multiple RMs via LASeR outperforms the sequential RM selection with Llama-3-8B.**
| Loss | Method | StrategyQA | GSM8K | MMLU | Avg. |
| --- | --- | --- | --- | --- | --- |
| NLL | Sequential | 82.75 | 71.80 | 65.41 | 73.32 |
|  | LASeR | 85.11 | 74.94 | 67.09 | 75.71 |
| DPO | Sequential | 83.26 | 71.94 | 65.38 | 73.53 |
|  | LASeR | 84.71 | 73.94 | 67.02 | 75.22 |
| KTO | Sequential | 83.62 | 73.07 | 69.02 | 75.24 |
|  | LASeR | 84.87 | 73.07 | 69.05 | 75.66 |
| NLL+DPO | Sequential | 83.90 | 72.94 | 68.02 | 74.95 |
|  | LASeR | 85.96 | 74.75 | 68.02 | 76.24 |

Generalization to Training Loss Functions.
In Sec. 3.1, we state that the choice of loss function used to train the LLM depends on the underlying task or domain.
Nevertheless, we always use the training loss as the MAB reward to update the MAB’s parameters. Here we study the performance of LASeR and baselines with 4 different loss functions, NLL, DPO, NLL + DPO (Pang et al., ), and KTO (Ethayarajh et al., ), in the reasoning domain.
Results in Table 6 show that training LLMs with multiple rewards using LASeR outperforms sequential RM selection by $2.4\%$, $1.7\%$, and $1.3\%$ when using NLL, DPO, NLL+DPO loss functions, respectively; while both methods yield comparable performance with KTO. Additionally, we found that the most effective training loss functions are NLL + DPO for StrategyQA, NLL for GSM8K, and KTO for MMLU.
However, irrespective of the choice of the underlying loss function, LASeR is more effective at balancing and adaptively selecting from multiple RMs.

**Table 7: Generalization performance of different models trained on StrategyQA, MMLU, and GSM8K, and evaluated on CSQA and MATH, respectively.**
| Method | Llama-3-8B | Mistral-7B |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| CSQA | MATH | Avg. | CSQA | MATH | Avg. |  |
| SFT | 65.64 | 29.13 | 47.39 | 59.06 | 16.38 | 37.72 |
| Best RM | 67.59 | 31.54 | 49.57 | 60.46 | 18.08 | 39.27 |
| Avg. RM | 67.16 | 30.36 | 48.76 | 60.06 | 17.25 | 38.66 |
| Random RM Selection | 68.31 | 30.21 | 49.26 | 60.19 | 16.96 | 38.58 |
| Sequential RM Selection | 67.73 | 30.25 | 48.99 | 60.56 | 17.96 | 39.26 |
| LASeR (Ours) | 69.26 | 31.67 | 50.47 | 61.65 | 18.97 | 40.31 |

#### Generalization to OOD Tasks.

We first assess the generalization ability of our method by training models on specific datasets and evaluating their performance on out-of-distribution reasoning tasks. Specifically, we train the model on the StrategyQA and MMLU datasets and evaluate its generalization on the CommonsenseQA (CSQA; Talmor et al., ) dataset. Similarly, we train on GSM8K and test on MATH (Hendrycks et al., ) to assess the model’s ability to generalize across different reasoning datasets. These tasks are designed to capture both general reasoning ability and OOD generalization across domains. We report the results in Table 7, where we find that across both Llama-3-8B and Mistral-7B models, models trained with LASeR yield the best average accuracy beating training with the best RM by roughly $2\%$ (absolute) on CSQA with Llama-3-8B. On Mistral-7B, training with LASeR outperforms both training with single best RM and sequential RM selection by slightly over $1\%$.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - Please give more insights on each of the RMs used. For example which RM performs well for which examples? Once that is known, could you cluster the batches for training based on all inputs in the same batch using the same RM? Would that not perform better? Currently a batch is a mix of different inputs, some of which may not be suitable to the RM chosen for the batch.
- If one is using multiple RMs, would one not like to tune the RM itself instead of choosing between RMs that are designed to be good general purpose models (i.e. can you train domain specific RMs rather than generic RMs)? I know the setting is that the reward models are already given to you. But if you are using multiple RMs, would you not like to train the RMs to be good at certain domains?

### Suggestions:

- Please consider adding more details on what the RMs are, and how they were chosen. Does the method generalize to more than 4 RMs?

### Limitations:
- The experiments did not make it clear to me that this is the best way to use multiple reward models. in fact none of the baselines are allowed to choose between multiple RMs for different batches. This points to a trade-off between the memory usage of multiple RMs vs accuracy on downstream tasks.

### Additional points
Happy to raise my score given a convincing rebuttal on some of the concerns raised.
soundness: 2
strengths: - The results are strong. They show improvements over their chosen baselines.
- The paper is written clearly, and is highly readable. All the points are well covered.
summary: The paper discusses one method to use multiple reward models during training. The method uses a linear adapter, called A_k corresponding to reward model k \in [K] which is used to obtain a certain "reward". This reward pertains to how good a reward model is for a given batch of inputs. The inputs are modelled as vectors obtained from the last token of a sentence, averaged over the sentences in the batch. The method strongly suggests using multiple reward models, but why use their method of LinUCB is not clear. The paper suggests that ensembling works, but there are many methods for ensembling. Why use their method is not clear from the given experiments.
weaknesses: - There is no baseline that selects more than 1 reward model over an epoch of training. The baselines are rather weak. What about a simple baseline that learns a classifier to choose the reward model based on the type of query. If Bandit algorithms perform better than such an approach, we can conclude that the method of using covariance works. Without any such baseline, we are left with the conclusion that one should choose the RM depending on the input c(t). Further to this, why use the authors' specific method dependent on c(t) is not clear.
- Why use MAB? why is there no discussion on the advantages gained by using the "exploration" that MABs provide?
- Adversarial training set choice. Imagine there are k RMs. Assume RM k is well coordinated with input_i where i%K = k. In other words, this RM k correlates well with human judgement for every i'th instance for i%K = k. For the remaining inputs, it performs poorly. Similarly, assume this is true for all RMs RM_i \in {RM_1,...RM_K}. Then assume a batch is of size K, where batches arrive without randomized order (i.e. sequential order). Then given a batch, no RM would correspond to the required outputs for all inputs in the batch.

- Why LinUCB? What is the "exploration" here? Why not simply train a classifier (or matrix A) to choose based on some criterion/classifier?
- Why not a linear combination of the reward functions? What if two are good? Why use arg-max over a weighted sum?

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Could the authors elaborate the definition of $e_m$ and the method for generating the sentence embedding from the model $\pi_m$? Which part of Figure 1 corresponds to this process?
2. Could the authors provide insights into why Random RM selection often outperforms the Offline RM Ensemble and why Sequential RM selection surpasses Random RM selection?
soundness: 2
strengths: 1. The proposed LASER iteratively trains LLMs using different RMs by dynamically selecting the most appropriate one for each training instance using a contextual bandit algorithm, specifically LinUCB, which effectively addresses the potential inefficiencies and conflicts present in ensemble methods that handle multiple RMs simultaneously.
2. The paper thoroughly evaluates LASeR across various datasets and tasks, showcasing its superior performance over baselines.
summary: This paper introduces LASeR, a novel approach for training large language models (LLMs) by adaptively selecting suitable reward models (RMs) for various tasks, framed as a multi-armed bandit problem. Unlike previous methods that ensemble multiple RMs, LASeR aims to enhance computational efficiency and avoids conflicts from using different RMs, by selecting and utilizing the most appropriate RM for each instance to rank outputs and generate preference data. Extensive experiments demonstrate that LASeR consistently outperforms baseline methods across multiple benchmarks.
weaknesses: 1. Fairness of Comparison: In Section 4, LASeR is compared with baselines that do not actively leverage information from interactions with the datasets. The RM selection in these baselines is either fully random or in offline fashion. Specifically, in Table 4, the "best RM" baseline (Zephyr-7b-alpha) is not the top performer on StrategyQA and MMLU, which contradicts claims in Appendix B. Additionally, the sequential RM selection shows better performance than other baselines across most datasets, yet the explanation for this surprising result is missing. It would be more appropriate to compare LASeR with ensemble methods where RM aggregation adapts based on training information, as LASeR leverages training loss to adapt RM selection.
2. Limited Comparison with Ensemble Methods: The paper only compares LASeR with a simple offline RM Ensemble using averaged scores, despite the existence of more sophisticated ensemble methods that incorporate training signals [A1] [A2]. It remains unclear whether LASeR outperforms such methods. Moreover, the results on reasoning are the only ones compared against Offline RM Ensemble. Due to these issues, it is insufficient to infer whether LASeR outperforms ensemble methods in terms of robustness to noisy rewards or different tasks like long-context understanding and instruction-following tasks.
3. Clarification of Unique Contributions: While the paper effectively incorporates LinUCB for RM selection in LLM training, there is insufficient discussion on the challenges encountered during this integration. Highlighting the unique contributions and the advantages over ensemble methods would strengthen the paper.
4. Efficiency Concerns Beyond Accuracy-Time Tradeoff: The paper does not address that, in high-dimensional embedding spaces or with a large number of RMs, LASeR’s need to compute the inverse of the covariance matrix A for each arm could be computationally intensive.
5. Minor issues: 
   1. Undefined/Misleading Notations: For example, in Equation (1), $x$ should be $x_i$, and $\sigma,\beta$ are not defined previously.  Footnote 2 mentions $P$ as the number of sampled pairs per prompt, while $n$ in the notation paragraph seems to have a similar purpose. In section 3.2, the notation should be $\mathbf{x}_{m,t}$ to represent the batch at the $t$-th step of the $m$-th iteration. Also, $\mathbf{x}_t$ should be defined clearly as the set of prompts, but is later used as the set of the $d$-dimensional embedding of the prompt. The left-hand side of equation (2) should be the index of the selected RM, while $R_t^{*}$ is the model itself.
   2. Grammar and Typos: Some sentences, such as those describing the stages of LASER, contain grammatical issues or are incomplete. The reviewer recommends a thorough review for consistency in symbols and text.


[A1] Zhang, Shun, et al. "Improving reinforcement learning from human feedback with efficient reward model ensemble." arXiv preprint arXiv:2401.16635 (2024).

[A2] Wang, Zihao, et al. "Transforming and Combining Rewards for Aligning Large Language Models." arXiv preprint arXiv:2402.00742 (2024).

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Can author address why chosing the reward as the negative train loss would help in mitigating bias of LLM?
soundness: 2
strengths: 1. Innovative approach to automatically selecting Reward Models using bandit algorithm. This MAB framework, allowing the model to dynamically choose the most suitable RM at instance level based on contextual information (embedding of the last token). This avoid the computational burden of previous approaches that based on RM ensemble.

2. Empirical results presented in the paper indicate that LASER achieves performance improvements over traditional methods, such as ensemble RM scores. LASER achieve improvement in wide range of tasks, suggest it's generalibility.
summary: The paper introduce the LASER framework, which enables LLMs to adaptively select the most appropriate Reward Models for different tasks. This approach addresses the challenges of using multiple RMs by optimizing the selection process, combining with bandit algorithm, such as LinUCB and Exp3, thereby enhancing the performance of LLMs across reasoning, QA and instruction following tasks.
weaknesses: A notable weakness of the paper is the approach taken to set the reward in the MAB problem as the negative training loss. While maximizing the reward is equivalent to minimizing the training loss, this method raises concerns about the alignment of selected preference pairs with the language model. Specifically, if the MAB algorithm consistently selects RM that align with the LM (means that the preference pair generate by the RM has low loss in LM), it's essentially just reinforce the LM's existing biases rather than addressing its mistakes.

This raises a critical question, why would selecting RM that align with the LM lead to improved performance? If the LM makes an error and the RM correctly identifies that mistake, the training loss would be high, which means the MAB algorithm might avoid selecting the RM that highlights the LM's shortcomings. Consequently, this could result in a scenario where the model becomes less capable of correcting its errors, as it may favor RMs that validate its outputs rather than challenge them.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: I have two questions:
1. Why does the ensemble baseline perform worse than the best RM? This is indicated in lines 95-97 and Table 1.
2. You mention that Table 3 shows results on long-context understanding without training and best-of-n sampling. Do you have any results where you use the RM to generate preference pairs and use those as in-context examples?
soundness: 3
strengths: - The paper is well written
- Using LinUCB is well motivated
- Solution makes intuitive sense
- Experiments are extensive
summary: This paper describes LASER, which:
- Uses an ensemble of reward models (specifically 4) to align language models
- Selects reward models with a multi-armed bandit (based on the LinUCB algorithm) to classify the (embedded) input into different RM tasks, and
- Uses the chosen RM with language model sampling to create and train on preference data.
weaknesses: Nitpicky things:
- (As far as I know) Figure 2 is not referred to in the main text
- Add a y-axis to Figure 4
- (As far as I know) Appendix B is not referred to in the main text
- I could not find where it mentions how many preference pairs are generated per iteration/

Meta-review: The paper proposes using a contextual bandit algorithm to select among multiple reward models
when performing mini-batch preference finetuning of LLMs.
Empirical results shown in the paper suggest that adaptively selecting reward models for different inputs/prompts
can yield substantial benefits over baselines that commit to a single reward model, or randomly select among models.

The reviewers agreed that the paper studies a well motivated problem,
but raised persistent questions about the motivation for the solution approach (i.e. the use of bandit algorithms like LinUCB).

Consider Figure 1.
Based on the reward that is revealed to the bandit (i.e. -L_m(t)), notice that the reward can also be observed for all the other arms not selected by the bandit algorithm.
For all of the responses generated by the LLM, each of the RMs can be used to create preference pairs and thereby compute -L_m(t) for each one of them.
This means that this is an online learning problem (learning from expert advice), with very effective algorithms like HEDGE, MULTIPLICATIVE WEIGHTS, etc.

Decision: Reject