Diffusion Reward Models
DRM models rewards with a diffusion transformer, capturing multimodal preferences and improving RLHF decisions.
Diffusion Reward Models treat reward modeling as conditional density estimation rather than a point estimate or a fixed parametric family. Conditioned on a frozen LLM encoder, a lightweight Diffusion Transformer turns Gaussian noise into a reward vector and supports both multi-attribute regression and pairwise preferences. Across five benchmarks, DRM matches or surpasses matched-scale baselines and remains competitive with much larger reward models. Uncertainty-aware rejection and lower-confidence-bound aggregation improve decisions, and using DRM as the RLHF reward improves downstream policy performance.
- Diffusion transformer estimates reward distributions without a fixed parametric family.
- One model handles multi-attribute regression and pairwise preference data.
- Matches or beats matched-scale baselines across five benchmarks.
- Rejection and lower-confidence bounds use distributional uncertainty.
- RLHF policies improve when DRM is the training-time reward.
Full article236 words · extracted from huggingface.co · click to collapse
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(rmid 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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2609.33803