Factored Causal Representation Learning for Robust Reward Modeling in RLHF

Abstract: A reliable reward model is essential for aligning LLMs 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.

• The paper introduces CausalRM, a method that decomposes a reward model's latent space into causal and non-causal components to mitigate reward hacking in RLHF.
• Experiments on mathematical reasoning and dialogue tasks show that CausalRM improves reward model accuracy and robustness to spurious features, outmatching baseline methods, especially in OOD benchmarks( on math OOD 85.6% vs. GoalRM 82.2% and InfoRM 82.5%).
• CausalRM enhances downstream RLHF performance, with consistent gains in final-answer accuracy and maintained alignment of predicted versus gold rewards during training.

CausalRM addresses reward hacking in RLHF by restructuring the reward model's latent space rather than post-hoc correcting known spurious features. The central idea is to decompose the backbone embedding of a prompt–response pair into two disentangled factors — a causal component that is minimally sufficient for predicting human preference, and a non-causal component that absorbs reward-irrelevant variation such as response length or sycophantic phrasing — and then structurally restrict reward prediction to depend only on the causal factor (2601.21350).

Motivation and causal framing

The paper's starting point is a causal graph over the standard reward modeling pipeline: the prompt–response pair $(x,y)$ generates both causal factors $z^c$ and spurious factors $z^{nc}$, both of which can influence the predicted reward $r$. The desired path $z^c \to r$ reflects genuine preference drivers; the path $z^{nc} \to r$ constitutes reward hacking, e.g., a reward model on mathematical tasks shifting its score when only answer length changes while solution quality is fixed. The authors formalize robustness as the invariance condition $r_\theta(x,y) \perp\!\!\!\perp z^{nc}$.

This framing distinguishes CausalRM from prior mitigation strategies. Methods such as ODIN (separate quality/length heads), length regularization for DPO, and MMD-based regularization require explicitly enumerating the spurious variables to be controlled. Representation-learning approaches such as GoalRM's contrastive goal-conditioned representations and InfoRM's variational information bottleneck filter irrelevant information but do not explicitly factorize the latent into causal and non-causal channels. CausalRM's claim is that an explicit factorization provides a stronger inductive bias than a single bottlenecked latent, since even capacity-constrained single latents can entangle spurious cues with reward-relevant features.

Architecture and training objective

CausalRM attaches a lightweight variational module to a frozen-recipe SFT backbone. The backbone embedding $h = f_\phi(x,y)$ is mapped to two diagonal-Gaussian posteriors $q_\alpha(z^c \mid h)$ and $q_\alpha(z^{nc} \mid h)$ with standard normal priors. Three heads operate on this factorization:

• Causal reward head: a linear head $z^c$0 predicts reward from $z^c$1 alone, implementing the structural bias toward invariance.
• Adversarial head: a head $z^c$2 predicts reward from $z^c$3 through a gradient reversal layer (GRL). The adversary minimizes its own preference loss; gradients reaching the encoder are negated, so the encoder maximizes adversarial loss, stripping reward-predictive information from $z^c$4.
• Reconstruction decoder: reconstructs $z^c$5 from $z^c$6, preventing degenerate solutions such as posterior collapse under the factorized bottleneck.

The overall objective is a minimax problem combining four terms: the Bradley–Terry pairwise preference loss plus a KL bottleneck on $z^c$7 (jointly derived as a variational lower bound on maximizing $z^c$8, i.e., sufficiency plus minimality), the reconstruction loss, and the adversarial loss with KL regularization on $z^c$9. The derivation follows the standard variational information bottleneck treatment: the preference log-likelihood lower-bounds $z^{nc}$0 via a Bradley–Terry variational likelihood, and KL terms to the prior upper-bound $z^{nc}$1. Latent dimensions are modest ($z^{nc}$2, $z^{nc}$3), and at inference the posterior mean $z^{nc}$4 replaces stochastic sampling for stable scoring.

Empirical results

Experiments cover mathematical reasoning (Qwen2.5-Math-7B, OpenMathInstruct-1 preferences, PPO via OpenRLHF) and open-ended dialogue (Qwen2.5-7B SFT'd on ShareGPT, Anthropic-HH preferences), evaluated in-distribution and across five OOD benchmarks per task against Standard RM, GoalRM, and InfoRM.

Reward model accuracy improves consistently. On math, CausalRM reaches 70.1% ID and 85.6% OOD average pairwise accuracy versus 68.3%/82.2% for GoalRM and 66.1%/82.5% for InfoRM; on dialogue it attains 72.3% ID and 62.3% OOD. Notably, the margin widens under distribution shift (+2.6 points over the second-best method on math OOD), which the authors interpret as evidence that explicit disentanglement reduces exploitation of spurious features precisely where such exploitation is most likely.

Downstream RLHF gains are consistent but more modest on math: final-answer accuracy improves by 1.5 points ID and 1.1 points OOD over the best baseline (74.0/89.6 vs. 72.5/88.5 for GoalRM). A diagnostic plot of predicted versus gold rewards during PPO training is more informative than the aggregate numbers: policies optimized with baseline reward models show gold-reward degradation as training proceeds — most severely for InfoRM, whose predicted reward climbs to 0.92 normalized while gold accuracy collapses from 0.65 to 0.46 — whereas CausalRM maintains aligned trajectories (predicted 0.89, gold 0.74 at the end of training). This divergence pattern is the classic signature of reward overoptimization and explains InfoRM's unexpectedly poor RLHF result (48.0% average accuracy, below the SFT reference of 66.9%).

On dialogue, judged by Qwen3-Max pairwise comparison, the CausalRM-trained policy wins 54.8% against Standard RM, 45.5% against InfoRM, and 42.3% against GoalRM on ID splits, retaining advantages OOD (51.3%, 38.7%, 31.6%). Win-rate dynamics against the SFT reference show baselines declining late in training (Standard RM falls to 44%) while CausalRM rises monotonically to 75%.

Invariance analyses

Two targeted probes support the causal-invariance claim directly:

• Length sensitivity: bucketing chosen responses by normalized length, CausalRM's mean predicted reward has a standard deviation of only 0.03 across bins, versus 0.12–0.22 for baselines, all of which exhibit pronounced negative sensitivity to long responses.
• Sycophantic artifacts: after constructing hacked preference data where a sycophantic prefix ("Sure, here is the response:") is prepended to chosen responses with probability 0.8, CausalRM trained on the perturbed set drops only −1.7 points ID / −1.1 points OOD relative to its unperturbed counterpart, whereas Standard RM drops −11.4 / −5.8. On several hacked OOD benchmarks CausalRM's accuracy actually improves slightly, indicating near-immunity to the injected artifact.

Ablations confirm each component contributes: removing factorization costs −4.0 ID points on math and raises length sensitivity from 0.03 to 0.14; removing the GRL costs −4.1 OOD points and raises length sensitivity to 0.13; removing reconstruction mainly hurts OOD (−2.5); and removing both KL regularizers produces the largest degradation (−11.0 math OOD, −24.6 dialogue-hacked ID). The KL terms are thus complementary to, not redundant with, the structural and adversarial mechanisms.

Limitations

The authors are explicit that CausalRM enforces invariance empirically rather than providing identifiability guarantees: the learned $z^{nc}$5 need not uniquely recover the true causal drivers of human preference, and formal identification would require interventional data, multiple environments, or confounder proxies absent from standard RLHF pipelines. Two further caveats bear on interpretation. First, the downstream RLHF gains on math, while consistent, are small (1.1–1.5 points), so the strongest evidence for the framework lies in the hacking diagnostics rather than headline accuracy. Second, the dialogue evaluation relies on Qwen3-Max as judge on 1,000 sampled instances per benchmark, introducing potential judge-model bias into win-rate comparisons. The paper also leaves open whether the invariance objective extends to process reward models and multi-turn settings with dynamic confounders, and how to connect empirical invariance objectives with identifiability-oriented formulations.

Conclusion

CausalRM demonstrates that factored causal representation learning — a structural restriction of the reward head to a minimal sufficient causal latent, combined with GRL-based adversarial suppression of reward leakage into a dedicated non-causal channel — yields reward models that are more accurate, generalize better under distribution shift, and are measurably less sensitive to length and sycophantic artifacts. The approach requires no enumeration of specific spurious variables and adds only lightweight modules atop the backbone, making it a practical template for incorporating causal principles into RLHF reward modeling.