Source-linked AI summary
Uncertainty Estimation for Language Reward Models
Adam Gleave, Geoffrey Irving
TL;DR
Human preference comparisons can train reward models, but collecting them remains expensive, and single point estimates do not capture model uncertainty that could support sample efficiency and RL robustness. The paper fine-tunes an ensemble of reward models from one pre-trained model, but ensemble active learning does not outperform random sampling and the uncertainty estimates explain at most 13% of model-error variance. The authors connect this weakness to limited ensemble diversity and suggest modifying pre-training or developing methods tolerant of that limitation.
Problem
Method
The authors fine-tune an ensemble of reward models initialized from one pre-trained model, with each member’s final layer randomly initialized.
Source: §3.2 · p. 4Results
Ensemble active learning provided no benefit over random sampling, and oracle-labeler uncertainty explained at most 13% of model-error variance.
Source: §5 · p. 8Takeaways & Limitations
The authors suggest modifying foundation-model training to include uncertainty estimation or developing methods that tolerate limited ensemble diversity.
Source: §5 · p. 8Fine-tuning from one foundation model can leave ensemble members insufficiently diverse, yielding poor-quality uncertainty estimates.
Source: §5 · p. 8
Abstract
from arXiv · showhide
Language models can learn a range of capabilities from unsupervised training on text corpora. However, to solve a particular problem (such as text summarization) it is typically necessary to fine-tune them on a task-specific dataset. It is often easier for humans to choose between options than to provide labeled data, and prior work has achieved state-of-the-art performance by training a reward model from such preference comparisons. However, collecting a large preference comparison dataset is still expensive -- and the learned reward models are unreliable out-of-distribution. We seek to address these problems via uncertainty estimation, which can improve sample efficiency and robustness using active learning and risk-averse reinforcement learning (RL). Specifically, we use bootstrap aggregating (bagging) to train an ensemble of reward models differing in the initialization of their final layer. Ensembles have proved successful in prior applications of active learning, but we find that in our setting ensemble active learning does not outperform random sampling. Further experiments show that while the aggregate predictions are well-calibrated, the ensemble's estimated epistemic uncertainty is only weakly correlated with model error. We suspect this is because the ensemble members are fine-tuned from a single model and so are similar to one another. This suggests current pre-training methods will need to be modified to support uncertainty estimation, e.g. by training multiple language models.
1 Introduction
Human preference comparisons offer an alternative to costly reference writing, while uncertainty estimates could support more sample-efficient learning and safer reward-model use. The paper finds that ensemble active learning does not outperform random sampling and that ensemble uncertainty is weakly related to error.
1 Introduction
- Human judgments of output quality are difficult to specify, as summaries must balance truthfulness, succinctness, and reader priorities.The challenge is even greater for open-ended conversational tasks.Source: §1 · p. 1
- Supervised learning commonly mimics human-written reference text [36] [31] [22], but writing examples is slow and expensive; existing data can be reused, though often low quality or an imperfect match for the task.An alternative learns a reward model from human feedback and uses reinforcement learning to fine-tune a language model; it has been applied to stylistic continuation [53] and document summarization.Source: §1 · p. 1
- Uncertainty estimates could improve sample efficiency through active learning, strengthen RL robustness, and support quality control.Active learning could reduce collection costs and improve reliability, while uncertainty could support risk-averse or risk-constrained RL.Source: §1 · pp. 1–2
- Ensemble active learning does not outperform random sampling, and on a synthetic dataset higher-confidence points are only slightly more accurate.The authors report that ensemble uncertainty is only weakly correlated with model error.Source: §1 · p. 2
- The authors speculate that fine-tuning ensemble members from one pre-trained model limits their diversity and weakens uncertainty estimates.They suggest training multiple distinct smaller pre-trained models or introducing uncertainty into pre-training, such as with dropout [12].Source: §1 · p. 2
2 Related Work
Prior work studies uncertainty estimation, active learning, preference-based reward learning, and fine-tuning foundation models. This paper addresses the gap at their intersection: actively selecting human preference feedback while fine-tuning language models.
2 Related Work
- Prior work on uncertainty in fine-tuned language models had focused, to the authors’ knowledge, on calibrating total uncertainty [10] [19] [21] [20].Active learning applications often depend on separately estimating aleatoric and epistemic uncertainty.Source: §2 · p. 2
- Ensembles estimate model uncertainty through disagreement [23], alongside hypermodels [11] and Bayesian neural networks [4] [12].Active learning may select maximally epistemically uncertain points, but other methods select low-confidence, expected-change, or low-parameter-variance points [24].Source: §2 · pp. 2–3
- Ensemble-based active learning improved sample efficiency in around three quarters of the continuous-control and Atari tasks tested by Christiano et al. [9].More sophisticated approaches have been demonstrated only in low-dimensional environments.Source: §2 · p. 3
- Foundation-model fine-tuning uses far fewer labels than training from scratch, but preference datasets can still require substantial human effort.Prior preference-based work includes completions preferred to the original model 88% of the time after 5,000 comparisons [53] and summaries outperforming human references after fewer than 100,000 comparisons. The authors identify actively selecting human feedback during language-model fine-tuning as an unfilled gap.Source: §2 · p. 3
3 Uncertainty Estimation
The method learns a distribution over reward models from human preference comparisons and a pre-trained language model. It approximates that distribution with an ensemble and uses ensemble-based strategies to select comparisons.
3.1 Problem Formulation
- The problem is to learn a distribution over reward models from preference comparisons between document completions, using a next-token pre-trained language model as the base.The feedback dataset is small relative to unsupervised corpora, motivating initialization from and fine-tuning of the pre-trained model.Source: §3.1 · p. 3
3.2 Ensemble Models
- Each ensemble reward model copies the pre-trained weights, replaces the final layer with a randomly initialized reward-output layer, and fine-tunes the full network on preference comparisons.Members differ through final-layer initialization and bootstrap sampling; each data point receives an independently sampled weight of 2 or 0 with equal probability, reused across epochs.Source: §3.2 · p. 4
3.3 Active Learning
- For active learning, the method either samples an ensemble member and selects its highest-reward comparison or selects the comparison with maximal ensemble probability variance.The variance between members’ probabilities that the first summary is preferred estimates epistemic uncertainty.Source: §3.3 · p. 4
4 Experiments
The experiments find no significant active-learning improvement over random sampling on real human preference data, while an oracle-labeler evaluation shows that estimated epistemic uncertainty is only weakly related to model error. The results also indicate that larger ensembles improve this relationship modestly, whereas bootstrapping does not yield a significant improvement.
4.1 Experimental Setup
- The evaluation used TL;DR preference comparisons for training, validation, and testing, plus CNN/DM comparisons as an out-of-distribution test set.The TL;DR splits contained 92,858 training, 33,083 validation, and 50,719 test comparisons; CNN/DM had 2,284 comparisons.Source: §4.1 · p. 4
4.2 Active Learning
- The active-learning study selected 4,096 labels from the TL;DR training pool, and its offline pool-based setup may underestimate gains from synthesizing novel summary pairs.The setup evaluated a 417M-parameter model and allowed selection from the unlabeled dataset.Source: §4.2 · p. 5
- Ensemble active learning did not significantly improve over random sampling on real human preference data.Uncertainty sampling, Thompson sampling, and variance-based selection also failed to outperform random sampling.Source: §4 · p. 4, §4.2 · p. 5, Figure 2 · p. 6
4.3 Oracle Labeler
- The oracle-labeler experiment used known oracle label distributions to assess how well ensemble disagreement tracked prediction error.It quantified prediction error with KL divergence and compared it with estimated epistemic uncertainty using Spearman correlation.Source: §4.3 · p. 6
- Epistemic uncertainty correlated only weakly with model error, reaching at most r = 0.36 in the oracle-labeler evaluation.The evaluation measured model error against the oracle distribution and estimated epistemic uncertainty from disagreement between ensemble members.Source: §4.3 · pp. 6, 8, Figure 3 · p. 7
- Correlation rose from r = 0.25 with 3 ensemble members to r = 0.36 with 42 members.The increase was modest, and correlation remained low overall.Source: §4.3 · p. 8
- Bootstrapping slightly improved mean uncertainty quality, but the difference was not statistically significant.Ensemble size mattered more; the reported confidence intervals for bootstrapped and non-bootstrapped results overlapped.Source: §4.3 · p. 8
5 Discussion
Ensembles provide no active-learning benefit in this summarization setting, and their uncertainty estimates explain at most 13% of model-error variance. The discussion identifies task scope and limited ensemble diversity as important boundaries, while noting untested approaches.
5 Discussion
- Ensemble active learning provides no benefit here, and uncertainty estimates explain at most 13% of model-error variance.The estimates are better than random but remain low quality when evaluated against an oracle labeler.Source: §5 · p. 8
- The evaluation covers only summary-quality discrimination, a narrow task where random sampling may be a strong baseline.Although the dataset spans topics such as business, sport, and politics, the task is narrower than open-ended dialogue.Source: §5 · p. 8
- Other uncertainty-estimation methods remain untested, including linear hypermodels; ensembles of thousands of models are only an extrapolated possibility.Fine-tuning only biases is another proposed way to fit more ensemble members into limited memory.Source: §5 · p. 8
- Staying close to a foundation model’s prior is associated with less ensemble diversity and poorer uncertainty estimates, motivating training changes or methods tolerant of that lack.The proposed directions include training ensembles of foundation models.Source: §5 · p. 8
A.1 Comparison to Prior Work
The language and reward models replicate prior work, and implementation differences appear not to change overall language-model perplexity or reward-model accuracy. The authors expect the results to apply to other pretrained models and preference-comparison implementations.
A.1 Comparison to Prior Work
- The language and reward models successfully replicate prior work.Source: §A.1 · p. 11
- Minor implementation differences do not seem to significantly change overall language-model perplexity or reward-model accuracy.Source: §A.1 · p. 11
- The authors expect these results to apply to other pretrained language models and preference-comparison algorithm implementations.Source: §A.1 · p. 11
A.2 Pre-trained Language Models
The experiments use small Gopher-family models trained on C4, with an architecture similar to GPT-2 but different normalization and positional encodings. Their pretrained-model performance is broadly similar to comparably sized GPT-2 models, with differences on WikiText 103 and LAMBADA.
A.2 Pre-trained Language Models
- The experiments use small Gopher-family models trained on the Colossal Clean Crawled Corpus (C4).Source: §A.2 · p. 11
- The architecture resembles GPT-2 but uses RMSNorm instead of LayerNorm and relative instead of absolute positional encodings.Source: §A.2 · p. 11
- LAMBADA accuracy is 8–13% lower than reported for GPT-2, while perplexity is lower (better) on WikiText 103.The authors attribute much of the LAMBADA difference to GPT-2’s stop-word filter and note that GPT-2 was evaluated on raw LAMBADA, unlike their standard-dataset evaluation.Source: §A.2 · p. 11
A.2.1 Reward Models
Reward models trained on the full human-feedback dataset achieve accuracy comparable to prior models, though a direct comparison is unavailable because the reported data subsets cannot be reproduced. The 1347M model’s accuracy falls between the cited lowest and highest prior figures.
A.2.1 Reward Models
- Reward models trained on the full Stiennon et al. human-feedback dataset achieve accuracy comparable to their models.Source: §A.2.1 · p. 11
- The reported comparison is not head-to-head because the subsets used for Stiennon et al.’s labeler-agreement results could not be reproduced.The authors compare accuracy across a range of dataset slices instead.Source: §A.2.1 · p. 11, Table 2 · p. 12
- 64.84% ± 1.63% accuracy for the 1347M model is above the reported 57.4% ± 2.0% low and below the reported 67.5% ± 1.1% high for Stiennon et al.’s 1.3B model.The values are not a direct head-to-head comparison because the reported labeler-agreement subsets could not be reproduced.Source: Table 2 · p. 12
A.3 Model Calibration
The models are generally well-calibrated, with calibration improving as model and training-dataset size increase.
A.3 Model Calibration
- The models’ predicted probabilities of summary preference are generally well-calibrated.The calibration curves lie close to the ideal y = x line.Source: §A.3 · p. 12
- Calibration improves with larger model parameter counts and more training data.Source: §A.3 · p. 12
- Calibration is assessed by comparing predicted probabilities with observed accuracy.The figure plots predicted probability on the x-axis and accuracy on the y-axis.Source: §A.3 · p. 12
A.4 Active Learning
No active-learning method significantly outperforms random sampling; the section also examines calibration and how ensemble uncertainty relates to model error.
A.4 Active Learning
- No active-learning method significantly outperformed random sampling across 1,024- and 4,096-point training sets and TL;DR Valid, Test, and CNN/DM evaluations.Source: §A.4 · p. 12, Figure 5 · p. 14
A.5 Oracle Labeler
- Preference predictions are well-calibrated, with calibration improving as dataset and model size increase.Figure 4 compares confidence-bin accuracy with predicted confidence across TL;DR datasets and out-of-distribution CNN/DM.Source: Figure 4 · p. 13
- The oracle-labeler analysis compares model error, measured by KL divergence, with ensemble-variance uncertainty using Spearman correlation across three test datasets.Figures cover TL;DR Validation, TL;DR Test, and CNN/DM Test, varying ensemble size with and without bootstrapping.Source: Figure 6 · p. 15, Figure 7 · p. 16, Figure 8 · p. 17