---
license: apache-2.0
library_name: pytorch
pipeline_tag: text-classification
base_model: sfairXC/FsfairX-LLaMA3-RM-v0.1
language:
- en
tags:
- reward-model
- diffusion
- rlhf
- alignment
- preference-modeling
---

# Diffusion Reward Models

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

## Links

- 📜 Paper — coming soon
- 💻 [Code](https://github.com/thunlp/DRM)
- 🤗 [Base encoder: FsfairX-LLaMA3-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-LLaMA3-RM-v0.1)

## Introduction

![DRM overview](https://raw.githubusercontent.com/thunlp/DRM/main/figures/fig01_drm-overview.png)

**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](https://huggingface.co/sfairXC/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`](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`](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`.

```bash
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:

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

Or use the pairwise-preference checkpoint:

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

See [`USAGE.md`](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.

![Reward-axis scaling on RewardBench v2](https://raw.githubusercontent.com/thunlp/DRM/main/figures/fig04_reward-axis-scaling-rewardbench-v2.png)

## 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

```bibtex
@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}
}
```