Diffusion Reward Models

Xiangyang Wang; Bingxiang He; Zeyuan Liu; Jiaze Wang; Ziqing Qiao; Yuxin Zuo; Huan-ang Gao; Cheng Qian

cs.LGSep 27, 2026

Diffusion Reward Models

Authors: Xiangyang Wang, Bingxiang He, Zeyuan Liu, Jiaze Wang, Ziqing Qiao, Yuxin Zuo, Huan-ang Gao, Cheng Qian, +8 more

Organizations: Tsinghua University · University of Illinois Urbana-Champaign

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.

Figures & tables

Figure 1 : Overview of DRM. A frozen encoder maps $(x,y)$ to a hidden state $\mathbf{h}$ that conditions a Diffusion Reward Head modeling $p(\mathbf{r}\mid x,y)$ . (a) Training. A single head supports both multi-attribute regression, supervised by a masked denoising loss $\mathcal{L}_{\mathrm{MSE}}$ , and pairwise preference data, supervised by an additional Bradley–Terry objective $\mathcal{L}_{\mathrm{BT}}$ on the denoised rewards. (b) Inference. Drawing $N$ samples via DDIM yields an empirical reward distribution that captures multimodal structure and is summarized by an aggregator into the final reward $R$ .

Reward Models	RewardBench v2	PPE Pref	PPE Corr	RMB Pairwise	RM-Bench	JudgeBench	Avg.
Discriminative Reward Models
internlm2-7b-reward Cai et al. [2024]	53.4	62.1	60.4	67.1	67.1	59.4	61.6
Eurus-RM-7b Yuan et al. [2024]	58.1	59.6	60.0	65.5	69.0	58.4	61.8
Starling-RM-34B Zhu et al. [2023]	45.5	62.8	60.3	72.0	67.1	63.8	61.9
RM-Mistral-7B Dong et al. [2023]	59.6	61.8	56.4	66.6	66.9	62.1	62.2
ArmoRM-Llama3-8B-v0.1 Wang et al. [2024a]	66.5	60.6	61.4	64.6	67.7	53.2	62.3

Table 1 : Overall benchmark results of different reward models. Baselines are taken from Liu et al. [2025a] where available; entries not reported therein are evaluated by us under the same protocol. Gray rows denote large-scale or closed baselines that are not directly comparable in model size or publicly reported training-data scale.

	DRM-Multi-8B	DRM-Pref-8B
Training Data	ArmoRM agg.	Tulu3 pref.
Reward Dim $K$	19	1
Diffusion Reward Head
Hidden Size	384	384
Blocks / Heads	3 / 6	3 / 6
Dropout	0.2	0.2

Table 2: Configurations of the two DRM variants.

Actor	Arena-Hard v2	MT-Bench
Tulu3-8B-SFT	1.0	59.4
+ ArmoRM RLHF	1.0	71.5
+ FsfairX RLHF	1.3	73.6
+ DRM-Multi RLHF	2.0	74.8

Table 3 : Downstream RLHF performance with different reward models. All methods use the same Tulu3-8B-SFT initialization and RLHF training setup.

Dataset	Dimension	Key Disagreement Statistic(%)	Separated Clusters(%)	Polarization(%)
HelpSteer2	Helpfulness	Range $\geq$ 2: 43.49	28.20	24.34
HelpSteer2	Correctness	Range $\geq$ 2: 41.99	27.79	24.24
MultiPref	Overall	Opposite-side: 45.64	37.23	11.40
MultiPref	Helpfulness	Opposite-side: 42.77	34.37	9.94

Table 4 : Structured disagreement in repeated human annotations on different datasets.

Model	Dimension	Wasserstein $\downarrow$	JS $\downarrow$	L1 $\downarrow$
DRM	Helpfulness	0.804	0.215	0.907
Empirical prior	Helpfulness	1.030	0.231	1.012
Global Gaussian	Helpfulness	1.032	0.254	1.054
Pointwise	Helpfulness	1.032	0.254	1.407
DRM	Correctness	0.846	0.251	1.006
Empirical prior	Correctness	1.019	0.224	0.994

Table 5 : Distributional distances to empirical human reward distributions on HelpSteer2-Disagreements. Lower is better for all metrics, and the best results are shown in bold.

Figure 2 : DRM multimodal ratio increases with human disagreement.

Task	100%	90%	80%	70%	Gain
GPQA (Best@k)	44.73	45.34	46.83	48.32	+3.59
MATH (Best@k)	50.78	53.15	55.37	56.98	+6.20
MMLU-Pro (Best@k)	60.94	62.26	63.66	62.85	+1.91
IFEval (Best@k)	62.89	62.69	62.93	64.25	+1.36
MBPP+ (Best@k)	69.43	69.96	70.20	70.42	+0.99
Average	-	-	-	-	+2.81

Table 6 : Selective prediction on PPE Correctness using DRM’s distributional uncertainty. Performance is reported at different coverage levels, with Gain denoting the improvement from 100% to 70% .

Setting	100%	90%	80%	70%	Gain
Helpfulness BoN	65.67	67.91	69.77	72.15	+6.48
Harmlessness BoN	61.24	64.41	66.54	68.23	+6.99
Helpfulness Pairwise	80.80	83.58	86.02	88.11	+7.31
Harmlessness Pairwise	73.97	76.33	77.61	78.52	+4.56

Table 7 : Selective prediction results on RMB at different coverage levels. Gain denotes the improvement from 100% to 70% coverage.

Setting	Candidates number	Mean	LCB-0.40	Gain
PPE BoN Macro	2	51.211	51.274	+0.063
PPE BoN Macro	4	54.328	54.445	+0.117
PPE BoN Macro	8	56.236	56.401	+0.165
PPE BoN Macro	16	57.381	57.556	+0.175
PPE BoN Macro	32	57.637	57.797	+0.160
RMB Harmlessness BoN	2	74.478	74.717	+0.239

Table 8 : Risk-sensitive Best-of-N selection using lower confidence bound (LCB) aggregation. LCB consistently improves over mean aggregation across PPE and RMB settings.

Figure 3 : Best-of- $N$ scaling curves on PPE Correctness, shown for the five-task average and each individual task.

Figure 4 : Reward-axis scaling on RewardBench v2.

Figure 5 : Reward distributions of diffusion, discriminative, and generative reward models.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Search range
Hidden size	$\{384,512,768\}$
Transformer blocks	$\{3,4,5\}$
Attention heads	$\{6,8,12\}$
Dropout	$\{0.0,0.1,0.2\}$
Learning rate	$\{0.3,0.5,1,2,4\}{\times}10^{-4}$
Batch size	$\{64,128,256\}$

Appendix

Table 9: Hyperparameter search space.

Hyperparameter	Search range
Hidden size	$\{384,512,768\}$
Transformer blocks	$\{3,4,5\}$
Attention heads	$\{6,8,12\}$
Dropout	$\{0.0,0.1,0.2\}$
Learning rate	$\{0.3,0.5,1,2,4\}{\times}10^{-4}$
Batch size	$\{64,128,256\}$

Appendix

Table 9: Hyperparameter search space.

Hidden	Depth	Heads	Beta schedule	RewardBench v2
384	3	6	sqcos	62.88
384	4	6	sqcos	59.83
384	5	6	sqcos	62.04
512	4	8	sqcos	62.38
512	4	8	linear	61.76
768	4	12	sqcos	54.80

Appendix

Table 10 : Maximum scores for each parameter combination. The final selected core architecture configuration is highlighted.

Model	Helpfulness (BoN)	Harmlessness (BoN)	Avg.
ArmoRM-Llama3-8B-v0.1	63.6	49.7	56.7
Skywork-Reward-Llama-3.1-8B-v0.2	60.5	56.8	58.7
internlm2-7b-reward	62.6	56.3	59.5
DeepSeek-GRM-27B	63.9	58.0	61.0
Eurus-RM-7b	67.9	54.3	61.1
Claude-3.5-Sonnet-20240620	70.5	51.8	61.2

Appendix

Table 11 : BoN split of RMB results of different reward models.

Model	Knowledge	Reasoning	Math	Code	Avg.
Arena-Hard/GPT-4o-20240513	51.6	52.6	68.4	49.1	54.3
Arena-Hard/Claude-3.5-Sonnet-20240620	52.3	59.6	61.0	48.3	54.6
DRM-Multi-8B	55.8	62.0	63.7	57.9	58.6
DRM-Pref-8B	54.9	57.0	62.2	60.3	57.1

Appendix

Table 12 : JudgeBench results of different reward models.

Model	PPE-MMLU	PPE-MATH	PPE-GPQA	PPE-IFEval	PPE-MBPP	Avg.
Skywork-Reward-Gemma-2-27B	53.9	62.7	52.7	53.8	59.3	56.5
Eurus-RM-7b	63.4	69.3	53.9	59.4	54.0	60.0
Starling-RM-34B	67.7	66.4	57.0	56.1	54.6	60.3
internlm2-7b-reward	66.7	72.6	54.6	64.0	43.9	60.4
Skywork-Reward-Llama-3.1-8B-v0.2	64.3	69.6	56.5	61.5	51.6	60.7
ArmoRM-Llama3-8B-v0.1	66.5	70.7	57.0	58.4	54.2	61.4

Appendix

Table 13 : PPE correctness results of different reward models.

Model	DDIM Steps $S$	RewardBench v2
DRM-Multi	5	64.9
DRM-Multi	10	65.6
DRM-Multi	50	51.3
DRM-Multi	100	49.4
DRM-Pref	5	64.0
DRM-Pref	10	65.7

Appendix

Table 14: Ablation on the number of DDIM sampling steps. We fix the guidance scale to

\omega=7

and the number of reward samples to

N=32

Model	DDIM Steps $S$	RewardBench v2
DRM-Multi	5	64.9
DRM-Multi	10	65.6
DRM-Multi	50	51.3
DRM-Multi	100	49.4
DRM-Pref	5	64.0
DRM-Pref	10	65.7

Appendix

Table 14: Ablation on the number of DDIM sampling steps. We fix the guidance scale to

\omega=7

and the number of reward samples to

N=32

Model	Guidance $\omega$	RewardBench v2
DRM-Multi	1	57.6
DRM-Multi	3.5	64.1
DRM-Multi	7	65.6
DRM-Multi	14	57.6
DRM-Pref	1	61.8
DRM-Pref	3.5	65.5

Appendix

Table 15: Ablation on the classifier-free guidance scale. We fix the number of DDIM sampling steps to

S=10

and the number of samples to

N=32

Inference Mask	RewardBench v2
Training masks	65.1
All-one mask	65.6

Appendix

Table 16: Ablation on the reward mask used during inference for DRM-Multi.

Model	Reward Dim.	Training Samples	RewardBench v2
UltraFeedback variant	5	240K	54.6
DRM-Multi-Half	19	284K	63.9

Appendix

Table 17: Analysis of reward-space dimensionality under comparable training scales.

Model	Supervision	Training Samples	RewardBench v2	PPE Pref	PPE Corr	RMB Pairwise	RM-Bench	JudgeBench	Avg.
DRM-Multi	Multi-attribute	569K	65.6	62.5	63.8	78.0	68.8	58.6	66.2
DRM-Multi-Half	Multi-attribute	284K	63.9	61.3	63.2	77.0	68.6	56.8	65.1
DRM-Pref	Pairwise preference	273K	65.7	63.0	62.5	78.2	68.1	57.1	65.8

Appendix

Table 18: Comparison of DRM variants under different training scales and supervision regimes. DRM-Multi-Half is trained on a stratified 50% subset of the DRM-Multi training corpus.

Benchmark	DRM-Multi-8B		DRM-Pref-8B		ArmoRM
Benchmark	Mean $\pm$ Std	95% CI	Mean $\pm$ Std	95% CI	ArmoRM
RewardBench v2	$65.133\pm 0.519$	$[64.761,65.504]$	$65.655\pm 0.197$	$[65.514,65.796]$	$66.5$
PPE Pref	$62.047\pm 0.138$	$[61.948,62.146]$	$63.047\pm 0.133$	$[62.952,63.142]$	$60.6$
PPE Corr	$63.653\pm 0.103$	$[63.580,63.727]$	$62.446\pm 0.078$	$[62.391,62.502]$	$61.4$
RMB	$77.966\pm 0.071$	$[77.915,78.017]$	$78.255\pm 0.042$	$[78.225,78.285]$	$64.6$
RM-Bench	$68.917\pm 0.156$	$[68.806,69.029]$	$68.060\pm 0.093$	$[67.994,68.127]$	$67.7$

Appendix

Table 19 : Stability of DRM across 10 inference seeds. We report the mean, standard deviation, and 95% confidence interval over seeds

\{0,\ldots,9\}

. All DRM evaluations use

S=10

N=32

, and

\omega=7

. ArmoRM is deterministic in evaluation mode and is reported as a fixed single-run reference.

Model	$N$	$S$	Encoder Lat. (ms)	RewardDiT Lat. (ms)	E2E Lat. (ms)	Throughput (samples/s)	Peak Mem. (GB)	Rel. Cost
FsfairX	–	–	38.983	–	38.983	25.652	14.104	$1.000\times$
QRM	–	–	48.099	–	48.099	20.791	14.121	$1.234\times$
DRM	1	10	38.834	21.197	62.403	16.025	14.152	$1.601\times$
DRM	8	10	38.834	21.799	63.259	15.808	14.152	$1.623\times$
DRM	32	10	38.834	22.461	63.559	15.734	14.152	$1.630\times$
DRM	1	5	41.246	10.501	53.231	18.786	14.156	$1.365\times$

Appendix

Table 20: Inference efficiency of DRM and baseline reward models. Latency is measured with batch size 1 on 256 randomly sampled RewardBench v2 examples. For DRM,

\omega=7

. Relative cost is normalized by the end-to-end latency of FsfairX.

Model	$S$	RBv2	PPE Pref.	PPE Corr.	RMB	RM-Bench	JudgeBench	Avg.
DRM-Multi	10	65.6	62.5	63.8	78.0	68.8	58.6	66.2
DRM-Multi	5	64.9	62.1	63.9	77.7	69.0	57.4	65.9
DRM-Pref	10	65.7	63.0	62.5	78.2	68.1	57.1	65.8
DRM-Pref	5	64.0	62.6	60.8	77.3	67.1	57.5	64.9

Appendix

Table 21: Reward-model performance under different DDIM sampling budgets. All DRM results use

N=32

and guidance scale

\omega=7

. The measured end-to-end latency is 63.559 ms for

S=10

and 53.755 ms for

S=5

, corresponding to

1.630\times

and

1.378\times

the latency of FsfairX, respectively.

Explore similar work

Latent Reward Registers for Diffusion Preference Alignment
Aug 4, 2026Yuanshen Guan, Zipeng Feng, Chengru Song +2Diffusion Alignment Reward Gradients
1Kling Team · University of Science and Technology of China
Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models
Jun 7, 2026Yawen Shao, Jie Xiao, Kai Zhu +6Large Language Model Reinforcement Learning Diffusion Language Models
University of Science and Technology of China · Tongyi Lab · Northeastern University
Reinforcement Learning from Denoising Feedback
May 25, 2026Qi He, Huan Chen, Ya Guo +3Large Language Model Reinforcement Learning Diffusion-Based Reinforcement Learning Methods
Fudan University · Ant Group · Hong Kong University of Science and Technology

Aug 4, 2026, cs.LGJ/K move · Enter open · S save

Latent Reward Registers for Diffusion Preference Alignment

Yuanshen Guan, Zipeng Feng, Chengru Song, Zhiwei Xiong+1

1Kling Team · University of Science and Technology of China

Diffusion Alignment Reward Gradients

Open paper AskPDF arXiv

Aligning diffusion models with human preferences usually relies on a sparse terminal reward evaluated on the final generated samples, presenting a severe temporal credit-assignment challenge across the multi-step denoising process. We propose Latent Reward Registers, a mechanism that estimates terminal preference directly from intermediate noisy latents by prepending learnable, position-free register tokens to the input sequence of a frozen Diffusion Transformer (DiT). This independent readout mechanism extracts latent reward evidence without altering the generator's hidden states or velocity field. The resulting dense, differentiable reward signal throughout the full denoising process facilitates two alignment strategies. For training, Reward-Gradient On-Policy Distillation (RG-OPD) distills reward-guided updates along on-policy trajectories, bypassing the computationally expensive rollouts of standard policy gradients. For inference, Reward-Guided Sampling (RGS) steers trajectories via magnitude-matched reward gradients without parameter updates. Empirically, at high noise levels (u = 0.8), the registers reach the highest pairwise accuracy among the evaluated latent reward models. Furthermore, RG-OPD outperforms online reinforcement learning baselines while reducing GPU hours by up to 33x, and RGS establishes a new state-of-the-art among training-free methods, strictly enhancing both alignment and perceptual metrics. Code and weights are available at https://github.com/Guanys-dar/latent-reward-register

Abstract

Explore similar work

Latent Reward Registers for Diffusion Preference Alignment

Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models

Reinforcement Learning from Denoising Feedback