Title: Reasoning with Language Model is Planning with World Model

Institute: UC San Diego, University of Florida

Authors: Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, Zhiting Hu

Link: https://arxiv.org/pdf/2305.14992.pdf

Shared Link: https://mp.weixin.qq.com/s/ZLktf__PCPsMYdX-Ikm_5A

Overview

Prior research has shown that humans possess an internal world model that lets them simulate actions and their effects on the state of the world, enabling deliberate planning for complex tasks such as motor control, imagination, reasoning, and decision making. LLMs, by contrast, can only reason autoregressively. The authors identify three main factors that limit LLM reasoning:

• The lack of an internal model that simulates the current state

• The lack of a reward mechanism to guide the model’s reasoning

• The inability to balance exploration of the future against exploitation of past experience

For the first problem, they borrow the idea of a world model from reinforcement learning and use the LLM itself to simulate the current state.

For the second problem, they design four LLM-based reward schemes.

For the third problem, they turn to Monte Carlo tree search.

World model

The world model here is the notion from reinforcement learning, covering the external environment, the simulation of actions, the outcomes of behavior, and so on. The authors implement it with prompting and an LLM.

Let the current state be $s_{t,t=0,1,2,3...T}$. The process in which the LLM, acting as an agent, infers action $a_t$ from $s_t$ and prompt $ c$ can be written as $a_t \sim p(a|s_t,c)$ . The LLM, now acting as the world model, then predicts the next state $s_{t+1}$ from $ a_t$ and $s_t$ .

Reward

Every reasoning step has to be assessed for feasibility so that the assessment can steer the reasoning. The authors design an evaluation function $r_t=r(s_t,a_t)\in\mathbb{R}$ that returns a reward for the state $s_t$ and action $a_t$ at time $t$.

Reasoning via Monte Carlo tree search

Monte Carlo tree search consists of four main steps:

• Selection: pick the leaf node with the highest expected return according to its UCB value

• Expansion: expand that leaf node

• Simulation: simulate from the expanded leaf node and judge the outcome

• Back-propagation: propagate the result back up the tree

The Selection and Simulation stages are worth a closer look.

Selection

Each candidate node gets a UCB value, computed as follows:

$a^*=\arg\max\limits_{a\in A(s)}\left[Q(s,a)+w\sqrt{\frac{\ln N(s)}{N(c(s,a))}}\right]$

• The first term, the Q value, is the node’s average reward over past visits; it represents exploitation of past experience.

• In the second term, $N$ is the visit count and $c(s,a)$ is the child of this node (the result of applying action $a$ to state $s$); this term represents exploration of the future.

• $w$ is the weight of the second term.

The formula gives the Selection policy two tendencies:

• It favors nodes that have yielded higher rewards in the past

• It favors nodes that have not yet been explored

The UCB formula in Monte Carlo tree search therefore resolves the exploration vs. exploitation dilemma.

Simulation

This stage uses the world model to simulate the future of the current node. Starting from that node, at each node $s_t$ we create an action $a_t$ following an expansion policy and use the world model to predict the next state $a_{t+1}$. Expansion continues until it reaches a terminal state.

See panel $(c)$ of the figure above.