Prior Constraints-based Reward Model Training for Aligning Large Language Models (2024)

Abstract: “Reinforcement learning with human feedback for aligning large language models (LLMs) trainsa reward model typically using ranking loss with comparison pairs. However, the training pro-cedure suffers from an inherent problem: the uncontrolled scaling of reward scores during rein-forcement learning due to the lack of constraints while training the reward model. This paperproposes a Prior Constraints-based Reward Model (PCRM) training method to mitigate thisproblem. PCRM incorporates prior constraints—specifically, length ratio and cosine similaritybetween outputs of each comparison pair—during reward model training to regulate optimiza-tion magnitude and control score margins. We comprehensively evaluate PCRM by examining itsrank correlation with human preferences and its effectiveness in aligning LLMs via RL. Exper-imental results demonstrate that PCRM significantly improves alignment performance by effec-tively constraining reward score scaling. As another bonus, our method is easily integrated intoarbitrary rank-based alignment methods, such as direct preference optimization, and can yieldconsistent improvement. The code is available at https://github.com/wangclnlp/DeepSpeed-Chat-Extension/tree/PCRM.”

## introduction reinforcement learning with human feedback (rlhf) has been proven to be an advanced technology to align large language models (llms) with human preferences @xcite @xcite . it builds upon preference data, which rates and compares different outputs given the same input, where this rating is conducted by either human annotators or llms @xcite @xcite . in practice, rlhf trains a reward model on the preference data with ranking loss for higher scores on preferred outputs than dispreferred ones. then, rl algorithms such as proximal policy optimization (ppo) @xcite are employed to fine-tune the llm with the aim of optimizing this reward. during rl training, the reward model will give scores as signals, and the llm will align the preference by increasing the probabilities of sampled outputs with higher reward scores. it is discernible that the alignment of the llm is significantly influenced by how well the reward model is trained. however, the training procedure of the reward model actually suffers from a failure mode: it learns from the preference data with the standard ranking loss yet provides a reward score for the sampled outputs during rl training. this training mode causes an inherent problem: the ranking loss increases the score margin between the outputs of this comparison without constraint during the training procedure, which makes an uncontrollable scale of scores in the process of rl training. for instance, the reward model trained by ranking loss can predict reward scores with a relatively large margin to two sampled outputs that do not differ very much @xcite . to address the problem, we propose a prior constraints-based reward model (pcrm) training method in this paper, which incorporates prior constraints while training the reward model. specifically, we select two features of the preference data for designing prior constraints: length ratio and cosine similarity of the outputs with the same prompt. the length ratio is computationally less demanding and directly reflects data differences, while the cosine similarity captures deeper semantic features despite being more computationally intensive. these constraints effectively regulate the optimization magnitude across different outputs for the same input, thereby controlling the margin of scores predicted by the reward model. we comprehensively evaluate the proposed pcrm in the following two ways. firstly, we test the rank correlation between the preferences predicted by pcrm and human preferences. it can assess the extent to which the reward model can serve as a surrogate for human-derived preference signals. secondly, we verify the effectiveness of the pcrm in aligning an llm via rl. this also demonstrates the influence of constraints on alignment performance straightly. notably, pcrm can yield a +2.48% improvement in the gpt-4 win rate for the dialogue task compared to the traditional rlhf. furthermore, we integrate our method into direct preference optimization (dpo), a rank-based alignment method. the results show that our method can also be effective in improving the rank-based alignment methods, e.g., a 2.95% increase in the gpt-4 win rate on the dialogue task compared to dpo. ## related work reinforcement learning with human feedback (rlhf) is a crucial technique to ensure that the behaviours of large language models (llms) are consistent with human preferences @xcite @xcite . recent works resorted to building a better reward model to improve the performance of rlhf on llms. these methods could be classified into three groups. the first group aimed to efficiently produce human preference data @xcite @xcite . for instance, @xcite directly employed an llm to annotate the comparison pairs. significantly, the reward model trained on the llm-annotated comparison pairs can achieve closer performance when applied to rlhf than the one trained on human-annotated comparison pairs. the second group tended to design fine-grained reward models to provide multiple reward scores for different reward criteria @xcite @xcite . notably, @xcite learned various evaluation models from an llm as reward models on the summarization task, including reward models scored for relevance, scored for fluency, scored for consistency, and scored for coherence. the third group that has attracted less attention generally explored how to merge multiple reward models in rlhf, such as reward model ensemble @xcite and learning reward weights @xcite . different from these methods, in this paper, we proposed a superior training schema that enables this reward model to give a more accurate score in rlhf. specifically, we refined the conventional ranking loss, typically utilized in reward model training, by incorporating prior constraints. these constraints are designed to regulate the margin between the scores of various outputs generated from the same input. to the best of our knowledge, there is almost no previous work in training a reward model with prior constraints. the few related one is llama2 @xcite , which added a margin component in ranking loss for training each comparison pair. although this margin simulates a constraint, it requires a manual annotation, which significantly increases the cost of training the reward model. additionally, this work did not provide a specialized exploration or analysis on the topic of constraining the margin of scores during training reward models. in comparison, all of our prior constraints were computed automatically and did not require any human annotation. we also provided sufficient theoretical analysis and experiments to constrain the margin of scores. furthermore, although there are some methods such as dpo @xcite and rrhf @xcite that circumvent the need for training reward models, they still suffer from an inherent limitation from ranking loss, which is mentioned in dpo @xcite without a solution. our proposed pcrm can be easily extended to these methods to yield some benefits (see section 5.7.2). ## experiments we evaluate the proposed pcrm in the following two ways. firstly, we test the reward model trained by pcrm. secondly, we analyze the performance of applying this trained reward model to rlhf. we conduct experiments on the commonly used generation tasks, including dialogue and summarization. ## conclusions rlhf plays an important part in applications based on llms for their ability to align language models with human preferences. as a proxy for human preferences, the reward model makes a significant contribution. we have introduced pcrm, a method of training reward models with prior constraints to better control the distribution of reward scores without sacrificing prediction accuracy. instead of optimizing the vanilla reward model with the goal of maximizing the margin of scores, pcrm restricts the optimization process with the provided prior information and fits the distribution of reward scores for better alignment. furthermore, this method is applicable to dpo, which treats the language model as a reward model and optimizes it directly. through experiments on downstream tasks, we have validated the effectiveness of this method.