使用 TRL 通过 Direct Preference Optimization (DPO) 微调 Llama 2

Hugging Face 已将 Direct Preference Optimization (DPO) 集成到 TRL 库中,使开发者能够使用简化的二元交叉熵损失来对齐如 Llama 2 这样的大型语言模型(LLMs)。此方法消除了对辅助奖励模型的需求以及通常用于 Reinforcement Learning from Human Feedback (RLHF) 的复杂强化学习 (RL) 机制。

DPO 与 PPO 的对比:简化对齐流程

Direct Preference Optimization (DPO) 通过跳过奖励建模和 RL 优化步骤来简化对齐过程。在传统的 RLHF 管道中,例如使用 Proximal Policy Optimization (PPO) 的管道,奖励模型被训练以估计人类偏好,并且会使用带有 KL 惩罚的冻结参考模型来防止策略模型偏离过远并产生无意义的输出。

DPO 用奖励函数到最优 RL 策略的解析映射来替换这一点。这使得奖励模型和参考模型上的 RL 损失可以直接转换为参考模型上的损失。因此,语言模型通过在偏好数据上的直接似然目标进行优化,从而消除了对单独奖励模型或繁琐的基于 RL 的优化的需求。

使用 TRL 库实现 DPO

TRL 库提供了实现 DPO 工作流的辅助函数。虽然传统的 RLHF 管道包含四个步骤——监督微调 (SFT)、偏好数据标注、奖励模型训练和 RL 优化——但 DPO 将后两个步骤合并为一个步骤。

数据要求

要在 TRL 中使用 DPOTrainer,偏好数据必须以包含三个特定键的字典格式提供:

  • prompt: 在推理时提供给模型的上下文提示。
  • chosen: 首选的生成响应。
  • rejected: 非首选的响应。

训练器配置

DPOTrainer 需要一个基础模型(来自 SFT 管道)、一个参考模型(通常是 SFT 训练后的基础模型的副本)以及一个温度超参数 beta(通常在 0.1 和 0.5 之间)。beta 参数控制参考模型的影响;较小的 beta 值意味着参考模型被更多地忽略。

案例研究:对齐 Llama 2 7B

Hugging Face 通过在 stack-exchange 偏好数据集上微调 Llama 2 7B 参数模型来演示 DPO 的应用,该数据集包含对问题的排名答案。

步骤 1:监督微调 (SFT)

模型首先通过使用 SFTTrainer 和通过 bitsandbytes 库的 QLoRA 技术进行 SFT。这涉及将基础模型以 4-bit 量化加载,并在目标模块 (q_proj, v_proj) 上添加 LoRA 层。

步骤 2:DPO 训练

SFT 之后,得到的模型被用作 DPO 训练的基础模型和参考模型。使用 AutoPeftModelForCausalLM,模型以 4-bit 配置加载并通过 QLoRA 方法进行训练。DPOTrainer 使用评估数据集评估进度并报告隐式奖励指标。

监控 DPO 性能

在训练和评估期间,DPOTrainer 跟踪四个关键奖励指标以衡量对齐进度:

  • rewards/chosen: 由策略模型和参考模型在所选响应上的对数概率之差的均值,按 beta 缩放。
  • rewards/rejected: 由策略模型和参考模型在被拒绝响应上的对数概率之差的均值,按 beta 缩放。
  • rewards/accuracies: 所选奖励高于对应被拒绝奖励的百分比。
  • rewards/margins: 所选奖励与被拒绝奖励之间的均值差。

成功的训练表现为 margins 和 accuracies 随着增加并趋近于 1.0,这意味着模型始终为首选响应分配更高的奖励。

Sources