[论文解读] Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training

Xidong Feng, Ziyu Wan+5 more|arXiv (Cornell University)|Sep 29, 2023

Natural Language Processing TechniquesComputer Science被引用 3

一句话总结

该论文提出 TS-LLM，一种类似 AlphaZero 的树搜索框架，利用学习到的值函数引导大语言模型（LLM）的解码与训练。通过整合深度树搜索、策略蒸馏与值函数学习，TS-LLM 在推理、规划、对齐与决策任务中达到最先进性能，搜索深度达 64，显著超越此前方法受限于 7–10 深度的局限。

ABSTRACT

Recent works like Tree-of-Thought (ToT) and Reasoning via Planning (RAP) aim to augment the reasoning capabilities of LLMs by using tree-search algorithms to guide multi-step reasoning. These methods rely on prompting a pre-trained model to serve as a value function and focus on problems with low search depth. As a result, these methods will not work in domains where the pre-trained LLM does not have enough knowledge to serve as an effective value function or in domains that require long-horizon planning. To address these limitations, we present an AlphaZero-like tree-search learning framework for LLMs (termed TS-LLM), systematically illustrating how tree-search with a learned value function can guide LLM decoding. TS-LLM distinguishes itself in two key ways. (1) Leveraging a learned value function and AlphaZero-like algorithms, our approach can be generally adaptable to a wide range of tasks, language models of any size, and tasks of varying search depths. (2) Our approach can guide LLMs during both inference and training, iteratively improving the LLM. Empirical results across reasoning, planning, alignment, and decision-making tasks show that TS-LLM outperforms existing approaches and can handle trees with a depth of 64.

研究动机与目标

• 为解决现有树搜索方法在 LLM 中的局限，如依赖提示工程的值函数与浅层搜索深度。
• 开发一种可泛化且可扩展的框架，适用于多样化任务、不同规模 LLM 与不同搜索深度。
• 使树搜索能够同时指导 LLM 的推理与端到端训练，通过迭代方式提升模型性能。
• 证明学习到的值函数在可靠性与有效性上优于基于提示的自我评估。

提出的方法

• 提出一种受 AlphaZero 启发的树搜索算法，其中值函数通过在 LLM 生成的轨迹上进行监督微调来学习。
• 使用蒙特卡洛树搜索（MCTS）结合学习到的值函数，在解码过程中引导动作选择，实现高达 64 深度的深度搜索。
• 采用策略蒸馏技术，将树搜索轨迹中的知识迁移至 LLM，以改进其策略。
• 利用树搜索轨迹中的真实标签学习值函数，实现策略与值函数的迭代优化。
• 将该框架应用于多种任务，包括推理（Game24）、问答（PrOntoQA）与 RLHF 对齐，使用统一的流水线。
• 采用混合搜索策略，结合树内与树间树搜索，以平衡探索与利用。

Figure 1: (a) Left: Two node expansion paradigms on Game24: sentence-level and token-level. We adopt sentence-level setting in this task. (b) Right: TS-LLM consists of an iterative process over tree-search and training. First, TS-LLM enhances LLM inference by tree-search to obtain improved trajector

实验结果

研究问题

• RQ1 在树搜索框架中，学习到的值函数是否能在 LLM 推理任务中优于基于提示的自我评估？
• RQ2 能否将类似 AlphaZero 的树搜索扩展至自然语言生成中的深度推理任务，实现 64 深度的搜索？
• RQ3 将树搜索整合到 LLM 训练中，是否能提升推理、规划与对齐等多样化任务的性能？
• RQ4 在准确率与可扩展性方面，TS-LLM 与 ToT 和 RAP 等现有方法相比表现如何？
• RQ5 树搜索与策略/值函数优化的迭代循环是否能为 LLM 带来持续的性能提升？

主要发现

• 在 PrOntoQA 任务中，TS-LLM 在深度 20 时达到 100% 准确率，显著优于所有基线方法，包括 MCTS-Rollout 与 CoT-SC-Tree。
• 在 Game24 任务中，TS-LLM 使用 MCTS-α-intra 树搜索在深度 100 时达到 84.53% 准确率，远超 CoT-SC（18.23%）与 BFS-V（72.65%）。
• 在 RLHF 对齐任务中，TS-LLM 使用 MCTS-Rollout 在 N=50 时达到平均奖励 2.491，最佳奖励 2.746，超越 CoT（0.387）与 BFS-V（-1.474）。
• 该框架成功扩展至 64 深度，远超 ToT 与 RAP 方法的 7–10 深度限制。
• TS-LLM 在所有评估任务中均实现性能提升——包括推理、规划、对齐与决策，展现出良好的泛化能力。
• TS-LLM 中学习到的值函数比基于提示的自我评估更可靠，即使在较小 LLM（如 125M 参数）上也表现出一致的性能增益。

Figure 2: Aggregation results for four tasks w.r.t. number of sequences(Path@ $N$ ) on the 1st row and the number of tokens on the 2nd row. TS-LLM benefits from aggregation but struggles to scale in small-scale problems.