📊 statistics

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

This paper analyzes a two-stage neural reward model within a Gaussian single-index framework to demonstrate how exponential reward weighting influences feature recovery in the first layer and establishes explicit temperature-dependent bounds on policy value gaps, thereby identifying an admissible range of deployment temperatures that balances optimization gains against learning costs.

Original authors: Rei Higuchi, Ryotaro Kawata, Akifumi Wachi, Shokichi Takakura, Kohei Miyaguchi, Taiji Suzuki

Published 2026-05-26

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Rei Higuchi, Ryotaro Kawata, Akifumi Wachi, Shokichi Takakura, Kohei Miyaguchi, Taiji Suzuki

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Mapmaker" and the "Explorer"

Imagine you are training an AI to be a great explorer. To do this, you need two things:

A Map (The Reward Model): A system that tells the AI which paths are good (high reward) and which are bad.
The Explorer (The Policy): The AI itself, which uses the map to decide where to go.

The problem this paper studies is a tricky feedback loop. Usually, we train a mapmaker on a static set of data (like a textbook). But in AI alignment (like making chatbots helpful), the mapmaker isn't just read; it's used to change the explorer's behavior. The explorer starts taking the "best" paths found on the map. This means the mapmaker is suddenly being judged on a new, different set of paths—the ones the explorer actually chooses.

The paper asks: If the mapmaker is a complex neural network (a "black box" with many layers), how does it learn to find the right "features" (the important details) when the data it's being tested on keeps shifting because the explorer is changing its mind?

The Setup: A Simple World (The Gaussian Single-Index Model)

To understand this complex problem, the authors create a simplified, controlled world.

The Terrain: Imagine a giant, smooth hill where the height represents the "reward."
The Secret: There is one specific direction (a hidden vector $\theta^*$ ) that points straight up the steepest part of the hill. If you know this direction, you know how to get the highest reward.
The Learner: A neural network is trying to figure out this secret direction.

The authors break the learning process into two distinct stages, like a two-step training camp.

Stage 1: Finding the Compass (Feature Recovery)

The Challenge:
At the start, the neural network's neurons are pointing in random directions, like a compass spinning wildly. The network needs to "lock on" to the secret direction ( $\theta^*$ ) to be useful.

The Twist (Exponential Weighting):
In standard training, the network looks at all data points equally. But in this "reward modeling" scenario, the training process is biased. It cares much more about the data points that have high rewards. It's like a student who only studies the questions they think will be on the test, ignoring the rest.

The Discovery:
The paper shows that this "bias toward high rewards" actually helps the network find the secret direction faster, but only if you tune a "temperature" knob ( $\beta_1$ ) correctly.

Too Cold (Low Temperature): The network gets obsessed with just a few "perfect" examples. It overfits and gets confused by noise, failing to learn the general direction.
Too Hot (High Temperature): The network ignores the high-reward examples and treats everything the same, losing the advantage of the reward signal.
Just Right: The authors prove there is a "Goldilocks" zone. If the temperature is above a certain constant level (independent of how huge the data is), the network's neurons will successfully align with the secret direction.

The Analogy:
Think of the neurons as a group of hikers trying to find the summit. The "reward weighting" is like a loudspeaker shouting, "Look at the top of the mountain!"

If the speaker is too quiet, the hikers wander aimlessly.
If the speaker is too loud and shrill, the hikers panic and only look at the very peak, missing the path up.
The paper proves that if the speaker is loud enough but not screaming, the hikers will successfully figure out which way is "up."

Stage 2: Drawing the Map (Policy Optimization)

The Challenge:
Once the network has found the secret direction (the compass is locked), it needs to draw the actual map. It does this by fitting a curve to the data. However, the final goal isn't just to draw a perfect map; it's to draw a map that, when used by the explorer, leads to the best possible outcome.

The Twist (Deployment Temperature):
The explorer uses a "temperature" knob ( $\beta_2$ ) to decide how aggressively to follow the map.

High $\beta_2$ : The explorer is cautious and explores many paths.
Low $\beta_2$ : The explorer is greedy and only takes the single "best" path.

The Discovery:
The paper analyzes the "Value Gap"—the difference between the reward the explorer could get with a perfect map and the reward it actually gets with the learned map.

They found two ways to fit the map:

Label-Weighted (Ideal): Using the true reward values to weight the data. This is the theoretical best case.
Surrogate-Weighted (Practical): Using a predicted reward (a guess) to weight the data. This is what happens in real life.

The Result:
The paper provides a formula for the "Value Gap." It shows that the gap is made of three parts:

Temperature Mismatch: How different the explorer's greediness is from the mapmaker's training.
Truncation: Errors caused by ignoring extreme, impossible paths (a safety measure).
Learning Error: How bad the map is.

The Catch with Surrogates:
When using the "Surrogate" method (the practical one), if the initial guess is wrong, the errors get amplified. It's like trying to navigate using a map drawn by someone who is also lost. The paper shows that this amplification depends heavily on the temperature. If you get too greedy (too low a temperature) with a bad map, the explorer gets stuck in a local trap, and the performance drops significantly.

The Main Takeaways

Reward Modeling is Different: You can't just treat reward modeling like a standard math problem. Because the reward model changes the data distribution (the explorer changes where it goes), you have to account for that shift.
Temperature is a Control Knob: There is a specific "temperature" setting for the training phase that ensures the neural network actually learns the underlying features (the secret direction) rather than just memorizing noise. This setting doesn't need to be impossibly high; a moderate, constant level is enough.
The "Proxy" Problem: If you use a rough guess to train your reward model (surrogate weighting), you are more fragile. You have to be careful not to be too greedy (too low a temperature) during deployment, or the errors in your guess will blow up and ruin the explorer's performance.
Balancing Act: The paper gives a mathematical recipe for choosing the right temperature. You want to be greedy enough to get high rewards, but not so greedy that you amplify the mistakes in your map.

Summary in One Sentence

This paper proves that for neural networks to successfully learn the "secret direction" of a reward landscape and then use it to guide an AI, the training process must carefully balance a "temperature" setting to avoid over-focusing on noise, and the deployment strategy must be cautious enough to prevent small errors in the map from causing the AI to crash.

Problem Statement

Reward modeling (RM) is a critical component in aligning large language models (e.g., via RLHF), yet it differs fundamentally from standard supervised prediction. In KL-regularized policy optimization, the learned reward function $r(x)$ is exponentiated to define the deployed policy $\pi_\beta(dx) \propto e^{r(x)/\beta}\pi_{ref}(dx)$ . Consequently, the downstream value of the system depends not merely on prediction accuracy under the reference distribution $\pi_{ref}$ , but on errors in regions where the reward is high and the policy concentrates. This creates a policy-induced distribution shift.

Existing theoretical analyses of offline alignment often assume fixed, known linear features, thereby bypassing the challenge of how neural networks discover relevant features (feature learning) and approximate nonlinear reward landscapes under such distribution shifts. This paper addresses the gap by analyzing how a two-layer neural network learns features and fits a nonlinear reward function when the training objective involves exponential reward weighting.

Methodology

The authors formulate the problem within a Gaussian Single-Index Model (SIM) setting:

Data Generation: Inputs $x \sim \mathcal{N}(0, I_d)$ , and labels $y = r^*(x) + \zeta$ , where the true reward is $r^*(x) = \sigma^*(\langle \theta^*, x \rangle)$ . Here, $\theta^*$ is a hidden target direction, $\sigma^*$ is an upper-bounded polynomial link function, and $\zeta$ is bounded noise.
Model Architecture: A two-layer neural network with $N$ neurons: $r_a(x) = \frac{1}{N} \sum_{j=1}^N a_j \sigma(\langle w_j, x \rangle + b_j)$ .
Two-Stage Training Protocol:
1. Feature Recovery (Stage 1): The first-layer directions $w_j$ are trained via online spherical SGD using a reward-weighted objective with temperature $\beta_1$ . The second-layer coefficients $a$ are fixed at initialization. The goal is to align $w_j$ with the hidden direction $\theta^*$ .
2. Readout Fitting (Stage 2): Once features are recovered, the directions $w_j$ $w_{j}$ and biases $b_j$ $b_{j}$ are frozen. The second-layer coefficients $a$ $a$ are fitted using weighted ridge regression on fresh samples. The weighting scheme uses a deployment temperature $\beta_2$ $β_{2}$ , which may differ from $\beta_1$ $β_{1}$ . Two weighting schemes are analyzed:
  - Label-weighted: Weights are $e^{y/\beta_2}$ .
  - Surrogate-weighted: Weights are $e^{r_{proxy}(x)/\beta_2}$ , where $r_{proxy}$ is a frozen proxy reward (e.g., from the feature recovery phase).

Key Contributions and Results

1. Feature Learning under Reward-Induced Distribution Shift
The paper characterizes the first-layer feature recovery under exponential reward weighting.

Signal Transformation: Exponential weighting acts as a nonlinear label transformation that alters the Hermite signal available to the first layer.
Exponent Reduction: The analysis shows that for any feature-learning temperature $\beta_1$ above a dimension-free $O(1)$ threshold, the effective learning bottleneck reduces from the Information Exponent (IE) to the Generative Exponent (GE).
Recovery Guarantees: Under upper-bounded polynomial assumptions, a constant fraction of neurons recover the hidden direction $\theta^*$ with high probability. The weak-recovery complexity is governed by the generative exponent and scales as $\tilde{O}(\beta_1^{2(I^*-1)} d^{(p_{gen}-1)\vee 1})$ , where $I^*$ is the power required to reach the GE.
Significance of $\beta_1$ : Crucially, the required lower bound on $\beta_1$ is independent of the dimension $d$ . This improves upon prior clipped-exponential analyses that required dimension-growing polylogarithmic temperatures, which would otherwise nullify the benefits of exponent reduction.

2. From Weighted Prediction to Tilted-Policy Regret
The authors move beyond standard prediction error bounds to evaluate the tilted-policy value gap:
$\mathcal{R} = \mathbb{E}_{x \sim \pi_{\beta^*}}[r^*(x)] - \mathbb{E}_{x \sim \hat{\pi}_{\beta_2}}[r^*(x)]$
where $\pi_{\beta^*}$ is the target policy and $\hat{\pi}_{\beta_2}$ is the policy induced by the learned reward.

Decomposition: The regret is decomposed into three terms:
1. Temperature Mismatch ( $T_{temp}$ ): Depends on the difference between the target temperature $\beta^*$ and deployment temperature $\beta_2$ .
2. Projected Truncation ( $T_{cut}$ ): Arises from restricting the analysis to a projected subspace to ensure integrability.
3. Learning Term ( $T_{learn}$ ): Depends on the weighted prediction error of the fitted reward.
Low-Temperature Behavior: The paper derives explicit bounds showing that the variance of the reward under the tilted policy scales as $\beta_2^2 / (2p_{max})$ , where $p_{max}$ is the local order of the reward maximum.

3. Label- and Surrogate-Weighted Regret Bounds
The paper provides explicit value-gap bounds for both fitting schemes:

Label-Weighted Bound: Serves as a theoretical benchmark. The bound separates temperature mismatch, truncation, and learning costs, linking policy value rates to the local shape of the nonlinear reward maximum.
Surrogate-Weighted Bound: More practical for plug-in pipelines. This bound introduces an additional factor $C_{0,R}(\beta_2)$ $C_{0, R} (β_{2})$ , which represents the density ratio between the surrogate-induced measure and the reference distribution.
- Proxy Mismatch Mechanism: When the proxy reward is imperfect, $C_{0,R}(\beta_2)$ can be large (scaling polynomially with $1/\beta_2$ ). This amplifies both approximation errors and finite-sample errors, identifying a statistical mechanism behind "proxy mismatch" and reward overoptimization.
Admissible Temperatures: The analysis yields an admissible set of deployment temperatures $\beta_2$ . Choosing a smaller $\beta_2$ improves the match to the target policy but amplifies learning costs via exponential weighting. The optimal choice balances these trade-offs.

Significance and Claims

The paper claims to provide a foundational characterization of the interaction between reward modeling and policy value within a controlled theoretical setting.

Feature Learning: It demonstrates that reward-induced exponential weighting can facilitate feature learning by reducing the effective learning exponent from IE to GE, provided the temperature $\beta_1$ is not too small. It establishes that a dimension-free $O(1)$ temperature is sufficient for recovery, correcting previous assumptions that required dimension-dependent temperatures.
Evaluation Metric: It argues that the relevant statistical target for reward modeling is the tilted-policy value, not reference-distribution prediction error. The derived bounds explicitly link the value gap to the local geometry of the reward landscape and the choice of deployment temperature.
Surrogate Risks: It identifies a specific statistical mechanism where using a proxy reward for weighting (surrogate weighting) can amplify errors if the proxy is imperfect, offering a theoretical explanation for reward overoptimization phenomena observed in practice.

The authors maintain that while the Gaussian SIM is a theoretical abstraction, it isolates the core mechanisms of feature learning and policy-induced distribution shifts, providing a basis for future work on richer input distributions and multi-index structures.