仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
来源|知乎
作者|IlikeShuhuaMilk
最近倒腾了一波RLHF，从ColossalAI到TRLX以及DeepSpeed-Chat，最后基于DeepSpeed-Chat成功训练上了自己的模型，最后效果也是肉眼可见的提升。对这一部分进行下总结，包括原理，代码以及踩坑与解决方案。
首先还是解释一下一些概念，从NLP的角度举一些例子。
首先是RL中的Policy，State，Action。
•Policyπ：这是我们需要学习的策略函数，我们使用语言模型来近似这个函数。
•StateS：模型接受的输入句子，这里的话状态是无限的，因为输入可能为任意句子。
•ActionA：根据所处的状态，策略函数做出动作。这里的话就是模型接受上文，预测下一个字符。这里的动作是有限集合，大小为vocabsize。
接下来介绍Reward，Return，Q，V。
•Reward（奖励）：在t时刻，给定上文，预测下一个字符，所获得的奖励。
•return（回报）：在t时刻，给定上文，持续预测下一个字符，获得奖励，直到生成结束符，到达最终状态，获得奖励，所获得的奖励和。
•Q（动作的价值）：在t时刻，给定上文，预测下一个字符，可以获得所有奖励的期望。。
•V（状态的价值）：在t时刻，基于给定上文，可以获得所有奖励的期望。。这里就是在上文状态下，求和所有下一个字符出现概率与字符对应价值的乘积。
PS：这里要注意区分价值和奖励：价值是未来累计奖励的期望。奖励是我们做出该动作后立即获取的收益。
整个过程主要是分为三步：SFT，TrainingRewardModel，RLHF。
这里主要介绍一下TrainingRewardModel和RLHF。
RewardModel相较于原始的SFTModel，在后面加上了一个valuehead，valuehead是一个Linear，输入维度为模型的hidden_dim，输出维度为1，输出表示模型预测每一字符获取的得分。DeepSpeed-Chat中使用最后一个字符的得分作为整个response的得分（当然也可以使用整个句子中每个字符的平均分作为整体的得分）。
训练RewardModel是一个排序任务，针对query，输入chosen和rejectedresponse，训练目标尽可能的使得chosen和rejected的差值更大，损失函数为：
以上就是第二步TrainingRewardModel的全部过程，基于rankloss训练了一个打分模型。在第三步强化学习中，reward模型将扮演环境的角色，针对模型预测的字符给出奖励分数。
首先来从整体上看一下这部分（这里就只介绍RL部分，PTX就是加上了预训练任务）：
RLHF基于A2C方法，这一步包含了四个模型:
•ActorModel：由SFT之后的模型初始化而来。作为策略（policy）模型，用于接收上文，做出动作，预测下一个字符。学习完毕之后，我们最终使用的就是这个模型。
•ReferenceModel：和ActorModel同样初始化自SFTModel，训练过程中冻结参数，用于和ActorModel做对比，保证模型不要偏离原始SFTModel太多。
•RewardModel：作为环境（env），训练过程中冻结参数，针对每一个状态，给出奖励分数。
•CriticModel：由RewardModel初始化而来，用于近似价值函数，输入为状态s，估计当前状态的价值V。
接下来梳理一遍训练过程。训练过程整体分为两步：makerexperience和learn。
首先是make_experience，首先在训练数据中抽取一部分query，然后ActorModel生成答案。然后我们依据这条答案获取我们所需要的经验：
•actor_logits:由ActorModel产生，包含对答案所有词的概率分布。
•reference_logits：由ReferenceModel产生，包含对答案所有词语的概率分布，用于和actorlogits进行对比，防止actormodel偏离SFTModel太远。
•reward_score:由RewardModel产生，为当前句子状态下，立即获取的收益分数。
•values：由CriticModel产生，估计当前句子状态下，到完成生成，可以获取的回报。
整体流程如下：
然后在learn的时候，通过所产生的经验进行学习。我们通过ActorModel与CriticModel近似策略函数和价值函数，整体流程如下：
关于learn这部分，详细介绍下CriticModel训练和ActorModel训练过程。
这里对下标new和old做一下说明（最近自己复习一下，看自己的文章，想了半天也没想起，最后被迫又去看代码了，可恶）
问：这里获得的actorprob和之前makeexperience中获得的有什么区别，为什么要区分old和new，包括critic也是？
答：因为模型会有dropout所以会有一些区别。此外如果ppo_epoch大于1，那么actor产生经验之后，在forward，backward之后，也会有不同[ 2]
CriticModel估计当前状态可以获取的价值，也就是我们前面所说的V值。模型的输入为状态s，也就是当前模型生成的句子，输出为状态价值V(s)。
这里有一个重要的近似（基于t时刻的状态s的价值=t时刻获取的奖励+折扣率时刻状态的价值）：
因此，我们在makeexperience的时候产生,使用其估计CriticModel的目标值。为TDtarget，是CriticModel在t+1时刻，对做出的估计。
在learn的时候，我们计算当前CriticModel的最新预测值,用其估计。由于中包含真实观测奖励，因此更为接近真实目标，我们需要将进行对齐，使得CriticModel预测的value更接近真实值。可以使用MSE损失来优化CriticModel，这部分损失表示为：
这里举一个例子：
早上起床去公司，地图预计通行时间是50min，然后走了10分钟到马连洼，地图预计用时还需要三十分钟。那么一开始预估的总体时间是50min，而基于一部分真实观测的总体时间是10+30=40min，那么我们最终的误差就是(10+30-50)**2啦。
Actormodel基于当前输入上文s，预测下一个字符的概率分布。
基于策略梯度算法，我们期望最大化状态价值，因此目标函数
但是基于这种方法，我们很多时候无论采取何种动作，我们得到的预期回报只要大于0，我们模型就会向这个动作方向更新。
举个例子，如果我们模型有“上，下，左，右”四个动作，分别有累计奖励“10，20，30，40”，我们做出任意动作，都会获取正向的累计奖励，因此模型也会向这个动作更新参数。而实际上，我们累计奖励的平均期望为25，对于动作“上，下”我们都应该受到惩罚。
因此我们需要引入一个基线来作为参考，在A2C方法中，引入基线为状态价值，为状态s的平均预期回报。因此我们的目标函数变为：
其中为优势值，
最后，我们在通过PPO算法来进一步优化。关于PPO详细讲解，可以看我的文章PPO详解。PPO的想法是希望限制训练过程中对policy进行的更改，避免policy进行过大的更新，来提升policy训练的稳定性。基于PPO算法，我们的目标函数变为：
其中。ratio用于比较当前策略和之前旧策略之间的变化，同时将其限制在区间之中。
以上就是第三步的核心内容，RL过程整体分为两步，makeexperience和learn。我们首先采样数据，然后生成结果，RewardModel给出环境的奖励，以及CriticModel对结果进行评判。之后我们依据所获取的经验来对模型进行更新。
这个步骤基本就是全部按照DeepSpeed-Chat代码来了，使用cohere-zh的数据，大约构造了2w条chosen-rejectedpair用于训练。最后训练训练了一个epoch，在验证集上准确率到了0. 79左右。
在这个步骤中，从跑通到收敛还是有不少麻烦，分享一些比较重要的点：
不过目前DeepSpeed-Chat也没有解决，需要关闭HybridEngine进行训练。
DeepSpeed-Chat还有一个很严重的问题就是，在makeexperience的时候，强制ActorModel生成到最大长度（设置max_length=min_length=max_min_length），这样子导致模型生成偏差很大。对于一个简单的问题，模型可能本来生成简单的一句话就可以完美回答了，但是却必须强制生成到最大长度，这样训练的模型和我们实际用起来的模型是有区别的。
对于这个问题，可以通过修改generate中的参数eos_token_id来解决。设置一个虚假的结束符，然后模型就可能生成正常的结束符。然后我们在构建attention_mask来遮蔽掉结束符后面的回答。例如：
通过以上两步基本就可以跑通流程了，但是训练过程中还遇到了一个比较大的问题，就是CriticLoss并不收敛，越来越大。
具体的原因是：训练后期，随着模型输出答案越来越好，我们的reward值也会越来越高，导致我们最终累计回报return的区间也会越来越大。而前面说过，我们CriticLoss是一个MSE损失，因此训练后期，随着return的估计范围越来越大，CriticModel就越难估计。在我训练的过程中，一开始return的范围是在3-4左右，训练后期涨到了18-20，因此我们需要想点办法来约束一下我们的return。
首先的超参的调节，γ参数为折扣回报率，DeepSpeed-Chat中初始设置为1，可以将其调小一些来缓解。其次的话，使用rewardscale这一trick帮助也非常大。
通过以上这些步骤，基本上我们正常训练了，模型最后也能看到一些效果，但是需要取得更好的效果，我们就需要引入一些trick了。
NormalizationofAdvantages
•将Advantage进行归一化：adv=(adv-mean)/std
•在mini-batch上进行
•没有指标的量化，从我个人看结果而言，感觉提升不大
OverallLossandEntropyBonus
•为了提高算法的探索能力，在actor的loss中增加一项策略熵，并乘以一个系数entropy_coef，使得在优化actor_loss的同时，让策略的熵尽可能大。一般我们设置entropy_coef=0. 01
•总体损失变为：loss=policy_loss-entropyentropy_coefficient+value_lossvalue_coefficient
•没有指标的量化，从我个人看结果而言，感觉没有太多提升
RewardScale
•对reward进行缩放，将reward除以标准差
•从训练log来看，对稳定criticloss效果很好，毕竟将reward进行缩放之后，降低了return的估计区间。
•没有指标的量化，从我个人看结果而言，提升很大，推荐使用。
关于其他的一些trick，如学习率衰减、梯度裁剪、ValueClipping等，本身框架就包含了，就不进行特别说明了。当然，这些trick在不同的数据或者模型上都会有不同的效果，需要自己进行探索。
通过以上这些方法，最后也顺利的完成了训练，希望对大家能有所帮助，也希望大家能分享自己有用的经验与理解，关于RL自己了解的也很少，还需要更多的学习。
