Deliberate Reasoning for LLMs as Structure-aware Planning with Accurate World Model

Authors: Siheng Xiong, Ali Payani, Yuan Yang, Faramarz Fekri
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models, multi-step reasoning, planning with world model, structured reasoning, generator-discriminator architecture

Abstract:
Enhancing the reasoning capabilities of large language models (LLMs) remains a key challenge, especially for tasks that require complex, multi-step decision-making. Humans excel at these tasks by leveraging deliberate planning with an internal world model to simulate the potential outcomes of various actions. Inspired by this, we propose a novel multi-step reasoning framework for LLMs, referred to as Structure-aware Planning with Accurate World Model (SWAP). Unlike previous approaches that rely solely on Chain-of-Thought (CoT) reasoning in natural language, SWAP incorporates structural information to guide the reasoning process via a world model and provides a soft verification mechanism over the steps. Moreover, SWAP overcomes the challenge of accurate world state predictions in complex reasoning tasks by introducing a Generator-Discriminator architecture, which enables more reliable world modeling. Specifically, the generator predicts the next state, and the discriminator ensures alignment with the logical consistency required by the problem context. SWAP also encourages the policy model to explore a broad range of potential actions to prevent premature convergence. By resolving the bottlenecks of generation diversity for both actions and states using diversity-based modeling (DBM) and improving discrimination accuracy through contrastive ranking (CR), SWAP significantly enhances the reasoning performance of LLMs. We evaluate SWAP across diverse reasoning-intensive benchmarks including math reasoning, logical reasoning, and coding tasks. Extensive experiments demonstrate that SWAP achieves substantial improvements over the baselines and consistently outperforms existing methods.

Decision: reject
Reviewer scores: 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1729859424680

thread_type:
official_review

<thread_object>
creator: reviewer_38qg

<content>
summary:
This paper proposes SWAP, which aims to leverage structural information to guide the reasoning process with a world model. The idea is to enable a policy model and a world model to generate diverse actions and states, leading to a higher possibility of getting the correct results. The authors use multiple benchmarks to demonstrate the effectiveness of their proposed method.

strengths:
1. The paper demonstrates that SWAP outperforms various baseline methods on several benchmarks, with ablation studies confirming the effectiveness of the different components proposed by the authors.

2. A key contribution of the paper is framing LLM reasoning as a graph, which allows the model to generate diverse actions and states. This approach helps avoid local minima and enhances verifiability in reasoning, effectively addressing the limitations of Chain-of-Thought (CoT) methods. This contribution is significant for advancing multi-step reasoning.

3. The overall structure of the paper is well-organized; however, there are some weaknesses in presentation that should be addressed.

weaknesses:
1. Many annotations are used before defining them makes the method hard to follow, for example, in line 235 $P_{ori}(z; T_{s_{t-1}})$, it's not clear what is "learned distribution", and the definition of $P_{sem}$ is vague. Annotations are not consistent as well, for example, in line 272 $M_{wm-G}$ is not consistent with $M_{wm}$ in line 175. $M_{\pi_D}$ and $M_{wm-D}$ in equations 8, and 9 are not explained in the paper.
2. One of the contributions is process supervision, however, a comparison to other process supervision methods is missing.
3. The paper proposes to generate diverse actions and states to enhance performance, however, there is a lack of comparison and analysis with diversity-seeking methods using reinforcement learning designed to enhance diversity, such as [1, 2].
4. There lacks a confidence interval for the results, it's not clear how the performance is robust to initialization and randomness.

[1] Vijayakumar, Ashwin K., et al. "Diverse beam search: Decoding diverse solutions from neural sequence models." arXiv preprint arXiv:1610.02424 (2016).

[2] Hu, Edward J., et al. "Amortizing intractable inference in large language models." The Twelfth International Conference on Learning Representations.

questions:
1. How do the semantic similarity is calculated? As mentioned, the method aims to improve the diversity of the reasoning process. However, multi-step reasoning problems can have different orders of actions, leading to semantically different but the same answers. As shown in Figure 1
2.  As shown in Table 1, fine-tuning a Llama-3-8b shows inferior or comparable results to CoT, which is counterintuitive. Could you briefly explain the reason?
3. What is the training cost for SWAP, could the author provide a comparison between SFT and the proposed method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730651235305

thread_type:
official_review

<thread_object>
creator: reviewer_zs5k

<content>
summary:
The paper discusses a framework called Structure-aware Planning with Accurate World Model (SWAP) that aims to enhance the reasoning capabilities of large language models (LLMs). Key designs include:

**Framework Overview**: SWAP integrates structural information into the reasoning process, providing a soft verification mechanism that guides LLMs through multi-step reasoning. The authors suggest this approach may improve upon traditional Chain-of-Thought (CoT) reasoning methods, which can lack effective verification mechanisms.

**Generator-Discriminator Architecture**: The framework employs a Generator-Discriminator architecture to enhance world modeling. The generator is responsible for predicting future states, while the discriminator evaluates these predictions to improve the accuracy of the reasoning process.

**Diversity-Based Modeling**: The paper introduces a method to encourage diversity in action generation and state prediction, which is intended to allow the model to explore a broader range of solutions and avoid premature convergence on suboptimal paths.

**Contrastive Ranking for Discrimination Accuracy**: The authors implement a contrastive ranking approach that focuses on relative comparisons to improve the discriminator's ability to identify valid and invalid reasoning steps.

strengths:
- The introduction of normalization metrics (Eq. 2 & 4) for Diversity-Based Modeling is novel and interesting. 
- The method is shown effective on diverse tasks (math, coding and logical reasoning) with Llama3-8b as backbone.

weaknesses:
- The paper writing lacks clarity. Especially about the structured search algorithm as mentioned in 4.1. The generator and discriminator framework only shows how to select the action when more than one choices are provided. However, when more than one state and action are selected, what is the next state to process among multiple parallel choices? Figure 2 illustrates the process as a linear step by step process and fails to present the tree search as shown in Figure 4. It would be better if Figure 2 is replaced by detailed pseudo code.

- The framework of this algorithm is very similar to related work (TOT, RAP and [1]). RAP also uses world model based tree search and the main difference is MCTS style search is used to rank choices while this work uses a discriminator to reject choices. [1] proposes a similar pipeline that also uses discriminator-aided tree search. Wonder if the authors could point out the main difference with these works and provide a elaborated compare of the main pipeline.

- Figure 3 plots the use of LORA tuning to improve generation diversity. However, only posthoc adjustment methods are proposed in the section. Wonder if the authors could further explain this figure.

- Table 1 lists a lot of redundant backbone models that may not be directly comparable with SWAP + Llama3-8b. It would be better if the authors compare algorithms while fixing backbones and add more comparison with different backbone models.

[1]Chen, Ziru, et al. "When is tree search useful for llm planning? it depends on the discriminator." arXiv preprint arXiv:2402.10890 (2024).

questions:
1. What is the efficiency of the SWAP tree search algorithm compared to TOT and RAP? SWAP uses discriminator to prune many options and may have an advantage in efficiency.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730685762682

thread_type:
official_review

<thread_object>
creator: reviewer_f56z

<content>
summary:
This paper proposes a framework for improving reasoning in LLMs. The framework consists of a graph representation of the multi-step reasoning of the problem, a generator to generate possible next steps, and a discriminator to rank the possible solutions generated by the generator. The paper proposes several improvements to the framework such as adding diversity to step generation, improving the discriminator via process-supervision, and adding meta-knowledge of the problem into the LLM. The paper then shows the effectiveness of the framework by showing improved results on various mathematical reasoning and coding benchmarks.

strengths:
* The paper is well-written and conveys the overall framework well.
* Representing step-by-step reasoning as a graph is an interesting idea.
* The ablation study shows the performance gain of each design choice.

weaknesses:
* The related work section is very thin. The paper should more thoroughly compare its framework with existing reward modeling and process supervision literature. For instance, how does this method (and its performance) compare with [1]? Especially since the authors take inspiration from such works, they should also include some process-supervision methods in their Llama3-8B baselines. 

* The reported benchmark numbers for Llama3-8B are significantly lower than what the official llama has reported (e.g., according to Llama's report, HumanEval zero-shot should be 62.2, and MATH 4-shot-CoT should be 30.0, and GSM8K 8-shot-CoT is 79.6). This casts doubt on the validity of the evaluation done by the authors, particularly for HumanEval, where the original Llama's reported performance is higher than what SWAP achieves.

* There is very little discussion on the advantage of having a graph (with node connectivities) rather than sequentially listing all the reasoning steps. Can the authors provide any ablations on this? How well is the LLM able to parse the dependencies through json node connectivity? Can the authors provide an example of how the LLM reasoning for a few problems, and what the final graph produced by the LLM would look like?


[1] Wang et al, Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations, 2023

questions:
* Current results consider Llama3-8B generalist model. Is the proposed method able to improve math-specialized (e.g., Qwen Math) or code-specialized (e.g., Deepseek Coder) models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730763359565

thread_type:
official_review

<thread_object>
creator: reviewer_p8ky

<content>
summary:
This paper proposes SWAP, a framework to prompt GPT-4 generate solution chain with reflection to solve problems, and fine-tune other models on it to improve performance.

strengths:
The writing is easy to follow. They show the prompts used and examples for different tasks in detail. They do ablation studies to show the effectiveness.

weaknesses:
1. The baseline methods only include inference-time techniques without additional fine-tuning. How do these results compare to methods that incorporate fine-tuning? Specifically, I'm referring to approaches where GPT-4 generates or rewrites content based on the training set, and a base model is subsequently fine-tuned on those outputs.

2. What is the actual novelty of this work? Numerous studies already prompt GPT-4 with various roles for tasks like analysis, reflection, and planning.

3. How is the semantic similarity probability calculated? Also, the current normalization method seems weird. Is there a reason for mapping negative values to zero? This design choice seems suboptimal to me.

questions:
What is the underlying connection to reinforcement learning in this approach? The process of constructing supervision data seems quite similar to the steps taken for creating training data for reward models in Math-Shepherd. Additionally, how does the discriminator here relate to the concept of a reward model?

Could you clarify what "world model" specifically refers to in this context? How is it defined, and why is it considered a world model?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732641699638

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Acknowledgment of Reviewer Contributions

comment:
We sincerely appreciate the valuable comments and insights provided by all the reviewers. We have carefully reviewed the feedback to ensure a thorough understanding of the concerns raised. In response, we have refined our presentation and conducted additional experiments to address these concerns comprehensively. We are pleased with the substantial improvements our paper has achieved as a result of this process.

We are particularly grateful to Reviewers F56z, zS5k, and 38qG for engaging with our rebuttal and providing additional discussion. We also look forward to further interactions with Reviewer p8KY and are eager to address any remaining concerns from all reviewers.

We recognize ICLR's high standards and have made every effort to refine our paper to meet the expectations of the community. We sincerely hope the reviewers could consider our revised submission and the progress made during this process.

Thank the reviewers again for their time, insights, and constructive feedback!
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a multi-step reasoning framework with methodological limitations. The Structure-aware Planning with Accurate World Model (SWAP) closely resembles existing approaches without presenting novel contributions. Reviewers highlighted persistent issues with presentation, notation clarity, and experimental validation. Performance improvements were marginal and inconsistent, with concerns about evaluation methodology and significant performance gaps. The implementation details remained vague, and authors incompletely addressed fundamental questions about the method's efficiency and generalizability. It is hard to say that the proposed approach substantially advances multi-step reasoning for large language models.