CPL: Critical Plan Step Learning Boosts LLM Generalization in Reasoning Tasks

📄 arXiv: 2409.08642v2 📥 PDF

作者: Tianlong Wang, Junzhe Chen, Xueting Han, Jing Bai

分类: cs.AI, cs.LG

发布日期: 2024-09-13 (更新: 2024-10-01)

💡 一句话要点

CPL：关键计划步骤学习提升LLM在推理任务中的泛化能力

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

关键词: 大型语言模型 强化学习 推理任务 泛化能力 蒙特卡洛树搜索

📋 核心要点

  1. 现有LLM推理方法侧重于特定任务，缺乏跨任务的泛化能力，难以应对复杂推理场景。
  2. 提出关键计划步骤学习（CPL），通过蒙特卡洛树搜索探索计划空间，并利用优势偏好优化学习关键步骤。
  3. 实验表明，CPL在GSM8K和MATH上训练后，显著提升了在多个领域外推理基准上的性能。

📝 摘要（中文）

后训练，特别是使用自博弈生成数据的强化学习（RL），已成为大型语言模型（LLM）的一种新的学习范式。然而，扩展RL以开发通用推理器仍然是一个研究挑战，因为现有方法侧重于特定于任务的推理，而没有充分解决跨更广泛任务的泛化问题。此外，与具有有限动作空间的传统RL不同，LLM在无限空间中运行，因此搜索有价值且多样化的策略以有效地解决问题至关重要。为了解决这个问题，我们提出在高层抽象计划中的动作空间内进行搜索，以增强模型的泛化能力，并引入关键计划步骤学习（CPL），包括：1）在计划上搜索，使用蒙特卡洛树搜索（MCTS）来探索多步推理任务中的多样化计划步骤；2）通过步骤级优势偏好优化（Step-APO）学习关键计划步骤，该优化将通过MCTS获得的步骤偏好的优势估计集成到直接偏好优化（DPO）中。这种组合有助于模型有效地学习关键计划步骤，从而提高推理能力和泛化能力。实验结果表明，我们的方法仅在GSM8K和MATH上训练，不仅显著提高了在GSM8K（+10.5%）和MATH（+6.5%）上的性能，还提高了超出领域的推理基准，如HumanEval（+12.2%）、GPQA（+8.6%）、ARC-C（+4.0%）、MMLU-STEM（+2.2%）和BBH（+1.8%）。

🔬 方法详解

问题定义：现有的大型语言模型（LLM）在推理任务中表现出色，但通常依赖于特定任务的训练数据，导致泛化能力不足。尤其是在复杂的多步推理问题中，模型难以找到有效的策略，并且容易受到训练数据偏差的影响。因此，如何提升LLM在推理任务中的泛化能力，使其能够适应更广泛的任务和场景，是一个重要的研究问题。

核心思路：CPL的核心思路是在高层抽象计划层面进行搜索，而不是直接在无限的动作空间中搜索。通过将推理过程分解为多个关键的计划步骤，并使用蒙特卡洛树搜索（MCTS）来探索不同的计划组合，从而找到更有效的推理路径。然后，利用步骤级优势偏好优化（Step-APO）来学习这些关键的计划步骤，使得模型能够更好地理解和利用这些步骤来解决问题。

技术框架：CPL主要包含两个阶段：计划搜索和关键步骤学习。在计划搜索阶段，使用MCTS来探索多步推理任务中的不同计划步骤，构建一个计划树。每个节点代表一个计划步骤，边代表步骤之间的转移。MCTS通过模拟和回溯来评估每个计划步骤的价值，并选择最有希望的步骤进行扩展。在关键步骤学习阶段，使用Step-APO来学习关键的计划步骤。Step-APO将MCTS获得的步骤偏好的优势估计集成到直接偏好优化（DPO）中，从而使得模型能够更好地学习和利用这些关键步骤。

关键创新：CPL的关键创新在于将计划搜索和关键步骤学习相结合，从而有效地提升了LLM在推理任务中的泛化能力。与传统的强化学习方法不同，CPL不是直接在无限的动作空间中搜索，而是在高层抽象计划层面进行搜索，从而降低了搜索的复杂度，并提高了搜索的效率。此外，Step-APO能够有效地利用MCTS获得的优势信息，从而使得模型能够更好地学习关键的计划步骤。

关键设计：在计划搜索阶段，MCTS的奖励函数设计至关重要，它需要能够准确地评估每个计划步骤的价值。在Step-APO中，优势估计的计算方式以及DPO的超参数设置都会影响模型的学习效果。此外，如何将推理任务分解为合适的计划步骤也是一个重要的设计问题。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，CPL方法在GSM8K和MATH数据集上进行了训练后，不仅在这些数据集上取得了显著的性能提升（GSM8K +10.5%，MATH +6.5%），而且在多个领域外推理基准上也取得了显著的提升，包括HumanEval (+12.2%)、GPQA (+8.6%)、ARC-C (+4.0%)、MMLU-STEM (+2.2%)和BBH (+1.8%)。这些结果表明，CPL方法能够有效地提升LLM在推理任务中的泛化能力。

🎯 应用场景

CPL方法具有广泛的应用前景，可以应用于各种需要复杂推理的任务，例如数学问题求解、代码生成、科学发现等。通过提升LLM的推理能力和泛化能力，CPL可以帮助人们更有效地解决各种实际问题，并推动人工智能技术的发展。未来，CPL还可以与其他技术相结合，例如知识图谱、符号推理等，从而进一步提升LLM的推理能力。

📄 摘要（原文）

Post-training, particularly reinforcement learning (RL) using self-play-generated data, has become a new learning paradigm for large language models (LLMs). However, scaling RL to develop a general reasoner remains a research challenge, as existing methods focus on task-specific reasoning without adequately addressing generalization across a broader range of tasks. Moreover, unlike traditional RL with limited action space, LLMs operate in an infinite space, making it crucial to search for valuable and diverse strategies to solve problems effectively. To address this, we propose searching within the action space on high-level abstract plans to enhance model generalization and introduce Critical Plan Step Learning (CPL), comprising: 1) searching on plan, using Monte Carlo Tree Search (MCTS) to explore diverse plan steps in multi-step reasoning tasks, and 2) learning critical plan steps through Step-level Advantage Preference Optimization (Step-APO), which integrates advantage estimates for step preference obtained via MCTS into Direct Preference Optimization (DPO). This combination helps the model effectively learn critical plan steps, enhancing both reasoning capabilities and generalization. Experimental results demonstrate that our method, trained exclusively on GSM8K and MATH, not only significantly improves performance on GSM8K (+10.5%) and MATH (+6.5%), but also enhances out-of-domain reasoning benchmarks, such as HumanEval (+12.2%), GPQA (+8.6%), ARC-C (+4.0%), MMLU-STEM (+2.2%), and BBH (+1.8%).