Representation-Aware Advantage Estimation: Your Reward Model Provides More Than A Scalar Output

2026-06-09

AI Summary

This work addresses a key limitation of traditional reinforcement learning from human feedback (RLHF)—its reliance solely on the scalar output of reward models, which discards rich semantic information embedded in their hidden representations, leading to noisy and coarse advantage estimates. To overcome this, the authors propose GraphAE, a representation-aware advantage estimation method that, for the first time, incorporates the reward model’s hidden states into advantage computation. By constructing a graph structure over responses and employing a lightweight graph propagation mechanism to aggregate contextual information, GraphAE substantially improves the quality of advantage estimation. The method seamlessly integrates into existing grouped RL algorithms and demonstrates consistent performance gains across benchmarks, achieving improvements of +6.3 on Arena-Hard-v0.1, +8.27 on AlpacaEval 2.0, and +0.22 on MT-Bench, thereby validating its superior sample efficiency, robustness, and generalization capability.

Abstract

Current reinforcement learning from human feedback (RLHF) methods primarily rely on scalar rewards from a trained reward model (RM). While effective, scalar rewards are often noisy and fail to capture fine-grained preference differences, whereas RM hidden states encode richer semantic and preference information. We introduce the representation-aware advantage estimation, which leverages RM hidden states and models them as auxiliary signals for better advantage estimation. Specifically, we propose the Graph-based Advantage Estimation (GraphAE), treat each sampled group as a graph, where nodes correspond to responses and edges capture their similarity in the RM hidden space. Then advantages are computed via graph propagation, enabling each sample to incorporate contextual information from its neighbors. GraphAE is lightweight and can be seamlessly integrated into existing group-based RL algorithms. We apply GraphAE to GRPO, GSPO and RLOO, and conduct extensive experiments on different models and benchmarks. Empirical results show consistent improvements across three benchmarks, with gains of up to + 6.3 on Arena-Hard-v0.1, + 8.27 on AlpacaEval 2.0, and + 0.22 on MT-Bench. These results demonstrate that leveraging RM representations leads to more sample efficient and robust RLHF.

Problem

• reinforcement learning from human feedback
• reward model
• scalar reward
• preference modeling
• advantage estimation

Innovation

• Representation-Aware Advantage Estimation
• Graph-based Advantage Estimation
• Reward Model Representations
• RLHF
• Graph Propagation