Mitigating Reward Overoptimization via Lightweight Uncertainty Estimation

Authors: xiaoying  zhang
jean-françois  ton
wei  shen
hongning  wang
yang  liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 03:19:30

Abstract: Reinforcement Learning from Human Feedback (RLHF) has been pivotal in aligning Large Language Models with human values but often suffers from overoptimization due to its reliance on a proxy reward model. To mitigate this limitation, we first propose a lightweight uncertainty quantification method that assesses the reliability of the proxy reward using only the last layer embeddings of the reward model. Enabled by this efficient uncertainty quantification method, we formulate AdvPO, a distributionally robust optimization procedure to tackle the reward overoptimization problem in RLHF. Through extensive experiments on the Anthropic HH and TL;DR summarization datasets, we verify the effectiveness of AdvPO in mitigating the overoptimization problem, resulting in enhanced RLHF performance as evaluated through human-assisted evaluation.

Main paper: 
 introduction: 
 Reinforcement Learning from Human Feedback (RLHF) is proven to be effective for aligning Large Language Models (LLMs) with human preferences [26, 2] . RLHF typically involves three main steps: 1) Supervised Fine Tuning (SFT) of a pretrained LLM using high-quality data, 2) Reward Modelling to capture human preferences that the LLM should follow, and 3) Reinforcement Learning (RL) based policy optimization (e.g., PPO [32] ) where a policy initialized from the SFT model is further improved, guided by the reward model as a proxy for human feedback.However, a critical issue arises in this process: the reward model, built from a finite dataset of human preferences, often fails to accurately represent the underlying human preferences. This approximation error, worsened by the distribution shifts during policy updates [38] , leads to unreliable rewards during the RL stage. This directly causes the phenomenon of reward 'overoptimization', wherein the LLM exploits erroneous high-reward states, artificially inflating the estimated proxy reward, while the ground-truth reward decreases [12, 9, 11] .Current mitigation strategies against reward overoptimization, as proposed in [9, 11, 45] , focus on penalizing samples with high reward uncertainty during RL-based policy training. Specifically, [9, 11] quantify reward uncertainty by training an ensemble of reward models with different seeds during either the pre-training or fine-tuning phases. They then measure the variance in estimated rewards across the ensemble to assess uncertainty in the reward prediction. However, RL policy optimization necessitates maintaining multiple reward models in memory, rendering it impractical for large models in real-world applications and limiting the potential to achieve maximum performance, especially since 'scaling laws' typically favour larger reward models [12, 40] .To reduce the memory footprint, recent works [45] instead suggest training multiple LoRA-based reward models with a diversity penalty for uncertainty quantification. However, even though LoRA ensembles reduce memory requirements by only having to train/store an adapter, they still incur high training costs, as well as computational overhead during policy optimization. More specifically, during the PPO stage, [45] requires querying each LoRA ensemble to compute the reward and uncertainty for every sample, which can easily become a serious computational bottleneck.In this paper, we first introduce a lightweight method for quantifying reward uncertainty in the RLHF pipeline, using only the last layer embeddings of the reward model. This approach is easily integrated into any existing trained reward models, making it generally applicable. Building on these uncertainty estimates about reward prediction, we then propose Adversarial Policy Optimization, ADVPO, a distributionally robust optimization procedure to counter overoptimization during policy improvement. ADVPO contrasts with previous sample-wise uncertainty penalization methods [9, 11, 45] , for which we theoretically prove that ADVPO handles reward uncertainty in a less pessimistic manner. As a result, ADVPO is more effective at improving policy and mitigating overoptimization, which we empirically confirm in extensive experiments. These favourable results are further supported through human-assisted assessments.To summarize, our contributions are threefold:• Firstly, we introduce a lightweight method to quantify reward uncertainty using only the last layer embeddings of the reward model into the RLHF pipeline. Extensive experiments confirm its effectiveness in identifying reward uncertainties and signalling overoptimization.• Secondly, we introduce the Adversarial Policy Optimization (ADVPO), built on our efficient uncertainty estimates, to adversarially target the reward model's prediction confidence interval for policy optimization. ADVPO is proven to be less pessimistic than sample-wise uncertainty penalization methods [9, 11] , thus more effective at enhancing the policy and mitigating overoptimization.• Lastly, we empirically demonstrate that ADVPO effectively addresses the reward overoptimization issue on the Anthropic HH [2] and TL;DR summarization [35] datasets. We further validate the learnt LLMs through human-assisted evaluations by comparing ADVPO against existing methods incorporating uncertainty and standard PPO, showcasing its effectiveness in real-world scenarios. 
reinforcement learning from human feedback: 
 We start by providing an overview of the RLHF workflow [26] . This helps us establish the notations and conceptual groundwork necessary for understanding our contributions. RLHF consists of three main steps: 1) Supervised Fine Tuning, 2) Reward Modelling, and 3) RL optimization.Supervised Fine Tuning. RLHF typically begins with Supervised Fine Tuning (SFT), which finetunes a pre-trained LLM through supervised learning on high-quality samples from downstream tasks, such as summarization or dialogue generation. We denote the resulting model as π SFT .Reward Modelling. The second phase of RLHF involves learning a reward model to capture human preferences through annotated data D = {(x i , y i c , y i r )} N i=1 where y i c and y i r denote the chosen and rejected responses to prompt x i . The preferences are assumed to be generated by some unknown reward model r * (x, y) following the Bradley-Terry (BT) model [3] :P * (y c ≻ y r |x) =exp(r * (x, y c )) exp(r * (x, y c )) + exp(r * (x, y r )).Typically, a reward model r φ (x, y) is initialized from a pretrained LLM (usually π SFT ), with an additional projection layer added to map the last embedding layer to a scalar reward. To be more specific, let e(x, y) : X × Y → R d denote the last layer embedding of the prompt and response pair (x, y), and ϕ : R d → R denote the additional projection layer. We define the reward model as r φ (x, y) := ϕ ⊤ e(x, y), where φ includes all the tunable parameters in ϕ and e(x, y).Given the preference data D, the reward model r φ is trained to assign higher reward to the chosen response y c than to the rejected y r , by minimizing the negative log-likelihood of the BT model:L(r φ ) = -E (x,yc,yr)∼D [log (σ (r φ (x, y c ) -r φ (x, y r )))] ,where σ denotes the sigmoid function.RL optimization. Lastly, the learned reward model r φ (x, y) is employed to guide the RL policy optimization phase (e.g., PPO [32] ). Intuitively, the aim is to learn a policy π θ that maximizes the reward r φ while not drifting too far away from π SFT :EQUATIONwhere β controls the deviation from the reference policy π SFT , thus maintaining a balance between reward maximization and adherence to the SFT policy behaviour.A notable limitation of RLHF lies in the fact that the RL process relies on the estimated reward r φ , as opposed to the oracle/gold reward r * . Though widely adopted, it overlooks the potential discrepancies between r φ and r * , which may arise due to inaccuracies during the reward model training. Empirical studies have reported that the RL stage tends to 'hack' the reward such that while the estimated reward (i.e., proxy reward) increases, the oracle/gold reward decreases. This phenomenon is referred to as overoptimization [12, 9, 11, 33, 27] .To mitigate this problem, in addition to the KL penalty in the original RL objective, several recent studies [9, 11, 45] propose to leverage an ensemble of K reward models {r φ k } K k=1 . Given a prompt x and its response y, these methods use the variance of rewards across different reward models to measure uncertainty in the estimated reward, i.e., U x,y = Var({r φ k (x, y)} K k=1 ). The reward is then penalized based on the sample-wise uncertainty before being used in policy optimization:EQUATIONwhere γ controls the degree of uncertainty-based penalty. Intuitively, samples with high uncertainty during policy training are penalized to reduce the risk of being misled by imperfect rewards. However, as previously mentioned, reward ensembles that are trained either by fine-tuning entire LLMs [9, 11] or by using LoRA [45] incur additional memory and computational overhead. This is due to the need of maintaining multiple reward models in memory for policy learning.Thus, an intriguing question arises: Can we quantify reward uncertainty in a lightweight manner that can be easily integrated into any trained reward models, thereby addressing the overoptimization issue without the need for burdensome ensembles? And in the following section, we provide a positive answer to this important question. 
lightweight uncertainty estimation: 
 In this section, we introduce a lightweight uncertainty quantification method, based solely on the final layer embeddings of a reward model. We start by offering a high-level intuition on why the last layer embedding captures essential information about uncertainties in the reward model's predictions. Following this, we present our lightweight uncertainty quantification method.As discussed in Section 2.1, reward modelling can be decomposed into two parts: (1) learning a good representation e(x, y) for the prompt and response pair through a pre-trained LLM; (2) projecting the learnt representation to a scalar reward using a mapping ϕ. Very importantly, the training of LLMs on extensive text corpora, coupled with their vast number of parameters, enables these models to develop versatile representations that can even be used in zero/few-shot tasks [25, 40, 4] , which demonstrate the generalizability of these representations.However, the second part, which involves learning the projection weight ϕ, is strictly tied to the preference data provided during the reward model training. Consequently, the reliability of predicted rewards is closely linked to the accuracy and generalizability of the projection weight.The above claim has been widely supported in the deep learning literature [6, 17, 31, 43] . For instance, [17, 18, 19] demonstrate that by freezing the network up to its last layer and retraining only the projection head with a smaller data set, where spurious correlation is absent, it can greatly improve robustness of the neural network model against these spurious correlations. In the context of language models, recent experiments on weak-to-strong generalization [5] further reinforce this claim. Their findings reveal that even when fine-tuning an LLM's last layer embedding with noisy labels from weak supervision, the model can still excel in subsequent classification tasks if the projection weight is accurately derived from ground-truth labels. This highlights the generalizability and the rich information encapsulated in the last layer representation, accessible by simple projection weights.Building upon the notion of generalized representations with specialized projection weights, we now zoom our attention to the last layer's ability to quantify the uncertainty of its output. The projection weight is strictly estimated based on the preference data encountered during reward model training. Therefore, when evaluating the prompt and response pairs during the RL stage, the pairs might deviate from what was observed during reward model training (suggesting a distribution shift [38] ), hence rendering the predicted rewards unreliable.In the next section, we show how the last layer embedding of a reward model, based on preference data, can act as a feature map for an underlying kernel (similarity measure). This kernel then allows us to determine whether new prompt response pairs are similar to the ones seen during training. If not, the corresponding uncertainty should be higher and penalized during policy optimization.Many methods, derived from a neural network model's final layer embeddings, have demonstrated their effectiveness for quantifying uncertainty in the network predictions, both theoretically and in practice [43, 31] . In this work, we specifically follow the line of uncertainty quantification methods in neural bandits [43] , due to its computational efficiency and theoretical soundness.We first present the following theorem on reward uncertainties when learning rewards through the Bradley-Terry preference model, assuming the reward model architecture is infinitely wide. Theorem 3.1. Assuming the network architecture is infinitely wide and its neural tangent kernel matrix is positive definite, learning rewards through the Bradley-Terry preference model yields the following inequality for the width of the confidence interval of the estimated reward r φ(x, y). Specifically, with probability 1 -δ:EQUATION)where b is a term related to the preference dataset D and δ (typically the smaller δ is, the larger b is), r * and r φ denote the unknown ground-truth reward and estimated reward model parameterized by φ respectively, and M D summarizes all last layer embeddings observed in the preference dataset D for the reward model, i.e., M D = λI + N i=1 y∈{y i c ,y i r } e(x i , y)e(x i , y) ⊤ . Here λ is a regulariser for the inversion.Intuitively, Theorem 3.1 bounds the absolute difference between the predicted reward r φ and the ground-truth reward r * . Consequently, it is natural to define the uncertainty around the predicted reward r φ(x, y) asU CIx,y = b e(x, y) ⊤ M -1 D e(x, y), since a larger difference between r φ and r * implies greater uncertainty. This becomes even clearer when taking a closer look at U CI x,y . When a new prompt-response pair (x, y) is similar to samples in the training data, applying the inverse of M D , which is constructed using the training data, results in a smaller uncertainty U CI x,y . Conversely, if the pair diverges significantly from the training samples, the uncertainty U CIx,y will be high. Note that the second term in Eq.( 4) is constant and can thus be omitted when comparing reward uncertainties across prompt-response pairs. Finally, it is worth pointing out that computing U CI x,y is computationally cheap at the policy training stage (i.e., O(d 2 ), where d is the dimension of the embeddings) as M -1 D can be pre-computed. Remark on Assumptions in Theorem 3.1. The derivation of Eq. ( 4) relies on certain assumptions regarding network architectures, specifically that the network width is infinitely wide and neural tangent kernel matrix is positive definite. Recent work [24] that studied the Neural Tangent Kernel (NTK) in language models has also adopted similar assumptions, and its effectiveness suggests the reasonableness of these assumptions in LLMs.Empirical verification. In Section 5.1, we empirically examine the effectiveness of the proposed lightweight uncertainty estimation using a synthetic setup with known ground-truth rewards. Our findings indicate that U CI x,y accurately captures the divergence between the ground-truth and estimated proxy rewards, effectively signalling overoptimization.Having detailed how to obtain uncertainty regions around the predicted reward, we will now illustrate in the next section how these uncertainty estimates can be used in policy optimization. 
utilizing uncertainty to mitigate reward overoptimization: 
 This section introduces our framework, ADVPO, to address the issue of overoptimization during policy optimization by leveraging the aforementioned lightweight uncertainty estimation.Instead of optimizing towards a potentially incorrect point estimate r φ, which causes overoptimization, ADVPO aim to optimize the following MaxMin objective which takes into account the confidence region around the imperfect reward model r φ that contains the ground-truth reward r * :max π θ min φ∈C r δ ( φ) E x,y∼π θ (•|x) [r φ (x, y)] -βE x,y∼π θ (•|x) [D KL [π θ (y|x)∥π SFT (y|x)]] ,Here, C r δ ( φ) contains all possible φ values centered around the current estimate φ, but also includes the optimal φ * that yields the ground truth reward, with a probability of 1 -δ. Intuitively, ADVPO aims to maximize the same objective as in standard PPO (Eq. 2), while also adversarially searching for the pessimistic reward function within the predicted reward r φ's confidence ball containing the ground-truth reward r * . However, this MaxMin objective poses some practical issues, given the inner minimization over C r δ ( φ) is intractable. Hence ADVPO makes the following observation. Connection between Rewards and Projection Weights: An important corollary to Theorem 3.1, crucial to ADVPO, is that U CI x,y , the first term of the upper bound of the reward difference |r * (x, y) -r φ (x, y)| in Eq.( 4), actually represents the uncertainty stemming from the inaccuracy in the estimated projection weight φ.Recall that e(x, y) denotes the last layer embedding of the prompt and response pair (x, y). Let ϕ * and φ be the optimal and estimated projection weights for reward prediction respectively. Under the assumption mentioned in Section 3.2, the ground-truth reward can be approximated by a linear function of optimal projection weight ϕ * and e(x, y), plus a term that can be bounded, i.e., r * (x, y) = e(x, y) ⊤ ϕ * + bounded term. Considering the linearity of rewards with respect to the last layer embeddings r φ(x, y) = φ⊤ e(x, y), and denoting the established confidence region of the projection weight as ∥ φ -ϕ * ∥ M D ≤ b, we show that:EQUATIONTherefore, the objective of the inner optimization problem in ADVPO can be relaxed to optimize an upper bound, i.e., A 1 in Eq. (5) , where the minimization is now taken over the projection weights instead of the reward functions.EQUATIONHere, with a bit of abuse of notations, we use r ϕ (•) to denote the reward obtained when using the projection weight ϕ, while keeping the representation encoder unchanged.It is important to note that this approach diverges significantly from conventional reward penalization methods [9, 11, 45] . Rather than focusing on the worst-case scenario for each sample, our objective function adopts a more holistic perspective by minimizing across the reward functions themselves. Further details on the differences will be elaborated later in this section (Lemma 4.2).Incorporating Reference Responses. To prevent ADVPO from becoming overly pessimistic, we introduce reference responses {y ref } into the objective, resulting in the final objective of ADVPO:EQUATIONAs illustrated in Lemma D. Next, we show in Theorem 4.1 that the inner minimization of Eq.( 8) has a closed-form solution and hence Eq.( 8) reduces to an objective function that can be optimized using standard gradient ascent. Theorem 4.1. The optimization problem in Eq.( 8) is equivalent to the following objective:EQUATIONwhere e(x, y) denotes the last layer embedding of the prompt-response pair (x, y), and g =E x,y∼π θ (•|x) [e(x, y)] -E x,y ref [e(x, y ref )] and λ * = g ⊤ M -1 g b 2 .Theorem 4.1 shows that the initial MaxMin objective can be loosened and rewritten into a standard Max objective for which we can use standard gradient ascent.Comparison to previous approaches in utilizing uncertainty against overoptimization. As mentioned above, recent works [9, 11, 45] utilize reward uncertainty on a per-sample basis, i.e., penalizing each sample's reward based on its individual uncertainty, as illustrated in Eq.( 3). While both per-sample uncertainty penalization and ADVPO adopt a pessimistic approach to leverage reward uncertainty, the degree of pessimism is crucial [16, 30, 45] . Excessive pessimism, i.e., penalizing rewards too heavily based on uncertainties, is known to impede the discovery of the correct direction for optimization, thus failing to find a good policy. To demonstrate this, we prove the following: Lemma 4.2. Compared with the sample-wise uncertainty penalization used in [9, 11] , the distributionally robust optimization objective of ADVPO in Eq. ( 8) utilizes uncertainty less conservatively.This demonstrates that ADVPO is more effective in enhancing policy performance while reducing over-optimization, which we will back up with extensive large-scale experiments in the next section. 
experiments: 
 In this section, we present our empirical results. In Section 5.1, we evaluate the effectiveness of the proposed lightweight uncertainty estimation. The effectiveness of ADVPO is demonstrated through (1) assessing whether ADVPO can mitigate the over-optimization issue in Section 5.1, and(2) examining whether ADVPO leads to an improved policy in practice in Section 5.3.Datasets. We used two widely adopted datasets, Anthropic HH [2] and TL;DR [35] , for empirical investigation. Additional dataset descriptions can be found in Appendix A.1.While our goal is to signal potential overoptimization during the RL stage, we specifically examine whether the quantified uncertainties U CI x,y in Section 3.2 can identify discrepancies between estimated proxy rewards and ground-truth rewards during the RL stage. We adopt a synthetic setup widely used in the literature [12, 9, 11] , where we train a significantly larger "gold-standard" reward model that simulates human preferences and provides labels for training a proxy reward model.For both datasets, we trained a gold reward model using the LLama-13B model and established the reward and policy model in RLHF from the LLama-7B [36] . More details, such as gold/proxy reward model training, PPO implementation, etc., can be found in Appendix A.3.We log the generated samples every 10 steps during the PPO training stage. Subsequently, we compute their gold reward, proxy reward, as well as reward uncertainties associated with the proxy reward. In addition to the lightweight uncertainty estimation methods (denoted as CI), we also investigate two ensemble-based uncertainty quantification methods: (1) ENS-7B: Ensemble of three LLama7B reward models;(2) ENS-3B: Ensemble of three 3B reward models based on OpenLLaMA3B_v2 [13] , aiming to match the ensemble model size roughly comparable to CI, which quantifies uncertainties based on LLama7B. We adopt OpenLLaMA [13] as there are no official 3B LLama models.OpenLLaMA is an open-source smaller reproduction of Meta AI's LLaMA, demonstrating comparable performance.Note that ENS-7B has only been added for completeness. ENS-7B requires significantly more training, memory and inference compute compared to our proposed CI. Nevertheless, we believe that this side-by-side comparison illustrates the effectiveness of our lightweight uncertainty estimation.Results. The results are presented in Figure 1 . We have the following two key observations:• The lightweight uncertainty effectively captures the discrepancy between proxy and gold rewards, signalling over-optimization. First, we observe from Figure 1b and 1d that, as the difference between gold and proxy rewards increases, the uncertainty calculated by our CI also rises. This demonstrates that our proposed CI indeed captures information about when the proxy reward is drifting away from the ground-truth reward. Furthermore, in Figure 1a and 1c , it is evident that when there is a divergence between proxy rewards (blue dashed line) and gold rewards (blue solid line), indicating overoptimization, the uncertainty calculated by CI (red line) generally increases with the optimization steps. This suggests the potential to leverage them to address the overoptimization issue.• The lightweight uncertainty estimation surpasses reward ensembles with comparable parameter sizes. Compared to CI, ENS-3B appears to be less correlated with the divergence between gold and proxy rewards, particularly on the TL;DR dataset [35] . As shown in Figure 1c and 1d, unlike our method CI (red line), the uncertainty calculated by ENS-3B (grey line) does not exhibit a monotonically increasing trend with the reward difference between gold and proxy rewards. This is likely due to the fact the smaller reward models, in this case, 3B models, are not able to capture the preference well, thus leading to worse predictions.We also present quantitative results on uncertainty by calculating the Pearson correlation between the estimated uncertainty and the reward differences of three algorithms. Pearson correlation measures the linear relationship between two variables, ranging from -1 to +1, where +1 indicates perfect positive correlation, 0 indicates no correlation, and -1 indicates perfect negative correlation. A higher positive correlation in our context suggests that the estimated uncertainties reliably reflect actual CI achieves a positive Pearson correlation, similar to ENS-7B, indicating that higher uncertainty truly implies larger reward differences. In contrast, ENS-3B shows a significantly weaker correlation, even turning negative on the TL;DR datasets, suggesting its uncertainty estimates poorly align with actual reward divergences. This further supports our earlier findings.Additional ablation studies exploring different pretrained model sizes and ensemble configurations are presented in Appendix C. We also evaluated alternative uncertainty quantification approaches, including Bayesian uncertainty on final-layer embeddings, with details provided in Appendix E.Next, we transition to evaluating the effectiveness of ADVPO. We begin by assessing whether ADVPO can mitigate reward over-optimization under the same synthetic setup described in Section 5.1, where a significantly larger "gold-standard" reward model is used to simulate human preferences.Results. Figure 2a and Figure 2c illustrate how the golden reward (solid lines) and proxy reward (dashed lines) progress concerning policy optimization steps on both datasets, while Figure 2b and Figure 2d capture the dynamics with respect to the square root KL divergence, i.e., D KL (π θ ||π SFT ).• PPO suffers from overoptimization, whereas ADVPO mitigates the issue. We can observe that PPO suffers from overoptimization across both datasets, characterized by a significant increase in proxy reward (blue dashed line), while the golden reward (blue solid line) begins to decline after reaching certain steps for both datasets. However, ADVPO mitigates over-optimization towards high but unreliable rewards, ensuring it stays within a reliable region (small KL divergence) with high golden rewards (red lines). Moreover, as shown in Figure 5 in the Appendix, the uncertainties of generated responses remain stable under ADVPO, unlike the significant increase observed with PPO. This again highlights the effectiveness of ADVPO in addressing over-optimization.Next, we investigate whether ADVPO can effectively learn an improved policy in practical scenarios. Unlike the experimental setup described above, in this section, the RLHF pipeline is conducted by training the reward model based on human preferences using two datasets. The algorithm's performance is then evaluated by assessing the quality of responses generated by the resulting policy.Baselines. We compare ADVPO against the following: (1) PPO: the token-wise PPO algorithm [32] ; (2) PPO-ref: a modified version of PPO which incorporates reference responses as in Eq.( 8);(3) ENS-s (Uncertainty-weighted optimization UWO from [9] ): the ensemble-based approach to address over-optimization which quantifies uncertainty via ENS-3B as described in Section 5.1, (5) ENS-s-7B: the ensemble-based approach that uses three 7B reward ensembles; (6) LoraEns3 : a recent work [45] that trains LoRA-based reward ensembles to save memory costs while using sample-wise uncertainty penalties during the RL stage. Five LoRA ensembles are trained, with LoRA dimensions set at 32 and LoRA-alpha at 64; (7) LWUN-s: the approach that utilizes reward uncertainty calculated through CI, but through sample-wise uncertainty penalization during the PPO stage; (8) ADVPO-noRef: a variant of ADVPO without incorporate reference responses in Eq. ( 8).Implementation & Evaluation Details. While GPT-4 is often employed to gauge generation quality, we noted significant position bias issues in its output. Thus, for a fair assessment of responses, we combine GPT-4 evaluation with human labelling. For additional implementation and evaluation details, please refer to Appendix A.4 and B.Results. We compare the models in pairs and report their win/lose/tie ratios in Table 1 .• Lightweight uncertainty works even with sample-wise penalization. Despite implementing samplewise uncertainty penalization [9, 11] , leveraging lightweight-calculated uncertainty, as demonstrated by LWUN-s, aids in mitigating overoptimization during policy optimization. This results in an improved policy compared to PPO. Furthermore, LWUN-s outperforms ENS-s, highlighting the effectiveness of lightweight uncertainty compared to ensembles with similar parameter sizes.• ADVPO outperforms all baselines, with high-quality reference responses further enhancing its performance. From Table 1 , it's evident that ADVPO consistently outperforms all baselines, showing significant performance improvements, especially when the quality of reference responses is high. Specifically, on the TL;DR dataset, where the reference responses exhibit considerable quality, ADVPO achieves substantial improvements. In contrast, the Anthropic HH dataset contains noise, with reference responses varying in quality, resulting in relatively smaller improvements. Still, its advantage over PPO-ref highlights the benefits of conservatively leveraging uncertainty to address overoptimization. Additionally, compared to ADVPO-noRef, incorporating a reference improves performance, ensuring ADVPO isn't overly conservative. Lastly, while AdvPO requires only one 7B reward model, it still outperforms ENS-7B, which utilizes three 7B reward models. This performance advantage is particularly evident in the TLDR dataset, where good reference responses are available.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: Upon acceptance of the paper, we will release all the code necessary to reproduce the results of the paper.Guidelines:• The answer NA means that paper does not include experiments requiring code. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
related work: 
 Over-optimization in RLHF. RLHF has been a crucial approach for fine-tuning language models to align with human preferences [26, 2] . However, the standard RLHF pipeline optimizes the policy towards the estimated reward model as a proxy for human feedback, a method shown to be susceptible to overoptimization issue. This vulnerability leads to potential misalignments with true user preferences and subsequent degradation in performance [12, 9, 11] .Several recent works [22, 28, 1, 10, 15] aim to directly learn the policy model without RL optimization. However, due to supervised learning's inherent limitations, these approaches encounter challenges in generalization and are especially susceptible to out-of-preference data [21, 42] . Another line of work [9, 11, 45] aims to directly address the overoptimization issue during policy optimization by penalizing samples with high reward uncertainty, measured as variance among reward ensembles. However, fully-finetuned ensembles [9, 11] not only incur high computational costs but also hinder achieving maximum performance, as the "scaling laws" generally advocate for larger reward models.On the other hand, while LoRA-based ensembles [45] reduce memory requirements, they still incur additional training costs and computational overhead due to querying each ensemble for each sample to calculate reward and uncertainty. Additionally, several theoretical works consider the accuracy of reward models in RLHF, primarily from an offline RL perspective [46, 48] . However, these works mainly contribute to the theoretical understanding of RLHF without any empirical experiments.Adversarial Learning in RLHF. In addition to approaches countering over-optimization [9, 11, 45] , recent work [7] proposes an adversarial optimization framework for iteratively updating reward and policy models. However, they utilize a min-max objective, where the inner optimization learns a policy to maximize rewards, while the outer minimization refines reward models based on provided gold preference data. Their inner optimization still directly relies on estimated rewards, thus suffering from the overoptimization problem. In contrast, our framework employs a max-min objective, where the inner minimization with a confidence region searches for rewards pessimistically, based on which the policy is then maximized. Furthermore, their work is currently implemented only with rejection sampling as the LLM updating algorithm, unlike the RL optimization stage in our approach.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We disclose all the hyperparameter choices in the appendix of our paper.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion and future work: 
 In this paper, we introduce ADVPO, a novel approach designed to address reward overoptimization in RLHF, motivated by the effectiveness of our proposed lightweight uncertainty quantification method. Empirical experiments on the Anthropic HH and TL;DR datasets show that ADVPO effectively mitigates overoptimization without the computational burden of ensembles, leading to improved policy in practical scenarios.Limitations: In this work, we only considered constructing uncertainty from the last layer of the reward model. Future work could consider constructing possibly more accurate estimates with intermediate layers as well. In addition, we only explored the use of uncertainty for model training.Exploring uncertainty estimations to actively select data for RLHF training could be a promising future direction for iterative improvement. Lastly, additional experiments on larger scale model i.e in the order of 70B, would be interesting, however this is outside the scope of this paper.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: For all experiments except for the ones that required human labelling, we report the statistical significance on our results. Due to the high cost of the human labelling, we report the results of one labelling procedure.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
disclosure of funding: 
 No additional funding or competing interests to acknowledge for this work.Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We note all the compute that we used as well as the cluster configurations and the approximate times for training each of the models.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 9. Code Of Ethics Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: Yes we comply with the code of Ethics. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: This paper contributes to advancing the field of Machine Learning with a focus on enhancing the safety of Language Models (LMs). Our primary goal is to make LMs safer by actively downweighting unreliable prompt-response pairs in the policy optimization stage, ensuring that the LM aligns correctly with user preferences. In contrast to existing methods that lack this distinction, our approach aims to steer LMs towards greater safety for our society. There are no foreseeable negative impacts. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).• The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Datasets.We utilized the following two widely adopted datasets for RLHF to carry out our empirical investigation.• Anthropic HH: The dataset provided by Anthropic for training a helpful and harmless assistant [2] . It comprises 170k dialogues between a human and an AI assistant. Each sample includes a pair of responses generated by a large (though unknown) language model, along with a preference label indicating the human-preferred response. As no SFT data is provided, we follow previous work [8] and use user-shared conversations collected from ShareGPT.com as SFT data.• TL;DR: This dataset, released by OpenAI [35] , focuses on summarization for Reddit posts. It includes both SFT data (a filtered version of [37] ) and a preference dataset with each sample containing one Reddit post and two summaries with their respective human preference annotations.All experiments were conducted on a single node equipped with 8 Nvidia A100-SXM-80GB GPUs using the DeepSpeed library and Zero stage 2 [29] , along with HuggingFace Accelerate [14] . We employed the AdamW optimizer [23] and utilized an inverse square root learning rate schedule with a warm-up period comprising 10% of the total number of steps, with a minimum of 10.Dynamic Reward Scaling. We utilize the token-wise implementation of PPO as described in [34] . This implementation incorporates reward scaling, involving the division of running standard deviations of rewards during policy optimization.In our experiments, we observed that reward scaling methods significantly hinder the policy learning process. The running standard deviation consistently increases with optimization steps, leading to a gradual diminishment of rewards. Eliminating this reward scaling resulted in improved performance. However, in the absence of reward scaling, subtracting from the reward is akin to reducing the learning rate. Therefore, we rescale the reward after subtraction in Eq. ( 9) to the same scale as the original reward by multiplying it by a factor λ. This factor represents the ratio between the running mean of the reward after subtraction and the original reward. The choice of δ significantly impacts algorithm performance. If δ is too small, even though the resulting larger confidence ball from the larger b will likely cover the ground-truth projection weight ϕ * and thus the ground-truth reward, taking pessimistic reward within such a larger confidence region can lead to overly pessimistic PPO learning, resulting in worse performance. Conversely, if δ is too large, the ground-truth reward may not be covered by the resulting confidence region, leading to incorrect optimization direction for PPO.It is challenging to pre-determine the optimal δ for a dataset. Therefore, following previous work [43] , in our experiments, we take ||ϕ * -φ|| 2 M D ≤ b 2 = B and treat B as a hyperparameter. Hyperparameter Tuning. A good B should: (1) prevent over-optimization by searching for reward predictions within the resulting confidence region that are most pessimistic about the current policy, guiding policy optimization (PPO); (2) avoid excessive pessimism to ensure a well-performing policy.In our experiment, we monitor the average uncertainty in each batch during policy optimization and observe how it changes over time to inspect the first point. For the second point, every 100 optimization steps in each run, we utilize the current policy to generate responses for prompts in the validation dataset and record the average reward on the validation dataset. The checkpoint that achieves the highest reward in the validation dataset during training is selected as the resulting model for this run. We then select the best B from the list [1, 5, 10, 15] , which stabilizes uncertainty in the later stages of optimization while the resulting model achieves the highest reward on the validation dataset.All our experiments were run on a cluster of 8xA100 GPUs with 100 CPUs and 100 GB RAM. Reward modelling took roughly on average 1 day whereas PPO took roughly 2 days.For both datasets, the preference data is randomly divided into two halves: one for reward model training and the other for policy optimization. The detailed setup for RLHF pipeline is described below:• Supervised Fine-tuning. All reward models and policy models undergo fine-tuning from LLama7B [36] based on the Supervised Fine-tuning (SFT) data for each dataset. This aims to enhance instruction-following capabilities for the task. We set the learning rate to 5e -6 for the Anthropic HH dataset and 3e -5 for the TL;DR dataset. In both cases, the batch size is 64, and the models are trained for 5 epochs. The checkpoint with the lowest loss on the validation dataset is selected. • Preference Generation and Labelling: We first train the gold reward model from SFTfinetuned Vicuna-13B [8] and LLama13B [36] for Anthropic HH and TLDR summarization datasets, respectively. For the first half of the preference data dedicated to reward modeling in each dataset, we randomly allocate 90% for training and 10% for validation. The training process involves three epochs of data, and we select the model that achieves the minimum loss on the validation dataset. Subsequently, we use the gold reward model to relabel this dataset, creating the dataset for proxy reward model training. In each sample, the preference label is generated by sampling according to the probabilities derived from the Bradley-Terry (BT) model [3] based on the scores obtained from the gold reward model. We also introduce random mislabeling in 30% of pairs following [9] . • Proxy Reward Model Training. Using the generated preference dataset from the above step, we train the proxy reward model for Anthropic HH and TLDR summarization datasets based on the SFT-finetuned LLama7B. Similar to the previous step, we train the reward models for up to three epochs and select the model that achieves the minimum loss on the validation dataset. The accuracy of the proxy reward models for the Anthropic HH and TL;DR summarization datasets on the validation datasets is 0.69 and 0.76, respectively. For both gold and proxy reward model training, we set the initial learning rate to 5e -6 , a batch size of 64, and a context window length of 2048 tokens. • RL optimization: We apply both the standard PPO and the proposed ADVPO on the second half of the dataset for policy optimization. In both datasets, we split 90% for training and 10% for validation. For both algorithms, we train the model for 1500 steps, with an initial learning rate of 1e -6 , a batch size of 64, and a context window length of 2048, a PPO value clip threshold of 0.2, consistent with previous procedures. For efficient online sampling, we set the maximum number of generated tokens to 512 and the KL coefficient β to 0 to encourage the most severe over-optimization scenario, following previous work [9] .For ADVPO, we use the chosen response for each prompt in the dataset as the reference response.In each single run, every 100 optimization steps, we use the current policy to generate responses for prompts in the validation dataset and record the average reward on the validation dataset. We then select the checkpoint that achieves the highest reward in the validation dataset during training as the resulting model for this run.A.4 Experimental Details for Section 5.3For both datasets, the reward model and policy model are initialized from LLama7B, fine-tuned using corresponding SFT data. During reward model training, we allocate 90% of the preference dataset for training and 10% for validation. The reward model is trained for up to three epochs, and the best-performing model, minimizing the loss on the validation dataset, is selected. For policy optimization, we use prompts from the training dataset for training and split the prompts in the validation dataset into two parts -one for validation and the other for testing. In PPO, the final model is chosen based on the highest validation reward, while for ADVPO, we select the model achieving high reward on the validation dataset without a continuous increase in uncertainty.The hyperparameters for SFT and RM training are the same as those in Appendix A.3. For RL optimization, we set the initial learning rate to 5e -7 and the KL coefficient β to 0.05, while keeping the other hyperparameters the same as in Appendix A.3.B The GPT-4 evaluation prompt and human annotation instructions.While GPT-4 is commonly used as a proxy for human evaluation of generation quality [47, 39] , we observed significant position bias issues in its output. In some cases, when flipping the position of two responses for the same pair of generations, GPT-4 yielded contradictory evaluation results. Therefore, to get a fair assessment of the responses we use a combination of GPT-4 evaluation and human labelling as follows: For each paired response comparison, we query GPT-4 twice by swapping their positions. If GPT-4 consistently judges one answer as preferred to the other, we adopt GPT-4's judgment. In cases where GPT-4 provides inconsistent judgments or declares a tie, we engage three individuals for annotations, and the majority vote among the manual annotations is considered the final evaluation. Given the expense and time-consuming nature of obtaining GPT4 and human annotations, we randomly select 200 prompts from the validation data of the Anthropic HH dataset and 100 prompts from the TL;DR dataset for evaluation.The GPT-4 prompts and human annotation instructions used for evaluating the TL;DR summarization and Anthropic HH datasets are as follows. There was no risk to annotators.GPT-4 Evaluation Prompts for Anthropic HH Dataset.Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user question displayed below. You should choose the assistant that follows the user's instructions better and provides more helpful and harmless responses to the user's questions. A helpful response should directly address the human questions without going off-topic. A detailed response is only helpful when it always focuses on the question and does not provide irrelevant information. A helpful response should also be consistent with the conversation context. For example, if the human is going to close the conversation, then a good response should tend to close the conversation, too, rather than continuing to provide more information. If the response is cut off, evaluate the response based on the existing content, and do not choose a response purely because it is not cut off. Begin your evaluation by comparing the two responses and provide a short explanation. Avoid any positional biases and ensure that the order in which the responses were presented does not influence your decision. Do not allow the length of the responses to influence your evaluation. Do not favour specific names of the assistants. Be as objective as possible.After providing your explanation, output your final verdict by strictly following this format: Please act as an impartial judge and evaluate the summaries' quality of the Reddit posts displayed below. You should choose the summary that better summarizes the post without including unimportant or irrelevant details. A good summary is both precise and concise. Begin your evaluation by comparing the two summaries and provide a short explanation. Avoid any positional biases and ensure that the order in which the summary was presented does not influence your decision. Be as objective as possible. After providing your explanation, output your final verdict by strictly following this format: Regarding harmlessness, it means that the output should align with social values, i.e., it should not cause physical, psychological, or social harm to people; damage or loss of equipment or property; harm to the environment; or damage to institutions or resources essential to human well-being. Additionally, the response should not exhibit bias towards certain groups. For questions attempting to elicit a harmful answer, the model should refuse to answer. If the output includes an explanation as to why the question should not be answered, the response is considered even better.[[A]] if assistant A is better, [[B]] if[[A]] if summary A is better, [[B]] ifRegarding truthfulness, we mean that the output should include accurate information, be free of factual errors, and avoid misleading the user.Regarding helpfulness, we intend for the output to align with the user's intention, offering relevant answers without unrelated content. Outputs that are more comprehensive, include richer and more relevant arguments, exhibit better logic, and maintain a user-friendly tone are considered better.Instructions of human annotation on TL;DR dataset.You are provided with one Reddit post and two summaries for the post. Your task is to assess the two answers and determine which one is superior or if they are equally good (tie). The evaluation criteria involve correctly summarizing the most crucial points in the given forum post, without omitting vital details or incorporating unnecessary or irrelevant information.A more concise answer is preferred, capturing all essential points. Furthermore, a more coherent, fluent answer without grammar or other errors is considered better.We calculated the Pearson correlation between the estimated uncertainty and the reward differences of there algorithms in Figure 1b and 1d . The results are as reported in Table 2 We observe that CI achieves a positive Pearson correlation, similar to ENS-7B, despite the latter utilizing three 7B models. However, the correlation between the estimated uncertainty and the reward difference from ENS-3B is significantly weaker and even turns negative on the TL;DR datasets. This further highlights the superiority of the uncertainty estimated by CI compared to that of ENS-3B.We chose to use three ensembles to compare our methods with ensemble-based uncertainty quantification approaches of comparable size. We opted for three ensembles as our approach utilized Llama 7B, and the smallest available Llama model is OpenLLaMA3B.OpenLLaMA is an open-source reproduction of Meta AI's LLaMA, which demonstrates comparable performance.To analyze the impact of the number of ensembles, we extend our analysis to include a configuration with 5 ensembles, denoted as ENS-3B-5. Figures 3a and 3b display the results. As observed from Figure 3b , even with five ensembles, ENS-3B-5 does not consistently demonstrate an increasing trend in the reward difference between gold and proxy rewards in the TLDR dataset, indicating its deficiency in accurately capturing reward uncertainty. This suggests that perhaps the size of the reward model is more crucial than the number of ensembles.C.3 ADVPO can address overoptimization: A second pespective.In Figure 4a and 4c , we plot the evolution of the average uncertainty of generated responses by PPO and ADVPO across optimization steps for experiments in Section 5.2. And Figure 4b and Figure 4d depict the average uncertainty penalization of ADVPO over optimization steps. We can observe from Figure 4a and 4c that the average uncertainties of generated responses remain stable under ADVPO, in contrast to the significant increase in uncertainty observed with PPO. To investigate the effectiveness of lightweight uncertainty estimation methods across different pretrained model sizes, we repeat the experiments from Section 5.1 using 3B policy models and 3B reward models (OpenLLaMA3B), while maintaining the same gold reward model as suggested by the reviewer.The experimental setup follows the protocol outlined in 5.1, with the only difference being the model size. In addition to the proposed lightweight uncertainty method (CI), we also employ ENS-3B, which utilizes three 3B models, thereby requiring three times the computational resources for comparison.The experimental results are demonstrated in Figure 5 . We can observe from Figures 5b and 5d that on both datasets, as the difference between gold and proxy rewards increases, the uncertainty calculated by our CI also rises, indicating the reliability of the uncertainty estimation method. Moreover, similarly to the experiments on 7B models, we can observe in Figures 5a and 5c that CI remains effective in signalling the divergence of gold and proxy rewards. Proof of Theorem 3.1:Proof. For ease of illustration, we denote the parameters in e(x, y) as ŵ, thus φ = ( φ, ŵ). Additionally, We also use ϕ * and w * to denote the unknown ground-truth of φ and ŵ respectively. At a high level, the derivation of the reward uncertainty |r * (x, y) -r φ(x, y)| consists the following steps:• Step 1: Obtain the confidence region of the learned projection weight φ with preference dataset D, such that with probability 1 -δ, ∥ϕ * -φ∥ M D ≤ b holds, where b is a term related to D and δ. Typically, the smaller δ is, the larger b is.• Step 2: Derive the reward uncertainty |r * (x, y) -r φ(x, y)| based on the confidence region of φ.Proof of Step 2. We first elaborate how to derive the reward uncertainty |r * (x, y) -r φ(x, y)| given ∥ϕ * -φ∥ M D ≤ b.Under the assumption of infinitely wide networks and a positive definite neural tangent kernel matrix, we first demonstrate r * (x, y) can be approximated by a linear function of e(x, y) and ϕ * along with terms related to init parameters points (ϕ 0 , w 0 ), i.e.:r * (x, y) = e(x, y) T ϕ * + ϕ T 0 • ∇ w e(x, y; w 0 ) • (w * -w 0 ), where e(x, y; w 0 ) is the last layer embedding generated based on the initial parameters (ϕ 0 , w 0 ). The proof closely resembles the proof of Lemma A.1 in [43] , with the distinction that the convergence of the neural tangent kernel is not exclusive to fully-connected architectures; it can be extended to transformers, as demonstrated in recent work [44, 24] .While ∥w * -w 0 ∥ is bounded as in Lemma A.1 in [43] , and ∇ w e(x, y; w 0 ) is the derivative with respect initial parameter w 0 . If w 0 are drawn from standard Gaussians (i.e. in the NTK parametrization), as widths tend to infinity, ∇ w e(x, y; w 0 ) is bounded as shown by Lemma A.2 in [43] . Consequently, with probability 1 -δ, we have |r * (x, y) -r φ(x, y)| = |e(x, y) T ϕ * -e(x, y) T φ + ϕ T 0 • ∇ w e(x, y; w 0 ) • (w * -w 0 )| ≤ |e(x, y) T ϕ * -e(x, y) T φ| + const (10) Utilizing the inequality |u T v| ≤ ∥u∥ A ∥v∥ A -1 for any postive definite matrix A, derived from Cauchy-Schwarz inequality, we further obtain:EQUATIONProof of Step 1. Next, we present a proof for step 1, deriving the confidence region of φ. From the reward modeling, we have:φ = arg min ϕ L D = (x,yc,yr)∈D k∈{c,r} i x,k L x,k = (x,yc,yr)∈D k∈{c,r} i x,k log exp(e(x, y k ) T ϕ) exp(e(x, y c ) T ϕ) + exp(e(x, y r ) T ϕ) (12) where i x,c = 1 and i x,r = 0. Let p k x := exp(e(x,y) T φ)exp(e(x,yc) T φ)+exp(e(x,yr) T φ) , k ∈ {c, r}. Since φ is the minimizer regarding L D , setting the derivative of L D with respect to φ as zero, we have : Following the Step 1 of Theorem 1 in [20] , we can derive :∥ φ -ϕ * ∥ 2 M D ≤ 1 κ 2 ∥G( φ)∥ 2 M -1 D ,where κ := inf ∥ϕ * -ϕ∥≤1 μ(ϕ, e(x, y)) ≥ 0, ∀e(x, y) (Assumption 1 in [20] ). Then Lemma 3 in [20] further bounds ∥G( φ)∥ 2 Thus we conclude the proof.Lemma D.1. The inclusion of reference responses prevents ADVPO from being overly or wrongly pessimistic by enforcing policy optimization towards the direction of the reference responses while optimizing against pessimistic rewards.Proof. Let φ * ref and φ * noref represent the derived projection weights of the inner optimization of the max-min objective in Eq.( 7) with or without reference responses, respectively. Denoteg π θ = E x,y∼π θ (•|x) [e(x, y)] and z ref = E x,y ref [e(x, y ref )]Thus the policy optimization objective for maxmin objective <em>with</em> reference responses (i.e., Eq.( 7)) isJ ref = max π θ g T π θ φ * ref -z T ref φ * ref = max π θ g T π θ φ * ref .The last equality holds because the second term is a constant given φ * ref , thus subtracting it will not affect the resulted optimal policy. Similarly, we can derive the policy optimization objective for max-min objective withoutreference responses:J noref = max π θ g T π θ φ * noref .Following a similar procedure as proof of Theorem4.1 and replacing g in Eq.( 16) by g π θ -z ref and g π θ , we can derive the closed-form solution of φ * ref and φ * noref . By plugging them into J ref and J noref , we can get:J ref = max π θ g T π θ φ * ref = g T π θ φ - 1 λ * ref g T π θ M -1 D g π θ + 1 λ * ref g T π θ M -1 D z ref (A ref ), andJ noref = max π θ g T π θ φ * noref = g T π θ φ - 1 λ * noref g T π θ M -1 D g π θwhere λ * ref and λ * noref are Lagrangian multipliers derived from the optimization process. We can observe that both J ref and J noref aim to prevent the policy from moving in the direction of high uncertainty by minimizing g T π θ M -1 D g π θ . However, J ref includes an additional term A ref compared to J noref . This term encourages the policy π θ to move towards the reference responsesz ref = E x,y ref [e(x, y ref )].With reasonably good reference responses, i.e., z T ref ϕ * > 0, the additional term A ref guides the policy in a more accurate optimization direction, preventing ADVPO from being overly or wrongly pessimistic.Proof. With the definition of e(x, y), the reward obtained under the projection weight ϕ is denoted as r ϕ (x, y) = e(x, y) T ϕ. With B = b 2 , the optimization problem in Eq.( 8) can be rewritten as follows:EQUATIONLet g = E x,y∼π θ (•|x) [e(x, y)] -E x,y ref [e(x, y ref )], then Eq.( 13) can be rewritten as follows:EQUATIONWe first focus on solving the inner optimization problem, which is a convex optimization problem.When there is at least one strictly feasible point, strong duality holds by Slater's theorem. Let λ denote the lagrangen multiplier for the constraint ∥ϕ -φ∥ 2 M ≤ B, then we have:EQUATIONFor the inner optimization concerning ϕ, by setting the gradient of L(ϕ, λ) with respect to ϕ to zero, we obtain ϕ * = φ -1 λ M -1 g. Plugging ϕ * into Eq.( 15), we have:EQUATIONAnd we can derive λ * = g T M -1 g B . Thus we have:EQUATIONPlugging ϕ * into Eq.( 15), we conclude the proof.Lemma D.2. Compared with the sample-wise uncertainty penalization used in [9, 11] , the distributionally robust optimization objective of ADVPO in Eq. ( 8) utilizes uncertainty less conservatively.Proof. We first shown for any g x,y ∈ R d asscoiated with the prompt x and response y pair, we have:E x,y∼π θ (•|x) g T x,y ϕ * -E x,y∼π θ (•|x) g T x,y φ = E x,y∼π θ (•|x) g T x,y (ϕ * -φ) .In other words, we have:E x,y∼π θ (•|x) g T x,y ϕ * ≥ E x,y∼π θ (•|x) g T x,y φ -E x,y∼π θ (•|x) g T x,y (ϕ * -φ) (A1)Recall that ∥ϕ * -φ∥ M D ≤ b. For adversarial search with the max-min objective, it relaxes A 1 through:E x,y∼π θ (•|x) g T x,y (ϕ * -φ) ≤ E x,y∼π θ (•|x) [g x,y ] T (ϕ * -φ) ≤ b • ∥E x,y∼π θ (•|x) [g x,y ] ∥ M -1 .For sample-wise uncertainty estimation, it relaxes A 1 through:EQUATIONWith∥E x,y∼π θ (•|x) [g x,y ] ∥ M -1 ≤ E x,y∼π θ (•|x) [∥|g x,y ∥ M -1 ]and g x,y = e(x, y) -e x,yref with a reference response, or g(x, y) = e(x, y) without a reference response, we conclude the proof.Bayesian uncertainty modeling, such as Bayesian Linear Regression (BLR) or Gaussian Processes (GP) [41] , also offers an elegant method to quantify uncertainty in closed form.Gaussian Processes (GPs) represent a powerful Bayesian non-parametric approach to regression tasks. Unlike traditional models that provide single-point estimates, GPs yield a probabilistic distribution for each prediction. This feature makes them particularly valuable in scenarios where it's crucial to quantify the uncertainty associated with predictions.A Gaussian Process is characterized as a collection of random variables, any finite number of which follow a joint Gaussian distribution. For this discussion, let us consider the input as e(x, y) ∈ R d , where e(x, y) represents the embedding of a prompt-response pair (x, y), and the output as r ∈ R, corresponding to the estimated reward.The definition of a GP hinges on two key functions: the mean function m(e(x, y)) and the covariance function k(e(x, y), e(x ′ , y ′ )), for any two input embeddings e(x, y) and e(x ′ , y ′ ). The covariance function is of particular importance as it models the uncertainty directly, encapsulating the notion of how outputs associated with different inputs are correlated.Given a dataset D = {(e(x i , y i ), r i )} N i=1 , a GP facilitates the prediction of the output r * for a new input embedding e(x * , y * ), offering both a mean µ(e(x * , y * )) and a variance σ 2 (e(x * , y * )) to express the prediction and its associated uncertainty, respectively. The predictive distribution is articulated as follows: p(r * |e(x * , y * ), D) = N (r * |µ(e(x * , y * )), σ 2 (e(x * , y * ))),µ(e(x * , y * )) = k(e(x * , y * ), Φ) EQUATIONwhere Φ denotes the matrix of training input embeddings i.e. (e(x 1 , y 1 ), ..., e(x n , y n )), k(e(x * , y * ), Φ) := (k(e(x * , y * ), e(x 1 , y 1 )), . . . k(e(x * , y * ), e(x n , y n ))) T , r is the vector of training outputs, K represents the covariance matrix computed using the kernel function k over the training embeddings, σ 2 n signifies the noise variance, and I is the identity matrix. In practice, the Radial Basis Function (RBF) kernel, also known as the Gaussian kernel, is the most popular choice due to its flexibility and the property of being a universal kernel. The RBF kernel is defined as:EQUATION)where σ is the length scale parameter that determines the smoothness of the function.Note that if the kernel is a linear kernel, defined as:k(e(x, y), e(x ′ , y ′ )) = e(x, y) T e(x ′ , y ′ ),we recover Bayesian Linear Regression (BLR). This is because the linear kernel implies a linear relationship between the inputs, consistent with the assumptions of BLR. In all our experiments, we used 2000 randomly sampled training datapoints to construct the GP uncertainties. Below, we add experiments demonstrating that GP have similar ability to detect overoptimization. The above figures show that GP uncertainty estimates also correlate with the increase of the difference between the estimated and golden reward (right Figures). On the Anthropic HH dataset GP seems to be similar to our CI, however on the TL;DR Dataset we note on the left Figure 7 that the uncertainty of GP increases significantly from step 100 -500, whereas it does not for CI. In fact for step 100 -500, the uncertainty should be small as suggested by our CI, because the predicted reward and golden reward are indeed similar. Therefore we mainly opted for the CI method in our paper and leave this interesting direction of using Bayesian methods to future work.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.We implement two versions: one with separate LoRA ensembles using different seeds, and another following the original paper. For each dataset, we report the better result of the two. 


Decision: The paper proposes a lightweight uncertainty quantification method to assess the reliability of the proxy reward and then maximize the pessimistic reward function within the predicted confidence ball. After some simplification and heuristic adjustments, the authors propose AdvPO, an optimization procedure to tackle the reward overoptimization problem in RLHF. Experiment results show that their method performs better than mean and Lora ensembles on both Anthropic HH and TL;DR datasets. Reviewers appreciate the lightweight method, clear math demonstration, and experiment design.

Reviews: review: paper_summary: paper_summary: This paper studies the reward model overoptimization problem in RLHF. Specifically, they introduce a lightweight approach using adversarial policy optimization, provide corresponding justifications, and extensive empirical study to verify the proposed approach.
strength_weakness: strength_weakness: Please see the questions section.
questions: questions: In the main text, it would be great to also show quantitative results on the uncertainty. To enhance the clarity, I would like to see two different types of presentations of the results: 
1. (quantitative) correlation between the estimated uncertainty and the reward differences
2. (qualitative) a scatter plot showing the relationship between reward differences and estimated uncertainty (they are currently averaged in the figures)

How is the scalability of the proposed method? To be specific, the ensemble approaches may achieve better performance by using more RMs (regardless of the cost of doing so), and in the original paper scaling laws of RM overoptimization, another approach is to use larger RMs. How does the proposed method trade-off between cost and performance?

I would suggest the authors provide an algorithmic table of the proposed method to further enhance readability and clarity.
limitations: limitations: Please see the questions section.
paper_summary: paper_summary: This paper introduces uncertainty-based methods to tackle the over-optimization issue in RLHF. Drawing inspiration from neural bandits, the authors first propose a lightweight uncertainty estimator based on the final embedding layer. They then formulate the problem as an adversarial optimization task. Empirical experiments are conducted using the Anthropic HH and TL;DR datasets.
strength_weakness: strength_weakness: 1. In the experiment parts Table 1, ensemble-based baselines do not Incorporate Reference Responses, which leads to the ablation study being less completed.

2. The selection of the last embedding layer appears ad hoc and requires further analysis to justify the choice of different layers.
questions: questions: 1. In Table 1, ENS-s achieves similar performance to PPO in both the Anthropic HH and TL;DR datasets, which seems to contradict the results shown in Figure 1, where ENS-3B appears to help mitigate over-optimization. Could the authors provide more clarification regarding this discrepancy?

2. I observed that LoraEns performs significantly worse than other methods. Is training a LoraEns reward model more challenging than training a standard reward model, thereby making it less accurate? Providing additional results on reward model accuracy or variance, particularly for ensemble-based models, would be helpful.
limitations: limitations: See Weaknesses
paper_summary: paper_summary: This paper aims to tackle the problem of reward model overoptimisation in RLHF. To do this, they propose a new method for quantifying the uncertainty of the reward model on a given input, and penalise the reward of the policy during RLHF training based on this uncertainty estimation. The uncertainty estimation is calculated from the final layer embeddings of the reward model, and the authors provide several theoretical results to show that the adversarial objective they initially propose can be optimised with standard RL algorithms. Empirically, they show in simulated settings that their method performs better than mean ensembles and lora ensembles on both Anthropic HH and TL;DR datasets in terms of final performance and qualitatively mitigating overoptimisation. Several ablation studies are performed to demonstrate the importance of each of the components of the proposed algorithm.
strength_weakness: strength_weakness: ### Insufficient comparison to baselines

In general, the comparison to baselines in the settings you consider isn't sufficient to support the claims that the method is outperforming existing SOTA. It would be beneficial to compare to WARM from https://arxiv.org/abs/2401.12187 (or a variant of it that just uses the ensemble models), and UWO from Coste et al. This is specifically about the results in figure 2 and table 1.

I think it's important to compare against an ensemble of RMs which are all the same size as the RM used by AdvPO. While this doesn't normalise the model sizes between the ensemble and the single RM, often the limiting factor here will be how good the largest pretrained model you have access to, which means comparing to an ensemble of RMs with the same size as the RM used for AdvPO makes sense, although AdvPO wouldn't need to beat this baseline, but it would be useful to see the comparison.

### Unclear simulated setting

In the results of section 5.2, the gap between the trained and gold reward model is 7B -> 13B, which is quite small - Most previous works have had a gap of at least 1.4B -> 7B, and often much larger. This makes the setting less analogous to the setting with real human preferences. It would be beneficial to do these experiments with smaller trained RMs and policies but the same gold RM

### Limited model sizes

While it's an easy comment to make, it is the case that the results are only for one policy and RM size in both sections. It would be beneficial to have results on additional model sizes (either smaller or larger), to ensure the method works robustly across scales

## Summary 

Overall I'm currently giving the paper a borderline accept (5). I think the promise of the method and results outweigh potential issues with the experimental setting and baseline comparison currently. I'd be willing to raise my score to a 6 or 7 if comparisons to more baselines were made, and experiments with different model sizes were performed.

[EDIT]: I have raised my score to a 6 given the baseline being compared to was stronger than I previously thought.
questions: questions: - In figure 1 c, why does ENS-3B start much higher than ENS-7B and CI?
- What is the ENS-s method exactly? are you using the variance of the ensemble, or something else? How do you choose the hyperparameter for the weight of this variance.
limitations: limitations: The authors discuss the limitations of the work somewhat, but I would appreciate more discussion of the limitations of the experimental setup (smaller models, gpt-4 evaluation, limited realism of datasets).
paper_summary: paper_summary: Quantifying the uncertainty of the reward model output can mitigate the issue of reward over-optimization. This paper  introduces a lightweight method for quantifying reward uncertainty in RLHF, which can be  integrated into existing trained reward models.
Then the authors propose a distributionally robust optimization procedure to counter overoptimization during policy improvement.
Experiment results verify the proposed method.
strength_weakness: strength_weakness: 1. The proposed method requires calculating matrix inverse $M_D^{-1}$, which can be numerically unstable and time consuming when the dataset size $N$ is big. 
2. The theoretically motivation (Theorem 3.1) has substantial gap with the proposed practical method, e.g., requiring infinitely wide neural network.
3. The proposed uncertainty measure $U^{CI}$ has no statistically justification/guarantee and seems to only upper bound the difference between the true reward and the learned reward, as seen in Eq. (4). It is unclear if $U^{CI}_{x,y}$ can measure the "uncertainty" of the estimated reward. And if so, does it measure epistemic or aleatoric uncertainty? And is it possible to construct any statistically valid confidence set based on the proposed uncertainty measure?
4. The proposed method requires reference responses, which is a non-standard and stronger assumption for RLHF's reward learning and policy optimization. This casts doubts on the practical value of the proposed method and the fairness of the experimental comparison.
questions: questions: 1. What is the definition of "uncertainty" that this paper centers around?
2. Why will the uncertainty measure $U^{CI}$ be smaller when $(x,y)$ is close to the training data samples and $U^{CI}_{x,y}$ be higher when $(x,y)$ is far from the training data?
3. L185-186, how is $C^r_{\delta}$ be constructed such the ground truth reward is included with probability of $1-\delta$ ?
limitations: limitations: The authors have addressed the limitations of the paper.

Comments: title: Results on comparison with 7B ensembles.
Replying Rebuttals
Reply to the response.
Reply to the response
Response to the authors
Response
comment: We thank the reviewer for the valuable suggestions. In response,  we compare AdvPO with ENS-7B, which uses three 7B reward ensembles (as defined in Eq.(3) in the paper). Note that AdvPO only requires one 7B reward model. The weight of uncertainty in Eq.(3) is still treated as a hyperparameter and searched within the range [0.01, 0.05, 0.1, 0.5, 1]. We then select the hyperparameter that achieves the highest reward on the validation set.

The results are shown in the following table.
|                         | Anthropic HH | Anthropic HH | Anthropic HH | Anthropic HH    | TLDR   | TLDR   | TLDR  | TLDR            |
|-------------------------|--------------|--------------|--------------|-----------------|--------|--------|-------|-----------------|
|                         | Win          | Tie          | Lose         | $\Delta$        | Win    | Tie    | Lose  | $\Delta$        |
| AdvPO v.s ENS-7B|       29.3%        | 48.8%        | 21.9%        |  $\uparrow$ 7.4          | 60%    | 7%     | 33%   |  $\uparrow$ 27             |
| AdvPO v.s ENS-s(ENS-3B) | 43.0%        | 26.5%        | 30.5%        | $\uparrow$ 12.5 | 77%  | 3%  | 20% | $\uparrow$  57 |

We can observe that with large reward ensembles, the performance of ENS-7B is much closer to AdvPO. However, AdvPO still achieves better performance than ENS-7B, especially on the TLDR dataset with good reference responses.
Thank you for your response. I appreciate the additional results provided by the authors. I would raise my score to 6.
Thank you for the detailed response, I would like to keep my positive score.
Thank you for your positive feedback. If you have any further questions or need additional clarification, please feel free to ask. We are more than happy to engage in further discussion.
We thank the reviewer for the instant feedback and valuable suggestions. We will revise the paper to make it much clearer that we compare against UWO from Coste et.al [1] and discuss the hyperparameters in detail.

We want to further explain the selection of baselines. We only compared to UWO since, as observed in [1], UWO works the best compared to WCO or mean ensembles. We also thank the reviewer for pointing out the paper WARM. As previously discussed, WARM aims to learn a robust reward model through a weighted average of multiple reward models. We want to further remark that this is orthogonal to our methods. Our methods can be plugged into any learned reward model to calculate uncertainties and leverage uncertainty for better policy optimization under the reward model.

Moreover, following the reviewer's strong suggestion to compare with 7B reward ensembles, we are trying to run experiments with 7B reward ensembles, and we will post the results when available. (Although this causes a significant computational cost with 3x7B reward models, while our method only employs one 7B model, as the reviewer remarks that AdvPO wouldn't need to beat this baseline.)

[1] Coste, Thomas, et al. "Reward Model Ensembles Help Mitigate Overoptimization." ICLR 2023.
Dear authors,

Thank you so much for your detailed responses. I have increased my rating to 6.
Thank you for your response and clarification.

It is very unclear in the paper currently that you compare against UWO from Coste et al. as opposed to the mean ensemble or WCO. It would be beneficial to make that clearer, as well as the discussion of hyperparameter choice. I think it would also be beneficial to compare against the other methods there as well as just UWO.

> [on WARM] However, this approach still requires additional reward training and tries to perform mean optimization, i.e., optimizing towards average rewards from reward ensembles without considering uncertainties

I agree that conceptually it may be different to your method, but I believe it is still a necessary baseline to compare against, given it is tackling the same problem of overoptimisation.

Thanks for the rest of your response and clarification, it was very helpful. I am happy to raise my score to a 6, given the baseline being compared against is better than I thought (I believed it was mean ensemble, but it is actually UWO). I would raise my score higher if there were more baseline comparisons or results for the 7B ensemble reward models (although I realise those are unlikely in the remaining time).