🚀 HH - RLHF奖励模型
本项目提供了一个基于LMFlow框架训练的奖励模型,该模型针对[HH - RLHF数据集](Dahoas/full - hh - rlhf)(仅有用部分),以[openlm - research/open_llama_3b](https://huggingface.co/openlm - research/open_llama_3b)为基础模型进行训练,具有良好的泛化能力,可用于多种数据集。
🚀 快速开始
本奖励模型是基于LMFlow框架训练的,用于[HH - RLHF数据集](Dahoas/full - hh - rlhf)(仅有用部分),基础模型为[openlm - research/open_llama_3b](https://huggingface.co/openlm - research/open_llama_3b)。
如需模型卡元数据参考,请查看规范:https://github.com/huggingface/hub - docs/blob/main/modelcard.md?plain = 1
文档/指南:https://huggingface.co/docs/hub/model - cards
⚠️ 重要提示
感谢您对本奖励模型的关注!我们推荐您使用 [weqweasdas/RM - Gemma - 2B](https://huggingface.co/weqweasdas/RM - Gemma - 2B) 。
✨ 主要特性
数据集预处理
HH - RLHF - Helpful数据集的训练集包含112K个比较样本,测试集包含12.5K个比较样本。我们首先将数据集中的 \n\nHuman
和 \n\nAssistant
分别替换为 ###Human
和 ###Assistant
。
然后,按以下方式划分数据集:
- SFT数据集:112K个训练样本 + 测试集中的前6275个样本,仅使用选中的回复;
- 奖励建模训练集:112K个训练样本 + 测试集中的前6275个样本,使用选中和拒绝的回复;
- 奖励建模测试集:原始测试集中的最后6226个样本。
训练
为更高效地使用数据,我们将文本拼接并分割成1024大小的块,而非按最长文本(每批)进行填充。然后在SFT数据集上对基础模型进行两个轮次的微调,学习率为2e - 5,采用线性衰减调度。
我们以5e - 6的学习率进行1个轮次的奖励建模,并采用线性衰减调度,因为模型在超过1个轮次时容易过拟合。我们丢弃长度超过512个标记的样本,因此奖励建模的训练集约有106K个样本,测试集约有5K个样本。
两个阶段均使用bf16,且不使用LoRA。
最终模型的评估损失为0.5,评估准确率为75.48%。(请注意,[HH - RLHF数据集](Dahoas/full - hh - rlhf)可能存在数据泄露问题。)
泛化能力
我们进一步测试了奖励模型的泛化能力,在另一个研究项目中进行了另一轮训练(超参数相同)。我们在开放助手数据集和聊天机器人数据集上测试了准确率,并将该奖励模型与直接在这两个数据集上训练的奖励模型进行了比较。结果如下:
数据集训练/测试情况 |
开放助手 |
聊天机器人 |
hh_rlhf |
开放助手 |
69.5 |
61.1 |
58.7 |
聊天机器人 |
66.5 |
62.7 |
56.0 |
hh_rlhf |
69.4 |
64.2 |
77.6 |
可以看出,在HH - RLHF上训练的奖励模型在开放助手和聊天机器人数据集上达到了相当甚至更好的准确率,即使它并非直接在这些数据集上训练。因此,该奖励模型也可用于这两个数据集。
💻 使用示例
基础用法
rm_tokenizer = AutoTokenizer.from_pretrained("weqweasdas/hh_rlhf_rm_open_llama_3b")
rm_pipe = pipeline(
"sentiment-analysis",
model="weqweasdas/hh_rlhf_rm_open_llama_3b",
device="auto",
tokenizer=rm_tokenizer,
model_kwargs={"torch_dtype": torch.bfloat16}
)
pipe_kwargs = {
"return_all_scores": True,
"function_to_apply": "none",
"batch_size": 1
}
test_texts = [
"###Human: My daughter wants to know how to convert fractions to decimals, but I'm not sure how to explain it. Can you help? ###Assistant: Sure. So one way of converting fractions to decimals is to ask “how many halves are there?” and then write this as a decimal number. But that's a little tricky. Here's a simpler way: if a fraction is expressed as a/b, then it's decimal equivalent is just a/b * 1.0 So, for example, the decimal equivalent of 1/2 is 1/2 * 1.0 = 0.5.",
"###Human: I have fresh whole chicken in my fridge. What dish can I prepare using it that will take me less than an hour to cook? ###Assistant: Are you interested in a quick and easy recipe you can prepare with chicken you have on hand, or something more involved? In terms of both effort and time, what are you looking for?"]
pipe_outputs = rm_pipe(test_texts, **pipe_kwargs)
rewards = [output[0]["score"] for output in pipe_outputs]
RAFT示例
我们使用RAFT测试奖励模型,并以EleutherAI/gpt - neo - 2.7B作为起始检查点。
每次迭代时,我们从HH - RLHF数据集中采样2048个提示,对于每个提示,当前模型生成K = 8个回复,并选择奖励最高的回复。然后,在这个选中的集合上微调模型以得到新模型。我们报告的学习曲线如下:

我们还使用LLaMA - 7B模型进行了实验,但首先使用HH - RLHF数据集中选中的回复对基础模型进行了1个轮次的微调,学习率为2e - 5。RAFT的超参数与GPT - Neo - 2.7B相同,奖励曲线如下:

📚 详细文档
参考
如果您发现此模型有用,请使用以下BibTeX引用我们的框架和论文:
@article{diao2023lmflow,
title={Lmflow: An extensible toolkit for finetuning and inference of large foundation models},
author={Diao, Shizhe and Pan, Rui and Dong, Hanze and Shum, Ka Shun and Zhang, Jipeng and Xiong, Wei and Zhang, Tong},
journal={arXiv preprint arXiv:2306.12420},
year={2023}
}
@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}
}
如果您对该奖励模型或奖励建模有任何疑问,请随时发送电子邮件至wx13@illinois.edu,我很乐意交流!