Diffusion Reward Models

2026-09-27

AI Summary

This work addresses the limitation of conventional reward models, which output point estimates or fixed distributions and thus fail to capture the multimodal structure of human preferences. We propose the Diffusion Reward Model (DRM), which reformulates reward modeling as conditional density estimation. By leveraging a frozen LLM encoder alongside a lightweight diffusion Transformer, DRM generates non-parametric, multimodal reward distributions, unifies regression and preference data, and supports uncertainty-aware aggregation via a Lower Confidence Bound (LCB) strategy. As the first approach to introduce diffusion models for reward prediction, DRM achieves performance on five benchmarks comparable to significantly larger baselines despite its smaller scale. It successfully recovers multimodal reward structures and substantially improves policy performance in downstream reinforcement learning from human feedback (RLHF).

Abstract

Reward models underpin the alignment of large language models, yet the dominant designs reduce each prompt--response pair to a point estimate or to a distribution from a fixed parametric family. This is at odds with human preference, which is inherently multimodal: the same response can be reasonably judged in many ways, and no single family covers all of them. To better fit this structure, we introduce DRM, a Diffusion Reward Model that recasts reward modeling as conditional density estimation over $p(\mathbf{r}\mid x,y)$. Conditioned on a frozen LLM encoder, a lightweight Diffusion Transformer denoises Gaussian noise into a reward vector, placing no parametric assumption on the output distribution and naturally representing its multimodal structure. A single architecture handles both multi-attribute regression and pairwise preference data, and at inference $N$ samples form an empirical reward distribution that can be aggregated into a scalar, a variance, or quantiles. Across five benchmarks, DRM matches or surpasses baselines under matched data and backbone, stays competitive with much larger discriminative, distributional, and generative RMs despite its modest training scale, and recovers multimodal reward structure where conventional heads collapse to a point. Uncertainty-aware rejection and lower-confidence-bound (LCB) aggregation further demonstrate that DRM can exploit distributional information beyond a scalar reward to improve reward-model decisions. Downstream RLHF experiments additionally show that using DRM as the training-time reward leads to improved policy performance, directly validating the practical benefit of diffusion-based reward modeling for RLHF training.

Problem

• Reward Modeling
• Human Preference
• Multimodal Distribution
• Large Language Model Alignment
• RLHF

Innovation

• Diffusion Reward Model
• Conditional Density Estimation
• Multimodal Preference
• Diffusion Transformer
• RLHF