DRM / README.md

Add the bibtex and paper link

9245926 verified 26 days ago

6.13 kB

---
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](https://huggingface.co/papers/2609.33803)
- 💻 [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
@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}, 
}
```

	---
	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](https://huggingface.co/papers/2609.33803)
	- 💻 [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
	@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},
	}
	```