用Step-KTO优化大模型数学推理能力，效果显著提升！

用Step-KTO优化大模型数学推理能力，效果显著提升！

原文链接：https://arxiv.org/abs/2501.10799
论文标题：Step-KTO: Optimizing Mathematical Reasoning through Stepwise Binary Feedback
作者：Lin等人

📖阅读时长：25分钟

🕙发布时间：2025-02-09

近日热文：全网最全的神经网络数学原理（代码和公式）直观解释
欢迎关注知乎和公众号的专栏内容
LLM架构专栏
知乎LLM专栏
知乎【柏企】
公众号【柏企科技说】【柏企阅文】

大语言模型（LLM）常用的思维链提示、自洽抽样等方法，在数学推理任务上取得了一定成果，但它们往往只看重最终答案是否正确，却忽视了推理过程是否连贯、可靠。在这篇论文中，作者提出了“逐步Kahneman-Tversky启发优化法”（Stepwise Kahneman-Tversky-inspired Optimization，简称Step-KTO）。这是一种全新的训练框架，结合了过程级和结果级的二元反馈机制，能引导大语言模型生成更值得信赖的推理过程。

主要贡献

1. 提出Step-KTO微调框架：该框架创新地融合了过程级与结果级反馈，不仅追求最终答案正确，还注重推理过程的准确性和逻辑性。
2. 验证训练有效性：研究显示，利用Step-KTO进行迭代训练，能让模型性能持续提升，充分证明了这种结合两种反馈的训练方式，对优化大语言模型推理能力的有效性。

概述

整个流程可参考下图及后续步骤说明：

1. 给定一组数学问题数据集，大语言模型（LLM）会针对每个问题，输出推理步骤以及最终答案。
2. 每个中间推理步骤由流程奖励模型（Process RM）进行评估，最终答案则由结果奖励模型（Outcome RM）来评判。
3. 将结果级的正确性反馈信号$c_{o}$、过程级的逐步正确性反馈信号$c_{sh}$，连同输入的问题$x$、模型给出的回答$y$一起记录下来。
4. 利用这些信号计算Step-KTO损失值，以此引导大语言模型，使其不仅要给出正确的最终答案，推理步骤也要连贯、正确。经过多轮迭代训练，模型的逐步推理能力和最终答案的准确率都会不断提高。

方法论

问题详情和符号

由参数$\theta$确定的策略模型$\pi_{\theta}$，在面对问题$x$时，会自回归地生成一个响应序列$y = (s_1, s_2, \ldots, s_{|y|})$ 。这里每个步骤$s_h$都是一个推理步骤，它们之间用特殊标记（比如 “## Step”）隔开。

最终答案$y$的正确性，可以通过基于规则的正确性函数$Regex(y, y^{\star}_x) \in \{0, 1\}$自动判断。这个函数会把模型给出的最终答案，和真实的标准答案进行对比。

二进制信号的两个级别

1. 结果反馈：$Regex(y, y^{\star}_x) \in \{0, 1\}$，用来表明从$y$得出的最终答案是否正确。
2. 逐步反馈：$Prm(x, y^{\star}_x, s_h) \in \{0, 1\}$，用于判断中间推理步骤$s_h$是否正确。

我们的目标是，把这两个反馈信号都融入到策略模型$\pi_{\theta}$的训练目标中，引导模型不仅给出正确答案，整个推理过程也要正确、连贯、可靠。

KTO背景

KTO旨在借助受Kahneman-Tversky启发的值函数，让策略$\pi_{\theta}$与二元反馈保持一致。它没有采用最大化偏好输出的对数似然，或者直接使用强化学习的方式，而是定义了一个逻辑斯蒂价值函数。这个函数在面对收益时表现出风险规避，面对损失时则呈现出风险寻求的特性。

原始的KTO损失主要关注最终答案层面。

其中，$\pi_{ref}$是一个参考策略（通常是初始模型的检查点），作为对比的基线；$\sigma$是逻辑斯蒂函数；$\beta > 0$用于控制风险规避程度；$\lambda_D$、$\lambda_U$是加权系数。$z_0$项中的$y′$代表任意一个输出序列，它作为参考点，保证优化过程的平衡。

结果级别的KTO损失计算公式为：

步骤KTO

Step-KTO对KTO进行了扩展，引入了逐步二元反馈$Prm(x, y^{\star}_x, s_h)$，以此来评估每个推理步骤的质量。

首先，在步骤级别定义隐含奖励：

它表示与参考策略$\pi_{ref}$相比，在策略$\pi_{\theta}$下生成步骤$s_h$的增量优势，反映了在相同上下文$(x, s_{<h})$下，选择$s_h$相较于参考模型基线可能性所带来的奖励增减。

接着，引入逐步KL基线：

给定二元逐步反馈$Prm(x, y^{\star}_x, s_h)$，定义一个与结果级情况类似的值函数。如果步骤$s_h$被认为是理想的，模型应相对于$z_0^{step}$增加其隐含奖励$r_{\theta}(x, s_h)$；反之，如果$s_h$不理想，则应降低该隐含奖励。具体公式为：

这里，$\lambda_D^{step}$、$\lambda_U^{step}$和$\beta_{step}$与结果级的对应参数作用类似，用于在单个步骤的粒度上控制奖励或惩罚的强度。

逐步损失为：

将逐步目标与结果级KTO损失相结合，得到最终的Step-KTO目标：

通过同时优化结果级和逐步级的反馈，Step-KTO确保模型从最初的推理步骤到最终得出答案的整个过程，都是正确且连贯的。

迭代训练

1. 对于Llama-3.3–70B-Instruct模型：直接将其作为种子模型$M_0$ 。

2. 对于Llama-3.1系列模型：先在训练数据上进行有监督的微调，之后再当作$M_0$使用。从$M_0$开始，按照以下流程迭代优化，得到$M_1$、$M_2$、…、$M_T$ ：

   • 候选人生成：对于数据集中的每个问题$x \in D$，用温度$T = 0.7$和核采样（$p = 0.95$）的方式，从策略$\pi_{M_t}(· | x)$中采样8个候选解$y_k$ 。这种随机解码策略能产生多样化的候选解，有助于筛选出正样本和负样本。
   • 结果评估：用结果正确性函数$Regex(y_k, y^{\star}_x)$，将每个候选解$y_k$与真实答案$y^{\star}_x$进行对比评估。如果所有采样得到的解都不正确，就把真实解$y^{\star}_x$当作正样本加入；要是所有采样解都正确，那么在当前迭代中就舍弃这个问题，优先从模型还有提升空间的问题中学习。
   • 逐步评估：针对筛选出来的解，运用逐步正确性函数$Prm(x, y^{\star}_x, s_h)$，评估每个推理步骤的质量，得到一组二进制信号，判断中间步骤是否符合理想的推理模式。
   • 数据集构建：把这些标注好的样本汇总成$D_{M_t} = \{(x, y, c_{out}, c_{step1}, \ldots, c_{step_{S - 1}}) | y \in D\}$ 。其中，$c_{out} = Regex(y, y^{\star}_x)$是结果级别的正确性标识，$c_{step_h} = Prm(x, y^{\star}_x, s_h)$是解$y$的$S - 1$个中间步骤的逐步正确性指标。
   • 参数更新：利用数据集$D_{M_t}$，根据选定的对齐目标（可以是Step-KTO损失，或者像IRPO这样的基线方法）来更新模型参数。
   • 迭代：重复上述过程$T$次，每次都基于$M_t$优化生成新模型$M_{t + 1}$ 。

实验

指标

1. Pass@1：指从策略$\pi_{\theta}$得到的单个贪婪完成解$y$正确的比例，反映了确定性解码下模型的直接准确率。
2. Maj@8：通过在温度$T = 0.7$的条件下，生成8个候选解$y_k \sim \pi_{\theta}(· | x)$，然后选择出现次数最多的答案，以此得到的准确率。该指标体现了模型在多个采样解下的稳健性。

基线方法

1. RFT（抑制微调）：对筛选后的数据集$\{(x, y) \in D_{M_t} | c_{out} = 1\}$进行有监督微调，只保留最终答案正确的解。与Step-KTO不同，RFT没有引入明确的偏好或奖励信号，只是直接模仿真实答案。
2. IRPO（迭代推理偏好优化）：在结果级别应用成对偏好的迭代训练，并添加额外的负对数似然（NLL）损失项来稳定训练。IRPO只关注结果的正确性，不使用逐步反馈来优化模型。
3. KTO（Kahneman-Tversky优化）：运用基于Kahneman-Tversky价值函数推导的HALO损失，为模型赋予类似人类的风险厌恶特性，以及对收益和损失的不对称加权。和其他只关注结果级别的方法一样，KTO没有考虑逐步正确性信号。
4. SimPO和IPO：二者都是直接偏好优化（DPO）的变体，仅在结果级别根据成对偏好进行优化。与依赖Bradley-Terry模型、可能过度强调确定性偏好的DPO不同，SimPO和IPO采用更简单的变换，直接利用偏好概率。
5. Step-DPO：这是DPO的另一种变体，它优化的是逐步偏好，而非结果级偏好。Step-DPO通过识别并纠正具体的错误推理步骤，为长链推理任务提供更细致的监督。

实施细节

1. 优化器设置：使用AdamW优化器（$\beta_1 = 0.9$，$\beta_2 = 0.95$，权重衰减为$0.1$） 。在前100步采用线性预热，之后使用余弦退火学习率调度，将学习率降到初始值的$0.1$倍。
2. 训练参数：有效全局批量大小设为约100万个令牌，训练约2000步。在训练过程中，定期使用MATH测试集中的保留数据集评估模型，为每种方法挑选最佳的检查点。

结果

1. 模型性能对比：研究对比了不同规模的Llama模型和专有模型在数学问题求解上的性能。整体来看，Step-KTO的表现始终优于那些仅依赖结果级正确性的基线方法，像KTO、IRPO、SimPO、IPO，还有RFT这类较为简单的方法。例如，在使用8B Llama-3.1-Instruct模型的MATH-500测试中，Step-KTO的Pass@1达到了63.2%，比基线KTO模型的60.6%有所提升，大幅超越了IRPO和RFT 。当模型扩展到70B时，效果进一步提升。Llama-3.1–70B-Instruct搭配Step-KTO在MATH-500上的Pass@1达到76.2% ，在AMC23（60.0%）和AIME24（16.7%）测试中也表现优异。
2. 迭代训练效果：以同一种子模型M0（Llama-3.1–8B-Instruct）为例，研究观察了模型在多次迭代训练（M1、M2、M3）中的性能变化。结果显示，Step-KTO不仅最终性能更优，而且在每次迭代中的提升也更为稳定。比如在MATH-500测试中，Step-KTO的Pass@1从M1的59.4%提升到M3的63.2% ，超过了同一检查点下IRPO和KTO的提升幅度。

3. 推理质量评估：对比在MATH-500测试集中，包含错误中间步骤但最终答案正确的解决方案比例，以此分析推理质量。结果发现，Step-KTO和KTO在迭代过程中都降低了这类错误的发生率。其中，Step-KTO的改善更为明显且稳定，从M0的27.3%降至M3的19.9% ，而KTO仅改善到21.1%。

局限性

1. 结果级反馈信号存在噪声：在数学推理任务里，即便最终答案是经过自动验证的，也可能无法完全涵盖正确性的所有细节。像是细微的格式差异，或者一些不常见但合理的表达方式，都可能导致误判，进而影响训练信号的准确性，限制模型性能的进一步提升。
2. 依赖真实解决方案：Stepwise-KTO方法目前在判断最终答案和中间推理步骤的正确性时，依赖于真实的标准答案。但在一些场景中，可能没有高质量的真实推理路径，或者正确的中间推理步骤本身就难以明确界定，这就使得定义有效的逐步反馈变得困难。如何让模型从质量欠佳或有噪声的参考中学习，甚至在没有明确真实推理的情况下，仅通过基于偏好的评估进行学习，仍是有待解决的问题。
3. 实验场景的局限性：当前实验主要集中在那些至少能得出部分正确结果，或者存在部分正确推理步骤的情况。如果模型总是给出错误的结果，甚至连部分有效的中间步骤都难以生成，那么Step-KTO能否有效提升模型性能还不确定。在这种困难情况下，可能需要借助课程学习、更强力的初始化策略，或者专门设计的探索策略等技术，辅助逐步反馈机制，才能引导模型得出正确答案，优化推理过程。

结论

论文提出的Step-KTO训练框架，结合了结果级和过程级的二元反馈，能引导大语言模型生成更连贯、可解释且可靠的推理过程。通过将逐步正确性信号融入训练过程，Step-KTO在保证甚至提高最终答案准确率的同时，提升了中间推理步骤的质量。在具有挑战性的数学推理基准测试中，该方法展现出持续的性能提升，特别是在迭代训练和处理复杂推理任务时效果显著。研究结果充分表明，优化模型的整个推理轨迹，而不仅仅是最终结果，具有重要价值。

本文由mdnice多平台发布