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损失函数

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)