Q*: Improving Multi-step Reasoning for LLMs with Deliberative Planning

## Abstract

Abstract Large Language Models (LLMs) have demonstrated impressive capability in many natural language tasks. However, the auto-regressive generation process makes LLMs prone to produce errors, hallucinations and inconsistent statements when performing multi-step reasoning.
In this paper, by casting multi-step reasoning of LLMs as a heuristic search problem, we aim to alleviate the pathology by introducing Q*, a general, versatile and agile framework for guiding LLMs decoding process with deliberative planning. By learning a plug-and-play Q-value model as heuristic function for estimating expected future rewards, our Q* can effectively guide LLMs to select the most promising next reasoning step without fine-tuning LLMs for the current task, which avoids the significant computational overhead and potential risk of performance degeneration on other tasks. Extensive experiments on GSM8K, MATH and MBPP demonstrate the superiority of our method, contributing to improving the reasoning performance of existing open-source LLMs.

## 1 Introduction

Large Language Models (LLMs) have exhibited impressive capabilities in solving various reasoning tasks encoded in natural languages, including math word problems [, , , , , ], code generation [, , ] and planning [, , ]. Unfortunately, even the most advanced LLMs still face significant challenges and are prone to introduce errors, hallucinations and inconsistent statements as the number of reasoning steps grows due to their auto-regressive nature [, ]. In fact, the auto-regressive generation process of LLMs can be characterized by “System 1" [], a mode of thought which is fast, instinctive but less accurate. Most of recent works focus on improving LLMs’ “System 1” capability by (1) constructing sophisticated prompts with extensive expertise to trigger the potential capacities of LLMs without modifying their parameters [, , , ], (2) fine-tuning LLMs with massive task-specific corpus at the price of significant computational burdens and the potential risk of performance degeneration on other tasks [, , , ], or (3) training reward models to rank the candidate responses [, , , ].

On the other hand, solving complex reasoning problems requires more in-depth, deliberative and logical thinking steps, *i.e.*, the “System 2" mode []. Taking solving math word problems as an example, any incorrect intermediate reasoning step (*e.g.*, calculation errors, mis-interpretations) can potentially lead to incorrect final answers. Prior attempts [, , , ] for enhancing “System 2” reasoning capability includes performing deliberation with basic tree search algorithms (*e.g.*, BFS or DFS), Monte Carlo Tree Search (MCTS) [], and A* []. Nonetheless, the utility functions used in these methods often require laborious expertise to design for each specific task, which are difficult to be extended to new scenarios. Furthermore, deliberation with MCTS would require significant number of rollouts before finding high-quality responses when solving the problems with many reasoning steps, which substantially slows down the overall decoding process.

In light of this, we propose Q*, a general, versatile and agile framework for improving the multi-step reasoning capability of LLMs with deliberative planning. Different from the existing deliberation methods, our method does not rely on domain knowledge to design the heuristic function. Besides, by leveraging plug-and-play Q-value models as heuristic function, our Q* can effectively solve various tasks via guiding LLMs to select the most promising next step without fine-tuning LLMs beforehand, which avoids the significant computational overhead and potential risk of performance degeneration in other tasks. Finally, Q* considers only one single step when performing deliberation, which is much cheaper than completing rollouts in MCTS. Specifically, the main contributions of our work are summarized as follows:

- •
We formalize the multi-step reasoning of LLMs as a Markov Decision Process (MDP) where the state is the concatenation of input prompt and the reasoning steps generated so far, the action is the next reasoning step and the reward measures how well the task is solved.
- •
We present several general approaches to estimate the optimal Q-value of state-action pairs, *i.e.*, offline reinforcement learning, the best sequence from rollouts, and completion with stronger LLMs. It is noteworthy that our methods only need the ground-truth of training problems and can be easily applied to various reasoning tasks without modification.
- •
We cast solving multi-step reasoning tasks as a heuristic search problem, where the objective is to find the most proper reasoning trace with maximum utility. Built upon A* search, our deliberation framework, Q*, leverages plug-and-play Q-value models as heuristic function and guides LLMs to select the most promising next reasoning step in best-first fashion.
- •
We conduct extensive experiments on math reasoning and code generation tasks, demonstrating that Q* can significantly improve the multi-step reasoning capability of existing open-sourced LLMs.

## 2 Related Works

#### LLM alignment.

Alignment has become an important technique to prevent the output of LLMs deviates from human’s expectation. Supervised Fine-Tuning (SFT) is probably the most fundamental alignment approach that directly minimizes the cross-entropy loss between the output and ground-truth. Reinforcement learning from Human Feedback
(RLHF) [], on the other hand, firstly learns a reward model (RM) from human preferences and then optimizes the SFT model with reinforcement learning algorithms to maximize the cumulative rewards from RM. Direct Preference Optimization (DPO) [] aligns LLMs directly according to the ranking information from human feedback without explicitly learning RM. Recently, Aligner [] came out as a model-agnostic alignment method by learning to re-write LLMs’ output. Compared to these methods, our Q* achieves the goal of alignment with distinct merits. Different from SFT and Aligner, Q* does not rely on massive human annotated preference pairs which are expensive to collect; different from RLHF and DPO, our Q* does not modify the parameters of LLMs, which avoids the potential risk of performance degeneration on other tasks.

#### Enhancing LLMs with planning.

Tree-of-thoughts (ToT) [] improves the LLMs’ reasoning capability by exploring the intermediate steps towards problem solving with basic tree-search algorithms. In the same vein, A* search and MCTS are applied to serve as a planning technique to enhance the performance of LLMs when solving challenging complex reasoning problems [, , , ]. Unfortunately, the utility function used in these methods is either constructed from LLMs’ feedback (*e.g.*, [, ]), which could be highly-inaccurate in complex problems, or specific to each individual task (*e.g.*, [, ]). Moreover, planning with MCTS often requires to perform costly rollout, which can significantly slow down the overall decoding process. In contrast, our Q* solely relies on training a Q-value model to guide LLMs to select the most promising next reasoning step and the pipeline can be easily applied to various reasoning tasks without modification. Besides, we consider only a single step each time in Q*, which is much cheaper than complete rollout in MCTS-based methods.

#### LLMs for math reasoning & code generation.

Math reasoning and code generation require LLMs to perform multi-step reasoning on relations, quantities and logics which are inherently challenging. Current techniques include: 1) prompt engineering which triggers the potential capacities of LLMs with sophisticated prompts [, , , , , ]. However, constructing such prompt needs extensive expertise and case-by-case tuning, which is difficult to generalize to different tasks; 2) Fine-tuning LLMs with massive math/code corpus [, , , , , , ], which usually comes at the price of significant computational burden and may compromise the performance on other tasks; 3) training RMs/verifiers to rank the candidate solutions without providing any guidance in intermediate steps [, , , ]. Differently, our Q* leverages a plug-and-play Q-value model to direct the deliberation process of LLMs, which effectively provides guidance for each intermediate step without modifying the parameters of LLMs. Moreover, by casting multi-step reasoning of LLMs as a heuristic search problem, our method can be generalized to various reasoning tasks without laborious prompt engineering.

## 3 Preliminary

### 3.1 Formulate the Multi-step Reasoning of LLMs as an MDP

Taking the question $q$ as input, the answer generation process of LLMs can be broken down into multiple reasoning steps, where the final answer sequence $\mathbf{a}$ can be treated as the concatenation of these $T$ single-step reasoning steps, formulated as $\mathbf{a}=[a_{1};a_{2};\dots;a_{T}]$. Each step can be a single line or fixed number of tokens outputted by LLMs.
Under this perspective, we can conceptualize the multi-step reasoning process of LLMs as a Markov Decision Process (MDP) $\langle\mathcal{S},\mathcal{A},\mathcal{T},\mathcal{R},\gamma\rangle$, where the state $s_{t}\in\mathcal{S}$ denotes the concatenation of the input question and the partial reasoning trace already generated by timestep $t-1$ (*i.e.*, $s_{t}=[q;a_{1};\dots;a_{t-1}]$) with the special case $s_{1}=q$, the action $a_{t}\in\mathcal{A}$ denotes the next reasoning step generated by LLMs taking the current state $s_{t}$ as input, the deterministic state transition $\mathcal{T}$ from the current state $s_{t}$ to the next state $s_{t+1}$ is accomplished through a simple operation of
concatenation, $\mathcal{R}$ is the reward function to measure how well the question is solved and $\gamma$ is the discount factor. The reward function is often outcome-based. That is, it gives reward by comparing the final results with ground-truth:

$$ $\mathcal{R}(s_{t},a_{t})=\begin{cases}1&t=T\land[s_{t};a_{t}]\;\text{matches the ground-truth}\\ 0&\text{otherwise}\end{cases},$ (1) $$

In particular, we assign a reward of 1 if the generated code passes all test cases (for code generation tasks) or the final answer matches the ground-truth (for math reasoning tasks) which is a common practise in previous studies [, ].
Finally, the policy $\pi_{\theta}$ is embodied by an LLM, which produces reasoning sequence conditioned on the input question:

$$ $\displaystyle\pi_{\theta}(a_{t}|s_{t})=\text{LLM}(a_{t}|s_{t}),~{}\pi_{\theta} (\mathbf{a}|q)=\prod_{t=1}^{T}\pi_{\theta}(a_{t}|s_{t}).$ (2) $$

Given the MDP and LLM policy $\pi_{\theta}$, the value of state-action pair $(s_{t},a_{t})$ is given by a Q-function $Q^{\pi_{\theta}}(s_{t},a_{t})=\mathbb{E}_{\pi_{\theta}}\left[\sum_{t^{\prime}=
t}^{T}\gamma^{T-t^{\prime}}\mathcal{R}(s_{t^{\prime}},a_{t^{\prime}})\right]$. The Q-function of an optimal policy $\pi^{*}$ is called optimal Q-function and satisfies the Bellman optimality equation:

$$ $Q^{*}(s_{t},a_{t})=\mathcal{R}(s_{t},a_{t})+\gamma\max_{a_{t+1}\in\mathcal{A}} Q^{*}(s_{t+1},a_{t+1}),$ (3) $$

which gives the value of starting state $s_{t}$, taking action $a_{t}$ and then following the optimal policy $\pi^{*}$.

### 3.2 A* Search

A* [] is an important heuristic search algorithm in deliberative planning [], multi-agent pathfinding [], and constraint reasoning []. Originally, A* is proposed for finding the shortest path from source $s$ to goal $g$ in path planning problems. It associates each frontier vertex $n$ with a value $f(n)=g(n)+h(n)$, where $g(n)$ is the accumulated path cost from source $s$ and $h(n)$ is a heuristic value that estimates the cost of the shortest path
from $n$ to goal $g$. The algorithm adopts a best-first search strategy, *i.e.*, in each iteration it always picks the vertex with minimum $f$-value to explore until reaching the goal. When the heuristic $h(\cdot)$ is admissible [], A* guarantees to find the optimal path.

Figure: Figure 1: Overview of Q*. (a): the deliberation process of Q*. Each state is associated with an $f$-value which is the weighted sum of the aggregated utility (cf. Eq. (5)) and the heuristic value (cf. Eq. (6)). (b-d): estimating optimal Q-value with fitted-Q-iteration, rollout and completion with stronger LLMs.
Refer to caption: extracted/5747045/fig/fig112.png

## 4 Q*: A General, Versatile and Agile Deliberation Framework for LLMs

Most of modern LLMs generate natural languages in an auto-regressive way, *i.e.*, predict the next token in a sequence given the previously generated tokens (cf. Eq. (2)). Therefore, when applied to multi-step reasoning problem, LLMs can potentially introduce errors, hallucinations and inconsistent statements in the subsequent reasoning trace if any previous step is incorrect, which may fail to solve the current problem. Indeed, given the fact that LLMs produce each token with limited computation resources, there is no way to devote more computational efforts to solve difficult problems. In short, LLMs cannot perform in-depth deliberation which is essential for solving complex multi-step reasoning problems.

We address this issue by presenting Q*, a general, versatile and agile deliberation framework based on A* to effectively guide LLMs to select the most promising next step when performing multi-step reasoning without costly fine-tuning LLMs for each task beforehand. In more detail, we cast finding the most proper reasoning sequence for a given problem as a heuristic search process, where each state $s_{t}$ is associated with a $f$-value estimating how much utility will be attained if we expand $s_{t}$:

$$ $\displaystyle f(s_{t})=g(s_{t})+\lambda h(s_{t}),$ (4) $$

where $g(s_{t})$ denotes the aggregated utility from the initial state $s_{1}$; $h(s_{t})$ is the heuristic value for measuring the probability of reaching the correct answer derived from $s_{t}$; $\lambda$ is a coefficient to balance the importance of $g(s_{t})$ and $h(s_{t})$ terms.

Specifically, we propose to use process-based reward function $\mathcal{R}_{P}$ that encodes the prior knowledge or preference of the reasoning task to compute the aggregated utility $g(s_{t})$. That is,

$$ $g(s_{t})=\text{Agg}(\mathcal{R}_{P}(s_{1}),\dots,\mathcal{R}_{P}(s_{i}),\dots, \mathcal{R}_{P}(s_{t})),$ (5) $$

where Agg $\in\{\min,\max,\sum,[-1]\}$, with
$[-1]$ standing for assigning the reward of last state as the utility, is the aggregation function to summarize the rewards in the path from $s_{1}$ to $s_{t}$, and $s_{i-1}$ is the prefix of $s_{i},\;1<i\leq t$. Such process-based reward function $\mathcal{R}_{P}$ could be learned from human feedback [, , ], ground-truth [, ], rules, or simply be the logits of a reasoning step which reflects the confidence of the LLM. Furthermore, we use the optimal Q-value of state $s_{t}$ (cf. Eq. (3)) as the heuristic value $h(s_{t})$. In other words, the $f$-value is given by:

$$ $\displaystyle f(s_{t})=g(s_{t})+\lambda\max_{a_{t}\in\mathcal{A}}Q^{*}(s_{t},a _{t}).$ (6) $$

Since enumerating all possible next reasoning steps is intractable, in practice we restrict the alternatives to the top-K of all step candidates returned by LLM, and thus Eq. (6) can be written as $f(s_{t})=g(s_{t})+\lambda\max_{a_{t}\in\text{top-K}(\pi_{\theta}(\cdot|s_{t}))
}Q^{*}(s_{t},a_{t})$.

### 4.1 Estimation of Optimal Q-value

A critical challenge of implementing Q* is to estimate the optimal Q-value of state-action pairs (cf. Eq. (6)) with a frozen LLM policy $\pi_{\theta}$ which could be suboptimal on the given reasoning problems. Therefore, we aim to learn a proxy Q-value model $\hat{Q}$ to approximate $Q^{*}$ from a dataset $D=\{q_{i},\{\mathbf{a}_{i_{j}}\}_{j=1}^{M}\}_{i=1}^{N}$, where $q_{i}$ is a training problem and $\mathbf{a}_{i_{j}}\sim\pi_{\theta}(\cdot|q_{i})$ is a trajectory sampled from the LLM policy $\pi_{\theta}$. Formally:

$$ $\hat{Q}=\arg\min_{Q}\frac{1}{NMT}\sum_{i=1}^{N}\sum_{j=1}^{M}\sum_{a_{t}\in \mathbf{a}_{i_{j}}}\left(Q(s_{t},a_{t})-\hat{y}(s_{t},a_{t})\right)^{2},$ (7) $$

where $s_{t}=[q_{i};a_{1};\dots;a_{t-1}]$ is the partial reasoning trace by timestep $t-1$ in $\mathbf{a}_{i_{j}}$ and $\hat{y}(s_{t},a_{t})$ is the label that approximates the true optimal Q-value, specifically $Q^{*}(s_{t},a_{t})$.

In more detail, we effectively construct Q-value labels $\hat{y}(s_{t},a_{t})$ for question $q_{i}$ in the following ways:

#### Offline reinforcement learning.

Given the offline dataset $D$, we learn the proxy Q-value model $\hat{Q}$ using Fitted Q-iteration []. Specifically, for each iteration $\ell$, we construct Q-value label as:

$$ $\hat{y}_{\ell}(s_{t},a_{t})=\begin{cases}\mathcal{R}(s_{t},a_{t})&t=T\\ \mathcal{R}(s_{t},a_{t})+\gamma\max_{a_{t+1}\in\text{top-K}(\pi_{\theta}(\cdot |s_{t+1}))}\hat{Q}_{\ell-1}(s_{t+1},a_{t+1})&\text{otherwise}\end{cases},$ (8) $$

where $\hat{Q}_{\ell-1}$ is the proxy Q-value model learned in iteration $\ell-1$. After that, we train a new proxy model $\hat{Q}_{\ell}$ according to Eq. (7). Such two phases will be alternated for $L$ iterations, and we use $\hat{Q}_{L}$ as the proxy Q-value model when performing deliberation.

#### Learning from rollout.

Staring from the state-action pair $(s_{t},a_{t})$, we first perform random rollout or MCTS [, ] with policy $\pi_{\theta}$ for enough iterations, forming a trajectory pool $P$. After that, we use the best reasoning sequence with the highest accumulated rewards to construct the Q-value label:

$$ $\hat{y}(s_{t},a_{t})=\mathcal{R}(s_{t},a_{t})+\max_{\begin{subarray}{c}\tau \sim P\\ (s_{t^{\prime}},a_{t^{\prime}})\in\tau\end{subarray}}\left[\sum_{t^{\prime}=t+ 1}^{T}\gamma^{T-t^{\prime}}\mathcal{R}(s_{t^{\prime}},a_{t^{\prime}})\right].$ (9) $$

#### Approximating optimal policy with stronger LLMs.

By exploiting the fact that the optimal Q-values are derived from an optimal policy, one can estimate the optimal Q-value of state-action pair $(s_{t},a_{t})$ by completing the trajectory with another stronger LLM $\pi_{\theta^{*}}$ (*e.g.*, GPT-4) if available:

$$ $\hat{y}(s_{t},a_{t})=\mathcal{R}(s_{t},a_{t})+\sum_{t^{\prime}=t+1}^{T}\gamma^ {T-t^{\prime}}\mathcal{R}(s^{*}_{t^{\prime}},a^{*}_{t^{\prime}}),$ (10) $$

where $s^{*}_{t^{\prime}}=[s_{t};a_{t};a^{*}_{t+1};\dots;a^{*}_{t^{\prime}-1}]$ and $a^{*}_{t^{\prime}}\sim\pi_{\theta^{*}}(\cdot|s^{*}_{t^{\prime}})$.

### 4.2 Deliberative Planning with A*

Once obtaining the proxy Q-value model $\hat{Q}$, we can plug it to Eq. (6) to compute the $f$-value of each state and perform best-first search with A*. Alg. 1 illustrates the deliberative planning process.

Figure: Algorithm 1 Deliberative planning for LLMs with A*

Specifically, we maintain a set for storing state candidates to be explored, denoted as $unvisited$, which initially only contains the input question $q$, and another set $visited$ to record the visited states. Each step we pick the state $s$ with the maximum $f$-value from the set $unvisited$ and expand it by querying the top-K alternatives with the LLM policy $\pi_{\theta}$. After that, both $visited$ and $unvisited$ sets will be updated and this process repeats until the terminal state (a complete trajectory) is reached. Finally, we extract the answer part of the terminal state $s$ as the result.

## 5 Experiments

### 5.1 Experimental Settings

Datasets. We evaluate the effectiveness of Q* on two math reasoning and one code generation tasks, where the dataset statistics have been summarized in Table 1.
1) GSM8K [] is a dataset of grade school math problems, where the solution is given in a one-line-per-step format with an exact numerical answer in the last line;
2) MATH [] is a dataset consisting of math problems of high school math competitions, where the solutions are given in a format that mixes latex code and natural language;
3) MBPP [] is an entry-level Python programming dataset, where the questions are coding challenges along with a test case that defines the function format.
The solutions are Python code that is excepted to pass the pre-collected test cases of each question.

**Table 1: Statistics of datasets.**
| Dataset | GSM8K | MATH | MBPP |
| --- | --- | --- | --- |
| Domain | Math Reasoning | Math Reasoning | Code Generation |
| Training | 5000 | 8000 | 374 |
| Testing | 1319 | 5000 | 500 |
| Average Steps | 4.5 | 11.0 | 7.0 |

Implementation Details. The implementation of Q* method mainly includes three steps:
1) Q-value estimation. As discussed in Section 4.1, we propose several ways for estimating Q values and will illustrate the implementation details in the next paragraph;
2) Utility aggregation for each step (cf. Eq. (5)).
For GSM8K dataset, we adopt a process reward model (PRM) trained on PRM800K [] to model $\mathcal{R}_{P}$ to provide an intermediate signal for each reasoning step, and use $\min$ as the aggregation function;
For MATH dataset, we set $g(s_{t})=0$ for all passed states $\{s_{i}\}_{i=1}^{t}$ in each trajectory for fairness, because PRM800K contains data samples constructed from MATH testing set and there is a potential risk of data leakage;
For MBPP dataset, we tokenize the code generated so far with function tokenize.generate_tokens and give a penalty of -0.5 if TokenError is raised, which is often the case that there are mismatched delimiters (e.g., parentheses, quotation marks) and invalid indention in the code. We use $[-1]$ as the aggregation function to cancel the previous penalties since the code is generated on-the-fly and mismatched delimiters may be fixed in subsequent steps.
3) A* planning. For GSM8K and MATH datasets, we treat a single line outputted by the LLM as an action, while the action in MBPP is defined as a code snippet with 24 tokens when planning. Besides, in all experiments, we set $\lambda=1$ when computing $f$-values and expand a state with $K=6$ actions at each reasoning step. Finally, following the common practice of Best-of-$N$, we perform planning to collect $N=6$ trajectories for each question, and select the one with the maximum $f$-value as the final result for evaluation.

For Q-value estimation, in our practice, we find that learning from rollout could be the most effective and robust way to collect precise Q-value labels.
Specifically, given the prompt of questions, we will firstly perform random rollout to obtain complete trajectories with LLM, denoted as $\pi_{\theta}$, under the setting of high temperature, *e.g.*, $\tau=0.9$ for math reasoning and $\tau=0.2$ for code generation,
and split each trajectory into a series of states according to the newline token “$\textbackslash n$”.
Then, for each state-action pair in a trajectory, denoted as $(s_{t},a_{t})$, we can perform random rollout or MCTS with the same LLM to collect a pool $P$ of trajectories, and then select the best reasoning path with the highest accumulated rewards to construct the corresponding Q-value label of the current state-action pair.
Note that the reward $\mathcal{R}(s_{t},a_{t})$ is given as 1 only if the obtained math numerical answer is correct or the program passes all test cases, indicating that the Q value of a state-action pair can be 1 only if it has the potential to generate a trajectory containing the correct answer.

### 5.2 Experimental Results

GSM8K. For the comparison on GSM8K dataset, we select Llama-2-7b [] as our base model, whose accuracy can achieve 65.2% after finetuning on MetaMath [].
Then, we treat Llama-2-7b finetuned on MetaMath as policy $\pi_{\theta}$, and perform random rollout to collect Q-value labels for training Q-value model (QVM).
For utility aggregation, we train a process reward model (PRM) on PRM800K [] to provide intermediate signal for each reasoning step.
With PRM and QVM in hand, traditional methods tend to treat either of them as a verifier to select the Best-of-$N$ trajectory or utilize them to perform PPO training of RLHF.
As the results shown in Table 2, we can find that with the same PRM/QVM, using it for verification performs significantly better than using it for alignment.
Further, in the comparison of planning-based methods, we can find that with the same QVM, Q* method with constant aggregated utility can still outperform Best-of-$N$ method.
With the PRM trained on PRM800K determining whether the intermediate reasoning steps are correct, Q* method that combines PRM and QVM achieves the best performance among all methods based on the same LLM, helping Llama-2-7b surpass the performance of close-sourced ChatGPT-turbo [] and reaching an accuracy of 80.8%.

MATH. As the results shown in Table 3, considering the weak performance of Llama-2-7b fine-tuned with MetaMath for the MATH dataset, we seek for two other stronger LLMs to evaluate the effectiveness of our Q* method.
One is Llama-2-7b fine-tuned on Synthetic Data [], which is constructed following the instruction of scaling up the SFT data, and achieves 41.9% accuracy on MATH dataset, approaching the performance of GPT-4 [].
The other base model is DeepSeek-Math-7b [], which could be the most powerful open-source 7b model for math reasoning on MATH dataset, achieving 50.8% accuracy in our evaluation.
From the results shown in the second and third blocks of Table 3, we can find that Q* can still lead to further performance improvement compared to the Best-of-$N$ method on either of base models. Additionally, it is noteworthy that the performance of DeepSeek-Math-7b enhanced with Q* has already surpassed a series of closed-source models on the leaderboard of MATH dataset (^1^11[https://paperswithcode.com/sota/math-word-problem-solving-on-math](https://paperswithcode.com/sota/math-word-problem-solving-on-math)), such as Gemini Ultra (4-shot), reaching an accuracy of 55.4% .

**Table 2: Comparison of Q* and other baselines on GSM8K dataset.**
| Base Model | SFT | Alignment | Verification | Accuracy |
| --- | --- | --- | --- | --- |
| GPT-3.5 (5-shot) [] | Unknown | PPO (RM) [] | - | 57.1% |
| ChatGPT-instruct (0-shot) [] | Unknown | PPO (RM) [] | - | 71.3% |
| ChatGPT-turbo (0-shot) [] | Unknown | PPO (RM) [] | - | 77.7% |
| GPT-4 (0-shot) [] | Unknown | PPO (RM) [] | - | 91.9% |
| GPT-4 (5-shot) [] | Unknown | PPO (RM) [] | - | 92.0% |
| Llama-2-7b (0-shot) | - | - | - | 49.5% |
| Llama-2-7b (0-shot) | MetaMath[] | - | - | 65.2% |
| Llama-2-7b (0-shot) | MetaMath[] | PPO (PRM) [] | - | 67.2% |
| Llama-2-7b (0-shot) | MetaMath[] | PPO (QVM) [] | - | 67.6% |
| Llama-2-7b (0-shot) | MetaMath[] | - | Best-of-$N$ (PRM) [] | 72.1% |
| Llama-2-7b (0-shot) | MetaMath[] | - | Best-of-$N$ (QVM) [] | 74.5% |
| Llama-2-7b (0-shot) | MetaMath[] | - | Q* (QVM) | 78.8% |
| Llama-2-7b (0-shot) | MetaMath[] | - | Q* (PRM+QVM) | 80.8% |

**Table 3: Comparison of Q* and other baselines on MATH dataset.**
| Base Model | SFT | Alignment | Verification | Accuracy |
| --- | --- | --- | --- | --- |
| GPT-3.5 (0-shot) [] | Unknown | PPO (RM) [] | - | 23.5% |
| GPT-4 (0-shot) [] | Unknown | PPO (RM) [] | - | 42.5% |
| Gemini Ultra (4-shot) [] | Unknown | PPO (RM) [] | - | 53.2% |
| Llama-2-7b (0-shot) | - | - | - | 7.9% |
| Llama-2-7b (0-shot) | MetaMath[] | - | - | 20.0% |
| Llama-2-7b (0-shot) | Synthetic Data[] | - | - | 41.9% |
| Llama-2-7b (0-shot) | Synthetic Data[] | PPO (QVM) [] | - | 42.5% |
| Llama-2-7b (0-shot) | Synthetic Data[] | - | Best-of-$N$ (QVM) [] | 46.8% |
| Llama-2-7b (0-shot) | Synthetic Data[] | - | Q* (QVM) | 49.1% |
| DeepSeek-Math-7b (0-shot) | Unknown | PPO (QVM) [] | - | 50.8% |
| DeepSeek-Math-7b (0-shot) | Unknown | PPO (QVM) [] | Best-of-$N$ (QVM) [] | 54.3% |
| DeepSeek-Math-7b (0-shot) | Unknown | PPO (QVM) [] | Q* (QVM) [] | 55.4% |

MBPP. As for the comparison on MBPP dataset, we also choose the most powerful open-source LLM in the aspect of code generation, specifically CodeQwen1.5-7b-Chat, as our base model for evaluating the effectiveness of Q*.
Following a similar procedure of math reasoning, we train a QVM for Q-value estimation and manually
construct the utility function as described in the previous part of implementation details (tokenize the code generated so far with function tokenize.generate_tokens and give a penalty of -0.5 if TokenError is raised).
From the results shown in Table 4, we can find that Q* can still outperform Best-of-$N$ method in the aspect of code generation, and help CodeQwen1.5-7b-Chat to achieve 77.0% accuracy on MBPP dataset, which is also a promising performance in the leaderboard of MPBB (^2^22[https://paperswithcode.com/sota/code-generation-on-mbpp](https://paperswithcode.com/sota/code-generation-on-mbpp)).

**Table 4: Comparison of Q* and other baselines on MBPP dataset.**
| Base Model | SFT | Alignment | Verification | Accuracy |
| --- | --- | --- | --- | --- |
| GPT-3.5 Turbo (self-debug) [] | Unknown | PPO (RM) [] | - | 72.8% |
| GPT-4 (self-debug) [] | Unknown | PPO (RM) [] | - | 80.2% |
| CodeQwen1.5-7b-Chat (0-shot) | Unknown | PPO (QVM) [] | - | 74.6% |
| CodeQwen1.5-7b-Chat (0-shot) | Unknown | PPO (QVM) [] | Best-of-$N$ (QVM) [] | 75.0% |
| CodeQwen1.5-7b-Chat (0-shot) | Unknown | PPO (QVM) [] | Q* (PRM+QVM) [] | 77.0% |

## 6 Conclusion

Solving challenging multi-step reasoning problems requires LLMs to perform in-depth deliberation beyond auto-regressive token generation. In this paper, we present Q*, a general, versatile and agile deliberation framework for LLMs. Unlike existing deliberation methods which need extensive expertise to design a utility function for each specific task, Q* relies on ground-truth solely to train value model and can be easily applied to various reasoning tasks without modification. Moreover, by leveraging plug-and-play Q-value models as the heuristic function, Q* can effectively guide LLMs to solve various tasks without fine-tuning LLMs beforehand, which avoids potential performance degeneration on other tasks. Finally, our Q* is agile because we consider only a single step each time rather than complete rollouts (*e.g.*, simulation in MCTS). Extensive empirical evaluations on math reasoning and code generation tasks confirm the superiority of our method.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 3
presentation: 3
questions: - In Algorithm 1, for the first iteration of the while loop, what would s be in line 3? The unvisited set only has the question, so what would the argmax over q result in? Perhaps you meant to initialize the unvisited set as the set of all the states in line 2?

- Similarly in line 3, if Q(s') were used instead of f(s'), how would the solution change? In the traversed path so far, g(s') remains the same for the upcoming paths and therefore doesn't influence the argmax. Thus, I think the solution would remain the same whether you use f(s'), or either of Q(s'), h(s').
soundness: 2
strengths: - The paper is well written with clear figures that explain the algorithm and process.
- The algorithm is conceptually interesting, and the comparison of three different techniques to estimate the Q-value function is valuable.
summary: The paper proposes a framework (called Q*) that guides LLM decoding toward better final solutions using an estimated optimal Q-value function. If the Q-value function is known a priori (somehow), then the method can optimize (search) solutions during test time without the need for fine-tuning. The proposed algorithm appears quite interesting at first glance for reasoning tasks. However, arguably, the formulation presented in the paper does not necessarily require planning (there is no state change/feedback for any intermediate step). Theoretically, the problem can be formulated as a single action prediction a = a_{1:T}.
weaknesses: - In my opinion, the major limitation is that the algorithm is conceptually designed for a different set of problems than those presented in the formulation (section 3.1) and experiments.
- The presented formulation starts with state s (question) and appends actions (autoregressively) until reaching the terminal state. There are no new observations/states or rewards for the intermediate states. The problem can be cast as a single-step process where, given a question, one needs to search over a large action space. I would argue that there is no need for multi-step reasoning.
- The presented algorithm is interesting, and I would see its value in applying it to multi-turn processes where, after applying an action (sequence of tokens), the system provides a new state, and then, conditioned on this new state/observation, one can take the next action.
- I see this work as aligning more with a beam search approach on how to get better output from an LLM when conditioning on the Q-value.
- In your experiments, could you please add uncertainty values (given that LLMs are quite stochastic)? Do the current values represent the mean or the best run?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please check weakness above especially the second point and the third point.
soundness: 2
strengths: 1.Clarity 
The paper is well-structured and written in a clear, accessible way. The authors employ several advanced methods, presented logically and effectively, enhancing the reader's understanding of complex techniques. 

2.Motivation
The motivation for this research is both evident and well-founded. Developing an effective, reliable framework for multi-step reasoning is a significant challenge in language model research. 

3.Originality
The originality of this work is noteworthy. Few studies have investigated integrating heuristic search into LLMs to improve multi-step reasoning. By introducing heuristic search methodologies into LLMs, this paper provides a novel approach with strong experimental results.
summary: This paper casts the multi-step reasoning process of large language models (LLMs) as a heuristic search problem and proposes Q*, a framework based on the heuristic A* algorithm that incorporates Q-value evaluation to estimate expected future rewards. Q* aims to guide LLMs in selecting promising next reasoning step without requiring fine-tuning for the targeted task.
weaknesses: 1.Novelty: This work appears to be an incremental advancement based on the Reasoning as Planning (RAP) framework utilizing Monte Carlo Tree Search (MCTS). The primary contributions are as follows:

(1) Integration of the A* Algorithm: The authors incorporate the A* algorithm as the foundational framework for path searching by defining appropriate evaluation functions to guide the search process.

(2) Modification of the Heuristic Function: The traditional heuristic function h(n) in A* is replaced with Q-values. The authors employ three distinct methods to evaluate these Q-values, as detailed in Section 4.1.
While these contributions extend the existing RAP-MCTS framework, it remains some problems as below: 

2.Choice of A* Algorithm
The paper utilizes the A* algorithm for heuristic search within the proposed framework. However, alternative heuristic algorithms such as Particle Swarm Optimization (PSO) [1] or Ant Colony Optimization (ACO) [2] could also be considered for heuristic search. The authors have not provided sufficient theoretical justification or empirical evidence to demonstrate that A* is the optimal choice for the path search framework in this context. Although line 194 mentions, "When the heuristic h(⋅) is admissible [3], A* guarantees to find the optimal path," this statement does not address why A* was chosen over other heuristic methods or provide comparative analysis to support its selection since no any guarantee about h(st) is admissable. 

3.Reward Function Design
The design of the reward function, as described in line 240, raises concerns regarding its ability to ensure that the heuristic h(⋅) is admissible. For the A* algorithm to guarantee the optimal reasoning path, the heuristic must not overestimate the true cost from the current state to the goal. The current reward function design does not convincingly demonstrate that h(st) meets the admissibility condition, thereby casting doubt on the claim that the heuristic search can reliably produce optimal reasoning paths as stated in the authors' contributions.

4.Additional Experimental Evaluations
The paper would benefit from additional experiments to evaluate the scalability and performance of the Q* framework in larger-scale search and optimization problems. Specifically:

(1) Scalability of Q*: An assessment of how the Q* framework performs as the complexity and size of the reasoning tasks increase would provide valuable insights into its practical applicability.
(2) Comparison with Other Heuristic Algorithms: Including comparative experiments with algorithms such as PSO or ACO would help determine the relative strengths and weaknesses of using A* within this context. Such comparisons could validate the choice of A* and highlight any advantages or limitations of the Q* framework relative to other heuristic search methods.

Reference:

[1]Dorigo M, Maniezzo V, Colorni A. Ant system: optimization by a colony of cooperating agents[J]. IEEE transactions on systems, man, and cybernetics, part b (cybernetics), 1996, 26(1): 29-41.

[2]Kennedy J, Eberhart R. Particle swarm optimization[C]//Proceedings of ICNN'95-international conference on neural networks. ieee, 1995, 4: 1942-1948.

[3]Russell S J, Norvig P. Artificial intelligence: a modern approach[M]. Pearson, 2016.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 1
questions: Some suggestions on writing and grammatical mistakes (I did not mark them all):

1. "leveraging plug-and-play Q-value models as heuristic function ..." -> "as a heuristic function", or use singular form throughout.

2. "can effectively solve various tasks via guiding LLMs to select..." -> by guiding LLMs

3. "We conduct extensive experiments on math reasoning and code generation tasks, which demonstrates ..." -> demonstrate

4. "Moreover, planning with MCTS often requires to perform costly rollout, which can significantly slow down the overall decoding process. In contrast, Q* solely relies on training a Q-value model to guide LLMs to select the most ..." -> I believe even for MCTS, a Q-value model or Q-table is learnt to guide the action selection after all the rollouts (e.g., in alphaGo). The writing needs to be refined here to reflect that the key difference is that performing MCTS during test-time is costly, while Q* learns the Q-value _offline_, and can be used directly during test-time.

5. "The solutions are Python code that is excepted ..." -> expected
soundness: 3
strengths: 1. The idea behind the paper is sound. By training a Q-function in an offline manner (possibly one for different task), we can deploy them during test-time to search for correct solutions without model re-training. However, I'd like to point out that this idea is not that novel because prior works in offline RL has used the same approach (and even demonstrated on LLMs).
2. Section 5.4 demonstrates that the paper's idea is preferable over fine-tuning the model over a specific task, because fine-tuning over a specific task might cause it to do worse in other tasks.
3. The paper acknowledges that its method raises the performance of smaller model but cannot outperform GPT-4 variants (Table 3,4,5).
summary: The paper presents Q*, where a Q-value function is learnt offline to estimate the "value" associated to each state in a reasoning problem. These Q-values are learnt with 3 methods: offline RL, Best of K sampling and offline MCTS. Combined with the aggregated utility function which labels the instantaneous reward associated with each state, they form the reward signal in A* search framework. During test-time, the paper performs A* search with the Q* values as the heuristic, producing the correct solution.
weaknesses: There are several weaknesses in this paper.

1. The first issue comes in the form of difference between offline data and test-time data distributions. The idea of learning a Q-function from offline data has been well studied in RL problems. However, it is well agreed in the RL domain that the problems faced during test-time might be different from those trained offline. For example, in Go and Chess, we cannot simply learn the Q-value function entirely from offline data (due to the enormous number of state space). Instead, we need to perform some variant of test-time search to find the best moves. Similarly, one might be faced with an entirely unseen math question during test-time and the Q-value, learnt from a different set of questions, is a poor heuristic during test-time. Could the authors comment about this?

2. The concept of aggregated utility g(st) is not clearly explained. I hope the author can explain clearly how the utility is learnt from static data because it forms the first half of the search heuristic (the second half uses Q-values). Furthermore, it seems difficult to learn the instantaneous utility associated with each state (e.g., a segment of code). The authors presented some ways to do so in the actual experiments, but do not explain how it is learnt clearly from offline data.

3. There are no error bars or repeated trials for experiments. Since this paper is mostly empirical, I think it is necessary to expect repeated trials and error plots.

4. There are quite a lot of grammatical mistakes and writing issues in the paper. I have marked some of them out in the next section.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Add some OOD experiments to show the generalization of the Q* framework. (for example, livecodebench or other math/code/agent benchmarks)

2. Remove the overstatements in the paper. 

3. Compare with TOOLCHAIN* in the experiment. (Including the efficiency and the performance)

4. Section 5.4 is unrelated to this paper, why did the authors add this section? (maybe LLAMA3.1 + MetaMath + Q* is a reasonable setting)

(updated, this is not the deduct point) 5. I have realized that OpenAI has a project named Q*, so I highly recommend the author change its title or it may mislead some people.
soundness: 2
strengths: 1. The paper contains enough background and related work. The authors give a clear explanation of the Q* framework.

2. Q* seems more efficient than the other methods (like MCTS).
summary: The author presents Q*, a framework for guiding the LLMs decoding process. This framework has 2 steps: 1. Training a Q-value estimator. 2. Using A* search to find the optimal state and action. The A* search contains 2 parts: 1. g(n) is the cost of the path from the start state to the current state, using an aggregation function to calculate the cost. 2. h(n) is the cost of the path from the current state to the goal state, using a Q-value estimator to calculate the cost. The experiment shows that the Q* can improve the performance of the LLMs' planning in the GSM8k, MATH, and MBPP benchmarks.
weaknesses: 1. Bad contribution statement. Most of the paper "formalizes the multi-step reasoning of LLMs as MDP" (like RAP[1]), so this is not a contribution. 

2. Lack of novelty. Both Q-value estimation[2][3]  and A*[4] search are mentioned in the related work, so the Q* is not a novel framework.

3. Lack of OOD experiment. The authors claim that the Q* is a "general framework", but the experiment only shows the result in GSM8k, MATH, and MBPP benchmarks, and the Q-value estimator is trained on these benchmarks. If the Q* can only work on these benchmarks, "the existing deliberation methods" (lines 76-77)[4] are still useful, and even better than Q* in GSM8k, so adding the OOD benchmark is needed.

4. Still need domain knowledge and manual selection. The authors claim that the Q* "does not rely on domain knowledge to design the heuristic function", however, the domain knowledge is still used in the aggregation function selection and the Q-value estimator training. And even the aggregation function selection is chosen manually and deliberately.



[1] Reasoning with Language Model is Planning with World Model

[2] Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning

[3] Step-level Value Preference Optimization for Mathematical Reasoning

[4] TOOLCHAIN*: EFFICIENT ACTION SPACE NAVIGATION IN LARGE LANGUAGE MODELS WITH A* SEARCH