PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Reward Models

Xiangyang Wang, Bingxiang He, Zeyuan Liu, Jiaze Wang, Ziqing Qiao, Yuxin Zuo, Huan-ang Gao, Cheng Qian, Wenbin Zhang, Ran Li, Youbang Sun, Ning Ding, Yuanchun Shi, Zhiyuan Liu, Chaojun Xiao, Chun Yu

arXiv:2609.33803v2cs.LGcs.AI

TL;DR

Scalar scores and fixed distribution families do not capture the multimodal structure of human reward judgments. DRM uses conditional diffusion to model reward distributions, and across five benchmarks it matches or surpasses matched baselines while supporting distribution-aware decisions and improved downstream RLHF policy performance.

  • Problem

    Conventional reward models use point-valued scores or fixed output-distribution families, although human judgments can disagree and induce multimodal reward structure.

    Source: §1 · pp. 1–2

  • Method

    DRM conditions a diffusion head on a frozen LLM representation to estimate p(r | x, y) without a fixed parametric output family, supporting both multi-attribute and preference training.

    Source: §1 · p. 2

  • Results

    Across five benchmarks, DRM matches or surpasses matched-data and matched-backbone baselines, remains competitive with larger reward models, and recovers multimodal reward structure.

    Source: §1 · pp. 2–3, §6 · p. 14

  • Takeaways & Limitations

    Uncertainty-aware rejection and lower-confidence-bound aggregation improve reward-model decisions, while downstream RLHF experiments report improved policy performance using DRM as the training-time reward.

    Source: §1 · p. 2

    DRM is trained on only a few hundred thousand open-source samples, does not yet match the strongest open-source scalar RMs in absolute terms, and has been studied with one 8B encoder and fixed data scale.

    Source: Limitations · p. 15

Abstract

from arXiv · show

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.

1. Introduction

Conventional reward models reduce prompt–response pairs to scalar scores or fixed distribution families, despite human preferences being multimodal. DRM models conditional reward distributions with diffusion and reports competitive benchmark results plus benefits from distributional information.

1. Introduction

  • Human judgments can disagree over values, rubric interpretation, and helpfulness–harmlessness trade-offs, making a stable point-valued reward an incomplete account of preference.The passage reports approximately 63% inter-annotator agreement on Anthropic-HH and persistent within-rubric disagreement in HelpSteer3-Preference.Source: §1 · pp. 1–2
  • Existing approaches remain tied to fixed output families: Gaussian heads are unimodal, categorical heads depend on bin granularity, and quantile heads use a fixed grid.Multi-objective models also require a predefined attribute schema, while generative judges still produce a single verdict.Source: §1 · p. 2
  • DRM replaces the conventional value head with a diffusion head that estimates p(r | x, y) without imposing a parametric output family.A frozen LLM representation conditions a lightweight Diffusion Transformer that denoises Gaussian noise into a reward vector.Source: §1 · p. 2
  • One architecture supports multi-attribute regression and pairwise preference data, while inference samples form an empirical distribution that can be aggregated for reward scoring or risk-sensitive use.The overview depicts masked denoising supervision for attribute data and an additional Bradley–Terry objective for preference pairs.Source: §1 · p. 2, Figure 1 · p. 3
  • Across five benchmarks, DRM matches or surpasses matched-data, matched-backbone baselines, remains competitive with larger reward models, and improves reward-model decisions using distributional information.The reported comparisons span discriminative, distributional, and generative models; downstream RLHF experiments also report improved policy performance.Source: §1 · pp. 2–3

2. Method

DRM uses a frozen encoder and a conditional diffusion head to model reward distributions, with shared architecture across multi-attribute and preference training. At inference, repeated samples preserve distributional information for aggregation and test-time scaling.

2.1. Architecture

  • A frozen LLM encoder represents each prompt–response pair, and a lightweight Diffusion Transformer iteratively denoises reward noise conditioned on that representation.The denoising process is motivated by diffusion models’ mode coverage and capacity to approximate continuous densities.Source: §2 · p. 3, §2.1 · pp. 3–4

2.2. Training

  • The shared diffusion parameterization supports multi-attribute regression and pairwise preference data by changing how supervision targets are constructed.For multi-attribute examples, a mask excludes unannotated dimensions from the denoising loss; preference-only training adds a Bradley–Terry ranking objective on denoised reward estimates.Source: §2.2 · pp. 4–5
  • For preference-only data, symmetric pseudo-reward targets centered at zero use a fixed margin, while a BT-style auxiliary loss enforces pairwise ordering.The ranking objective is applied to denoised reward estimates, alongside denoising supervision.Source: §2.2 · p. 5

2.3. Inference

  • At inference, N reverse-process samples approximate the conditional reward distribution, and aggregation can yield a scalar compatible with standard RLHF or distribution-aware decision signals.The paper studies uncertainty-aware rejection and risk-sensitive aggregation beyond a single scalar score.Source: §2.3 · pp. 5–6
  • Increasing reward samples for a fixed response tightens the estimate of its mean, creating a test-time scoring lever unavailable to deterministic scalar reward models.This reward-axis scaling is distinct from Best-of-N response selection, which selects among candidate responses.Source: §2.3 · p. 6

3. Main Evaluation

The evaluation tests DRM under multi-attribute and pairwise-preference supervision, then examines benchmark performance, reward attributes, correctness, and downstream RLHF. Under matched data and backbone, DRM-Multi-8B outperforms ArmoRM and QRM, with further analyses reporting balanced attribute scores and improved RLHF outcomes.

3.1. Experimental Setup

  • Two 8B DRM variants use a frozen FsfairX-LLaMA3 encoder and are evaluated on five benchmarks spanning chat, instruction following, math, code, reasoning, factuality, and safety.DRM-Multi-8B uses 569K multi-attribute examples; DRM-Pref-8B uses 273K chosen/rejected pairs.Source: §3.1 · p. 6

3.2. Main Results

  • 66.2 average: DRM-Multi-8B exceeds matched ArmoRM (62.3) and QRM (64.1), with its largest gain on RMB (78.0).The comparison holds training data and backbone fixed, isolating the reward head.Source: Table 1 · p. 7, §3.2 · pp. 7–8
  • DRM remains competitive across discriminative, distributional, and generative reward models, including much larger models and judges.The reported examples include Llama-3.1-Nemotron-70B (67.8), DeepSeek-GRM-27B (65.6), and GPT-4o (67.7).Source: Table 1 · p. 7, §3.2 · p. 8
  • DRM-Multi-8B and DRM-Pref-8B score similarly (66.2 vs. 65.8) under multi-attribute and pairwise-preference supervision.This supports transfer of the reward-space diffusion framework across the two supervision regimes.Source: §3.2 · p. 8

3.3. Fine-Grained Analysis of DRM

  • DRM maintains closely matched Helpfulness and Harmlessness scores on RMB Best-of-N, with an average difference below 5 points.The passage contrasts this balance with reward models that score higher on one dimension than the other.Source: §3.3 · p. 8

3.4. DRM as a Training-Time Reward for RLHF

  • DRM-Multi-RLHF improves Arena-Hard v2 from 1.3 to 2.0 and MT-Bench from 73.6 to 74.8 versus the cited scalar-reward baselines.The comparison uses the same actor initialization and RLHF training setup.Source: Table 3 · p. 8, §3.4 · p. 9

3.3. Fine-Grained Analysis of DRM

  • 58.6 on JudgeBench and 63.8 on PPE Correctness show DRM-Multi-8B can identify correctness differences across factual, knowledge, math, and code tasks.It outperforms GPT-4o and Claude-3.5-Sonnet on JudgeBench and surpasses multiple open-source reward models on PPE Correctness.Source: §3.3 · p. 8

4. Distributional Analysis of DRM

Repeated human ratings reveal clustered, polarized disagreement, and DRM’s reward distributions reflect these patterns. Experiments also show that distribution-aware decisions and additional test-time sampling can improve reward-model decisions or scores, while a single example illustrates multimodal rewards.

4.1. Distributional Validation with Human Disagreement

  • Repeated ratings show substantial clustering and polarization, with MultiPref disagreement more pronounced for subjective attributes than harmlessness.On HelpSteer2, 43.49% of helpfulness examples have rating range at least 2; MultiPref fully consistent harmlessness annotations occur in 87.11% of examples.Source: §4.1 · p. 9
  • On HelpSteer2-Disagreements, DRM achieves the lowest helpfulness distances and best correctness Wasserstein distance against empirical human ratings.Helpfulness distances are 0.804 Wasserstein, 0.215 JS divergence, and 0.907 L1; correctness Wasserstein distance is 0.846.Source: §4.1 · p. 10, Table 5 · p. 10
  • DRM’s multimodal ratio rises with stronger human disagreement, reaching 63.2% for low-high polarized helpfulness examples.For correctness, the multimodal ratio rises from 37.6% for rating range at most 1 to 62.0% for low-high polarized examples.Source: §4.1 · p. 10, Figure 2 · p. 10

4.2. Distribution-Aware Decision Making

  • Rejecting uncertain decisions improves retained accuracy as coverage falls, while LCB outperforms mean aggregation across evaluated PPE and RMB Best-of-N settings.At 70% coverage, PPE Correctness gains 2.81 percentage points on average across five subtasks; at N=32, PPE macro score rises from 57.637 to 57.797 with LCB.Source: §4.2 · pp. 11–12, Table 6 · p. 11, Table 7 · p. 11

4.3. Test-Time Scaling: Two Axes

  • DRM variants scale best among evaluated models as Best-of-N grows, and reward sampling raises RewardBench v2 scores without retraining.DRM-Multi-8B rises from 56.5% with one sample to 65.6% with 32 samples; DRM-Pref-8B rises from 64.4% to 65.7%.Source: Figure 4 · p. 12, Figure 3 · p. 12, §4.3 · p. 13

4.4. Qualitative Illustration of Multimodal Rewards

  • On one prompt-response pair, DRM produces well-separated reward regions, whereas ArmoRM collapses to a point and DeepSeek-GRM remains sharply unimodal.The comparison uses 100 scores per model on a common reward dimension normalized to [−1, 1]; it is a qualitative single-example illustration.Source: §4.4 · p. 13, Figure 5 · p. 13

5. Related Work

Related work situates reward modeling within RLHF and documents multimodal human preferences. It contrasts scalar, generative, and fixed-family distributional reward models with DRM’s diffusion-based approach.

Reinforcement Learning from Human Feedback

  • In general-domain alignment, learned reward quality bounds alignment quality, unlike verifiable-reward settings where ground-truth signals can support RLVR-style training.RLHF methods rely on learned rewards to fit human feedback; RLVR-style training is an exception in domains with ground-truth answers.Source: §5 · p. 13

Multimodal Structure of Human Preference

  • Human preference varies across annotators, cultures, and demographics, while reducing attributes to a scalar can average away non-degenerate preference structure.Reported inter-annotator agreement across several RLHF datasets is 63–73%, with most disagreement described as genuine individual preference rather than noise.Source: §5 · pp. 13–14

The Paradigm of Reward Modeling

  • Reward-model paradigms include deterministic discriminative scores, generative verdicts, and distributional heads that commit to fixed families; DRM instead uses a diffusion reward head.Examples of fixed-family heads predict Gaussian, categorical, or fixed-quantile distributions.Source: §5 · p. 14

6. Conclusion

The paper presents DRM as a non-parametric conditional reward-distribution model with one architecture for multi-attribute and preference data. Across benchmarks and downstream experiments, it reports competitive reward-model performance, useful distributional decision signals, and improved RLHF policy performance.

6. Conclusion

  • DRM’s DiT head models conditional reward distributions without a parametric family and handles both multi-attribute and preference data.At inference, reward samples provide distributional statistics and a reward-axis scaling lever unavailable to scalar reward models.Source: §6 · p. 14
  • Across five benchmarks, DRM matches or surpasses matched baselines, remains competitive with much larger reward models, and recovers multimodal structure conventional heads cannot represent.The reported comparisons include discriminative, distributional, and generative reward models.Source: §6 · p. 14
  • Uncertainty-aware rejection and risk-sensitive aggregation use learned distributions as decision signals beyond their mean, and DRM improves policy performance when used as the RLHF reward model.Source: §6 · p. 14

Limitations

The study identifies data and scale boundaries for DRM, and describes how its preference objective relates to full distribution-level likelihood. The diffusion head uses masked reward inputs and conditioning to predict denoising noise.

Limitations

  • DRM uses only a few hundred thousand open-source samples, does not match the strongest open-source scalar RMs in absolute terms, and is evaluated with one 8B encoder at fixed data scale.Its behavior across backbone sizes, model families, and larger data regimes remains unexplored.Source: Limitations · p. 15

A. Relationship to Distribution-Level Preference Likelihood

  • The denoised-estimate Bradley–Terry objective is an upper-bound-style surrogate that encourages higher distribution-level win probability, but is not equivalent to direct distributional-likelihood optimization.Source: A. Relationship to Distribution-Level… · p. 20

B. Implementation Details of Diffusion Reward Head

  • The head takes noisy rewards, a diffusion timestep, frozen-backbone text features, and a reward mask, then predicts noise; masked training excludes unannotated dimensions from supervision.At inference, an all-one mask activates all reward dimensions for joint generation in one DDIM process.Source: B. Implementation Details of Diffusion… · pp. 20–21

C. Benchmark Details

The five benchmarks span preference and correctness judgments across language, reasoning, safety, and other capabilities. Their formats include best-response selection, preference pairs, and objectively labeled response comparisons.

RewardBench v2

  • RewardBench v2 has 1,865 test samples and six capability categories, using best-of-4 selection to assess whether a model can choose the best response.Source: C. Benchmark Details · p. 21

PPE Preference + PPE Correctness

  • PPE combines 16,038 preference pairs with 2,555 prompts scored across 32 responses for correctness on verifiable tasks including MMLU-Pro, MATH, GPQA, IFEval, and MBPP-Plus.Source: C. Benchmark Details · p. 21

RMB

RM-Bench

  • RM-Bench tests sensitivity to subtle response differences across easy, normal, and hard examples, including interference from length, formatting, and style.It contains 1,327 test samples spanning chat, code, math, and safety domains.Source: C. Benchmark Details · p. 22

JudgeBench

  • JudgeBench uses response pairs labeled by objective correctness, emphasizing factual and logical correctness in knowledge, reasoning, mathematics, and coding.It contains 350 GPT-4o-generated and 270 Claude-3.5-Sonnet-generated pairs.Source: C. Benchmark Details · p. 22

D. Hyperparameter Tuning

The selected DRM configuration was tuned on RewardBench v2, and further ablations examine sampling, guidance, masking, and reward-space breadth. The results motivate the defaults used in the experiments.

D. Hyperparameter Tuning

  • 384_3_6_sqcos with lr = 5e-5, Dropout = 0.2, and Batch = 64 was selected as the best-performing configuration while maintaining low training and inference cost.The search tuned architecture and optimization settings on RewardBench v2; DRM-Pref selected λBT = 0.5 on validation.Source: Table 10 · p. 22, D. Hyperparameter Tuning · pp. 22–23

E.1. DDIM Sampling Steps and Guidance Scale

  • S = 10 and guidance scale ω = 7 performed best for both DRM variants; more steps gave no benefit and could degrade ranking, while weak or overly strong guidance also hurt.Reducing steps from 10 to 5 caused only a modest drop; ω = 7 was the best tested guidance scale.Source: §E.1 · pp. 23–24, Table 14 · p. 24, Table 15 · p. 24

E.2. Number of Reward Samples

  • RewardBench v2 scores rose from 56.5 to 65.6 for DRM-Multi and from 64.4 to 65.7 for DRM-Pref as reward samples increased from N = 1 to N = 32.The study fixes S = 10 and ω = 7; N = 32 becomes the default.Source: §E.2 · p. 24

E.3. Reward Mask

  • The all-one inference mask raises DRM-Multi's RewardBench v2 score from 65.1 to 65.6, supporting joint generation of all 19 reward dimensions in one sampling process.Training examples have labels for only a subset of dimensions, so the training loss masks unannotated dimensions.Source: §E.3 · pp. 24–25, Table 16 · p. 24

E.4. Reward Dimension

  • At comparable training scales, the 19-dimensional DRM-Multi-Half scores 63.9 on RewardBench v2, versus 54.6 for the 5-dimensional UltraFeedback variant.The comparison tests reward-space breadth using a 5-dimensional model trained on the largest single UltraFeedback subset.Source: §E.4 · p. 25, Table 17 · p. 25

E.5. Training Scale and Supervision Regime

Training-scale comparisons distinguish data quantity from supervision type, while seed tests assess stochastic inference stability. The efficiency analysis measures latency and throughput under a specified evaluation setup.

E.5. Training Scale and Supervision Regime

  • DRM-Pref scores slightly higher on preference benchmarks, while DRM-Multi performs better on correctness benchmarks; both supervision regimes remain effective and competitive.Preference scores are 65.7 on RewardBench v2, 63.0 on PPE Preference, and 78.2 on RMB Pairwise for DRM-Pref; DRM-Multi reaches 63.8 on PPE Correctness and 58.6 on JudgeBench.Source: §E.5 · pp. 25–26
  • Halving DRM-Multi data reduced its score from 66.2 to 65.1, while DRM-Pref averaged 65.8 versus 65.1 for DRM-Multi-Half at comparable scale.The authors attribute the full DRM-Multi versus DRM-Pref gap largely to training scale rather than an inherent weakness of pairwise supervision.Source: §E.5 · p. 25, Table 18 · p. 26

E.6. Stability Across Inference Seeds

  • Across 10 inference seeds, DRM-Multi-8B averages 65.927 ± 0.183 and DRM-Pref-8B 65.771 ± 0.100, exceeding ArmoRM's fixed 62.300 score.The reported gains over ArmoRM are +3.627 and +3.471, respectively, larger than the variation from stochastic sampling.Source: Table 19 · p. 26, §E.6 · p. 26

F. Inference Efficiency; F.1. Measurement Setup

  • The efficiency study compares DRM with FsfairX and QRM, measuring encoder, head, and end-to-end latency alongside sample and token throughput and peak GPU memory.It uses 256 RewardBench v2 examples at batch size 1; end-to-end latency includes encoding, DDIM sampling, reward aggregation, and overhead.Source: F. Inference Efficiency · p. 27, §F.1 · p. 27, Table 20 · p. 27

F.2. Results

The appendix quantifies DRM’s inference costs and shows how reducing sampling steps trades latency against modest score changes. It also reports that cached encoder representations make repeated scoring rely on the lightweight denoiser.

F.2. Results

  • At default N=32, S=10, DRM costs 63.559 ms end-to-end (1.63× FsfairX) and adds 0.048 GB peak memory, mainly from RewardDiT denoising.Source: §F.2 · p. 27
  • At S=10, increasing N from 1 to 32 raises RewardDiT latency only from 21.197 to 22.461 ms, whereas S=5 cuts it to 11.054 ms.Source: §F.2 · p. 27

F.3. Accuracy-Latency Trade-off

  • 15.4% lower end-to-end latency at S=5 (63.559 to 53.755 ms) accompanies DRM-Multi's 66.2-to-65.9 and DRM-Pref's 65.8-to-64.9 score drops; DRM-Multi still beats QRM, 65.9 vs. 64.1.The S=5 setting reduces relative latency cost from 1.63× to 1.38×; the table reports 1.630× and 1.378×, respectively.Source: §F.3 · pp. 27–28, Table 21 · p. 28

F.4. Offline Encoder Caching

  • For 256 samples, precomputing frozen FsfairX embeddings takes 9.385 s total (36.662 ms per sample; 4.360 MB), leaving repeated scoring to RewardDiT at about 22 ms per example.Source: §F.4 · p. 28

G. Other Experiments Details

The appendix details benchmark reporting, model and PPO configurations, compute requirements, and artifact and data practices. It also records how LLMs were used during manuscript preparation.

G. Other Experiments Details

  • RMB results in Table 11 use Best-of-N splits and are not directly comparable to the RMB pairwise scores in Table 1.Source: G. Other Experiments Details · p. 28
  • The full pipeline has approximately 7.52B parameters, including a 7.50B frozen FsfairX encoder and an approximately 12.0M-parameter RewardDiT; only RewardDiT is trained.Source: G. Other Experiments Details · p. 28
  • PPO uses full-parameter actor optimization from the same Llama-3.1-Tulu-3-8B-SFT initialization, with 62,880 UltraFeedback prompts per run and 131 rollout-update steps.Source: G. Other Experiments Details · p. 28
  • For PPO reward computation, DRM-Multi's 32 samples over 19 dimensions are averaged across samples and then dimensions; policies are evaluated on Arena-Hard v2 and MT-Bench using GPT-4.1.Source: G. Other Experiments Details · pp. 28–29
  • RewardDiT checkpoint training took approximately 1.54 GPU-hours, excluding the one-time cost of generating embeddings with the frozen encoder.Source: G. Other Experiments Details · p. 29

H. Responsible NLP and Artifact Use

  • The study uses public datasets, benchmarks, pretrained models, and baselines under their licenses and access conditions, and does not redistribute restricted artifacts.Source: H. Responsible NLP and Artifact Use · p. 29
  • No new personal data are collected or individuals identified; potentially sensitive preference and safety data are used only for research and aggregate evaluation.Source: H. Responsible NLP and Artifact Use · p. 29
  • LLMs assisted with manuscript language polishing, but not with hypotheses, methodology, experiments, analysis, or conclusions; authors reviewed edits and take responsibility for the manuscript.Source: H. Responsible NLP and Artifact Use · p. 29

Related papers

Loading 2609.33803v2…