Uncertainty-aware Reward Model: Teaching Reward Models to Know What is Unknown

2024-10-01

arXiv.org

AI Summary

Existing reward models (RMs) struggle to capture both aleatoric uncertainty—arising from inherent stochasticity in human preferences—and epistemic uncertainty—stemming from model’s limited knowledge and capacity—while lacking mechanisms to quantify reward prediction reliability. This work introduces the first unified two-level framework for jointly modeling both uncertainties, proposing a single-model variant (URM) and an ensemble-based variant (URME). URM employs a probabilistic value head and disentangled preference attribute modeling; URME further leverages inter-model discrepancy analysis for fine-grained uncertainty estimation and automatic identification of unreliable reward predictions. Evaluated on RewardBench, our methods significantly outperform state-of-the-art RMs. When integrated with Best-of-N sampling and iterative DPO/PPO optimization, they substantially enhance LLM generation quality, robustness to distributional shifts (e.g., across domains or prompts), and decision reliability—where low-uncertainty predictions consistently correlate with higher output quality and stronger generalization.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Multiagent Systems: Multiagent Systems under Uncertainty

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Reward models (RMs) are essential for aligning large language models (LLM) with human expectations. However, existing RMs struggle to capture the stochastic and uncertain nature of human preferences and fail to assess the reliability of reward predictions. To address these challenges, we introduce the Uncertainty-aware Reward Model (URM) and its ensemble variant, URME. URM employs a probabilistic value head to capture aleatoric uncertainty by modeling the distribution of disentangled human preference attributes. URME further quantifies epistemic uncertainty by examining discrepancies among individual URMs within the ensemble, enabling identification of unreliable evaluations. Our empirical evaluations demonstrate that URM achieves strong performance on RewardBench, outperforming competitive large-scale models. Additionally, extensive experiments, including best-of-n sampling (BoN), iterative direct preference optimization (iterative DPO), and proximal policy optimization (PPO), demonstrate that URM and URME significantly enhance LLMs' generation quality. Notably, reward predictions with lower uncertainty are far more reliable, demonstrate significantly higher quality, and result in substantially improved alignment.

Problem

• Improving reward models for human preferences
• Quantifying uncertainty in reward predictions
• Enhancing large language models alignment

Innovation

• Uncertainty-aware Reward Model
• Probabilistic value head
• Ensemble variant URME