Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs

Summary

This paper is concerned with the overoptimization issue in reward modeling: When optimizing a policy against a reward model, this leads to a distributional shift that can lead to an increase in the proxy score while the true score decreases. This paper addresses these issues by regularizing the hidden states in reward models utilizing an additional output head that (with some variations) predicts the text of the winning trajectories in the comparison data. The idea is that this prediction task (a) helps keep the intermediate, strong representations from pretraining, and (b) adapts them to text that is more similar to what's encountered in post-training, thus leading to stronger generalization to the data encountered when optimizing the policy.

Strengths

My impression is that this is a strong paper dealing with the important overoptimization issue in RLHF. The motivation is strong and the experimental evaluation considers many baselines, settings, and datasets to be checked.

Weaknesses

No strong weaknesses come to mind, though see the questions and comments below for other experiments that could be interesting to gain more confidence (though imho not necessary for this submission to be accepted).

Questions

I'm using this section for questions, suggestions, and minor weaknesses. a. "*In the RLHF stage, various policy optimization methods can be applied, with two frequently used methods being Best-of-n Sampling (BoN) and Proximal Policy Optimization (PPO).*" --- This paper seems to view RLHF as the stage *after* reward modeling, so just policy optimization. In contrast, I think it's more typical to view reward modeling as *part of* RLHF. b. "*While straightforward, the DPO regularization requires a reference model during training.*" --- Don't you also need a reference policy during the PPO stage, given by the KL regularization with the SFT model? So it seems like this is no additional cost here? c. Equation (7): Did you also consider replacing $\log \circ \sigma \circ \beta \circ \log$ by just $\log$? Do you expect this to perform better/worse? d. Did you consider to interleave reward modeling with pretraining on a pretraining dataset to regularize the hidden states? I wonder if this performs better or worse than essentially doing pretraining on the preference dataset (as you do in your SFT regularization method). e. 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? I could imagine that the former performs better since it brings the reward model closer to understand a distribution that the PPO stage *will steer towards*, but I'm curious if that prediction holds up empirically. f. For HHH-Alignment, MT-Bench, and RewardBench, could you give more details on what the benchmarks test for and what the scores mean? g. I'd recommend putting the limitations and broader impact sections into the main paper.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

Limitations are described in Appendix D.

Summary

This paper introduces a method that retains the base model's language model head while incorporating text-generation losses to preserve the hidden states' text generation capabilities.

Strengths

* The paper is well-written, and the idea is straightforward. * The code is easy to understand and implement. * This method appears to be efficient and readily integrable with existing alignment approaches.

Weaknesses

* 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. * The results demonstrate limited advantages. It is recommended to validate the method's benefits with a larger dataset.

Questions

* The training details mention that the reward models are trained for only 2 epochs. This choice may not be optimal, considering the risk of overfitting in the full fine-tuning setting. It would be helpful to discuss the reasoning behind this decision or the impact of learning rate and convergence speed. * The simplicity of the regularization method is intriguing. Can it effectively mitigate reward hacking, potentially surpassing reward model ensemble techniques? It would be valuable to explore and discuss this aspect.

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

* 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. * The inclusion of manually synthetic data in the dataset may not accurately reflect real-world data, limiting the generalizability of the findings.

Summary

This paper proposes generalizable reward model (GRM) which modifies the standard reward-learning objective by adding an auxiliary task with a separate language modeling head. The auxiliary loss is either DPO or SFT. Experiments and ablations are conducted using mistral and gemma models and data from unified-feedback and evaluating on HHH, MT-bench, and rewardbench showing GRM improves OOD performance. Finally, GRM reward models are used to train new policies.

Strengths

1. The idea of adding an auxiliary loss is simple and elegant. The paper provides a good intuitive argument for why it may be useful and clearly presents the potential ways to implement the auxiliary loss (DPO and SFT). 2. The experiments and ablations are thorough and well-motivated. The paper considers a variety of baseline techniques to make reward models more robust like adding a margin term, label smoothing, or an ensemble. The paper also does a good job of comparing the low-data and high-data regimes. 3. The results seem to show a consistent, if sometimes modest, improvement from GRM over baseline methods. And the improvements are larger in the low-data regime.

Weaknesses

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

Questions

None

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

There could be more discussion of limitations

Summary

The paper addresses the limitations of current reward models used in the reinforcement learning from human feedback (RLHF) framework, specifically their generalization capabilities to unseen prompts and responses. This limitation often leads to reward over-optimization, where the excessive optimization of rewards results in a decline in actual performance. The study proposes an approach to enhance the reward model's generalization ability against distribution shifts by regularizing the hidden states.

Strengths

The motivation for the study is sound, and the experiments validate the effectiveness of the proposed method.

Weaknesses

1. The proposed method lacks innovation as it combines DPO and SFT loss into the RM training phase as a regularization term. The inclusion of SFT loss in RM training has already been explored in previous works, such as InstructGPT and Anthropic's RM training. 2. In the introduction, the authors mention that a randomly initialized head can distort pre-trained features, negatively impacting out-of-distribution (OOD) performance. Inspired by this finding, they propose to regularize feature distortion during fine-tuning for preference learning. However, there is no experimental evidence provided to support that this motivation holds true for RM. The improvements from adding regularization alone do not sufficiently prove that this motivation is solid. 3. The results for label smoothing are missing in Figures 2 and 3, which should be addressed. 4. The experimental section should include the alignment results after RL.

Questions

Please refer to the Weaknesses section for questions regarding the paper.

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes