Diffusion Reward Models

Diffusion Reward Models. By Xiangyang Wang and fourteen others, from Tsinghua University, the Chinese University of Hong Kong, and one other institution. Published on the twenty-seventh of September. Its code repository is public but slight: six source files.

Reward models sit underneath the alignment of large language models. The dominant designs boil each pairing of prompt and response down to a single point estimate, or at best to a distribution drawn from one fixed parametric family. Human preference does not look like that. It is inherently multimodal: one response can reasonably be judged in several different ways, and no single family of distributions covers all of them. The model's output has a simpler shape than the thing it is trying to capture.

DRM, the Diffusion Reward Model, is built to fit that structure. It recasts reward modeling as conditional density estimation: learning the full distribution of the reward given a prompt and a response, rather than one number or one fixed shape.

The architecture pairs a frozen language model encoder with a lightweight Diffusion Transformer. Conditioned on the encoder, the Diffusion Transformer denoises Gaussian noise into a reward vector. Because the output distribution carries no parametric assumption, its multimodal structure is represented naturally. One architecture handles two kinds of training data: regression over multiple attributes, and pairwise preferences.

At inference, the model draws several samples, and together they form an empirical reward distribution. That distribution can then be summarized in whatever form a decision needs: a scalar, a variance, or quantiles. A scalar keeps compatibility with anything that expects an ordinary reward. A variance says how settled the judgment is. Quantiles describe its shape, including a cautious lower end that a decision can lean on.

Across five benchmarks, DRM matches or surpasses baselines trained on the same data and backbone. It stays competitive with much larger discriminative, distributional and generative reward models, despite a modest training scale. And where conventional reward heads collapse to a point, DRM recovers the multimodal structure of the reward.

The authors also show the distribution being put to work. Rejection that is aware of uncertainty, and aggregation by a lower confidence bound, both use information beyond a scalar reward to improve the reward model's decisions. Finally, in downstream RLHF experiments, using DRM as the reward during training leads to better policy performance, which the authors take as direct validation of diffusion-based reward modeling for RLHF.

The shift is from asking a reward model for a verdict to asking it for a spread of opinion. A single number can only say how good a response is. A distribution can also say how much judges would disagree, and whether they would split into camps. The paper's case is that a diffusion model, which assumes no particular shape for its output, is a natural tool for producing that spread, and that downstream training benefits from having it.