Secrets of RLHF in Large Language Models Part II: Reward Modeling

Abstract: Reinforcement Learning from Human Feedback (RLHF) has become a crucial technology for aligning LLMs with human values and intentions, enabling models to produce more helpful and harmless responses. Reward models are trained as proxies for human preferences to drive reinforcement learning optimization. While reward models are often considered central to achieving high performance, they face the following challenges in practical applications: (1) Incorrect and ambiguous preference pairs in the dataset may hinder the reward model from accurately capturing human intent. (2) Reward models trained on data from a specific distribution often struggle to generalize to examples outside that distribution and are not suitable for iterative RLHF training. In this report, we attempt to address these two issues. (1) From a data perspective, we propose a method to measure the strength of preferences within the data, based on a voting mechanism of multiple reward models. Experimental results confirm that data with varying preference strengths have different impacts on reward model performance. We introduce a series of novel methods to mitigate the influence of incorrect and ambiguous preferences in the dataset and fully leverage high-quality preference data. (2) From an algorithmic standpoint, we introduce contrastive learning to enhance the ability of reward models to distinguish between chosen and rejected responses, thereby improving model generalization. Furthermore, we employ meta-learning to enable the reward model to maintain the ability to differentiate subtle differences in out-of-distribution samples, and this approach can be utilized for iterative RLHF optimization.

• The paper introduces ensemble voting to measure preference strength and identify noisy data, enabling effective label flipping and improved model performance.
• It employs an adaptive margin and label smoothing in the loss function to stabilize training and prevent overfitting on strong preference data.
• Contrastive learning and meta-learning techniques are applied to enhance feature differentiation and maintain robust performance on out-of-distribution data.

Secrets of RLHF in LLMs Part II: Reward Modeling

Introduction

Reinforcement Learning from Human Feedback (RLHF) is pivotal in enhancing AI systems by aligning their behavior with human preferences. This alignment is typically facilitated through reward models that serve as proxies for human intent, guiding reinforcement learning optimization. However, RLHF faces challenges such as incorrect, ambiguous preference pairs and limited generalization across diverse data distributions. This paper presents methods to address these challenges by refining reward models, using techniques like preference strength measurement and contrastive learning.

Preference Strength Measurement

The paper introduces a mechanism to measure preference strength based on ensemble voting from multiple reward models. This metric helps distinguish data into incorrect, ambiguous, and strong preference categories, allowing optimization strategies that mitigate the impact of noisy data.

Figure 1: Mean and standard deviation of preference differences derived from 10 reward models for all paired data, revealing potential incorrect preferences and low model consistency.

When analyzing preference strength, substantial incorrect or ambiguous data can substantially degrade model performance. The research suggests flipping labels for these subsets to improve validation outcomes significantly, showcasing this technique as robust for noise mitigation.

Figure 2: Label flipping for incorrect preferences significantly improves model performance, indicating these metrics efficiently identify and leverage useful information from incorrect data.

Adaptive Margin and Label Smoothing

To enhance reward model efficiency, an adaptive margin component based on preference strength is incorporated into the loss function. This enhancement aligns with findings indicating that diverse preference data types require tailored handling strategies.

Figure 3: Adding an adaptive margin component enhances model performance by accommodating different preference strengths.

Label smoothing further contributes by penalizing overconfidence, thus stabilizing learning on strong preference data. It prevents overfitting and encourages the model to derive general, robust features, beneficial in broad application contexts.

Contrastive Learning for Reward Modeling

Contrastive learning approaches, such as SimCSE and SwAV, are adapted to reward modeling to improve feature extraction. These methods focus on the inherent feature similarity challenge between chosen and rejected responses, utilizing contrastive losses to refine model differentiation capabilities.

Figure 4: t-SNE feature distribution illustrates reduced overlap between chosen and rejected responses with SimCSE-enhanced reward modeling.

MetaRM: Meta-Learning for Generalization

MetaRM, leveraging meta-learning principles, aligns reward modeling with domain shifts introduced during policy training. By optimizing reward models against the environment's shifted distribution, this approach maintains robust differentiation capabilities in varied data distributions.

Figure 5: MetaRM enhances the reward model's discriminative power under new distributions, showcasing effective alignment strategies using existing preference pairs.

Evaluation on OOD Data

The research evaluates models on out-of-distribution (OOD) data, substantiating that the strategies not only improve performance within training domains but also ensure effectiveness in new contexts without additional preference annotations. The results demonstrate substantial improvements over baseline models, advocating for the method's broad applicability.

Figure 6: Experimental results on OOD data confirm the superior effectiveness of the method compared to baseline models.

Conclusion

This paper presents advanced strategies for refining reward models within RLHF frameworks, addressing significant challenges related to preference data noise and generalization. By leveraging innovative techniques such as label flipping, adaptive margins, contrastive learning, and meta-learning, the research illustrates substantial improvements in model alignment and stability across diverse data contexts. Future developments may further explore cross-domain application potential, enhancing RLHF's role in AI alignment.