Deliberate Reasoning in Language Models as Structure-Aware
Planning with an Accurate World Model

We formulate the planning task as a Markov Decision Process (MDP) represented as $({\mathcal{S}},{\mathcal{A}},{\mathcal{P}},{\mathcal{R}})$ where:

• •

  State $s_{t}\in{\mathcal{S}}$: Represents the current graph, capturing all known or inferred information along with entailment relations.

• •

  Action $a_{t}\in{\mathcal{A}}$: Corresponds to an expansion step, where new information is derived or inferred from the current state, leading to a state transition.

• •

  Transition probability ${\mathcal{P}}(s_{t+1}|s_{t},a_{t})$: Defines the probability of transitioning to $s_{t+1}$ after taking $a_{t}$ in $s_{t}$.

• •

  Reward ${\mathcal{R}}(s_{t},a_{t})$: Measures the quality of $a_{t}$ given $s_{t}$. While reward functions are typically used to select the best action, our framework instead directly compares different actions using a discriminator.

This MDP formulation establishes a foundation for applying planning-based methods to enhance the sequential decision-making capabilities of LMs. By iteratively updating their parameters, the models progressively learn an optimal policy, thereby improving their reasoning performance.

The key innovation that distinguishes our approach from related work is conceptualizing the multi-step reasoning process as entailment graph construction, which outlines how the premises in $s_{0}$ lead to intermediate conclusions, ultimately validating the final answer in $s_{T}$. Formally, let $\mathcal{G}=(\mathcal{V},\mathcal{E})$ represent the structure, where $\mathcal{V}$ is the set of nodes, with each node $v\in\mathcal{V}$ representing a statement, e.g., evidence, an assumption, or a lemma/rule; $\mathcal{E}$ is the set of directed (hyper)edges, where each (hyper)edge $e=(\mathcal{V}_{\text{src}},\mathcal{V}_{\text{tgt}})\in\mathcal{E}$ represents an entailment relation from a source node set $\mathcal{V}_{\text{src}}\subseteq\mathcal{V}$ (the premises) to a target node set $\mathcal{V}_{\text{tgt}}\subseteq\mathcal{V}$ (the conclusions).

Given $s_{0}$, the world model ${\mathcal{P}}_{\text{wm}}$ first builds the initial graph $\mathcal{G}_{0}$ by extracting key statements and their relations. During the reasoning, ${\mathcal{P}}_{\text{wm}}$ incrementally grows the graph by adding new nodes and edges, ultimately forming $\mathcal{G}_{T}$, which includes the final answer. Additionally, symbolic verification is introduced to ensure the graph quality. For simplicity, let us denote the state (in natural language) with structural information as $(s,\mathcal{G})$. Incorporating this structure provides two main benefits: 1) the policy model can make more informed decisions using the structural information; and 2) the world model can predict more accurate next state.

In this section, we present our framework (SWAP) that enables LMs to systematically construct and utilize an entailment graph for solving a wide range of reasoning tasks. We use ${\mathcal{P}}_{\pi}$ to denote the policy, ${\mathcal{P}}_{\text{wm}}$ to denote the world model, ${\mathcal{P}}_{\text{d}}$ to denote the discriminator, and ${\mathcal{P}}_{\text{c}}$ to denote the controller all based on pre-trained LMs. We consider $Q,A,G,H,s,\mathcal{G},a$ as language sequences, i.e., $Q=(Q[1],\cdots,Q[L])$ where each $Q[l]$ is a token, so that ${\mathcal{P}}(Q)=\prod_{l=1}^{L}{\mathcal{P}}(Q[l]|Q{[1..l-1]})$. We use $(s,\mathcal{G})$ to denote the state (in natural language) with structural information, and $c=(G,H,s,\mathcal{G})$ to denote the context of goal $G$, plan $H$ and $(s,\mathcal{G})$.

For notational convenience, we define the generation process as $\texttt{gen}(\texttt{model},\texttt{input},N)$ where $N$ is the number of generations, the symbolic verification process as $\texttt{symCheck}(\texttt{input})$, and the discrimination process as $\texttt{dis}(\texttt{model},\texttt{input},b)$ where $b$ is the number of preserved candidates. To search potential plans and actions, we simulate the future situations of the current state using the world model. Specifically, we use $\texttt{sim}(c,t)$ to denote the simulation starting from $(s,\mathcal{G})$ up to step $t$ given the goal $G$ and plan $H$ from the context $c$.

Algorithm 1 outlines the workflow. Given a reasoning question $Q$, the world model ${\mathcal{P}}_{\text{wm}}(G,s_{0},\mathcal{G}_{0}|Q)$ first extracts the goal $G$ and the initial state $(s_{0},\mathcal{G}_{0})$. The policy then proposes a set of plans $H$ by sampling $N_{H}$ times from ${\mathcal{P}}_{\pi}(H|G,s_{0},\mathcal{G}_{0})$. The top $b$ candidate plans are selected by the discriminator ${\mathcal{P}}_{\text{d}}$ based on the simulation results $(s_{T},\mathcal{G}_{T})$ under each plan. Given the goal $G$, selected plan $H$ and current state $(s_{t-1},\mathcal{G}_{t-1})$, expansion at step $t$ begins with the policy sampling $N_{a}$ times from ${\mathcal{P}}_{\pi}(a_{t-1}|G,H,{s}_{t-1},\mathcal{G}_{t-1})$ as the next action pool. The discriminator ${\mathcal{P}}_{\text{d}}$ then evaluates and selects the top $b$ candidate contexts ($G$, $H$, $s_{t-1}$, $\mathcal{G}_{t-1}$, $a_{t-1}$) based on simulated states $(s_{t},\mathcal{G}_{t})$.

The accurate state prediction (Algorithm 2) is then performed in parallel for each selected context ($G$, $H$, $s_{t-1}$, $\mathcal{G}_{t-1}$, $a_{t-1}$). Specifically, the world model predicts the next state $(s_{t},\mathcal{G}_{t})$ by sampling $N_{s}$ times from ${\mathcal{P}}_{\text{wm}}(s_{t},\mathcal{G}_{t}|s_{t-1},\mathcal{G}_{t-1},a_{t-1})$. Then the discriminator ${\mathcal{P}}_{\text{d}}$ selects the top candidate state. Based on the selected $(s_{t},\mathcal{G}_{t})$, the controller determines whether to continue reasoning. If the goal is achieved, the controller ${\mathcal{P}}_{\text{c}}(A|G,s_{t},\mathcal{G}_{t})$ generates the final answer $A$, stores $(G,s_{t},\mathcal{G}_{t},A)$ in the completed pool $\mathcal{D}$, and reduces $b$ by 1. Otherwise, $(G,H,s_{t},\mathcal{G}_{t})$ will be added to the context pool $\mathcal{C}$ for the next step. The process continues until the step limit $T$ is reached or $b$ becomes 0. Finally, the top answer $A^{*}$ is selected by the discriminator ${\mathcal{P}}_{\text{d}}$ based on the completed states in $\mathcal{D}$.

We identify two fundamental bottlenecks in planning-based reasoning: generation diversity and discrimination accuracy. Enhancing diversity is crucial for expanding the solution space, increasing the likelihood of discovering globally optimal solutions. To address this, we propose a Diversity-based Modelling (DM) approach (Figure 2). The core idea is to encourage the policy to generate multiple candidate options for the same step, ensuring that each option differs from the previous ones, thereby mitigating self-bias.

Given the current state $(s_{t},\mathcal{G}_{t})$, we define ${\mathcal{P}}^{\text{ori}}_{\pi}\left(a_{t}|G,H,s_{t},\mathcal{G}_{t}\right)$ as the original distribution learned from positive trajectories during training. For $n$-th generation, we aim to introduce diversity by considering an additional distribution ${\mathcal{P}}^{\text{sem}}_{\pi}(a^{n}_{t}|a_{t}^{1..n-1})$, which captures steps that are semantically equivalent to those generated previously $a_{t}^{1..n-1}$. Specifically, the probability of $l$-th token $a^{n}_{t,l}$ in the $n$-th generation $a_{t}^{n}$ is

where $a_{t}^{j}$ denotes the $j$-th generation, and for notational simplicity, we move the token index to a subscript, so that $a^{n}_{t,1..l-1}$ denotes the preceding tokens of the $l$-th token $a^{n}_{t,l}$. The distribution ${\mathcal{P}}^{\text{sem}}_{\pi}(a^{\prime}_{l}|a,a^{\prime}_{1..l-1})$ is learned from training data generated by GPT-4o, where each pair $(a,a^{\prime})$ consists of semantically equivalent actions.

The final distribution is given by:

where the decay factor, $\gamma_{l}=\gamma_{0}\cdot\alpha^{l}$ with $\alpha\leq 1$, is introduced to emphasize diversity in early stages of generation while gradually reducing this effect. This ensures that the deduplication effect is stronger initially to explore different paths but weakens over time to avoid drifting too far from plausible solutions, thereby maintaining accuracy.

The normalization function

is applied to discard negative-valued tokens (that either resemble previous responses or deviate from the intended progression of reasoning) and ensure a diverse and relevant response. The selection of this function is further discussed in Appendix A.

As highlighted in recent works (Huang et al., 2023; Chen et al., 2024b), discrimination accuracy is a critical aspect of planning-based methods. However, training an effective process reward model (PRM) is challenging, as it reduces each candidate option to a single numerical score, oversimplifying complex decision-making aspects. Moreover, those raw scores may not be inherently comparable without proper calibration and additional context.

To address this issue, our discriminator employs Contrastive Ranking (CR) to evaluate multiple candidate options within prompts. By directly comparing these options, the discriminator can distinguish nuanced differences, particularly identifying erroneous components, thereby simplifying the evaluation process. To illustrate the automatic annotation process (Figure 3), consider a positive trajectory $\left[(s_{0},\mathcal{G}_{0}),a_{0},\cdots,(s_{T},\mathcal{G}_{T})\right]$ that leads to the correct final answer. We randomly select an intermediate step $t$, and generate $K$ alternative reasoning trajectories: $\{[a^{j}_{t},\cdots,(s^{j}_{T_{j}},\mathcal{G}^{j}_{T_{j}})]\}_{j=1}^{K}$, where $T_{j}$ represents the length of the $j$-th trajectory. Among these $K$ trajectories, we first filter out invalid trajectories through symbolic verification. Then the first erroneous steps in negative trajectories are identified by semantically comparing with the corresponding steps in the positive trajectory. We further perform $N_{\text{veri}}$ completions for outcome verification, i.e., if none of the completions result in the correct answer, we confirm the identified steps as erroneous.

Given the process annotations, we define the inputs and outputs of the discriminator while incorporating meta knowledge ${\mathcal{K}}_{\text{meta}}$, such as common pitfalls and errors, to improve ranking quality. Specifically,

where $E$ denotes an explanation, highlighting differences among the $K$ future states before making a decision. The superscript ’best’ indicates the final selected option, i.e., $a_{t}^{\text{best}}$ and $(s_{t+1}^{\text{best}},\mathcal{G}_{t+1}^{\text{best}})$. For action ranking, we use the simulated immediate next state $(s_{t+1},\mathcal{G}_{t+1})$, as $a_{t}$ only needs to align with the plan $H$. For plan ranking, we use the simulated terminal state $(s_{T},\mathcal{G}_{T})$, that is, ${\mathcal{P}}_{\text{d}}(E,H^{\text{best}}\mid{\mathcal{K}}_{\text{meta}},G,s_{0},\mathcal{G}_{0},\{H^{j},s^{j}_{T_{j}},\mathcal{G}^{j}_{T_{j}}\}_{j=1}^{K})$. For training data, we select a pair of correct and incorrect options from the positive and negative trajectories of the same question. We then fine-tune the discriminator using this data, along with meta-knowledge and explanations bootstrapped from GPT-4o. Further discussions on the methodology and implementation details can be found in Appendix A, D.

During inference, given a set of candidate options, we apply points-based ranking to select the top $b$ candidates denoted as dis(model, input, b) (in Algorithm 1, 2). To manage computational complexity, we randomly sample comparison groups, ensuring that the total number of comparisons remains linear with respect to the number of candidates. To further enhance robustness, we randomly reorder group members before evaluation. Additionally, we perform symbolic verification prior to invoking the discriminator. Further details on these strategies are provided in Appendix A.

In this section, we demonstrate the versatility and effectiveness of our framework by applying it to a diverse set of reasoning tasks. Dataset statistics and examples are provided in Appendix B. The key parameter settings in SWAP are as follows: DM: $\gamma_{0}=0.7$, $\alpha=0.95$. CR: $N_{\text{veri}}=3$; training group size: $K=2$; inference group size: less than $3$. To enhance training effectiveness, we employ advanced training strategies such as curriculum learning and self-improving training for supervised fine-tuning, followed by preference learning for further optimization. We evaluate SWAP against popular reasoning strategies, including: CoT and Self-consistency (Wang et al., 2023b), ToT (Yao et al., 2023) and RAP (Hao et al., 2023), Supervised fine-tuning (SFT) on CoTs and verification with PRMs (Lightman et al., 2023; Wang et al., 2023a). We conduct experiments using two different base models: LLaMA3-8B-Instruct (Dubey et al., 2024) and Mistral-7B-Instruct (Jiang et al., 2023). Implementation settings are as follows: Self-consistency and PRMs use 32 rollouts. For SWAP, ToT, and RAP: the generation limit is 3, with a breadth limit of 3.The total number of rollouts is 32. The step limit is 10 for MATH500, and 6 for all other datasets. Temperature for these methods is set as 0.7. More implementation details are provided in Appendix C, D.

The overall performance is shown in Table 1, with fine-grained results and example analyses provided in Appendix E and F, respectively. The key findings are summarized as follows:

We examine the impact of total rollouts and breadth limit on overall accuracy (Figure 4), providing insights into optimal parameter selection and demonstrating inference-time scaling.

We analyze the impact of the key components introduced in this paper (Table 2). The complete framework achieves the highest performance across all tasks, demonstrating that each component positively contributes to overall accuracy. Notably, planning has the most significant impact by effectively selecting optimal actions. These improvements are consistently observed across different benchmarks, highlighting the robustness of SWAP.

We evaluate the average number of tokens generated by different methods on GSM8K dataset using Llama-3-8B-Instruct. The results are summarized in Table 3. We observed that while the theoretical time complexity of SWAP is comparable to ToT (BFS with pruning), it generates more tokens in practice due to the incorporation of a world model and the construction of entailment graphs. On the other hand, SWAP is significantly more efficient than RAP, which requires simulations until terminal states to estimate the expected future rewards.