How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis

arXiv:2605.24749v1 Announce Type: cross Abstract: Reward modeling is not only a prediction problem: in KL-regularized policy optimization, the learned reward is exponentiated to define the deployed policy, so downstream value depends on errors in reward-tilted regions. We study this feedback in a Gaussian single-index model with r^(x) = sigma^(langle heta^, xrangle) and x sim N(0, I_d). We analyze a two-stage neural reward model that first learns the hidden direction heta^ from reward-weighted samples and then fits the readout layer by weighted ridge regression. Exponential reward weighting changes the Hermite signal available to the first layer; for any feature-learning temperature eta_1 above a dimension-free O(1) threshold, a constant fraction of neurons recover the hidden direction, with weak-recovery complexity governed by the generative exponent. After feature recovery, we derive tilted-policy value-gap bounds for an idealized label-weighted fit with weights e^{y/eta_2} and a more practical surrogate-weighted fit with weights e^{r_{a_0}(x)/eta_2}. Keeping the eta_2-dependence explicit yields an admissible set of deployment temperatures, balancing the gain from lowering eta_2 against the learning cost amplified by exponential weighting; in the surrogate-weighted case, proxy-dependent factors shrink this admissible set.

Source: arXiv cs.LG | 2026-05-26