Teburile/DRM

Diffusion Reward Models

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

Links

• 📜 Paper — coming soon
• 💻 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

Model	File	Reward dimensions	Supervision	Training data
DRM-Multi-8B	DRM-Multi-8B/model.pth	19	Masked multi-attribute denoising	RLHFlow/ArmoRM-Multi-Objective-Data-v0.1
DRM-Pref-8B	DRM-Pref-8B/model.pth	1	Denoising + Bradley–Terry preference loss	allenai/llama-3.1-tulu-3-8b-preference-mixture

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.

Reward Model	RewardBench v2	PPE Pref	PPE Corr	RMB Pairwise	RM-Bench	JudgeBench	Avg.
ArmoRM-Llama3-8B-v0.1	66.5	60.6	61.4	64.6	67.7	53.2	62.3
QRM-Llama3.1-8B-v2	70.7	57.2	60.3	61.1	72.5	62.6	64.1
URM-LLaMa-3.1-8B	73.9	60.2	60.4	65.7	72.0	64.1	66.1
DRM-Multi-8B	65.6	62.5	63.8	78.0	68.8	58.6	66.2
DRM-Pref-8B	65.7	63.0	62.5	78.2	68.1	57.1	65.8

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

@article{drm2026,
  title  = {Diffusion Reward Models},
  author = {Wang, Xiangyang and He, Bingxiang and Liu, Zeyuan and Wang, Jiaze and Qiao, Ziqing and Zuo, Yuxin and Yu, Tianyu and Chen, Qianyu and Gao, Huan-ang and Qian, Cheng and Zhang, Wenbin and Li, Ran and Sun, Youbang and Ding, Ning and Shi, Yuanchun and Liu, Zhiyuan and Xiao, Chaojun and Yu, Chun},
  year   = {2026}
}