Learning Latent Reasoning Traces for Scalar Reward Models End-to-End

2026-07-31

AI Summary

This work addresses key limitations in existing reward models—namely, the poor generalization of scalar approaches, the lack of numerical flexibility and probabilistic interpretability in generative methods, and the difficulty hybrid models face in ensuring that intermediate reasoning effectively supports reward prediction. To overcome these challenges, the paper introduces LatentRM, a novel framework that, for the first time, formulates intermediate reasoning trajectories as discrete latent variables and aligns them explicitly with scalar reward estimation via end-to-end on-policy optimization grounded in maximum likelihood. LatentRM thereby preserves the robustness of generative modeling while inheriting the probabilistic interpretability of scalar reward models. Empirical results demonstrate that LatentRM significantly outperforms current scalar, generative, and hybrid reward models across both in-distribution and out-of-distribution settings, as well as under RLHF training, achieving superior preference modeling and policy alignment in open-ended dialogue and complex reasoning tasks.

Abstract

Reward models (RMs) are central to aligning large language models with human preferences via reinforcement learning. Although traditional scalar RMs enable efficient and probabilistic reward modeling, they rely on superficial cues that fail to generalize to complex or out-of-distribution (OOD) tasks. Conversely, generative RMs leverage extensive reasoning to improve robustness on challenging tasks, but their natural language-based scores lack the numerical flexibility and probabilistic interpretability that scalar RMs offer. While recent approaches combine both paradigms through off-policy multi-task learning, such parallel optimization does not guarantee that generated reasoning traces actively align with or benefit downstream scalar reward prediction. To address this mismatch, we propose LatentRM, a reward modeling framework that learns intermediate reasoning traces as discrete latent variables to explicitly maximize the likelihood of downstream scalar rewards. Through on-policy optimization of the latent reasoning space end-to-end, LatentRM tightly couples deep reasoning-based evaluation with precise scoring. Extensive validations on in-distribution and OOD datasets and RLHF show that LatentRM outperforms scalar, generative, and hybrid RMs on preference modeling and policy alignment across tasks ranging from open-ended conversation to complex reasoning.

Problem

• reward modeling
• scalar reward
• reasoning traces
• out-of-distribution generalization
• preference alignment

Innovation

• Latent Reasoning
• Reward Modeling
• Scalar Reward
• On-policy Optimization
• Discrete Latent Variables