Search papers, labs, and topics across Lattice.
This paper addresses the problem of spurious correlations in reward models used in Reinforcement Learning from Human Feedback (RLHF) by proposing a factored representation learning framework. The framework decomposes contextual embeddings into causal factors sufficient for reward prediction and non-causal factors capturing reward-irrelevant attributes, constraining the reward head to depend only on the causal component. Experiments on mathematical and dialogue tasks demonstrate improved robustness and downstream RLHF performance compared to baselines, with analyses showing mitigation of reward hacking behaviors like exploiting length and sycophantic bias.
Stop reward hacking: disentangling causal and non-causal factors in reward models makes RLHF more robust.
A reliable reward model is essential for aligning large language models with human preferences through reinforcement learning from human feedback. However, standard reward models are susceptible to spurious features that are not causally related to human labels. This can lead to reward hacking, where high predicted reward does not translate into better behavior. In this work, we address this problem from a causal perspective by proposing a factored representation learning framework that decomposes the model's contextual embedding into (1) causal factors that are sufficient for reward prediction and (2) non-causal factors that capture reward-irrelevant attributes such as length or sycophantic bias. The reward head is then constrained to depend only on the causal component. In addition, we introduce an adversarial head trained to predict reward from the non-causal factors, while applying gradient reversal to discourage them from encoding reward-relevant information. Experiments on both mathematical and dialogue tasks demonstrate that our method learns more robust reward models and consistently improves downstream RLHF performance over state-of-the-art baselines. Analyses on length and sycophantic bias further validate the effectiveness of our method in mitigating reward hacking behaviors.
