Uncertainty-aware reward models: how sure is the score?

One number cannot say whether the model is confidently right or guessing.

A scalar reward model gives

For an answer like those in the training data, 0.82 may be accurate. For an odd answer unlike anything the model was trained on, 0.82 may be an accident of extrapolation. Both look the same. An uncertainty-aware reward model reports two things: a reward, and how uncertain that reward is.

In RLHF the policy maximises the learned reward \(r_\phi\) while a KL penalty keeps it near a reference policy:

The learned reward is a proxy for what we actually want, the gold reward. Early in training, raising the proxy also raises the gold reward. Later the policy finds answers the proxy overrates, and the gold reward falls while the proxy keeps rising. This is Goodhart's law applied to reward models, and it is called reward over-optimisation or reward hacking.

Gao et al. measured this in a controlled setting: a large reward model played the role of the gold labeller, and smaller proxy reward models were trained on its labels. They plotted gold reward against distance from the initial policy,

and found that simple curves fit well for best-of-\(n\) sampling and for RL:

Both rise, peak and fall. The fitted coefficients changed smoothly with the size of the proxy reward model, with larger proxies over-optimised less. The lesson for this chapter: the proxy is least reliable exactly where the policy pushes hardest, far from the data it was trained on. That is where we most need to know how uncertain it is.

The simplest way to measure uncertainty is to train several reward models and see whether they agree.

Members that were trained differently agree where the data pin the answer down and drift apart where they do not. A large \(\sigma\) therefore signals low confidence.

Four reward models score answer P:

High reward, high confidence. The same four models score answer Q:

The mean of Q is middling, but the members range from 0.22 to 0.95. One model loves Q and another rejects it. A system that saw only 0.595 would not know to be cautious.

An ensemble is only useful if its members make different mistakes. Ways to make them differ, roughly from cheapest to most diverse:

• Different fine-tuning seeds. Same pretrained backbone, different head initialisation and data order.
• Bootstrapped data. Each member trains on a resample of the preference pairs.
• Shared backbone, several heads. One Transformer with \(K\) linear heads. Almost free at inference, but the heads share every feature, so they disagree little.
• Different pretraining seeds. Members start from separately pretrained backbones. Most expensive, and the most diverse.

A related idea averages the weights rather than the outputs. Fine-tune \(K\) reward models from the same pretrained start, then use

This is WARM, weight-averaged reward models (Ramé et al.). The result is a single model with single-model inference cost that was reported to be more robust to distribution shift and label noise than its members. It gives a better point estimate, but no \(\sigma\): a single model cannot disagree with itself.

The policy needs one number. The ensemble gives \(K\). Besides the mean of Equation (5), three rules are common.

Coste et al. compare mean, worst-case and a variance-penalised rule of the form (12), which they call uncertainty-weighted optimisation, for RLHF against an ensemble. A further option keeps \(\mu\) as the reward but down-weights high-\(\sigma\) samples in the policy update. Applied to answers P and Q of Section 4, with \(\lambda=1\):

The rules barely touch P and cut Q sharply. Because \(\sigma<1\) here, \(\sigma^2<\sigma\) and the variance penalty is the mildest of the three at the same \(\lambda\).

Using \(R_\lambda\) in place of \(\mu\) makes the reward pessimistic where the ensemble is unsure. Compare an answer with \(\mu=0.90,\ \sigma=0.20\) and one with \(\mu=0.75,\ \sigma=0.02\), at \(\lambda=1\):

By the mean the first answer wins; by the conservative reward the second does. An uncertain high reward is discounted. If the policy finds a strange answer that one member loves and the others do not, the disagreement raises \(\sigma\) and the exploit stops paying. This targets reward hacking, out-of-distribution exploitation and overconfident reward-model mistakes, the failures that produce bad RL updates.

Disagreement detects errors the members do not share. Model member \(k\)'s reward as the true reward plus a shared bias \(b\) plus its own error \(\varepsilon_k\), with the \(\varepsilon_k\) independent with variance \(s^2\):

Averaging shrinks the independent part, but \(b\) passes straight into \(\mu\), and \(\sigma\) does not depend on \(b\) at all. If every member has learned that longer answers are better, a policy that pads its answers gets a high \(\mu\) and a small \(\sigma\), and no penalty on \(\sigma\) will stop it.

Eisenstein et al. found exactly this pattern: reward-model ensembles reduced reward hacking but did not remove it, because members tended to share error patterns. Ensembles whose members differed in pretraining seed generalised better than ensembles differing only in fine-tuning seed, since they shared less.

Ensemble disagreement measures the epistemic kind. For a comparison label \(\ell\in\{0,1\}\), let member \(k\) predict \(p_k=P_k(y_2\succ y_1)\) and pick a member uniformly at random. The law of total variance splits the label's variance into two parts:

If annotators split 50/50 on a pair, a well-trained ensemble learns \(p_k\approx0.5\) in every member: the aleatoric part is 0.25 and the epistemic part is 0. The ensemble is confident, correctly, that people disagree. If instead the members predict \(p_k=[0.9,0.1,0.9,0.1]\), the mean is again 0.5, but the split is 0.09 aleatoric and 0.16 epistemic: the models do not know. Only the second case should trigger the penalty of Section 7. Disagreement between people is the subject of the previous chapter on conditional reward models.

Ensembles are the most common estimate, not the only one.

• MC dropout. Keep dropout on at inference and run the same model \(M\) times with different random masks \(m_j\). The spread of the outputs is a cheap stand-in for an ensemble (Gal and Ghahramani).

• Predicting a distribution. Give the head two outputs, a mean \(\mu_\theta\) and a standard deviation \(\sigma_\theta\), and treat each reward as Gaussian. The preference probability then accounts for both answers' spreads:

where \(\Phi\) is the standard normal CDF. A single model trained this way mostly learns aleatoric spread from the labels; it has no built-in way to know what it has not seen.

Uses of \(\sigma\) beyond the conservative reward:

• Active labelling. Send the pairs with the highest disagreement to human annotators first.
• Out-of-distribution flags. Route high-\(\sigma\) answers to a stronger judge or a human.
• Best-of-\(n\) selection. Pick the candidate with the highest \(\mu-\lambda\sigma\) rather than the highest \(\mu\).

• Cost. Independent members multiply training and inference compute by \(K\).
• Diversity is hard to get. Members fine-tuned from one backbone agree too much, so \(\sigma\) understates the real uncertainty.
• Shared errors are invisible. A bias common to all members enters \(\mu\) and leaves \(\sigma\) unchanged (Equation 14).
• \(\sigma\) is not calibrated. Nothing guarantees that \(\sigma=0.1\) means the same thing on different prompts.
• Choosing \(\lambda\). Too small and hacking continues; too large and the policy avoids anything novel, including genuinely better answers.
• Epistemic only. Disagreement does not measure genuine disagreement between people.

Use the lab's default scores. Members score A as \([0.66,0.63,0.67,0.64]\), B as \([0.99,0.55,0.97,0.69]\) and C as \([0.79,0.66,0.81,0.62]\). Rank the answers by the mean, by \(\mu-\sigma\), and by the worst case.

Ensembles make a learned reward more cautious, but every member is still a learned model with the same blind spots. The next chapter combines learned rewards with rules that cannot be fooled in the same way: exact-match checks, unit tests and format constraints.