Reasoning with Language Model is Planning with World Model

coticoger·2025년 9월 16일

Motivation

LLM은 다양한 작업에서 추론 능력을 보여주었다. 최근 접근 방식은 Chain-of-Thought나 하위 질문을 해결하도록 유도하여 추론 능력을 강화했다. 하지만 LLM은 여전히 사람들이 쉽게 해결할 수 있는 과제들에서 어려움을 겪고 있다.

사람은 internal world model을 가지고 있으며, 이를 통해 사람들이 행동을 시뮬레이션하고 그 행동이 world's state에 미치는 영향을 예측하여 복잡한 작업에서 계획을 세울 수 있다. 예를 들어, 목표에 대한 행동 계획을 수립할 때, world model을 사용한 계획은 다양한 대체 행동 경로를 탐색하고, 가능한 미래 시나리오를 통해 예상되는 결과를 평가한 후, 그 평가를 바탕으로 계획을 반복적으로 개선하는 과정을 수행한다.

이러한 방법은 현재 LLM 추론 방식과 대비된다. LLM은 autoregressive manner로 순차적으로 과거의 결과를 활용하여 새로운 추론을 이어간다. 그러므로 미래에 대한 시나리오를 직접적으로 고려하지 않는다.

사람은 미래를 예측하고 계획을 수정하는 데 집중하는 반면, LLM은 과거의 결과에 의존하여 점차적인 추론을 이어간다

현재 LLM 추론의 핵심적인 한계는 다음과 같다
1. 세계의 상태를 시뮬레이션할 internal world model이 없음
2. 추론이 원하는 상태로 나아가도록 평가하고 유도할 수 있는 보상 메커니즘의 부재
3. 방대한 추론 공간을 효율적으로 탐색하기 위한 탐색과 활용의 균형을 맞추는 능력의 부족

이러한 한계를 극복하기 위해서 이 논문은 RAP(Reasoning Via Planning)을 제시한다.

RAP는 LLM에 world model을 추가하고, MCTS를 통해 효율적인 탐색 후 고보상 추론 경로를 생성한다.

적절한 프롬프트를 사용하여 LLM을 재활용함으로써 world model을 획득함. 추론 과정에서 LLM이 가장 유망한 action을 순차적으로 고려하여 추론 트리를 구축하고, world model을 사용하여 미래의 결과를 예측한다. 추정된 미래 보상은 backpropagate되어 LLM의 현재 추론 단계에 대한 신뢰를 업데이트하고, 이를 통해 더 나은 대안을 탐색하여 추론을 개선하도록 유도함.

Reasoning with LLMs

방법	설명
CoT	복잡한 질문들을 순차적인 중간 단계들로 분해한 후 최종 답을 도출 CoT는 체인을 한 번에 생성하며, 단계 수가 증가함에 따라 오류를 유발할 수 있음
Self-Consistency	여러 체인을 샘플링하여 다수결 투표를 통해 최적의 답을 선택
Least-to-most-prompting	질문을 더 간단한 하위 질문으로 나눈 뒤, 이를 순차적으로 답함
self-evaluation	중간 단계들에 대한 피드백을 제공하여, 각 단계에서 발생할 수 있는 오류를 수정하고 최적의 방향으로 추론을 수정
CoRe	수학 서술형 문제를 풀기 위해 MCTS를 사용하여 추론 단계 생성기와 검증기를 파인튜닝함

⚠️ 위 방법들은 world model을 공식적으로 도입하거나, 보상과 상태를 통합한 프레임워크로 구현하지 않음

RAP은 world model과 보상, advanced planning을 결합하는 더 원초적인 프레임워크임

Planning with LLMs

planning은 intelligent agent의 핵심 능력으로, 특정 목표를 달성하기 위한 일련의 행동을 생성하는 과정임.

최근에는 LLM을 사용하여 계획을 생성하는 접근이 주목받고 있음. 이 방식은 자연어 처리 능력을 활용하여 인간처럼 계획을 수립하는 데 유용함.

기존 코드 기반 계획은 도메인에 제한이 많고, 특정 환경에 종속됨. PDDL을 사용하여 계획을 수립하지만, 좁은 도메인에서만 효과를 갖음.

RAP은 수학적 추론이나 논리적 추론과 같은 오픈 도메인 문제를 처리할 수 있음.

RAP FRAMEWORK

1. Language Model as World Model

RAP은 주어진 문제에 따라 상태와 행동의 일반적인 개념을 다양한 방식으로 구체화함

추론 과정은 MDP로 설명할 수 있다.

주어진 현재 상태 $s_{t,t = 0,1,...,T}$ (초기 상태 $s_0$ ). LMM은 생성 분포 $p(a \mid s_t, c)$ 에서 샘플링하여 행동 공간을 생성함.(여기서 $c$ 는 적절한 프롬프트이다). 행동이 선택되면, world model은 다음 추론 상태 $s_{t+1}$ 를 예측한다. 구체적으로, 이전에 사용한 LLM을 재활용하여 상태 전이 분포 $p(s_{t+1}\mid s_t,a_t,c')$ 를 얻으며, $c'$ 는 LLM이 상태를 생성하도록 안내하는 또 다른 프롬프트이다.

즉, 현재 상태 $s_t$ 에서 LLM이 적절한 프롬프트 $c$ 를 바탕으로 확률적으로 행동 $a_t$ 를 선택함. 그런 다음, 이 선택된 행동 $a_t$ 와 또 다른 프롬프트 $c'$ 를 사용하여 LLM이 다음 상태 $s_{t+1}$ 을 예측함.
이 과정이 반복되면서, LLM은 현재 상태와 선택된 행동을 바탕으로 다음 상태를 예측하고, 이를 통해 계획을 세움

이 과정을 계속 진행하면 추론 경로가 생성되며, 이 경로는 상태와 행동이 교차된 시퀀스 $(s_0,a_0,s_1,...,a_{T-1},s_T)$ 로 구성됨. 이는 CoT와 같은 기존의 추론 방법과 다르다. CoT는 중간 추론 단계가 행동들의 시퀀스로만 구성됨 $(a_0,a_1,...)$

추론에 예측된 world 상태를 추가하는 것은 LLM이 더 구체적이고 일관된 추론을 하는 데 도움이 됨. 중요한 점은, 전체 추론 경로가 LLM 자체로 시뮬레이션되며, 외부 실제 환경과는 상호작용하지 않는다는 것이다.

계획 알고리즘이 통합되어 있기 때문에, LLM은 단순히 주어진 정보에서 다음 단계를 생성하는 것을 넘어서, 복잡한 추론 경로를 효율적으로 탐색하고 최적의 경로를 선택할 수 있음. 이 과정에서 LLM은 다양한 가능성을 반복적으로 평가하고 수정하는 방식으로, 마치 계획을 세우는 인간의 사고처럼 최적의 결과를 도출하려 함.

Reward Design

그렇다면 우리는 어떻게 각 추론 단계의 실행 가능성과 유용성을 평가할까? 각 추론 단계의 평가는 보상 함수 $r_t = r(s_t,a_t) \in R$ 에 의해 수행된다. 보상함수는 다른 작업에 대한 지식이나 선호를 반영하기 위해 여러 가지 방식으로 지정될 수 있음.

Likelihood of the action

LLM이 in-context demonstration(이전의 경험)과 현재 상태에서 조건화된 행동을 생성할 때, 특정 행동의 확률은 LLM의 선호를 반영함. 그러므로 행동의 로그 확률을 보상으로 포함시킬 수 있음. 이 보상은 에이전트로서의 LLM의 "instinct"를 반영하며, 어떤 행동을 탐색할지에 대한 사전 정보로도 사용될 수 있음

Confidence of the state

world model로 부터 여러 샘플 답변을 뽑아내고, 가장 빈번하게 나타난 답변의 비율을 신뢰도로 사용함. 더 높은 신뢰도는 상태 예측이 LLM의 세계 지식과 더 일치함을 나타내며, 이는 일반적으로 더 신뢰할 수 있는 추론 단계로 여겨짐.

Self-evaluation by the LLM

추론에서 오류를 피하는 것보다 알아차리는 것이 더 쉬운 경우가 있음. 따라서 LLM이 "이 추론 단계가 올바른가?"라는 질문으로 스스로를 평가하도록 하는 것이 유익함. 그리고 "Yes"라는 토큰의 다음 단어 확률을 보상으로 사용함. 이 보상은 추론의 올바름에 대한 LLM의 자체 추정을 평가함

Task-specific heurisitics

보상 함수에 과제에 맞는 특화 휴리스틱을 유연하게 추가할 수 있도록 함. 예를 들어 Blockworld에서 계획을 생성할 때, 블록의 예측된 현재 상태를 목표와 비교하여 보상을 계산함. 이 보상은 목표를 향해 움직임.

⭐️ Planning with Monte Carlo Tree Search

World model과 Reward를 갖추면, LLM은 계획 알고리즘과 함께 추론을 진행할 수 있음. 이 논문은 MCTS를 채택하였고 이는 추론 트리의 공간을 전략적으로 탐색하고, 탐색과 활용의 균형을 적절하게 맞추어 고보상 추론 경로를 효율적으로 찾는 강력한 계획 알고리즘임.

각 노드는 상태를 나타내고, 각 엣지는 행동과 현재 상태에서 다음 상태로의 전이를 나타냄. LLM 에이전트가 트리에서 가장 유망한 노드를 확장하고 탐색하도록 유도하기 위해, 알고리즘은 상태-행동 가치 함수 $Q : S \times A \rightarrow R$ 를 유지함. $Q(s,a)$ 는 상태 $s$ 에서 행동 $a$ 를 취했을 때 예상되는 미래 보상을 추정함.

위 이미지는 각 반복에서 트리를 확장하고 $Q$ 값을 갱신하는 네 가지 연산을 보여줌. 이 과정은 지정된 계산 예산에 도달할 때까지 계속되며, 결과적으로 추론 경로는 트리에서 획득됨

Selection

첫 번째 단계는 다음 단계에서 확장하기에 가장 유망한 기존 트리의 일부를 선택하는 것. 루트 노드에서 시작하여, 다음 노드로 자식 노드를 선택함. 이 단계는 리프 노드에 도달하면 끝이 난다.

탐색과 활용 사이의 균형을 맞추기 위해,UCT(Upper Confidence bounds applied to Trees)를 사용하여 각 자식 노드를 선택함. 구체적으로, 노드 $s$ 에서, 우리는 $Q$ 값(활용)과 불확실성(탐색)을 고려하여 트리 내의 행동을 선택함.

a^*= arg \max_{a \in A(s)}Q(s,a) + w \sqrt{\frac{ln N(s)}{N(c(s,a))}}

$N(s)$ : 노드 $s$ 가 이전 반복에서 방문된 회수
$c(s,a)$ : 상태 $s$ 에서 행동 $a$ 를 적용했을 때의 자식 노드
$w$ : 탐색과 활용 간의 균형을 조절

자식 노드가 이전에 덜 방문되었을수록(해당 자식 노드의 불확실성이 클수록), 식의 두 번째 항은 더 크게 됨 ➡️ 탐색을 유도하는 역할

Expansion

이 단계에서는 Selection에서 선택된 리프 노드에 새로운 자식 노드를 추가하여 트리를 확장함. 리프 노드의 상태가 주어졌을 때, LLM을 사용하여 $d$ 개의 가능한 행동을 샘플링하고, LLM을 사용하여 각각의 다음 상태를 예측함. 그 결과 $d$ 개의 자식 노드가 생성됨

⚠️ Selection에서 선택된 리프 노드가 terminal 노드라면, 확장 단계는 건너뛰고 back-propagation 단계로 이동함.

Simulation

예상되는 미래 보상( $Q$ values)을 추정하기 위해, 아 단계에서는 world model을 사용해 현재 노드의 미래 상황을 시뮬레이션함. 현재 노드에서 시작하여, 각 노드 $s$ 에서 roll-out policy에 따라 행동을 생성하고, world model을 사용하여 다음 상태를 예측함. roll-out 과정은 terminal 상태에 도달할 때까지 계속됨.

roll-out policy을 정의하는 방법에는 여러 가지가 있을 수 있음. 실험에서 단순화와 잡음을 줄이기 위해 Expansion과 유사한 과정을 따름 . 즉, $d$ 개의 후보 행동을 생성한 뒤, 가장 큰 보상을 가진 행동을 선택함

a' = arg \max_{a'}r(s,a)

Back-propagation

이전 단계들에서 terminal state에 도달하면, 루트 노드에서 terminal 노드까지의 추론 경로를 얻게됨. 이제 경로상의 각 상태-행동 쌍에 대해 보상을 역전파하여 $Q$ 값을 업데이트함. 상태 $s$ 에서 모든 미래 단계의 보상을 집계하여 $Q(s,a)$ 를 업데이트함.

사전에 정해진 MCTS 반복 횟수에 도달하면 알고리즘을 종료하고, 구성된 트리에서 최종 추론 경로를 선택하여 평가함. 선택 방법에는 여러 가지가 있음
1. 루트 노드에서 시작하여 $Q$ 값이 가장 높은 행동을 반복적으로 선택하다가 terminal state에 도달하는 방법
2. 가장 높은 보상을 얻은 반복에서 나온 경로를 직접 선택
3. 가장 많이 방문된 리프 노드를 선택
⭐️ 실제로는 2번째 방법이 최상의 결과를 도출하는 경우가 많음

RAP Aggregation

수학적 추론은 여러 MCTS 반복에서 생성된 추론 경로와 답을 집계하여 최종 답을 생성함. 하지만 계획 생성이나 논리적 추론에서는 추론 경로가 필요하므로 RAP-Aggregation을 사용하지 않음

Experiments

RAP과 CoT, Least-to-most prompting, Self-consistency, GPT-4와 비교
기본 모델 : LLaMA-33B, Sampling temperature : 0.8

Plan Generation

주어진 목표를 달성하기 위해 일련의 행동을 생성하는 것을 목표로함

Task setup

Blockworld 벤치마크에서 RAP을 평가함. 여기서 에이전트는 블록들을 특정 순서로 탑으로 재비치하는 임무를 수행함.

상태 : 블록들의 현재 배치로 정의
행동 : 블록을 이동시키는 지시 (4개의 동사 STACK, UNSTACK, PUT, PICKUP)
행동 공간 : 행동에 대한 도메인 제약과 블록들의 현재 배치를 고려하여 현재 가능한 행동들을 생성함.
상태 전이 : 현재 행동을 취하고 LLM을 사용하여 관련 블록들의 상태 변화 예측을 요청하는 방식으로 이루어짐.

새로운 블록 조건을 추가하고 더 이상 유효하지 않은 조건을 제거하여 현재 상태를 갱신함. 하나의 상태가 목표 조건을 모두 만족하거나 트리의 깊이 제한에 도달하면 해당 노드의 탐색을 종료함.

이 도메인에서의 행동 품질을 평가하기 위해, 우리는 두 가지 별도의 보상을 사용함.

LLM에 몇 가지 예시 테스트 케이스와 그 해답을 제공한 후, 현재 상태에서 특정 행동의 로그 확률을 계산함. 이 보상은 LLM의 직관을 반영하며, 목표에 몇 단계만 남았을 때는 유의미하지만, 멀리 있는 목표에서는 신뢰하기 어려움.

행동을 수행한 후 새로운 상태를 목표와 비교하여 보상을 제공하며, 이 보상은 만족된 조건의 수와 비례함. 구체적으로, 모든 조건이 만족되었을 경우에는 매우 큰 보상을 부여하여 해당 계획이 최종 해답으로 선택되도록 함

Results

Blockworld 데이터셋의 테스트 케이스를 사용하고, 이를 필요한 최소 행동 수에 따라 그룹화함. 그 결과, 2단계 이내에 해결 가능한 30개, 4단계 이내에 해결 가능한 57개, 6단계 이내에 해결 가능한 114개의 케이스가 구성됨. 각 테스트 케이스에는 최대 5개의 블록이 포함됨.

베이스라인 방법으로는 LLM에 해답이 함께 제공된 4개의 테스트 케이스를 프롬프트로 제시하고, 새로운 질문에 대해 계획을 생성하도록 요청함(CoT). RAP의 경우, 동일한 프롬프트를 사용하여 LLM이 $r_1$ 을 계산하도록 함

LLaMA-33B와 CoT는 2단계 문제에 대해서만 몇 가지 성공적인 계획을 생성할 수 있으며, 더 어려운 문제에서는 실패하였음.

RAP은 모든 문제를 4단계 이내에 해결하며, 일부 6단계 문제도 해결할 수 있어 평균 64%의 성공률을 기록함.

Case Study

CoT와 RAp의 추론 경로를 비교하며, 개선된 이유를 아래와 같이 요약할 수 있음
1. 추론하는 동안 world state를 유지함으로써 RAP은 현재 상태에 대한 유효한 행동을 인식하고, 불법적인 계획 생성을 피할 수 있음
2. RAP은 첫 번째 직관에서 실패했을 때 다른 해결책을 시도하고 백트래킹을 통해 새로운 경로를 탐색할 수 있음.
3. 자기 평가 보상을 계산할 때, LLM에 현재 상태만 제공하고 이전 히스토리는 숨김. 예를 들어 위 이미지에서 $a_2$ 의 보상을 계산하기 위해 LLM에 새로운 테스트 케이스를 제공하며, 여기서 $s_2$ 는 초기 상태임.

Math Resoning

Task setup

최종 질문에 답하기 위해서 문제의 맥락을 바탕으로 다단계 수학 계산을 수행해야함. 그러므로 최종 질문을 더 작은 하위 질문들로 분해하는 것이 자연스러움

상태 : 중간 변수들의 값
행동 : 알려지지 않은 중간 변수에 대한 하위 질문을 제시하는 것

세계 모델은 중간 변수들과 문제 설명을 사용하여 하위 질문에 답하고, 새로운 중간 변수 값을 다음 상태에 추가함

LLM의 유용성에 대한 자기 평가 $r_{t,1}$ 과 상태의 신뢰도 $r_{t,2}$ 를 결합하여 가중 기하 평균 $r_t = r^{\alpha}_{t,1} \times r^{1 - \alpha}_{t,2}$ 를 보상 함수로 사용함. 이 보상은 더 관련 있고 유용한 하위 질문을 유도함. 추론 경로의 길이가 보상에 미치는 영향을 고려하기 위해, 미래 단계들의 평균 보상 중 최대값을 사용하여 $Q$ 값을 계산함

Q^*(s_t,a_t) = \max_{s_t,a_t,r_t,...,s_l,a_l,r_l,s_{l+1}}avg(r_t,...,r_l)

Results

GSM8k를 사용하여 프레임워크를 평가하였음. CoT, Least-to-Most prompt, Self-consistency 변형들을 베이스라인으로 평가하였음. 모든 방법에 대해 4-shot 예시를 사용함

RAP 프레임워크는 48.8%의 문제를 정확하게 풀었으며, CoT, Least-to-Most Prompt + Self consistency를 모두 능가했음.

또한, MCTS 반복 회수와 Self-consistency 샘플의 개수를 달리하여 정확도를 계산하였음

모든 반복 횟수 / 샘플 수에 대해 RAP-Aggregation이 베이스라인을 지속적으로 능가함을 발견하였음. 이는 반복 회수나 샘플 수가 적을 때, RAP 프레임워크가 보상에 의해 유도된 신뢰할 수 있는 추론 경로를 찾는 데 훨씬 더 우수함을 나타냄

Logical Reasoning

Task setup

사실 집합과 논리적 규칙을 제공하며, 모델은 주어진 사실에 논리적 규칙을 적용하여 가설 사실이 참인지 거짓인지를 검증해야함. 이러한 과제는 최종 답(참/거짓)만 요구하는 것이 아니라, 결과를 증명하는 상세한 증거도 요구함

상태 : 집중하는 사실
행동 : 사실 집합에서 규칙을 선택하는 것

world model은 한 번의 추론 단계를 수행하여 새로운 사실을 다음 상태로 생성함
보상은 self-evaluation을 통해 계산됨
구체적으로, LLM에 몇 가지 예시와 해당 라벨을 제공하여 추론 단계의 품질을 더 잘 이해할 수 있도록 도움. 미래 단계들의 평균 보상을 사용하여 $Q$ 함수를 업데이트함

Results

Analysis

Complex problems

더 강력한 LLM이 더 복잡한 문제를 해결하는 데 도움이 되는지 추가적으로 연구하기 위해, Llama-2 70B를 사용하여 전체 Blockworld 데이터셋에 대한 실험을 수행하였음. 전체 Blockworld는 602개의 테스트 케이스로 구성됨. 각 테스트 케이스에 필요한 최소 행동 수를 기준으로 이들을 그룹화하였고, 두 가지 다른 설정으로 실험을 진행하였음

Easy 설정 : 각 케이스에 필요한 최소 행동 수에 대한 사전 지식을 가정함. 이 정보를 활용하여, 테스트 케이스와 동일한 최소 행동 수를 가지는 예시 케이스들을 사용함. 각 그룹의 케이스에 대해, 10개의 케이스를 무작위로 선택하여 예시 케이스 풀을 만들고, 나머지 케이스들을 테스트 세트로 남김. 추론 중에는 이 풀에서 무작위로 4개의 예시 케이스를 샘플링하여 프롬프트를 형성하는데 사용함
Hard 설정 : 전체 데이터셋에서 10개의 케이스를 무작위로 선택하여 예시 풀을 만들고, 그 후 이 케이스들을 테스트 케이스에서 제외함. 추론 중에는 이 전역 풀에서 4개의 예시 케이스를 무작위로 샘플링하여 테스트 케이스의 최소 행동 수와 관계없이 프롬프트를 만듦

coticoger

이전 포스트

ML-Agents

다음 포스트

Reasoning with Language Model is Planning with World Model