Docs: refresh model card and usage

README.md

@@ -1,109 +1,139 @@
 ---
 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
-  - rewarddit
-  - preference-modeling
-  - text-generation-evaluation
-base_model:
-  - sfairXC/FsfairX-LLaMA3-RM-v0.1
+- reward-model
+- diffusion
+- rlhf
+- alignment
+- preference-modeling
 ---
 
-# DRM Checkpoints
+# Diffusion Reward Models
 
-This repository hosts two RewardDiT checkpoints released with our paper.
+This repository provides the released **DRM-Multi-8B** and **DRM-Pref-8B** RewardDiT checkpoints.
 
-Paper: TODO: paste arXiv or project page link here.
+## Links
 
-Code: TODO: paste GitHub repository link here.
+- 📜 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
 
-| Name | File | Reward Dim | Training Data | Intended Use |
-| --- | --- | ---: | --- | --- |
-| DRM-Multi-8B | `DRM-Multi-8B/model.pth` | 19 | `RLHFlow/ArmoRM-Multi-Objective-Data-v0.1` | Multi-objective reward generation and scalar scoring by mean aggregation |
-| DRM-Pref-8B | `DRM-Pref-8B/model.pth` | 1 | `allenai/llama-3.1-tulu-3-8b-preference-mixture` | Pair-preference reward scoring |
+| 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` |
 
-Both checkpoints use FsfairX-LLaMA3-RM-v0.1 embeddings as the text condition.
+### Shared architecture and inference defaults
 
-## Recommended Inference Settings
+- 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
 
-Use the following settings for both checkpoints:
+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.
 
-```text
-mask_split = false
-num_steps = 10
-guidance_scale = 7.0
-num_samples = 32
-gate = off
-debias = off
-```
+## How to Use
 
-The public `ScoreGenerator` does not include gate or debias modules. Scores are computed by sampling RewardDiT rewards, averaging over samples, and then averaging over reward dimensions.
+DRM uses a custom reward head and should be loaded through the released repository rather than `AutoModelForCausalLM`.
 
-## Usage
+```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
+```
 
-After downloading the code repository and this model repository, run:
+Score a prompt–response pair with DRM-Multi-8B:
 
 ```bash
-python src_upload/score_generator.py \
-  --ckpt path/to/DRM-Multi-8B/model.pth \
+python score_generator.py \
+  --ckpt checkpoints/DRM-Multi-8B/model.pth \
   --prompt "User prompt" \
   --response "Assistant response"
 ```
 
-For the preference checkpoint:
+Or use the pairwise-preference checkpoint:
 
 ```bash
-python src_upload/score_generator.py \
-  --ckpt path/to/DRM-Pref-8B/model.pth \
+python score_generator.py \
+  --ckpt checkpoints/DRM-Pref-8B/model.pth \
   --prompt "User prompt" \
   --response "Assistant response"
 ```
 
-The default command-line inference settings in `score_generator.py` are already set to `num_steps=10`, `guidance_scale=7.0`, and `num_samples=32`.
+See [`USAGE.md`](USAGE.md) for the explicit inference arguments.
 
-## Model Details
+## Training Details
 
 ### DRM-Multi-8B
 
-- Reward dimension: 19
-- Text embedding dimension: 4096
-- Hidden size: 384
-- Depth: 3
-- Attention heads: 6
-- Dropout: 0.2
-- Beta schedule: `squaredcos_cap_v2`
-- Prediction type: `epsilon`
+- 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
 
-- Reward dimension: 1
-- Text embedding dimension: 4096
-- Hidden size: 384
-- Depth: 3
-- Attention heads: 6
-- Dropout: 0.2
-- Beta schedule: `squaredcos_cap_v2`
-- Prediction type: `epsilon`
-- Pair loss: denoising loss + Bradley-Terry loss + reward L2 regularization
-- `reward_reg_weight`: 0.001
+- 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`
 
-## Citation
+Experiments were conducted on NVIDIA A800-SXM4-80GB GPUs.
 
-TODO: paste BibTeX citation here after the paper is available.
+## Evaluation
 
-```bibtex
-@article{TODO,
-  title = {TODO},
-  author = {TODO},
-  journal = {arXiv preprint},
-  year = {TODO}
-}
-```
+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
 
-These checkpoints are research artifacts intended for reward-modeling experiments. They inherit the limitations of their training data and text encoder, and should not be treated as calibrated absolute measures of response quality.
+- **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}
+}
+```

USAGE.md

@@ -1,12 +1,21 @@
 # Usage
 
-Download the checkpoints from this repository and use them with the released code.
+Clone the released code, install its dependencies, and download the checkpoints:
+
+```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
+```
 
 ## DRM-Multi-8B
 
 ```bash
-python src_upload/score_generator.py \
-  --ckpt path/to/DRM-Multi-8B/model.pth \
+python score_generator.py \
+  --ckpt checkpoints/DRM-Multi-8B/model.pth \
   --prompt "User prompt" \
   --response "Assistant response" \
   --num_steps 10 \
@@ -17,8 +26,8 @@ python src_upload/score_generator.py \
 ## DRM-Pref-8B
 
 ```bash
-python src_upload/score_generator.py \
-  --ckpt path/to/DRM-Pref-8B/model.pth \
+python score_generator.py \
+  --ckpt checkpoints/DRM-Pref-8B/model.pth \
   --prompt "User prompt" \
   --response "Assistant response" \
   --num_steps 10 \
@@ -26,4 +35,4 @@ python src_upload/score_generator.py \
   --num_samples 32
 ```
 
-The scorer uses `mask_split=False`, gate off, and debias off.
+The scorer uses `mask_split=False`, with the gate and reward-debiasing transform disabled. It averages over samples and then over reward dimensions to return one scalar per input.