Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs

We thank the reviewer for the valuable comments, and we provide clarification to your concerns as follows. We appreciate it if you have any further questions or comments.

Q1: This paper seems to view RLHF as the stage after reward modeling ... I think it's more typical to view reward modeling as part of RLHF.

A1: Thanks for the question. We agree with the reviewer that reward modeling is one step of RLHF. We will revise Section 2 Background to reflect this more accurately.

Q2: "The DPO regularization requires a reference model during training." --- Don't you also need a reference policy during the PPO stage ... So it seems like this is no additional cost here?

A2: The cost includes both memory usage and computation time. Although DPO regularization does not increase memory usage for PPO, it still extends the computation time required for reward modeling. Furthermore, for other policy optimization methods such as BoN and reject sampling, the introduction of a reference model results in additional memory usage. Therefore, we claim that DPO regularization is more costly for reward modeling and prefer SFT and DPO w/o reference regularization according to the experimental results.

Q3: Equation (7): Did you also consider replacing $log \sigma (\beta log)$ by just $log$? Do you expect this to perform better/worse?

A3: In ideal situations, the two forms should perform similarly. We also tried the $log$ form but found that it requires different hyperparameter tuning for $\alpha$ due to changes in the loss scale. In the tables below, "GRM logreg" outperforms the baseline reward model and matches or even slightly exceeds the performance of our GRM on OOD tasks when $\alpha$ is appropriate.

We found that the current form of SFT regularization can directly use the same hyperparameters as our DPO regularization. Therefore, we opted for this solution to maintain coherence with these regularizations and avoid the need for hyperparameter adjustments.

Q4: Did you consider to interleave reward modeling with pretraining on a pretraining dataset to regularize the hidden states?

A4: Thanks for the insightful quesiton. One challenge of the described approach is that we don't have the pretraining dataset for many language models. We think that other data formats can also be useful as a regularization, but preference data is more advantageous because it allows us to avoid using external datasets during reward modeling by leveraging the preference dataset, which may better match the distribution of prompts and responses. Please refer to the global response B for the additional results.

Q5: Did you check whether in SFT regularization it matters whether you predict the text of the winning responses, instead of trying to predict the losing responses?

A5: In our SFT regularization, we only use the winning responses to regularize the hidden states. During the rebuttal, we also tested regularizing with the negative responses, referred to as "GRM neg." The results show very close IID evaluation accuracy but relatively lower OOD accuracy compared to our default setting, potentially due to the generally lower quality of rejected responses. The impact of this reward model on PPO training requires further examination. However, we agree with the reviewer's conjecture that the rejected responses may lead to unsatisfactory representations for the desired distribution.

Q6: For HHH-Alignment, MT-Bench, and RewardBench, could you give more details on what the benchmarks test for and what the scores mean?

A6: In Section 4, we briefly introduce three benchmarks. The HHH-Alignment dataset evaluates language models on helpfulness, honesty, and harmlessness. The MT-Bench dataset contains 3.3K human preferences for model responses generated by LLMs in response to MT-Bench questions. RewardBench is a new benchmark designed to evaluate the reward model's ability to select human-preferred responses for chat, reasoning, and safety tasks. In summary, these benchmarks are used to assess the reward model's alignment with human preferences and generalization ability across different prompt-response distributions. The scores represent the average accuracy across each benchmark, with higher accuracy indicating better performance.

Q7: I'd recommend putting the limitations and broader impact sections into the main paper.

A7: We agree with the reviewer and will move the two sections to the end of the main paper in the revised version.

Thank you for the valuable comments, and we provide clarification to your concerns as follows.

Q1: Figure 3 (b) appears unusual, as the gold score decreases at the beginning of training. This could indicate suboptimal hyperparameter tuning or potential drawbacks in the pipeline.

A1: We found that this issue is influenced by the learning rate of PPO. By using a smaller learning rate (1e-5), we observed an initial improvement in the gold score during the first few steps, followed by subsequent drops (see Figure 2 of the supplementary PDF). Although the trend remains the same, the smaller learning rate delays the drop. In contrast, even with a larger learning rate, using our GRM allows PPO to learn stably, further emphasizing the importance of our proposed method for reward modeling and alleviating the burden of tuning hyperparameters for PPO.

Moreover, we carefully followed the pipeline commonly adopted by the community [2][5] and used the open-source implementation of PPO from the Huggingface/trl package to ensure a correct and reliable pipeline.

Q2: The results demonstrate limited advantages. It is recommended to validate the method's benefits with a larger dataset.

A2: We would like to argue that the advantages of GRM are significant. As illustrated in Table 1 and 2 of our paper, increasing the training data from 40K to 400K enhances the OOD scores from 70.3/69.1 to 73.4/71.2. In contrast, GRM with just 40K training data achieves OOD scores of 78.7/73.0, showing a much greater improvement than a 9x increase in the dataset size. This is particularly important in real-world applications where there are limited samples for fine-tuning.

Additionally, in global response A, we present further results on RewardBench, demonstrating the strong potential of scaling GRM to larger models and datasets. Our findings indicate that GRM even outperforms a 34B reward model and GPT-4 as a judge.

Moreover, as shown in Table 1 and 2, GRM outperforms a strong ensemble baseline (n=3) while using only about 1/3 of the computational cost. In Section 5.2, GRM significantly outperforms other reward models for RLHF in terms of the gold score, alleviating the reward over-optimization problem of RLHF.

Q3: The training details mention that the reward models are trained for only 2 epochs ... It would be helpful to discuss the reasoning behind this decision or the impact of learning rate and convergence speed.

A3: In our experience, the training design depends on the dataset's quality, diversity, and size, as well as the base model and optimizer hyperparameters used for training. There is no one-size-fits-all hyperparameter setting. Prior work [2] trains reward models for 5 epochs and suggests determining this number based on a nearly converging validation loss.

Following this insight, we reserve 1% of the training set for validation (e.g., 4K for 400K training data) and found that 2 epochs are sufficient for reward modeling with LoRA in our setting. As shown in the table below, we observed convergence in the validation loss during the second epoch, with no further improvement in the third epoch. For full-parameter training experiments added during rebuttal, which are more prone to overfitting, we train the reward model for only one epoch.

Q4 The simplicity of the regularization method is intriguing. Can it effectively mitigate reward hacking, potentially surpassing reward model ensemble techniques?

A4: Yes, in Section 5.2, we already explored the effectiveness of GRM for RLHF, particularly focusing on two commonly used policy optimization methods: BoN and PPO. As shown in Figure 3 of our main paper and 3 figures in the supplementary PDF, GRM consistently excels in both the PPO and noisy label experiments. Unlike the baselines (including the ensemble techniques), which show an increasing proxy score but a declining gold score, GRM demonstrates its effectiveness in mitigating reward hacking.

Q5: The experiments may lack solidity due to the heavy reliance on the reward model in the PPO algorithm. If the reward model is not properly trained, the results can be significantly impacted, introducing randomness to the work.

A5: We would like to emphasize that our main focus is on improving reward modeling (Section 5.1), subsequently enhancing RLHF (Section 5.2). Our reward model is not limited to the PPO algorithm; it can be applied to any policy optimization method based on explicit reward modeling, such as BoN or reject sampling. As the reviewer mentioned, an improperly trained reward model can be hacked by RL algorithms, as demonstrated in Section 5.2. In contrast, GRM shows better robustness against reward hacking, underscoring its importance for robust preference learning.

Q6: The inclusion of manually synthetic data in the dataset may not accurately reflect real-world data, limiting the generalizability of the findings.

A6: In Section 5.2, synthetic data is used in the RLHF experiment to allow the gold reward model to annotate preference labels without human labor. This approach is also adopted by prior works [2][3][5], and even research from companies like OpenAI and Google uses this synthetic method to reduce human labor. We have also mentioned this in the limitations section (Appendix D). In other experiments in Section 5.1, we use the UnifiedFeedback dataset, which includes datasets based on human annotators, such as Anthropic/hh-rlhf, lmsys/chatbot_arena_conversations, openai/summarize_from_feedback. Therefore, we believe our setting also reflects real-world data.

Thank you for the valuable comments, and we provide clarification to your concerns as follows.

Q1: While the paper is generally comprehensive, there could be additional exploration of just using language modeling as an auxiliary task. It is not clear if the benefit of GRM is coming from language modeling the data from the preference dataset or if it would benefit from language modeling on any data (or even data from an even broader distribution, controlling for more OOD data). A further exploration of this would be good.

A1: Thanks for the insightful question. We believe that other data formats can also be useful, but preference data is more advantageous because it allows us to avoid using external datasets during reward modeling by leveraging the preference dataset, which may better match the distribution of prompts and responses. To illustrate this, we provide an experiment below using GRM with text-generation regularization on an open-source pretraining dataset, togethercomputer/RedPajama-Data-1T-Sample (including text from Commoncrawl, Arxiv, and books), referred to as 'GRM pretrain reg'. The results in global response B suggest that 'GRM pretrain reg' outperforms the baseline reward model and matches the performance of GRM when the dataset size is large (400K). However, when the dataset size is small, using a pretraining dataset is less effective than using the preference dataset.

Q2: There could be more discussion of limitations

A2: Thank you for the question. In addition to the limitations mentioned in Appendix D, another constraint is that we did not test reward models with parameter sizes exceeding 10B due to computational limitations. In the main paper, we tested models with sizes of 2B and 7B. During the rebuttal period, we also provided results for 8B reward models in global response A, which demonstrate promising potential for scaling GRM to larger reward models and datasets. Further efforts to extend our method to even larger reward models could be highly promising.

Thank you for the valuable comments, and we provide clarification to your concerns as follows.

Q1: The proposed method lacks innovation ... The inclusion of SFT loss in RM training has already been explored in previous works, such as InstructGPT and Anthropic's RM training.

A1: We would like to clarify that no previous work has incorporated text-generation loss as a regularization term for reward modeling. InstructGPT and Anthropic's RM training optimize the vanilla reward loss (Eq 1 in [1]) based on the Bradley-Terry model, and their SFT regularization is typically applied to language model training rather than reward modeling.

Our method is novel in that it integrates text-generation loss directly into the reward modeling process, a fundamentally different task from text generation. The enhancement effect of text-generation loss on reward modeling has not been explored in prior work. Additionally, our reward model structure is distinct to enforce this regularization. We maintain both a language head and a reward head within the reward model and apply text-generation loss to the language head during reward training.

In our experiments, compared to previous methods that use ensemble techniques [2][3] to enhance reward model capabilities, our method significantly reduces computational costs while achieving superior performance. This demonstrates promising results for more reliable and cost-effective reward modeling.

Q2: No experimental evidence supporting the claim that a randomly initialized head can distort pre-trained features in RM, consequently leading to a negative impact on OOD performance.

A2: The phenomenon is well-documented by [4], both theoretically and empirically (across a range of computer vision tasks). It is also easy to validate in the preference learning setting when using a smaller dataset size. We included a baseline, "Classifier (Frozen)", which fixes the base model’s features and only fine-tunes the classification head. When the dataset size is 8K (see the table below), the OOD evaluation results of the baseline reward model (without freezing the backbone) are worse than those of the frozen one, demonstrating the negative effect of distorting pre-trained features. However, we would like to note that when the dataset size is sufficiently large, this negative effect can be mitigated, and the baseline reward model can surpass the frozen reward model due to having more trainable parameters to fit the large dataset.

In contrast, by regularizing the hidden states, our GRM can achieve the regularizing effect while fine-tuning all parameters, showing strong performance with both large and small dataset sizes.

Q3: The results for label smoothing are missing in Figures 2 and 3, which should be addressed.

A3: We found that reward models with label smoothing are easily hacked by BoN and PPO, resulting in worse performance compared to other baselines. Therefore, we omitted them from the policy optimization experiments. To address the reviewer's concern, we have included the label smoothing results for the BoN and PPO experiments using the 7B base model in Figure 1 of the supplementary PDF provided during the rebuttal.

Q4: The experimental section should include the alignment results after RL.

A4: In Section 5.2, we evaluated different reward models for both RL training and Best-of-N sampling, using the gold score as the measurement for alignment, as adopted by prior studies [2][5] for reward modeling. Additionally, to demonstrate the advantage of GRM over vanilla reward modeling, we evaluated the win rate of models after PPO training with GRM against those with the vanilla reward model. The evaluation was conducted using GPT-4o on 100 randomly selected prompts from the test set in UnifiedFeedback, with the order of responses randomly flipped to avoid order bias. The results below show a significantly higher win rate for GRM than the vanilla reward model across two different base reward models.

Dear Reviewer tSHJ,

We deeply appreciate your time, effort, and valuable feedback. We hope that our responses, including detailed explanations about the differences between the prior reward modeling paradigm and our approach, as well as additional experimental results on feature distortion and alignment after RL, have addressed your concerns.

If our responses have resolved your issues, we would greatly appreciate it if you could consider raising the score. If you have any additional concerns, please feel free to post them, and we would be happy to discuss them further before the discussion deadline on Aug 13th.

Thank you once again for your thoughtful review.

Best, The Authors

Dear Reviewer tSHJ,

We deeply appreciate your time, effort, and valuable feedback. Since the discussion deadline is less than 10 hours away and we have not yet received an acknowledgment from you, we kindly request you to review our replies. We hope that our responses, including detailed explanations about the differences between the prior reward modeling paradigm and our approach, as well as additional experimental results on feature distortion and alignment after RL, have addressed your concerns.

If our responses have resolved your issues, we would greatly appreciate it if you could consider raising the score. If you have any additional concerns, please feel free to post them, and we would be happy to discuss them further before the discussion deadline

Thank you once again for your constructive review.

Best,

The Authors