Real-Time Aligned Reward Model beyond Semantics

TL;DR

The paper tackles reward misalignment and overoptimization in RLHF by introducing Real-Time Aligned Reward Model (R2M), which couples the evolving internal states of the policy with a lightweight reward-model head through a Sequence-to-Token Cross Attention mechanism. By integrating policy feedback and a GREBT loss that blends Bradley–Terry supervision with a Group Reward Entropy term, R2M achieves real-time alignment to distribution shifts while keeping computational overhead minimal. Empirical results across dialogue and summarization tasks show that R2M substantially improves reward accuracy and policy performance, with ablations confirming the necessity of policy feedback and the GREBT components. The approach offers a practical path to robust RLHF that mitigates reward overoptimization and enhances stability during long-horizon training.

Abstract

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.

Figures (9)

• Figure 1: The average hidden state similarity of the same-preference pair set and the different-preference pair set across transformer layers. Each pair consists of two query-response samples with respective preference labels.
• Figure 2: Negative correlation between absolute difference of reward scores allocated by the RM and hidden state similarity. Each data point corresponds to a query-response pair labeled with either identical or differing human preferences.
• Figure 3: Overview of R2M. We first aggregate the last-layer hidden states $h_i$ from the policy with the LLM part output of the reward model. This aggregated representation is then fed into the scoring head for reward prediction. When the policy updates, we get the real-time feedback $h'_i$ and utilize it to construct preference pairs. Finally, we optimize the reward model by jointly minimizing the Bradley-Terry loss and the Group Reward Entropy loss.
• Figure 4: The structure of R2M. Building on the dataflow based on solely surface semantic information (left), R2M introduces an additional dataflow based on the policy feedback (right).
• Figure 5: We compare RLOO and RLOO+R2M in terms of loss, reward and KL divergence during RL optimization, using Qwen2.5-3B-Instruct and LLaMA3-8B-Instruct as policy models, and Skywork-Reward-V2-Llama-3.1-8B as the reward model. For KL divergence, we calculate it as the average of log probability differences between the reference model and the policy model for each token.

Theorems & Definitions (18)

• Theorem 3.1
• Theorem 4.1
• Definition 1.1: Distribution Shift Degree
• Definition 1.2: Reward Misalignment Error
• Definition 1.3: Post-Fusion Hidden State Alignment Quality
• Definition 1.4: Lipschitz Constants
• Definition 1.5: Hidden State Norm Bound
• Definition 1.6: Maximum Hidden-State Semantic Deviation
• proof
• Corollary 1.7