探索LiteSearch:提升LLM数学推理效率的新算法

LiteSearch: Efficacious Tree Search for LLM

摘要

本文介绍了一种名为LiteSearch的新型引导树搜索算法,旨在提高大型语言模型(LLM)在复杂数学推理任务中的性能,同时显著降低计算成本。传统的树搜索算法如蒙特卡洛树搜索(MCTS)虽然能提升LLM的性能,但往往需要超过贪婪解码10倍的计算资源。LiteSearch通过动态节点选择和节点级探索预算(最大子节点数)计算,有效地解决了这一问题。该算法结合搜索历史和价值网络(未来)的指导,迭代选择最有希望的树节点进行扩展,确保在分配的计算预算内进行。实验结果显示,LiteSearch在GSM8K和TabMWP数据集上不仅提供了竞争性的性能,而且计算成本显著低于基线方法。

原理

LiteSearch算法的核心在于其动态节点选择和节点级探索预算计算机制。算法通过价值网络评估每个节点的价值,结合搜索进度(历史)和价值网络(未来)的指导,选择最有希望的节点进行扩展。每个节点的探索预算根据其价值和深度动态计算,确保在搜索过程中有效地平衡探索与利用。具体来说,节点的探索预算与其价值成反比,即价值高的节点分配较少的计算资源,以避免不必要的计算,而价值低的节点则分配更多的计算资源以确保充分的探索。这种机制不仅促进了高效的利用,加速了最终答案的收敛,还保证了足够的探索,以覆盖足够的状态空间,维持性能。

流程

LiteSearch的工作流程主要包括两个主要步骤:选择和扩展。在选择阶段,算法根据节点的价值和进度(通过与贪婪解码的比较计算得出)选择最有希望的节点。在扩展阶段,算法根据节点的价值和深度动态计算其探索预算,并在预算内扩展该节点。这一过程迭代进行,直到生成的答案达到预期的价值阈值或达到最大迭代次数。例如,在处理一个数学推理问题时,LiteSearch首先初始化一个以问题为根的搜索树,然后迭代地选择和扩展节点,直到找到满足条件的答案。

应用

LiteSearch算法在数学推理任务中的高效性能和低计算成本使其在教育技术、自动化问题解决系统、以及需要复杂推理能力的AI助手等领域具有广泛的应用前景。随着LLM在处理更复杂任务的能力不断提升,LiteSearch的动态搜索策略和计算效率优化将使其成为提升LLM推理能力的重要工具。此外,该算法的设计理念和方法也可以扩展到其他需要树搜索策略的应用场景,如游戏AI、规划和调度问题等。