Source-linked AI summary

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

Chaojie Wang, Yanchen Deng, Zhiyi Lyu, Liang Zeng, Jujie He, Shuicheng Yan, Bo An

arXiv:2406.14283v4cs.AI

TL;DR

Multi-step autoregressive reasoning can accumulate errors, while LLMs may lack the computation for in-depth deliberation. Q* uses plug-and-play Q-value heuristics to guide next-step selection without fine-tuning; experiments report improvements on math reasoning and code generation tasks.

  • Problem

    Autoregressive LLMs can introduce errors, hallucinations, and inconsistent statements across reasoning steps, and may not devote enough computation to difficult problems.

    Source: §1 · p. 1, §4 · p. 4
  • Method

    Q* casts reasoning as heuristic search and uses plug-and-play Q-value models to guide LLMs toward promising next steps without task-specific fine-tuning.

    Source: §1 · p. 2, §4 · p. 4
  • Results

    Q* improves multi-step reasoning for existing open-source LLMs on math reasoning and code generation, including 77.0% accuracy on MBPP with CodeQwen1.5-7b-Chat.

    Source: §1 · p. 2, §5.2 · p. 7
  • Takeaways & Limitations

    The authors conclude that Q* can be applied to various reasoning tasks without modification or fine-tuning the LLM beforehand.

    Source: §1 · p. 2, §6 · p. 8

Abstract

from arXiv · show

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

The paper motivates deliberative reasoning because autoregressive LLMs can accumulate errors across steps, while existing approaches rely on task-specific prompting, costly fine-tuning, reward models, or expensive search. It introduces Q*, which uses plug-and-play Q-value models to guide next-step selection without fine-tuning the LLM.

1 Introduction

  • Autoregressive LLMs can produce errors, hallucinations, and inconsistent statements as reasoning steps accumulate.The paper characterizes this fast, less accurate mode as “System 1”.Source: §1 · p. 1
  • Existing approaches improve “System 1” through sophisticated prompts, task-specific fine-tuning, or reward models that rank candidate responses.The paper notes that fine-tuning can require substantial computation and risk performance degeneration on other tasks [5] [6] [20] [21].Source: §1 · p. 1
  • Prior search-based deliberation can require laborious task-specific utility design, and MCTS may need many rollouts, slowing decoding.The paper identifies BFS, DFS, MCTS, and A* as prior approaches to enhancing “System 2” reasoning [25–30].Source: §1 · pp. 1–2
  • Q* uses plug-and-play Q-value models as heuristics in an A*-based search to guide LLMs toward promising next reasoning steps without fine-tuning.The framework formalizes reasoning as an MDP and estimates optimal Q-values using offline reinforcement learning, rollout sequences, or completion with stronger LLMs.Source: §1 · p. 2

2 Related Works

The related work covers alignment methods and planning approaches for improving LLM outputs. The paper positions Q* as using a learned Q-value model for intermediate-step guidance without changing LLM parameters or relying on task-specific utility design.

2 Related Works

  • Alignment methods include SFT, RLHF, DPO, and Aligner, while Q* avoids massive human preference-pair collection and does not modify LLM parameters.The paper says this avoids the potential risk of performance degeneration on other tasks.Source: §2 · p. 2
  • Planning methods such as ToT, A*, and MCTS guide reasoning through search, but their utility functions may rely on inaccurate LLM feedback or task-specific design.MCTS can also require costly rollouts that slow decoding [26] [27] [28].Source: §2 · p. 2
  • Prompt engineering requires expertise and case-by-case tuning, while fine-tuning can be costly and risk other-task performance; verifiers rank solutions without intermediate-step guidance.Q* instead uses a plug-and-play Q-value model to direct deliberation without modifying LLM parameters.Source: §2 · p. 3

3 Preliminary

The preliminary section models multi-step reasoning as an MDP whose states contain the question and generated reasoning so far, and whose actions are the next reasoning steps. It also introduces A* search, which prioritizes frontier states using accumulated utility and a heuristic estimate.

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

  • The MDP state concatenates the question with the partial reasoning trace, and the action is the next reasoning step generated from that state.The state transition concatenates the new step; the reward measures how well the question is solved, and the policy is embodied by an LLM.Source: §3.1 · p. 3
  • For code generation, reward is 1 when code passes all test cases; for math reasoning, reward is 1 when the final answer matches ground truth.The passage describes this as an outcome-based reward.Source: §3.1 · p. 3

3.2 A* Search

  • A* ranks frontier vertices by accumulated path cost plus estimated remaining cost, then explores the vertex with the lowest value first.With an admissible heuristic, A* guarantees an optimal path.Source: §3.2 · p. 3
  • Figure 1(a) depicts each state’s f-value as a weighted sum of aggregated utility and heuristic value; panels (b–d) show Q-value estimation routes.The routes shown are fitted-Q-iteration, rollout, and completion with stronger LLMs.Source: Figure 1 · p. 4

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

Q* casts multi-step reasoning as heuristic search, using a Q-value model to help an LLM choose promising next steps without task-specific fine-tuning. It estimates Q-values through three approaches and uses them to guide A* deliberative planning.

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

  • Q* scores each reasoning state with accumulated process utility plus a weighted heuristic based on the probability of reaching the correct answer.The heuristic is the optimal Q-value; the process reward can encode task preferences or prior knowledge.Source: §4 · p. 4

4.1 Estimation of Optimal Q-value

  • The proxy Q-value model is estimated through offline reinforcement learning, best-reward rollout sequences, or trajectory completion with a stronger LLM.Offline reinforcement learning uses Fitted Q-iteration; stronger-LLM completion is an option when such a model is available.Source: §4.1 · p. 5

4.2 Deliberative Planning with A*

  • A* repeatedly expands the unvisited state with the highest f-value using top-K LLM step alternatives, stopping at a complete trajectory and returning its answer.The proxy Q-value model supplies the heuristic used to calculate each state's f-value.Source: §4 · p. 4, §4.2 · p. 5

5 Experiments

Across GSM8K, MATH, and MBPP, Q* outperforms Best-of-N and achieves reported task-specific accuracies. The experiments evaluate rollout-based Q-value estimation alongside task-specific utility signals and A* planning.

5.1 Experimental Settings

  • The evaluation covers grade-school math, high-school competition math, and entry-level Python coding, using lines as math actions and 24-token snippets as MBPP actions.Planning expands K = 6 actions per step, uses λ = 1, collects N = 6 trajectories per question, and selects the trajectory with the highest f-value.Source: §5.1 · pp. 5–6
  • Utility aggregation is task-specific: GSM8K uses a PRM with minimum aggregation, MATH sets passed-state PRM signals to zero to avoid potential test-set leakage, and MBPP penalizes TokenError by -0.5.For MBPP, aggregation uses [-1] to cancel earlier penalties because later code steps may fix mismatched delimiters.Source: §5.1 · p. 6
  • Q-value labels are built by selecting the highest-reward reasoning path from rollouts or MCTS for each state-action pair.Rollouts use high temperatures (τ = 0.9 for math and τ = 0.2 for code); reward is 1 for a correct math answer or code passing all tests.Source: §5.1 · p. 6

5.2 Experimental Results

  • 80.8% accuracy on GSM8K is reached by Llama-2-7b with Q* combining PRM and QVM, surpassing ChatGPT-turbo.With the same PRM/QVM, verification performs better than alignment, and Q* with constant aggregated utility outperforms Best-of-N.Source: §5.2 · p. 7, Table 2 · p. 7
  • 55.4% accuracy on MATH is reached by DeepSeekMath-7b with Q*, which improves over Best-of-N on both evaluated base models.The reported result also surpasses closed-source models including Gemini Ultra (4-shot).Source: §5.2 · p. 7, Table 3 · p. 8
  • 77.0% accuracy on MBPP is reached by CodeQwen1.5-7b-Chat with Q*, outperforming Best-of-N.The task-specific utility penalizes tokenizer errors while code delimiters or indentation are mismatched.Source: §5.2 · p. 7, Table 4 · p. 8

6 Conclusion

Q* guides LLM reasoning with a plug-and-play Q-value model and single-step deliberative planning. The paper reports improvements across math and code tasks without fine-tuning the LLMs.

6 Conclusion

  • Q* uses plug-and-play Q-value models as heuristic functions to guide LLMs across reasoning tasks without modifying the LLMs.The framework relies on ground-truth to train value models and avoids task-specific utility design.Source: §6 · p. 8
  • Q* considers one reasoning step at a time rather than complete rollouts such as MCTS simulations.The conclusion characterizes this as making the framework agile.Source: §6 · p. 8
  • Evaluations on math reasoning and code generation tasks report that Q* improves reasoning performance without fine-tuning the LLM beforehand.The conclusion notes that avoiding fine-tuning also avoids potential performance degeneration on other tasks.Source: §6 · p. 8
Loading 2406.14283v4…