Offline Reinforcement Learning for LLM Multi-Step Reasoning

📄 arXiv: 2412.16145v2 📥 PDF

作者: Huaijie Wang, Shibo Hao, Hanze Dong, Shenao Zhang, Yilin Bao, Ziran Yang, Yi Wu

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-12-20 (更新: 2024-12-25)

💡 一句话要点

提出OREO：一种用于LLM多步推理的离线强化学习方法

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 离线强化学习 多步推理 大型语言模型 软贝尔曼方程 价值函数

📋 核心要点

  1. 现有DPO方法在多步推理任务中面临配对数据稀缺和信用分配困难两大挑战。
  2. OREO通过优化软贝尔曼方程，联合学习策略模型和价值函数，从而减少对配对数据的依赖。
  3. 实验结果表明，OREO在数学推理和具身智能体控制任务上优于现有离线学习方法。

📝 摘要（中文）

本文提出OREO（Offline Reasoning Optimization），一种用于增强大型语言模型（LLM）多步推理能力的离线强化学习方法，旨在快速适应复杂任务。直接偏好优化（DPO）虽然在对齐LLM与人类偏好方面表现出潜力，但不太适合多步推理任务，原因在于：(1) DPO依赖于配对偏好数据，而多步推理任务通常难以获得此类数据；(2) DPO对所有token一视同仁，对于多步推理任务中的信用分配（通常伴随稀疏奖励）效果不佳。OREO基于最大熵强化学习的思想，通过优化软贝尔曼方程联合学习策略模型和价值函数。理论上，它减少了对配对数据的需求，并实现了更好的信用分配。实验表明，OREO在多步推理基准测试中超越了现有的离线学习方法，包括数学推理任务（GSM8K、MATH）和具身智能体控制（ALFWorld）。该方法可以扩展到多迭代框架，并且学习到的价值函数可以用于指导树搜索，从而进一步提高测试时的性能。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）在多步推理任务中的优化问题。现有方法，如直接偏好优化（DPO），依赖于配对偏好数据，而此类数据在多步推理任务中难以获取。此外，DPO对所有token进行统一处理，无法有效进行信用分配，尤其是在奖励稀疏的情况下。

核心思路：OREO的核心思路是借鉴最大熵强化学习的思想，通过优化软贝尔曼方程，联合学习一个策略模型和一个价值函数。这种方法旨在减少对配对数据的需求，并实现更有效的信用分配。价值函数可以评估中间步骤的质量，从而更好地指导模型的推理过程。

技术框架：OREO的整体框架包括以下几个关键部分：首先，使用离线数据集进行训练。然后，通过优化软贝尔曼方程，联合学习策略模型和价值函数。策略模型负责生成推理步骤，而价值函数负责评估这些步骤的质量。在测试阶段，可以使用学习到的价值函数来指导树搜索，从而进一步提高性能。

关键创新：OREO的关键创新在于其联合学习策略模型和价值函数的方式，以及利用软贝尔曼方程进行优化。与DPO等方法相比，OREO不需要配对偏好数据，并且能够更好地进行信用分配。此外，利用学习到的价值函数指导树搜索也是一个重要的创新点。

关键设计：OREO的关键设计包括：(1) 软贝尔曼方程的具体形式，它决定了价值函数的更新方式；(2) 策略模型和价值函数的网络结构，需要根据具体的任务进行选择；(3) 训练过程中的超参数设置，例如学习率、批量大小等；(4) 如何利用价值函数指导树搜索的具体算法。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果表明，OREO在GSM8K、MATH等数学推理任务以及ALFWorld具身智能体控制任务上均取得了显著的性能提升，超越了现有的离线学习方法。具体的数据提升幅度在论文中有所展示，表明OREO在多步推理能力方面具有显著优势。此外，利用学习到的价值函数指导树搜索可以进一步提高性能。

🎯 应用场景

OREO具有广泛的应用前景，可以应用于各种需要多步推理能力的场景，例如数学问题求解、代码生成、机器人控制、游戏AI等。该方法可以帮助LLM更好地理解复杂任务，并生成更准确、更可靠的推理结果。此外，OREO还可以用于开发更智能的对话系统和智能助手。

📄 摘要（原文）

Improving the multi-step reasoning ability of large language models (LLMs) with offline reinforcement learning (RL) is essential for quickly adapting them to complex tasks. While Direct Preference Optimization (DPO) has shown promise in aligning LLMs with human preferences, it is less suitable for multi-step reasoning tasks because (1) DPO relies on paired preference data, which is not readily available for multi-step reasoning tasks, and (2) it treats all tokens uniformly, making it ineffective for credit assignment in multi-step reasoning tasks, which often come with sparse reward. In this work, we propose OREO (Offline Reasoning Optimization), an offline RL method for enhancing LLM multi-step reasoning. Building on insights from previous works of maximum entropy reinforcement learning, it jointly learns a policy model and value function by optimizing the soft Bellman Equation. We show in principle that it reduces the need to collect pairwise data and enables better credit assignment. Empirically, OREO surpasses existing offline learning methods on multi-step reasoning benchmarks, including mathematical reasoning tasks (GSM8K, MATH) and embodied agent control (ALFWorld). The approach can be extended to a multi-iteration framework when additional resources are available. Furthermore, the learned value function can be leveraged to guide the tree search for free, which can further boost performance during test time.