Ask a Strong LLM Judge when Your Reward Model is Uncertain

TL;DR

This work addresses the challenge that reward models (RMs) trained from human preferences often generalize poorly to out-of-distribution inputs, while strong LLM judges, though more accurate, incur prohibitive inference costs for online RLHF. It introduces an uncertainty-based routing framework that uses a spectral-normalized Gaussian process-based pairwise PM (SNGP-PM) to quantify uncertainty in reward differences and route only uncertain pairs to a costly LLM judge, while confident pairs are evaluated by the fast PM RM. The approach provides a principled blend of PM and LLM judgments to construct advantage estimators for policy gradient updates (via RLOO/GRPO) and demonstrates improved out-of-distribution generalization on RewardBench and RM-Bench, as well as better downstream alignment on multiple instruction-following benchmarks with fewer judge calls. Overall, the framework offers a practical path to leverage strong LLM reasoning in online RLHF by smartly trading off accuracy and computational cost, guided by calibrated uncertainty estimates. The uncertainty routing hinges on key equations such as p(x,y1,y2) = g(h)/u(x,y1,y2) with u(x,y1,y2) = sqrt(1 + lambda * phi(h)^T Sigma phi(h)) and uses these to determine when to consult the judge and how to form the final policy update.

Abstract

Reward model (RM) plays a pivotal role in reinforcement learning with human feedback (RLHF) for aligning large language models (LLMs). However, classical RMs trained on human preferences are vulnerable to reward hacking and generalize poorly to out-of-distribution (OOD) inputs. By contrast, strong LLM judges equipped with reasoning capabilities demonstrate superior generalization, even without additional training, but incur significantly higher inference costs, limiting their applicability in online RLHF. In this work, we propose an uncertainty-based routing framework that efficiently complements a fast RM with a strong but costly LLM judge. Our approach formulates advantage estimation in policy gradient (PG) methods as pairwise preference classification, enabling principled uncertainty quantification to guide routing. Uncertain pairs are forwarded to the LLM judge, while confident ones are evaluated by the RM. Experiments on RM benchmarks demonstrate that our uncertainty-based routing strategy significantly outperforms random judge calling at the same cost, and downstream alignment results showcase its effectiveness in improving online RLHF.

Figures (3)

• Figure 1: Left: Average accuracy of pairwise RM at uncertainty quantiles on RewardBench (RB) and RM-Bench (RM). Each scatter point corresponds to 10% of the data corresponding to a quantile. The accuracy shows a negative Spearman's rank correlation with uncertainty, with p-values less than $10^{-29}$ and $10^{-136}$ on RewardBench and RM-Bench, respectively, suggesting that uncertain pairs are more likely to get wrong rankings. Middle: Uncertainty gap between ID (HelpSteer2-Preference, train and validation sets) and OOD datasets (RewardBench and RM-Bench), where the uncertainty scores given by our RM range from 1. Uncertainty scores are averaged within 10 bins divided by reward difference. The marker size indicates the portion of data contained within the bin, and the shaded area represents the range within a standard deviation. Right: Overall averaged uncertainty scores. The uncertainty scores are consistently higher on OOD data.
• Figure 2: Uncertainty-based routing outperforms random routing with the same number of LLM judge calls on preference accuracy and downstream alignment. See \ref{['fig:acc-improvement-remaining']} for other benchmarks.
• Figure 3: Uncertainty-based routing outperforms random routing with the same number of strong LLM judge calls on preference accuracy (RM-Bench) and downstream alignment (Arena-Hard-v0.1).

Theorems & Definitions (2)

• Remark 1
• Remark 2