想训练ChatGPT？得先弄明白Reward Model怎么训（附源码）
©作者 | 潘柯宇
研究方向 | 内容理解、信息抽取
prompt: 刚收到货，感觉
output 1: 刚收到货，感觉 有 点 不 符 合 预 期 ，不 好
output 2: 刚收到货，感觉 挺 无 奈 的 送 货 速 度 不 太 行
..现在，我们将利用强化学习（PPO）的方式来对生成模型进行「好评生成」的训练。
每当模型生成一个句子，我们就给出一个相应的得分（reward），用于表征该条生成评论是否是「正向好评」，如下所示：
output 1: 刚收到货，感觉有 点 不 符 合 预 期 ，不 好 -> 0.2 分
output 2: 刚收到货，感觉有 挺 无 奈 的 送 货 速 度 不 太 行 -> 0.1 分
output 3: 刚收到货，感觉有 些 惊 喜 于 货 物 质 量 -> 0.9 分
...
引入判别模型代替人工打分
如果依靠人工为每一个输出打分，这将是一个非常漫长的过程。
如果我们能找到一个判别模型：接收一个句子作为输入，输出这个句子是好评的概率。
那么我们就可以直接利用这个判别模型的输出作为生成句子的 reward。
因此，我们引入另一个「情绪识别模型」来模拟人工给出的分数。
该模型基于网络评论数据集训练，能够对句子进行「正向、负向」的情绪判别，如下所示：
为了保证生成句子的多样性，我们设定了一个 prompt 池，模型会从中随机选择一个 prompt 来进行答案生成：
# prompt池
prompts = [
'刚收到货，感觉',
'这部电影很',
'说实话，真的很',
'这次购物总的来说体验很'
]
...
for _ in range(config['batch_size']):
random_prompt = random.choice(prompts) # 随机选择一个prompt
tokens = gpt2_tokenizer.encode(random_prompt)
batch['tokens'].append(tokens)
batch['query'].append(random_prompt)
query_tensors = [torch.tensor(t).long().to(device) for t in batch["tokens"]]
...
for i in range(config['batch_size']):
gen_len = config['gen_len']
response = gpt2_model.generate(query_tensors[i].unsqueeze(dim=0), # 利用当前选择的prompt生成句子
max_new_tokens=gen_len, **gen_kwargs)
response_tensors.append(response.squeeze()[-gen_len:])
[
'刚收到货，感觉 很 一 般',
'这部电影很 俗 而 且 很 无 趣',
'这次购物总的来说体验很 烂 不 是 我 想 要 的',
...
]
2.2 Reward 评估（Evaluation）
# 情绪识别模型初始化
senti_tokenizer = AutoTokenizer.from_pretrained('uer/roberta-base-finetuned-jd-binary-chinese')
senti_model = AutoModelForSequenceClassification.from_pretrained('uer/roberta-base-finetuned-jd-binary-chinese')
sentiment_pipe = pipeline('sentiment-analysis', model=senti_model, tokenizer=senti_tokenizer, device=pipe_device)
...
texts = [q + r for q,r in zip(batch['query'], batch['response'])] # 将 prompt 和生成的 response 做拼接
pipe_outputs = sentiment_pipe(texts) # 计算正向/负向情感得分
[
0.4,
0.3,
0.3,
...
]模型迭代阶段我们会利用 PPO 进行模型参数的更新，更新代码只用一行：
ppo_trainer.step(query_tensors, response_tensors, rewards) # PPO Update
▲ 模型迭代示意图
loss_p, loss_v, train_stats = self.loss(logprobs, values, rewards, query, response, model_input)
loss = loss_p + loss_v
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
...pg_loss
其中，importance ratio 是指产生同样的 token，在 active actor model 和 reference actor model 下的概率比值，这也是 PPO 模型中的 Importance Sampling 系数。
for t in reversed(range(gen_len)):
nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
delta = rewards[:, t] + self.ppo_params['gamma'] * nextvalues - values[:, t] # 优势函数：r + Vnext - V
lastgaelam = delta + self.ppo_params['gamma'] * self.ppo_params['lam'] * lastgaelam # GAE, 用于平衡 bias 和 variance
advantages_reversed.append(lastgaelam)
advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)
logits, _, vpred = self.model(model_input) # 跑一遍模型，得到句子中每个token被选择的概率
logprob = logprobs_from_logits(logits[:,:-1,:], model_input[:, 1:]) # 将概率取log对数
ratio = torch.exp(logprob - old_logprobs) # log相减，等同于概率相除
pg_losses = -advantages * ratio
value_loss
value_loss 是 PPO 中 critic 的 loss 函数，其目的在于评判每一个 token 被生成后的 value 是多少。
这是因为在 PPO 中需要有一个 critic 网络，为了实现这个效果，我们需要对 GPT 模型进行改造。
我们在 GPT 中加入一个 Value Head，用于将 hidden_size 向量映射到一个 1 维的 value 向量：
class GPT2HeadWithValueModel(GPT2PreTrainedModel):
"""The GPT2HeadWithValueModel class implements a GPT2 language model with a secondary, scalar head."""
def __init__(self, config):
super().__init__(config)
config.num_labels = 1
self.transformer = GPT2Model(config)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
self.v_head = ValueHead(config) # 添加 Value Head
self.init_weights()
...
class ValueHead(nn.Module):
"""The ValueHead class implements a head for GPT2 that returns a scalar for each output token."""
def __init__(self, config):
super().__init__()
self.summary = nn.Linear(config.hidden_size, 1) # (hidden_size -> 1)
...value_loss 就应该等于 Value Head 产生的预测值 v_pred 和真实值 r + v_next 之间的差值：
returns = advantages + values # r + v_next - v + v => r + v_next
logits, _, vpred = self.model(model_input) # 跑一遍语言模型，得到每个 token 的 v_pred
vf_losses1 = (vpred - returns) ** 2 # MSE▲ InstructGPT Reward Model 训练流程
如上图所示，ChatGPT 并不是直接让人工去标注每一句话的真实得分是多少（尽管模型最终要预测的就是每句话的得分），而是让人去对 4 句话按照好坏程度进行「排序」。
通过这个「排序序列」，模型将会学习如何为每一个句子进行打分。
听起来很绕对吧？ 既然最终目的是训练一个句子打分模型，为什么不让人直接打分，而是去标排序序列呢？
今天我们就来好好聊一聊这个非常巧妙的思想。
视频讲解在这里：
1. 香蕉是一种黄色的水果，通常长在树上，是猴子非常喜爱的水果。
2. 香蕉很酸，富含矿物质元素。但如果我们只让标注员对这两个答案进行好坏排序，就能得到统一的结果：
▲ 「相对排序」容易统一
loss = r(A) - r(B) + r(A) - r(C) + r(A) - r(D) + r(B) - r(C) + ... + r(C) - r(D)
loss = -loss
1.买过很多箱这个苹果了，一如既往的好，汁多味甜～ 2.名不副实。 3.拿过来居然屏幕有划痕，顿时就不开心了 4.什么手机啊！一台充电很慢，信号不好！退了！又买一台竟然是次品。
1.一直用沙宣的洗发露！是正品！去屑止痒润发护发面面俱到！ 2.觉得比外买的稀，好似加了水的 3.非常非常不满意，垃圾。 4.什么垃圾衣服，买来一星期不到口袋全拖线，最差的一次购物
...
class RewardModel(nn.Module):
def __init__(self, encoder):
"""
init func.
Args:
encoder (transformers.AutoModel): backbone, 默认使用 ernie 3.0
"""
super().__init__()
self.encoder = encoder
self.reward_layer = nn.Linear(768, 1) # reward layer 用于映射到 1 维 reward
def forward(
self,
input_ids: torch.tensor,
token_type_ids: torch.tensor,
attention_mask=None,
pos_ids=None,
) -> torch.tensor:
"""
forward 函数，返回每句话的得分值。
Args:
input_ids (torch.tensor): (batch, seq_len)
token_type_ids (torch.tensor): (batch, seq_len)
attention_mask (torch.tensor): (batch, seq_len)
pos_ids (torch.tensor): (batch, seq_len)
Returns:
reward: (batch, 1)
"""
pooler_output = self.encoder(
input_ids=input_ids,
token_type_ids=token_type_ids,
position_ids=pos_ids,
attention_mask=attention_mask,
)["pooler_output"] # (batch, hidden_size)
reward = self.reward_layer(pooler_output) # (batch, 1)
return reward
def compute_rank_list_loss(rank_rewards_list: List[List[torch.tensor]], device='cpu') -> torch.Tensor:
"""
通过给定的有序（从高到低）的ranklist的reward列表，计算rank loss。
所有排序高的句子的得分减去排序低的句子的得分差的总和，并取负。
Args:
rank_rewards_list (torch.tensor): 有序（从高到低）排序句子的reward列表，e.g. ->
[
[torch.tensor([0.3588]), torch.tensor([0.2481]), ...],
[torch.tensor([0.5343]), torch.tensor([0.2442]), ...],
...
]
device (str): 使用设备
Returns:
loss (torch.tensor): tensor([0.4891], grad_fn=<DivBackward0>)
"""
if type(rank_rewards_list) != list:
raise TypeError(f'@param rank_rewards expected "list", received {type(rank_rewards)}.')
loss, add_count = torch.tensor([0]).to(device), 0
for rank_rewards in rank_rewards_list:
for i in range(len(rank_rewards)-1): # 遍历所有前项-后项的得分差
for j in range(i+1, len(rank_rewards)):
diff = F.sigmoid(rank_rewards[i] - rank_rewards[j]) # sigmoid到0~1之间
loss = loss + diff
add_count += 1
loss = loss / add_count
return -loss
...
global step 10, epoch: 1, loss: -0.51766, speed: 0.21 step/s
global step 20, epoch: 1, loss: -0.55865, speed: 0.22 step/s
global step 30, epoch: 1, loss: -0.60930, speed: 0.21 step/s
global step 40, epoch: 1, loss: -0.65024, speed: 0.21 step/s
global step 50, epoch: 1, loss: -0.67781, speed: 0.22 step/s
Evaluation acc: 0.50000
best F1 performence has been updated: 0.00000 --> 0.50000
global step 60, epoch: 1, loss: -0.69296, speed: 0.20 step/s
global step 70, epoch: 1, loss: -0.70710, speed: 0.20 step/s
...
▲ loss、acc 曲线图
texts = [
'买过很多箱这个苹果了，一如既往的好，汁多味甜～',
'一台充电很慢，信号不好！退了！又买一台竟然是次品。。服了。。'
]
>>> tensor([[10.6989], [-9.2695]], grad_fn=<AddmmBackward>)
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：[email protected]
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
微信扫码关注该文公众号作者
戳这里提交新闻线索和高质量文章给我们。
来源: qq
点击查看作者最近其他文章