Rethinking Reward Models for Multi-Domain Test-Time Scaling

This repository contains the codebase for our paper, "Rethinking Reward Models for Multi-Domain Test-Time Scaling."

---

Quick Start

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

Synthetic Verfication Rationale Generation for gORM/gPRM (optional)

# 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

# Training dORM / dPRM
# Use the appropriate config file:
# ./configs/dORM-14B.yaml
# ./configs/dPRM-14B.yaml
# ./configs/dORM-8B.yaml
# ./configs/dPRM-8B.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

accelerate launch -m generative.train \
  --config ./configs/gORM-14B.yaml \
  --output_dir ./[TRAINING_RESULTS]/gORM-14B \
  --per_device_batch_size 4 \
  --category all

Inference (reward)

# TEST can be one of:
# test (CoTs generated by Llama3.1-8B-Instruct)
# test_smollm (CoTs generated by Smollm3)
# test_qwen (CoTs generated by Qwen2.5-7B-Instruct)
# test_gemma (CoTs generated by gemma2-9B-it)
# test_llama (CoTs generated by Llama3.1-80B-Instruct)

# 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:
# dongboklee/gORM-14B-merged, TASK_TYPE=gORM
# dongboklee/gPRM-14B-merged, TASK_TYPE=gPRM

python -m generative.get_reward \
  --data_path dongboklee/[TEST] \
  --model_id dongboklee/gORM-14B-merged \
  --output_dir ./[REWARD_RESULTS]/gORM-14B-[TEST] \
  --task_type gORM \
  --category all

# Inference for gORM / gPRM (for your own trained models)
# Use the appropriate model checkpoint:
# [LOCAL_DIR]/gORM-14B, TASK_TYPE=gORM
# [LOCAL_DIR]/gPRM-14B, TASK_TYPE=gPRM

# Merge LoRA for vLLM inference
python -m generative.merge_lora \
  --input_dir [LOCAL_DIR]/gORM-14B # saved to [LOCAL_DIR]/gORM-14B/tmp

python -m generative.get_reward \
  --data_path dongboklee/[TEST] \
  --model_id [LOCAL_DIR]/gORM-14B/tmp \
  --output_dir ./[REWARD_RESULTS]/gORM-14B-[TEST] \
  --task_type gORM \
  --category all

Evaluation

# TEST can be one of:
# test (CoTs generated by Llama3.1-8B-Instruct)
# test_smollm (CoTs generated by Smollm3)
# test_qwen (CoTs generated by Qwen2.5-7B-Instruct)
# test_gemma (CoTs generated by gemma2-9B-it)
# test_llama (CoTs generated by Llama3.1-80B-Instruct)
TEST=[choose_one_above]

# Or use your own reward dirs instead of HF hubs:
# [MODEL_NAME]/[TEST]/[CATEGORY]_reward.json
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]

---

Assets

Please find the assets of this repo below, including training and test datasets, and model checkpoints.

Datasets

Name	Description
MMLU-Pro_Llama-3.1-8B-Instruct_train	multi-domain training dataset for dORM/dPRM (mostly adapted from VersaPRM).
MMLU-Pro_Llama-3.1-8B-Instruct_gORM_train	multi-domain training dataset for gORM generated by QwQ-32B.
MMLU-Pro_Llama-3.1-8B-Instruct_gPRM_train	multi-domain training dataset for gPRM generated by QwQ-32B.
MMLU-Pro_Llama-3.1-8B-Instruct_test	multi-domain test dataset with CoTs (N=128) generated by Llama-3.1-8B-Instruct (mostly adapted from VersaPRM).
MMLU-Pro_SmolLM3-3B_test	multi-domain test dataset with CoTs (N=16) generated by SmolLM3-3B.
MMLU-Pro_Qwen2.5-7B-Instruct_test	multi-domain test dataset with CoTs (N=16) generated by Qwen2.5-7B-Instruct.
MMLU-Pro_gemma-2-9b-it_test	multi-domain test dataset with CoTs (N=16) generated by gemma-2-9b-it.
MMLU-Pro_Llama-3.1-70B-Instruct_test	multi-domain test dataset with CoTs (N=16) generated by Llama-3.1-70B-Instruct.
GPQA-diamond_Llama-3.1-8B-Instruct_test	GPQA-diamond test dataset generated by Llama-3.1-8B-Instruct.
GPQA-diamond_SmolLM3-3B_test	GPQA-diamond test dataset generated by SmolLM3-3B.
GPQA-diamond_Qwen2.5-7B-Instruct_test	GPQA-diamond test dataset generated by Qwen2.5-7B-Instruct.
GPQA-diamond_gemma-2-9b-it_test	GPQA-diamond test dataset generated by gemma-2-9b-it.
GPQA-diamond_Llama-3.1-70B-Instruct_test	GPQA-diamond test dataset generated by Llama-3.1-70B-Instruct.
MedQA_SmolLM3-3B_test	MedQA test dataset generated by SmolLM3-3B.
MedQA_gemma-2-9b-it_test	MedQA test dataset generated by gemma-2-9b-it.
LEXam_SmolLM3-3B_test	LEXam test dataset generated by SmolLM3-3B.
LEXam_gemma-2-9b-it_test	LEXam test dataset generated by gemma-2-9b-it.

Model Checkpoints

Name	Backbone	Trained On
dORM-14B	14B backbone	MMLU-Pro_Llama3.1-8B-Instruct_train
dPRM-14B	14B backbone	MMLU-Pro_Llama3.1-8B-Instruct_train
gORM-14B	14B backbone	MMLU-Pro_Llama3.1-8B-Instruct_gORM-train
gPRM-14B	14B backbone	MMLU-Pro_Llama3.1-8B-Instruct_gPRM-train
dORM-8B	8B backbone	MMLU-Pro_Llama3.1-8B-Instruct_train
dPRM-8B	8B backbone	MMLU-Pro_Llama3.1-8B-Instruct_train
gORM-8B	8B backbone	MMLU-Pro_Llama3.1-8B-Instruct_gORM-train
gPRM-8B	8B backbone	MMLU-Pro_Llama3.1-8B-Instruct_gPRM-train

---