[논문 리뷰] Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training

Xidong Feng, Ziyu Wan외 5명|arXiv (Cornell University)|2023. 09. 29.

Natural Language Processing TechniquesComputer Science인용 수 3

한 줄 요약

이 논문은 대규모 언어모델(LM)의 디코딩과 훈련을 안내하는 데 학습된 가치 함수를 사용하는 AlphaZero 유사 트리 탐색 프레임워크인 TS-LLM을 제안한다. 깊은 트리 탐색과 정책 분산, 가치 함수 학습을 통합함으로써 TS-LLM은 추론, 계획, 일치성, 의사결정 작업에서 최신 기술 수준의 성능을 달성하며, 이는 이전 방법이 제한된 깊이 7–10에 비해 훨씬 높은 깊이 64까지 확장된 성능이다.

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가 생성한 트레이젝터리에 기반해 지도 미세조정을 통해 학습된 가치 함수를 사용하는 AlphaZero 유사 트리 탐색 알고리즘을 도입한다.
• 학습된 가치 함수를 사용한 몬테카를로 트리 탐색(MCTS)을 통해 디코딩 중 행동 선택을 안내함으로써 깊이 64까지의 깊은 탐색을 가능하게 한다.
• 트리 탐색 트레이젝터리에서의 지식을 LLM에 전달하기 위해 정책 분산을 활용하여 LLM의 정책을 향상시킨다.
• 트리 탐색 트레이젝터리의 진짜 레이블을 사용해 가치 함수를 학습함으로써 정책 및 가치 함수의 반복적 향상이 가능하도록 한다.
• 동일한 파이프라인을 사용해 추론(Gamma24), 질문 응답(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 정확도와 확장성 측면에서 기존 방법들인 ToT와 RAP에 비해 TS-LLM의 성능은 어떻게 비교되는가?
• 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)를 뛰어넘었다.
• 이 프레임워크는 기존 ToT와 RAP가 제한된 깊이 7–10을 훨씬 초월해 깊이 64까지 확장되는 데 성공했다.
• 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.