使大型语言模型与 BRAIn 对齐

社区文章发布于2024年6月11日

Gaurav

gauravpandey1

TL;DR：我们介绍了 BRAIn——一种用于 RLHF 的分布匹配方法，在 Anthropic HH 和 TL;DR 摘要任务上取得了最先进的性能，优于 DPO 和其他 RLHF 方法！

注：这项工作已被 ICML-2024（主会议）接受发表

ArXiv 版本链接：： https://arxiv.org/pdf/2402.02479

大型语言模型训练的3个阶段

在过去的几年里，大型语言模型（LLMs）在多轮对话、创意写作、数学和逻辑推理等各种任务上展现出了巨大的能力。这些大型语言模型通常分三个阶段进行训练：

自监督预训练
监督指令微调（SFT）
基于人类反馈的强化学习（RLHF）

自监督预训练阶段使模型具备了语言理解和生成能力，而监督指令微调阶段则教会模型遵循自然语言指令。在 RLHF 阶段，模型被鼓励遵循我们期望的行为。期望的概念可以是显性的（例如，输出中不含脏话）或隐性的人类偏好某些输出文本而非其他。

与人类偏好对齐

那么，如何使模型与人类偏好对齐呢？PPO-RLHF，GPT-3.5 和 GPT-4 背后的 RLHF 方法，通过首先训练一个奖励模型来模拟人类偏好，即人类偏好的输出应比其他输出获得更高的奖励。然后，对大型语言模型（也称为策略）进行微调，以生成由奖励模型确定的高奖励输出。我们还确保对齐后的 LLM 与 SFT LLM（监督指令微调的 LLM）接近，从而防止它忘记在前两个阶段获得的能力。

PPO-RLHF 最近已被离线对比技术取代，例如序列似然校准（SLiC）、直接偏好优化（DPO）及其变体。这些方法训练 LLM 以对比偏好/高奖励输出和被拒绝/低奖励输出。DPO 已成为对齐高性能模型（如 Zephyr、Mixtral 和 LLama-3）的事实标准方法。

尽管 PPO-RLHF 和 DPO 算法之间存在巨大差异，但这两种方法具有相同的最终目标，即 PPO 最优策略。另一组鲜为人知的方法使用分布匹配来将 LLM 与这种最优策略对齐。理想情况下，这将需要从最优策略中采样，而这已被证明具有挑战性。因此，分布匹配方法（DPG、GDC、GDC++）改为从提议分布中采样，并根据其重要性权重对这些样本进行加权。尽管分布匹配背后的直觉很清晰，但这些方法在与人类反馈对齐方面并未成功。

我们的贡献——BRAIn

在研究分布匹配方法缺乏成功的原因时，我们发现分布匹配方法（GDC、GDC++）中的梯度估计具有高方差。这意味着每个时间步的更新方向都会根据从大型语言模型中采样的输出而发生很大变化。下面通过一个简单的玩具示例进行演示：

假设我们试图达到的目标分布是标准的一维正态分布 $\mathcal{N}(0,1)$ 。令当前模型分布为 $\mathcal{N}(1,1)$ ，提议分布为 $\mathcal{N}(\theta,1)$ ，其中 $\theta$ 从 $0$ 变化到 $1$ 。下面，我们绘制了不同分布匹配目标相对于模型分布均值参数的梯度估计方差。样本从提议分布中抽取。可以看出，当提议分布与目标分布不同时，分布匹配方法（GDC、GDC++）的梯度估计方差很高。

这项研究促使我们创建了 BRAIn - Bayesian reward-conditioned Augmented Inference（贝叶斯奖励条件增强推理），它通过以下方式扩展了分布匹配方法：

我们通过使用贝叶斯规则将奖励建模假设纳入其中，泛化了 PPO-RLHF、DPO 和分布匹配方法中的目标分布。
我们提出了一种自归一化基线，可显著降低分布匹配中梯度估计的方差，如上图所示。通过引入自归一化基线，我们在 TL;DR 摘要和 Anthropic-Helpful & Harmless 响应生成任务上实现了 SOTA 性能，并将 DPO 确立为 BRAIn 的一个特例。

BRAIn 的目标

将后验作为目标

给定输入提示 $x$ ，不同的 RLHF 算法试图达到输出集 $y$ 上的目标分布 $p_{T} (y ∣ x)$ 。这个目标分布取决于两个因素：

基本分布。这通常是一个 SFT 模型，表示为 $p_{SFT}$
奖励函数 $r (x, y)$

BRAIn 使用贝叶斯规则结合了上述两个因素的信息。具体来说，SFT 模型作为先验，而奖励函数用于定义似然项。由此产生的后验被称为目标 $p_{T} (y ∣ x)$ 。

使用重要性权重进行训练

令 $q_\theta$ 是我们希望对齐到目标 $p_{T}(y|x)$ 的模型。理想情况下，可以通过从目标中采样并使用这些样本训练 $q_\theta$ 来实现，如下所示。

然而，由于从目标中采样可能具有挑战性，我们改用提议分布 $q (y ∣ x)$ 进行采样，并根据 $\frac{p_T(y|x)}{q(y|x)}$ 重新加权这些样本。由于 $p_{T}$ 的归一化常数难以处理，我们如下所示对权重进行自归一化：

关于提议分布的注意事项：生成样本的理想分布是什么？显然，由于我们正在努力达到目标，理想情况下，我们应该从目标 $p_{T}$ 中采样。然而，由于这具有挑战性，我们选择从最接近目标的分布中采样。在训练开始时，我们从 SFT 模型 $p_{SFT}$ 中采样。但是，随着训练的进行，我们将包含来自最新策略的样本。

自归一化基线以降低方差

上述目标的梯度由下式给出 $\nabla_\theta \mathcal{L}(\theta) = \sum_{i=1}^n \hat{\alpha}_{y_i} \nabla_\theta \log p_\theta(y_i|x)$ 。该梯度估计已用于 GDC 的 LLM 对齐，不同之处在于权重未进行自归一化。正如我们之前所示，GDC 梯度估计具有高方差，这导致性能不佳。