maywell / RM-Gemma-2B

Reward Model Overview

The reward model is trained from the base model google/gemma-2b-it . See the 7B version RM-Gemma-7B .

The training script is available at https://github.com/WeiXiongUST/RLHF-Reward-Modeling .

Model Details

If you have any question with this reward model and also any question about reward modeling, feel free to drop me an email with [email protected] . I would be happy to chat!

Dataset preprocessing

The model is trained on a mixture of

• HH-RLHF
• SHP
• UltraFeedback
• Capybara
• HelpSteer
• Orca

The total number of the comparison pairs is 250K, where we perform the following data selection and cleaning strateges:

• HH-RLHF: we use all the base, rejection sampling, and online subsets but delete the samples whose chosen == rejected, leading to 115547;
• SHP: we only use the samples with score ratio > 2, for each prompt, we only take 1 comparison, leading to 55916;
• Ultrafeedback: similar to UltraFeedback-Binarized , we use the fine-grained score instead of the overall one to rank samples. Meanwhile, for each prompt, we take the best one v.s. random chosen one in the remaining samples. Finally, we delete the selected pairs with equal scores, leading to 62793.
• HelpSteer: we use the mean of helpfulness and correctness to rank samples. Meanwhile, we take the best sample v.s. the random chosen one in the remaining samples. Finally, we delete the selected pairs with equal scores, leading to 8206;
• Capybara: we delete the pairs whose chosen and rejected samples are of the same rating, leading to 7562;
• Orca: we delete the pairs whose chosen and rejected samples are of the same rating, leading to 6405.

Training

We train the model for one epoch with a learning rate of 1e-5, batch size 256, cosine learning rate decay with a warmup ratio 0.03. We present the training curve as follows.

Uses

  from transformers import AutoTokenizer, pipeline
  rm_tokenizer = AutoTokenizer.from_pretrained("weqweasdas/RM-Gemma-2B")
  device = 0 # accelerator.device
  rm_pipe = pipeline(
      "sentiment-analysis",
      model="weqweasdas/RM-Gemma-2B",
      #device="auto",
      device=device,
      tokenizer=rm_tokenizer,
      model_kwargs={"torch_dtype": torch.bfloat16}
  )

  pipe_kwargs = {
      "return_all_scores": True,
      "function_to_apply": "none",
      "batch_size": 1
  }

  chat = [
   {"role": "user", "content": "Hello, how are you?"},
   {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
   {"role": "user", "content": "I'd like to show off how chat templating works!"},
  ]

  test_texts = [tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=False).replace(tokenizer.bos_token, "")]
  pipe_outputs = rm_pipe(test_texts, **pipe_kwargs)
  rewards = [output[0]["score"] for output in pipe_outputs]

Results

We collect the existing preference datasets and use them as a benchmark to evaluate the resulting reawrd model.

Note that for MT-Bench dataset (lmsys/mt_bench_human_judgments), we delete the samples with tie as the comparison results. The Alpaca data is from Here .

Reference

To be added. The reward model may be readily used for rejection sampling finetuning (

@article{dong2023raft,
  title={Raft: Reward ranked finetuning for generative foundation model alignment},
  author={Dong, Hanze and Xiong, Wei and Goyal, Deepanshu and Pan, Rui and Diao, Shizhe and Zhang, Jipeng and Shum, Kashun and Zhang, Tong},
  journal={arXiv preprint arXiv:2304.06767},
  year={2023}
}