A Training Data Recipe to Accelerate A* Search with Language Models

📄 arXiv: 2407.09985v2 📥 PDF

作者: Devaansh Gupta, Boyang Li

分类: cs.AI, cs.LG

发布日期: 2024-07-13 (更新: 2024-10-23)

备注: Accepted to Findings of EMNLP, 2024

🔗 代码/项目: GITHUB

💡 一句话要点

提出一种训练数据选择方法，加速基于语言模型的A*搜索

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: A*搜索 语言模型 启发式学习 数据选择 强化学习

📋 核心要点

  1. 现有方法在学习启发式函数时，较少考虑搜索算法与机器学习模型之间的交互，导致训练效率低下。
  2. 该论文的核心思想是，通过分析A*算法和LLM的需求，发现二者对靠近目标的搜索节点有共同的需求，从而设计数据选择分布。
  3. 实验结果表明，该方法在经典规划问题上显著减少了迭代次数和搜索时间，提高了训练效率和搜索速度。

📝 摘要（中文）

本文研究了利用大型语言模型（LLM）进行启发式搜索算法（如A）时，加速训练和降低计算需求的训练数据选择问题。现有方法较少考虑搜索算法与机器学习模型之间的交互。本文从经验上分离了A搜索算法的需求与LLM泛化的需求，发现二者存在重叠：A*需要更准确地预测靠近目标的搜索节点，而LLM也需要相同的节点来实现有效的泛化。基于此，本文推导出一种用于学习基于LLM的启发式方法的数据选择分布。在迷宫导航、推箱子和滑块拼图三个经典规划领域，该技术将找到解决方案所需的迭代次数减少高达15倍，并将搜索的实际运行速度提高高达5倍。

🔬 方法详解

问题定义：论文旨在解决如何高效地训练基于LLM的启发式函数，以加速A搜索算法的问题。现有方法在训练LLM时，通常采用随机采样或简单启发式方法选择训练数据，忽略了A算法本身的特性，导致训练效率低下，模型泛化能力不足。现有方法的痛点在于没有充分利用A*算法的搜索过程信息来指导训练数据的选择。

核心思路：论文的核心思路是，A算法在搜索过程中，对靠近目标状态的节点预测的准确性要求更高，而LLM的泛化能力也依赖于对这些关键节点的学习。因此，可以通过优先选择靠近目标状态的节点作为训练数据，来同时满足A算法和LLM的需求，从而提高训练效率和模型性能。这种数据选择策略能够使LLM更专注于学习对A*搜索过程至关重要的信息。

技术框架：该论文的技术框架主要包括以下几个步骤：1. 使用A算法在给定的环境中进行搜索，记录搜索过程中的节点信息。2. 分析搜索过程中的节点，根据节点到目标状态的距离（或其他相关指标）来评估节点的重要性。3. 基于节点的重要性，构建一个数据选择分布，优先选择靠近目标状态的节点作为训练数据。4. 使用选择的训练数据来训练LLM，使其学习启发式函数。5. 使用训练好的LLM作为启发式函数，指导A算法进行搜索。

关键创新：该论文的关键创新在于，它将A算法的搜索过程与LLM的训练过程相结合，提出了一种基于A算法需求的数据选择策略。这种策略能够充分利用A算法的搜索信息，选择对A算法至关重要的节点作为训练数据，从而提高训练效率和模型性能。与现有方法相比，该方法更加关注A算法的特性，能够更有效地训练LLM，并加速A搜索过程。

关键设计：论文的关键设计在于数据选择分布的设计。具体来说，论文根据节点到目标状态的距离，设计了一个概率分布，使得距离目标状态越近的节点，被选择作为训练数据的概率越高。此外，论文还考虑了其他因素，如节点的访问频率、节点的启发式值等，来进一步优化数据选择分布。损失函数采用标准的回归损失函数，用于衡量LLM预测的启发式值与真实启发式值之间的差距。网络结构方面，可以使用各种常见的LLM架构，如Transformer等。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，在迷宫导航、推箱子和滑块拼图三个经典规划领域，该方法将找到解决方案所需的迭代次数减少高达15倍，并将搜索的实际运行速度提高高达5倍。与传统的随机采样方法相比，该方法能够显著提高训练效率和搜索速度，证明了其有效性和优越性。

🎯 应用场景

该研究成果可应用于各种需要启发式搜索的领域，如游戏AI、机器人路径规划、物流优化、自动驾驶等。通过结合LLM的强大泛化能力和A*算法的高效搜索能力，可以解决复杂环境下的规划问题，提高决策效率和智能化水平。未来，该方法有望扩展到更复杂的任务和环境，并与其他AI技术相结合，实现更高级别的智能。

📄 摘要（原文）

Combining Large Language Models (LLMs) with heuristic search algorithms like A holds the promise of enhanced LLM reasoning and scalable inference. To accelerate training and reduce computational demands, we investigate the coreset selection problem for the training data of LLM heuristic learning. Few methods to learn the heuristic functions consider the interaction between the search algorithm and the machine learning model. In this work, we empirically disentangle the requirements of A search algorithm from the requirements of the LLM to generalise on this task. Surprisingly, we find an overlap between their requirements; A* requires more accurate predictions on search nodes near the goal, and LLMs need the same set of nodes for effective generalisation. With these insights, we derive a data-selection distribution for learning LLM-based heuristics. On three classical planning domains, maze navigation, Sokoban and sliding tile puzzles, our technique reduces the number of iterations required to find the solutions by up to 15x, with a wall-clock speed-up of search up to 5x. The codebase is at https://github.com/devaansh100/a_star.