摘要

本文提出了一种新的训练数据采样方法，旨在加速结合语言模型（LLM）的A搜索算法。传统的A搜索结合LLM的方法存在计算成本高且性能提升不显著的问题。本文通过识别和利用搜索树中对加速A*搜索贡献最大的节点，提出了一种新的数据分布来下采样训练数据，从而在保持计算成本受限的同时，学习到更有效的启发式函数。实验结果表明，该方法在迷宫导航和推箱子（Sokoban）两个经典规划领域中，能够将找到解决方案所需的迭代次数减少高达13倍，实际运行时间加速高达5倍。

原理

本文的核心创新在于通过量化搜索树中每个节点对减少A搜索长度的贡献，从而设计了一种新的数据采样分布D(n, τ)。该分布优先选择那些对搜索过程贡献大的节点进行训练，特别是那些接近目标的节点。通过这种方式，语言模型能够更有效地学习到启发式函数，从而在A搜索中更快地引导搜索方向。此外，该方法还考虑了语言模型对训练数据的需求，确保了模型在低数据 regime 下的泛化能力。

流程

1. 数据准备：从已解决的问题中提取搜索树节点作为训练数据。
2. 节点贡献量化：使用公式C(n) = log( |π∗| / (|π∗| − g(n))) 量化每个节点n对搜索长度的贡献。
3. 数据采样：根据量化结果，使用D(n, τ)分布对训练数据进行下采样。
4. 模型训练：使用下采样后的数据训练语言模型，学习启发式函数。
5. A*搜索集成：在A*搜索的扩展步骤中，利用训练好的语言模型计算启发式值，引导搜索方向。
6. 性能评估：在迷宫导航和推箱子两个领域中，通过IID和OOD测试集评估搜索效率和解决方案质量。

应用

该方法不仅限于迷宫导航和推箱子游戏，还可以扩展到其他需要复杂规划和搜索的领域，如机器人路径规划、物流优化等。通过提高搜索效率，该方法有望在实际应用中大幅减少计算时间和资源消耗，特别是在需要快速响应和高效率的场景中。