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

arXiv: https://arxiv.org/abs/2605.24749 · HTML: https://mlsanity.com/paper/2605.24749/

## Key facts

- No citing papers in the map yet (data through 14 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.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)

## Builds on (most cited first, 26 of 26)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Principled Reinforcement Learning with Human Feedback from Pairwise or K-wise Comparisons](https://mlsanity.com/paper/2301.11270/) (2023)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [Aligning Language Models with Preferences through f-divergence Minimization](https://mlsanity.com/paper/2302.08215/) (2023)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting](https://mlsanity.com/paper/2206.00761/) (2022)
- [Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks](https://mlsanity.com/paper/1910.01619/) (2019)
- [Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimization](https://mlsanity.com/paper/2407.13399/) (2024)
- [Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models](https://mlsanity.com/paper/2305.10633/) (2023)
- [Follow the Wisdom of the Crowd: Effective Text Generation via Minimum Bayes Risk Decoding](https://mlsanity.com/paper/2211.07634/) (2022)
- [Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit](https://mlsanity.com/paper/2406.01581/) (2024)
- [Reinforcement Learning from Human Feedback with Active Queries](https://mlsanity.com/paper/2402.09401/) (2024)
- [Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions](https://mlsanity.com/paper/2405.15459/) (2024)
- [Is a Good Foundation Necessary for Efficient Reinforcement Learning? The Computational Role of the Base Model in Exploration](https://mlsanity.com/paper/2503.07453/) (2025)
- [Nonlinear spiked covariance matrices and signal propagation in deep neural networks](https://mlsanity.com/paper/2402.10127/) (2024)
- [Optimal Design for Reward Modeling in RLHF](https://mlsanity.com/paper/2410.17055/) (2024)
- [Beyond NTK with Vanilla Gradient Descent: A Mean-Field Analysis of Neural Networks with Polynomial Width, Samples, and Time](https://mlsanity.com/paper/2306.16361/) (2023)
- [Generating Diverse and High-Quality Texts by Minimum Bayes Risk Decoding](https://mlsanity.com/paper/2401.05054/) (2024)
- [From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD](https://mlsanity.com/paper/2510.21020/) (2025)

## 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.