DRM Recasts Reward Modeling as Conditional Density Estimation

A new arXiv preprint reframes reward modeling for LLM alignment as conditional density estimation rather than a single-number regression, using a Diffusion Transformer to sample from p(r|x,y) instead of collapsing every response to one score.

The authors' argument for the switch is stated plainly in the abstract: reward heads today "reduce each prompt–response pair to a point estimate or to a distribution from a fixed parametric family," which they call "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." DRM sits a "lightweight Diffusion Transformer" on top of a frozen LLM encoder and denoises Gaussian noise into a reward vector, with "no parametric assumption on the output distribution."

One architecture is reported to cover both multi-attribute regression and pairwise preference data. At inference, N samples "form an empirical reward distribution that can be aggregated into a scalar, a variance, or quantiles," which the paper leverages for "uncertainty-aware rejection and lower-confidence-bound (LCB) aggregation."

On results, the abstract is measured rather than triumphant. Across five unnamed benchmarks, DRM "matches or surpasses baselines under matched data and backbone" and "stays competitive with much larger discriminative, distributional, and generative RMs despite its modest training scale." The authors add that "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." The abstract publishes no per-benchmark numbers, no sample count N, and does not name the frozen encoder.