Uncertainty-Aware Reward Modeling for Stable RLHF
For practitioners of RLHF, this addresses the critical problem of reward hacking by making reward models uncertainty-aware, leading to more stable and reliable alignment.
RLHF faces reward hacking due to deterministic reward models and uniform advantage weighting in group-based policy optimization. UARM uses conformal prediction for uncertainty estimation and heteroscedastic variance decomposition to reweight advantages, improving calibration and alignment across three benchmarks.
Reinforcement learning from human feedback (RLHF) aligns large language models by training reward models on preference data and optimizing policies to maximize predicted rewards. However, this pipeline faces two fundamental challenges: (1) reward models cannot signal when their predictions are unreliable, since they usually act as deterministic point estimators; and (2) modern group-based policy optimization can amplify unreliable reward signals, as exemplified by GRPO's uniform treatment of rewards during advantage computation. As policies explore increasingly diverse responses, these two limitations create a critical vulnerability: unreliable reward estimates may be granted disproportionate influence, triggering severe reward hacking. We propose Uncertainty-Aware Reward Modeling (UARM), which equips reward models with calibrated uncertainty via quantile-based conformal prediction and reweights GRPO advantages through heteroscedastic variance decomposition. Experiments across HelpSteer, UltraFeedback, and PKU-SafeRLHF demonstrate that UARM significantly improves reward model calibration, reduces reward hacking, and enhances downstream alignment quality compared to standard GRPO and uncertainty-agnostic baselines.