Forecast how reward model changes affect rlhf before training begins
Forecasting How RLHF Responds to Reward-Model Edits Before Training
Machine Learning
Summary
When training AI with reinforcement learning using human feedback (RLHF), small changes to how the AI's reward system scores responses can change what the AI learns to produce. The authors found a way to predict how an AI's behavior will change after altering the reward model without running the full training process. They use the idea that reward models act like a set of prices, and they measure how different features in the AI’s responses relate to changes in these prices. Their prediction method explains a large part of the changes seen after training across different setups and reward models.
What this means in practice
- •For machine learning engineers: Predict how changing a reward model will affect a policy’s output before expensive reinforcement learning training.
- •For ai safety teams: Rank reward model edits quickly to identify which changes improve AI behavior in early training steps without full retraining.
- •For online content moderation teams: Adjust reward models used to train content filters and forecast their impact on filter behavior before deployment.$Commercial implications: Enables development of predictive tools to tune moderation models efficiently, improving product reliability and reducing time-to-market.
Authors
Zixiao Wang, Duzhen Zhang, Xingbo Du, Dongxia Wu, Le Song
Abstract
Responses with headers, lists and greater length receive high scores from reward models (RMs), and RLHF policies learn to produce these features even when true quality does not improve. Practitioners therefore edit their RMs, but they learn whether an edit works only after an RL run. We show that the response of a policy to an RM edit can be forecast before training, from samples of the current policy and the scores of the old and new RM. If we read the RM as a set of prices, the first-order theory of KL-regularized RL implies that each feature moves in proportion to its within-prompt covariance with the change in reward. We implant controlled price changes into retrained RMs under a 2^(4-1) fractional factorial design with a self-distilled placebo, run 54 RL runs from six parent policies, and fix every forecast before the runs. With the covariance predictor, we predict the direction of all 20 changes in supply whose forecast has the same sign at every parent and explain 82% of the standardized change. With counterfactual prices, the classical alternative, we explain 50%. The forecast also holds under GRPO and for real RM upgrades, where we explain 66% out of sample with a coefficient frozen from the factorial runs. For rent, the gap between the proxy reward and a gold reward from stronger judges, we rank RM edits correctly with the same covariance over the first 30 steps (Spearman 0.94). This measures the horizon of the forecast: by step 200, the ranking of the factorial edits reverses under both RLOO and GRPO, while the real upgrades keep their order. When we vary the KL coefficient beta, the initial response does not depend on beta, and the response saturates below the KL equilibrium. In one experiment, the long-run response to a price rise is consistent with an additional implicit regularizer beta_0 of about 0.003 to 0.007.