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

• URL: http://arxiv.org/abs/2605.24749v1
• Date: Sat, 23 May 2026 22:00:38 GMT
• Title: How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis
• Abstract summary: We study the feedback in a Gaussian single-index model with $r*(x) = *(langle *, xrangle)$ and $x sim N(0, I_d)$.<n>We analyze a two-stage neural reward model that first learns the hidden direction $*$ from reward-weighted samples and then fits the readout layer by weighted ridge regression.
• Score: 53.063298916923976
• License: http://creativecommons.org/licenses/by/4.0/
• 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) = σ^*(\langle θ^*, x\rangle)$ and $x \sim N(0, I_d)$. We analyze a two-stage neural reward model that first learns the hidden direction $θ^*$ 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 $β_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/β_2}$ and a more practical surrogate-weighted fit with weights $e^{r_{a_0}(x)/β_2}$. Keeping the $β_2$-dependence explicit yields an admissible set of deployment temperatures, balancing the gain from lowering $β_2$ against the learning cost amplified by exponential weighting; in the surrogate-weighted case, proxy-dependent factors shrink this admissible set.

Related papers

• When Does $\ell_2$-Boosting Overfit Benignly? High-Dimensional Risk Asymptotics and the $\ell_1$ Implicit Bias [15.113649527486276]
  We show that benign overfitting fails at the linear rate.<n>We show that while this localization mechanism should persist in the presence of signals, the exact signal-noise decomposition remains an open problem.
  arXiv  Detail & Related papers  (2026-05-07T14:14:09Z)
• Post-Training with Policy Gradients: Optimality and the Base Model Barrier [27.674563695368665]
  We study post-training linear autoregressive models with outcome and process rewards.<n>We prove that variants of policy gradient (PG) can achieve likelihood $1 - varepsilon$ with an essentially minimax optimal number of reward queries.
  arXiv  Detail & Related papers  (2026-03-07T00:25:53Z)
• Optimal Unconstrained Self-Distillation in Ridge Regression: Strict Improvements, Precise Asymptotics, and One-Shot Tuning [61.07540493350384]
  Self-distillation (SD) is the process of retraining a student on a mixture of ground-truth and the teacher's own predictions.<n>We show that for any prediction risk, the optimally mixed student improves upon the ridge teacher for every regularization level.<n>We propose a consistent one-shot tuning method to estimate $star$ without grid search, sample splitting, or refitting.
  arXiv  Detail & Related papers  (2026-02-19T17:21:15Z)
• Amortized Simulation-Based Inference in Generalized Bayes via Neural Posterior Estimation [1.096028999747108]
  Generalized Bayesian Inference (GBI) tempers a loss with a temperature $>0$ to mitigate overconfidence and improve under model misspecification.<n>We give the first fully amortized variational approximation to the tempered posterior family $p_(mid x) propto (),p(x mid )$ by training a single $(x,)$-conditioned neural posterior estimator $q_(mid x,)$ that enables sampling in a single forward pass.
  arXiv  Detail & Related papers  (2026-01-29T22:20:47Z)
• Retire: Robust Expectile Regression in High Dimensions [3.9391041278203978]
  Penalized quantile and expectile regression methods offer useful tools to detect heteroscedasticity in high-dimensional data. We propose and study (penalized) robust expectile regression (retire) We show that the proposed procedure can be efficiently solved by a semismooth Newton coordinate descent algorithm.
  arXiv  Detail & Related papers  (2022-12-11T18:03:12Z)
• Restricted Strong Convexity of Deep Learning Models with Smooth Activations [31.003601717265006]
  We study the problem of optimization of deep learning models with smooth activation functions. We introduce a new analysis of optimization based on Restricted Strong Convexity (RSC) Ours is the first result on establishing geometric convergence of GD based on RSC for deep learning models.
  arXiv  Detail & Related papers  (2022-09-29T21:24:26Z)
• Nonparametric regression with modified ReLU networks [77.34726150561087]
  We consider regression estimation with modified ReLU neural networks in which network weight matrices are first modified by a function $alpha$ before being multiplied by input vectors.
  arXiv  Detail & Related papers  (2022-07-17T21:46:06Z)
• Bounding the Width of Neural Networks via Coupled Initialization -- A Worst Case Analysis [121.9821494461427]
  We show how to significantly reduce the number of neurons required for two-layer ReLU networks. We also prove new lower bounds that improve upon prior work, and that under certain assumptions, are best possible.
  arXiv  Detail & Related papers  (2022-06-26T06:51:31Z)
• High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation [89.21686761957383]
  We study the first gradient descent step on the first-layer parameters $boldsymbolW$ in a two-layer network. Our results demonstrate that even one step can lead to a considerable advantage over random features.
  arXiv  Detail & Related papers  (2022-05-03T12:09:59Z)
• Breaking the Sample Complexity Barrier to Regret-Optimal Model-Free Reinforcement Learning [52.76230802067506]
  A novel model-free algorithm is proposed to minimize regret in episodic reinforcement learning. The proposed algorithm employs an em early-settled reference update rule, with the aid of two Q-learning sequences. The design principle of our early-settled variance reduction method might be of independent interest to other RL settings.
  arXiv  Detail & Related papers  (2021-10-09T21:13:48Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.