LiteSearch: LLM을 위한 효율적인 트리 탐색

LiteSearch: Efficacious Tree Search for LLM

June 29, 2024

저자: Ante Wang, Linfeng Song, Ye Tian, Baolin Peng, Dian Yu, Haitao Mi, Jinsong Su, Dong Yu

cs.AI

초록

최근 연구에 따르면 트리 탐색 알고리즘(예: 몬테카를로 트리 탐색)이 복잡한 수학적 추론 과제에서 대형 언어 모델(LLM)의 성능을 극적으로 향상시킬 수 있다고 합니다. 그러나 이러한 알고리즘은 비효율적인 탐색 전략으로 인해 탐욕적 디코딩보다 10배 이상의 계산 자원을 필요로 하기 때문에 실제 애플리케이션에 배포하기가 어렵습니다. 본 연구에서는 이러한 문제를 해결하기 위해 동적 노드 선택과 노드 수준 탐색 예산(최대 자식 노드 수) 계산을 포함한 새로운 가이드 트리 탐색 알고리즘을 제안합니다. 최종 답변을 향한 탐색 진행 상황(과거)과 단계별 주석 없이 훈련된 값 네트워크(미래)의 지도를 고려함으로써, 이 알고리즘은 할당된 계산 예산 범위 내에서 가장 유망한 트리 노드를 반복적으로 선택한 후 이를 확장합니다. GSM8K 및 TabMWP 데이터셋에서 수행된 실험 결과, 우리의 접근 방식은 경쟁력 있는 성능을 제공할 뿐만 아니라 기준 방법 대비 훨씬 낮은 계산 비용을 자랑하는 것으로 나타났습니다.

English

Recent research suggests that tree search algorithms (e.g. Monte Carlo Tree Search) can dramatically boost LLM performance on complex mathematical reasoning tasks. However, they often require more than 10 times the computational resources of greedy decoding due to wasteful search strategies, making them difficult to be deployed in practical applications. This study introduces a novel guided tree search algorithm with dynamic node selection and node-level exploration budget (maximum number of children) calculation to tackle this issue. By considering the search progress towards the final answer (history) and the guidance from a value network (future) trained without any step-wise annotations, our algorithm iteratively selects the most promising tree node before expanding it within the boundaries of the allocated computational budget. Experiments conducted on the GSM8K and TabMWP datasets demonstrate that our approach not only offers competitive performance but also enjoys significantly lower computational costs compared to baseline methods.