Robust Reward Modeling for Large Language Models via Causal Decomposition

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Lu, Yunsheng, Yang, Zijiang, Pan, Licheng, Chu, Zhixuan
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917412784308224
author Lu, Yunsheng
Yang, Zijiang
Pan, Licheng
Chu, Zhixuan
author_facet Lu, Yunsheng
Yang, Zijiang
Pan, Licheng
Chu, Zhixuan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2604_13833
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Robust Reward Modeling for Large Language Models via Causal Decomposition
Lu, Yunsheng
Yang, Zijiang
Pan, Licheng
Chu, Zhixuan
Computation and Language
F.2.2; I.2.7
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.
title Robust Reward Modeling for Large Language Models via Causal Decomposition
topic Computation and Language
F.2.2; I.2.7
url https://arxiv.org/abs/2604.13833