“Offline Reinforcement Learning for LLM Multi-Step Reasoning”由Huaijie Wang、Shibo Hao等人撰写。论文提出了OREO(Offline REasoning Optimization)这一离线强化学习方法,用于增强大语言模型(LLMs)的多步推理能力,在数学推理和实体代理控制任务中展现出良好性能,为提升LLMs的推理能力提供了新途径。
1 研究背景与目的
背景
LLMs在复杂多步推理任务中的应用日益广泛,但传统RL算法的在线数据收集成本高昂,限制了其实际应用。离线RL方法虽更实用,但现有的如DPO等方法在多步推理任务中存在不足,如依赖成对偏好数据且在信用分配上效果不佳。
多步推理任务的数据特点是奖励稀疏,通常只有在任务成功或失败的终端步骤才有非零奖励,且数据多为独立轨迹。
目的:提出OREO方法,通过优化软Bellman方程联合学习策略模型和价值函数,减少对成对数据收集的需求,实现更细粒度的信用分配,以提升LLMs的多步推理能力。
2 相关工作
-
LLM的强化学习:RLHF是使LLM响应与人类偏好对齐的常用方法,传统RL方法如PPO被广泛用于LLM的后训练,近期基于拒绝采样、偏好和REINFORCE的RL方法也受到关注。最大熵RL与RLHF目标有强关联,为微调LLM提供了新方向。
-
LLM推理:LLMs可通过生成中间推理步骤解决复杂问题,近期研究通过监督微调、拒绝采样和RL算法等增强其推理能力,但DPO在推理任务中的应用存在局限。另一种方法是训练过程奖励模型(PRM)提供更细粒度反馈,而本文方法训练的价值函数可用于测试时搜索以提升推理性能。
3 OREO方法
-
基于MDP定义推理任务:将LLM推理过程定义为马尔可夫决策过程(MDP),状态为token序列,动作是生成新token,奖励函数在终端步骤反映推理链正确性。在与环境交互的任务中,状态和转移函数稍有不同。
-
软Bellman方程及与DPO的关系:软Bellman方程描述了最优策略和价值函数的关系,为OREO方法提供理论基础。DPO通过放宽软Bellman方程要求和引入Bradley-Terry偏好模型来推导,但存在不必要的成对数据收集和信用分配问题。
-
OREO学习目标:采用类似PCL的方法微调LLM,通过优化软Bellman方程属性来训练策略模型和价值模型。价值网络采用MSE损失,策略目标函数包含与价值函数相关的项和正则化项。
-
损失变体:包括步级OREO(将动作视为整个推理步骤)和响应级OREO(模仿DPO行为),并根据实验选择合适的参数。
-
迭代OREO:可迭代应用,每次迭代后用更新的策略模型收集新数据用于进一步训练。
-
测试时搜索策略:利用训练得到的价值函数指导数学推理任务的步级束搜索或在实体代理控制任务中选择最佳动作,提高模型性能。
4 实验设置与结果
-
实验设置
-
数据集和评估指标:采用GSM8K和MATH数据集进行数学推理任务评估,以准确率为指标;使用ALFWorld数据集进行实体代理控制任务评估,以成功率为指标。
-
基线方法:与监督微调、拒绝采样、DPO和KTO等方法进行比较。
-
实验细节:包括数据集构建、推理步骤分割和超参数设置等。
-
-
实验结果
-
主要结果:OREO在数学推理和实体代理控制任务中均优于基线方法,在不同模型大小上表现一致。在数学推理任务中,相对于SFT有显著提升;在ALFWorld任务中,能有效利用失败轨迹,提高在未见环境中的成功率。
-
迭代OREO:随着迭代次数增加,OREO的性能持续提升,而拒绝采样等基线方法出现饱和迹象。
-
隐式与显式价值函数对比:OREO通过显式参数化价值函数,在区分正确和错误推理步骤上比DPO中基于策略模型的隐式价值函数更有效。
-
测试时搜索价值函数:在数学推理任务中,OREO利用价值函数进行步级束搜索,随着计算预算增加,准确率显著提高;在ALFWorld任务中,选择最佳-K动作可提高成功率。
-
5 研究结论与局限
-
结论:OREO算法通过优化软Bellman方程训练策略和价值函数,在多步推理任务中表现优异,减少了对成对偏好数据的需求,实现了更有效的信用分配,且价值函数可用于测试时搜索进一步提升性能。
-
局限:部分实验受计算资源限制,仅使用1.5B模型,未来计划在更大规模上进行实验;目前主要在数学推理和实体代理控制任务中评估,未来将扩展到更多类型任务。