license: cc-by-nc-4.0
language:
- en
pipeline_tag: text-classification
tags:
- pytorch
- reward_model
- transformers
- RLHF
library_name: transformers
This is part of the Chai reward-model series, using the GPT2 architecture with a classification head, optimising for a user accepting the completion generated by the base model.
Its training dataset consists of purely user-generated content retry_and_continue_50m_reward_model, where a user has the option to decline the generated response via the retry button or end the conversation.
Model details
- Developed by Chai Research
- Model type: Transformer-based Classification Model
- Language: English
- License: cc-by-nc-4.0
- Contact: for general correspondence, please email [email protected]
Uses and limitations
Intended use
This reward model was developed primarily for commercial purposes. It learns an inner representation of response quality rated by humans that can be used to conduct best-of-N sampling and Reinforcement Leanring with the PPO framework.
In addition to scientific uses, you may also further fine-tune and adapt this reward model for deployment, as long as your use is in accordance with the Creative Commons Attribution Non Commercial 4.0 (cc-by-nc-4.0) license, i.e. non-commercial use. This model works with the Transformers Library. If you decide to this pre-trained reward model as a basis for your fine-tuned model, please note that you need to conduct your own risk and bias assessment.
Out-of-scope use
This reward model is not intended for deployment as-is. It is not a product and cannot be used for human-facing interactions without supervision.
This model has not been optimised for common reward-model objectives such as harmfulness, truthfulness and helpfulness, it is only trained based on user actions present on the Chai mobile app platform. Therefore, this model will not rank responses appropriately when evaluating on common open-sourced datasets. All base model responses within the training data were generated using an in-house variant of GPT-J, therefore the model performance may degrade when the input is generated using other language models.
How to use
This reward model can be loaded using the AutoModelForSequenceClassification functionality, with a GPT2 tokenizer where the pad_token_id is set to the EOS token id, padding sides need to be set according to the configurations used during model training.
from transformers import AutoTokenizer, AutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForSequenceClassification.from_pretrained("ChaiML/gpt2_base_retry_and_continue_12m_reward_model")
tokenizer.pad_token_id = 50256
tokenizer.truncation_side = ‘left’
tokenizer.padding_side = ‘right’
tokens = self.eval_tokenizer(candidates, return_tensors='pt', return_attention_mask=True, padding='longest', truncation=True, max_length=256)
reward = model(**tokens).logits
Model training
Training dataset
This model was trained by randomly sampling 12 million rows out of the ChaiML/retry_and_continue_50m_reward_model dataset. The original dataset contains over 50 million rows of completions (chatbot responses), along with number of remaining user messages within their corresponding conversations and whether the user pressed the "retry" button (where the completion is rejected and resampled). The model which was used to generate these completions is a in-house variant of GPT-J, with the following sampling parameters:
| Parameters | Value |
|---|---|
| temperature | 0.72 |
| repetition_penalty | 1.13125 |
| max_new_tokens | 64 |
| top_p | 0.725 |
| top_k | 0 |
| eos_token_id | 198 |
| do_sample | True |
Training procedure
The gpt2_base_retry_and_continue_12m_reward_model was trained using a gpt2 base model and a classification head with single output. Binary Cross Entropy loss was used. The model was trained on 4xA40 GPUs, 16 per device batch size and gradient accumulation of 1 (therefore the effective batch size is 64), with 1e-5 learning rate for 2 epochs for a total of 375,000 steps. Tensor parallelism and pipeline parallelism were used to distribute the model across GPUs. For evaluation metrics used during training, please see our Weights and Biases Log.
BibTeX entry
To cite this model:
@misc{
author = {Chai Research, Irvine, Boubert, Raina, Liusie, Mudupalli, Korshuk, Liu, Cremer, Assassi, C. Beauchamp, Lu, Rialan, W. Beauchamp},
title = {{Rewarding chatbots for real-world engagement with millions of users}},
howpublished = {\url{https://arxiv.org/abs/2303.06135}},
year = 2023,
month = Mar
}
If you use this model, we would love to hear about it! Reach out on correspondence email or Discord.
Acknowledgements
This project would not have been possible without the support from members of Seamless Capital
We thank the following authors from the Machine Intelligence Laboratory for their collaboration: