🔗9.3.6 Q*

Paper link: Improving Multi-step Reasoning for LLMs with Deliberative Planning
⭐
During the autoregressive generation process, LLMs can have hallucinations and inconsistent statements when performing multi-step reasoning. Researchers from Skywork AI and Nanyang Technological University proposed the Q* framework, a general and flexible framework that improves the multi-step reasoning ability of LLMs by treating multi-step reasoning as a heuristic search problem.
By leveraging deliberate planning, Q* framework formalizes the multi-step reasoning process of LLMs as a Markov Decision Process (MDP) and uses the A* search algorithm and Q-value modeling to guide the model in selecting the most promising reasoning steps, thereby enhancing multi-step reasoning performance.
Extensive experiments on GSM8K, MATH, and MBPP demonstrate the effectiveness of this approach, helping to improve the reasoning abilities of existing open-source large language models.

Background

Problems Faced by LLMs in Reasoning:
  • Errors and Inconsistency: Although LLMs perform well in many natural language tasks, during the autoregressive generation process in multi-step reasoning tasks, as reasoning steps increase, errors, hallucinations, and inconsistent statements easily occur.
  • Lack of Deep Thinking: Current improvement methods mainly focus on enhancing LLMs’ “System 1” ability (i.e., fast, instinctive, but not very accurate thinking mode). Examples include:
    • Using extensive domain knowledge to construct complex prompts
    • Finetuning LLMs with large task-specific datasets
    • Training reward models to rank candidate responses
  • Limited Method Generalization: However, solving complex reasoning problems requires more advanced “System 2” thinking (i.e., deeper, deliberate, and logical reasoning steps), where LLMs fall short. Prior attempts to enhance “System 2” ability include asic tree search algorithms (BFS, DFS), MCTS (See Section 9.3.4 and Section 9.3.5), A*, etc. However, the utility functions these methods use typically require carefully designed domain knowledge for each specific task, making it difficult to generalize to new scenarios. Moreover, applying MCTS for deliberative planning in problems with many reasoning steps requires massive amounts of simulations.

Preliminary

📌
Representing LLM multi-step reasoning as a Markov Decision Process (MDP, see Section 3.2.2 Markov Process):
  • Given a question as input, the answer generation process of LLMs can be decomposed into multiple reasoning steps. The final answer sequence can be regarded as the concatenation of reasoning steps . Each step can be an entire line of output from the LLM or a fixed number of tokens. Thus, the multi-step reasoning process of LLMs can be formalized as an MDP process , where:
    • State : represents the concatenation of the input question and the reasoning trajectory generated up to time step , i.e.
    • Action : represents the next reasoning step generated by the LLM given the current state .
    • State transition : deterministic transition, i.e. is directly achieved by concatenating .
    • Reward function : usually based on the final outcome:
    • Discount factor .
  • The policy of an LLM, , can be expressed as:
  • Given the MDP and the LLM’s policy, we can define the state-action value function (Q-function) and the optimal strategy’s corresponding Q-function:
📌

A* Search Algorithm:

The A* algorithm is a commonly used heuristic search algorithm for pathfinding and graph traversal. It is used to find the shortest path from a source to a target . The A* algorithm calculates the priority of each node using the following function:
  • : the overall priority of a node. When selecting the next node to expand, the algorithm always chooses the node with the highest priority (smallest value).
  • : the accumulated cost from the start node to the current node.
  • : the estimated cost from the current node to the target node, also known as the heuristic function of A*.
For grid-based graphs, the following heuristic functions can be used:
  • If movement is only allowed up, down, left, and right, use Manhattan distance.
  • If movement is allowed in eight directions, use the Diagonal distance.
  • If movement is allowed in any direction, use Euclidean distance.
The A* algorithm uses two sets to represent nodes that are waiting to be expanded and nodes that have already been expanded: open_set and close_set. The algorithm flow is as follows:

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

This study treats the search for the most appropriate reasoning sequence for the problem as a heuristic search process. The utility function for each state in the search is defined as:
where represents the aggregated utility from the initial state to the current state:
Each state uses the process reward function for the reasoning step, and then the aggregation function combines them in different methods. For example, indicates taking the reward of only the final state as . The process reward function can be obtained from a PRM trained on preference data, or it can be a rule-based model, ground truth, or the logits of the model’s reasoning steps (reflecting the model’s confidence).
The heuristic term in this paper directly uses the optimal Q-value function as the heuristic function. At the same time, since the action space of reasoning steps is too large, this paper directly restricts the action space to the top-K candidate steps returned by the LLM. Therefore, the final utility function is:
✨

Estimation of the Optimal Q-value:

This study learns a proxy Q-value model to approximate the optimal Q-function. Specifically, it is learned from a dataset of LLM-sampled trajectories:
The objective is:
where:
  • is a trajectory sampled from the LLM (containing multiple actions, and means one action in this trajectory)
  • is the partial reasoning trace
  • is the label for the state-action value .
This paper proposes three methods to construct these labels:
  • Offline RL Method: Given an offline dataset , Fitted Q-iteration algorithm is used to iteratively learn the Q-value model. For each iteration , the label is constructed as follows:
      • After obtaining the labels for the current iteration, it can optimize to get the Q-value function of this iteration, and then proceed to obtain the labels for the next iteration. This process is repeated until convergence or until reaching the specified number of iterations.
  • Learning from rollout: Use the current LLM to perform random rollout or MCTS rollout, collecting a large number of rollout trajectories to form a trajectory pool . Then, from the subsequent trajectories, compute the label for the current state using the trajectory with the maximum cumulative reward:
  • Learning from stronger LLMs: Directly use trajectories generated by a stronger model, such as GPT-4 to compute labels:
      • where
(a): the deliberation process of Q*. Each state is associated with an value which is the weighted sum of the aggregated utility and the heuristic value. (b-d): estimating optimal Q-value with fitted-Q-iteration, rollout, and completion with stronger LLMs.
(a): the deliberation process of Q*. Each state is associated with an value which is the weighted sum of the aggregated utility and the heuristic value. (b-d): estimating optimal Q-value with fitted-Q-iteration, rollout, and completion with stronger LLMs.
✨

Deliberative Planning with A*

Once the proxy Q-value model is obtained, it can be used to calculate the value of each state, and then A* search is applied to perform optimal priority search.
Specifically, this study maintains a set to store unexplored state candidates, denoted as the unvisited set, which initially contains only the input problem . There is also a visited set to record the states that have already been explored. At each step, the algorithm selects the state with the largest value from the unvisited set, and then uses the LLM to generate top-K reasoning steps to expand this state into new candidates. This process is repeated until a terminal state is reached (i.e., a complete trajectory is obtained). The answer part of the terminal state is then returned as the final result.
notion image

Experiment

🔬

Experimental Setup

  • Datasets: Evaluations are conducted on three datasets: GSM8K, MATH, and MBPP (entry-level Python Programming dataset).
  • Implementation details:
    • Q-value estimation: The experiments in this paper show that Learning from rollout gives the best results for estimating Q-values.
    • Utility aggregation (Agg):
      • For GSM8K, a process reward model (PRM) trained on PRM800K is used to model , with min as the Agg function.
      • For MATH, to ensure fairness, all states in each trajectory are set to , because PRM800K contains samples constructed from the MATH test set, which poses a potential data leakage risk.
      • For MBPP, Agg is set to [-1].
    • A* planning:
      • For GSM8K and MATH, one line of reasoning is treated as one reasoning step.
      • For MBPP, a 24-token code snippet is treated as one reasoning step.
✨

Experimental Results

  • GSM8K dataset: Using Llama-2-7b as the base model, combined with the Process Reward Model (PRM) and the Q-value model (QVM), the Q* method achieved an accuracy of 80.8%, outperforming other methods based on the same LLM, and even surpassing the closed-source ChatGPT-turbo.
      • notion image
  • MATH dataset: Using Llama-2-7b fine-tuned on synthetic data and DeepSeek-Math-7b as base models, the Q* method achieved further improvements compared to the Best-of-N method. With DeepSeek-Math-7b, the enhanced performance even surpassed closed-source models such as Gemini Ultra, reaching an accuracy of 55.4%.
      • notion image
  • MBPP dataset: Using CodeQwen1.5-7b-Chat as the base model, the Q* method performed excellently in code generation, achieving an accuracy of 77.0%, outperforming the Best-of-N method.
      • notion image

Conclusion

Q* is a general, versatile, and agile framework designed to improve the multi-step reasoning ability of LLMs through careful and deliberate planning.
  • Q* formalizes multi-step reasoning as an MDP and uses A* search with a Q-value model, which effectively guides LLMs in selecting the most promising next step during multi-step reasoning tasks. This avoids the need to fine-tune the LLM for each task, reducing computational overhead and the risk of performance reduction.
  • The framework is general, flexible, and easy to implement, relying only on value models trained on real data. Moreover, it considers only single-step reasoning at each stage, making it more efficient than traditional approaches. Experimental results demonstrate the superiority of Q* in mathematical reasoning and code generation tasks.
 
Prev
CivitAI’s Payment Issue
Next
Awesome-AI-Tutorials
Loading...
Article List
LLM Learning Roadmap
✨ Awesome-Anything
Lazy loaded image Awesome-AI-Tutorials
Lazy loaded image Awesome-AI-Researchs
Lazy loaded image Awesome-GitHub-Repos
Lazy loaded image Awesome-AI-Examples
🖼️ Digital Image Processing
Lazy loaded image DIP I Image Demosaicing and Histogram Manipulation
Lazy loaded image DIP II Image Denoising
Lazy loaded image DIP III Edge Detection
Lazy loaded image DIP IV Digital Half-toning
Lazy loaded image DIP V Geometric Image Modification
🍃 LLM Components
Lazy loaded image 1.1 Tokenization
Lazy loaded image 1.2 Embedding
Lazy loaded image 1.3 Attention
Lazy loaded image 1.4 FFN, Residual Addition, and LN
Lazy loaded image 1.5 Positional Encoding
Lazy loaded image 1.6 Architecture & Decoding Policy
🌱 LLM Pre-training
Lazy loaded image 2.1 Pre-training Definition
Lazy loaded image 2.2 Pre-training Data
Lazy loaded image 2.3 Pre-training Workflow
Lazy loaded image 2.4 Continued Pre-training (CPT)
Lazy loaded image 2.5 Pre-training Evaluation
☘️ LLM Post-Training
Lazy loaded image 3.1 Supervised Fine-Tuning (SFT)
Lazy loaded image 3.2 Reinforcement Learning
Lazy loaded image 3.3 RLHF: Reinforcement Learning from Human Feedback
Lazy loaded image 3.4 Optimizing PPO-Based Algorithms
Lazy loaded image 3.5 Direct Preference Optimization (DPO)
🍀 LLM Popular Models
Lazy loaded image 4.1 BERT & Variants
Lazy loaded image 4.2 PaLM Series
Lazy loaded image 4.3 GPT Series
Lazy loaded image 4.4 LLaMA Series
Lazy loaded image 4.5 GLM Series
Lazy loaded image 4.6 Qwen Series
Lazy loaded image 4.7 DeepSeek Series
Lazy loaded image 4.8 MOE Series
Lazy loaded image 4.9 Other
🪴 LLM Applications
Lazy loaded image 5.1 Prompt Engineering
Lazy loaded image 5.2 LLM-based Agents
Lazy loaded image 5.3 Retrieval-Augmented Generation (RAG)
Lazy loaded image 5.4 Enterprise-level RAG Practice (I)
Lazy loaded image 5.5 Enterprise-level RAG Practice (II)
🌿 LLM Optimization
Lazy loaded image 6.1 GPU Usage Analysis
Lazy loaded image 6.2 Latency Analysis for Inference
Lazy loaded image 6.3 FlashAttention
Lazy loaded image 6.4 PagedAttention
Lazy loaded image 6.5 Packing for Efficient Training
Lazy loaded image 6.6 PEFT
🌾 LLM Compression
Lazy loaded image 7.1 Quantization
Lazy loaded image 7.2 Pruning
Lazy loaded image 7.3 Knowledge Distillation
🌵 LLM Hands-on Practice
Lazy loaded image 8.1 nanoGPT: Training a GPT from Scratch
Lazy loaded image 8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example
Lazy loaded image 8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis
Lazy loaded image 8.4 OpenRLHF: Reward Model Training and Code Analysis
Lazy loaded image 8.5 OpenRLHF: PPO Tutorial and Code Analysis
🌴 LLM Must-read Papers
Lazy loaded image 9.1 Technical Reports Explained
Lazy loaded image 9.2 MedLLM
Lazy loaded image 9.3 o1 Techinical Paper
Lazy loaded image 9.4 Prompting
Lazy loaded image 9.5 LLM-as-Judge
🌳 LLM Q&A
Lazy loaded image 10.1 Why GRPO training loss starts at 0
Lazy loaded image 10.2 VLM-R1 Reproduction Results
Lazy loaded image 10.3 KL(P||Q) or KL(Q||P)?
Lazy loaded image 10.4 Left Padding or Right Padding?
Lazy loaded image 10.5 What is QK-Norm?
Lazy loaded image 10.6 Why LLM training rarely adopts Dropout?
Lazy loaded image 10.7 Dropout and DropPath
🐝 VLM Image Encoders
Lazy loaded image 11.0 Introduction
Lazy loaded image 11.1 Vision Transformer
Lazy loaded image 11.2 CLIP
Lazy loaded image 11.3 BLIP
Lazy loaded image 11.4 EVA
Lazy loaded image 11.6 BLIP-2
Lazy loaded image 11.7 SigLIP
📝 MISC.
🤔 Large Language Models Pass the Turing Test.
Lazy loaded image OpenAI ChatGPT Progression
⏳ Claude Sonnet 4 Test with Sandtris
🔍 What DeepSeek-R1's Surge Teaches Us About SEO
💰 CivitAI’s Payment Issue
🎓 arXivSub, an Academia Assistant