Robust Reward Modeling for Large Language Models via Causal Decomposition

2026-04-15

AI Summary

This work addresses the susceptibility of existing reward models to spurious cues—such as response length and excessive deference—during alignment of large language models, which often impedes accurate capture of the true intent behind user prompts. To mitigate this, the authors propose an intent reconstruction regularization method grounded in causal disentanglement: a decoder maps candidate responses back to latent intent embeddings of the input prompts, and the reconstruction error serves as a regularizing signal during reward model training, explicitly suppressing shortcut features unrelated to the prompt. This approach, the first to incorporate intent reconstruction into reward modeling, significantly enhances robustness and generalization on Gemma-2–based models, improving RewardBench accuracy from 0.832 to 0.868, yielding shorter yet more frequently preferred outputs in Best-of-N selection, and maintaining performance under length perturbations and mild topic drift.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Cognitive Modeling & Cognitive Systems: Computational Creativity

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Reward models are central to aligning large language models, yet they often overfit to spurious cues such as response length and overly agreeable tone. Most prior work weakens these cues directly by penalizing or controlling specific artifacts, but it does not explicitly encourage the model to ground preferences in the prompt's intent. We learn a decoder that maps a candidate answer to the latent intent embedding of the input. The reconstruction error is used as a signal to regularize the reward model training. We provide theoretical evidence that this signal emphasizes prompt-dependent information while suppressing prompt-independent shortcuts. Across math, helpfulness, and safety benchmarks, the decoder selects shorter and less sycophantic candidates with 0.877 accuracy. Incorporating this signal into RM training in Gemma-2-2B-it and Gemma-2-9B-it increases RewardBench accuracy from 0.832 to 0.868. For Best-of-N selection, our framework increases length-controlled win rates while producing shorter outputs, and remains robust to lengthening and mild off-topic drift in controlled rewrite tests.

Problem

• reward modeling
• spurious correlations
• prompt intent
• alignment
• overfitting

Innovation

• causal decomposition
• reward modeling
• intent reconstruction
• spurious correlation
• alignment