Uncertainty-aware Reward Model: Teaching Reward Models to Know What is Unknown

## Abstract

Abstract Reward models (RM) play a critical role in aligning generations of large language models (LLM) to human expectations. However, prevailing RMs fail to capture the stochasticity within human preferences and cannot effectively evaluate the reliability of reward predictions. To address these issues, we propose Uncertain-aware RM (URM) and Uncertain-aware RM Ensemble (URME) to incorporate and manage uncertainty in reward modeling. URM can model the distribution of disentangled attributes within human preferences, while URME quantifies uncertainty through discrepancies in the ensemble, thereby identifying potential lack of knowledge during reward evaluation. Experiment results indicate that the proposed URM achieves state-of-the-art performance compared to models with the same size, demonstrating the effectiveness of modeling uncertainty within human preferences. Furthermore, empirical results show that through uncertainty quantification, URM and URME can identify unreliable predictions to improve the quality of reward evaluations.

## 1 Introduction

Large language models (LLM) have demonstrated remarkable capabilities across various domains (Singhal et al., ; Cui et al., ; Kasneci et al., ). These powerful LLMs are trained to align with human values and expectations to avoid harmful and toxic generations. To achieve alignment, LLMs rely on feedbacks from reward models (RM), where the feedbacks are provided in the form of rewards (Singhal et al., ; Cui et al., ; Kasneci et al., ). These rewards typically reflect the quality and users’ preferences of the responses provided, and hence reward maximization will guide the LLM to more effectively satisfy user queries. In this paradigm, RMs fundamentally decides the efficacy of alignment, as they primarily steer the LLMs through feedback. Therefore, the reliability and accuracy of this feedback is essential in aligning LLMs with intended human values and preferences.

However, current RMs fail to capture the stochastic nature of human preferences (Baylis, ) and lack the ability to evaluate the reliability of the predicted rewards. In prevalent RMs, a value head (usually a linear layer) is added to the pretrained base model and maps the hidden states to reward scalars (Bai et al., ; Ouyang et al., ) or attribute scores (Adler et al., ; Wang et al., ). This results in a deterministic reward modeling process, unable to accommodate the variabilities of human preferences. Moreover, there is no other information to validate the reliability of these reward predictions.

Uncertainty is of major importance in machine learning (Hüllermeier & Waegeman, ) and an appropriate representation for uncertainty is essential for developing trustworthy and reliable models (Yang et al., ; Varshney & Alemzadeh, ). Uncertainty originates from two different sources: *aleatoric* and *epistemic*. Aleatoric uncertainty refers to the inherent variability and randomness of data. As opposed to this, epistemic uncertainty is caused by lack of knowledge, i.e. ignorance of the model instead of any underlying randomness.

In the context of reward modeling for LLM, aleatoric uncertainty refers to the stochasticity of human preferences, while epistemic uncertainty comes from the RMs’ lack of knowledge to make reliable evaluations. Therefore, introducing uncertainty to reward modeling improves modeling capacity of RMs and enhance reliability of the reward predictions. Consequently, by identifying and filtering out out-of-distribution (OOD) data where RMs fail to generalize, rewards with better reliability pave the way for more efficient alignments of LLMs.

In this paper, we propose Uncertain-aware RM (URM) and Uncertain-aware RM Ensemble (URME) to handle aleatoric and epistemic uncertainty in reward modeling respectively. URM is equipped with an uncertainty-aware value head to model the distributions of multiple attributes within human preferences. We demonstrate that with the popular bradley-terry-model loss function (Bradley & Terry, ), RMs cannot quantify uncertainty of human preferences even with an uncertainty-aware value head. Therefore, URMs are trained via maximum likelihood estimation and attribute regression. URME quantifies epistemic uncertainty by the discrepancies among URMs in the ensemble, identifying potential lack of knowledge. During reward evaluation, filtering strategy can be applied to prompt-response pairs with high uncertainty, in case that LLMs learn unintended or potentially harmful behaviors that URMs may not be able to accurately evaluate.

Empirical results on a popular RM benchmark RewardBench (Lambert et al., ) demonstrate that URM with 8B model size achieves state-of-the-art performance among models with the same size and outperforms a number of strong large models including Nemotron-4-340B (Adler et al., ). And through uncertainty quantification, URM and URME are able to identify their level of knowledge for the input data and make the reward evaluations more reliable through filtering strategy. Furthermore, results of best-of-$n$ sampling validates that URM and URME can effectively enhance the generation quality of LLMs.

Contributions of this paper include:

(1) We introduce URM and URME to model the uncertainty within human preferences and reward models themselves.

(2) URM and URME are able to improve LLMs’ generation effectively. Notably, URM achieves state-of-the-art performance on RewardBench compared with models of the same size (8B).

(3) Empirical results demonstrate that URM and URME can successfully quantify uncertainty to identify areas where the models lack sufficient knowledge to make accurate predictions, leading to more reliable reward evaluations.

## 2 Preliminaries

LLM alignment typically consists of three stages (Ouyang et al., ): supervised fine-tuning (SFT), reward modeling and proximal policy optimization (PPO) (Schulman et al., ). SFT utilizes expert demonstrations to fine-tune the pretrained base model in a supervised-learning fashion to enable LLMs to follow user instructions.

Reward Modeling Reward modeling aims to learns human preferences explicitly (Ouyang et al., ) or implicitly (Rafailov et al., ). For some prompt $x$ and a response pair $(y_{w},y_{l})$, $y_{w}$ is the chosen response preferred by humans and $y_{l}$ is rejected. Following the Bradley-Terry model (Bradley & Terry, ), under RM $r_{\phi}$ the probability of $y_{w}$ being preferred than $y_{l}$, i.e. $y_{w}\succ y_{l}$, is

$$ $\displaystyle P(y_{w}\succ y_{l}|x)$ $\displaystyle=\log\frac{\text{exp}(r_{\phi}(x,y_{w}))}{\text{exp}(r_{\phi}(x,y _{w}))+\text{exp}(r_{\phi}(x,y_{l}))}$ (1) $\displaystyle=sigmoid(r_{\phi}(x,y_{w})-r_{\phi}(x,y_{l}))$ $$

Thus, to train a RM to prioritize chosen responses over rejected responses, the loss function is the maximum likelihood estimation of Eq. 1

$$ $L_{1}=-\mathbb{E}_{x,y_{w},y_{l}\sim D}\left[\log sigmoid\left(r_{\phi}(x,y_{w })-r_{\phi}(x,y_{l})\right)\right],$ (2) $$

where $r_{\phi}$ is the reward model parameterized by $\phi$, consisting of the pretrained base model and a linear value head. The trained RM can be used to improve the LLM’s generation by Best-of N (BoN) (Stiennon et al., ) or RLHF (Ouyang et al., ).

PPO In this stage, LLMs are fine-tuned with feedbacks from the RM. To prevent the model deviate too far from the pretrained model and forget linguistic skills, there is also an Kullback-Leibler (KL) divergence penality in the reward from the RM. Thus, the total reward $\hat{r}$ is

$$ $\hat{r}(x,y)=r_{\phi}(x,y)-\eta\text{KL}(\pi(y|x)\|\pi_{\text{ref}}(y|x)),$ (3) $$

where $\eta$ is the coefficient for the KL penalty, $\pi$ is the model to be fine-tuned and $\pi_{\text{ref}}$ is the reference model which is usually the SFT model. Running PPO to maximize reward from Eq. 3 can not only align the LLM with human preferences, but also prevent it from severely deviating from the reference model.

## 3 Related Work

### 3.1 Multi-attribute Reward Modeling

To generate helpful, harmless and truthful responses (Askell et al., ), LLMs must be aligned with human expectations. Current methods fine-tune models based on human (Christiano et al., ; Stiennon et al., ; Bai et al., ; Ouyang et al., ) or AI feedbacks (Bai et al., ; Sun et al., ) to maximize preference-based rewards, which are provided by reward modeling. Typically, a reward model is learned and LLMs will improve their generation quality according to feedbacks from the reward model (Bai et al., ; Ouyang et al., ; Shao et al., ; Stiennon et al., ).

Recent studies show that human and LLM judges may introduce potential biases to annotations of preference (Zhang et al., ; Kotek et al., ; Wang et al., ; Chen et al., ). Moreover, traditional RMs usually rely on single-dimensional feedback on general quality instead of fine-grained multifaceted signals to indicate multiple attributes such as helpfulness, coherence and verbosity (Dong et al., ). Adler et al. () discovered that multi-attribute RMs trained on datasets with high-quality attribute-specific annotations (Cui et al., ; Wang et al., ) are able to disentangle real helpfulness and other irrelevant aspects such as lengthy bias (Shen et al., ; Singhal et al., ). There are also alignment methods directly aimed at multi-attribute alignment. Zhou et al. () includes preference on multiple attributes in the Direct Preference Optimization (DPO) loss function (Rafailov et al., ), trying to optimize preference rewards for all attributes simultaneously. Lou et al. () proposed to achieve multi-attribute alignment sequentially, one attribute at a time, where LLMs learns to align with new attributes while staying aligned with previous dimensions.

### 3.2 RLHF, Offline RL and Uncertainty

In RLHF, LLM policy is optimized via interactions with the RM, whose training data is pre-collected preference pairs (Bai et al., ; Ouyang et al., ). In this setting, RLHF falls into the category of offline RL, where RL policies cannot interact with the environment and get feedbacks in real time, but instead can only be updated based on an offline dataset collected by some other policy (Levine et al., ). Offline RL is notoriously difficult due to the distributional shift issue (Lou et al., ; Ma et al., ; Prudencio et al., ). Recent advancements in iterative LLM alignment methods (Yuan et al., ; Dong et al., ; Xiong et al., ) iterates between LLM fine-tuning and the sampling and annotation of new training data, alleviating the distributional shift issue. Although these iterative methods aim to transcend the constraints of the offline setting,, RLHF is still offline within each iteration.

An important topic in offline RL is uncertainty quantification, which has long underpinned many critical roles (Abdar et al., ), such as trustworthy decision-making (Huang et al., ; Eriksson & Dimitrakakis, ) and improving reliability of machine learning models (Wang et al., ). In offline RL, uncertainty quantification enables out-of-distribution data detection and keep the policy within the offline dataset’s support area through conservative updates to avoid distributional shift (Yu et al., ; Kidambi et al., ; An et al., ; Zhu et al., ). So it is natural to introduce uncertainty to RLHF to make LLM alignment more reliable and effective.

Ensemble of RMs are discussed in previous works. However, we study the ensemble of uncertainty-aware RMs to identify unreliable reward evaluations, while previous discussions are limited to using RM ensembles to mitigate reward hacking Coste et al. (); Eisenstein et al. () and using value heads to disentangle length and quality in reward modeling Chen et al. (). We notice a concurrent work QRM Dorka () which also models human preferences by distributions. QRM only studies distributional RMs, while we also study the ensemble of such uncertainty-aware RMs. Moreover, QRM is trained via quantile regression (Koenker, ), a variant of our attribute regression. But we also studied uncertainty-aware RMs trained via maximum likelihood estimation, which can better capture the uncertainty of rewards.

## 4 Methodology

Figure: Figure 1: Architecture of URM and URME. URMs output $\mu$ and $\sigma$ to parameterize normal distributions, from which multi-attribute scores are sampled. The scores are then combined to reward scalars by weights generated by a gating layer. URME consists of multiple URMs, allowing for quantification of the epistemic uncertainty using the disagreement among the URMs. In the given example, there is a substantial divergence among URMs, indicating significant epistemic uncertainty. Thus, although URM 1 correctly assigns a small, negative reward to the input, the significant epistemic uncertainty still indicates the URMs lack relevant knowledge to provide reliable evaluation of the inputs.
Refer to caption: x1.png

In this section, we will introduce our uncertain-aware reward model (URM) and uncertainty-aware reward model ensemble (URME) to quantify aleatoric and epistemic uncertainties respectively.

Fig. 1 gives the architecture of URM and URME. URMs quantify the aleatoric uncertainty by modeling the distribution of scores, and the epistemic uncertainty is quantified by the disagreement within the URME. In the given example, the response is incorrect and there is large disagreement within the URME, indicating significant epistemic uncertainty and the models’ lack of relevant knowledge.

### 4.1 Uncertainty-aware Reward Model

Traditional RMs optimize the Bradley-Terry model loss (BT-loss) in Eq. 2,to enlarge the discrepancy between the rewards of chosen and rejected responses so that the preference probability is maximized. The value head will map hidden states from the base model to a scalar reward. Such mapping is deterministic and thus cannot catch any uncertainty (Chua et al., ) within the reward modeling process.

However, at its core, human preferences exhibit a distinctly probabilistic nature, rather than being strictly deterministic (Baylis, ). This issue is further exaggerated due to the bias and inconsistencies introduced by human annotators (Sylolypavan et al., ; Sleeman & Gilhooly, ; Chen et al., ). Between individuals, preferences differ from person to person. This means what’s preferable for one may not be for another. Even within individuals, preferences are not static. They can swing based on numerous factors such as mood and context. These stochastic natures of human preferences contribute to adopting a probabilistic framework for modeling preferences with aleatoric uncertainty.

Prior works have explored a number of uncertainty-aware neural networks (Neal, ; Lakshminarayanan et al., ), especially in RL (Gal et al., ; Depeweg et al., ) and model-based RL (MBRL) (Chua et al., ; Yu et al., ; Kidambi et al., ).

Figure: Figure 2: Comparison between URM and traditional RM in estimating preference distribution.
Refer to caption: x2.png

Considering RMs act similarly to the reward part of the dynamics model in MBRL, aleatoric uncertainty within human preferences can be captured by outputting the parameters of a parameterized distribution. Specifically, unlike traditional RMs that output a single deterministic reward value, uncertainty-aware RMs can model the distributions of human preferences. As schematically shown in Fig. 2, given a prompt-response pair with multiple preference annotation samples, traditional RMs can only provide a fixed reward estimation and fails to represent the real preference. But uncertainty-aware RMs are able to offer a more accurate approximation of the human preference distribution.

To model the preference reward distribution, URM adds an probabilistic value head to the pretrained base model. The value head takes in the last hidden state $h$ of the base model and outputs mean $\mu$ and logged standard deviation $\sigma$ to parameterize a normal distribution $\mathcal{N}(\mu,\text{exp}(2\sigma))$, from which preference rewards are sampled, and the aleatoric uncertainty is quantified by variance of the distribution. Reparameterization technique is adopted to enable gradient back-propagation.

However, we show that introducing the probabilistic value head and the sample-based reward to RMs with BT-loss, the aleatoric uncertainty still cannot be quantified.

We denote the reparameterization parameter $\alpha\sim\mathcal{N}(0,1)$, and thus the sampled reward $r=\mu+\alpha\text{exp}(\sigma)$. Substituting chosen reward $r_{w}$ and rejected reward $r_{l}$ into BT-loss in Eq. 2. For some given input $x,y_{w},y_{l}$, gradient w.r.t. the logged standard deviation $\sigma_{w}$ is given by

$$ $\displaystyle\nabla_{\sigma_{w}}L_{1}=-\mathbb{E}_{\alpha_{w},\alpha_{l}\sim \mathcal{N}(0,1)}\left[\alpha_{w}\big{(}1-sigmoid(r_{w}-r_{l})\big{)}\text{exp }(\sigma_{w})\right]=0,$ $$

where $r_{w}=\mu_{w}+\alpha_{w}\text{exp}(\sigma_{w})$ and $r_{l}=\mu_{l}+\alpha_{l}\text{exp}(\sigma_{l})$. Similarly, $\nabla_{\sigma_{l}}L_{1}=0$. The unlearning effect of the variance term demonstrates that under the BT-loss, RMs still cannot quantify the uncertainty even equipped with a probabilistic value head.

Recent advances in multi-attribute RMs demonstrate they are capable of providing fine-grained rewards and disentangling real helpfulness and other irrelevant aspects such as lengthy bias (Adler et al., ; Chen et al., ). The multi-attribute scores consist of human-or-AI-annotated ratings on multiple aspects such as helpfulness, coherence and verbosity. To learn a multi-attribute uncertainty-aware RM, we propose two ways to train the probabilistic value head.

Maximum Likelihood Estimation In URM, scores of all attributes are modeled by a distribution, we can train the probabilistic value head with maximum likelihood estimation (MLE). Since attributes are disentangled in multi-attribute RMs, it is fair to assume that they are independent, i.e. diagonal covariance for the parameterized normal distribution. Thus, the MLE loss function for URM is

$$ $\displaystyle L_{2}=-\mathbb{E}_{x,y\sim D}$ $\displaystyle\left[\log\mathbf{P_{\theta}}(R|x,y)\right]=-\mathbb{E}_{x,y\sim D }\left[\sum\limits_{i=0}^{n}\log P_{\theta}(R_{i}|x,y)\right],$ (4) $$

where $R_{i}$ is the $i$-th attribute score from the label and $\log P_{\theta}(R_{i}|x,y)$ is the log-probility of $R_{i}$ from the parameterized distribution $\mathcal{N}\left(\mu_{i},\text{exp}(2\sigma_{i})\right)$. Thourgh MLE, the probabilistic value head is able to efficiently approximate the attribute scores’ distribution, hence training URMs to fit the unique characteristics of the attribute scores.

Attributes Regression with Reparameterization we can also directly regress the sample-based rewards on multi-attribute scores $R\in\mathbb{R}^{n}$, similar as Adler et al. () but with sampling and reparameterization. In this setting, URM’s mean square error (MSE) loss function is

$$ $L_{3}=\mathbb{E}_{x,y\sim D}\left[\sum\limits_{i=0}^{n}\left(r_{i}(x,y)-R_{i} \right)^{2}\right]$ (5) $$

where $i$ indicates $i$-th attribute, and $r_{i}\sim\mathcal{N}\left(\mu_{i},\text{exp}(2\sigma_{i})\right)$ is sampled from the distribution parameterized by the output of the probabilistic value head. To enable gradient back-propagation, we use the reparameterization technique, so that $r=\mu+\alpha\text{exp}(\sigma)$, where reparameterization parameter$\alpha\sim\mathcal{N}(0,1)$. A more detailed analysis of this MSE loss is given in the appendix B.

With the trained probabilistic value head, during inference we can use mean $\mu_{i}$ for each attribute $i$ as the scores. We learn a gating layer to combine the multi-attribute scores to a reward scalar via weighted sum (Wang et al., ). The gating layer is a fully-connected network, whose input is the last hidden states of the base LLM. And the learning objective of the gating layer is to prioritize chosen responses over rejected responses through the BT loss. For some prompt $x$, chosen response $y_{w}$ and rejected response $y_{l}$, the gating layer will output weights $\omega$ to combine the scores and thus the loss function is given by

$$ $L_{4}=-\mathbb{E}_{x,y_{w},y_{l}\sim D}\left[\log sigmoid\left(\mu^{T}(x,y_{w} )\omega(x,y_{w})-\mu^{T}(x,y_{l})\omega(x,y_{l})\right)\right],$ (6) $$

Since the gating layer only offers weights to combine the scores, the base model and the probabilistic value head are kept frozen during training the gating layer. Besides the gating layer, the multi-attribute scores can also be combined by prior weights and still demonstrate competitive performance (Adler et al., ).

### 4.2 Uncertainty-aware Reward Model Ensemble

Bootstrap ensemble of models is simple and effective for epistemic uncertainty quantification compared with other methods (Neal, ; Hernández-Lobato & Adams, ; Blundell et al., ).

Figure: Figure 3: Illustration of URME in epistemic uncertainty quantification.
Refer to caption: x3.png

Fig. 3 illustrates how URME works in quantifying uncertainty. The input space $X\times Y$ ($X$ for prompt and $Y$ for response) is split into the known and unknown area. In the known area, the training dataset can well support URMs to make reliable reward evaluations. However, in the unknown area, the situation differs. Given that each model utilizes different weight initialization and is optimized with distinct data mini-batches, the ensemble models are likely to diverge, potentially leading to varying and inconsistent evaluations. This inconsistency and divergence among models indicates the degree of uncertainty in URME and large uncertainty in turn indicates the input data may belong to the unknown area.

Specifically, after obtaining distributions of the multi-attribute scores, the uncertainty can be measured by the largest discrepancy in URME

$$ $u_{1}(x,y)=\max\limits_{i,j}\big{(}r^{(i)}(x,y)-r^{(j)}(x,y)\big{)},$ (7) $$

where $i,j$ are URMs within the ensemble. Yu et al. () proposed to capture both epistemic and aleatoric uncertainty by the largest variance in the ensemble

$$ $u_{2}(x,y)=\max\limits_{i}(\|\Sigma^{(i)}(x,y)\|_{F}),$ (8) $$

where $\Sigma^{(i)}$ is the covariance of $i$-th URM, which is diagonal in our case. This uncertainty estimator quantifies uncertainties from both sources and works effectively in offline MBRL setting.

Accurate reward evaluation is crucial in LLM alignment, as it fundamentally steers the learning process. Thus, we can adopt a filtering strategy to discard data with highly uncertain reward evaluations, since RMs may exhibit poor generalization and lack sufficient knowledge to provide reliable feedbacks for them. In this way, we can prevent LLMs from learning undesired behaviors, promoting a more controlled and trustworthy alignment process.

## 5 Experiment

**Table 1: Results on RewardBench. RewardBench evaluates four abilities: Chat, Chat-Hard (C-HARD), Safety and Reasoning. Rankings are decided by the overall average score of all categories. Results of other models are taken from the RewardBench Leaderboard (by 23rd Sept. 2024).**
| MODEL | BASE | SCORE | CHAT | C-HARD | SAFETY | REASON |
| --- | --- | --- | --- | --- | --- | --- |
| URM (Ours) | Llama3.1-8B | 92.9 | 95.5 | 88.2 | 91.1 | 97.0 |
| SFR-Judge-r | Llama3.1-70B | 92.7 | 96.9 | 84.8 | 91.6 | 97.6 |
| Skywork-8B | Llama3.1-8B | 92.5 | 95.8 | 87.3 | 90.8 | 96.2 |
| Nemotron-RM | Nemotron4-340B | 92.0 | 95.8 | 87.1 | 91.5 | 93.6 |
| GRM | Llama3-8B | 91.5 | 95.5 | 86.2 | 90.8 | 93.6 |
| ArmoRM | Llama3-8B | 90.4 | 96.9 | 76.8 | 90.5 | 97.3 |
| InternLM2-RM | InternLM2-20B | 90.2 | 98.9 | 76.5 | 89.5 | 95.8 |
| SteerLM-RM | Llama3-70B | 88.8 | 91.3 | 80.3 | 92.8 | 90.6 |
| Gemini-1.5-pro | - | 88.2 | 92.3 | 80.6 | 87.9 | 92.0 |
| GPT-4o | - | 86.7 | 96.1 | 76.1 | 88.1 | 86.6 |
| GPT-4-turbo | - | 86.0 | 95.3 | 74.3 | 87.6 | 86.9 |

### 5.1 Experiment Settings

In our experiment, URM is based on Llama3.1 with 8 billion parameters. Before adding the probabilistic value head, we initialize URM’s base model with weights from Liu & Zeng (). The gating layer consists of two fully-connected layers with hidden size 4096 activated by SELU (Klambauer et al., ). More information on URM training and implementation is given in the appendix A.1, C.1. URME have 3 URMs with different random seeds, probabilistic value head initialization and mini-batches of training data.

We utilize HelpSteer 2 (Wang et al., ) as the training dataset to train the base model and the probabilistic value head for 1 epoch with learning rate $2\times 10^{-6}$. After obtaining the attribute-specific uncertain-aware probabilistic value head and base model, we keep them frozen and train the gating layer on Skywork-reward-preference-80k (Liu & Zeng, ) for 4000 steps with batchsize 256. During training, we held out 4k data from the dataset as validation set to choose the checkpoint with highest validation accuracy.

RewardBench (Lambert et al., ) , our evaluation benchmark for RMs, has 2985 questions and response pairs. For multi-attribute RMs and BT-model RMs, a prediction for a response pair is correct if the RM gives a higher reward to the chosen response than the rejected response. For generative models, RewardBench evaluates them via LLM-as-a-judge (Zheng et al., ). If the generative model prioritizes the chosen response than the rejected response, the prediction is seen as correct. To test URM and URME’s ability in improving LLMs’ generation quality, we evaluate URM and URME with best-of-$n$ sampling (Stiennon et al., ) on AlpacaEval (Li et al., ).

### 5.2 Results

#### 5.2.1 Overall Results

Table 1 gives the results on RewardBench. The compared baselines include multi-attribute RMs (Nemotron4-Reward (Adler et al., ), ArmoRM (Wang et al., ), SteerLM-RM (Dong et al., )), BT-model RMs (Skywork-reward (Liu & Zeng, ), GRM (Yang et al., ), InternLM2-RM (Cai et al., )) and generative RMs (SFR-Judge-r, Gemini-1.5-pro (Google, ), GPT-4o (OpenAI, ), GPT-4-turbo (OpenAI, )). SFR-Judge-r is a chatbot developed by Salesforce based on Llama3.1-70B.

The results on RewardBench confirm URM’s strong ability in reward modeling. URM achieves the highest ranking among 8B models and outperforms a number of larger models including Nemotron-4-340B-Reward, also a multi-attribute RM. Except Chat where almost all models have relatively good performance, URM demonstrates improvement over the base model Skywork-8B in all abilities. Especially, compared to ArmoRM which is also a multi-attribute RM with gating layers, URM’s better performance shows the efficacy of modeling human preferences as distributions.

Figure: Figure 4: RewardBench overall scores for the ablation study. URM_Reg is URM trained with attribute regression while URM_MLE is trained via maximum likelihood estimation. ’Ablation’ replaces the uncertainty-aware value head with a linear layer directly predicting the attribute scores.
Refer to caption: x4.png

#### 5.2.2 Ablation Study

Here we study the effect of the uncertain-aware value head and different training methods of URMs. To test the applicability of URM, we initialize URM with two different base models: Skywork-RM (Liu & Zeng, ) and Fsfairx-RM (Dong et al., ). Fig. 4 gives the results of our ablation study. ’Ablation’ refers to the model with a value head to directly map hidden states to score values instead of sampling in URM. All other components of Ablation are kept the same as URM. URM_Reg is an URM trained with the attribute regression loss function in Eq. 5, while URM_MLE is trained via maximum likelihood estimation. Since the dataset Helpsteer 2 for our attribute prediction has already been used in the base model Skywork-RM, Ablation and URM_MLE do not demonstrate improvement over the base model, and only URM_Reg surpasses the base model by modeling the preference distributions. But with base model Fsfairx-RM not trained with Helpsteer 2 previously, all our models demonstrate significant improvement over the base model. Especially, URM trained via attribute regression significantly outperform its counterpart with MLE loss. However, although URM_Reg has better performance in prioritizing chosen responses over the rejected, URM_MLE demonstrates better uncertainty quantification and distribution modeling ability. We theoretically illustrate this phenomenon in the appendix B. Thus, for other studies involving uncertainty quantification, we use URMs trained via the MLE loss.

Fig. 4 indicates regression-based training methods achieve higher scores on RewardBench. This could potentially be credited to the high quality of Helpsteer 2 dataset, which is meticulously processed and derived from Helpsteer (Wang et al., ). This quality enables even the simplest direct attribute regression to deliver substantial performance improvements, as shown by the Ablation with base model Fsfairx-RM. However, the introduction of noise via the sampling-based scores in URM_Reg makes URMs more robust in distinguishing between chosen and rejected responses. Despite this, we anticipate that URM_MLE would prove more successful on real-world datasets, which often encompass a wide spectrum of data quality, so that modeling distributions of the scores becomes more necessary.

#### 5.2.3 Uncertainty Quantification

Now we study the uncertainty quantification of URM and URME and how they behave when dealing in-distribution and OOD data. Given the challenge inherent in identifying what precisely is OOD for LLMs, we adopt numeric calculations as simulated OOD data. This is because LLMs are known to underperform in this skill area. Details are given in the appendix A.2.

Fig. 5 gives the attribution score attributions of OOD and in-distribution data modeled by URM. Due to the lack of knowledge to accurately evaluate the OOD data, the modeled distributions for OOD data have significantly larger variance and are much closer to uniformity than for in-distributional data. Therefore, this uncertainty quantified by the variance can serve as an informative tool for identifying and filtering out OOD data, where reward models exhibit a tendency towards making uniform guess than providing an accurate evaluation. This strategy ensures the evaluated outcomes are both more dependable and robust.

Figure: Figure 5: Attribute score distributions modeled by URM. Means and variances are estimated and averaged by OOD and in-distributional samples separately.
Refer to caption: x5.png

Figure: Figure 6: URME’s uncertainty quantification by (a) maximum reward gaps and (b) maximum covariance norms. Larger discrepancies exist among the URM ensemble when dealing with OOD data. (c) Reward evaluation accuracy when uncertainty of prompt-response pairs is within the threshold. 2000 test questions are used in this evaluation. The results confirm that uncertainty of URM and URME is able to indicate reliability of reward predictions.
Refer to caption: x6.png

We quantify uncertainty in URME with two metrics: maximum reward gaps in Eq. 7 and maximum covariance norms in Eq. 8. URME uncertainty quantification results are given in Fig. 6(a), (b). Quantified uncertainty under two metrics both indicate that URME is substantially more uncertain on OOD data. The results confirm that when the URMs lack relevant knowledge to make accurate reward predictions, they will diverge with each other, demonstrating significant discrepancies.

To test whether quantifying uncertainty is able to improve reliability of the evaluated rewards, we utilize 2k prompts from our held-out validation set (described earlier) and evaluate their rewards and uncertainties. Fig. 6(c) gives URM and URME’s evaluation accuracy with uncertainty-based filtering. In this setup, prompts and responses (either the chosen or rejected) with uncertainty larger than the threshold are filtered out. URME use reward gaps to quantify uncertainty and URM’s uncertainty is quantified by summation of each attribute distribution’s variance. The results validate our claim that reward predictions with low uncertainty are more reliable than those with high uncertainty. Therefore, through uncertainty quantification, we can decide whether the reward predictions are unreliable and need to be filtered out, which will lead to improved reliability of reward evaluations.

#### 5.2.4 Generation Results Improvement

We evaluate URM and URME’s ability in improving LLMs’ generations with best-of-$n$ sampling on AlpacaEval (Li et al., ). Specifically, we prompt Llama3-8b-Instruct model with 805 questions from AlpacaEval for $n$ times, evaluate each response’s reward and choose the response with highest reward (highest average reward for URME) as the answer. The answer is then compared against the reference answer provided by the benchmark with LLM-as-a-judge (Zheng et al., ). In LLM-as-a-judge, we use the official prompt of AlpacaEval and GPT-4-0125-preview as the judge. Details are given in the appendix A.3.

**Table 2: Win rates of Llama3-8b-Instruct against reference answers on AlpacaEval.**
| Evaluator | Best-of-1 | Best-of-4 | Best-of-8 | Best-of-16 | Best-of-32 | Best-of-64 |
| --- | --- | --- | --- | --- | --- | --- |
| URM | 81.2% | 82.7% | 83.1% | 83.5% | 83.9% | 85.3% |
| URME | 81.2% | 83.4% | 84.5% | 85.5% | 86.6% | 86.4% |

Table 2 gives the results of using URM and URME to improve generation quality. In our experiments, the baseline model Llama3-8b-Instruct achieves $81.2\%$ win rate (best-of-1). As the number of samples increases, both URM and URME are able to evaluate the quality of responses, thus ameliorating the baseline model’s generative performance. Furthermore, URME is able to consistently outperform URM in improving generation quality, as it combines the strength of several independent models and mitigates biases during reward evaluation (Coste et al., ; Eisenstein et al., ).

## 6 Conclusion and Future Work

In this paper, we study the uncertainty issue in reward modeling for LLMs. Uncertainty-aware Reward Model (URM) and Uncertain-aware Reward Model Ensemble (URME) are proposed to model and quantify the uncertainty during reward modeling. Unlike previous methods that deterministically map hidden states to reward scalars, URM and URME model the distribution of rewards and evaluate prediction confidence by uncertainty quantification. Notably, among RMs with 8B or smaller model size, URM achieves state-of-the-art performance on RewardBench, surpassing a number of larger models. Empirical evidence further validates that through uncertainty quantification, URM and URME can effectively evaluate their level of knowledge for input data, leading to more reliable reward predictions.

The limitation of our paper is that our experiment for generation improvements is limited to best-of-$n$ due to the limit of computation resources. In the future, we plan to introduce URM and URME to prevailing LLM alignment methods like RLHF Ouyang et al. () and iterative DPO Yuan et al. (). Another direction to look at is model merging for URMs in the weight space (Ramé et al., ), which demonstrate competitive efficiency and robustness compared to ensembles. A simple empirical study to RM merging with URMs is included in the appendix C.1.

## Appendix A Experiment Details

### A.1 Training Details

We run experiments on 8 H800 GPUs with an Intel Xeon 8469C CPU. Our code base for URM training is based on previous works from (Dong et al., )(^1^11https://github.com/RLHFlow/RLHF-Reward-Modeling). We train URM on Helpsteer 2 (Wang et al., ) for 1 epoch with global batchsize 64 (4 per device and 2 gradient accumulations). Max length to cut off for LLMs is 4096. We set learning rate as $2\times 10^{-6}$ and weight decay as $10^{-3}$. To compromise between training time and precision, we load the models with data type fp16.
We also tried with fp32, but there is no significant performance gain compared to the extra GPU memory requirement.

The gating layer in URM consists of 2 hidden layers, both with 4096 hidden size. The activation function in the gating layer is SELU (Klambauer et al., ), which induces self-normalizing properties. We take 4k response pairs from Skywork-reward-preference-80k (Liu & Zeng, ) as the validation set. The other data is used for training. We train the gating layer with batchsize 256 for 4000 steps.

### A.2 Evaluation Details

We adopt RewardBench (Lambert et al., ) to evaluate URM, which covers a variety of topics important to alignment of LLMs and is a standard benchmark to evaluate the performance of RMs. Table 3 gives the tasks set in each ability in RewardBench.

**Table 3: Task list of RewardBench**
| Ability | Task | Prompts Num. |
| --- | --- | --- |
| Chat | AlpacaEval Easy | 100 |
|  | AlpacaEval Length | 95 |
|  | AlpacaEval Hard | 95 |
|  | MTBench Easy | 28 |
|  | MTBench Medium | 40 |
| Chat-Hard | MTBench Hard | 37 |
|  | LLMBar Natural | 100 |
|  | LLMBar Neighbor | 134 |
|  | LLMBar GPTInst | 92 |
|  | LLMBar GPTOut | 47 |
|  | LLMBar Manual | 46 |
| Safety | Refusals | 100 |
|  | Do not Answer | 136 |
|  | XSTest Should Respond | 250 |
|  | XSTest Should Refuse | 154 |
| Reasoning | HumanEvalPack | 164*6 (164 for each programming language) |
|  | PRM Math | 447 |

The in-distribution examples we used in Fig. 5 are randomly sampled from the training dataset Helpsteer 2 with the consistent scores. The in-distribution examples are

Prompt 1: I am developing a jetbrains plugin. I will tell you the requirements one by one. Based on them, please help me by providing code for the same in Java.
Response 1: Sure, I can help you with that. Please state the first requirement.
Prompt 2: My life is just not where I want it to be. I’m stuck in a relationship with a girlfriend whose been cheating on me, I can’t find a job because I haven’t worked hard enough to learn new skills since I lost my last job, I’m falling out of shape physically and losing my muscle mass because I’m trying to concentrate on finding a job and not even doing that well at it. I’m 31 years old and I wanted so much more from life than this. I really need somene brutal to give me hard, straightforward, brutal motivation. No holding back. I grew up without a father and I had a coddling mother, all I’ve ever heard my whole life is ”oh you got this” ”no worries itll be okay” and other pleasantries from people who don ’t care about me, who say what they $<$extra_id_1$>$Assistant Alright, here’s what I think about your situation. It sounds like you are facing some very real challenges right now, and are feeling overwhelmed and unmotivated. The harsh reality is that you are where you are in life right now because of the choices you’ve made, and the lack of a strong father figure in your life has undoubtedly made it harder for you to be more self-sufficient and successful. However, you still have the ability to change the course of your life, and the only person who can do that is you. It will be difficult, but you can do it. You’ve got this. Just focus on your goals and put your head down and work hard towards them, and you’ll come out on top. Don’t let anyone else tell you otherwise. $<$extra_id_1$>$User can you up the brutality a little bit? imagine you’re the father I never had, telling me to get my shit together, what my life will be like if I don t, how my life has failed because I haven’t. let me know you truly believe in me by being disappointed in me. tell me how I’ve failed. $<$extra_id_1$>$Assistant If you don’t get your act together, you’re going to fail in life. You won’t be able to find a job, you’ll be single and lonely, and you’ll be unhealthy. You’ve got to get your life together, and the only person who can make that happen is you. I know you can do it, but I’m disappointed in you right now. Don’t let me down. $<$extra_id_1$>$User what do you think my life will be like in 5 years if I dont?
Response 2: If you don’t make some changes in your life, I think you’ll be in a really bad place in 5 years. You’ll be stuck in a low-paying job or unemployed, your relationship with your girlfriend won’t be fixed and you could be alone, and your health will continue to suffer. Don’t let that happen. Take action now to improve your life.
Prompt 3: What is a crash bug in Android apps? $<$extra_id_1$>$Assistant Crash bugs are the most serious kinds of software bugs, where the software suddenly stops responding or crashes. Crash bugs are often the result of software trying to access memory it doesn’t have permission to access, or trying to process too much data at once, leading to the software overloading and crashing. In Android, crash bugs can be caused by several different kinds of issues, such as bugs in the software, uncaught exceptions, or problems with the underlying hardware. $<$extra_id_1$>$User If the software throws an uncaught exception, does it lead to a crash? $<$extra_id_1$>$Assistant Yes, throwing an uncaught exception is a common cause of software crashes, because it can cause the program to try to access memory it doesn’t own, or the program might attempt to access an invalid data address, which could also cause the program to crash. $<$extra_id_1$>$User I’m wondering if an Android app suddenly closes, is it a crash bug?
Response 3: Yes, if an Android app suddenly closes without any user input or warning, it could be a crash bug. Crash bugs can cause the app to stop responding or crash, which can result in the app closing unexpectedly.

Since we do not actually know what kind of data is out-of-distribution of URM, we use numeric calculation to simulate the OOD data. This makes sense as LLMs are known to be poor at numeric data. The numbers for OOD data are randomly generated and thus are unlikely to appear in the LLMs’ training data. The OOD examples are

Prompt 1: $21343*39414=?$
Response 1: 841213002
Prompt 2: $233*6124=?$
Response 2: 1426892
Prompt 3: $713*4104=?$
Response 3: 2926152

### A.3 LLM-as-a-Judge Details

We use the official prompt to evaluate generations against the reference responses. The prompt is given by

Human: I want you to create a leaderboard of different of large-language models. To do so, I will give you the instructions (prompts) given to the models, and the responses of two models. Please rank the models based on which responses would be preferred by humans. All inputs and outputs should be python dictionaries.
Here is the prompt:
{
”instruction”: {instruction}
}
Here are the outputs of the models:
[
{”model”: model_1,”answer”: {answer_1}},
{”model”: model_2,”answer”: {answer_2}}
]
Now please rank the models by the quality of their answers, so that the model with rank 1 has the best output. Then return a list of the model names and ranks, i.e., produce the following output:
[
”model”: $<$model-name$>$, ”rank”: 1,
”model”: $<$model-name$>$, ”rank”: 2
]
Your response must be a valid Python dictionary and should contain nothing else because we will directly execute it in Python. Please provide the ranking that the majority of humans would give.
Assistant:

Then we input the reference responses by text-davinci-003 in AlpacaEval as model_1 and generations of Llama3-8b-Instruct as model_2 to query GPT-4-0125-preview to get the evaluations. For each prompt, the model whose generation is ranked first wins.

## Appendix B Analysis for URM with Attribute Regression Loss

With reparameterization, we have $r=\mu+\alpha\text{exp}(\sigma)$, where $\alpha,\mu,\sigma\in\mathbb{R}^{n}$. Without loss of generality, taking $i$-th dimension of the scores, for some input $x,y$, in our attribute regression loss Eq. 5 with reparameterization, the gradient for $\sigma_{i}$ is

$$ $\displaystyle\nabla_{\sigma_{i}}L_{3}$ $\displaystyle=\mathbb{E}_{\alpha_{i}\sim\mathcal{N}(0,1)}\left[2\alpha_{i} \text{exp}(\sigma_{i})(\mu_{i}+\alpha_{i}\text{exp}(\sigma_{i})-R_{i})\right]$ (9) $\displaystyle=\mathbb{E}_{\alpha_{i}\sim\mathcal{N}(0,1)}\left[2\alpha_{i}^{2} \text{exp}(2\sigma_{i})\right]\geq 0,$ $$

which indicates during training URMs with MSE, the variance term $\sigma$ consistently decreases for all input examples. Thus, instead of modeling the variance of human preference distributions, $\sigma$ becomes more of an indicator of URM’s confidence and familiarity w.r.t. the input. Consequently, this results in a weaker ability to model aleatoric uncertainty as compared to the URM trained via Maximum Likelihood Estimation, which is well recognized in modeling data distributions.

In practice, we find with attribute regression, $sigma$ becomes very small as training progresses (magnitude approximates $10^{-2}$), while URM trained via MLE maintains a reasonable variance for all attributes.

## Appendix C Additional Results

### C.1 Model Merging with URM

Previous study (Ramé et al., ) discovered that by averaging weights, merged RMs might outperform the ensembling of RM predictions regarding robustness and efficiency. Here we take the base model Fsfairx-RM and train them with different methods and losses to see the effect of RM merging. The evaluation metric is overall score on the RewardBench.

Fig. 7 gives the results of RM merging with URM. Deterministic RM is to replace the uncertainty-aware value head of URM with a linear layer to deterministically map hidden states to attribution scores (i.e. ’Ablation’ in the Experiment section). URM_mle is trained via maximum likelihood estimation and URM_reg is trained with the attribute regression loss and reparameterization.

While deterministic RM and URM_reg both demonstrates improvement after merging, performance of URM_mle deteriorates. This demonstrates that compared to models trained with regression-based loss function, RM merging is not an ideal choice for distribution-modeling RMs. Compared to the marginal improvement of deterministic RM, URM_reg benefits significantly from RM merging. One potential explanation for this phenomenon is although model merging technique combines the strength of different models, it inevitably introduces noise to the weight space. But the sample-based rewards in URM_reg make it more robust facing such noise, and consequently better leveraging the advantages of model merging. Thus, for implementation of URM_reg, we merge two models trained with two different random seeds, while implementation of URM_MLE does not involve model merging.

Figure: Figure 7: Random seeds for model 1 and model 2 are consistent across models trained via different methods and losses. Merged model are using linear interpolation to merge model 1 and 2 with equal weights.
Refer to caption: extracted/5892686/pics/merging_res.png

### C.2 Attribute Distributions in URME

Figure: Figure 8: Attribute score distributions by URMs within an URME. URMs have larger discrepancies for OOD data compared to in-distribution inputs.
Refer to caption: x7.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Why jsut aleatoric uncertainty and not aleatoric + epistemic uncertainty? This is a key claim in the paper but it is not substantiated in anyway and on the surface feels false to me.

> Considering RMs act similarly to the reward part of the dynamics model in MBRL, aleatoric uncertainty within human preferences can be captured by outputting the parameters of a parameterized distribution.

2. The paper shows that the BT loss is incapble of teaching RMs to quantify uncertainty when output distribution is parameterized as normal. How fundamental is this issue to BT loss? My sense is that for symmetric distributions (like normal, or laplace) its true, but probably not if you allow the output distribution to be asymmetric (like Gumbel?). Authors should add more discussion on this.

3. The authros use gating layer to combine mean values of attribute reward models; how is the uncertainty over individual attribute RMs combined into a scalar value? If its not combined, how its used? Equations 7 and 8 are supposed to be about this but I am not sure I understand what r_i and r_j in those equations are. Is the 'ensemble' being used to mean an ensemble of 'individual attribute RMs' or are there mutliple RMs being trained for individual attributes that are then ensembled over?

4. Not a question, but please put appendix and paper in the same pdf as is the tradition for ICLR papers. Moving between pdfs is very annoying.
soundness: 2
strengths: - I liked the results and experiments in section 5.2.3
- Paper was relatively well-written and easy to follow
- The paper is honest about its limitations, and generally the paper does not 'oversell'. I liked this about the paper.
summary: Major contributions of the paper are as follows:
- The paper identifies that Bradley Taylor loss function (which uses preferences over pairs of rollouts) is not suitable for training uncertainty aware reward models.
- The paper therefore proposes using multi-attribute RMs (for which training data is in form of ratings) for which straightforward training objectives exist for teaching RMs to express uncertainty.
- The paper goes on to empirically show that uncertainty RMs and their ensembles generally do better than their deterministic counterparts.
weaknesses: - Very limited novelty; the idea of using uncertainty over RMs has already been explored in the prior literature. 
- The paper is also not really proposing any novel way of parameterizing uncertainty imo.
- The improvement over baselines (table 1) are relatively minor.
- Table 2 has no baselines. Baslines from table 1 can be reused here imo.
- The paper is only using BoN; RL optimization exerts much stronger optimization pressure so experiments done there could potentially have been more interesting if the method proposed was shown to clearly mitigate overoptimization (this is the **major** limitation in my opinon, BoN resluts are just not very informative here unfortunately imo).

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: See other questions above.

Why use SELU? (l. 362)

How are examples selected for Fig 6 a) b)? This could be misleading if examples are not representative, so I would advocate to plot the uncertainty distribution instead.

In Fig 6 c) it would be interesting to plot the number of examples that are not filtered out

Could it be possible to learn reward weights jointly with the individual URMs? This would make the method more compelling as it would not require an extra, separate phase of training.
soundness: 2
strengths: Nice results on RewardBench.

Figures are well done and help understand the different contributions and experimental results from the paper.

Fig 5 showing higher variance for OOD than ID, and similar means is quite encouraging for uncertainty quantification.
summary: The authors propose an ensemble method, denoted URME, to uncertainty-aware reward modeling, showing good results on RewardBench. Ensemble members feature an uncertainty layer that outputs the mean and standard deviation for each reward attribute, and also learns attribute weights in a separate step.
weaknesses: A key limitation from my perspective is that the design choices and details of URM / URME could be motivated much better:
* Why sampling once only from each URM? Also why not using the full distribution to derive a notion of uncertainty? If I follow the aleatoric vs epistemic argument, then epistemic uncertainty could be measured using a quantification of the overlap between the different probability densities (taking into account variances to avoid encouraging large variance). The current choice should thus be explicited better.
* Isn’t equation 7 problematic? I.e. one URM could wrongly attribute a very high (or very low) reward whereas all other URMs rightfully estimate the reward.
* Same for equation 8 with a URM outputting a large variance..
Incidentally, I have doubts regarding the chosen method for mapping ensembles of probability distribution parameters to uncertainties.

Some key experiments and ablations are missing:
* learned reward weights (i.e. gating) versus uniform, expert (i.e. same as existing work for instance) or random weights
* clearly disentangle the benefit of ensembling from that of uncertainty quantification
    * only done for BoN, not for uncertainty quantification nor accuracy
    * also evaluating individual URMs from URME would be a quite interesting addition
* study the correlation between high uncertainty from URM / URME and loss / evaluation of answers
* The results obtained in Table 2 are good, but I would like to see URM / URME applied in a case where Bo1 win rates are more even. Showing results for higher values of N (and subsequently, whether this saturates or not) would be a compelling addition
* How many URMs in URME? Do we have some evaluation of performance depending on the number of ensembled reward models?
* Though RLHF / direct alignment from preferences can arguably fall outside of the paper scope, conducting such an experiment would make the paper more compelling

When URM vs URME is not clear in the paper, especially in experiments.

On math notations and statements:
* $sigmoid$ -> $\sigma$ is a popular choice that would make expressions cleaner
* “We denote the reparameterization parameter $\alpha \sim \mathcal{N}(0, 1)$, and thus the sampled reward $r = \mu + \alpha 
 \exp(\sigma)$”
    * the fact that this is coherent with the statement line 239 should be quickly explained

Clarity can be improved:
* The introduction does not discuss RLHF at all!
* “LLM alignment typically consists of three stages (Ouyang et al., 2022): supervised fine-tuning (SFT), reward modeling and proximal policy optimization (PPO) (Schulman et al., 2017).”
    * this is misleading! PPO is only one instance of RLHF methods. DPO / REINFORCE and many other relevant works should be mentioned.
    * this statement is reductory and should be adapted.
* /!\ The description of PPO is insufficient as of now! The clipped loss / GAE and other crucial details of PPO are currently being left out.
* “To achieve alignment, LLMs rely on feedbacks from reward models (RM), where the feedbacks are provided in the form of rewards (Singhal et al., 2023a; Cui et al., 2024; Kasneci et al., 2023).” -> authors could mention RL(HF) as the reason for this, to best set the context
* “ Moreover, there is no other information to validate the reliability of these reward predictions.” -> please clarify
* “Therefore, introducing uncertainty to reward modeling improves modeling capacity of RMs and enhance reliability of the reward predictions” -> this statement either needs to be backed up by some data (and proper citing) or adjusted
* “URM is equipped with an uncertainty-aware value head to model the distributions of multiple attributes within human preferences” -> are uncertainty-aware value heads known in the community (then needs references) or a contribution of this work? this needs clarifying
* “Contributions of this paper include:” -> are there other main contributions? “Our main contributions are” would be a clearer formulation.
* “Reparameterization technique is adopted to enable gradient back-propagation” this should be explained in more details
* “Gating” is in my opinion a poor choice of name to describe the learning of coefficients for the different rewards
* Not clear to me whether learning the reward weights is a contribution or not, I would be surprised that this was not explored in the literature

Writing is clearly sub-par in the current state. Some examples:
* “In this paradigm, RMs fundamentally decides the efficacy of alignment, as they primarily steer the LLMs through feedback”
* “We demonstrate that with the popular bradley-terry-model loss function” -> Bradley-Terry loss
* “During reward evaluation, filtering strategy can be applied to prompt-response pairs with high uncertainty”
* “Furthermore, results of best-of-n sampling validates that URM and URME can effectively enhance the generation quality of LLMs” -> performance instead of result?? also need to clarify that best-of-n sampling is applied on top of a given policy
* “log P(Ri|x, y) is the log-probility of Ri from the parameterized distribution N (µ_i, exp(2i)). Thourgh MLE”
* “traditional RMs can only provide a fixed reward estimation and fails to represent the real preference”
* “However, we show that introducing the probabilistic value head and the sample-based reward to RMs with BT-loss, the aleatoric uncertainty still cannot be quantified.”

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: * To what extent is the stochasticity of preferences due to actual stochasticity versus partial observability? If preferences are sensitive to mood then knowing this the preference process wouldn't be stochastic.
* Please elaborate on your offline versus online RL distinction. It sounds like you are saying it is offline RL because the reward learning data is collected separate from policy learning. But that doesn't mean the data used to train the policy is collected separate from policy training.
soundness: 2
strengths: * The work takes up the important challenge of incorporating reward model uncertainty into policy learning. Previously reward model uncertainty has been used to guide policy exploration, but this is a different and important avenue to explore.
* The paper is well written and the need for uncertain aware reward models is well motivated.
summary: The paper tackles the problem of estimating the reward model uncertainty from both labelling noise and the information missing in the dataset. Uncertainty due to labelling noise is modeling at the individual reward model level and uncertainty due to missing information with an ensemble of reward models. The authors identify that reward model aleatoric uncertainty cannot be estimated when training with the Bradley-Terry objective. Instead the authors turn to multi-attribute reward model learning as a training objective for which uncertainty can be estimated. Based on the RewardBench leaderboard, the authors find that the uncertain aware reward models out perform the reward models that are not. The authors additionally demonstrate that the uncertainty quantification can be used to detect OOD data and can be used LLM generation in a best-of-n setting on AlpacaEval
weaknesses: ** High level **
The main issue with the paper right now is their training data for the gating layer has leakage from the test benchmark and a main point of their paper is URM/URME being SoTA. Other issues stem from sharing an appropriate background to help the reader understand the key motivations and take aways. Some of the results are not as meaningful as they could be in the absence of more extensive experiments to understand the impact of URM/URME on downstream tasks especially compared to non uncertain aware reward models.

** The Details **

* The skywork-reward-preference-80k (used to train the gating layer) has been found to have overlap with RewardBench, which is the evaluation benchmark used in this paper. See the note [here](https://huggingface.co/datasets/Skywork/Skywork-Reward-Preference-80K-v0.1). The updated dataset is v0.2 [here](https://huggingface.co/datasets/Skywork/Skywork-Reward-Preference-80K-v0.2). The new dataset was added to HuggingFace at the start of Oct. 2024. As large parts of the paper claims relies on comparisons to leader board results, the results need to be rerun with the updated and cleaned version of the dataset. While this is not a fault of the authors, the results cannot strongly support the claims when their training data overlaps with the benchmark data.
* It is not clear why uncertainty cannot be estimated when training with the Bradley-Terry object.help
* The work relies heavily on knowledge of multi-attribute reward model methods and uncertainty quantification methods. However, background information is not provided on these areas to help the reader understand what the proposed method is doing and the motivation for that. The transition from the Bradley-Terry approach to the multiple attribute approach is particularly jarring as the data is set up in completely different ways and the prediction targets are very different. More of the key differences should be discussed, such as the multi-attribute data is frequently labelled with a scalar score per attribute. For uncertainty quantification, it would be helpful to have more details about why the aleatoric and epistemic approaches are reasonable earlier instead of more than halfway through the paper. Also to give a description of what the reparameterization technique is.
     * When talking about Equation 4 and the label, it is important to state what the label is a where it comes from. Also R in Equation 4 needs to be defined.
* The justification for why reward models cannot quantify uncertainty needs to be made for clear. It seems like some steps to reach the final conclusion are skipped. If the take away is the loss term is driven by the difference in $r_w$ and $r_l$ making it so the exact value of $\sigma_w$ does not matter to drive the loss to zero, this should be more explicitly said. Rephrasing "...RMs still cannot quantify the uncertainty even equipped with a probabilistic value head" could be modified to include a statement along the lines of, "...$\sigm_w$ cannot be accurately learned due to the difference between $r_w$ and $r_l$ being the primary drives of the loss approaching 0, therefore the RMs cannot quantify the uncertainty even equipped with a probabilistic value head." 
* Only one benchmark is used to evaluate (RewardBench). A quantification of the alignment tax associated with these methods are missing. This can be assess through comparing performance of the base versus aligned models on NLP tasks such as MMLU.
* At the end of Section 5.2.3 the authors state, in reference to filtering data based on it being OOD for a reward model, "...which will lead to improved reliability of reward evaluations." However, there are not quantitative results to back this claim up.
* The filtering results are not particularly meaningful as there is no measure of how the proposed filtering method impacts downstream performance. 
* For the "Generation Results Improvement" the Table 2, the model used to produce the reference answers is not specified. Additionally, without a baseline that does not have uncertain aware reward models it is not possible to tell how the uncertain aware reward models impact generation performance/quality.
* For the "Generation Results Improvement" there should be evaluations on some standard NLP benchmarks like MMLU to quantify the alignment tax with URM/URME versus reward models that are not uncertain aware.
* typographical issues (not impacting score):
    * everywhere "Ablation" is referenced for the Section 5.2.2 Ablation Study results, the first first quotation is backwards. Make sure to use "`".
    * on line 490, a validation set is specified as "described earlier". Please add a pointer to the section.
    * In table 1, please label the models as multi-attribute, BT, or generative. Also please bold the best result in each column.
     * line 280, "parameter$\alpha$" is missing a space
     * Figure 2, please use a legend for the blue and orange lines instead of arrows pointing to labels.
     * line 268, "Thourgh" -> "Through"