LLM面面观之RLHF平替算法DPO
1. 背景
最近本qiang~老看到一些关于大语言模型的DPO、RLHF算法，但都有些云里雾里，因此静下心来收集资料、研读论文，并执行了下开源代码，以便加深印象。
此文是本qiang~针对大语言模型的DPO算法的整理，包括原理、流程及部分源码。
2. DPO vs RLHF
RLHF vs DPO
上图左边是RLHF算法，右边为DPO算法，两图的差异对比即可体现出DPO的改进之处。
1. RLHF算法包含奖励模型(reward
model)和策略模型(policy model，也称为演员模型，actor model)，基于偏好数据以及强化学习不断迭代优化策略模型的过程。
2. DPO算法不包含奖励模型和强化学习过程，直接通过偏好数据进行微调，将强化学习过程直接转换为SFT过程，因此整个训练过程简单、高效，主要的改进之处体现在于损失函数。
PS：
1. 偏好数据，可以表示为三元组(提示语prompt, 良好回答chosen, 一般回答rejected)。论文中的chosen表示为下标w(即win)，rejected表示为下标l(即lose)
2. RLHF常使用PPO作为基础算法，整体流程包含了4个模型，且通常训练过程中需要针对训练的actor model进行采样，因此训练起来，稳定性、效率、效果不易控制。
1) actor model/policy
model: 待训练的模型，通常是SFT训练后的模型作为初始化
2) reference
model: 参考模型，也是经SFT训练后的模型进行初始化，且通常与actor model是同一个模型，且模型冻结，不参与训练，其作用是在强化学习过程中，保障actor model与reference model的分布差异不宜过大。
3) reward model:奖励模型，用于提供每个状态或状态动作对的即时奖励信号。
4) Critic model:作用是估计状态或状态动作对的长期价值，也称为状态值函数或动作值函数。
3. DPO算法仅包含RLHF中的两个模型，即演员模型(actor
model)以及参考(reference model)，且训练过程中不需要进行数据采样。
4. RLHF可以参考附件中的引文
3. DPO的损失函数
DPO的损失函数
如何将RLHF的Reward model过程简化为上式，作者花了大量篇幅进行了推导，感兴趣的读者可以参考附件DPO的论文。
DPO算法的目的是最大化奖励模型(此处的奖励模型即为训练的策略)，使得奖励模型对chosen和rejected数据的差值最大，进而学到人类偏好。
上式的后半部分通过对数函数运算规则，可以进行如下转化。
Loss公式转化
转化后的公式和源代码中的计算函数中的公式是一致的。
其中左半部分是训练的policy模型选择chosen优先于rejected，右半部分是冻结的reference模型选择chosen优先于rejected，二者的差值可类似于KL散度，保障actor模型的分布与reference模型的分布不会有较大的差异。
4. 微调流程
DPO微调流程
上图展示了DPO微调的大致流程，其中Trained
LM即为策略模型，Frozen LM即为参考模型，二者均是先进行SFT微调得到的模型进行初始化，其中Trained LM需要进行训练，Frozen LM不参与训练。
两个模型分别针对chosen和rejected进行预测获取对应的得分，再通过DPO的损失函数进行损失计算，进而不断的迭代优化。
5. 源码
源码参考代码：https://github.com/eric-mitchell/direct-preference-optimization
5.1 DPO损失函数
def preference_loss(policy_chosen_logps: torch.FloatTensor,
policy_rejected_logps: torch.FloatTensor,
reference_chosen_logps: torch.FloatTensor,
reference_rejected_logps: torch.FloatTensor,
beta: float,
label_smoothing: float = 0.0,
ipo: bool = False,
reference_free: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
# policy_chosen_logps:训练模型对于chosen经过log后logits
# policy_rejected_logps:训练模型对于rejected经过log后logits
# reference_chosen_logps:训练模型对于chosen经过log后logits
# reference_rejected_logps:训练模型对于rejected经过log后logits
# beta: policy和reference的差异性控制参数
# actor模型选择chosen优先于rejected
pi_logratios = policy_chosen_logps - policy_rejected_logps
# reference模型选择chosen优先于rejected
ref_logratios = reference_chosen_logps - reference_rejected_logps
if reference_free:
ref_logratios = 0
#差值可类似于KL散度，保障actor模型的分布与reference模型的分布不会有较大的差异
logits = pi_logratios - ref_logratios # also known as h_{\pi_\theta}^{y_w,y_l}
if ipo:
losses = (logits - 1/(2 * beta)) ** 2 # Eq. 17 of https://arxiv.org/pdf/2310.12036v2.pdf
else:
# Eq. 3 https://ericmitchell.ai/cdpo.pdf; label_smoothing=0 gives original DPO (Eq. 7 of https://arxiv.org/pdf/2305.18290.pdf)
# label_smoothing为0，对应的DPO论文的算法
losses = -F.logsigmoid(beta * logits) * (1 - label_smoothing) - F.logsigmoid(-beta * logits) * label_smoothing
# chosen和rejected的奖励
chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps).detach()
rejected_rewards = beta * (policy_rejected_logps - reference_rejected_logps).detach()
return losses, chosen_rewards, rejected_rewards
5.2 批次训练过程
def get_batch_metrics(self, batch: Dict[str, Union[List, torch.LongTensor]], loss_config: DictConfig, train=True):
"""Compute the SFT or DPO loss and other metrics for the given batch of inputs."""
if loss_config.name in {'dpo', 'ipo'}:
# policy模型针对chosen和rejected进行预测
policy_chosen_logps, policy_rejected_logps = self.concatenated_forward(self.policy, batch)
with torch.no_grad():
# reference模型针对chosen和rejected进行预测
reference_chosen_logps, reference_rejected_logps = self.concatenated_forward(self.reference_model, batch)
if loss_config.name == 'dpo':
loss_kwargs = {'beta': loss_config.beta, 'reference_free': loss_config.reference_free, 'label_smoothing': loss_config.label_smoothing, 'ipo': False}
elif loss_config.name == 'ipo':
loss_kwargs = {'beta': loss_config.beta, 'ipo': True}
else:
raise ValueError(f'unknown loss {loss_config.name}')
#损失计算
losses, chosen_rewards, rejected_rewards = preference_loss(
policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps, **loss_kwargs)
reward_accuracies = (chosen_rewards > rejected_rewards).float()
elif loss_config.name == 'sft':
policy_chosen_logits = self.policy(batch['chosen_input_ids'], attention_mask=batch['chosen_attention_mask']).logits.to(torch.float32)
policy_chosen_logps = _get_batch_logps(policy_chosen_logits, batch['chosen_labels'], average_log_prob=False)
losses = -policy_chosen_logps
return losses.mean()
5.3 LM的交叉熵计算
def _get_batch_logps(logits: torch.FloatTensor, labels: torch.LongTensor, average_log_prob: bool = False) -> torch.FloatTensor:
#经模型后的logits进行批量计算logps
assert logits.shape[:-1] == labels.shape
#基于先前的token预测下一个token
labels = labels[:, 1:].clone()
logits = logits[:, :-1, :]
loss_mask = (labels != -100)
# dummy token; we'll ignore the losses on these tokens later
labels[labels == -100] = 0
#交叉熵函数
per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
if average_log_prob:
return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
else:
return (per_token_logps * loss_mask).sum(-1)
5.4 其他注意
1. hugging face设置代理
源码会从hugging face中下载英文语料和模型，由于网络限制，因此设置代理映射，将HF_ENDPOINT设置为https://hf-mirror.com，即设置: os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
2. 如果仅想要熟悉DPO整体流程，可以下载较小的生成式模型，如BLOOM 560M，GPT2等
6. 总结
一句话足矣~
本文主要针对大语言模型的DPO算法的整理，包括原理、流程及部分源码。
此外，建议大家可以针对源码进行运行，源码的欢迎大家一块交流。
7. 参考
(1) RLHF：https://blog.csdn.net/v_JULY_v/article/details/128579457
(2) DPO论文: https://arxiv.org/pdf/2305.18290v2.pdf
(3) DPO代码: https://github.com/eric-mitchell/direct-preference-optimization
(4) DPO理解1：https://medium.com/@joaolages/direct-preference-optimization-dpo-622fc1f18707
(5) DPO理解2: https://zhuanlan.zhihu.com/p/669825918
