PaperlayerView exact revision on arXiv

Source-linked AI summary

Real-Time Aligned Reward Model beyond Semantics

Zixuan Huang, Xin Xia, Yuxi Ren, Jianbin Zheng, Xuefeng Xiao, Hongyan Xie, Li Huaqiu, Songshi Liang, Zhongxiang Dai, Fuzhen Zhuang, Jianxin Li, Yikun Ban, Deqing Wang

arXiv:2601.22664v4cs.AI

TL;DR

RLHF reward overoptimization arises when policies exploit spurious reward patterns and a fixed reward model fails to track policy distribution shifts. R2M incorporates evolving policy hidden states as feedback, and experiments report improved dialogue and summarization performance with minimal computational cost.

  • Problem

    A fixed reward model can misalign with an evolving policy, allowing reward overoptimization through spurious patterns rather than faithful capture of human intent.

  • Method

    R2M dynamically incorporates evolving policy hidden states into the reward model to adapt to policy distribution shifts without additional labeled data or environmental feedback.

  • Results

    R2M improves AlpacaEval 2 win rate by 5.2% - 8.0%, length-controlled win rate by 2.9% - 6.1%, and TL;DR win rate by 6.3% over baselines.

  • Takeaways & Limitations

    R2M provides a lightweight way to strengthen reward models and mitigate reward overoptimization during RLHF.

  • Takeaways & Limitations

    The paper states that R2M maintains ethical and societal implications consistent with standard RLHF practices and identifies no specific additional societal risks requiring special consideration.

Abstract

from arXiv · show

Reinforcement Learning from Human Feedback (RLHF) is a pivotal technique for aligning large language models (LLMs) with human preferences, yet it is susceptible to reward overoptimization, in which policy models overfit to the reward model, exploit spurious reward patterns instead of faithfully capturing human intent. Prior mitigations primarily relies on surface semantic information and fails to efficiently address the misalignment between the reward model (RM) and the policy model caused by continuous policy distribution shifts. This inevitably leads to an increasing reward discrepancy, exacerbating reward overoptimization. To address these limitations, we introduce R2M (Real-Time Aligned Reward Model), a novel lightweight RLHF framework. R2M goes beyond vanilla reward models that solely depend on the semantic representations of a pretrained LLM. Instead, it leverages the evolving hidden states of the policy (namely policy feedback) to align with the real-time distribution shift of the policy during the RL process. This work points to a promising new direction for improving the performance of reward models through real-time utilization of feedback from policy models.

1. Introduction

RLHF reward overoptimization is worsened when a fixed reward model cannot track the policy’s changing distribution. R2M addresses this by incorporating policy hidden-state feedback into a lightweight, dynamically adapting reward model.

  • Motivation: Reward overoptimization occurs when policy models exploit spurious patterns, including response length, formatting, n-grams, or emojis, instead of human intent.A fixed reward model’s approximation errors grow as the policy distribution shifts during RLHF.
  • Prior limitations: Iteratively retraining the reward model could adapt to policy behavior, but direct retraining at every iteration is computationally prohibitive.
  • R2M: R2M reinforces the reward model iteratively by dynamically adapting to the policy’s internal states without additional labeled data or environmental feedback.
  • R2M: R2M redesigns the reward-model scoring head to integrate evolving policy hidden states and aggregate policy feedback without retraining the underlying model.
  • Results: 5.2% - 8.0% AlpacaEval 2 win rate, 2.9% - 6.1% length-controlled win rate, and 6.3% TL;DR win rate gains are reported over vanilla RLOO baselines.The improvements occur on dialogue and summarization tasks with minimal computational cost.

2. Preliminary

The RLHF optimization stage samples multiple responses, scores them with a scalar reward model, estimates within-group advantages using RLOO, and optimizes the policy with importance sampling. R2M is designed to integrate into this REINFORCE-based workflow without changing the underlying RL algorithm.

  • RLHF workflow: RLHF comprises supervised fine-tuning, reward modeling, and RL optimization, with R2M directly integrated into the third stage.
  • Trajectory Sampling: At each training step, the old policy samples K responses for each query, producing query-response groups for optimization.
  • Trajectory Sampling: A scalar reward model scores each query-response pair, and RLOO estimates advantages within each sampled response group.
  • Policy Optimization: The policy performs a forward pass for each query-response pair and is optimized with importance sampling using an objective with ε and β hyperparameters.
  • R2M integration: R2M can be deployed across REINFORCE-based RLHF frameworks while retaining RLOO as a primary baseline.

3. Motivation

Deep-layer policy hidden states correlate with human preference labels and reward-score differences, motivating their use as policy feedback for reward-model alignment. Theorem 3.1 further states that R2M provides a tighter reward-misalignment bound than a vanilla reward model when γ(t) > 0.

  • Motivation: Deep-layer policy hidden states contain information closely correlated with both human preferences and reward scores, supporting their use to enhance vanilla reward models.
  • Hidden states and preferences: Different-preference pairs have lower average hidden-state similarity than same-preference pairs, with the gap widening at greater layer depth.
  • Hidden states and rewards: Higher hidden-state similarity corresponds to smaller absolute reward-score differences, indicating a strong negative correlation between the two quantities.
  • Motivation: The correlation between policy hidden states and reward scores suggests potential alignment between the policy model and reward model.
  • Theoretical motivation: When γ(t) > 0, R2M yields a tighter upper bound on reward misalignment ϵ than the vanilla reward model.

4. Method

R2M incorporates evolving policy hidden states into reward prediction and iteratively optimizes the reward model without retraining its LLM component. Its design combines cross-attention, time-dependent representation weighting, and GREBT optimization to adapt reward allocation to policy distribution shifts.

  • Reward Model Structure: R2M adds a policy-feedback data flow that integrates policy hidden states with the reward model’s semantic representation before scalar reward prediction.The framework collects policy last-layer hidden states and aggregates them with the reward model’s Reward Token Embedding.
  • Reward Model Structure: Sequence-to-Token Cross Attention extracts information from policy hidden-state sequences and injects it into the reward token representation.The cross-attention module uses the reward token as a query and policy hidden states as keys and values.
  • Reward Model Structure: Time-Step-Based Weighted Combination balances the original reward-token embedding with policy feedback, gradually increasing feedback influence as training proceeds.Early training prioritizes the existing embedding, while later updates progressively emphasize policy-feedback representations to adapt to distribution shift.
  • Iterative Reward Model Lightweight Optimization: R2M refreshes policy hidden states during policy updates and uses them in an additional lightweight reward-model optimization phase.The reward model is updated from real-time policy feedback without retraining the reward model’s LLM component.
  • Iterative Reward Model Lightweight Optimization: GREBT constructs preference pairs from the highest- and lowest-scoring responses and jointly uses Bradley-Terry and group reward entropy objectives.The entropy term addresses group degeneration by encouraging greater reward diversity within response groups.
  • Iterative Reward Model Lightweight Optimization: Theorem 4.1 states that GREBT reduces group degeneration and that the reduction increases as the weight coefficient α increases.For α ∈ [0, 1], the theorem gives Ci(φα) < Ci(φ0) and larger reductions for larger α.

5. Experiments and Analyses

Experiments evaluate R2M across dialogue and summarization tasks, baselines, reward-model updates, training dynamics, computational cost, and ablations. R2M consistently improves performance while using lightweight iterative updates that incorporate policy feedback.

  • Experimental Setup: R2M is integrated into RLOO and GRPO and compared with vanilla RL, pretrained-RM, frozen-feedback, and iterative-RM-head baselines.The experiments cover dialogue and document summarization settings.
  • Main Results: R2M achieves the best or second-best performance across evaluation metrics and significantly outperforms baseline methods.This pattern holds for both RLOO+R2M and GRPO+R2M models.
  • Ablation Study: Ablations show that using policy feedback without training provides no improvement and can degrade vanilla RL, while removing BT loss decreases LC by 3.0 and WR by 2.5.The results support updating R2M and using the mixed GREBT objective rather than a single loss.
  • Main Results: Compared with RLOO, RLOO+R2M increases LC win rate by 2.9% to 6.1%, raw win rate by 5.2% to 8.0%, and TL;DR win rate by 6.3%.These gains are attributed to replacing the frozen reward model with an iteratively updated model that allocates rewards using policy feedback.
  • Main Results: R2M outperforms GRPO with a pretrained reward model by 2.9% to 4.1% on Alpaca-Eval LC and by 2.7% to 5.1% on Alpaca-Eval WR.The paper attributes this enhancement to real-time policy alignment and additional information from policy feedback.
  • Reward-Model Analysis: R2M consistently outperforms Iterative RMHead because it recomputes reward signals with policy feedback before updating the reward model.The reconstructed reward distribution supplies additional information used by LGREBT to update R2M.
  • Training Dynamics: RLOO+R2M produces higher rewards and lower loss than RLOO, with denser high-KL, high-reward points during dialogue-task optimization.The analysis interprets this pattern as more effective policy updates and larger parameter shifts toward higher rewards.
  • Efficiency: R2M’s lightweight updates incur negligible additional computational cost because they aggregate policy feedback with lightweight attention without updating the reward model’s LLM.The method’s cross-attention module and scoring head remain relatively lightweight.

6. Discussion

R2M addresses reward noise and distributed-training overhead through lightweight communication and a GRE loss that amplifies genuine quality differences without creating artificial distinctions.

  • Multi-node communication: A lightweight cross-attention copy on the policy node reduces per-step transfer from B × S × Dp to 2×B×Dp.Only the reward-token embedding is transmitted during rollout.
  • Multi-node communication: Asynchronous synchronization avoids network blocking because reward-model updates are faster than policy optimization.The policy can continue training while the cross-attention module is synchronized.
  • Multi-node communication: R2M reduces rollout communication to 2×B×Dp, while training-time communication remains identical to standard RL.The reduction no longer scales with sequence length S.
  • GRE loss: When responses are identical, GRE produces a uniform distribution and zero parameter updates, preventing artificial distinctions.The loss activates only when meaningful quality differences exist.
  • GRE loss: GRE amplifies genuine quality differences while expected gradients wash out noise-induced variations across the data distribution.When true rewards differ, the symmetry breaks and the expected gradient points in the correct direction.

7. Conclusion

R2M is a lightweight RLHF framework that aligns the reward model with policy distribution shifts by incorporating evolving policy hidden states, improving performance with marginal computational cost.

  • 7. Conclusion: R2M incorporates evolving policy hidden states to achieve real-time alignment with policy distribution shifts.It integrates into existing RLHF algorithms without modifying their core mechanisms.
  • 7. Conclusion: R2M improves performance while introducing only marginal additional computational costs.The framework enhances the vanilla reward model through policy feedback.

Impact Statement

R2M is presented as practical for research and real-world use because it integrates with current RLHF algorithms and adds minimal computational overhead. The paper reports no specific societal risks beyond standard RLHF implications.

  • Practical impact: R2M enables iterative policy alignment and more accurate reward allocation with minimal computational overhead.It integrates with current RLHF algorithms without altering their core mechanism.
  • Societal impact: The paper identifies no specific societal risks requiring special consideration beyond those associated with standard RLHF practices.This scope statement concerns R2M's focus on reward-model enhancement within RL optimization.
  • Theoretical support: The theoretical analyses support R2M's hidden-state incorporation and GRE loss as mechanisms for enhancing vanilla reward models while maintaining robustness against reward overoptimization.These analyses provide theoretical support for the framework's core components.

A.1. Proof of Theorem 3.1

The proof compares reward-misalignment bounds for vanilla and hidden-state-aware reward models. Under positive post-fusion alignment quality, R2M has a strictly tighter bound, and perfect fusion leaves distribution shift as the remaining term.

  • Theorem comparison: R2M strictly tightens the reward-misalignment upper bound relative to the vanilla reward model when γ(t) > 0.The comparison assumes the trained fusion provides non-trivial post-fusion alignment.
  • Hidden-state fusion: R2M uses policy hidden states as key/value inputs and reward-model features as queries in a cross-attention module that produces fused features.The fused representation is written as h_f = M_ca(h^(t)) in the reward model's representation space.
  • Definitions: The distribution-shift degree is defined as ΔD(t) = TV(D(t)∥D(0)), while reward-misalignment error is evaluated over D(t).The policy induces the drifted distribution during training.
  • Proof strategy: The proof separately bounds vanilla and R2M misalignment errors before comparing them in the fused representation space.The vanilla model is treated as lacking access to fused hidden-state information.
  • Corollaries: When hidden states are perfectly aligned after fusion, the R2M bound simplifies to ϵ_R2M(t) ≤ ΔD(t) · L.In that case, misalignment is solely controlled by distribution shift.
  • Corollaries: As distribution shift grows, hidden-state fusion yields a larger relative improvement in the total misalignment bound.The analysis links this reduction to observed long-horizon stability under significant distribution drift.

A.2. Proof of Theorem 4.1

The proof shows that adding the Group Reward Entropy term to GREBT strictly reduces group degeneration, with stronger mitigation as α increases. The reduction vanishes as α approaches zero and is maximized at α = 1.

  • Strict Mitigation of Group Degeneration: GREBT strictly reduces the group degeneration degree relative to the pure Bradley-Terry minimizer.For a degenerate pure-BT solution, C(φα) < C(φ0) and the reduction is strictly positive.
  • Monotonic Increase of Mitigation Strength: The degeneration reduction ΔC(α) increases strictly with the weighting coefficient α ∈ (0, 1].The proof establishes that C(φα) decreases strictly as α increases, so ΔC(α) = C(φ0) − C(φα) increases.
  • Definitions: Group degeneration degree C(φ) is defined as the Group Reward Entropy loss and ranges from 0 to log K.C(φ) = log K corresponds to identical rewards across responses, while C(φ) approaching 0 corresponds to one response dominating the reward distribution.
  • GREBT Objective and Assumptions: The GREBT loss combines pure BT loss with GRE using α ∈ (0, 1] to control mitigation strength.The theorem assumes differentiability and a positive-definite Hessian at the GREBT minimizer.
  • Limiting Cases: At α = 1, GREBT becomes pure GRE and achieves maximum degeneration reduction, whereas α → 0+ recovers pure BT with no reduction.These limiting cases identify the full range of mitigation controlled by α.
  • Application to R2M: The theorem addresses group degeneration that emerges late in R2M training under shared reward-model and policy-feedback conditions.The GRE term provides a tunable theoretical mitigation mechanism for this degeneration.

G.1. Experimental Settings of Section 3

The experiments validate policy hidden states as feedback for R2M by measuring their relation to preference labels, reward differences, and policy-state changes across training. They also specify the datasets, models, sampling procedures, and evaluation protocols used for these validations.

  • Preference-Label Validation: The study samples chosen and rejected UltraFeedback query-response pairs to compare hidden-state similarity within and across preference categories.Pairs sharing labels form intra-category sets, while pairs with different labels form cross-category sets.
  • Reward-Score Validation: The experiments compare hidden-state similarity with reward-score differences for 300 UltraFeedback query-response pairs scored by Skywork-Reward-V2-Llama-3.1-8B.This directly relates policy representations to reward-model outputs.
  • Policy Feedback Choice: R2M uses last-layer hidden states as policy feedback because they serve as sequence representations and encode information through transformer forward propagation.The selected states are intended to combine semantic information with policy-state information.
  • Internal-State Validation: Internal-state validation extracts last-layer states from the same query-response pair at training steps 60, 120, 180, and 240 and compares them with pairwise cosine similarity.The procedure tests whether the representations vary across policy training steps.
  • Semantic Validation: Semantic validation manually labels 100 UltraFeedback queries’ policy responses as hacking or non-hacking before comparing their hidden states.Each query contributes eight hacking and eight non-hacking responses.
  • RLHF Experiments: Dialogue experiments use UltraFeedback-trained policies and evaluate Alpaca-Eval-related reward-model accuracy with and without policy feedback.The evaluation samples 1,024 preference pairs and compares rewards assigned to chosen and rejected responses.
  • RLHF Experiments: Summarization experiments use TL;DR with 2,048 queries per step, 50-token generation limits, and GPT-4 win-rate evaluation against reference summaries.The setup includes Pythia-2.8B-TL;DR-SFT and Pythia-2.8B-TL;DR-RM.

H. More Method Details of R2M

The method follows the standard RLHF pipeline of supervised fine-tuning, reward modeling, and policy optimization. R2M retains the reward model’s language-model representations while updating its scoring head and adding iterative reward-model optimization.

  • RLHF Workflow: RLHF proceeds through supervised fine-tuning, reward-model training on preference data, and policy optimization against the learned reward model.The reward model is trained to score chosen responses above rejected responses under the Bradley-Terry model.
  • Reward Modeling: A vanilla reward model maps the final-token hidden state of a pretrained LLM through a scoring head to produce a scalar reward.This representation is called the Reward Token Embedding.
  • RL Optimization: Policy optimization maximizes reward while constraining deviation from the supervised fine-tuned reference policy through β.The coefficient β balances reward maximization with adherence to reference-policy behavior.
  • Lightweight Optimization: R2M avoids full reward-model retraining by freezing the reward model’s language-model component and updating only its projection head.The approach is motivated by the robustness gains obtainable from retraining a projection head on smaller datasets.
Loading 2601.22664v4…