RM Gemma 2B

🚀 報酬モデルの概要

この報酬モデルは、ベースモデル google/gemma-2b-it から学習されています。7Bバージョンの RM-Gemma-7B も参照してください。

学習スクリプトは、https://github.com/WeiXiongUST/RLHF-Reward-Modeling で入手できます。

✨ 主な機能

モデルの詳細

この報酬モデルや報酬モデリングに関する質問があれば、wx13@illinois.edu までメールを送ってください。喜んでお話しします！

データセットの前処理

このモデルは、以下のデータセットの混合データで学習されています。

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

比較ペアの総数は250Kで、以下のデータ選択とクリーニング戦略を実施しています。

• HH-RLHF: すべての基本、棄却サンプリング、オンラインサブセットを使用しますが、選択されたサンプルと棄却されたサンプルが等しいサンプルを削除し、115547件になります。
• SHP: スコア比が2を超えるサンプルのみを使用し、各プロンプトについて1つの比較のみを取り、55916件になります。
• Ultrafeedback: UltraFeedback-Binarized と同様に、全体のスコアではなく細粒度のスコアを使用してサンプルをランク付けします。また、各プロンプトについて、最良のサンプルと残りのサンプルからランダムに選択したサンプルを比較します。最後に、選択されたスコアが等しいペアを削除し、62793件になります。
• HelpSteer: 有用性と正確性の平均を使用してサンプルをランク付けします。また、最良のサンプルと残りのサンプルからランダムに選択したサンプルを比較します。最後に、選択されたスコアが等しいペアを削除し、8206件になります。
• Capybara: 選択されたサンプルと棄却されたサンプルの評価が同じペアを削除し、7562件になります。
• Orca: 選択されたサンプルと棄却されたサンプルの評価が同じペアを削除し、6405件になります。

学習

このモデルは、学習率1e-5、バッチサイズ256、ウォームアップ率0.03のコサイン学習率減衰で1エポック学習されています。学習曲線は以下の通りです。

💻 使用例

基本的な使用法

  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]

📚 ドキュメント

結果

既存の嗜好データセットを収集し、ベンチマークとして使用して、結果として得られた報酬モデルを評価します。

なお、MT-Benchデータセット (lmsys/mt_bench_human_judgments) の場合は、比較結果が引き分けのサンプルを削除しています。Alpacaデータは こちら から取得しています。

参考文献

この報酬モデルは、棄却サンプリング微調整にすぐに使用できます。

@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}
}