Diffusion Reward Models

This repository provides the released DRM-Multi-8B and DRM-Pref-8B RewardDiT checkpoints.

Links

• 📜 Paper
• 💻 Code
• 🤗 Base encoder: FsfairX-LLaMA3-RM-v0.1

Introduction

DRM (Diffusion Reward Model) models the conditional reward distribution p(r | x, y) instead of reducing every prompt–response pair to a single point estimate or a fixed parametric family. A frozen LLM encoder conditions a lightweight Diffusion Transformer (RewardDiT), which denoises Gaussian noise into reward vectors. At inference time, multiple samples form an empirical distribution that can provide a scalar score, uncertainty estimate, or risk-sensitive statistic. DRM diffuses reward vectors, not text.

Both checkpoints use the frozen 7.5B-parameter FsfairX-LLaMA3-RM-v0.1 encoder and train only an approximately 12M-parameter RewardDiT head.

Checkpoints

Shared architecture and inference defaults

• Text embedding dimension: 4096
• RewardDiT hidden size: 384
• Depth: 3 blocks
• Attention heads: 6
• Dropout: 0.2
• Diffusion schedule: squaredcos_cap_v2, 1000 training steps
• Prediction type: epsilon
• Inference: 10 DDIM steps, guidance scale 7, 32 reward samples

The released scorer does not use a gate model or reward-debiasing transform. It averages over sampled rewards and then over reward dimensions to produce one scalar per input.

How to Use

DRM uses a custom reward head and should be loaded through the released repository rather than AutoModelForCausalLM.

git clone https://github.com/thunlp/DRM.git
cd DRM
pip install -r requirements.txt

hf download Teburile/DRM DRM-Multi-8B/model.pth --local-dir checkpoints
hf download Teburile/DRM DRM-Pref-8B/model.pth --local-dir checkpoints

Score a prompt–response pair with DRM-Multi-8B:

python score_generator.py \
  --ckpt checkpoints/DRM-Multi-8B/model.pth \
  --prompt "User prompt" \
  --response "Assistant response"

Or use the pairwise-preference checkpoint:

python score_generator.py \
  --ckpt checkpoints/DRM-Pref-8B/model.pth \
  --prompt "User prompt" \
  --response "Assistant response"

See USAGE.md for the explicit inference arguments.

Training Details

DRM-Multi-8B

• Training data: ArmoRM aggregated multi-attribute preferences, 19 attributes
• Objective: masked denoising; unlabeled reward dimensions are excluded from the loss
• Learning rate: 5e-5
• Batch size: 64
• Head training cost: approximately 1.11 GPU-hours, excluding encoder embedding generation

DRM-Pref-8B

• Training data: Tulu3 pair-preference mixture
• Objective: denoising loss + Bradley–Terry loss + reward L2 regularization
• Bradley–Terry coefficient: 0.5
• Reward regularization weight: 0.001
• Learning rate: 5e-5
• Batch size: 64

Experiments were conducted on NVIDIA A800-SXM4-80GB GPUs.

Evaluation

The table reports results across five benchmarks and six metrics. For ArmoRM, QRM, URM, and DRM-Multi-8B, the training data and FsfairX backbone are matched; only the reward head differs.

DRM-Multi-8B improves the six-metric average by 3.9 points over ArmoRM and performs on par with parametric distributional reward heads without assuming an output family. It does not lead on every benchmark; for example, its RewardBench v2 score is 65.6, compared with 66.5 for ArmoRM.

Limitations

• Not the strongest reward model in absolute terms. These models use a modest amount of open-source data and do not match the strongest reward models trained at larger, non-comparable scales.
• Sampling steps are sensitive. Ten DDIM steps work well, while 50–100 steps degrade ranking accuracy; image-diffusion step counts should not be transferred directly.
• Scaling remains untested. The released results use one 8B encoder and fixed data scales.
• Standard reward-model risks apply. DRM may inherit biases from its training data and may be vulnerable to reward hacking when optimized without oversight.

Citation

@misc{wang2026diffusionrewardmodels,
      title={Diffusion Reward Models}, 
      author={Xiangyang Wang and Bingxiang He and Zeyuan Liu and Jiaze WangZiqing Qiao and Yuxin Zuo and Huan-ang Gao and Cheng Qian and Wenbin Zhang and Ran Li and Youbang Sun and Ning Ding and Yuanchun Shi and Zhiyuan Liu and Chaojun Xiao and Chun Yu},
      year={2026},
      eprint={2609.33803},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2609.33803}, 
}