Stanford CS329A: AI 에이전트의 계획 및 다단계 추론

개요

AI 에이전트의 다단계 추론은 추론(무엇을 할지 계획), 행동(도구나 검색을 통해 단계 실행), 탐색(피드백을 기반으로 경로를 다듬기)의 통합이 필요합니다. 핵심 과제는 단순 순차 생성에서 벗어나 모델이 올바른 해답에 도달하는 방식을 다양화·탐색·효율적으로 만드는 것입니다.

Language Agent Tree Search (LATS)

LATS는 언어 모델 경로에 몬테카를로 트리 탐색(MCTS)을 적용해 추론·행동·계획을 하나로 통합합니다. 선형 사유 흐름을 트리 기반 탐색으로 전환해 최종 답변으로 가는 경로를 최적화합니다.

LATS 프레임워크

LATS는 여섯 단계로 구성됩니다:

  1. Selection(선택): Upper Confidence Bound applied to Trees(UCT)를 사용해 새로운 경로 탐색과 기존 고가치 경로 활용 사이의 균형을 맞추며 확장할 노드를 선택합니다.
  2. Expansion(확장): 모델이 선택된 노드에서 여러 가능한 행동을 샘플링합니다.
  3. Evaluation(평가): 각 결과 상태를 LLM-as-a-Judge(모델에 0‑1 점수를 요청)와 self-consistency score(특정 행동이 샘플링된 빈도)의 가중 평균으로 점수화합니다.
  4. Simulation(시뮬레이션): 모델이 가장 높은 점수의 상태를 탐욕적으로 확장해 종료 상태(성공 또는 실패)에 도달하거나 연산 예산을 초과할 때까지 진행합니다.
  5. Backpropagation(역전파): 경로의 결과를 사용해 해당 결과에 이르는 모든 상태의 값을 업데이트합니다.
  6. Reflection(반성): 모델이 경로를 분석해 성공·실패 원인을 파악하고, 이를 컨텍스트에 추가해 향후 탐색을 개선합니다.

성능 및 트레이드오프

LATS는 HotPotQA(다중 홉 추론)와 WebShop(실제 전자상거래 작업)에서 인간 전문가 수준에 근접한 성능을 보였으며 파인튜닝 없이도 좋은 결과를 얻었습니다. 그러나 트리를 반복적으로 확장·역전파해야 하므로 추론 비용이 크게 증가합니다. 또한 현재는 시도와 오류 탐색이 불가능한 금융 거래와 같은 되돌릴 수 없는 행동을 다루지 못합니다.

SPRINT: 병렬 추론 트레이스

SPRINT는 프론티어 모델(예: DeepSeek‑R1, Gemini)에서 순차적인 "생각"의 비효율성을 해결합니다. 긴 추론 체인이 정확도와 상관관계가 높지만 지연 시간과 비용을 증가시키는 문제를 완화합니다.

병렬 계획 및 실행

SPRINT는 많은 추론 단계가 독립적이며 병렬로 실행될 수 있다는 관찰에 기반합니다. 이 프레임워크는 사후 학습/파인튜닝 접근법을 사용해 모델이 플래너와 실행자 역할을 동시에 수행하도록 교육합니다:

  • Data Creation(데이터 생성): DeepSeek‑R1과 같은 모델의 추론 트레이스를 GPT‑4o가 주석 달아 계획 단계와 실행 단계를 구분합니다. 그런 다음 이를 Directed Acyclic Graph(DAG)로 정리해 동시에 실행 가능한 단계들을 식별합니다.
  • Fine-tuning(파인튜닝): 모델을 이러한 재구성된 트레이션에 대해 감독‑파인튜닝(SFT)하여 병렬 계획(예: "Plan 1", "Plan 2")과 해당 실행을 위한 특정 태그를 출력하도록 학습시킵니다.
  • Inference(추론): 테스트 시 모델이 여러 독립적인 계획을 생성하고, 이 계획들을 병렬로 실행(예: 여러 Python 도구 호출)한 뒤 결과를 컨텍스트에 다시 합쳐 다음 계획 단계로 진행합니다.

주요 발견

  • 효율성: SPRINT는 답변에 도달하기 위해 필요한 순차 토큰 수를 크게 줄여, 일부 작업에서 순차 토큰 수를 약 40% 감소시켰습니다.
  • 정확도: 단순 효율성 향상이 목표였음에도, 구조화된 병렬 사고 과정이 MATH, Countdown, GPQA Diamond와 같은 벤치마크에서 정확도를 실제로 향상시켰습니다.
  • 일반화: MATH 데이터로 훈련된 모델은 추가 학습 없이도 도메인 외 작업에서 병렬 사고 능력을 발휘했습니다.

SWiRL: 합성 다단계 RL

SWiRL은 도구 사용과 다단계 추론을 가르치면서, 강화학습(RL) 과정에서 도구를 실시간으로 실행하는 비용과 불안정성을 없애는 데 초점을 맞춥니다.

SWiRL 방법론

SWiRL은 데이터 수집과 모델 최적화를 분리합니다:

  1. Offline Synthetic Generation(오프라인 합성 생성): 모델이 다단계 트레이션(추론 → 도구 호출 → 환경 응답)을 반복적으로 생성합니다.
  2. Process Labeling(프로세스 라벨링): LLM-as-a-Judge가 각 개별 단계에 대해 제안된 행동/쿼리의 품질에 기반해 보상을 할당합니다(도구 실제 출력과 무관).
  3. Multi-Step RL(다단계 RL): 모델은 RL을 사용해 이전 컨텍스트가 주어졌을 때 단일 행동의 기대 보상을 최적화하도록 학습합니다. 핵심은 훈련 중 도구를 실행하지 않고, 사전에 수집된 오프라인 데이터에서 "좋은" 행동을 제안한 것에 대해 보상을 받는다는 점입니다.

데이터 필터링 및 일반화

  • Process vs. Outcome Filtering(프로세스 vs. 결과 필터링): 연구진은 process‑filtered data(LLM이 "좋다"고 판단한 단계 유지)가 outcome‑filtered data(정답만 남긴 데이터)보다 훈련에 더 효과적임을 발견했습니다. 이는 올바른 중간 추론을 학습하는 것이 단순히 최종 정답을 모방하는 것보다 가치가 크다는 것을 의미합니다.
  • Cross-Domain Transfer(교차 도메인 전이): SWiRL은 놀라운 일반화를 보였습니다. GSM8K(계산기 도구)에서 훈련된 모델이 HotPotQA(검색 도구)에서 성능이 향상되었고, 그 반대도 마찬가지였습니다. 이는 모델이 특정 도구 API가 아니라 다단계 도구 호출 및 추론이라는 일반 메타 스킬을 학습하고 있음을 시사합니다.
  • RL vs. SFT: 다단계 RL은 감독 파인튜닝(SFT)보다 크게 앞섰습니다. RL은 모델이 모방 데이터에 존재하는 잘못된 경로를 벗어나 탐색할 수 있게 해줍니다.

Sources