This repository contains the code related to "Reward Model Ensembles Help Mitigate Overoptimization". In particular, it provides the following:

• An easy way to perform LLM instruction supervised fine-tuning (SFT) as laid out in the paper
• An easy way to create and train one or more reward models to be used in reward model ensembles (or on their own)
• A best-of-n inference pipeline, both with individual reward models and ensembles
• A PPO-based RLHF training pipeline, both for individual reward models and ensembles
• Some crucial models and datasets to facilitate future work and experiments akin to those in the paper

We hope you can find the code, models, and datasets provided helpful for your own research.

Note: SFT, RM, and PPO training backends are based on the Open-Assistant and trlx libraries. Base models are taken from the open-source Pythia suite.

git clone https://github.com/tlc4418/llm_optimization.git
cd llm_optimization
pip install -e .

We provide the following models and datasets on HuggingFace to promote and faciliatte reproducing experiments and future work.

• tlc4418/pythia_1.4b_sft_policy: 1.4B Pythia model after SFT on the AlpacaFarm "sft" split. Used as the initial policy model in our experiments.
• tlc4418/pythia_70m_sft: 70M Pythia model after SFT on the AlpacaFarm "sft" split. Used as the base model for most of our reward model experiments (before RM training). Multiple reward models can be created from this during reward model training and it is relatively inexpensive to do so.

• tlc4418/1.4b-policy_preference_data_gold_labelled: Preference dataset using labels from the AlpacaFarm dataset, generated answers from a 1.4b fine-tuned Pythia policy model, and labelled using the AlpacaFarm "reward-model-human" as a gold reward model. Used to train reward models.
• tlc4418/gold_labelled_gens: Dataset of 12600 answer generations from the 1.4B Pythia SFT model (provided above), using the AlpacaFarm dataset "val" split, and labelled with the AlpacaFarm "reward-model-human" to give "gold" scores. Used for best-of-n inference. This dataset is particularly expensive to geenrate and we hope it can help other with future work.
• We provide wrappers and functionality for using these datasets, as well as different parts of the AlpacaFarm dataset. See the dataset guide for details on this and how to use your own datasets.

You can easily perform SFT on any HuggingFace or local language models by creating a new entry in the SFT config. An example is given for a 70M Pythia model. The following are crucial fields:

Any default hyperparameters can also be overwritten. The 70M Pythia model sample config entry shows a few examples.

Once the config has been set, training can be started with the following command, using the new entry name you created (e.g. "pythia-70m"):

accelerate launch --config_file configs/accelerate_config.yaml src/sft/trainer_sft.py --configs defaults {your_sft_config_entry}

Very similar to SFT, you can perform reward model training on any registered HuggingFace or local reward model by creating a new entry in the RM config. The following are crucial fields:

Again, default hyperparameters can be overwritten.

Reward model training can then be started with this new config entry (e.g. "rm-pythia-44m"). The --rng_seed argument is optional and will otherwise be sourced from the RM config. Command to launch training:

accelerate launch --config_file configs/accelerate_config.yaml src/reward_modeling/training/trainer_rm.py --configs defaults_rm {your_rm_config_entry} --rng_seed {your_choice_seed}

Best-of-n inference (also known as re-ranking) can be performed using a base policy model (post-SFT) and one or more trained reward models, as follows:

python src/bon/run_bon_pipeline.py {your_reward_models_path}

This command is customizable with the following arguments and options:

A help function is also provided (--help) which will display a condensed version of the above, including the appropriate flag names for specifying each parameter.

Our implementation uses an unbiased estimator (see paper for details) for robust and unbiased results.

The relevant result of running this command will be a "bon_sampled_results.json" results file for each run (one for each individual seed and one for each ensemble if desired). This file contains a list of dictionary entries, where each entry contains the sampled index n, the proxy reward model score at that n, and the corresponding gold reward model score at that n, as follows:

[
    {
        "n": {some_int},
        "proxy_score": {proxy_score},
        "gold_score": {gold_score}
    },

    ...
]

These data points can then be used to plot the BoN performance of different policies according to both proxy and gold reward model score, as a function of n.

RL training with PPO can be performed using a base policy model (post-SFT) and one or more trained reward models. The policy will be trained using PPO and the reward models as a reward function. We provide options for both single reward models and ensembles, along with conservative optimization techniques. Example configs entries for this are provided in the RL config. We detail some of the following important fields (which will likely be changed the most) for a new config entry:

Many default hyperparameters can also be overwritten, both in the RL config and more specific PPO ones in the PPO config file

PPO training (and subsequent gold reward model evaluation) can be started with the following command:

accelerate launch --config_file configs/accelerate_config.yaml src/ppo/trainer_rl.py --configs defaults defaults_rlhf {your_rm_config_entry}

Gold reward model evaluation is performed after training is completed so as to avoid loading the large gold reward mdoel during training. KL divergence from the initial policy is also recorded at each evaluation step, such that the evaluation results in the "eval/" output folder can be used to plot both proxy and gold reward model score either as a function of steps or of KL divergence. Please see the AlpacaFarm respository for instructions on downloading the 7B reward model used as the gold reward model in our paper (reward-model-human)

Simply use the provided* tlc4418/pythia_1.4b_sft_policy and tlc4418/pythia_70m_sft models.

Alternatively, you could run your own SFT training for the policy model and base reward model, following the SFT instructions above.

Train a reward model from the base reward model (post-SFT), with any rng seed (here "1"):

accelerate launch --config_file configs/accelerate_config.yaml src/reward_modeling/training/trainer_rm.py --configs defaults_rm rm-pythia-44m --rng_seed 1

Run BoN sampling for the trained reward model:

python src/bon/run_bon_pipeline.py models/rm-pythia-44m_seed{seed} --seeds 1

For PPO, run the following command:

accelerate launch --config_file configs/accelerate_config.yaml src/ppo/trainer_rl.py --configs defaults defaults_rlhf pythia_rlhf_individual

Here we set the config to already contain the trained reward model under the rank_config's model_names, but you will want to change this to match yours when customizing this process.

Simply use the provided tlc4418/pythia_1.4b_sft_policy and tlc4418/pythia_70m_sft models.

Alternatively, you could run your own SFT training for the policy model and base reward model, following the SFT instructions above.

Train as many (k) reward models as you want from the base reward model (post-SFT). Just run the following k times, changing the seed each time (e.g. 1-5):

accelerate launch --config_file configs/accelerate_config.yaml src/reward_modeling/training/trainer_rm.py --configs defaults_rm rm-pythia-44m --rng_seed {seed}

Assuming you have trained 5 reward models above, this will run BoN sampling for all 5 models and the 3 ensembles types (mean, WCO, UWO):

python src/bon/run_bon_pipeline.py models/rm-pythia-44m_seed{seed} --seeds 1,2,3,4,5 --ensembles

For PPO, you will want to run the following command for each ensemble you want to train, making sure the appropriate conservative optimization objective and uwo_weight (if applicable) are set in the RL config:

accelerate launch --config_file configs/accelerate_config.yaml src/ppo/trainer_rl.py --configs defaults defaults_rlhf pythia_rlhf_ensemble

Here we set the config to already contain the 5 trained reward models under the rank_config's model_names, but you will want to change these to match yours when customizing this process.

If you use the models, data, or code in this repo, please consider citing our work.

@article{coste2023reward,
  title={Reward model ensembles help mitigate overoptimization},
  author={Coste, Thomas and Anwar, Usman and Kirk, Robert and Krueger, David},
  journal={arXiv preprint arXiv:2310.02743},
  year={2023}
}