mlsanityData to 7 Oct 2026

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

arXiv 2605.24749 ↗First posted 23 May 2026stat.ML cs.LGcited by 0

On the map: LLM alignment / RLHF / preference optimisation → Reward modeling and RLHF

Abstract (arXiv)

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 θ^*, xrangle) and x ~ N(0, Id). 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 ey/β2 and a more practical surrogate-weighted fit with weights e^ra0(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.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation. Builds on 26 of its 34 references that are in the map. First posted on arXiv 23 May 2026.

Builds on

26 of 34 references.

Closest ancestors · 6

Same direction, newest first

  1. 2025Is a Good Foundation Necessary for Efficient Reinforcement Learning? The Computational Role of the Base Model in Exploration10 Mar
  2. 2024Optimal Design for Reward Modeling in RLHF22 Oct
  3. 2024Reinforcement Learning from Human Feedback with Active Queries14 Feb
  4. 2023Principled Reinforcement Learning with Human Feedback from Pairwise or K-wise Comparisons26 Jan
  5. 2022Scaling Laws for Reward Model Overoptimization19 Oct
  6. 2017Deep reinforcement learning from human preferences12 Jun

Other references · 20

  1. 2022Training language models to follow instructions with human feedbackcited by 15,712
  2. 2019The Curious Case of Neural Text Degenerationcited by 2,881
  3. 2020Learning to summarize from human feedbackcited by 2,562
  4. 2019Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learningcited by 750
  5. 2022High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representationcited by 146
  6. 2023SGD learning on neural networks: leap complexity and saddle-to-saddle dynamicscited by 142
  7. 2020Online stochastic gradient descent on non-convex losses from high-dimensional inferencecited by 119
  8. 2023Aligning Language Models with Preferences through f-divergence Minimizationcited by 118
  9. 2022Learning Single-Index Models with Shallow Neural Networkscited by 92
  10. 2022On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgettingcited by 80
  11. 2019Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networkscited by 79
  12. 2024Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimizationcited by 78
8 more
  1. 2023Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Modelscited by 60
  2. 2022Follow the Wisdom of the Crowd: Effective Text Generation via Minimum Bayes Risk Decodingcited by 54
  3. 2024Neural network learns low-dimensional polynomials with SGD near the information-theoretic limitcited by 45
  4. 2024Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functionscited by 35
  5. 2024Nonlinear spiked covariance matrices and signal propagation in deep neural networkscited by 25
  6. 2023Beyond NTK with Vanilla Gradient Descent: A Mean-Field Analysis of Neural Networks with Polynomial Width, Samples, and Timecited by 20
  7. 2024Generating Diverse and High-Quality Texts by Minimum Bayes Risk Decodingcited by 11
  8. 2025From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGDcited by 2

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026Reward Models Are Secretly Value Functions: Temporally Coherent Reward Modeling
  2. 2026Are we really tilting? The mechanics of reward guidance in flow and diffusion models
  3. 2025Information-Theoretic Reward Modeling for Stable RLHF: Detecting and Mitigating Reward Hacking
  4. 2026Fine-Tuning on Self-Generated and Reward-Weighted Data: Learning Dynamics, Convergence Rates, and Benefits of Off-Policyness
  5. 2026Binary Rewards and Reinforcement Learning: Fundamental Challenges
  6. 2025Reward Models in Deep Reinforcement Learning: A Survey
  7. 2025Learning a Pessimistic Reward Model in RLHF
  8. 2024Confronting Reward Overoptimization for Diffusion Models: A Perspective of Inductive and Primacy Biases