LATS: Language Agent Tree Search — 추론, 행동, 계획을 하나의 프레임워크로 통합

Tree of Thoughts 이후에 무엇이 올지 고민해 왔습니다. 추론 단계를 검색할 수 있다면, 왜 행동도 검색하지 못할까요? 이것이 바로 LATS(Language Agent Tree Search)가 수행하는 일이며, 제가 지금 이 논문을 읽고 있는 이유입니다. Andy Zhou, Kai Yan, Michal Shlapentokh-Rothman, Haohan Wang, Yu-Xiong Wang의 논문(ICML 2024, arXiv:2310.04406)은 단일 에이전트 프레임워크에서 추론, 행동, 계획을 가장 명확하게 합성했으며, 그 결과는 정말 무시하기 어렵습니다.

논문 요약

LATS가 해결하는 핵심 문제는 기존 에이전트 연구의 구조적 격차입니다. ReAct는 추론과 행동을 교차시키지만, 궤적이 잘못되었을 때 뒤로 돌아가 다른 경로를 시도할 메커니즘이 없습니다. Tree of Thoughts는 추론 단계에 대한 분기를 가능하게 하지만 내부 LM 지식에만 의존합니다. 즉, 탐색 중에 도구를 호출하거나 외부 피드백을 받을 수 없습니다. Reflexion은 언어적 자기 수정을 추가하지만, 선형적인 재시도 루프 특성상 대안을 탐색하지 않고 새로운 궤적에 전념하게 됩니다. LATS는 이 세 가지 아이디어를 적절한 몬테카를로 트리 탐색(MCTS) 백본과 융합하여, LLM 에이전트가 여러 분기를 탐색하고, 실제 환경 피드백을 받으며, 경로가 실패할 경우 백트래킹(되돌아가기)을 할 수 있도록 합니다.

기술적 메커니즘은 6단계의 MCTS 루프로 구성됩니다: 선택(Selection, UCT 공식을 통해 탐색할 다음 노드 선택), 확장(Expansion, LM에서 n개의 후보 행동 샘플링), 평가(Evaluation, 하이브리드 가치 함수로 각 노드 점수 매기기), 시뮬레이션(Simulation, 터미널 상태까지 실행), 역전파(Backpropagation, 조상 노드 가치 업데이트), 성찰(Reflection, 실패 시 잘못된 점에 대한 언어적 요약을 생성하여 컨텍스트로 저장). 가치 함수는 주목할 만합니다: V(s) = λ·LM(s) + (1−λ)·SC(s). 여기서 LM(s)는 환경 피드백을 받은 후 궤적의 품질에 대한 LM 자체의 추정치이며, SC(s)는 형제 노드들 사이에서 해당 행동이 얼마나 자주 샘플링되는지에 기반한 자기 일관성 점수입니다. 이것은 훈련된 보상 모델이 아니라 전적으로 프롬프트 기반으로 작동하는 가치 함수입니다.

핵심 아이디어

• HumanEval에서 GPT-4 + LATS는 92.7%의 pass@1을 달성했으며, 이는 GPT-4 + Reflexion의 91.0%, GPT-3.5 + ReAct 단독의 56.9%와 비교됩니다. GPT-3.5 + LATS는 83.8%까지 상승합니다.
• HotPotQA에서 LATS(CoT + ReAct)는 0.71의 Exact Match를 달성하여 ReAct 베이스라인의 0.32보다 두 배 이상의 정확도를 보였습니다.
• WebShop(웹 탐색 및 구매)에서 LATS는 75.9점(성공률 38.0%)을 기록하여 Reflexion의 64.2점(35.0%)을 앞섰습니다. 이는 여러 페이지에 걸쳐 상태를 관리해야 하는 작업에서 의미 있는 차이입니다.
• Game of 24(순수 추론 퍼즐)에서 LATS는 동일한 GPT-4 백본을 사용했음에도 ToT의 0.20보다 높은 0.44의 성공률을 기록했습니다.
• 놀랍게도 LATS는 HotPotQA(k=50)에서 솔루션을 찾기 위해 ToT보다 적은 노드를 확장하고(평균 66.65 vs 84.05), 더 적은 토큰(173,290 vs 210,215)을 사용했습니다. 이론적으로는 더 비싸 보임에도 불구하고 말이죠.

성과와 한계점

벤치마크 수치는 실질적이며 프레임워크는 개념적으로 깔끔합니다. UCT 공식은 ToT의 임시적인 BFS/DFS에는 없는 원칙적인 탐색-이용(exploration-exploitation) 균형을 제공합니다. 순수한 LM 내성 대신 외부 환경 피드백을 가치 함수에 통합한 것은 올바른 방향이며 결과가 이를 증명합니다.

하지만 논문에 담겨 있고 저자들도 인정하지만 충분히 강조되지 않은 치명적인 가정이 있습니다: LATS는 환경을 이전 상태로 되돌리는 기능(state reversion)을 필요로 합니다. 체크포인트 기능 없이는 트리를 분기할 수 없습니다. 한 번 행동을 취하면 되돌릴 수 없기 때문입니다. 저자들은 LM 작업에서 "과거 텍스트 입력을 복사하여 붙여넣기"를 통해 이를 관리할 수 있다고 언급했지만, 실제 행동 환경(데이터베이스, 파일 시스템, 부수 효과가 있는 API)에서 이는 많은 프로덕션 시스템이 충족하기 어려운 엄격한 요구 사항입니다. WebShop 결과가 베이스라인보다 낫긴 하지만, 복잡한 환경에서 자기 성찰은 구체적이기보다 일반적인 내용이 되는 경향이 있습니다. 에이전트는 겉으로는 다르지만 구조적으로는 동일한 실수를 반복하며 정체될 수 있습니다. 논문은 이를 언급하지만 해결책은 제시하지 않습니다.

또한 MCTS 구조 대 가치 함수 설계의 기여도를 분리하는 어블레이션 연구가 없습니다. 동일한 하이브리드 가치 함수를 사용하는 더 간단한 분기 방식도 격차를 상당히 줄일 가능성이 높지만, 저자들은 이를 직접 테스트하지 않았습니다.

금융 AI에서 이것이 중요한 이유

Beancount 장부는 한 가지 핵심적인 이유 때문에 LATS 스타일의 트리 탐색을 위한 거의 이상적인 환경입니다: 모든 장부가 git 리포지토리로 관리된다는 점입니다. 실제 환경에서 LATS를 비실용적으로 만드는 '상태 복원' 요구 사항이 git checkout이나 git stash를 통해 아주 쉽게 충족됩니다. 쓰기 권한을 가진 에이전트는 여러 분기에서 후보 분개(journal entry)를 제안하고, 대차대조표 제약 조건(가치 함수)에 대해 평가하며, 가장 높은 점수를 받은 경로만 커밋할 수 있습니다. 실패한 분기에는 "계정 유형이 잘못 분류되어 자산 = 부채 + 자본 항등식을 위반했습니다"와 같은 언어적 성찰이 제공됩니다.

하이브리드 가치 함수 설계도 직접 적용 가능합니다. 장부 에이전트의 경우, LM(s)는 의미론적 적합성(이것이 올바른 카테고리처럼 보이는가?)을 바탕으로 제안된 항목의 점수를 매기고, SC(s)는 에이전트가 유사한 과거 거래를 얼마나 일관되게 분류하는지 추적하는 장부 자체의 이력에 기반한 자기 일관성 체크 역할을 할 것입니다.

다만 상태 복원 가정에 대해서 금융 비유를 들 때 주의할 점이 있습니다. 실제 장부는 종종 다운스트림 효과를 가집니다. 기록된 항목이 인보이스를 발행하고, 이것이 결제 워크플로우를 트리거할 수 있습니다. 이러한 경우 LATS의 가정은 깨집니다. 장부가 git 제어 하에 있는 일반 텍스트 파일이고 변경 사항이 다운스트림 트리거 전에 로컬에서 이루어지는 Beancount의 경우에만 이 가정이 유효하며, 이는 명시적으로 유지해야 할 설계 제약 조건입니다.

다음 읽을거리

• 환경 모델 없는 MCTS 기반 계획: "Reasoning with Language Model is Planning with World Model" (Hao et al., 2023, arXiv:2305.14992) — LATS가 직접적으로 기반을 두고 개선한 RAP 연구입니다.
• LM 가치 함수는 얼마나 잘 일반화되는가? "Let's Verify Step by Step" (Lightman et al., 2023, arXiv:2305.20050) — 프롬프트 기반 가치 함수의 대안으로서의 프로세스 보상 모델(PRM).
• 비가역적 상황에서의 안전한 다단계 계획: "Decision-Making with Language Models via Successive Prompting" (Creswell et al., 2023) — 상태 복원 요구 사항을 피하는 더 간단한 계획 방식입니다.