Skip to content

tomekkorbak/pretraining-with-human-feedback

Repository files navigation

Pretraining Language Models with Human Preferences

This repo contains the code accompanying the paper Pretraining Language Models with Human Preferences. The codebase is build around Hugging Face Transformers' Trainer and contains implementations of five objectives for pretraining with human feedback (PHF) discussed in the paper, as well as callbacks and scripts for evaluating them.

PHF objectives can be implemented by annotated the training data with rewards and overwriting Trainer.compute_loss to use them as additional training signal. Rewards are provided by an instance of apo.scorers.Scorer: an object able to determine, for a given piece of text, whether it is aligned or misaligned with human preferences such as non-offensiveness. The scorer is also used for evaluating samples from PHF-trained LMs.

The codebase is built around Hugging Face ecosystem and wand (for monitoring and experiment management).

Quickstart

We assume Python 3.9+. To run the training script for MLE on the toxicity task, do:

pip install -r requirements.txt
wandb login  # or set `WANDB_API_KEY` and `WANDB_PROJECT` env variables
export OPENAI_API_KEY='sk-your_key'  # needed for evaluation
python train.py --task configs/toxicity/pretrain.yml --method configs/toxicity/mle.yml

Configuration

The train.py scripts requires paths to two config files: for task and for method. Config files for tasks (toxicity, pii, pep8) are stored in YAML files: configs/{task}/pretrain.yml (for pretraining experiments) and configs/{task}/finetuning.yml (for finetuning). Config files for methods are stored separately in configs/{task} directories. Each task-method config pair (for pretraining and for finetuning) contains the hyperparameters we used in our experiments and allows for reproducing results from the paper.

Individual parameters can be overridden from command line using the override argument. For instance:

python train.py --task configs/toxicity/pretrain.yml --method configs/toxicity/mle.yml --override training.per_device_train_batch_size=8

Tasks

Name Config files Training data Scorer Description
Toxicity configs/toxicity tomekkorbak/pile-detoxify DetoxifyToxicityScorer Misalignment score is the probability of toxicity according to detoxify
PII configs/pii tomekkorbak/pile-pii-scrubadub PIIScorer Misalignment score is the number of PIIs (e.g. names, URLs) per character, according to scrubadub
PEP8 configs/pep8 kejian/codeparrot-train-more-filter-3.3b-cleaned PEP8Scorer Misalignment score is the number of PEP8 violations per character, according to pycodestyle

Objectives

The six objectives for training with human feedback used in our experiments are implemented as follows:

Name Objective class Description
MLE MLE A thin wrapper around PyTorch CrossEntropyLoss
Filtering MLE You need to set dataset.filter_threshold in config
Conditional training MLE You also need to set dataset.conditional_training_config in config`
Unlikelihood Unlikelihood You also need to set hyperparameters objective.score_threshold and objective.alpha
AWR AWR You also need to set hyperparameters objective.alpha and objective.beta
RWR AWR A special case of AWR with objective.alpha=1

Pretrained models

The models pretrained in our experiments are available on HugginFace Hub:

Objective Toxicity PEP8 PII
MLE tomekkorbak/goofy_pasteur kejian/mighty-mle tomekkorbak/nervous_wozniak
Filtering median tomekkorbak/amazing_shannon kejian/mighty-filtering tomekkorbak/cocky_carson
Conditional tomekkorbak/hungry_saha kejian/mighty-conditional tomekkorbak/boring_mcclintock
UL tomekkorbak/nifty_banach kejian/mighty-ul tomekkorbak/affectionate_wescoff
AWR tomekkorbak/upbeat_ramanujan kejian/vigor-awr tomekkorbak/confident_knuth
RWR tomekkorbak/keen_clarke kejian/mighty-rwr tomekkorbak/gifted_hugle

Metrics

On each evaluation step, apo.callbacks.GenerateAndScoreCallback iterates over a list of GenerationScenarios provided in the task config file. For each scenario, num_samples samples are generated and the following wandb metrics are computed:

  • score, average misalignment (across num_samples samples) of the generated samples assigned by the scorer
    • score_max@25, average maximum score in 25 samples (similar to expected maximum toxicity in the RealToxicityPrompts paper)
  • current_samples, a wandb.Table of samples together with their prompts (if any) and scores

In addition to scoring LM samples, we use apo.callbacks.KLGPT3Callback to estimate KL of the current LM from GPT-3. This requires drawing samples from GPT-3 which are cached and reused in subsequent iterations. |

Codebase structure

.
├── apo
│   ├── callbacks.py  # callbacks implementing the evaluation pipeline 
│   ├── dataset_wrappers.py  # an iterable for streaming blocks of tokens for training
│   ├── kl_gpt3.py  # logic for measuring KL from GPT-3
│   └── metrics.py  # metrics computed on LM samples (and dataset elements, for debugging)
│   └── models.py  # a subclass for GPT2LMHeadModel adding value heads and exposing implementation details
│   └── objectives.py  # classes implementing loss functions
│   ├── scorer_utils.py
│   ├── scorers.py  # classes for scoring LM samples and dataset elements
│   └── trainer.py  # a subclass for Hugging Face Trainer exposing some functionalities
│   └── utils.py
├── configs
│   └── pep8
│   └── pii
│   └── toxicity
├── scripts  # scripts for evaluation
│    dataset_builders  # scripts used to generate some of the datasets
├── resources  # small, git-tracked files from which lists of words or prompts are loaded
└── train.py  # the main training script

Citing

@misc{https://doi.org/10.48550/arxiv.2302.08582,
  doi = {10.48550/ARXIV.2302.08582},
  url = {https://arxiv.org/abs/2302.08582},
  author = {Korbak, Tomasz and Shi, Kejian and Chen, Angelica and Bhalerao, Rasika and Buckley, Christopher L. and Phang, Jason and Bowman, Samuel R. and Perez, Ethan},
  keywords = {Computation and Language (cs.CL), Machine Learning (cs.LG), FOS: Computer and information sciences, FOS: Computer and information sciences},
  title = {Pretraining Language Models with Human Preferences},
  publisher = {arXiv},  
  year = {2023},
  copyright = {Creative Commons Attribution 4.0 International}
}