LoRe is a lightweight, modular codebase for learning personalized reward models from preference data in multi-user environments. It supports both joint reward learning and few-shot personalization, and is built with extensibility in mind.
LoRe currently supports experiments on three benchmark datasets:
- Reddit TLDR: headline preference modeling
- PRISM: multi-turn dialogue response preferences
- PersonalLLM: user-personalized reward modeling for open-ended language model responses
- Low-rank reward model learning across users
- Few-shot personalization with new users
- Evaluation on seen/unseen users and prompts
- Modular dataset and optimizer configuration
LoRe requires Python 3.8+ and PyTorch. To install dependencies:
pip install -r requirements.txtLoRe/
βββ utils.py # Core training, optimization, and evaluation helpers
βββ RedditTLDR/ # TLDR dataset scripts
β βββ prepare.py # Preprocess the dataset
β βββ train_basis.py # Train shared reward model and user weights
β βββ vary_fewshot.py # Evaluate few-shot personalization
βββ PRISM/ # PRISM dataset scripts
β βββ prepare.py
β βββ train_basis.py
β βββ vary_fewshot.py
βββ PersonalLLM/ # PersonalLLM dataset scripts
β βββ prepare.py
β βββ train_basis.py
The following functions are central to training and evaluating personalized reward models. You may want to modify these if youβre extending LoRe:
LoRe(...): Class modeling shared reward modelV(linear transformation on fixed embeddings) and user-specific weightsWLoRe_regularized(...): Class modeling shared reward modelV(linear transformation on fixed embeddings), cosine similarity regularization to base model, and user-specific weightsWPersonalizeBatch(...): Class to model weights for new users
run(...): Runs the entire pipeline with 1. Learning the basis rewards, 2. Evaluation on seen users, 3. Fewshot learning on new users, 4. Evaluation on new users. The input K_list can be modified to specify the number of basis. 0 is the reference model, and 1 is the BT model.run_regularized(...): Runs the entire pipeline as run(...) but with regularization on the final layer
Below are instructions for each dataset folder, following a consistent workflow:
- Prepare the dataset if required.
- Train the reward model basis using joint learning.
- Evaluate few-shot personalization with unseen users.
Inside the RedditTLDR/ directory:
prepare.py: preprocesses the TLDR datasettrain_basis.py: trains the shared reward model and user weightsvary_fewshot.py: evaluates few-shot personalization performance
Example usage:
cd LoRe/RedditTLDR
python prepare.py # only needed once
python train_basis.py
python vary_fewshot.pyInside the PersonalLLM/ directory:
prepare.py: prepares model response data and user splitstrain_basis.py: learns reward basis across usersvary_fewshot.py: evaluates few-shot generalization
Example usage:
cd LoRe/PersonalLLM
python prepare.py # only needed once
python train_basis.py
python vary_fewshot.pyInside the PRISM/ directory:
prepare.py: prepares PRISM dialogue in chat formatgenerate-prepare-embeddings.py: prepares PRISM embeddingstrain_basis.py: runs reward model training and evaluationeval_rm2.py: runs learn reward basis models on the reward bench 2 dataset
Example usage:
cd LoRe/PRISM
python prepare.py # only needed once
python generate-prism-embeddings.py # only needed once
python train_basis.py # train the model for a list of ranks, default regularization is specified
python eval_rb2.py --rm_head "path to saved final layer (V) weights" # evaluate learnt reward basis on RewardBench2 to avoid overfitting to PRISMExperiments on a much larger community alignment dataset for scalable, multi-user preference learning.
See the CONTRIBUTING file for how to help out.
CC-BY-NC 4.0 licensed, as found in the LICENSE file.
If you use this codebase, please cite us:
@misc{bose2025lorepersonalizingllmslowrank,
title={LoRe: Personalizing LLMs via Low-Rank Reward Modeling},
author={Avinandan Bose and Zhihan Xiong and Yuejie Chi and Simon Shaolei Du and Lin Xiao and Maryam Fazel},
year={2025},
eprint={2504.14439},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2504.14439},
}