Toward Self-Improvement of LLMs via Imagination, Searching, and Criticizing

Despite the impressive capabilities of Large Language Models (LLMs) on various tasks, they still struggle with scenarios that involves complex reasoning and planning. Recent work proposed advanced prompting techniques and the necessity of fine-tuning with high-quality data to augment LLMs' reasoning abilities. However, these approaches are inherently constrained by data availability and quality. In light of this, self-correction and self-learning emerge as viable solutions, employing strategies that allow LLMs to refine their outputs and learn from self-assessed rewards. Yet, the efficacy of LLMs in self-refining its response, particularly in complex reasoning and planning task, remains dubious. In this paper, we introduce AlphaLLM for the self-improvements of LLMs, which integrates Monte Carlo Tree Search (MCTS) with LLMs to establish a self-improving loop, thereby enhancing the capabilities of LLMs without additional annotations. Drawing inspiration from the success of AlphaGo, AlphaLLM addresses the unique challenges of combining MCTS with LLM for self-improvement, including data scarcity, the vastness search spaces of language tasks, and the subjective nature of feedback in language tasks. AlphaLLM is comprised of prompt synthesis component, an efficient MCTS approach tailored for language tasks, and a trio of critic models for precise feedback. Our experimental results in mathematical reasoning tasks demonstrate that AlphaLLM significantly enhances the performance of LLMs without additional annotations, showing the potential for self-improvement in LLMs.

한 줄 요약

AlphaLLM은 MCTS와 LLM을 결합한 자가 개선 프레임워크로, 수학적 추론 과제에서 Llama-2-70B 기반 모델의 정확도를 57.8%에서 92.0%까지 향상시킨다.

핵심 기여도

AlphaLLM은 MCTS와 LLM을 결합한 자가 개선 프레임워크로, 추가 라벨 데이터 없이도 모델 성능을 향상시킨다.
ηMcts는 언어 과제에 최적화된 효율적인 탐색 알고리즘으로, 상태 병합과 적응적 분기 인수를 도입하여 탐색 효율성을 높인다.
세 가지 평가 모델(Value Function, PRM, ORM)을 사용하여 정확한 피드백을 제공하며, ORM은 복잡한 수학 문제에 외부 도구를 결합하여 정확도를 높인다.
GSM8K에서 57.8% → 92.0%, MATH에서 20.7% → 51.0%의 성능 향상이 실험적으로 검증되었다.

핵심 아이디어

AlphaLLM은 LLM의 자가 개선을 위해 상상(Imagination), 탐색(Searching), 비판(Criticizing)의 세 단계를 결합한 프레임워크이다. 이는 AlphaGo의 성공 요소를 언어 모델에 적용한 것으로, 특히 복잡한 추론 과제에서 효과적이다. ηMcts는 언어 생성 과정을 MDP 옵션으로 모델링하여, 각 옵션이 특정 하위 과제에 대한 토큰 집합을 생성하도록 설계되어 탐색 깊이를 줄이고 효율성을 높인다. 또한, 상태 병합(State Merge)과 적응적 분기 인수(Adaptive Branching Factor)를 도입하여 탐색 공간의 너비와 깊이를 균형 있게 조절한다. 평가 모델인 ORM은 외부 도구(예: Python SymPy)를 활용하여 복잡한 수학 문제의 정확도를 판단하며, 이는 GPT-4 수준의 성능을 달성하는 데 기여한다.

기술적 접근법

**모델 구성**: AlphaLLM은 Prompt Synthesis, ηMcts, 세 가지 평가 모델(Value Function, PRM, ORM)로 구성된다.
**ηMcts**: 언어 과제에 최적화된 MCTS 알고리즘으로, 옵션 기반 MDP 모델링, 상태 병합, 적응적 분기 인수를 도입하여 탐색 효율성을 높인다.
**평가 모델**: Value Function은 기대 보상을 추정, PRM은 노드 정확도를 평가, ORM은 전체 경로를 평가하며, MATH 데이터셋에서는 SymPy 도구를 활용하여 정확도를 높인다.
**학습 설정**: 학습률 1e-6, 1 에포크, 배치 사이즈 128, 최소 학습률 1e-6, 선형 워밍업 및 감소 적용.
**데이터셋**: GSM8K(전체 테스트셋), MATH(부분 집합).

주요 결과

**GSM8K**: Llama-2-70B 기반 AlphaLLM은 57.8% → 92.0%로 향상 (GPT-4 수준).
**MATH**: WizardMath-70B 기반 AlphaLLM은 20.7% → 51.0%로 향상.
**ηMcts 성능**: GSM8K에서 88.9%, MATH에서 48.7%의 정확도를 달성.
**성분별 분석**: Adaptive Branching → 84.9%, PRM 추가 → 85.9%, ORM 추가 → 86.5%, 상태 병합 → 87.0%로 성능 향상.

의의 및 한계

AlphaLLM은 추가 라벨 데이터 없이도 LLM의 성능을 향상시키는 자가 개선 프레임워크로, 복잡한 추론 과제에서 GPT-4 수준의 성능을 달성하는 데 기여한다. 특히, ηMcts와 세 가지 평가 모델의 결합은 언어 과제의 대규모 탐색 공간과 주관적 피드백 문제를 해결하는 데 효과적이다. 그러나 MATH 데이터셋의 경우 외부 도구 없이 ORM이 38.8%의 정확도로 저하되는 한계가 있으며, 이는 도구 의존성의 중요성을 드러낸다. 또한, 계산 자원 제약으로 MATH에서는 단 한 번의 자가 개선만 수행되었으며, 이는 성능 향상의 잠재력을 완전히 발휘하지 못한 측면이 있다.

실용적 활용

AlphaLLM은 수학 문제 해결, 코드 생성, 복잡한 문서 분석 등 정확한 추론이 필요한 산업 분야에서 활용 가능하다. 특히, 라벨 데이터가 제한된 상황에서 자가 개선을 통해 모델 성능을 지속적으로 향상시킬 수 있어, 교육, 금융, 의료 분야의 전문가 시스템 개발에 기여할 수 있다.