RewardRank: Optimizing True Learning-to-Rank Utility
Gaurav Bhatt*, Kiran Koshy Thekumparampil, Tanmay Gangwani, Tesi Xiao, Leonid Sigal
$Abstract$ . Traditional ranking systems optimize offline proxy objectives that rely on oversimplified assumptions about user behavior, often neglecting factors such as position bias and item diversity. Consequently, these models fail to improve true counterfactual utilities such as such as click-through rate or purchase probability, when evaluated in online A/B tests. We introduce RewardRank, a data-driven learning-to-rank (LTR) framework for counterfactual utility maximization. RewardRank first learns a reward model that predicts the utility of any ranking directly from logged user interactions, and then trains a ranker to maximize this reward using a differentiable soft permutation operator. To enable rigorous and reproducible evaluation, we further propose two benchmark suites: (i) Parametric Oracle Evaluation (PO-Eval), which employs an open-source click model as a counterfactual oracle on the Baidu-ULTR dataset, and (ii) LLM-as-User Evaluation (LAU-Eval), which simulates realistic user behavior via large language models on the Amazon-KDD-Cup dataset. RewardRank achieves the highest counterfactual utility across both benchmarks and demonstrates that optimizing classical metrics such as NDCG is sub-optimal for maximizing true user utility. Finally, using real user feedback from the Baidu-ULTR dataset, RewardRank establishes a new state of the art in offline relevance performance. Overall, our results show that learning-to-rank can be reformulated as direct optimization of counterfactual utility, achieved in a purely data-driven manner without relying on explicit modeling assumptions such as position bias
To create the Conda environment from the env.yaml file:
-
Make sure you have Conda installed.
-
Run the following command in your terminal:
# Replace 'my_env_name' with your desired environment name # Replace '/your/custom/path' with the path where you want to create the environment $conda env create --file env.yaml --name my_env_name --prefix /your/custom/path
We use the dataset provided by ULTR-reranking (arXiv:2404.02543), which can be downloaded from HuggingFace.
By default, the dataset is cached in the .cache directory. We use this dataset and distill soft labels using the IPS model from the baidu-bert-model repository.
To set up the IPS model:
$ cd RewardRank
$ git clone [email protected]:philipphager/baidu-bert-model.git
$ mv baidu-bert-model bbmWe will use src/create_kdeval_data.py to generete the pseduo labels. Pass the path of downloaded data (by default in the ~/.cache/huggingface/philipphager___baidu-ultr_uva-mlm-ctr/clicks/0.1.0/...). Also provide the path of the output folder.
$ python src/create_kdeval_data.py --data_root='' --out_dir=''Download the KDD Cup dataset from the official ESCI GitHub repository. The dataset consists of the following three files:
shopping_queries_dataset_examples.parquetshopping_queries_dataset_products.parquetshopping_queries_dataset_sources.csv
We preprocess the dataset by filtering out non-English query groups (QGs) and sampling 8 items per QG.
To run the preprocessing script, provide the paths to the .parquet and .csv files and use the --process_esci flag:
$ python src/create_llm_data.py --process_esci \
--esci_examples='path/to/shopping_queries_dataset_examples.parquet' \
--esci_products='path/to/shopping_queries_dataset_products.parquet' \
--esci_sources='path/to/shopping_queries_dataset_sources.csv'To estimate purchase_probability using an LLM, we assume the model from Hugging Face has been downloaded locally at the path model_dir/model_id. You can configure inference parameters such as maximum context length (--max_tokens) and temperature as needed.
Run the inference script with the following command:
$ python src/create_llm_data.py --llm_inf \
--model_dir='path/to/model_dir' \
--model_id='model_id' \
--max_tokens=4096For Claude 3.5 Sonnet v2, we perform batch inference using Amazon Bedrock. Similar batch inference capabilities are also available for GPT-4o and other commercial LLMs through their respective APIs.
To train the reward model, first configure the appropriate flags in the train_reward.sh script based on your evaluation strategy. Specify the correct path to the preprocessed dataset.
Set the following flag:
ips_train=1Set the following flags:
llm_train=1
ips_train=0
org_train=0You can adjust other training parameters (e.g., learning rate, batch size, n_gpus, etc) within the train_reward.sh script as needed. Once configured, run the script to start training:
$ bash train_reward.shTo train the rankers, first configure the appropriate flags in the train_ranker.sh script based on your evaluation strategy (by setting ips_train and llm_train flag). Specify the correct path to the preprocessed dataset.
$ bash train_ranker.shUse the following flags (to be run using train_ranker.sh) to train RewardRank:
python main.py --n_gpus=$n_gpus --use_wandb \
--batch_size=$batch_size --output_path=$output_path \
--output_folder=$model --data_path=$data_path \
--save_epochs=10 --epochs=$epochs --lr=$lr --lr_drop=$lr_drop \
--num_workers=2 --eval_epochs=4 --train_ranker \
--load_path_reward=$load_path_reward --ultr_models=$ultr_mod \
--use_doc_feat --residual_coef=0.5 --reward_correctionUse the following flags (to be run using train_ranker.sh) to train PG-RANK*:
python main.py --n_gpus=$n_gpus --use_wandb \
--batch_size=$batch_size --output_path=$output_path \
--output_folder=$model --data_path=$data_path \
--save_epochs=10 --epochs=$epochs --lr=$lr --lr_drop=$lr_drop \
--num_workers=2 --eval_epochs=4 --train_ranker \
--load_path_reward=$load_path_reward --ultr_models=$ultr_mod \
--use_doc_feat --pgrank_loss --pgrank_disc -mc_sample=5Use the following flags (to be run using train_ranker.sh) to train URCC*:
python main.py --n_gpus=$n_gpus --use_wandb \
--batch_size=$batch_size --output_path=$output_path \
--output_folder=$model --data_path=$data_path \
--save_epochs=10 --epochs=$epochs --lr=$lr --lr_drop=$lr_drop \
--num_workers=2 --eval_epochs=4 --train_ranker \
--load_path_reward=$load_path_reward --ultr_models=$ultr_mod \
--use_doc_feat --urcc_lossBe sure to set the appropriate flags for LLM_exp in the train_ranker.sh script.
To evaluate the rankers, first configure the appropriate flags in the eval_ranker.sh script based on your evaluation strategy (by setting ips_eval and llm_eval flag). Specify the correct path to the preprocessed dataset.
$ bash eval_ranker.shEvaluate the trained model using PO_Eval by using the following flags (to be run using eval_ranker.sh):
python main.py --batch_size=$batch_size --output_path=$output_path \
--output_folder=$model --data_path=$data_path \
--save_epochs=5 --eval_epochs=1 --epochs=$epochs \
--lr_drop=$lr_drop --n_gpus=$n_gpus \
--load_path_reward=$load_path_reward --gain_fn=exp \
--train_ranker --use_doc_feat --load_path=$load_path \
--eval --eval_ultrGenerate inference-based ranking data for online LLM evaluation by using the following flags (to be run using eval_ranker.sh):
python eval_llm.py --batch_size=$batch_size --output_path=$output_path \
--data_path=$data_path --n_gpus=$n_gpus --load_path=$load_path \
--eval_llm --llm_exp --output_folder=$model \
--eval_online --train_rankerEvaluate LLM responses for a given ranking strategy by using the following flags (to be run using eval_ranker.sh):
python eval_llm.py --batch_size=$batch_size --output_path=$output_path \
--data_path=$data_path --n_gpus=$n_gpus --load_path=$load_path \
--eval_llm --llm_exp --output_folder=$model \
--eval_offline --train_rankerDon’t forget to set the required flags and paths in the eval_ranker.sh script before running evaluations.
If you find this repo useful, please cite:
@article{bhatt2025rewardrank,
title={RewardRank: Optimizing True Learning-to-Rank Utility},
author={Bhatt, Gaurav and Thekumparampil, Kiran Koshy and Gangwani, Tanmay and Xiao, Tesi and Sigal, Leonid},
journal={arXiv preprint arXiv:2508.14180},
year={2025}
}