Step-DPO 论文——数学大语言模型理解

news2024/9/21 14:41:14

论文题目：STEP-DPO: STEP-WISE PREFERENCE OPTIMIZATION FOR LONG-CHAIN REASONING OF LLMS

翻译为中文就是：“LLMs长链推理的逐步偏好优化”

论文由港中文贾佳亚团队推出，基于推理步骤的大模型优化策略，能够像老师教学生一样优化大模型。

以Qwen2-72B-Instruct模型作为基础模型进行微调优化后，其数学成绩超越了GPT-4、Gemini1.5-Pro、Claude3-Opus等闭源模型。

论文链接：https://arxiv.org/pdf/2406.18629

代码仓库：https://github.com/dvlab-research/Step-DPO

1. 介绍

大语言模型（LLMs）在数学推理上具有重大挑战，这是由于数学需要精确的推理链。然而，直接偏好优化（DPO）对长链数学推理的益处有限，因为采用DPO的模型难以识别错误答案中的详细错误。

所以作者提出了Step-DPO方法，它将整个答案划分多个步骤作答（Step1, Step2, Step3, ...），大大提高的模型的精度。

在MATH数据集上，在Qwen2-7B-Instruct上准确率从53.0% 提升到58.6%，GSM8K数据集，准确率从85.5%提升到87.9% 。使用 Qwen2-72B-Instruct模型，在MATH和GSM8K上分别取得 70.8% 和 94.0%的准确率。

1.1 像教育学生一样训练大模型

数学推理被认为是大语言模型（LLMs）中一种关键的长链推理能力。由于需要广泛的思维链（CoT），这项任务尤其具有挑战性，其中可能包括许多推理步骤，这些步骤中的任何错误都可能导致最终得不到正确答案。

（1）首先，最常用的方法就是监督微调（SFT），使用各种数据增强对齐来微调模型。然而，当SFT数据达到一定数量时，模型经常出现幻觉，性能也随之趋于饱和。一个潜在的原因是，随着首选输出的概率增加，不理想输出的概率也会增加。这种现象使得模型在长链推理中更容易出错。

（2）最近，直接偏好优化（DPO）（Rafailov et al., 2024）被提出用于使用偏好对数据进行对齐（每个偏好对都包含一个输入提示、偏好输出及非偏好输出），因其简单性而广受欢迎。尽管DPO在Chat聊天任务中很有效，但它对长链（long-chain）数学任务效果不明显。如下图2所示。

（3）于是作者提出了Step-DPO，基于推理步骤的直接偏好优化。Step-DPO 逐步检查每个步骤的答案是否正确，这使得模型能够轻松定位错误Step，以进行有效的优化，显著增强了长链推理。

2. STEP-DPO 公式

2.1 DPO

我们先看到DPO的优化目标函数：

$\begin{aligned} L_{DPO}(\theta)=-E_{(x,y_{win},y_{lose})\backsim D}[log \sigma (\beta log \frac {\pi_{\theta} (y_{win} \mid x)}{\pi_{ref}(y_{win \mid x})} - \beta log \frac{\pi_{\theta}(y_{lose} \mid x)}{\pi_{ref}(y_{lose} \mid x)})] \end{aligned}$

其中， $\ x$ 是输入提示， $\ y_{win}, y_{lose}$ 分别表示正确的回答、错误的回答， $\ D$ 是偏好对数据集。 $\sigma$ 表示 sigmoid 函数， $\pi_{\theta}$ 与 $\pi_{ref}$ 分别表示当前要优化的微调模型以及训练过程中保存不变的参照模型， $\beta$ 是一个超参数用来控制距离。

2.2 Step-DPO

我们再看到Step-DPO，它不再像DPO从整体对比答案，而是将每个推理步骤视为一个基本单元，对比单个推理步骤，更精细地提升模型的推理能力。目标优化公式：

$\begin{aligned} L(\theta)=-E_{(x,s_{1 \backsim k-1},s_{win}, s_{lose})\backsim D}[log \sigma (\beta log \frac {\pi_{\theta} (s_{win} \mid x; s_{1 \backsim k-1})}{\pi_{ref}(s_{win} \mid x; s_{1 \backsim k-1})} - \beta log \frac{\pi_{\theta}(s_{lose} \mid x; s_{1 \backsim k-1})}{\pi_{ref}(s_{lose} \mid x; s_{1 \backsim k-1})})] \end{aligned}$

回答 $\ y$ 可以分解为多个步骤 $\ y=s_{1}, ..., s_n$ ， $\ x$ 表示输入提示。Step-DPO 优化目标就是最大化正确的下一个推理步骤 $\ s_{win}$ 的概率，最小化错误步骤 $\ s_{lose}$ 的概率，如图3所示。

3. 分布式数据构建

Step-DPO 的训练数据集是怎样的呢？每个数据样本中应该包含下面4项：

1）prompt $\ x$ ；

2）初始推理步骤 $\ s_{1 \backsim k-1}$ ；

3）首选推理步骤 $\ s_{win}$ ；

4）不需要（错误）的推理步骤 $\ s_{lose}$

如下图所示：

（1）错误收集

首先，我们收集数学问题问答的数据集 $\ D_0 = \{ (x, \hat{y}) \}$ ，x 是数学问题， $\ \hat{y}$ 是真实答案。

然后，使用初始（参照）模型 $\ \pi_{ref}$ 来得到每个数学问题 x 的答案。

在进行模型推理之前，添加思维链（CoT）前缀作为提示，比如：“Let‘s think step by step. Step 1:”，以确保模型的推理结果被结构化为多个推理步骤。

模型推理完成之后可得到每个数学问题x的推理结果y，然后选择与真实答案 $\ \hat{y}$ 不一致的那些结果，汇总得到数据集 $\ D_1$ ：

$\begin{aligned} D_1 = \{ (x, \hat{y}, y) \mid x \in D_0 \} \end{aligned}$

（2）错误步骤定位

假设每个错误的推理结果都被明确地表示为推理步骤序列 $\ y = s_1, s_2, ..., s_n$ ，随后需要人工或利用GPT-4验证每个推理步骤的正确性，直到找到第一个错误步骤 $\ s_k$ ，选择 $\ s_k$ 作为错误的推理步骤 $\ s_{loss}$ 。这样得到一个包含错误步骤的数据集 $\ D_2$ ：