Learning Goal-Conditioned Representations for Language Reward Models

We thank the reviewer for their useful comments and insightful feedback. We respond to the reviewers comments and questions. We plan on incorporating all these responses in the final work.

"..the mechanism for getting the prototype.." Please see our shared response to the reviewers. Additionally, 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.

"..updating the reward model and training.." Thank you for your comment. While 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.

"..statistical significance by simply citing the size of the evaluation set.." Thank you for pointing this out. For 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.

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.

"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. For example, for reasoning tasks such as GSM8k and MATH, correct solutions depend on the final answer and intermediate reasoning. Similarly, 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. By 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).

"...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. For 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.

Still, 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.

Finally, we would like to emphasize we take the average across many preferred completions, to more accurately capture the concept they represent. As demonstrated in the results provided in the shared reviewer response, averaging across more preferred completions produces a better representation, which improves performance.

"..What is the meaning of averaging representations of different preferred responses?" We refer the reviewer to the shared response.

"..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.

"..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. Concretely, for each completion, we first format using the Nemo format template and use the Reward Model to predict a reward score. In 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.

"..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.

"..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. With the selected sample, we perform majority vote and also note the average proportion of the sample K that are correct solutions. The results are provided below.

The results show that our Q-Function RM clearly outperforms the baseline.

"..CI over multiple seeds of PPO training.." Thank you for pointing this out. These results are the average accuracy across 4 independent runs. We 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. We will include them in the main text in the final version of our paper.

"..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). Our work focuses on improving the representations of RMs, so that we can compute Q values in a extremely lightweight manner. For instance, in 4.2.2, we use 20 examples. Hence, these experiments focuses on comparing with baselines in this low data setting.

"..on policy reward model.." We 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.

[6] Toshniwal, Shubham et al. “OpenMathInstruct-1: A 1.8 Million Math Instruction Tuning Dataset.” ArXiv abs/2402.10176 (2024).

We thank the reviewer for their response. We will incorporate these clarifications into the final version.

We are grateful to the reviewer for their thoughtful review and insightful suggestions. We are pleased to know they acknowledge the novelty of the goal-conditioned approach with LLMs and found our experiments to be thorough. The reviewer brings up the interesting point of comparing our proposed method with DPO. We compare the performance of DPO with our proposed method on the mathematical reasoning with code execution tasks (Section 4.1). Furthermore, we intend to incorporate this response and discussion into the final version of the paper.

In 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. We compare performance of DPO with PPO training using the baseline reward model and our contrastive reward model. We present average accuracy across 4 independent runs for PPO and 2 independent runs for DPO. The base model results are also shown as a reference point presented in [6].

Interestingly, 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. These 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.

[1] Xu, Shusheng et al. “Is DPO Superior to PPO for LLM Alignment? A Comprehensive Study.” ArXiv abs/2404.10719 (2024).

[6] Toshniwal, Shubham et al. “OpenMathInstruct-1: A 1.8 Million Math Instruction Tuning Dataset.” ArXiv abs/2402.10176 (2024).

We thank the reviewer for useful comments and for appreciating the originality and performance improvements of the method. We address the reviewers questions and comments in the following sections. Additionally, we plan on incorporating all our responses and discussions into the final version of the paper.

Correlation between Q-values and reward scores Thank you for your comment. To 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. The reward scores and response lengths are shown in Figure 2 within the additional supplementary material page and the performance of the model at the end of each episode is given here.

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.

"..the functional form of the Q function.." The Q function is parameterized by a scoring function $f$ and encoder. In this paper, we use the cosine similarlity as $f$ and the encoder is a causal LM, such as Llama. In 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.

"..choice of goal states during inference" We refer the reviewer to the shared response.

"..investigating approaches such as Upside Down RL.." Our work focuses on improving the representations learned by reward models (RMs) for aligning LMs. Our improved methods of training RMs leads to improvements in RM performance, downstream policy learning, and model steerability. Approaches 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.