LLM面面观之RLHF平替算法DPO

1. 背景

最近本qiang~老看到一些关于大语言模型的DPO、RLHF算法，但都有些云里雾里，因此静下心来收集资料、研读论文，并执行了下开源代码，以便加深印象。

此文是本qiang~针对大语言模型的DPO算法的整理，包括原理、流程及部分源码。

2. DPO vs RLHF

上图左边是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的损失函数

如何将RLHF的Reward model过程简化为上式，作者花了大量篇幅进行了推导，感兴趣的读者可以参考附件DPO的论文。

DPO算法的目的是最大化奖励模型(此处的奖励模型即为训练的策略)，使得奖励模型对chosen和rejected数据的差值最大，进而学到人类偏好。

上式的后半部分通过对数函数运算规则，可以进行如下转化。

转化后的公式和源代码中的计算函数中的公式是一致的。

其中左半部分是训练的policy模型选择chosen优先于rejected，右半部分是冻结的reference模型选择chosen优先于rejected，二者的差值可类似于KL散度，保障actor模型的分布与reference模型的分布不会有较大的差异。

4. 微调流程

上图展示了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损失函数

 1 def preference_loss(policy_chosen_logps: torch.FloatTensor,
 2                     policy_rejected_logps: torch.FloatTensor,
 3                     reference_chosen_logps: torch.FloatTensor,
 4                     reference_rejected_logps: torch.FloatTensor,
 5                     beta: float,
 6                     label_smoothing: float = 0.0,
 7                     ipo: bool = False,
 8                     reference_free: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
 9     # policy_chosen_logps: 训练模型对于chosen经过log后logits
10     # policy_rejected_logps: 训练模型对于rejected经过log后logits
11     # reference_chosen_logps: 训练模型对于chosen经过log后logits
12     # reference_rejected_logps: 训练模型对于rejected经过log后logits
13     # beta: policy和reference的差异性控制参数
14
15     # actor模型选择chosen优先于rejected
16     pi_logratios = policy_chosen_logps - policy_rejected_logps
17     # reference模型选择chosen优先于rejected
18     ref_logratios = reference_chosen_logps - reference_rejected_logps
19
20     if reference_free:
21         ref_logratios = 0
22
23     # 差值可类似于KL散度，保障actor模型的分布与reference模型的分布不会有较大的差异
24     logits = pi_logratios - ref_logratios  # also known as h_{\pi_\theta}^{y_w,y_l}
25
26     if ipo:
27         losses = (logits - 1/(2 * beta)) ** 2  # Eq. 17 of https://arxiv.org/pdf/2310.12036v2.pdf
28     else:
29         # Eq. 3 https://ericmitchell.ai/cdpo.pdf; label_smoothing=0 gives original DPO (Eq. 7 of https://arxiv.org/pdf/2305.18290.pdf)
30         # label_smoothing为0，对应的DPO论文的算法
31         losses = -F.logsigmoid(beta * logits) * (1 - label_smoothing) - F.logsigmoid(-beta * logits) * label_smoothing
32
33     # chosen和rejected的奖励
34     chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps).detach()
35     rejected_rewards = beta * (policy_rejected_logps - reference_rejected_logps).detach()
36
37     return losses, chosen_rewards, rejected_rewards

5.2 批次训练过程

 1 def get_batch_metrics(self, batch: Dict[str, Union[List, torch.LongTensor]], loss_config: DictConfig, train=True):
 2     """Compute the SFT or DPO loss and other metrics for the given batch of inputs."""
 3
 4     if loss_config.name in {'dpo', 'ipo'}:
 5         # policy模型针对chosen和rejected进行预测
 6         policy_chosen_logps, policy_rejected_logps = self.concatenated_forward(self.policy, batch)
 7         with torch.no_grad():
 8             # reference模型针对chosen和rejected进行预测
 9             reference_chosen_logps, reference_rejected_logps = self.concatenated_forward(self.reference_model, batch)
10
11         if loss_config.name == 'dpo':
12             loss_kwargs = {'beta': loss_config.beta, 'reference_free': loss_config.reference_free, 'label_smoothing': loss_config.label_smoothing, 'ipo': False}
13         elif loss_config.name == 'ipo':
14             loss_kwargs = {'beta': loss_config.beta, 'ipo': True}
15         else:
16             raise ValueError(f'unknown loss {loss_config.name}')
17         # 损失计算
18         losses, chosen_rewards, rejected_rewards = preference_loss(
19             policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps, **loss_kwargs)
20
21         reward_accuracies = (chosen_rewards > rejected_rewards).float()
22
23     elif loss_config.name == 'sft':
24         policy_chosen_logits = self.policy(batch['chosen_input_ids'], attention_mask=batch['chosen_attention_mask']).logits.to(torch.float32)
25         policy_chosen_logps = _get_batch_logps(policy_chosen_logits, batch['chosen_labels'], average_log_prob=False)
26
27         losses = -policy_chosen_logps
28
29     return losses.mean()

5.3 LM的交叉熵计算

 1 def _get_batch_logps(logits: torch.FloatTensor, labels: torch.LongTensor, average_log_prob: bool = False) -> torch.FloatTensor:
 2     # 经模型后的logits进行批量计算logps
 3
 4     assert logits.shape[:-1] == labels.shape
 5
 6     # 基于先前的token预测下一个token
 7     labels = labels[:, 1:].clone()
 8     logits = logits[:, :-1, :]
 9     loss_mask = (labels != -100)
10
11     # dummy token; we'll ignore the losses on these tokens later
12     labels[labels == -100] = 0
13
14     # 交叉熵函数
15     per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
16
17     if average_log_prob:
18         return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
19     else:
20         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