Stabilizing RLHF through Advantage Model and Selective Rehearsal

TL;DR

This work tackles RLHF instability in LLM alignment by introducing two innovations: Advantage Model (AM) to balance reward score distributions and reduce reward hacking, and Selective Rehearsal to prevent catastrophic forgetting by curating PPO training data and rehearsing SFT knowledge. AM directly models the advantage of a response over expected rewards and enforces calibrated, bounded scores, leading to improved calibration and stable training. Selective Rehearsal uses representative clustering (via SimCSE) and AM-guided ranking to select high-skill data for rehearsal, boosting performance while maintaining expert-aligned capabilities. Experiments on English and Chinese data, including public HH-RLHF and proprietary datasets, show AM improves stability and rewards; adding selective rehearsal further increases win rates and mitigates forgetting, demonstrating practical efficacy for robust RLHF in diverse settings.

Abstract

Large Language Models (LLMs) have revolutionized natural language processing, yet aligning these models with human values and preferences using RLHF remains a significant challenge. This challenge is characterized by various instabilities, such as reward hacking and catastrophic forgetting. In this technical report, we propose two innovations to stabilize RLHF training: 1) Advantage Model, which directly models advantage score i.e., extra reward compared to the expected rewards and regulates score distributions across tasks to prevent reward hacking. 2) Selective Rehearsal, which mitigates catastrophic forgetting by strategically selecting data for PPO training and knowledge rehearsing. Our experimental analysis on public and proprietary datasets reveals that the proposed methods not only increase stability in RLHF training but also achieve higher reward scores and win rates.

Figures (6)

• Figure 1: Left: The distribution of reward scores for both the QA and Code Generation tasks. There is a noticeable disparity in the learned reward score distributions between the two tasks, despite the expectation that the distributions should be similar. Right: The win/loss rate over the SFT model on the forget set exhibits a significant decline. This drop in the win rate can be attributed to reward hacking and the phenomenon of catastrophic forgetting.
• Figure 2: Ranking accuracy is shown as a function of the difference in scores between higher and lower ranked responses. The orange lines indicate the calibrated prediction of accuracy $1 / (1 + e^{-\Delta})$ in which $\Delta$ denotes the score difference. On the left, we show calibration of RM and AM on HH-RLHF data while on the right we show results for our proprietary data. We observe that AM calibration is better than RM's.
• Figure 3: Distributions of RM and AM scores for pairs of good and bad examples from the proprietary data.
• Figure 4: Mean and standard variance for each task categorized by a task spectrum on the in-house data.
• Figure 5: PPO training curves on the Main Test Set with different scoring models. RM-PPO and AM-PPO denote PPO trained with Reward Model and Advantage Model, respectively. AM-PPO-SER additionally equips with Selective Rehearsal.