Improving Discriminative Capability of Reward Models in RLHF Using Contrastive Learning (2024)

Abstract: Reinforcement Learning from Human Feedback (RLHF) is a crucial approach to aligning language models with human values and intentions. A fundamental challenge in this method lies in ensuring that the reward model accurately understands and evaluates human preferences. Current methods rely on ranking losses to teach the reward model to assess preferences, but they are susceptible to noise and ambiguous data, often failing to deeply understand human intentions. To address this issue, we introduce contrastive learning into the reward modeling process. In addition to supervised ranking loss, we introduce an unsupervised contrastive loss to enable the reward model to fully capture the distinctions in contrastive data. Experimental results demonstrate that the proposed contrastive learning-based reward modeling method effectively enhances the generalization of the reward model, stabilizes the reinforcement learning training process, and improves the final alignment with human preferences.

## introduction in the evolving landscape of artificial intelligence, particularly within language models, achieving alignment between ai behaviors and human intentions emerges as a critical goal @xcite . the alignment ensures that ai systems operate in harmony with the expectations and values of their users and designers. reinforcement learning from human feedback (rlhf) is a critical alignment technique encompassing two main steps @xcite . firstly, a reward model is trained using preference data gathered from numerous crowdworkers to discern outputs more aligned with human preferences. secondly, the language model is optimized through reinforcement learning (rl) to maximize the reward. the reward model plays a pivotal role in the rlhf process, and our aim is to establish it as a trustworthy proxy of human preferences. however, current reward modeling faces several challenges that hinder the accurate understanding and modeling of human intentions by the reward model. firstly, the presence of noisy and ambiguous data can impair the generalization capabilities of the reward model when over-reliant on manually annotated labels @xcite . secondly, existing supervised ranking losses may cause the model to rely on easy features to distinguish human preferences, rather than deeply understanding the underlying differences in the data @xcite @xcite . to address these issues, we investigate the integration of contrastive learning into the reward modeling process. specifically, we combine unsupervised contrastive losses based on clustering and instance-level differences with the existing supervised losses in reward modeling. this approach aims to enhance the model's understanding of differences among various responses while maintaining its scoring capabilities. additionally, we explore the construction of contrastive data for reward modeling to better integrate contrastive learning with reward modeling. experimental results demonstrate that the contrastive learning-based reward modeling effectively improves the generalization ability of the reward model, stabilizes the subsequent reinforcement learning training process, and enhances the final alignment performance. ## preliminaries the rlhf framework is structured around three core stages: supervised fine-tuning (sft), pref-erence sampling alongside reward model (rm) training, and rl fine-tuning through proximal policy optimization (ppo) @xcite . initially, a language model is refined through supervised training on a curated dataset tailored for specific tasks, yielding a model we refer to as π sft . ## method in reward modeling, the challenge of differentiating "chosen" from "rejected" responses stems from high feature similarity (figure 1 ), which hinders model performance. contrastive learning addresses distribution obtained through t-sne reveals a significant overlap in the features of chosen and rejected responses in the baseline model. however, when simcse is introduced into the reward modeling, this overlap between chosen and rejected responses decreases.(see appendix b for details) this by improving feature extraction and generalization through the use of varied samples, thus enhancing models' adaptability to new data. its integration into rlhf requires careful selection of contrastive samples. contrastive data construction inspired by simcse (gao et al., 2021), we enhance data variance by applying dropout after the attention mechanism outputs and in the feed-forward layers of the llama model. we propose two approaches for constructing contrastive data: 1) single representation pairs: performing contrastive learning with single representations from preference data. formally, @xmath0, y (i) )} 2n @xmath1, where 2n is the total number of chosen and rejected responses. dropout is applied to create pairs of representations, @xmath2, h (i) s } 2n @xmath3, where h (i) t and h (i) s are semantically related representations. ## experiments and results experiments we use llama 2, a model with 7 billion parameters, across all setups. our evaluation focuses on general dialogue tasks and summarization. for more details, see appendix a. we also provide case study, see appendix d. evaluation we compare our contrastive learning methods against sft, vanilla ppo, and dpo across tasks of helpfulness, harmlessness, and summarization (methodology in appendix c). results in to further validate the effectiveness of our method and the robustness of our evaluation, we assess two additional benchmarks for llm alignment: mt-bench @xcite and arena-hard @xcite . these benchmarks are widely recognized and utilized in the field of llm alignment evaluation. both have undergone extensive consistency checks and evaluation techniques that compare human assessments with gpt-4 evaluations, demonstrating a high degree of alignment. the evaluation results are presented in ## conclusion this study advances rlhf by integrating contrastive learning into reward modeling. this novel approach, which combines unsupervised contrastive losses with supervised ranking losses, better captures human preference distinctions. our findings indicate that this method not only improves the reward model's generalization but also ensures a more stable rl training process, align-ing ai behaviors more closely with human values. the enhanced model performance supports the effectiveness of contrastive learning in addressing current challenges of rlhf, paving the way for more reliable ai-human alignment.