Learning Goal-Conditioned Representations for Language Reward Models
Authors: Vaskar Nath
Dylan Z Slack
Jeff Da
Yuntao Ma
Hugh Zhang
Spencer Whitehead
Sean M. Hendryx
Venue: NeurIPS 2024 poster (2024)
Keywords: Goal-Conditioned Q-functions
Contrastive Learning
Reinforcement Learning from Human Feedback
Representation Learning
Reward Model
Primary area: natural_language_processing
TL;DR: We propose to enhance the learned representations of LLM reward models via a goal-conditioned contrastive learning objective, which we show improves reward model performance and downstream LLM alignment.

Abstract: Techniques that learn improved representations via offline data or self-supervised objectives have shown impressive results in traditional reinforcement learning.
Nevertheless, it is unclear how improved representation learning can benefit reinforcement learning from human feedback on language models.
In this work, we propose training reward models (RMs) in a contrastive, $\textit{goal-conditioned}$ fashion by increasing the representation similarity of future states along sampled preferred trajectories and decreasing the similarity along randomly sampled dispreferred trajectories.
This objective significantly improves reward model performance by up to 0.09 AUROC across challenging benchmarks, such as MATH and GSM8k. These findings extend to general alignment as well -- on the Helpful-Harmless dataset, we observe 2.3\% increase in accuracy.
Beyond improving reward model performance, we show this way of training RM representations enables improved steerability because it allows us to evaluate the likelihood of an action achieving a particular goal-state (e.g. whether a solution is correct or helpful).
Leveraging this insight, we find that we can filter up to 55\% of generated tokens during majority voting by discarding trajectories likely to end up in an "incorrect" state, which leads to significant cost savings.
We additionally find that these representations can perform fine-grained control by conditioning on desired future goal-states.
For example, we show that steering a Llama 3 model towards helpful generations with our approach improves helpfulness by $9.6$\% over a supervised-fine-tuning trained baseline.
Similarly, steering the model towards complex generations improves complexity by $21.6$\% over the baseline.
Overall, we find that training RMs in this contrastive, goal-conditioned fashion significantly improves performance and enables model steerability.

Learning Goal-Conditioned Representations for Language Reward Models
Vaskar Nath*
Dylan Slack
Jeff Da
Yuntao Ma
Hugh Zhang
Spencer Whitehead
Sean Hendryx
Scale AI
Abstract
Techniques that learn improved representations via offline data or self-supervised objectives have shown impressive results in traditional reinforcement learning. Nevertheless, it is unclear how improved representation learning can benefit rein- forcement learning from human feedback on language models. In this work, we propose training reward models (RMs) in a contrastive, goal-conditioned fashion by increasing the representation similarity of future states along sampled preferred trajectories and decreasing the similarity along randomly sampled dispreferred trajectories. This objective significantly improves reward model performance by up to 0.09 AUROC across challenging benchmarks, such as MATH and GSM8k. These findings extend general alignment as well on the Helpful- Harmless dataset, we observe 2.3% increase in accuracy. Beyond improving reward model performance, we show this way of training RM representations enables improved steerability because it allows us evaluate the likelihood of an action achieving a particular goal-state (e.g., whether a solution correct or helpful). Leveraging this insight, find that voting by discarding trajectories likely to end up in an "incorrect" state, which leads to significant cost savings. We additionally find that these representations can perform fine grained control by conditioning on desired future goal-states. For example, show that steering Llama 3 model towards helpful generations with our approach improves helpfulness by 9.6% over a supervised seneedd-iiettuue finn--tuningg trained baseline. Similarly, steering the model towards complex generations improves complexity by 21 21.6% over the baseline. Overall, we find that training RMs this contrastive, goal-conditioned fashion significantly improves performance and enables model steerability.
filter up to 55% of generated tokens during majority
1 Introduction
Aligning Language Models (LMs) with human preferences has proven to be an essential step for the adoption and safe use of these models, with the dominant paradigm being Reinforcement Learning from Human Feedback (RLHF) [50]. To accomplish this, standard setup is to collect labels from humans for generated responses (e.g.. preferences, quality ratings)
(50). These labels can then be used to train a reward model to produce a ranking/scoring of a given sequence or set of sequences.
"Equal contribution
Corresponding author: vaskar nathdecale.com
I Equual senior authorship
Code available at httppp/s//tthhcccco//aaaaaaaaaae
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Phomet
If Alice has 10 apples and gives away 2, how many apples does Alice have?
Preferre Respense
Disprelomed Respo see
Since Alice started off with tan] apples, we need to remove from 10. Hence, Alice has 10-2=8 apples.
Alice began with a total of 10 apples, so we need to [add] (wo to ten. Hence, Alice has 10+2=12 apples.
RM
RM
Hildaen Stane Impresartations
Hittden Shvie Regresentiaions
Peolle Cne
hopplhe
P
Macimiza Smnllrily
Miiiizu Simblaiiy
Figure 1: Overview of contrastive goal-conditioned learning for text. Pictured is prompt with preferred and dispreferred response. Both source state tokens (ten) for the positive and negative trajectory are sampled from the preferred response. For illustrative purposes, the positve and negative source states are sampled as the same token, in practice they can different. The positive goal state is sampled as some future token (subtract) from the preferred response, and the negative goal state sampled from any token (add) from the dispreferred response. The corresponding representations are retrieved from the last hidden state of the reward model. The training objective then to maximize and minimize the similarity of the positive and negative representation pairs, respectively.
The policy LM is then trained to maximize the expected returns from this reward model using a Reinforcement Learning (RL) algorithm.
High-quality representations have been shown to be an important piece for the success of RL algorithms 16 [2). Although such representations can be learned during end-to-end training, many efforts have found it important to integrate more explicit representation learning components into RL algorithms, such via data augmentation [39 auxiliary losses [21]. Some work even casts certain RL algorithms as representation learning methods where using the similarity between state representations serves value function, demonstrating success manipulation and navigation tasks [20). Despite these successes in different areas, representation learning for aligning LMs has been less explored, while more emphasis has been placed on, e.g., pre-training reward models 37
or learning from different types of rewards [15 38
In this paper, present simple yet effective approach improve the representations learned by reward models for aligning LMs. We train LM-based reward models to learn representations that capture the expected reward likelihood of achieving goal state (e.g., correct solution to problem, helpful response) at intermediate steps of the input sequence, inspired by goal-conditioned contrastive objective and apply the reward model's 's hidden representations from desirable and undesirable sequences. Enforcing this loss on representations from intermediate steps of the sequence helps encode dense signal to which trajectories promising at different points in the sequence, which we show offers several useful properties, such as helping to localize errors evaluating partial completed sequences. This method is flexible enough to support different kinds of alignment data and does not require further annotations beyond common sequence-level annotations.
RL 13 20). To do so,
more
We find that this approach improves the reward model's ability to identify correct/incorreet solutions in mathematical reasoning, boosting the AUROC on the task by up to 0.09 over standard preference ranking training. Towards natural language alignment, we find this method is able to increase the reward model's 's accuracy of identifying helpful versus harmless responses by 2.3% the Helpful- Hannless dataset [8].
We also demonstrate the utility of the learned representations themselves, e.g., for filtering solutions to improve accuracy and steering the outputs towards responses with certain attributes in guided decoding. On mathematical reasoning, show that are able to filter up to 55% of generated

---

tokens by discarding trajectories that are likely to lead to incorrect solutions as deemed by the learned representations while achieving similar or better performance. Similarly, using these representations to steer Ilama 3 [T] model by conditioning on desired future goal-states, we improve helpfulness by 9.6% correctness by 12.2%, coherence by 16.5%, and complexity by 21. 21.6% over the supervised- fine-tuning trained baseline.
In summary, our contributions are follows: 1) We explore improving the learned representations of reward models and its effect on LM alignment. Towards this, we present simple and effective representation learning method based on a goal-conditioned contrastive objective. 2) We demonstrate that training reward models with this method can improve reward model performance on mathematical reasoning and helpfulnesshaamlessnees
benchmarks. 3) We show that simply utilizing a reward model trained with this method can improve policy LM alignment on math reasoning benchmarks. 4) We investigate using these representations as a mechanism to evaluate the likelihood of achieving a desired goal state by filtering generations in majority-vott scheme and guided decoding, showing that they can be used increase accuracy and control policy LM outputs.
2 Related Work
Language model (LM) alignment and Reinforcement Learning from Human Feedback (RLHF). The goal of RLHF [50 71 (6] for aligning LMs is to train a policy LM that generates responses that are preferred by humans. This framework has proven instrumental for increasing the utility and generalization of LMs, improving their capabilities on a wide variety of tasks, and their safety [8] g
45 61 49.. Critically, RLHF improves the policy LM along scalar reward signal, which indicates the value of a generated piece of text, and their efforts have focused on different means of providing such signals 77 56). Typically, researchers train reward model that learns to score generations based on piir-wise human preference data. Subsequently, the model provides reward signals during reinfocement learning (RL) training of the policy, often via Proximal Policy Optimization (PPO) (59] While reward models have often been trained to score text with just single scalar value, several works have shown that training the reward model to provide fne-grained reward signals over spans of text can help improve model performance by providing intermediate rewards for the policy LM to learn from during RL costly human annotations on the sentence or token level to train the model, which can be prohibitive.
43 688. Nevertheless, fine-grained reward methods often require extensive and
Representation learning for RL. Reinforcement learning (RL) methods often learn representations as part of policy learning from reward signals. However, learning representations from solely reward signals can be highly sample inefficient, leading to poor generalization, and struggles in high-dimension state and action settings 41 63 74). Consequently, techniques that learn improved representations have been shown to significantly improve RL performance
[73 19 25
24).
Typically. these methods construct additional supervision via offline data or self-supervised objectives to drive representation learning [12 I1. These methods have achieved significant success in
variety of settings, such as vision-based control and multi-task RL [10] [4]] (30]. One such way of constructing supervision for representation learning is by leveraging future states and actions sampled from offline trajectories [48] 18, (6)) In particular, prior work has shown that increasing the similarity between state-action pairs sampled from the same trajectory can serve as a simple yet powerful way to learn improved representations for RL, that additionally acts as a type of goal condiiindiionnng by making representations more similar along the same trajectory 22 22 Still, representation and goal-conditioned RL. largely have not been considered for natural language domains.
3 Method
Drawing on recent efforts on goal-conditioned reinforcement learning 55] 15 and representation
learning 20 31 in other areas, we show that we can learn representations that approximate a Q- function within language reward model using contrastive learning. Since the Q-function quantifies the expected cumulative future rewards of taking a specific action from given state [47]. forcing the reward model's output scores to be entangled with an approximated Q-function has the potential to improve the reward model's credit assignment and preference ranking accuracy. We first present preliminaries on reward modeling and contrastive learning Section 3.1 and present our method in Section 3.2
3

---

3.1 Preliminaries
Preference ranking reward modeling for LMs. Typically, a reward model parameterized r(z,y) + I to return scalar reward, given prompt and completion sequence of tokens = yo....,rr) Vr]. Given a dataset 'D consisting of triples (x, ,0""",""" v²) of preferred and dispreferred completions, yn and y², respectively, the reward model optimized via paired loss:
L
[D] (aay" log (a(r(z, y²*) - (()))))))
Note, the reward model, r(z y), is trained to provide scalar feedback on the entire completion y.
Approximating Q-functions via contrastive learning. In deep reinforcement learning, the Q- function quantifies the expected cumulative future rewards of taking an action from a given state and parameterized by a neural network [7 20 Let & be the current state from which model can take an action oN, Taking action at the current state and continuing on that trajectory will lead to some future state Sg. The future states reached along this trajectory can be positive, of. by reaching a target goal state or negative, $g. . by reaching a different or undesired state. A critic function f can approximate a Q-function up to multiplicative constant by optimizing a contrastive learning objective [20]:
(f ae s))) log(1 of at 5k 1)).
Here, the critic function f is scoring function between a state-action encoder and state encoder.
3.2 Learning Goal-Conditioned Representations in Language Reward Models
In this section, we describe how to efficiently learn goal-conditioned representations using decoder- only transformer models [64] 54]. Additionally, we describe how to use these representations as Q-values that can be leveraged to detect generation errors and in guided decoding [33].
Reward modeling through the bottleneck of a Q-function. Given the prompt and a corresponding completion sequence y, we seek to learn a language reward model that both scores y and has representations that can be used to approximate Q7, (V(o.....)) 1t++), E [o, captures the expected reward for a policy π TBb+ (((o....t): 2), for some goal state Yo and next token yt+11 where rewards are defined as likelihood of achieving the goal at the next token generation. Intuitively, the goal state could be an optimal generation from the policy that a human would find satisfactory, such as the correct solution to a math problem or a helpful response. The Q-function captures the expected rewards from choosing token yt+1 the current step, considering the prompt x and tokens generated so far ((O.....t).
goal-conditioned Q-function. This Q-function
We initialize our reward model, r, from an instruction-tuned causal LM. We consider this reward model to be decomposed into two components: 1) feature extractor é, which produces a representation for a given token; 2) a reward projection head r' that predicts reward score based on that token representation. We use & to obtain representations of our state-action pair, h2+1, and arbitrary future states, hy:
h+41 = < 1
and h = -VJel11o...
(3)
Hence, we parameterize our reward model as
((I,8) r' (o(orllloo......-
(4)
In practice, the feature extractor is the LM decoder and the representations that the hidden token representations prior to predicting an output token. Meanwhile, the reward projection is either a linear layer or multi-layered perceptron (MLP) in our experiments.
With this reward model, we train the representations from the feature extractor é using a contrastive loss (20), which becomes:
C = lo (f((o........) Ft+1, (1 + log (1 (f M141, )));
(5)
where I (3 1141, )x) is the cosine similarity between the representations of the state-action
pair (1) Vt+1) and a future state ye. In Equation y and y7 are our positive and negative

---

goal states, which we discuss how to obtain in practice later in this section. By forcing the reward scores predicted by r' depend on the representations from ø. which is learned via this contrastive objective, the reward score incorporates signal as to which actions are closer to desired goal state.
We jointly train the reward model on the contrastive objective in Eq. 5 and the preference ranking objective optimized by Eq. Thus, our final objective is given by
R + x LC
=
where A is a hyperparameter to balance the tradcoff between optimizing the contrastive loss and the rewards.
In the context of RLHF, given that we have preference labels between two or more generations for the same prompt and recalling that L approximates Q-function, we can improve the term's alignment with human preferences by sampling positive goal states from preferred generations y and negative goal states from dispreferred generations yf.
Computing contrastive loss in practice. Given a preference- ranking instance from a training batch, let the preferred completion, y", and dispreferred completion. y", have lengths T and T", respectively. The representations of the positive state-action and goal state pair, (ys, ). are given by sampling two points, j.J ~ T). < j. from the preferred completion and retrieving hidden states h." and hy. Similarly, the representations of the negative pair, (ys', Vg ). are given by sampling any index m ["] from the preferred completion and any index n ~ ["] from the dispreferred completion and then retrieving hidden states h. and hrg We note that source states h²" and h% are sampled independently and are not necessarily equal. The cosine similarities between positive and negative pair representations are passed through the sigmoid function and averaged across the training batch. We will refer to this method as the Single Goal State (SGS) contrastive loss.
Computing Q-values. Once the reward model has been trained, we can use it to compute an additional Q-value, which gives us a measure proportional to the expected future rewards given the current state and a specific action. The Q-value at a token ł is the cosine similarity between the trained reward model's hidden state, h4, and goal state, hy. The the choice of goal state here is flexible For our experiments in Section [4.1] unless specified otherwise, we set the goal state to be the mean representation of all the preferred completion in the reward model training set. More explicity, the goal state is computed as:
hg =
Σ h(yw) MNED
N
where h(.) represents the reward model's last hidden layer of the final token of the completion. In Appendix C we ablate this choice and find this be the best method for computing the goal state during inference. Thus, we do not require any additional annotations aside from the starting preference ranking dataset. In experiments the goal state. The prototype is constructed using generations labeled as helpful, complex, coherent, and correct by annotators Nvidia's 's HelpSteer dataset [67]. Out of the highly-scored generations, we sample a subset and use the mean represetations of the sample as the goal state for the reward
LM steering, we construct a prototype to be used as
model. We ablate these choices in Appendix D.2
4 Experiments
To explore our method, we experiment with two settings: First, mathematical reasoning with code (Section 4.I. where LMs must use reasoning and coding to answer mathematical questions. The ability to code has become increasingly important for LLMs as it enables a broad set of capabilitiess for reasoning (e.g., tool use (58). This line of experiments elucidates the ability of our method, and the resulting representations, to guide and improve step-by-step reasoning capabilities. Second, natural language alignment (Section (42)). which an LM's responses must align with certain desired properties, such as helpfulness and harmlessness. These experiments show the capacity of reward models trained with our method to steer LM outputs more open-ended setting. In particular, we focus on helpfulness, harmlessness, and other desirable characteristics (e.g., correctness, coherence).
5

---

Cosallama 7b Reward
Q-Function 76 Reward
In-distribulion
Oul-of-distribulion
6.35
0.852
11:003
1.61
0.333
0.520
0.765
6761
C.75
AUROC
0.702
0.70
6.05
GSMBK
MATH
algebra222
GSM-Hard
Asdiv
mawps
svamp
Figure 2: AUROC scores comparing the baseline Codellama 7b Reward vs. our proposed method Q-Function 7b Reward on the rewards attributed to the base-model greedy generations across several math benchmarks.
4.1 Mathematical Reasoning with Code
4.1.1 Experimental Setup
For training reward models, we use the Open/aathInstruct-- dataset (62) math instruction tuning dataset with 1.8M problem-solution pairs on the training sets of GSM8k and MATH, which are generated by the Mixtral-8x -B model [32]. To form reward model training set, create pairs of solutions for the same problem. where one solution is correct (i.e., preferred) and the other is incorrect (i.e., dispreferred).
Our base model is OpenMath- Pa--ytammm-----yhh--hhhh 62), a strong decoder-only LM that has been instruction tuned on Ipstuuthhnstrrutt-11 to use Python interpreter in order to solve mathematical problems, which we add a reward head to predict reward scores (Section (.22). This model trained with the traditional preference ranking objective serves as our baseline (Codellama 7 Reward), which we compare with the same model trained using our method (Q-Function 7b Reward).
We evaluate on several popular math benchmarks. Since the problems in the preference ranking dataset come from the training sets GSM8k and MATH, we consider their respective test splits to be in-distribution (ID). We also evaluate on test sets we consider to be out-of -iitt--disttibttionn (OOD), namely algebra222 (27, GSM-Hard 23 Asdiv 466 mawps 35 and svamp 52
to test for further
generalization.
Detailed information on data, models, and training can be found in Appendix A
4.1.2 Reward Modeling Evaluations
We obtain completions from the base model by greedy decoding, following prior work [2 66), and measure the reward score AUROC. This quantifies the ability of model's reward score to discern between correct/incorrect sequences across different thresholds. In Appendix E we additionally assess Q-Function 7b Reward's 's ability to rank completions.
Learned representations help reward scores discern correct/incorrect solutions. In Figure 2 we see that the reward model trained with our representation learning approach consistently achieves higher AUROC scores compared to the model trained with only the standard preference ranking loss. On the ID datasets, we see a gain of approximately 0.05 and 0.09. Meanwhile, on the OOD datasets, the average improvement is 0.04, with maximum of roughly 0.08 on svamp. Although these datasets are math problems, they vary in the types and difficulties of questions asked. Despite being OOD, we see that our method improves the performance when judging these prompts and generations. Overall, these results suggest that the representations learned using our approach help improve the reward score predictions such that they are more accurate and generalizable.

---

GSMBK
MATH
Codellara Reward
Codellama 7b Reward
D-Function 7b Reward
Q-Function 7b Reward
0.0-
679
AUROC Score
AUROC Score
3.75
0.55
Percentile Model Generation Considered
Percentile Model Generation Considered
Figure 3: AUROC scores the rewards attributed partial base-model generations across 50 samples on GSM8k and MATH. The error bars depict the 95% confidence intervals (with sample size 50) at each percentile of generation considered. The Q-Function reward model has incremental increase in performance with more information, whereas, the traditional reward model's performance lot more varied in attributing intermediate rewards.
These representations also help reward scores judge which partial solutions are promising. Part of the motivation and design of our approach is encode the likelihood of reaching certain goal state at intermediate steps in the generated sequence. To examine how well our method accomplishes this, measure reward score AUROC when the model is only shown up to certain percentage of the generated sequence (i.e., percentile). Here, we show this on the ID datasets, where we consider the 50 sample generations from the base model and plot the mean AUROC across the 50 samples. Plots for other benchmarks can be found in Appendix [F Looking at Figure 3 we find that the AUROC scores for both reward models tend to increase they see more of generated sequence, but Q-Function 7b Reward largely maintains higher AUROC across different percentiles and generally exhibits fewer significant drops between percentiles. These improvements across percentiles imply that the reward scores do indeed better capture which generated sequence are likely to be correct or incorrect.
4.1.3 Directly Using the Learned Representations
Our experiments thus far illustrate the effect of the learned representations on the reward scores. Here, we turn our attention to utilizing these representations directly to explore some of the encoded properties. In these experiments, we filter generations from the base model using the Q-values (i.e., cosine similarities between intermediate and goal state representations). We take 50 base model generations for all the GSM8k and MATH test problems, which are provided by Toshiniwal et al., and we filter any generation that has a Q-value less than zero (62). We compare this against baseline of majority voting over the 50 generations.
Q-values can prune incorrect sequences. We find that filtering based on Q-values provides meaningful improvements in accuracy compared to majority voting over all 50 samples, based on Table ! Specifically, the majority vote accuracy improves by 1.2% and 3.9% on GSM8k and MATH, respectively. Looking at the increased proportion of correct sequences after filtering, it appears that sequences with the negative Q-values generally correspond to incorrect sequences and are filtered out in this setting, thereby improving the accuracy. On the other hand, we also see that the average percentage sequences that are filtered out can be quite large, such as 65.6% on the MATH dataset, and some correct solutions can also be filtered out. Nevertheless, using the Q-values based these learned representations can help improve the accuracy of such majority voting schemes.
Q-values can detect errors in reasoning. Experiments on rewards scores for partial sequences (Section 4.122) suggest that the learned representations help capture the likelihood that a sequence will reach goal state from an intermediate step. We can also qualitatively explore this property using Q-values and the representations themselves. based on the Q-value filtering process above. In the left example, the Q-value drops when there is mistake of multiplying the "donuts ddonuts_prices ppwpceeeeeeeee pateepeerdzzeeedd dozen" by "donuts peenutspper__e _er__e__d____ dozen".. The Q-value in the right example similarly shifts once the "jazz_dance" www.dddaeeeceeeeeeee
Figure we show examples that are filtered
is naively added to the
7

---

GSMBK
Accuracy (%)
Prop. Correct (%)
Avg. Filtered (%)
Avg. Tokens Saved (%)
Maj. @ 50
84.8
74.9
Filtered
86.0
84.0
35.9
25.5
MATH
Accuracy (%)
Prop. Correct (%)
Avg. Filtered (%)
Avg. Tokens Saved (9)
Maj. 50
55.6
40.2
Filtered
59.6
52.0
65.6
55.6
Table I: Majority vote versus Q-value filtering performance. Accuracy shows the final majority voting accuracy on the respective benchmarks. Prop. Correct is the proportion of the correct class in the final sample set considered. Avg. Filtered is the average percentage (out of 50 total generations) that were discarded across all problems the benchmark. Avg. Tokens Saved is the average percentage of tokens that are saved with the assumption that the remaining tokens after the first negative Q-value discarded across all the problems in the benchmark. Both the difference in accuracy and proportion of correct class are t-test significant.
hikery hought types of postrics, bought dozen donuts which vost per dozen, dozen mini cupcakes which
Problem: Toula
Problem: in dance class 20 students, 20% enrolled in contemparary dance.
25% remsining earolled jazz dance. and the curolled hip hop
cost 60 doren, and 6 dozen mini cheesecakes for 55 per dozen. Hew much the rotal court
drace. percentage
students ensolled in hip-hop dasce?
and are preton tane Pyako cosa
ap pes fuated po Lanso
excoe
Ilr Lp ane prenee
dcm
ccm
4m
"is oo jodel ti Ja
proce
mine
Figure 4: Two examples from the GSM8K dataset that was filtered via the Q-value pruning. The token level Q-values are portrayed as heat map where the colors red and blue represents scores close to -1 and 1, respectively. Both examples illustrate that the Q-values pinpoint the exact logical error in reasoning. The full version of these examples can be found in Appendix G
"contemporary_dan dance peaancceepeeeeaaee..... Overall, these examples illustrate the potential effectiveness of using the representations and resulting Q-values for localizing errors and identifying promising sequences.
4.1.4 Aligning Policy Models with RLHF
We investigate whether the reward model trained with our proposed method provides a better reward signal when performing RLHF, which is an important usage of these reward models. We compare the math problem solving performance of policy model that is aligned using preference-ranking reward model versus our Q-Function reward model in the Proximal Policy Optimization (PPO) (59 algorithm. We additionally compare the performance against DPO as it serves as a strong baseline [55]. The prompt dataset given by problems in the training sets of GSM8k and MATH. To introduce diversity in the prompts, which has been shown to improve PPO, we also add problems from the MathQA dataset a large scale dataset with real world math problems [4. We keep the settings and hyperparameters during PPO constant, and further details are in Appendix A
Table shows that, given the same settings, utilizing reward model trained with our method improves policy model accuracy across several benchmarks. We see an average gain, weighted by benchmark size, 1.7%. Intuitively, this shows that the improvements we observe in reward scores (Section 4.2 do translate to improvements in the policy model via RL, which further supports the value of the representations learned via our approach. Although, the gains in accuracy are not as significant compared to those observed with the reward scores. A possible explanation for this is that, due to experimental constraints such as additional data collection, the preference-ranking dataset used for reward model training are off-policy Mixtral-8x71 model [32] generations. Hence, clear next step for this setting would be to retrain the reward model with on-policy preference-ranking and iterate. Interestingly, observe that DPO performs par with our method for the ID tasks

---

(GSM8k and MATH), However, the policy trained with the contrastive RM performs much better on OOD tasks (algebra222, GSM-Hard, Asdiv, mawps, and svamp) when compared DPO. These findings could be explained by related works that analyze DPO and find it can perform poorly in OOD settings 772.. All in all, these results indicate that PPO training with the contrastive RM leads to better generalization than training with DPO, particularly in OOD settings.
Model
ID
OOD
GSM8K
MATH
algebra222
GSM-Hard
Asdiv
mawps
svamp
Codellama Base
75.9
43.6
65.6
60.1
77.7
93.5
79.6
DPO
80.1
44.8
64.8
59.9
76.9
90.3
76.6
Codellama PPO
79.3
43.4
65.8
61.1
77.4
91.6
78.5
Q-Function PPO 80.5
45.1
70.9
62.7
79.5
93.6
81.2
Table Accuracy of the policy model trained via PPO using preference ranking reward model vs our Q-Function reward model. We present average accuracy across 4 independent runs. The base model results are also shown as a reference point presented in [62] and DPO results are shown as a strong baseline. The full table with confidence intervals is in the appendix (Ta
4.2 Natural Language Alignment
While fine-grained feedback (and subsequently, Q-values) are most sensible in reasoning environ- ments with explicitly incorrect logical steps, natural extension of the experiment to apply our approach to reward models in the natural language setting, where we expect models to learn improved representations with respect their alignment towards human preferences. We first explore this for training a reward model for helpfulness and harmlessness [8]. Then, we explore if the learned representations can be further used for steering language model inference time, via guided decoding. to generate sequences that are more helpful, complex, or coherent. The backbone of our reward model Llama I-st-unstt and detailed information can be found in Appendix B.2
4.2.1 Improving Reward Modeling for Helpfulness and Harmlessness
We use the Helpful-Harmles dataset g for training and evaluation. We select this dataset because of it's large scale, established representation in the RLHF community, and general purpose nature which is helpful avoid restricting our natural language reward model towards a specific domain. Similar to our previous experiments, we compare training with our approach versus standard reward model training. Table 3 (left) shows that training with our method also improves reward scores on this natural language data, increasing the ability of the reward model to discern between helpful and unhelpful responses.
4.2.2 Reward Model Representations can Steer Language Models Towards Helpful Generations
A practical extension of learning stronger natural language representations is using the extensions for steering the LM. To test this, we experiment with guided decoding [3] 70 with
In--tuutttt which defines a decoding objective function to guide inference during beam search. In traditional guided decoding (and our SFT baseline), the confidence score of the model is extracted from an additional inference step on each beam. In our method (Q-Function Prototype), the confidence score is the cosine similarity between the sequence embedding from the reward model and the hidden state 20 2). We construct the prototype by embedding 20 examples that score highly in all categories (helpfulness, correctness, coherence, complexity) using HelpSteer, a model alignment and stecring dataset [67]. The reward model is trained using HelpfulHarmless per previous sections [ Details for guided decoding are in the Appendix
Table 3 (right) shows our results, and evaluation is done via LLM-as judge (GPT-4) (7). Using the representations fine-tuned by the contrastive loss improves across all categories and generally steers generations towards more helpful generations. We find that complexity improves the most, suggesting that the generations produced by the using the prototype helps the model avoid beams that are too simple. Suprisingly, the correctness of the responses is increased also, suggesting that

---

although the SFT prototype usually prefers beams that are simple and correct, the prototype succeeds picking beams with added complexity without sacrificing correctness or helpfulness.
Model
HH
Model
Helpful
Correct.
Coherent
Complex.
Llama 8b Reward 68.2
SFT
45.2
43,9
41.8
39.2
Q-Function 8b Reward
70.5
Prototype 54.8
56.1
58.3
60.8
Table Comparison of model performances different settings. (a) Reward model accuracy on the Helpful-Harmless test set. Our result highlights that the representations learned by the Q-function are helpful for discerning between desirable and undesirable sequences in the natural language setting. (b) Winrate (%) vs. SFT model. In the Q-Function prototype setting, we use the cosine similarity of the sequence and the prototype (made from 20 examples that score high on all 4 categories) to score each beam. In the SFT setting, the model is fine-tuned on the same examples as the prototype and guided decoding is run via model scores. Evaluation is done via GPT-4.
5 Limitations and Future Work
Deriving more informative goal states. In Appendix D.2 we demonstrate the importance of having a precise and meaningful goal state. At train time, our SGS method achieves this by looking at future tokens of the preferred and dispreferred completions. During inference, however, we take the mean representation of multiple goal state completions from the training data. Although we see that the SGS method generalizes despite this discrepancy. future work could investigate other methods to derive more precise inference- -errenece--eeeeeee goal states.
Partial Q-values correlate with partial reward scores. An interesting observation, shown in Appendix H. that the reward scores for partial completions correlated with the Q-values. While our experiments show that this can benefit reward scores, especially for math reasoning, disentangling these values may be desirable in some instances. For example, one may want to incorporate the completeness of the generated sequence in the reward score, such as for creative writing. meaning incomplete sequences should strictly get low reward scores. However, currently, partial sequences may still score highly if they are likely to reach goal state. This is expected as the final reward score derived by a linear projection from the hidden representation, while the Q-values are derived as the cosine similarity between that same representation and goal state. In Appendix [H] we explore the effects of learning an MLP projection for the rewards to help decouple the two signals. Further disentangling these signals in the representations could be an interesting direction for future work.
Further advancing policy models. A clear line of future work here is to integrate the additional Q-value signal that we gain from this method into the RL algorithm to further improve policy model training. Future work should also investigate the impacts of iteratively retraining the reward model on on-policy completions and the policy model on the updated reward model.
6 Conclusion
Learning proper representations in the reward model essential for the model's generalization and ability to produce accurate rewards during RLHF We introduce a method of applying goal-conding condiioiit Q-functions to learn representations that capture the expected reward for a goal-conditioned policy. On math benchmarks, such as GSM8k, we improve performance and show that the reward model has a greater ability to discern between correct and incorrect solutions. On natural language alignment experiments, we show improvement in reward model performance and that the embedding can be used for natural language steering. Our findings show that using Q-values during reward model training can improve the representations of the reward model and suggest promising future directions for further advancing language models via R..HE.
References
[1] Meta AI. Introducing meta Ilama 3: The most capable openly available Um to date. https mwwa/ 1lama-3. 3/
//ai meta. com/blog/ wwwbb/llgg//at....eee
10

---

[2] Arwa Alanqary, Gloria Lin, Joie Le. Zhi-Xun Rua- Tan, Vikash K. Mansinghka, and Joshua B. Tenenbaum. Modeling the mistakes of boundedly rational agents within a bayesian theory of mind. ArXiv, abs/2106.1322222 2021.
[3] Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, and Yuxiong He. Deep- speed inference: Enabling efficient inference of transformer models at unprecedented scale, 2022.
[4] Aida Amini, Saadia Gabriel, Shanchuan Lin, Rik Koncel-Kedziorske Yejin Choi, and Hannaneh Hajishirzi. Mathqa: Towards interpretable math word problem solving with operation-based formalisms. CoRR, abs/1905. 13b319,9..,339,, 2019.
[5] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay, Advances in neural information processing systems, 30, 2017.
[6] Raghuram Mandyam Annasamy and Katia Sycara. Towards better interpretability in deep g-networks. In Proceedings of the AAAI conference artificial intelligence, volume 33, pages 4561-4569, 2019.
[7] Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant laboratory for alignment. arXiv preprint 12Xiv:21 2021.
[8] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.0580 2022.
[9] Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, et al. Constitutional ai: Harmlessness from ai feedback. arXiv preprint arXivv2212.08077 2022.
[10] Andrea Banino, Adria Puigdomenech Badia, Jacob C Walker, Tim Scholtes, Jovana Mitrovic, and Charles Blundell. CoBERL: Contrastive BERT for reinforcement learning. In International Conference on Learning Representations, 2022.
[11] André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, David Silver, and H. V. Hasselt. Successor features for transfer in reinforcement learning. ArXiv, abs/1606.00 15s/1660605122 2016.
[12] Diana Borsa, André Barreto, John Quan, Daniel Jaymin Mankowitz, Rémi Munos, H. Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. ArXiv, abs/1812.07626, 2018.
[13] Elliot Chane-Sane, Cordelia Schmid, and Ivan Laptev. Goal-conditioned reinforcement learning with imagined subgoals. In International Conference on Machine Learning, pages 1430-1440. PMLR, 2021.
[14] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1597-1607. PMLR, 2020.
[15] Ching-An Cheng, Andrey Kolobov, Dipendra Misra, Allen Nie, and Adith Swaminathan. LIf-bench: Benchmark for interactive learning from language feedback. arXiv preprint arXiv:2312.06853, 2023.
[16] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14166 14168, 2021.
[17] Tri Dao. F Faster attention with better parallelism and work partitioning, 2023.

---

[18] Alexey Dosovitskiy and Vladlen Koltun. Learning to act by predicting the future. ArXiv, abs/161 2016.
[19] Simon S. Du, Sham M. Kakade, Ruosong Wang, and Lin F. Yang. Is a good representation suffi- cient for sample efficient reinforcement learning? In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[20] Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Ruslan Salakhutdinov. Contrastive learning as goal-conditin condiiiodt reinforcement learning, Advances in Neural Information Process- ing Systems, 2022.
[21] Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pages 512-519. IEEE, 2016.
[22] Carlos Florensa, Jonas Degrave, Nicolas Manfred Otto Heess, Jost Tobias Springenberg, and Martin A. Riedmiller. Self-supervised learning of image embedding for continuous control. 2019.
ArXiv,
[23] Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. Pal: Program-aided language models. arXiv preprint arXiv:2211. 10435, 2022.
[24] Carles Gelada, Saurahh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. Deep- MDP: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Confer- ence on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2170-2179. PMLR, 09-15 Jun 2019.
[25] Dibya Ghosh and Marc G. Bellemare. Representations for stable off-policy reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3556-3565 PMLR, 13-18 Jul 2020.
[26] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus. Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fer- nández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-362, September 2020.
[27] Joy He- Yueya, Gabriel Poesia, Rose Wang, and Noah D Goodman. Solving math word prob- lems by combining language models with symbolic solvers. arXiv preprint arXiv:2304.0911 09100, 2023.
[28] Dan Hendrycks, Collin Burns, Sauray Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. Advances in neural information processing systems, 33, 2021.
[29] Jian Hu, Xibin Wu, Weixun Wang, Xianyu, Dehao Zhang, and Yu Cao. Openrlhf: An easy-to- use, scalable and high-performance rlhf framework, 2024.
[30] Max Jaderberg. Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017.
[31] Ashish Jaiswal, Ashwin Ramesh Babu, Mohammad Zaki Zadeh, Debapriya Banerjee, and Fillia Makedon. A survey on contrastive self-supervised learning. Technologies, 9(1):2, 2020.
[32] Albert Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, al. Mixtral of experts. arXiv preprint 400.v:2401.0440 2024.
12

---

[33] Minbeom Kim, Hwanhee Lee, Kang Min Yoo, Joonsuk Park, Hwaran Lee, and Kyomin Jung. Critic-guided decoding for controlled text generation. arXiv preprint arXiv:2212.10998, 2022.
[34] Hannah Rose Kirk, Alexander Whitefield, Paul Rottger, Andrew M. Bean, Katerina Margatina, Juan Ciro, Rafael Mosquera, Max Bartolo, Adina Williams, He He. Bertie Vidgen, and Scott A. Hale. The prism alignment project: What participatory, representative and individualised human feedback reveals about the subjective and multicultural alignment of large language models.
ArXiv, abs/2404. 16019, 2024.
[35] Rik Ckoncel-Kedziorss Subhro Roy, Aida Amini, Nate Kushman, and Hannaneh Hajishirzi. Mawps: A math word problem repository, pages 1152-1157, 01 2016.
[36] Andreas Kopf, Yannic Kilcher, Dimitri von Rutte, Sotiris Anagnostidis, Zhi Rui Tam, Keith Stevens, Abdullah Barhoum, Nguyen Minh Duc. Oliver Stanley, Rich ard Nagyfi, ES Shahul, Sameer Suri, David Glushkov, Arna Dantuluri, Andrew Maguire, Christoph Schuhmann, Huu Nguyen, and Alexander Mattick. Openassistant conversations democratizing large language model alignment. ArXiv, abs/2304.07327, 2023.
[37] Tomasz Korbak, Kejian Shi. Angelica Chen, Rasika Vinayak Bhalerao. Christopher Buckley, Jason Phang. Samuel R Bowman, and Ethan Perez. Pretraining language models with human preferences. In International Conference on Machine Learning, pages 17506-17555 PMLR, 2023.
[38] Minae Kwon, Sang Michael Xie, Kalesha Bullard, and Dorsa Sadigh. Reward design with language models. arXiv preprint arXiv:2303.00011, 2023.
[39] Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. 33: 1388448--19895, 2020.
[40] Tor Lattimore, Csaba Szepesvari, and Gellert Weisz. Learning with good feature represen- tations in bandits and in RL with generative model. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 1 119 of Proceedings of Machine Learning Research, pages 5662-5670. PMLR, 13-18 Jul 2020.
[41] Charline Le Lan, Stephen Tu, Adam Oberman, Rishabh Agarwal, and Marc G. Bellemare. On the generalization of representations in reinforcement learning. In Gustau Camps- Valls, Francisco J. R. Ruiz, and Isabel Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 4132 4157. PMLR, 28-30 Mar 2022.
[42] Yitao Liang, Marlos C Machado, Erik Talvitie, and Michael Bowling. State of the control of atari games using shallow reinforcement learning. arXiv preprint arXiv: 2015.
[43] Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step, 2023.
[44] Guoqing Liu, Chuheng Zhang, Li Zhao, Tao Qin, Jinhua Zhu, Li Jian, Nenghai Yu, and Tie- Yan Liu. Return-based contrastive representation learning for reinforcement learning. In International Conference on Learning Representations, 2021.
[45] Jacob Menick, Maja Trebacz, Vladimir Mikulik, John Aslanides, Francis Song, Martin Chad- wick, Mia Glaese, Susannah Young. Lucy Campbell-Gillinn Gillingham, Geoffrey Irving, et al. Teaching language models to support answers with verified quotes. arXiv preprint urXiv:2203.1114. 2022.
[46] Shen- Yun Miao, Chao-Chun Liang, and Keh-Yin Su. A diverse corpus for evaluating and developing english math word problem solvers, 2021.
[47] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015.
[48] Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. In Neural Information Processing Systems, 2018
13

---

[49] Rejichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christo- pher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, et al. Webept: Browser-assisted question-answeringg with human feedback. arXiv preprint arXiv:21 2021.
[50] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
[51] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach Vito,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019.
[52] Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are nlp models really able to solve simple math word problems?, 2021.
[53] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, Prettenhofer, R. Weiss, V. Dubourg. J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2888 28830, 2011.
[54] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[55] Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model, 2024.
[56] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2024.
[57] Joshua Robinson, Ching- Yao Chuang, Suvrit Sra, and Stefanie Jegelka. Contrastive learning with hard negative samples. arXiv preprint arXiv:2010.04552 2020.
(58) Timo Schick, Jane Dwivedi- Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. Advances in Neural Information Processing Systems, 35, 2023.
[59] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint w:Xiv:1707006477 2017.
[60] Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, and Sergey Levine. contrasttt networks: Self-supervised learning from video. 2018 IEEE International Conference on Robotics and Automation (ICRA). pages 1 1134-1141, 1141, 2017.
[61] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-302) 2020.
[62] Shubham Toshniwal, Ivan Moshkov, Sean Narenthiran, Daria Gitman, Fei Jia, and Igor Gitman. Openmathinstruct-I1 1: A 1.8 million math instruction tuning dataset. arXiv preprint arXiv: 1001 2024.
[63] Frederik Träuble, Andrea Dittadi, Manuel Wuthrich, Felix Widmaier, Peter Vincent Gehler, Ole Winther, Francesco Locatello, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. Representation learning for -ittr--distriittion generalization in reinforcement learning. In ICML 202/ Workshop on Unsupervised Reinforcement Learning. 2021.
[64] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
14

---

[65]
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cour- napeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nel- son, Eric Jones, Robert Kern, Eric Larson, CJ Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, António H. Ribeiro, Fabian Pedregosa, Paul van Mul- bregt, Aditya Vijaykumar, Alessandro Pietro Bardelli, Alex Rothberg, Andreas Hilboll, Andreas Kloeckner, Anthony Scopatz, Antony Lee, Ariel Rokem, C. Nathan Woods, Chad Fulton, Charles Masson, Christian Häggström, Clark Fitzgerald, David A. Nicholson, David R. Hagen, Dmitrii V. Pasechnik, Emanuele Olivetti, Eric Martin, Eric Wieser, Fabrice Silva. Felix Lenders, Florian Wilhelm, G. Young, Gavin A. Price, Gert-Ludwig Ingold, Gregory E. Allen, Gregory R. Lee, Hervé Audren, Irvin Probst, Jörg P. Dietrich, Jacob Silterra, James T Webber, Janko Slavië, Joel Nothman, Johannes Buchner, Johannes Kulick, Johannes L.. Schönberger, José Vinícius de Miranda Cardoso, Joscha Reimer, Joseph Harrington, Juan Luis Cano Rodríguez, Juan Nunez- Iglesias, Justin Kuczynski, Kevin Tritz, Martin Thoma, Matthew Newville, Matthias Kümmerer, Maximilian Bolingbroke, Michael Tartre, Mikhail Pak, Nathaniel J. Smith, Nikolai Nowaczyk, Nikolay Shebanov, Oleksandr Pavlyk, Per A. Brodtkorb, Perry Lee, Robert T. McGibbon, Roman Feldbauer, Sam Lewis, Sam Tygier, Scott Sievert, Sebastiano Vigna, Stefan Peterson, Surhud More, Tadeusz Pudlik, Takuya Oshima, Thomas J. Pingel, Thomas Robitaille, Thomas Spura, Thouis R. Jones, Tim Cera, Tim Leslie, Tiziano Zito, Tom Krauss, Utkarsh Upadhyay, Yaroslav O. Halchenko, and Yoshiki Vázquez- Baeza. Scipy 1.0: fundamental algorithms for scientific computing in python. Nature Methods, 17(3):261-21) February 2020.
[66] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203. 1171, 2022.
[67] Zhilin Wang, Yi Dong, Jiaqi Zeng, Virginia Adams, Makesh Narsimhan Sreedhar, Daniel Egert, Olivier Delalleau, Jane Polak Scowcroft, Neel Kant, Aidan Swope, and Oleksii Kuchaiev. Helpsteer: Multi-attribute helpfulness dataset for steerlm. ArXiv, 0tts/2311.09555 2023.
[68] Zeqiu Wu, Yushi Hu, Weijia Shi, Nouha Dziri, Alane Suhr, Prithviraj Ammanabrolu, Noah A. Smith, Mari Ostendorf, and Hannaneh Hajishirzi. Fine- Fine-grained human feedback gives better rewards for language model training, 2023.
[69] Yuxi Xie, Kenji Kawaguchi, Yiran Zhao, Xu Zhao, Min- Kan, Junxian He, and Qizhe Xie. Decomposition enhances reasoning via self-evaluation guided decoding, 2023.
[70] Yuxi Xie, Kenji Kawaguchi, Yiran Zhao, Xu Zhao, MingSung Kan, Junxian He, and Qizhe Xie. Self-evaluation guided beam search for reasoning. Advances in Neural Information Processing Systems, 2023.
[71] Jing Xu, Megan Ung, Mojtaba Komeili, Kushal Arora, Y-Lan Boureau, and Jason Weston. Learning new skills after deployment: Improving open-domain internetdrriven dialogue with human feedback. arXiv preprint arXiv:2208.0322 2022.
[72] Shusheng Xu, Wei Fu, Jiaxuan Gao, Wenjie Ye, Weilin Liu, Zhiyu Mei, Guangju Wang, Chao Yu, and Yi Wu. Is dpo superior to ppo for Ilm alignment? a comprehensive study, 2024.
[73] Jiaqi Yang, Wei Hu, Jason D. Lee, and Simon Shaolei Du Impact of representation learning in linear bandits. In International Conference on Learning Representations, 2021.
[74] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block MDPs. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pages 1211411 1224. PMLR, 13-18 Jul 2020.
[75] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Erie Xing, Haotong Zhang, Joseph Gonzalez, and Ion Stoica. Judging llm-as-a-judge with mt- mntbbench and chatbot arena. ArXiv, fals/2306.05685,, 2023.
15

---

[76] Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford. Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprinl arXiv: /99.088999 2019.
16

---

Appendix to Learning Goal-Conditioned Representations for Language Reward Models
A Mathematical Reasoning with Code Details
In this appendix, discuss further experimental details for results in Section 4.1
A.1 Reward Model Training and Evaluation Details
Preference ranking dataset. To construct the preference ranking dataset, we used the OpenMathinstruut-II a math instruction tuning dataset with 1,8M problem-solutin pairs gener- ated by the Mixtral-8x7B model. The questions come from a blend of the training subsets for GSM8K and MATH. We matched up generations that had the final answer output correct vs. incorrect as the preferred and dis-preferred generation. Some problems are used several times, but preferred and all dis dis-preferred solutions unique. In total, curated 257, 584 total preference rankings over total of 12, 518 unique problems.
Prompt format. To prompt the base model to answer a math question, we use the same Nemo prompt format that was used to fine-tune the model. This prompt format was used both for reward model training, PPO training, and all evaluations. The following is the prompt format:
System:
You're re an expert Python programmer and mathematician. Help the user solve this problem using code when necessary. Make sure to put the answer (and only answer) inside \boxed[].
User: {problem}
Assistant:
where {problem} is replaced with math problem from the dataset.
OpenRLHF library. We implement the reward model training with contrastive loss and the baseline of standard preference ranking using the OpenRLHF library 29 in combination with Py Torch [51]. We made the following modification the OpenRLHHF Library:
A new contrastive loss method that computes the contrastive objective
Optionally add auxiliary loss the reward model training objective experimenting with our reward model method
Optionally change the reward model reward head from a linear projection to multi-layer perceptron. The MLP layer is simple fully connect layer with ReLU activation in between, all with the same hidden size as the representation dimension.
Reward model training parameters. We train both the baseline reward model and the Q-Function reward model as well as any reward model ablations for epoch with batch size of 64.
The hyper-parameter for the contrastive lambda of 0.5 was chosen by doing a sweep of 0.1, 0.5, 0.9, and found that the 0.5 setting had the most stable to training loss as well as training rewards for preferred vs. unpreferred completions.
Reward model evaluation. We evaluate the Reward Model performance the test splits of several math benchmarks: GSM8K [16), MATH [4) algebra222 (27). GSM-Hard [23]. Asdiv [46]. mawps (35) and svamp [52]. To get set of completions on each problem in these set of benchmarks, we use the greedy generation of the base model along with annotations of whether each completion
provides the correct answer, which are provided by Toshiniwal al. [2] Then for each completion, we first format using the Nemo format template described earlier, and use the Reward Model to predict a reward score. predicting the reward score, we also retrieve the last hidden state of the
Reward Model well the predicted reward score for each token in the completion, not including
17

---

hyperparameter and settings
value
batch size
64
learning rate representation dimension
x 10-6
4096
contrastive lambda (A)
0.5
num epoch
preference ranking loss function optimizer seed
sigmoid Adam
42
Table 4: pppam--rrameterss for Reward Model training
the prompt tokens. Utilizing the Reward Model predicted reward score and the annotation of the correctness of each completion, we compute an AUROC score using the Python scikit-learn package plot the partial AUROC scores every tenth percentile of each
[53]. Using the same procedure, completion. The AUROC scores are compared between Q-Function 7b Reward (our method) and Codellama 7b Reward and the results summarized in Section
A.2 Learned Representations Experiment Details
Computing Q-values. In order to retrieve the Q-values, take the cosine similarity between the goal state and the hidden states for each token, which is retrieved during the forward pass when collecting the reward scores. The goal state in all the experiments is the mean representation of all the preferred completion the preference ranking dataset. More explicitly, we pass each preferred problem and response pair through the trained Reward Model and then we take the average across the last hidden state of the final token of each problem and response pair.
Filtering sampled generations. In order understand the value of these learned representations, we use these Q-values to filter generations. For each of the In-Distribution datasets, GSM8K [16] and MATH 4 each completion and we discard it if the Q-value for any the tokens less than O, which indicates that the expected future rewards from that state is low. This is performed for each problem and each generation sample. Using the final set of remaining completions, we then compute the majority vote accuracy, comparing against the baseline of majority vote 50. Furthermore, in order to visualize these Q-values, we plot heat map of these values on the corresponding tokens. The results and plots from this experiment is detailed Section 4.1.3
the 50 completion samples provided by Toshiniwal et al. [62]. We take
A.3 RLHF Training and Evaluation Details
In these set of experiments, we investigate whether the reward model trained with our proposed method provides a better reward signal when performing RLHE, Here we explain further experimental details for the results in Section
Prompt dataset. The prompt dataset for PPO comprises of the training problems from GSM8k and MATH, which constitutes a total of 12,518 problems [16] 28). We also incorporate all the problems from the MathQA dataset [4], which has 37.297 unique problems. Thus, in total, the prompt dataset for PPO has 49,815 unique problems. In addition, each training step of the policy model, we add small pretraining coefficient. The SFT dataset for the pre-training coefficient is comprised of all the preferred completions in the preference ranking dataset.
PPO training parameters. We also use the OLe package and PyTorch for PPO training 29 51
No major code changes were made to the OpenRLHF package, with the exception of the dataset loader for our custom prompt dataset. Table summarizes the hyperparameters used for all PPO training reported in this paper.
Policy model evaluation. We evaluate the policy-trained model on the test set on all the aforemen- tioned math benchmarks. For each problem in these benchmarks, we process the input to use the prompt format detailed above and perform zero shot evaluation. We perform 4 independent runs to account for variance. In order to perform this evaluation, we utilize the Nemo-Skan Skill [62] package. As
18

---

hyperparameter and settings
value
train batch size
128
number of episodes
1 512
rollout batch size training epochs
1
actor learning rate
3 x 10-6
critic learning rate
9 10-6
init kl coeff
0.02
normalize reward
true
generation max length
1024
GAE gamma
1.0
GAE lambda
0.95
policy model temperature
1
pretraining coefficient
0.05
learning rate
9 10 10-6
representation dimension
4096
contrastive lambda (A) optimizer
0.5
Adam
Table 5: parpee-rrraeeers for PPO training
a pre-requesite to performing evaluation on this package, we first had to convert the model to the Nemo model format. dispey the average accuracy scores across independent runs, well as the 95% confidence interval. The variability across the 4 runs comes from random seeding, which influences the initialization of the Value Model as well as dataset shuffling.
Reward Model
IID
OOD
GSM8K
MATH
algebra222
GSM-Hard
Asdiv
mawps
svamp
DPO Codellama
80.1
0.2
64.8 0.8
59.9 + 0.4
76.9
90.3 +
76,6 0.1
79.3-10.2
43.410.2
6..8
61.110..
77.410.2
91.6.10.3
78.510.9 81.2-004
Q-Function 80.5110.3
45.110.1
7..9
62.7
79,510.5
93.6
Table 6: Average accuracy and 95% confidene interval after the base policy model was trained via DPO, PPO using preference ranking reward model, and PPO using our Q-Function reward model, across 4 different independent runs.
Code execution setup. Our sandbox environment is constructed using Python 3. and includes necessary mathematical libraries such NumPy support the execution of LM-generated code [26]. This environment is containerized and deployed across Kubernetes cluster with 64 replicas, designed to handle the training and evaluation workload in parallel. Each replica is allocated CPU and 1 Gi of memory to efficiently manage the computational demands. Each sandbox establishes communication and accepts requests through Flask server.
Code execution security considerations. our current evaluation setup, the code generated by the Language Model is primarily focused solving mathematical questions and originates from trusted sources, allowing us to maintain a controlled and relatively secure environment. However, extending this methodology to broader applications, where code is executed from varied unknown sources, significantly increases security concers such as malware intrusion and unauthorized data access. To mitigate these risks, it is imperative to create sandbox environment to isolate the code exccution from the regular operational environment. Key protocols include enforcing least privilege in the sandbox limit permissions, minimizing external communications reduce attack surfaces, and rigorously validating both incoming code and outgoing artifacts to ensure they adhere to security best practices.
A.4 Compute and Statistical Significance
For the reward model training and PPO, we used a single node, 8 gpus, and 88 CPU cores, 80 Gi GPU memory and ITB system memory. The total reward model training time ranged between 6-7 hours
19

---

and the total PPO training was approximately 13-14 hours. We additionally used flash attention 2 and deepseed with zero_strrr stagee 3, for reward model training, and zero_ stagee 2 for PPO training [7
Moreover, all the experiment results presented are statistically significant, performed by the scipy library (55) Where applicable, all error bars, whether a table or plot, 95% confidence intervals, calculated via numpy package [26].
B Natural Language Alignment Details
In this appendix, we discuss further experimental details for our results in the Section 4.2
B.1 Training a Natural Language Reward Model
Dataset. To train our natural language reward model baseline and contrastive reward model, we use the training set of Helmluu- [8] Helpful-Harmles is a popular natural language alignment dataset created by Anthropic with -asistats dialog chains. Each pair contains chosen and rejected response, where the chosen response generally aligns on axes of helpfulness and harmlessness. This dataset contains 169,352 rows where each row contains pair of chosen and rejected samples.
Training. For training the reward model, we use the OpenRLHF [29] package and Pytorch. No major code changes were made to the OpenRLHF package, and we use the dataset loader given by OpenRL to load the Helpful-Harmless dataset for training. Table 7 shows the hyperparameters used during training for reproducibility.
hyperparameter and settings
value
batch size
64
learning rate representation dimension
10-4
4096
contrastive lambda (A)
0.5
num epoch preference ranking loss function
1
sigmoid Adam
optimizer seed
0
Table 7: Hyper- parameters for Natural Language Reward Model training
Evaluation. To evaluate the natural language reward model, we use the evaluation script in Open- RLHF (29). We evaluate on the test set of Helpful- Harmlesss. For each chosen, rejected pair, we count the reward model score as correct the chosen generation has a higher score than the reject pair.
B.2 Guided Decoding with Contrastive Reward Model
Dataset. For guided decoding, we consider several different datasets. We considered OpenAssistant (36), a large corpus of assistant dialog and Prism- Alismmmligment (34). a conversation alignment dataset. We also considered HelpSteer [67]. a dataset created by Nvidia with labels on how helpful, coherent, correct, verbose, and complex a given response is. We settled on HelpSteer as our dataset for experimentation due to the relatively high-quality of the dataset (human-collected vs automatic) and the fine-grained category labels. The dataset contains 37, 120 examples in the training set.
Guided decoding, We do not further train the model, but rather randomly sample 20 examples from each of the fine-grained category labels in HelpSteer: coherent, complexity, correctness, and helpfulness. We do not use the verbosity category because the hidden states are averaged at inference time, thus, there is a worry that we would lose the information needed to identify a sequence as verbose. We use SelfFxal Guided Decoding (69 for running the decoding experiments. Equation g shows the objective function I:T") used in guided decoding beam search.
(
II T
(8)
20

---

PLM is the language model distribution, C(s²) is the confidence score, and p is a hyperparameter.
For the steering experiments, temperature to 0, temperature as 0.2, the number of samples to 2, and the confidence ratio to e.g. equal weight to the confidence score and LM loss.
the number of tokens per beam 52, the beam search
Evaluation. We evaluate on the test set of HelpSteer, which contains 1790 examples. For evaluation, we use GPT-4, asking GPT-4 to compare the SFT model and the model with the decoding prototype which generation is more helpful, correct, coherent, and complex. We ask GPT-4 to give a single token (0 or 1). For statistical significance, the majority. In addition, we find that for a small percent of cases GPT-4 will not product a single token even when asked. In this case, throw away the example that did not produce properly formatted response when computing averages.
run GPT-4 times for each evaluation example and take
B.3 Compute and Statistical Significance
Natural language alignment training jobs were done on I core, GPUs, 88 CPU cores, 1TB system memory, and 80GB GPU memory. An average training time for episode of each prompt around 13-14 hours.
All experiments are statistically significant. For the Helpful-Harmless experiments, the experiments are evaluated on a test set 8.55k rows. For HelpSteer, the validation set is 1.79k rows.
C Inference-Time Goal State Ablation
In this appendix section, we explore different choices of inference-time goal states and compare their performance. To ablate this, we conduct new experiment under the mathematical reasoning with code setting (Section 4.1) 1. where goal state and (b) the number of examples used to construct the goal state. particular, we evaluate 3 settings for picking the sample completions for constructing the goal state:
evaluate the effect of (a) choosing poor examples to construct the
1. We vary the number of preferred completions used to construct the goal state evaluate the effect adding more examples of good reasoning.
2. We vary the number of dispreferred completions used to construct the goal state, to evaluate whether bad examples lead to poor performance.
3. We incrementally add more dispreferred completions to fixed sample of preferred comple- tion in order to measure the robustness of our goal state computation negative examples. We refer to the addition of dispreferred completions on top of all preferred completions as adding "coorupted"" examples.
The results are in Figure 5 Overall, these results show that negative or unhelpful completions degrade performance, indicating the importance of choosing relevant examples for computing the goal state. Additionally, they demonstrate that having more examples of the concept leads to better performance more generations filtered with comparable proportion of those remaining generations being correct.
We additionally ablate our choice of using the last token in the completion sequence to construct the goal state. In this experiment, we repeat our filtering experiments from Section 4.1.3] Except, we randomly sample token from in the completion sequence and use it as the goal state:
Sampling Method
GSMSk
MATH
Accuracy (%)
Prop. Correct (%)
Accuracy (%)
Prop. Correct (%)
Last Token
86.0
84.0
59.6
52.0
Random Token
85.9
81.8
57.3
45.3
Table Majority vote accuracy and average proportion correct when filtering examples using last
token and random token sampling for constructing the goal state.
21

---

24.13
05.5
2e0 Pep
10
10
353
10.64
13.00
OCspreferees Exampies
Preferies Examples
Commateed Exaripe
D.sarrefere an Examples
Preferied Examples
Comusted Exarip
(a) MATH
(b) GSM8K
Figure 5: We plot the average proportion of non-filtered responses that correct and the average number of completions (out total 50) that are filtered, as we vary the sample size (1, 10, 100, 1K, 10K, and 30K) of dis-preferred, preferred, and corrupted examples when constructing the goal state. All of the examples are retrieved from the preference ranking training dataset. We plot the performance metrics for the MATH and GSM8K test sets. Each of the test sets has 50 completions per problem that is generated by the base model and we filter any completion for which the Q-value for any token in the completion sequence is below O. At each sample size, run 3 independent random sample from the full set of preferred and dis-preferred completions to construct independent goal states, with 95% CI.
From the result of this ablation, see that last token sampling has superior performance, likely be- cause the last token on average carries the most information about what good complete generations look like.
D Train-Time Goal State Ablation
D.1 Exploring Sources of Goal States
natural question about our approach is: What the impact on the choice of goal state represen- tation? For sampling negative goal states, we experiment with randomly sampling any arbitrary completion (not necessasarily the same prompt) vs. the corresponding dispreferred completion to show the effect of picking completions that have similar contexts and invalid reasoning steps when constructing the contrastive loss. Furthermore, we experiment with taking the average of the hidden states of the positive and negative goal states across the train batch and then apply cosine similarity with the current token representation. The results of these experiments are in Table
Method
IID
OOD
GSM8k
MATH
algehra222
GSM-Hard
Asdiv
mawps
svamp
Pref. Rank
0.805
0.618
0.850
0.768
0.829
0.775
0.761
Q O-Function (SGS)
0.852
0.708
0.879
0.781
0.854
0.819
0.841
O-Function (RS)
0.801
0.626
0.832
0.716
0.723
0.760
0.802
O-Function (AVG)
0.835
0.707
0.851
0.765
0.839
0.773
0.863
Table 9: Comparison of AUROC scores for different training time contrastive goal states on several popular math benchmarks. Pref. rank is the reward model trained with the standard preference-ranking objective, Q-Function (SGS) is trained with the method described in Section 3.2 Q-Function (RS) trained with the randomly sampled goal-state and Q-Function (AVG) trained with the batch averaged goal state.
The ablation in Table shows that the positive and negative goal state representations we use are largely more effective than randomly sampled and batch averaged goal representations. Somewhat similar observations have been made for other contrastive learning approaches where negatives, especially hard-nege negatiess,,,,, can play an important role 114 57]. see that when we randomly sample positive and negative goal states, there is clear regression in performance on the OOD datasets,
22

---

even below the standard preference ranking trained reward model. Meanwhile, the batch averaged goal states offers competitive but largely lower performance compared our final setup. Given these results, we use the SGS contrastive loss method as our final setup.
D.2 Source State and Goal State Sampling Ablation
In this section, we ablate the sampling strategy for the source and goal states. In computing the contrastive loss during training. we need to sample positive source and goal state pair as well as negative source and goal state pair. To investigate how we should sample these tokens, we explore different approaches.
Method 1: Random Source and Goal State. In this approach, the positive and negative source states and sampled randomly from the preferred response, then the positive goal state is sampled randomly amongst any token after the positive source state, and last the negative goal state is randomly sampled anywhere from the dis-preferred response. This is the method we describe in Section 3.2
Method 2: Late Goal State. In this approach, we sample the positive and negative goals states only towards the end of the preferred and dis-preferred response, respectively. The intuition here that the token later in the sequence encode more information about the whole response and thus provide more information as a goal state. In practice, we restrict the goal state sampling to come after the 90th percentile with respect to the total number of tokens in the response.
Method 3: Late Source and Goal State. In this approach, we sample both the source state and goal state towards the end of the preferred and dis-preferred response (after 90th percentile of tokens generated).
Sampling Method
ITD
O0D
GSMSK
algebra222
GSM
Asdiv
mawps
svamp
Random Source and Goal State
0.852
0.708
0.879
0.781
0.854
0.819
0.841
Late Goal State
0. 0.841
0.696
0.821
0.715
0,860
0.729
0.761
Late Source and Goal State
0.829
0.643
0.867
0,779
0.829
0,831
0.833
Table 10: Comparison of AUROC scores for different source and goal state sampling method.
The ablation in Table shows that the Random Source and Goal State sampling on average performs the best across both in-distribution and out-of-distribution tasks. We also see that with Late Goal State sampling, we have similar performance to Random Source and Goal State sampling for the in-distribution tasks but we notice a performance drop across several out-of-distribution benchmarks. The opposite is true for the Late Source and Goal State sampling where performance on the in- distribution benchmark drop and the performance out-of-distribution benchmarks are comparable to the random sampling. Given these results, we use the Random Source and Goal State sampling method as our final setup, which is equivalent to the (SGS) method referred to in the rest of the paper.
E Reward Model Ranking Performance
In this appendix section, we compare the Q-Function 7b Reward to Codellama 7b Reward by judging their respective abilities to rank completions. Since we have a total of 50 generations per problem. we use the baseline reward mode and our reward model to select the Top 1, 5, 10, and 25 samples as ranked by the reward scores. With the selected sample, we perform majority vote and also note the average proportion of the sample K that correct solutions. The results are provided in Table II] We show the Majority Vote Accuracy and Average Proportion Correct when selecting Top- completions in order of reward model scores. The two reward models used here are the baseline Codellama RM and our Q-Function RM. The results show that our Q-Function RM clearly outperforms the baseline.
23

---

GSM8K
MATH
Model
Top-K
Accuracy (%)
Prop. Correct (%)
Accuracy (%)
Prop. Correct (%)
1
84.6
84,6
51.7
51.7
Q- Q-Function RM
5
86.2
84.5
59.7
51.1
10
86.2
83.8
59.5
50.3
25
85.5
81.8
57.8
47.0
80.8
80.8
43.8
43.8
Codellama RM
5
84.3
81.9
54.0
45.8
10
85.2
82.1
56.0
46,9
25
85.8
81.1
56.5
46.1
Table 11: Majority Vote Accuracy and Average Proportion Correct when selecting Top-K completions in order of reward model scores. The two reward model use here is the baseline Codellama RM and Q-Function RMM.
F Partial Completion Rewards
In this appendix, we show the plots for the remaining out-of-distributin datasets, namely, algebra222, Asdiv, GSM-Hard, mawps, and svamp. The general trend observed in Section[4.1.2) still holds, where, Q-Function 7b Reward model incrementally improves as percentile of model generation considered increases. On the other hand, there is a lot more variability Codellama Reward partial auroc scores. More importantly, the performance when considering partial solutions is generally more accurate for the Q-Function 7b Reward. This demonstrates that the improved judgement of partial solutions that observed for the in-distribution tasks extend to out- diitr---istiibuionnn tasks as well.
016
0.00
O-Function
AUROE Score
AUROC
Percentle Model Generation Considered
Percentlle of Model Generator Considered
manps
Function Reward
Reward
AUROC Score
SUROO Score
Percentil Model Generation Considered
Percentiee Model Generation Considered
guame
Codeliama QuPunction
3.70
Figure 6: AUROC scores on the rewards attributed to partial base-model greedy generations on the OOD datasets.
24

---

G Q-Value Visualization
In this appendix, display the full Q-value heat map for the examples shown Figure We also show additional examples that show the Q-values correctly scoring reasoning errors as well some examples where the Q-values are noisy, unable to explicitly pin-point the logical error.
G.1 Good Examples
Problem: Toula went the bakery and bought various types of pastries. She bought 3 dozen donuts which cost 68 per dozen, 2 dozen mini cupcakes which cost 80 per dozen, and dozen mini cheesecakes for 55 per dozen. How much the total cost?
Let
salve
this
problem
using
Python
code
=|
m code
don uts
per
do
don
price
per
uts
cost
1its
per
11S
price
per
zen
cup
alees
per
Jo
cup
akos
price
per
do
O
akees
nikes
peT
do
alkes
price
pCr
do
che
akees
do
price
per
per
do
pen
price
per
total
cost
don
uts
cost
cup
alees
cost
cbe
akees
cosE
total
cOst
</
11
m
cooee
coode
suulpul
m
code
outpat
1h
us
the
cotal
cost
of
past
rics
box
dollars
Problem: Marissa is hiking 12-mile trail. She took hour to walk the first 4 miles, then another hour to walk the next two miles. If she wants her average speed to be miles hour, what speed (in miles per hour) does she need to walk the remaining distance?
Let
salve this problem
using Python code
<|n1 conee
distance
tra ve
Ted
time
aken
speed
nining
distance
distancc distance
lime
sining
renninnng
average
speed
time
rem
nining
code
conde
batgrul
m
code
outpur
Th
us
Mar
issa
neetss
walk
IE
box
miles
per
hour
In
achieve
her
average
speed
milles
per
hour
25

---

Problem: dance class of students, 20% enrolled in contemporary dance, 25% of the remaining enrolled in jazz dance, and the rest enrolled in hip-hop dance. What percentage of the entire students enrolled hip-hop dance?
Let
solve
this
problem
usingg
Pythen
conbe
=12
rolled in conterporary
dance
cont empor Jryy [_d
ance peroent
age
= | of the remaining on solled
in
dance
ilazz
ance peecentt
1
the Irolicd [hip
|-1
dance
hip
1-|n|op
Jncc
percent
age
|
cont
empor
Jary
ancee
percent
age
juzzZ
J
ance
percent
sotal
number
[of
studdnts
total
stud
1-2100
number
srudents
cn rolled
hip
dance
hip
ance
slud ents
lotal
stud
emils
hip
up
sice
perceni
Agg
percentage
Lof
studenis
cm rolled in hip
dance
hip
op
ance
pereent
of
total
hip
ance
stud
ents
total
stud
ents
Q
op
ance
pereent
age
<3
code
m
code
outtput
1
|
<
code
ousput
us
students
rolled
hip
op
Problem: When Freda cooks canned tomatoes into sauce, they lose half their volume. Each 16 ounce can tomatoes that she uses contains three tomatoes. Freda's last batch of tomato sauce made 32 ounces of sauce. How many tomatoes did Freda use?
Let solve this
problem
using
Python
codc
code
cach
can
tom SEOCS
has
tom
aloes
tom
atoes
per
[each]
can
of
bonn
atoes
has
Oun
ces
oun
ces
per
Fred
made
CNETI
Des
cc
made
Fred
uscd
M
sau
cc
made
own
CCS
per
can
Fred
used
tom
stoess
per
CILTI
12
code
<
11
code
output
<
code
output
Fred
box
26

---

G.2
Noisy Examples
Problem:
Lee used to be able to run the 400-meter hurdles two seconds faster than Gerald would run the 400-meter hurdles, But Gerald changed his diet, which improved his speed by 10%. If Lee the 400-meter hurdles in 38 seconds, how fast can Gerald, with his improved diet, run the 400-meter hurdles, in seconds?
Let
solve
this
problem
using
Python
code
|code
Icc
time speed
|-13
le
ger
ald
speed
Lc
speed
speed
Ber
ald
speed
ger
old
time
SeT
uld
speed
with
im prov einent
ger
time
m
code
m
code
outpul
5
<2
pode
oulpul
us
Ger
ald
cam
run
the
meter
hUT
L0%
11
box
ed
seconds
John drives for 3 hours a speed of 60 mph and then turns around because realizes he forgot something very important at home. He tries to get home in 4 hours but spends the first 2 hours in standstill traffic. He spends the next half-hour driving at speed of 30mph, before being able to drive the remaining time of the 4 hours going at 80 mph. How far is he from home at the end of those
Problem:
4 hours?
Let
solve
this
problem
using
Python
codc
code
>
John
drives
hours
speed
ph
distance
tra
vc
led
first
part
6
then
Eorgot
something very important
distance
tra ve
second
part the First hours stand
tries
1o
get home
hours
but cnds
ill
traffiic
distance
tra
vc
1ed
second
part 0
ends
the
next
balf
hour
driving
speed of O ph
distance
tra
second
part
before boing
abll
10
drive
the
remaining
tumec of the hours
going
at
0
ph
distance
tra
led
second
purt
total distance
1ra
vc
led
total
distance
led
distance
led first distance
lod
sccond
total
distance
112
VC
<3
bode
m
outpul
m
code
output
Th
John
box
ed
miles
from
home the cad of
those
hours
27

---

H Exploring Correlation Between Q-values and Reward Scores
In this section, we explore the correlation between our learned Q-values and reward scores for partial completions, as illustrated in Table 12 and empirically investigate the possibility of reward hacking during PPO training. The partial completions are from the base model on the GSM8K test split [[6]. The correlation is expected since, by design, the reward model outputs are dependent on the learned features that approximate the Q-function. This presents a potential risk in practice that the policy model may exploit during PPO. We explore attempting to decouple the reward score from the Q-value for partial completions by projecting the hidden features which approximate the Q-value through multilayer perceptron. Table 12 and Table 13 compares the result of projecting the hidden features via a MLP layer versus a linear layer. The finding is we can reduce correlation between rewards and Q-values for partial completions while maintaining high reward model performance. Iterative retraining of the reward model on on-policy completions would further mitigate any potential reward model hacking given that our complete loss function also considers the preference ranking paired loss in
Method
IID
OOD
GSM8k
MATH
algebra222
GSM-Hard
Asdiv
mawps
svamp
Pref. Rank
0.805
0.618
0.850
0.768
0.829
0.775
0.761
Q-Function (Linear)
0.852
0.708
0.879
0.781
0.854
0.819
0.841
Q-Function (MLP)
0.832
0.693
0.871
0.770
0.871
0.780
0.841
Table 12: Comparison of AUROC scores for different reward projection strategy several popular math benchmarks. Both Q-Function (Linear) and Q-Function (MIP) are trained with the method
described in Section 3.2] the only difference being the reward head is a single linear layer versus 3-layer MLP followed by a linear layer.
Model
Percentile
0.2
0.4
0.6
0.8
Q-Function (Linear)
0.983
0.986
0.988
0.988
Q-Function (MLP)
0.870
0.908
0.921
0.915
Table 13:
Correlation between Q-value and reward score different percentile of model generations. We see that the Q-value and reward scores for Q-Function (Linear) has higher correlation than Q-Function (MIP).
To empirically evaluate the risk of the policy gaming the reward function, conduct exper- iment where we run PPO training for an extended time (5 episodes) order to evaluate whether reward hacking occurs in practice. The reward scores and response lengths shown in Figure 7 and the performance of the model at the end of each episode given in Table 14
train//eppree, Jingth
233
200
Figure 7: Train time average rewards and average response length for PPO using Codellama RM (purple) and PPO using our Q-Function RM (brown)
These results indicate training is fairly stable and the policy model does not game the reward function in practice. One explanation for why gaming does not occur is that typically during RLHF a relatively low learning rate is used, making large shifts in the models generations less likely to occur.
28

---

Model
Episode
GSM8k
MATH
algebra222
GSM-Hard
Asdiv
mawps
svamp
79.7
43.6
66.3
61.2
77.9
92.0
78.5
79.6
45.0
72.5
61.5
78.0
91.6
79.9
Codellama
80.5
45.3
69.8
62.0
78.3
92.3
80.9
80.5
45.5
62.5
78.3
92.0
79.5
80.5
45.2
70.7
62.5
78.6
92.0
80.5
80.2
45.9
67.6
62.0
79.2
81.6
80.6
46.5
74.3
61.9
79.5
93.1
80.3
Q-Function 81.7
46.6
73,0
63.5
79.2
93.2
81.6
81.0
46.5
73.4
63.2
79.2
93.4
81.0
81.1
46.5
74.3
63.0
79.8
93.6
81.0
Table 14: PPO-trained policy model accuracy across several math benchmarks comparing using baseline Codellama RM and our Q-Function RM for single long 5 episode training run.
29

---

I Examples of Prototype Guided Decoding
Question: Consider this reference information delimited in """:
There is man who circles the perimeter with a baby in his arms unmoving. Locusts burn with the silhouettes of saints at dusk. Saints are the cloud. We are in a dry storm. The man extends his circles pulling the baby through the cactus scrub. Look at his melting trainers in the heat, they arent what he asked for. There are black leather skids on the dry stone wall. People in black cloaks of your eye. A pig turns on a spit. The prairie is a terrarium for the blaze but the edge is dry of fire. It is the height of one season, bushes burn. A burnt five-year-on without eyelids turns quick cartwheels through the heat wave under the big pale sky, black and blue.
out of the corner
What type of poem is this text from?
SFT: This text is prose poem.
Prototype: This text is prose poem, which is a type of poetry that combines the elements of prose with those of poetry to create a unique and imaginative style of writing.
The guided decoding successfully produces response that is still correct while adding complexity to the response by describing that the poetry combines elements of prose and imaginative writing styles. We find that, in general, guided decoding responses produce more complex responses without sacrificing correctness.
Question: Reference: <start of reference> If you invested 1,000 in Facebook in 2012, here's how much you'd d have now Even given the stock market's 's recent volatility, a 1,000 investment in Facebook 10 years ago would still prove to be good bet: You could have quadrupled your money. An initial investment in Facebook in 2012, when it made initial public offering, would be worth slightly more than 4,600 of Wednesday, according to CNBC calculations. The investment would have seen lifetime total return of more than 360 percent, including price appreciation and dividend gains reinvested. The market's recent activity has taken toll, of course: In one month, Facebook's 's stock dipped from 5,051. Its current portfolio value is just above 30 percent now, down from its peak of more than 40 percent. And the market keeps ebbing and flowing. In the charts below, all data splits are adjusted and gain-loss figures do not include dividends, interest, distributions or fees except on cash accounts. The portfolio value represents current holdings and the comparison charts represent current and historical prices of individual benchmarks, stocks or exchange-traded traded funds. Mcanwhile, daily users in the United States and Canada dropped in 2017 184 million daily average users in the fourth quarter from 185 million in the third quarter. Facebook chief financial officer David Wehner, on the company's ' earnings call, said usage would "fluctuate" based on market penetration. The revenue increase despite the usage drop could be due to seasonal changes related to the advertising market. As the company points out, users the fourth quarter accounted for 6.39 billion of Facebook's 's revenue, compared to the 5.03 billion users accounted for in the third quarter. Facebook has retreated slightly from its record price reached earlier this month as investors debate whether changes to the platform to combat fake news and other problems will hurt advertising growth. But the stock is still up percent on the year. And some investors remain bullish about the stock. "News feed changes the power of the platform." Andy Hargreaves, analyst KeyBanc Capital Markets, wrote in note to clients. "Facebook offers reach, targeting and ad-unit quality that is difficult to match, and they are the underlying drivers of its growth. As long as [it] continues to improve its effectiveness, the foundation for strong growth should remain intact, regardless of adjustments to the news feed." <end of reference> Instruction: What is share of stock?
not impair
SFT: A share of stock is unit of ownership in company.
30

---

Prototype: A share of stock is unit of ownership in company that represents the right to proportionate share the company's assets and profits.
In this response, the prototype response is slightly more correct as it describes the value of a share of stock. In addition, the prototype response more complex than the SFT response.
Question:
What is the difference between pound and an ounce? Please explain with examples.
SFT: A pound is unit measurement for weight, while an ounce is unit measurement for volume.
Prototype: A pound is equal to 0.45359237 kilograms, while an ounce is equal to 28.3495 grams. A: A pound unit of weight in the metric system, while ounce is unit of weight in the English system.
Here, find that the SFT response is incorrect. The prototype response gives the correct response, but produces some additional information at the beginning of the response that does not answer the question. This can be seen as side effect of the guided decoding process since more complex response not always desired.
Question: What are some common misunderstandings about the causes and consequences of the 2008 global financial crisis?
SFT: One is that it was caused by lack of regulation in the financial sector, when in reality, the crisis was caused by a combination of factors, including excessive risk-taking by financial institutions, poor risk management practices, and lack A: There are few common misunderstandinns about the causes and consequences of the 2008 global financial crisis.
Some believe that the crisis was caused by lack of regulation or irresponsible lending practices. Some also believe that the crisis was caused by a lack of regulation or irresponsible lending practices. In reality, the crisis was the result of complex and interconnected events that took place over many years. There are many misconceptions about the 2008 financial crisis, including the belief that it was caused by lack of regulation or irresponsible lending practices.
Prototype:
The prototype response is more coherent than the SFT response in this example. In the SFT response, there lack of coherence because the response separated into two distinct sections, seperated by "A:". However, this issue does not exist in the prototype response. Each sentence builds off the last, for example "Some believe..." followed by "in-reality".
31

---

Question: Reference: <start of reference> As a casino player, it would be amazing to predict when a slot machine will next hit. You could time exactly when you play and how much to bet perfectly, maximising your winnings. The problem? If this was widely possible, no casino would exist. Theyd all go bankrup straight away. Instead, modern slot machines are designed to give determined results, that are not generated at random. It is incredibly difficult to guess what the result of spin will be. There may be slot strategis that might improv your chances of winning. But any legitimate methods won't 't' include being able to know when a slot is about to drop. Yet there's 's still lot of talk about this online, with "tips" and "triks" on how to predict when a slot machine will next payout. We'l IIl explain what really happens, along with how you can give yourself slightly better chance of winning legitimatly. You cannot tell when a modrn slot machine will hit because the outcome of each spin is random. No matter how many times machine has spun, and no matter what the outcomes of those spins were, the probability the next result remains the same. But people often believe otherwise. One famous example of why the gambler's 's fallacy. The gambler's fallacy is an erroneous way of thinking about the probability of independent events. When people commit the gambler's fallacy, they think that the likelihood of random event is impacted by the occurrence of previous events, or series of previous events. For example, believing that the result of previous spins on slot will impact the result of present or future spins is committing the gambler's s fallacy. The result is not affected by what's 's happened before. So if you win or lose one round, it has no impact on whether you'l 'Il win or lose the next round. <end of reference> Instruction: Rewrite the passage correcting any grammar and spelling mistakes that have been made. Also omit any false information that the passage may contain. Make sure to keep the same tone and message of the original passage.
SFT: Casino players believe they can predict when slot machine will next hit, but this is not possible.
You could time exactly when you play and how much to bet perfectly, maximising your winnings. The problem? If this was widely possible, no casino would exist. They'd 'd all go bankrup straight away. Instead, modern slot machines are designed to give determined results, that are not generated at random. It is incredibly difficult to guess what the result of a spin will be. There may be slot strategis that might improv your chances of winning. As a casino player, it would be amazing to predict when slot machine will next hit.
Prototype:
The prototype response gives more detail including additional reasoning: "modern slot machines are designed to give determined results". Interestingly, this response is generally incorrect, as modern slot machines will give random results that have deterministic rules (for example, no jackpot right after another jackpot). This shows that, while the prototype can help improve correctness, it can still suffer from hallucination.
Assets
In Table 15 we summarize assets used in this paper along with their corresponding URLs and licenses.
Asset Name
URL
License Type
Lhma
httpp:/// //g.th/bbbb/tt/////////./c con/nets Jl Lans]
Lluma Community License
OpenMaih Codel Jama
3+
hotpp: httpee///ggggtteee Saggggggtac ace, da wwwpvidia/ Upwnaats Lodel.laa. 7b- Pyshon
Lama Community License
OpenRLUF
20
bespo: / coe/ Oow/ /OpcaBLiE
Apache 2.0 Apuche 2.0
NeMo-Skill
5H
Stops:// gitbnb com/ con/tsspk//////// NeMo Se1l1s
OpenManhinstuuc-1
02
bttps: http///ggaaaaccccco/e co/datasets/ /ed CowwMaatiaattrucc-
NVIDIA License
GSMBK
[6
atage // inscinngece co/ dat atett Reo02
MIT License
MATII
23
bripe: //ithuu
MIT License
algehra222
httpp: Busy
Apache-2.0
GSM Hand
h:tpc higgeiigt co catasess/ree ooine Tachi AGE hard
MIT License
Asdiv 55
beepe (// 2ait com/ chancoum ii 3Pa1y ostaret
CC BY-NC 4.0
mawps 33
License
svamp
Attps:/ /hii:///iieeeaae corderacete 1h5 1h
License
Helptalness pad Harmlessness X
hc bc4p:: // /1gg are c0/ coneta An ooop las ria oatarett avaata/ Be Botget
MIT License
HelpSteer
x
https://hugg
CC CC-BY-4.
Table 15: Asset information and licences.
32

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We propose new method for training reward models via a contrastive objective. In Section we deseribe our method in detail. We also claim that such method helps reward model improve their performance and that the learned representation is useful signal for approximating future gains. In the experiments section, we provide two different settings, mathematical reasoning with python interpreter (Section 4.1 and natural language alignment (Section 4.2). where we explicitly state the experimental setup and results referred to in the abstract.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We meaningfully explore two limitations to our approach in the limitations section (Section Namely, we call out the discrepency of train time vs. test time settings for our method and that there correlation between the reward and Q-value signals. We also provide/explore some approaches to reconcile these limitations for further generalizability of our method.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they seale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
33

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: This paper does not present theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: For our proposed method, we provide detailed breakdown of the theoretical ideas as well as how implement it in practice. Moreover, in Table and Table we lay out all the relevant hyperparmenters required to reproduce the results. Also, we list out the minimal modification we make the OpenRLHF package and so prospective readers should be able to implement the exact methods and settings with relative ease.
Guidelines:
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 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 specific model and empirical evaluation, it may be necessary to either make 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., 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
34

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material? [Yes]
Justification:
The code is released
https:///gttubb com cttm:///aaaaaaattaaaaaaallaoooaa
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc.
Sew.adde//CoeSSmmsssssioooo los.suubiissooooocccccoooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https 7/nips. wwtp:ccccppbpuuueeueeeeee
SeroodeSumisssioooo soiobiiiiiiiooiiiiyyccoooo for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: For each experiment, explicitly mention how the training/evaluation datasets were constructed, the magnitude of the datasets, the base models used, as well as the experimental setup/methods use for all reported values.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
35

---

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: All the experiment results presented are statistically significant, performed by the scipy library [65]] Where applicable, all error bars, whether table or plot, are 95% % confidence intervals, calculated via numpy package [26].
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, confi- dence intervals, statistical significance tests, 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.
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 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 out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: In the Experimental Details section in the appendix, we explicitly state the compute details required for each training and evaluation job and approximately how long each run took.
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 compute required for each of the individual experimental runs well 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 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
Answer: [Yes]
Justification: The models and data used in our experiments are open source and used in accordance to their respective licenses.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
36

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We introduce new method for reward model training and show the validity of this approach on self-contained experiments that will not have outreaching effects on society.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (c.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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We are not releasing new dataset nor new model.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
37

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
The assets used in this work along with their respective licenses provided in Appendix []] this work, we exclusively use them for research purposes and in accordance with their license and usage agreements.
Justification:
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
This submission is currently not releasing new assets due to time and security constraints.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This work does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
38

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This work does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere 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.
39


Reviews:
[{"review_id": "VY2OZTjqYE", "reviewer": "Reviewer_U1X3", "summary": "This paper presents a method of applying goal-conditioned Q-functions to learn representations via contrastive learning to capture the expected reward. By incorporating an auxiliary contrastive loss for training the reward model, the performance of language model alignment obtains improvement. Experiments on GSM8k and MATH further validate the superior performance of the proposed method.", "questions": "- How is the performance of the proposed method compared with DPO? Since DPO is a well-known LLM alignment method, I encourage the authors to add a comparison and discussion between them.", "limitations": "Limitations have been discussed in this paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This paper adapts the goal-conditioned representation learning in RL to help boost the performance of the language-based reward model for LLM alignment, which is novel in the LLM area.\n- A thorough set of experiments demonstrates the superior performance of the proposed method.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to the reviewer for their thoughtful review and insightful suggestions.\nWe are pleased to know they acknowledge the novelty of the goal-conditioned approach with LLMs and found our experiments to be thorough.\nThe reviewer brings up the interesting point of comparing our proposed method with DPO. \nWe compare the performance of DPO with our proposed method on the mathematical reasoning with code execution tasks (Section 4.1). \nFurthermore, we intend to incorporate this response and discussion into the final version of the paper.\n\nIn particular, we use the GSM8K + MATH preference ranking dataset that we used for training both the baseline and contrastive reward models as the DPO training dataset. \nWe compare performance of DPO with PPO training using the baseline reward model and our contrastive reward model.\nWe present average accuracy across 4 independent runs for PPO and 2 independent runs for DPO. \nThe base model results are also shown as a reference point presented in [6].\n\n| Model           | GSM8k             | MATH              | algebra222        | GSM-Hard          | Asdiv             | mawps             | svamp             |\n|-----------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|\n| Base            | 75.9              | 43.6              | 65.6              | 60.1              | 77.7              | 93.5              | 79.6              |\n| DPO             | 80.1 ± 0.1        | 44.8 ± 0.2        | 64.8 ± 0.8        | 59.9 ± 0.4        | 76.9 ± 0.0        | 90.3 ± 0.2        | 76.6 ± 0.1        |\n| Codellama PPO   | 79.3 ± 0.2        | 43.4 ± 0.2        | 65.8 ± 1.6        | 61.1 ± 0.3        | 77.4 ± 0.2        | 91.6 ± 0.3        | 78.5 ± 0.9        |\n| Q-Function PPO  | **80.5 ± 0.3**    | **45.1 ± 0.1**    | **70.9 ± 1.7**    | **62.7 ± 0.5**    | **79.5 ± 0.5**    | **93.6 ± 0.3**    | **81.2 ± 0.4**    |\n\nInterestingly, we observe that DPO performs on par with our method for the In-Distribution tasks (GSM8k and MATH). However, the policy trained with the contrastive RM performs much better on OOD tasks (algebra222, GSM-Hard, Asdiv, mawps, and svamp), compared to DPO.\nThese findings could be explained by related works that analyze DPO and find it can perform poorly in OOD settings [1]. All in all, these results indicate that PPO training with the contrastive RM leads to better generalization than training with DPO, particularly in OOD settings.\n\n[1] Xu, Shusheng et al. “Is DPO Superior to PPO for LLM Alignment? A Comprehensive Study.” ArXiv abs/2404.10719 (2024).\n\n[6] Toshniwal, Shubham et al. “OpenMathInstruct-1: A 1.8 Million Math Instruction Tuning Dataset.” ArXiv abs/2402.10176 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hbWbnp9qSb", "reviewer": "Reviewer_twr2", "summary": "This work combines contrastive representation learning, goal-conditioned RL, and reward models used in RLHF for language model alignment. The authors introduce a new method that uses a contrastive loss to encourage the reward model to learn what they define as \"goal-conditioned representations.\" These representations essentially encode the expected reward at different steps of the generated text, supposedly helping the model predict which generations lead to good (or bad) outcomes. They test this out on math problem-solving tasks and a helpfulness/harmlessness dataset. They find some improvements in reward model accuracy and show that these representations can be used for things like filtering out bad solutions and steering the language model during generation.", "questions": "- Have you considered investigating approaches such as Upside Down RL [1, 2] given that conditioning in this case is not on a goal state but the desired rating?\n\n1. https://arxiv.org/abs/1912.02875\n2. https://arxiv.org/abs/2106.01345", "limitations": "The authors have adequately addressed the limitations of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Originality: The paper creatively combines ideas from goal-conditioned RL techniques to improving reward models for LM alignment. This results in a novel method for training reward models which further improves expressiveness of the reward model via the additional loss. \n- Significance: In addition to improving the overall reward models, the authors obtain an additional quantity that measures the $Q$ values of a given state (or more precisely state-action-goal), which enable a wide set of downstream use-cases, such as early misalignment detection. While these are preliminary results, they demonstrate exciting potential for further exploration.", "weaknesses": "- Correlation between Q-values and reward scores: The authors acknowledge the high correlation between Q-values and reward scores for partial sequences, raising concerns about the policy model potentially gaming the reward model during RLHF. While they suggest further exploration of decoupling these signals, this issue deserves more thorough analysis.\n- Clarity: The paper lacks clarity in multiple areas. In particular, the functional form of the $Q$ function is not entirely clear to me. Choosing the goal states is another source of obscurity. While it is clear that during training, positive and negative examples are chosen from the respective completions, in my opinion the more important part, which is the choice of goal states during inference is not entirely clear to me. The authors briefly mention that they \"we ake the mean representation of multiple goal state completions from the training data\", but given that this is an important part of the work, I believe it requires a lot more analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for useful comments and for appreciating the originality and performance improvements of the method.\nWe address the reviewers questions and comments in the following sections.\nAdditionally, we plan on incorporating all our responses and discussions into the final version of the paper.\n\n**Correlation between Q-values and reward scores** Thank you for your comment.\nTo evaluate the risk of the policy gaming the reward function, we conduct a new experiment where we run PPO training for an extended time (5 episodes) in order to evaluate whether reward hacking occurs in practice. \nThe reward scores and response lengths are shown in Figure 2 within the additional supplementary\nmaterial page and the performance of the model at the end of each episode is given here.\n\n| Model        | Episode | GSM8k | MATH | algebra222 | GSM-Hard | Asdiv | mawps | svamp |\n|--------------|---------|-------|------|------------|----------|-------|-------|-------|\n| Codellama    | 1       | 79.7  | 43.6 | 66.3       | 61.2     | 77.9  | 92.0  | 78.5  |\n|              | 2       | 79.6  | 45.0 | 72.5       | 61.5     | 78.0  | 91.6  | 79.9  |\n|              | 3       | 80.5  | 45.3 | 69.8       | 62.0     | 78.3  | 92.3  | 80.9  |\n|              | 4       | 80.5  | 45.5 | 70.7       | 62.5     | 78.3  | 92.0  | 79.5  |\n|              | 5       | 80.5  | 45.2 | 70.7       | 62.5     | 78.6  | 92.0  | 80.5  |\n| Q-Function   | 1       | 80.2  | 45.9 | 67.6       | 62.0     | 79.2  | 93.9  | 81.6  |\n|              | 2       | 80.6  | 46.5 | 74.3       | 61.9     | 79.5  | 93.1  | 80.3  |\n|              | 3       | 81.7  | 46.6 | 73.0       | 63.5     | 79.2  | 93.2  | 81.6  |\n|              | 4       | 81.0  | 46.5 | 73.4       | 63.2     | 79.2  | 93.4  | 81.0  |\n|              | 5       | 81.1  | 46.5 | 74.3       | 63.0     | 79.8  | 93.6  | 81.0  |\n\nThese results indicate training is fairly stable and the policy model does not game the reward function in practice.\nOne explanation for why gaming does not occur is that typically during RLHF a relatively low learning rate is used, making large shifts in the models generations less likely to occur.\n\n**\"..the functional form of the Q function..\"** The Q function is parameterized by a scoring function $f$ and encoder.\nIn this paper, we use the cosine similarlity as $f$ and the encoder is a causal LM, such as Llama.\nIn particular, to compute the Q-value for a goal-state and state-action pair, we embed both the goal state and the hidden state for the state-action pair, and take the cosine similarity between the two representations.\n \n\n**\"..choice of goal states during inference\"** We refer the reviewer to the shared response.\n\n**\"..investigating approaches such as Upside Down RL..\"** Our work focuses on improving the representations learned by reward models (RMs) for aligning LMs. \nOur improved methods of training RMs leads to improvements in RM performance, downstream policy learning, and model steerability.\nApproaches such as upside down RL do not focus on improving RM performance to better align LMs, so we did not consider investigating them in this work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G32yp7oPhX", "reviewer": "Reviewer_2p1n", "summary": "The paper frames the reward learning problem for LLM as a goal-conditioned RL. It uses the contrastive learning loss from Eysenbach et al. 2022 as an additional objective for reward training. The main innovation is the adjustment of goal-conditioned RL to pairwise preference datasets. The paper shows that adding that contrastive loss on the hidden representation of the reward model can lead to a better reward model and, consequently, better policies.", "questions": "* Regarding the evaluation of the reward model using the AUROC metric, can you elaborate on how you calculated it? Do you use the BT model output as classifier prediction? Where do the GT annotations come from? I looked at the references you provided (line 206), but it doesn't seem like AUROC was used there.\n* In the experiment described in section 4.1.3, was the filtration done using the Q-value of the full answer or the partial one? In addition, this experiment is missing a baseline of best-of-50 using the vanilla reward model and your own reward model. This is a standard baseline when improving decoding using reward functions.\n* For experiment 4.1.4, can you provide CI over multiple seeds of PPO training? It is a common practice since RL training is known to be unstable, and the performance can vary between experiments [1]. I also agree with the authors that seeing the results of an on-policy reward model will be interesting, although this can be expensive to train because of the need for annotations.\n* Regarding experiment 4.2.2, it is well established that using Q values during decoding can improve performance. Wouldn't it be more relevant to compare this with other methods that use Q functions during decoding [2]? A beam search over SFT seems to me to be too weak of a baseline.\n\n[1] Agarwal, Rishabh, et al. \"Deep reinforcement learning at the edge of the statistical precipice.\" Advances in neural information processing systems 34 (2021): 29304-29320.\n\n[2] Han, Seungwook, et al. \"Value Augmented Sampling for Language Model Alignment and Personalization.\" arXiv preprint arXiv:2405.06639 (2024).", "limitations": "The paper properly addresses its limitations.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* Although the proposed method is not new and relies heavily on Eysenbach et al. 2022, it has not been used before with LLMs.\n* The methods and experiments are described in a straightforward and easy-to-follow manner.\n* I've found the results of experiments 4.1.2 and 4.2.1 particularly interesting. There has been no change to the objective of RLHF reward models since its inception, and the proposed loss seems to be able to improve the reward model without any further annotations.", "weaknesses": "* The fact that the definition of a \"goal state\" in the language space is ambiguous since the state includes the entire generated response is ignored in the paper. Even in tasks that we have a clear \"goal\" in, like math, it is a bit weird. Given two responses with the same final solution but different intermediate reasoning, do you expect their hidden representation to be exactly the same? What is the meaning of averaging representations of different preferred responses? Is it supposed to be an approximation to the average cosine similarity to all vectors?\n* This leads me to the fact that the prototypes used for Q-value estimation during inference time seemed a bit arbitrary. Did you ablate this during your work on the paper?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewers insightful and useful comments. In the following sections, we address the questions and comments raised by the reviewer. We plan on incorporating all our discussions into the final version of this paper.\n\n**\"The definition of a \"goal state\" in the language space...\"** One of the benefits of our method is that it allows capturing goals which depend on previous information within a generation.\nFor example, for reasoning tasks such as GSM8k and MATH, correct solutions depend on the final answer and intermediate reasoning.\nSimilarly, for natural language tasks like helpfulness and harmfulness, a preferred solution with respect to a particular goal may depend on several parts of a generation.\nBy taking a set of responses and averaging, we are able to produce a more accurate representation of the goal (see the shared reviewer response where we find averaging on more examples leads to better performance).\n\n**\"...Given two responses with the same final solution but different intermediate reasoning, do you expect their hidden representation to be exactly the same?\"** We observe that the representations for correct completions to the same prompt can be different depending on the reasoning path. \nFor example, we run UMAP on 5K sample completions from the train set as well as a test set (GSM8k) and plot the resulting 2D plot for 7 prompts that have multiple solutions in Figure 3 within the additional supplementary material page.\n\nStill, we observe preferred and disprefferred completions are separated. In Figure 3, we also plot 5K preferred and 5K dispreffered base model completions from the train and test (GSM8k) datasets and plot the UMAP on the hidden representation. \n\nFinally, we would like to emphasize we take the average across many preferred completions, to more accurately capture the concept they represent.\nAs demonstrated in the results provided in the shared reviewer response, averaging across more preferred completions produces a better representation, which improves performance.\n\n**\"..What is the meaning of averaging representations of different preferred responses?\"** We refer the reviewer to the shared response.\n\n**\"..the prototypes used for Q-value estimation during inference time seemed a bit arbitrary. Did you ablate this during your work on the paper?\"** We refer the reviewer to the shared response.\n\n**\"..evaluation of the reward model using the AUROC metric, can you elaborate on how you calculated it..\"**  For each problem in the benchmarks, we take the greedy generation of the base model along with annotations of whether each completion provides the correct answer, which are provided by Toshiniwal et al. [6], to compute the AUROC score. \nConcretely, for each completion, we first format using the Nemo format template and use the Reward Model to predict a reward score. \nIn predicting the reward score, we also retrieve the predicted reward score for each token in the completion, not including the prompt tokens. Utilizing the predicted reward score and the annotation of the correctness of each completion, we compute an AUROC score using the Python scikit-learn package. Using the same procedure, we compute the partial AUROC scores at every tenth percentile of each completion.\n\n**\"..was the filtration done using the Q-value of the full answer or the partial one..\"** By definition, the Q-value is computed for each token in the sequence by taking the cosine similarity of the goal-state and the RM representation of the token, which is given by the last hidden layer. A model completion is filtered if the Q-value for any token in the completion sequence is less than the threshold, which was 0 for our experiments.\n\n**\"..baseline of best-of-50 using the vanilla reward model and your own reward model..\"** Since we have a total of 50 generations per problem, we use the baseline reward model (Codellama RM) and our reward model (Q-Function RM) to select the Top 1, 5, 10, and 25 samples as ranked by the reward scores. \nWith the selected sample, we perform majority vote and also note the average proportion of the sample K that are correct solutions. \nThe results are provided below.\n| Model| Top-K | GSM8k Accuracy (%) | GSM8k Prop. Correct (%) | MATH Accuracy (%) | MATH Prop. Correct (%) |\n|-|-|-|-|-|-|\n| Q-Function RM | 1| 84.6| 84.6| 51.7| 51.7|\n|| 5| 86.2| 84.5| 59.7| 51.1|\n|| 10| 86.2| 83.8| 59.5| 50.3 |\n|| 25| 85.5| 81.8| 57.8| 47.0|\n|Codellama RM| 1| 80.8| 80.8| 43.8| 43.8|\n|| 5| 84.3| 81.9| 54.0| 45.8|\n|| 10| 85.2| 82.1| 56.0| 46.9|\n|| 25| 85.8| 81.1| 56.5| 46.1|\n\nThe results show that our Q-Function RM clearly outperforms the baseline.\n\n**\"..CI over multiple seeds of PPO training..\"** Thank you for pointing this out. \nThese results are the average accuracy across 4 independent runs.\nWe included the 95\\% CI in Table 7 in the appendix (the CI's between the baseline and Q-Function are non-overlapping), and apologize for not including them in the main text.\nWe will include them in the main text in the final version of our paper.\n\n\"..other methods that use Q functions during decoding..\" While there are other methods that use Q functions during decoding (Seungwook, et. al.), these rely on training a reward model (RM) and subsequently training value networks from large offline datasets (30K-100K examples).\nOur work focuses on improving the representations of RMs, so that we can compute Q values in a extremely lightweight manner. \nFor instance, in 4.2.2, we use 20 examples.\nHence, these experiments focuses on comparing with baselines in this low data setting.\n\n**\"..on policy reward model..\"**\nWe chose to limit policy training in our experiments to a single iteration due to the significant labeling and computational costs associated with experimenting with further RLHF iterations.\n\n[6] Toshniwal, Shubham et al. “OpenMathInstruct-1: A 1.8 Million Math Instruction Tuning Dataset.” ArXiv abs/2402.10176 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OoYxV5Y1Oy", "reviewer": "Reviewer_o9bF", "summary": "The paper introduces an additional contrastive loss term for reward model training that targets the learning of goal-conditioned representations that encode expected reward for partially complete sequences. The results show that this has a positive impact on reward model accuracy, downstream RLHF with the reward model, and guided generation.", "questions": "In the appendix, the authors claim statistical significance by simply citing the size of the evaluation set. Given that confidence intervals are not reported (presumably due to computational constraints), is this claim grounded in anything more rigorous?", "limitations": "The authors effectively outline the limitations of their work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper shows novelty and insight in using representation learning and goal conditioning to improve reward models and experimentally demonstrates the utility of doing so.\n\nExperiments cover a number of useful metrics (spanning reward model accuracy, downstream utility in RL, downstream utility in guided generation) on appropriate benchmarks in two distinct settings (reasoning and alignment). \n\nThe paper is clearly written, with methods, results and analysis effectively communicated.", "weaknesses": "Whilst the paper demonstrates positive results on guided generation, the mechanism for getting the prototype seems quite arbitrary and it would be useful to present a few (maybe use case dependent) alternative approaches.\n\nAs the authors mention, the performance gains from using their reward model are less than the gains in reward model accuracy (and relatively small in general). Given the authors claim that this is most likely due to off-policy issues, it is unclear why the experiment of updating the reward model and training using the updated reward model was not run.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their useful comments and insightful feedback.\nWe respond to the reviewers comments and questions.\nWe plan on incorporating all these responses in the final work.\n\n**\"..the mechanism for getting the prototype..\"** Please see our shared response to the reviewers. \nAdditionally, we consider two other methods: (a) prompting the model to generate the prototype, and (b) using an auxiliary dataset (in our case, HelpSteer to construct the prototype). Both of these methods underperform the original method on Helpful-Harmless (a: $70.5 \\rightarrow 69.5$, b: $70.5 \\rightarrow 69.6$), asserting our original mechanism of getting the prototype.\n\n**\"..updating the reward model and training..\"** Thank you for your comment. \nWhile we state it could be possible to improve performance of the policy further by continuing to update the reward model and perform PPO (4.1.4), we choose to limit policy training in our experiments to a single iteration due to the significant labeling and computational costs associated with experimenting with further iterations of RLHF training.\n\n**\"..statistical significance by simply citing the size of the evaluation set..\"** Thank you for pointing this out.\nFor the natural language alignment experiments, we evaluated statistical significance by performing a Student's t-test. For both the experiments in Section 4.2.2 the p-values are significant, namely, the p-value for Llama 8b Reward experiment is 0.002 and the p-value for the Q-function 8B reward experiment is 0.001.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their valuable and insightful comments.
It is appreciated that the reviewers found our work to be well-written, novel, and insightful for improving RM and policy performance. 
In this section, we provide further elaboration on our method of constructing the goal state.
We will include this discussion and results in the final paper.

**Goal State Construction**  

Because reviewers brought up several questions about the goal state, we briefly review how we construct the goal state.
In order to construct the goal state, given a set of completions $\\{y_i = [y_{(i, 0)}, \dots, y_{(i, t_i)}] | i = 1, \dots, N\\}$ we take the last reward model hidden state of the final token, $h(y_{(i, t_i)})$, and then we average across all completions: $ \frac{1}{N} \sum_{i=1}^N h(y_{(i, t_i)}) $, where $h$ represents the last hidden layer of the reward model. 

**Choice of Goal State**  

Reviewers brought up several questions about why we chose to average across a set of preferred completions from the training dataset as the goal state during inference. 
We experimented with using the mean across a set of completions that contain certain desirable attributes (e.g., being correct or helpful) as the goal state.
Our choice of taking the mean across these representations is motivated by prior work, which has found taking averages across representations can more accurately encode concepts or relationships [2, 3].
For example, by averaging across the set of correct completions in Math, we better capture the notion of the solution being correct. 

To ablate this choice of goal state during inference, we conduct a new experiment under the mathematical reasoning with code setting (Section 4.1) where we evaluate the effect of (a) choosing poor examples to construct the goal state and (b) the number of examples used to construct the goal state.
In particular, we evaluate 3 settings for picking the sample completions for constructing the goal state. 
First, we vary the number of preferred completions used to construct the goal state to evaluate the effect of adding more examples of good reasoning. 
Second, we vary the number of dispreferred completions used to construct the goal state, to evaluate whether bad examples lead to poor performance. 
Third, we incrementally add more dispreferred completions to a fixed sample of preferred completion in order to measure the robustness of our goal state computation to negative examples.
We refer to the addition of dispreferred completions on top of all preferred completions as adding ``corrupted'' examples. 
The results are in Figure 1, within the additional supplementary material page.
Overall, these results show that negative or unhelpful completions degrade performance, indicating the importance of choosing relevant examples for computing the goal state.
Additionally, they demonstrate that having more examples of the concept leads to better performance as more generations are filtered with comparable proportion of those remaining generations being correct.

We additionally ablate our choice of using the last token in the completion sequence to construct the goal state.
In this experiment, we repeat our filtering experiments from 4.1.3. 
Except, we randomly sample a token from in the completion sequence and use it as the goal state:

| Sampling Method | GSM8k Accuracy (%) | GSM8k Prop. Correct (%) | MATH Accuracy (%) | MATH Prop. Correct (%) |
|-----------------|--------------------|-------------------------|-------------------|------------------------|
| Last Token      | 86.0               | 84.0                    | 59.6              | 52.0                   |
| Random Token    | 85.9               | 81.8                    | 57.3              | 45.3                   |


From the result of this ablation, we see that last token sampling has superior performance, indicating using the last token helps achieve better performance.


[2] Mikolov, Tomas et al. “Efficient Estimation of Word Representations in Vector Space.” International Conference on Learning Representations (2013).

[3] Le, Quoc V. and Tomas Mikolov. “Distributed Representations of Sentences and Documents.” International Conference on Machine Learning (2014).

Meta-review: This paper proposes a method to bias reward models in RLHF to align with goals at intermediate stages of the input sequence, using contrastive learning. This helps using these rewards to prune outputs during the generation process and to even improve on some benchmarks. Reviewers generally appreciated the paper’s presentation and contributions. Certain discussion points (e.g., averaging over completions, comparison with DPO in-distribution and OOD, clarifying the Q function, base-of-k baseline) are worth including in the paper.
Decision: Accept (poster)