This repository contains the codebase for our paper, "Rethinking Reward Models for Multi-Domain Test-Time Scaling."
conda create -n multi-rm python=3.10.14
conda activate multi-rm
pip install -r requirements.txt
pip install flash-attn --no-build-isolation
# TASK_TYPE can be one of:
# gORM / gPRM
TASK_TYPE=[choose_one_above]
# generate data
python -m data_generation.generate_data \
--output_dir [OUTPUT_DIR] \
--task_type ${TASK_TYPE}
# preprocess data
python -m data_generation.preprocess_data \
--output_dir [OUTPUT_DIR] \
--task_type ${TASK_TYPE}
# shorten critique (optional)
python -m data_generation.shorten_critique \
--output_dir [OUTPUT_DIR] \
--task_type ${TASK_TYPE}# Training dORM / dPRM
# Use the appropriate config file:
# ./configs/dORM-14B.yaml
# ./configs/dPRM-14B.yaml
# ./configs/dORM-8B.yaml
# ./configs/dPRM-8B.yaml
# ./configs/dORM-qwen.yaml
# ./configs/dPRM-qwen.yaml
accelerate launch -m discriminative.train \
--config ./configs/dORM-14B.yaml \
--output_dir ./[TRAINING_RESULTS]/dORM-14B \
--per_device_batch_size 4 \
--category all
# Training gORM / gPRM
# Use the appropriate config file:
# ./configs/gORM-14B.yaml
# ./configs/gPRM-14B.yaml
# ./configs/gORM-8B.yaml
# ./configs/gPRM-8B.yaml
# ./configs/gORM-qwen.yaml
# ./configs/gPRM-qwen.yaml
accelerate launch -m generative.train \
--config ./configs/gORM-14B.yaml \
--output_dir ./[TRAINING_RESULTS]/gORM-14B \
--per_device_batch_size 4 \
--category all# DATASET_NAME can be one of:
# MMLU-Pro / GPQA-diamond / MedQA / LEXam
#
# MODEL_NAME is the model whose CoTs the test set contains:
# MMLU-Pro / GPQA-diamond: Llama-3.1-8B-Instruct, SmolLM3-3B, Qwen2.5-7B-Instruct, gemma-2-9b-it, Llama-3.1-70B-Instruct
# MedQA / LEXam: SmolLM3-3B, gemma-2-9b-it
DATASET_NAME=[choose_one_above]
MODEL_NAME=[choose_one_above]
TEST=${DATASET_NAME}_${MODEL_NAME}_test
# Inference for dORM / dPRM
# Use the appropriate model checkpoint:
# dongboklee/dORM-14B
# dongboklee/dPRM-14B
# or use your own trained models
python -m discriminative.get_reward \
--data_path dongboklee/${TEST} \
--model_id dongboklee/dORM-14B \
--output_dir ./[REWARD_RESULTS]/dORM-14B-${TEST} \
--per_device_batch_size 8 \
--category all
# Inference for gORM / gPRM
# Use the appropriate model checkpoint (a LoRA adapter):
# dongboklee/gORM-14B, TASK_TYPE=gORM
# dongboklee/gPRM-14B, TASK_TYPE=gPRM
# or use your own trained models: [LOCAL_DIR]/gORM-14B, [LOCAL_DIR]/gPRM-14B
python -m generative.get_reward \
--data_path dongboklee/${TEST} \
--model_id dongboklee/gORM-14B \
--output_dir ./[REWARD_RESULTS]/gORM-14B-${TEST} \
--task_type gORM \
--category all# DATASET_NAME can be one of:
# MMLU-Pro / GPQA-diamond / MedQA / LEXam
#
# MODEL_NAME is the model whose CoTs the test set contains:
# MMLU-Pro / GPQA-diamond: Llama-3.1-8B-Instruct, SmolLM3-3B, Qwen2.5-7B-Instruct, gemma-2-9b-it, Llama-3.1-70B-Instruct
# MedQA / LEXam: SmolLM3-3B, gemma-2-9b-it
DATASET_NAME=[choose_one_above]
MODEL_NAME=[choose_one_above]
TEST=${DATASET_NAME}_${MODEL_NAME}_test
# Or use your own reward dirs instead of HF hubs:
# [REWARD_MODEL_NAME]/[TEST]/[CATEGORY]_reward.json
#
# Swap -14B for -8B or -qwen (e.g. dORM-8B, dORM-qwen) to evaluate the other backbones
python -m evaluation.evaluate \
--data_path dongboklee/${TEST} \
--output_dir [OUTPUT_DIR] \
--reward_dirs \
dongboklee/dORM-14B-${TEST} \
dongboklee/dPRM-14B-${TEST} \
dongboklee/gORM-14B-${TEST} \
dongboklee/gPRM-14B-${TEST} \
--model_names dORM-14B dPRM-14B gORM-14B gPRM-14B \
--strategies last min mean mean \
--num_runs 100
# CSV_FILE can be one of:
# [OUTPUT_DIR_FROM_ABOVE]/best_of_n.csv
# [OUTPUT_DIR_FROM_ABOVE]/weighted_vote.csv
CSV_FILE=[choose_one_above]
# [OUTPUT_FILE_PREFIX]=example
# -> example_legend.png / example_legend.pdf
# -> example.png / example.pdf
python -m evaluation.plot \
--input_file ${CSV_FILE} \
--output_file [OUTPUT_FILE_PREFIX]Please find the assets of this repo below, including training and test datasets and model checkpoints.
| Name | Description |
|---|---|
| MMLU-Pro_Llama-3.1-8B-Instruct_train | MMLU-Pro training dataset for dORM/dPRM, adapted from VersaPRM. |
| MMLU-Pro_Llama-3.1-8B-Instruct_gORM_train | MMLU-Pro training dataset for gORM. |
| MMLU-Pro_Llama-3.1-8B-Instruct_gPRM_train | MMLU-Pro training dataset for gPRM. |
@article{
lee2026rethinking,
title={Rethinking Reward Models for Multi-Domain Test-Time Scaling},
author={Dong Bok Lee and Seanie Lee and Sangwoo Park and Minki Kang and Jinheon Baek and Dongki Kim and Dominik Wagner and Jiongdao Jin and Heejun Lee and Tobias Bocklet and Jinyu Wang and Jingjing Fu and Sung Ju Hwang and Jiang Bian and Lei Song},
journal={Transactions on Machine Learning Research},
issn={2835-8856},
year={2026},
url={https://openreview.net/forum?id=PgouBhL7IR},
note={J2C Certification}
}