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

6. Performance of SFT on CoTs

The performance of SFT on CoTs slightly decreases for certain datasets (MATH, HumanEval, MBPP). For these datasets, we fine-tune the model using the provided reasoning processes (or code) written by humans. Language models typically benefit from training data with more steps (tokens) that detail the thinking process. However, we observed that human-written answers are often more concise, omitting some intermediate steps and details.

This results in a reduction in the number of reasoning tokens (during inference) after fine-tuning, which negatively impacts the model's performance. As noted by OpenAI, there exists an inference-time scaling law: models achieve higher accuracy with more inference tokens.

Additionally, it is important to note that these pre-trained models have likely already been fine-tuned on these datasets. Re-fine-tuning on the same datasets yields marginal benefits and can even degrade performance.

7. Training cost

As discussed in Section 4.2 and illustrated in Figure 2 (the revised version), our approach is implemented using SFT with LoRAs, with the additional cost arising from training a semantic equivalent LoRA. Specifically, starting from the original steps, we generate semantically equivalent alternatives by prompting GPT-4o. The semantic equivalent LoRA is then fine-tuned using the collected training data. For each trajectory, we sample some steps and generate two alternatives for each step. Consequently, the additional training cost of our method is no more than twice that of SFT with LoRA on the trajectory, resulting in a total cost of less than three times the standard SFT with LoRA.

Reference:

[1] Lightman, Hunter, et al. "Let's verify step by step." arXiv preprint arXiv:2305.20050 (2023).

[2] Wang, Peiyi, et al. "Math-shepherd: A label-free step-by-step verifier for llms in mathematical reasoning." arXiv preprint arXiv:2312.08935 (2023).

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

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

Thank you for your insightful comments, which are incredibly helpful in enhancing our work!

1. Notation

Thanks for your suggestion! We refine the notations throughout the framework, and add the detailed pseudo-code (Algorithm 1 and 2). Please refer to the revised version (Section 4.1, 4.2, and 4.3) for more details.

2. Comparison to process reward model

We add the following content in the related work section.

Although recent research has increasingly explored automatic process annotations using tree search, training an effective PRM remains challenging, as from a mathematical perspective, it assigns a numerical value within $[0, 1]$ to each state independently. To overcome this problem, we propose a novel strategy for automatic ranking annotation, i.e., given the current context and a set of candidate options, selecting the best option based on relative quality. Our ranking strategy offers significant advantages over traditional PRMs: 1) it emphasizes relative quality, making it more robust to noise; 2) it simplifies optimization and enhances generalization. Notably, our high-quality automatic ranking annotation method is non-trivial as it systemically incorporates three key factors: 1) structural information; 2) correctness; and 3) semantical equivalence.

We also add experiments with PRMs (PRM800k [1] and Math-Shepherd [2]) with the same baseline model in the revised version (Table 1). It shows that the performance of our method is better than those of PRMs.

3. Comparison to other diversity-seeking methods

Compared to related work [3,4], our approach offers the following advantages:

1. Diverse beam search [3] performs beam search in groups using a diversity-augmented objective. However, it has several limitations: 1) Beam search at the token level is computationally intractable; 2) Searching the most suitable strategies and hyperparameters for similarity calculation can be time-consuming; 3) For reasoning tasks involving special tokens (e.g., math reasoning or first-order logic reasoning), embedding-based similarity calculations may be unreliable.

   In contrast, SWAP (implemented as SFT with LoRA) employs an end-to-end learning paradigm, leveraging the world knowledge in pre-trained LMs. Our generation process is sampling-based, making it more efficient than token-level beam search.

2. GFlowNets fine-tuning [4] is a diversity-seeking reinforcement learning algorithm that uses amortized Bayesian inference. Although it demonstrates better performance compared to SFT with limited training data, it is unclear whether it can scale to large-scale datasets and complex reasoning tasks. As a reinforcement learning method, GFlowNets fine-tuning can be significantly more challenging and costly to train when dealing with large-scale datasets.

   In contrast, SWAP is more scalable and better suited for handling large-scale datasets and complex reasoning tasks efficiently.

We have summarized the above advantages and added to the revised version (Section 4.2).

4. Confidence interval

We have added confidence interval for the results in Table 1. Please refer to the revised version.

5. Calculation of semantic similarity

We have added the calculation process in the revised version (Section 4.2).

The effect of action order in multi-step reasoning is an interesting question. Currently, our approach focuses on increasing diversity step-by-step. In our experiments, we observed that the order of actions typically remains stable, since the actions are often interdependent.

As future work, this issue could be addressed through post-processing filtering based on similarity calculations. Specifically, if a generated trajectory is too similar to any existing trajectories in the stored pool, it will be discarded. To compare trajectories, the proposed graph representation provides a significant advantage, as graph similarity algorithms can be effectively utilized for this task. Additionally, we can explore data augmentation by reordering the actions based on the graph's dependency structure, further enhancing diversity and robustness.

4. Our approach on specialized models

This is an interesting question! Given the time constraints, we primarily conducted experiments using the DeepSeek-Math-7B-Instruct model and observed that SWAP also performs effectively with this specialized model. Expanding the evaluation to include additional specialized models is planned as part of future work.

Reference:

[1] Lightman, Hunter, et al. "Let's verify step by step." arXiv preprint arXiv:2305.20050 (2023).

[2] Wang, Peiyi, et al. "Math-shepherd: A label-free step-by-step verifier for llms in mathematical reasoning." arXiv preprint arXiv:2312.08935 (2023).

[3] https://github.com/meta-llama/llama3/blob/main/eval_details.md

[4] Dubey, Abhimanyu, et al. "The llama 3 herd of models." arXiv preprint arXiv:2407.21783 (2024).

[5] Touvron, Hugo, et al. "Llama: Open and efficient foundation language models." arXiv preprint arXiv:2302.13971 (2023).

Thank you for your insightful comments; they are extremely valuable for enhancing our work!

1. Related work

We add the following context in the revised version (related work section) to compare our framework with existing reward modeling and process supervision literature.

There are two primary types of reward models: Outcome Reward Model (ORM) and Process Reward Model (PRM). The ORM evaluates the fully generated solution by assigning a single scalar confidence score. Its training relies on outcome supervision by comparing generated answers with the ground truth. In contrast, the PRM provides stepwise rewards throughout the reasoning process, assigning a scalar confidence score to each intermediate steps. Empirical evidence shows that, compared with outcome supervision, process supervision ensures the correctness of each step, providing more benefits to multi-step reasoning. However, the training of PRM requires process supervision, which is hard to obtain, e.g., collecting process annotation from humans is inherently not scalable. Although recent research has increasingly explored automatic process annotations using tree search, training an effective PRM remains challenging, as from a mathematical perspective, it assigns a numerical value within $[0, 1]$ to each state independently. To overcome this problem, we propose a novel strategy for automatic ranking annotation, i.e., given the current context and a set of candidate options, selecting the best option based on relative quality. Our ranking strategy offers significant advantages over traditional PRMs: 1) it emphasizes relative quality, making it more robust to noise; 2) it simplifies optimization and enhances generalization. Notably, our high-quality automatic ranking annotation method is non-trivial as it systemically incorporates three key factors: 1) structural information; 2) correctness; and 3) semantical equivalence.

We also add experiments with PRMs (PRM800k [1] and Math-Shepherd [2]) with the same baseline model in the revised version (Table 1). It shows that the performance of our method is better than those of PRMs.

2. Performance gap

We checked the evaluation details of the LLama3 model [3,4,5]. The reported results actually come from majority voting rather than single-time inference. Specifically, for the same question, they generate 100 solutions for GSM8k, 256 solutions for MATH, and 200 solutions for HumanEval. In addition, it is possible that they perform some hyper-parameter optimization (e.g., seed, temperature, topK, topP). We also found that since the test set size of HumanEval is relatively small (164), the large performance gap is actually responding to a few test samples.

It would be too expensive to replicate the same results. Thus, for practical considerations, we use the default evaluation setting (single-time inference, low temperature=0.2, fixed seed) in our experiments for all methods to ensure a fair comparison.

3. Advantages of Graph Representation

We add the following content in the revised version (related work section).

Furthermore, we notice that although some reasoning processes are inherently non-linear, existing methods mainly follow a linear problem-solving manner. Language models are expected to implicitly infer the non-linear structure from the linear representation of the reasoning process, which proves challenging for complex reasoning tasks. To help the model, we integrate structural information into the reasoning process which explicitly represents the reasoning structure within the context. These structures provide the language model with additional guidance and control, enabling extra capabilities such as symbolic learning and verification.

We emphasize that our approach replaces the original state with a state-graph pair, consisting of a natural language description and its corresponding graph representation. The graph structure serves as an explicit representation of the non-linear dependency between steps. The effectiveness of integrating graph representations is already demonstrated through an ablation study (Table 2 'w/o structure info').

For parsing, we explored various formats and finalized using a JSON dictionary, where child indices serve as keys and lists of parent indices as values. During training data collection, we provided demonstrations in this format to GPT-4o and applied ad-hoc filtering to ensure high data quality. After fine-tuning, the base model demonstrated the ability to effectively generate and interpret this format. The good performance of our method (with ablation study) also demenstrate the parsing capability of LMs.

We already provide multiple example outputs in the paper (Appendix F).

Thank you for the follow-up questions.

1. Training on PRM800k

To maintain consistency with the original experiments (which were conducted on the complete test set of MATH), we also evaluate PRM (trained on PRM800k) on the full test set. We observed that the original PRM800k data [1] includes 4,500 test samples from the MATH dataset in its training data. To address this issue, we follow the approach in [2] and use their code [3] to fine-tune Llama3-8b-instruct on a filtered PRM800k dataset as a PRM. Specifically, we remove samples from the training set that belong to the test split of MATH.

2. Evaluation details of official Llama3

From our understanding, [4] uses the vague term "maj@1" to describe the evaluation process for GSM8k and MATH. If only one-time inference is conducted, the term "maj" would be unnecessary (as seen with other benchmarks in [4]). Furthermore, we observed that Llama3 largely follows the evaluation settings of Llama1 for many benchmarks (as mentioned in [4]). To verify, we referred to the Llama1 paper [5] and found that it employs majority voting with 256 samples for MATH and 100 samples for GSM8k (Table 7 in [5]). Based on our experimental results, we believe Llama3 uses the same majority voting settings. Also, we only use 4-shot CoT for all benchmarks (including GSM8k), in contrast to the 8-shot CoT utilized in the official report.

The CoT-SC in Table 1 refers to Chain-of-Thought reasoning with majority voting based on 8 samples.

3. Effect of temperature

We set the temperature to 0.2 as planning methods require diversity. We want to maintain this setting for a fair comparison.

Additionally, we conducted experiments on HumanEval using different temperatures, with the results as follows:

We observed that a lower temperature (0) leads to more stable results. However, the average accuracy remains at 52.4, which is close to 50.2 (reported in the paper). We believe that demonstrating the superiority of the proposed method under the same fair comparison setting is sufficient.

4. Performance gap on HumanEval

We provide the generated results on HumanEval from Llama3-8b-instruct (0-shot) for further verification.

Anonymous cloud storage: https://drive.google.com/drive/folders/1nSg15osa-dwBTYPuepyYT1NiyCXOSgie?usp=sharing

Reference

[1] https://github.com/openai/prm800k/tree/main

[2] Sun, Zhiqing, et al. "Easy-to-hard generalization: Scalable alignment beyond human supervision." arXiv preprint arXiv:2403.09472 (2024).

[3] https://github.com/Edward-Sun/easy-to-hard

[4] https://github.com/meta-llama/llama3/blob/main/eval_details.md

[5] Touvron, Hugo, et al. "Llama: Open and efficient foundation language models." arXiv preprint arXiv:2302.13971 (2023).

To avoid confusing the audience, we have added clarifications in Table 1 to better convey the implementation details. Specifically, we have updated method names as follows: "Few-shot CoT" has been changed to "Few-shot CoT (4-shot)", "CoT-SC" has been updated to "SC@maj8", and "PRM (PRM800K)" is now labeled as "PRM (PRM800K*)". Additionally, we have clarified in the table caption with the note: "We use the filtered PRM800K dataset [1] to evaluate performance on the full MATH test set." Please refer to the revised version.

Another potential issue is the comparison between low-temperature sampling (e.g., 0.2) and greedy decoding. Given the time constraints, our experiments were primarily conducted on HumanEval. We observed that while greedy decoding produces more stable results, the performance differences between the two methods are minimal. In the next version of this work, we plan to conduct experiments across all benchmarks and include this analysis in the appendix.

We hope these modifications effectively address your concerns!

[1] Sun, Zhiqing, et al. "Easy-to-hard generalization: Scalable alignment beyond human supervision." arXiv preprint arXiv:2403.09472 (2024).

Thank you for your insightful and constructive comments, which are greatly appreciated and extremely helpful in enhancing our work!

1. Clarity

Thanks for your suggestion! Figure 2 is mainly used to illustrate the input and output of different modules. We replace it with the detailed pseudo-code in the revised version (Algorithm 1 and 2). Additionally, we modify Figure 1 to depict the step-by-step process of graph construction. To enhance clarity, we also refine the notations throughout the framework. Please refer to the revised version (Section 4.1) for the updated details.

2. Difference from related work

The key innovation in our framework is viewing multi-step reasoning as the process of entailment graph construction, supported by a structure-aware planning approach tailored for this purpose.

Using graph representation offers several advantages: (1) It explicitly captures the non-linear structure of the reasoning process, providing LMs with enhanced guidance; (2) It allows for greater control over the reasoning process. For instance, graph representation facilitates structural verification during training data collection and enhances discrimination accuracy during inference. (3) It also enables exciting future work, such as exploring the impact of action order in multi-step reasoning. By capturing the dependency between steps, it becomes possible to identify interchangeable steps for data augmentation. Furthermore, for long-context reasoning tasks (e.g., OpenAI's o1 model), a graph structure can provide improved understanding and better control over the process.

Given the representation, we further notice that existing methods (ToT, RAP) mainly uses prompt engineering which totally or partially rely on self-evaluation. These strategies bring limited benefits in complex reasoning tasks since self-evaluation without external feedback can be unreliable [2,3]. To overcome this challenge, we propose using a generator-discriminator structure with fine-tuning, and further identify the bottlenecks of generation diversity and discrimination accuracy. We address these bottlenecks with architecture-level adaptation. This generator-discriminator structure also contributes to obtaining an accurate world model which is not considered in related work (RAP).

As for [1], it mainly focuses on code generation (text-to-SQL Parsing, Python-program-based math reasoning), and relies on external feedback (i.e., program execution results). Unlike [1], we do not constrain the model on code generation which could hurt its expressiveness on complex reasoning tasks [4]. It also does not consider architecture-level adaptation for generation diversity and discrimination accuracy.

All these strategies distinguish our framework from related work and contribute to substantial improvements (see TOT, RAP in Table 1).

3. Explanation of generation diversity calculation

We have included the detailed calculation process in the revised version (Section 4.2). Please refer to it for further details.

4. Comparison to existing methods with different base models

Thanks for your suggestion! We have delete the methods in Table 1 that uses different base models. We now test different methods with the same base model (Llama 3-8b-instruct). We also add PRM methods (PRM800k [5] and Math-Shepherd [6]). In addition, we change the base model into Mistral-7b-instruct and test all the methods. The results (Table 1) show that our method consistently outperform existing methods.

5. Efficiency

The time complexity of SWAP is $O(bNT)$, where $b$ is the breadth limit, $N$ is the generation number limit, and $T$ is the step limit. In contrast, the time complexity of RAP (using MCTS) is $O(N_{sim}NT)$, where $N_{{sim}}$ is the total simulation number limit. Typically, a large number of simulations $N_{{sim}} \gg b$ are required to reliably estimate $Q(s, a)$. For ToT, the time complexity depends on the implementation strategy: 1) Breadth-First Search (BFS): without pruning: $O(N^T)$; with pruning: $O(bNT)$. 2) Depth-First Search (DFS): The complexity depends on the state evaluator. The traversal continues until the state evaluator deems the final state satisfactory, making the complexity tied to the evaluation criteria.

In conclusion, SWAP is more efficient than RAP and ToT (BFS without pruning version). It is similar to ToT (BFS with pruning version).

Thanks to the authors for their efforts to improve the paper. The addition of pseudocode and the revised Table 1 make the paper much clearer than the original draft. Therefore, I would like to raise my rating to borderline.

I still have two remaining questions:

• Can these trained scores be used as reward models, similar to Math-Shepherd?
• The efficiency comparison may not be entirely accurate when based solely on bounds, as there could be differences in coefficients. Could the authors provide an estimate of the real-time ratio for different methods? For example, this could be based on the average number of tokens generated per test case.

Thank you for your thoughtful comments, all of which are very helpful for improving our work!

1. Baseline with fine-tuning

To investigate the impact of fine-tuning, we consider the following methods for comparison:

1. SFT on CoT (shown in Table 1): for datasets that provide reasoning process such as GSM8k and MATH, we directly fine-tune on it; for datasets without reasoning process such as FOLIO, ReClor, we use GPT-4o to generate the reasoning process data and fine-tune on CoTs that lead to the correct final answer. For coding datasets, we fine-tune on the completed code.
2. SWAP w/o discriminator (shown in Table 1): Since we simulate a Markov decision process (a sequence of interleaved states and actions) with structural information, which are different from CoTs, we consider the fine-tuned version of our framework with only the generator. Since there is no discriminator, we do not perform any planning during inference.

Our findings include that fine-tuning on trajectories with structural information brings benefits than CoTs, but more importantly, planning during inference further improves the performance. It is supported by the performance comparison between SWAP and SWAP w/o discriminator in Table 1.

2. Novelty

The key innovation in our framework is viewing multi-step reasoning as the process of entailment graph construction, supported by a structure-aware planning approach tailored for this purpose.

Using graph representation offers several advantages: (1) It explicitly captures the non-linear structure of the reasoning process, providing LMs with enhanced guidance; (2) It allows for greater control over the reasoning process. For instance, graph representation facilitates structural verification during training data collection and enhances discrimination accuracy during inference. (3) It also enables exciting future work, such as exploring the impact of action order in multi-step reasoning. By capturing the dependency between steps, it becomes possible to identify interchangeable steps for data augmentation. Furthermore, for long-context reasoning tasks (e.g., OpenAI's o1 model), a graph structure can provide improved understanding and better control over the process.

Given the representation, we further notice that existing agentic frameworks mainly use prompt engineering which totally or partially rely on self-evaluation. These strategies bring limited benefits in complex reasoning tasks since self-evaluation without external feedback can be unreliable [1,2]. To overcome this challenge, we propose using a generator-discriminator structure with fine-tuning, and further identify the bottlenecks of generation diversity and discrimination accuracy. We address these bottlenecks with architecture-level adaptation. All these strategies distinguish our framework from related work and contribute to substantial improvements (Table 1).

3. Semantic similarity probability calculation

We add the calulation process in the revised version (Eq. 1).

As for normalization, there are mainly two cases when tokens have negative values in the adjusted probability $P_{\pi{\text{G}}}$:

1. The token’s original probability $P_{\pi{G}}^{ori}$ is high, but its semantic similarity probability $P^{{sem}}_{\pi{{G}}}$ is even higher. This typically occurs for tokens that resemble previous generations, making them less desirable for diversity.
2. The token’s original probability $P^{{ori}}_{\pi{{G}}}$ is not high, indicating it likely represents a step that deviates from the intended progression of reasoning.

In both cases, setting negative values to zero allows us to discard these tokens and focus on maintaining a diverse and relevant generation. Other alternatives, such as Softmax, can distort the probability of irrelevant tokens by redistributing values across all tokens, even those that should ideally have very low or zero probability.

Consider an example with an intermediate step in a math problem:

For the first token, the original probability distribution $P_{\pi{G}}^{ori}$ is: {"Let": 0.6, "Simplify": 0.2, "Consider": 0.2, every other token: 0}. Given the previous generation "Let $x$ be 4.", the semantically similar probability distribution $P_{\pi{G}}^{sem}$ for the first token becomes: {"Let": 0.5, "Set": 0.2, "Consider": 0.3, every other token: 0}. Let $\gamma=1$, then $P_{\pi{G}}^{ori} - \gamma P_{\pi{G}}^{sem}$ becomes: {"Let": 0.1, "Simplify": 0.2, "Consider": -0.1, "Set": -0.2, every other token: 0}. Using the proposed method yields {"Let": 0.33, "Simplify": 0.66, every other token: 0}, focusing on high-relevance tokens. However, applying Softmax produces {"Let": 0.00011, "Simplify": 0.00012, "Consider": 0.00009, "Set": 0.00008, every other token: 0.0001}, which spreads probability mass thinly across all tokens, diluting the relevance of the original distribution and ultimately harming model performance.