Uncertainty Quantification for Large Language Model Reward Learning under Heterogeneous Human Feedback
Pangpang Liu, Junwei Lu, Will Wei Sun
TL;DR
This work tackles uncertainty in reward learning for LLM alignment under heterogeneous human feedback. It introduces a heterogeneous rationality model that jointly learns the reward and human rationality via an alternating gradient descent, and proves convergence with an asymptotic distribution, enabling confidence intervals for rewards and rationality. The paper further develops hypothesis tests for reward differences and a pessimistic BoN policy that uses lower confidence bounds to robustly select actions, with sublinear regret guarantees. Extensive simulations and real LLM experiments demonstrate accurate uncertainty quantification, statistically sound comparisons between models, and practical gains from uncertainty-aware BoN strategies in policy selection. Overall, the framework provides statistically principled tools for reliable reward modeling and decision-making in RLHF pipelines with diverse human feedback.
Abstract
We study estimation and statistical inference for reward models used in aligning large language models (LLMs). A key component of LLM alignment is reinforcement learning from human feedback (RLHF), where humans compare pairs of model-generated answers and their preferences are used to train a reward model. However, human feedback is inherently heterogeneous, creating significant challenges for reliable reward learning. To address this, we adopt a heterogeneous preference framework that jointly models the latent reward of answers and human rationality. This leads to a challenging biconvex optimization problem, which we solve via an alternating gradient descent algorithm. We establish theoretical guarantees for the resulting estimator, including its convergence and asymptotic distribution. These results enable the construction of confidence intervals for reward estimates. Leveraging these uncertainty quantification results, we conduct valid statistical comparisons between rewards and incorporate uncertainty into the best-of-$N$ (BoN) policy framework. Extensive simulations demonstrate the effectiveness of our method, and applications to real LLM data highlight the practical value of accounting for uncertainty in reward modeling for LLM alignment.