Factored Causal Representation Learning for Robust Reward Modeling in RLHF

TL;DR

This paper tackles reward hacking in RLHF by introducing CausalRM, a factored causal representation learning framework that decomposes the backbone embedding into a minimal sufficient causal latent $z^{c}$ and a non-causal latent $z^{nc}$. The reward predictor is constrained to depend only on $z^{c}$, while an adversarial head with gradient reversal discourages $z^{nc}$ from encoding reward-relevant information, aided by a reconstruction head to prevent degenerate factorization. The method optimizes a joint objective promoting sufficiency and minimality for $z^{c}$, invariance against non-causal variation, and non-degeneracy, with a variational information bottleneck guiding the factorization. Empirical results on mathematical reasoning and open-ended dialogue show that CausalRM improves reward-model accuracy and downstream RLHF performance, while substantially mitigating length and sycophantic-bias-based reward hacking, indicating stronger robustness and generalization for alignment tasks. These findings suggest that causal-inspired latent-factor disentanglement can be a practical and scalable way to harden reward models against spurious cues in RLHF systems.

Abstract

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.

Figures (10)

• Figure 1: Causal graph for standard reward modeling. The prompt--response pair $(x,y)$ encode both causal ($z^c$) and non-causal ($z^{nc}$) factors, which in turn affect the predicted reward $r$. While the path $z^c \to r$ is desired, the spurious path $z^{nc} \to r$ leads to reward hacking.
• Figure 2: Overview of CausalRM. The backbone embedding $h$ is factorized into causal latents $z^c$ and non-causal latents $z^{nc}$ via a variational encoder. Reward prediction is restricted to depend only on $z^c$, while an adversarial head trained through a gradient reversal layer (GRL) discourages $z^{nc}$ from encoding reward-predictive information. A reconstruction decoder prevents degenerate factorization by reconstructing $h$ from $[z^c;z^{nc}]$.
• Figure 3: Reward hacking diagnosis on mathematical reasoning. The dashed curve shows the average normalized reward predicted by each reward model on the ID test set, and the solid curve is the average gold score measured by final-answer accuracy.
• Figure 4: Average win rate against the SFT model on the ID test sets of open-ended dialogue benchmarks during RLHF.
• Figure 5: Sensitivity of predicted reward to response length on mathematical reasoning. Length is normalized to $[0,1]$ and rewards are averaged within length quantile buckets. $\sigma_{\text{len}}$ denotes the standard deviation of bucket-wise mean rewards.