Source-linked AI summary
Efficient Exploration at Scale
Seyed Mohammad Asghari, Chris Chute, Vikranth Dwaracherla, Xiuyuan Lu, Mehdi Jafarnia, Victor Minden, Zheng Wen, Benjamin Van Roy
TL;DR
The paper addresses limited understanding of RLHF scaling with preference data and develops an online algorithm that updates reward and language models as choices arrive. It combines affirmative nudges, epistemic reward uncertainty, and information-directed exploration, matching offline RLHF trained on 200K choices with fewer than 20K and projecting a 1,000x gain at 1M choices.
Problem
Systematic understanding of how overall RLHF performance scales with preference-data quantity remains elusive, motivating study of better-scaling RLHF techniques.
Method
The algorithm incrementally updates reward and language models from choice data, adding an affirmative nudge and using epistemic reward uncertainty with information-directed exploration.
Results
Fewer than 20K choices match offline RLHF trained on 200K choices, representing more than a 10x data-efficiency gain; 1M choices are projected to yield a 1,000x gain.
Takeaways & Limitations
Online learning and uncertainty-guided exploration offer enormous gains in RLHF data efficiency, according to the authors.
Takeaways & Limitations
The authors state that the potential 1,000x gain leaves substantial room for improvement through further research, including improved exploration algorithms.
Abstract
from arXiv · showhide
We develop an online learning algorithm that dramatically improves the data efficiency of reinforcement learning from human feedback (RLHF). Our algorithm incrementally updates reward and language models as choice data is received. The reward model is fit to the choice data, while the language model is updated by a variation of reinforce, with reinforcement signals provided by the reward model. Several features enable the efficiency gains: a small affirmative nudge added to each reinforcement signal, an epistemic neural network that models reward uncertainty, and information-directed exploration. With Gemma large language models (LLMs), our algorithm matches the performance of offline RLHF trained on 200K labels using fewer than 20K labels, representing more than a 10x gain in data efficiency. Extrapolating from our results, we expect our algorithm trained on 1M labels to match offline RLHF trained on 1B labels. This represents a 1,000x gain. To our knowledge, these are the first results to demonstrate that such large improvements are possible.
1. Introduction
The paper develops an online RLHF algorithm that updates reward and language models as human choices arrive, using several exploration mechanisms to improve data efficiency. It reports matching offline RLHF trained on 200K choices with fewer than 20K and projects a 1,000x gain at 1M choices.
- The algorithm incrementally fits the reward model to observed choices and updates the language model with a reinforce variant using reward-model signals.
- Three innovations support data efficiency: an affirmative reinforcement-signal nudge, an epistemic neural network for reward uncertainty, and information-directed exploration.
- Fewer than 20K choices match offline RLHF trained on 200K choices, representing more than a 10x data-efficiency gain.
- A projected 1,000x gain at 1M choices is reported as evidence that very large data-efficiency improvements are possible with LLMs.
2. Literature Review
The literature review identifies an unresolved question about how RLHF performance scales with preference data and presents online, active-exploration, and scaling-law work as relevant context. It argues that logarithmic plots can reveal scaling-law differences obscured by linear axes.
- Online Adaptation: Online methods sample responses on-policy, while offline methods use a fixed sampling distribution and face data-coverage and stationary-target challenges.
- Active Exploration: Active exploration strategies use uncertainty or informativeness to select examples, while related LLM alignment work formalizes preference collection as an active contextual dueling bandit.
- Scaling Laws: Systematic understanding of overall RLHF performance as preference-data quantity changes remains elusive.
- Scaling Laws: The review asks whether current RLHF techniques improve with more data or require techniques with better scaling properties.
- Scaling Laws: Logarithmic rather than linear data axes reveal qualitative differences between RLHF scaling laws.
3. Experiment Pipeline
The experiment pipeline gathers simulated preference choices between sampled responses, updates policies from batched feedback, and evaluates each policy against a baseline using out-of-sample win rate. It uses Gemma policies and a Gemini-based feedback simulator across diverse prompts.
- Baseline and Experimentation Policies: A 9B Gemma model provides the baseline, while top-K policies restrict token sampling to the K highest-probability tokens.
- Baseline and Experimentation Policies: Top-1 is the deterministic baseline, whereas stochastic top-5 policies diversify candidate responses for informative human choices.
- Feedback Simulation: The simulator converts two response rewards into a Bradley-Terry preference probability and samples the choice from a Bernoulli distribution.
- Gathering Feedback and Training: Each prompt produces two responses and one choice; after batches of 64, the algorithm can adjust parameters using the feedback.
- Performance Evaluation: Performance is the average simulator preference probability over 1K out-of-sample prompts comparing competing and baseline top-1 policies.
4. Algorithms
The study compares offline, periodic, online, and information-directed RLHF algorithms that differ in how they gather data, update models, and use reward uncertainty. Across methods, reward models and policies begin from shared parameterized components with broadly comparable sizes and update rules.
- Algorithm Alternatives: Offline RLHF gathers all data from the initial policy before fitting a reward model and optimizing the policy.
- Algorithm Alternatives: Periodic RLHF alternates data collection and offline-style updates over fixed periods, while online RLHF updates reward and policy parameters after each batch.
- Model Architecture: Reward models share a transformer backbone whose output is mapped to a scalar reward through a head, with an ensemble of MLP heads used for information-directed sampling.
- Common Elements: The algorithms update both reward-model and policy parameters, while using approximately similar parameter counts and similar update rules.
4.2. Update Rules
The update rules fit the reward model to observed choices and update the language model with a regularized policy-gradient variant of reinforce. Reward-model and policy gradients are aggregated, clipped, and applied iteratively.
- The reward model predicts which of two responses a rater will choose for a given prompt.
- Reward-model gradients from batches of prompts are clipped and used with AdamW to obtain updated reward parameters.
- The policy update maintains an exponential moving average of parameters that serves as an anchor for regularization.
- The policy gradient is a reinforce variant using the reward-model choice probability minus 1/2 as its reinforcement signal, with anchor regularization weighted by β.
- Policy gradients are summed over prompts and response pairs, clipped, and applied with AdamW to update policy parameters.
4.3. Alternatives
The alternatives compare offline, periodic, and online RLHF, with online learning incrementally updating models to avoid the computational burden of repeated retraining. Its strong performance depends on a reward model, an affirmative nudge, and uncertainty-guided information-directed sampling.
- Periodic RLHF: Periodic RLHF alternates policy updates and data collection using progressively updated policies, with period τ = 400 batches in the experiments.
- Online RLHF: Incrementally updating reward and language models avoids training new models and policies from scratch as often as periodic RLHF.
- Online RLHF: Reward-model-free policy updates improve over offline RLHF but are not competitive with online RLHF, which uses a reward model.
- Online RLHF: Prior online RLHF methods can tank after some batches, whereas an affirmative nudge added to every reinforcement signal avoids tanking without reducing the learning rate.
- Online RLHF: The online algorithm samples sixteen responses per prompt, queries a random pair, and uses four response pairs for policy-gradient updates.
- Information-Directed Exploration: Information-directed sampling adds a small-parameter reward-model head to estimate uncertainty and guide response selection for human-feedback queries.
Architecture
The reward model is an epistemic neural network that augments a point estimate with indexed ensemble components to represent reward uncertainty, while adding under 5% to the original model’s parameters.
- The architecture uses an MLP point-estimate head plus ensembles of 100 prior networks and 100 differential networks.
- The ENN reward model takes prompt X, response Y, and epistemic index Z, producing reward r_ϕ(Y|X,Z).
- For Z = 0, inference uses the point-estimate reward model; for Z > 0, the corresponding prior and differential networks are added.
- The ENN increases parameters by well under 5% relative to the original nine-billion-parameter model.
Queries
For each prompt, the algorithm samples response pairs and queries human feedback on the pair whose ensemble choice probability has the greatest variance.
- For each batch prompt, the policy samples sixteen responses and forms response pairs for evaluation.
- It selects the pair maximizing this variance and sends it to the human feedback simulator.
- Choice-probability variance measures informativeness, motivating selection of pairs expected to produce informative human choices.
Training
Training updates the point-estimate reward model and each differential network using choice-data losses, while leaving prior networks unchanged in the point-estimate update.
- The point-estimate reward model is updated from batches of 64 queries and choices using the online RLHF loss on r(Y|X,0).
- Prior networks remain untouched during the point-estimate update.
- Each differential network is then updated individually using losses computed from r(Y|X,Z) for Z = 1, ..., 100.
5. Results
Information-directed exploration improves RLHF performance with substantially less choice data than offline RLHF, with larger gains projected at higher data volumes. Examples illustrate that the method can elicit more informative comparisons and produce better responses.
- 5.1. Win Rates: Offline RLHF needs more than 200K choices to match the performance achieved by information-directed exploration with 20K choices.
- 5.1. Win Rates: Over 10x data efficiency is reported for information-directed exploration relative to offline RLHF.
- 5.2. Projected Gains: Projected data-efficiency gains increase with choice count, reaching about 1,000x after one million choices for the online algorithm.
- Qualitative Examples: In a math example, the offline RLHF response is convoluted and wrong, while the information-directed response is concise, coherent, and correct.
- Informative Feedback: The algorithm selects infomax response pairs by maximizing choice-probability variance and contrasts them with infomin pairs that minimize it.
- Informative Feedback: The informativeness of a pair changes as labels are collected and the uncertainty model is trained; the examples use initial ENN parameters θ0.
- Informative Feedback: Infomin pairs make equivalent statements, whereas infomax pairs provide substantially different reasoning forms and therefore elicit informative choices.
6. Conclusions and Future Work
The paper reports large data-efficiency gains from online learning and uncertainty-guided exploration, enabled in part by an affirmative nudge. It identifies several extensions, including improved exploration, prompt selection, multiturn dialogue, agents, and AI-assisted feedback.
- Online learning and uncertainty-guided exploration each offer enormous gains in data efficiency.
- An affirmative nudge was introduced to overcome tanking behavior commonly observed during RLHF.
- Improving the exploration algorithm: Improving exploration could involve deeper reward-model uncertainty, language-model uncertainty, and more efficient reward-informed optimization.
- Selecting prompts: The response-selection algorithm could be extended to select prompts expected to yield informative feedback.
- Multiturn dialog and agents: The approach could be extended to multiturn conversations and agents by incorporating value models for anticipated or delayed rewards.
- AI assisted feedback: The methods could support AI-assisted feedback that frames comparisons to elicit more informative human judgments.