DOTS: Learning to Reason Dynamically in LLMs via Optimal Reasoning Trajectories Search

Thank you very much for the feedback. We provide clarifications to your questions below.

1. W1 Q3: Please refer to our general reply about more baselines.

2. W2: We appreciate the reviewer's thoughtful feedback. (1) Generalization w.r.t. new LLM: We view the ability to adapt to the capabilities of different LLMs as a strength rather than a limitation for two key reasons. In practical applications, when developing a system, we typically begin by selecting a base LLM and then optimizing it. As a result, creating a planner that works universally across all LLMs is less meaningful than enhancing one particular LLM. Besides, for some questions, there may be no universal optimal trajectory. As demonstrated in Table 3, different LLMs exhibit varying performances even when using the same prompt engineering method. Thus, we argue that developing a universal planner may not be feasible. (2) Generalization w.r.t. new layers or actions: Our current pipeline cannot directly work with new layers or actions. When new layers or actions are added to the reasoning space, one will need to search for the optimal trajectories again, as now the search space has changed and the optimal trajectories in the past may not be optimal anymore. An LLM planner will then be trained based on the new set of optimal trajectories for the training instances. We argue that the additional work for this generalization is necessary (as we aim for optimality in planning), but the cost can be reduced. For example, the optimal trajectories from the past can be reused as results when the new layers or actions were not invoked when one decides whether the new layers or actions are helpful. Future work can also explore continual learning for efficiently adapting an existing LLM planner to plan with the new components.

3. W3: Please refer to our general reply about OOD task performance.

4. W4 Q4: We thank the reviewer for this suggestion and have added the full ablation result of all tasks in Appendix G. It consistently shows the importance of each module in our method.

Thank you very much for your insightful comment.

1. W1 and Q4: We thank the reviewer for this insightful question. For certain specific datasets, such as Game of 24, there exists a single reasoning path (PoT and verification) that consistently delivers the best performance. Therefore, the upper boundary for our method for this Game of 24 is PoT and verification (i.e., “self-refine” in our tables). Our method achieves comparable results to this upper boundary and outperforms other methods greatly, demonstrating our method can identify and select the best reasoning path for such problems, which is a satisfying result. Besides, in practical scenarios, questions from the user are not always in the same pattern and we cannot control the types of questions users may ask. Therefore, we argue that the average accuracy and the accuracy on datasets such as MATH, which encompass diverse questions, is a more significant indicator of our method’s utility. From our results, we see that dynamic reasoning, where an LLM flexibly adapts its reasoning trajectories, is absolutely necessary for the model to perform consistently and stably well across tasks. We kindly refer the reviewer to our General Response for a clarification of our goal in this work.

2. W2: For OOD results, please refer to our general reply.

3. W3 and Q3: We sincerely thank the reviewer for providing constructive feedback. To demonstrate near-distribution zero-shot settings, we selected two datasets from the big bench but not in BBH: auto-categorization and last-letter concatenation datasets. Both are reasoning datasets and represent near-distribution and zero-shot scenarios. We tested it with the Llama-3-8b-Instruct as the task-solving LLM.

From this table, we can find that our method performs better than prompting engineering and SFT methods in near-distribution zero-shot scenarios. It can select the appropriate action for each task (CoT for the auto categorization dataset and PoT for the last letter concatenation).

4. Q1: We apologize for the confusion and have revised the pdf.

5. Q2: Thank you for your comments. We trained the model using the few-shot datasets and the MATH dataset together. Regarding whether these few-shot datasets can be considered "near-distribution" to the MATH dataset, we believe this concept is somewhat vague, as mathematics and the few-shot datasets, such as symbolic reasoning and games, are not entirely related. However, we will revise the experiment section to provide a clearer statement of our data used in training.

6. Q3 and Q4 are answered above.

7. Q5: The idea of trying multiple reasoning attempts during inference and selecting the best result is very interesting. Our method can be seen as a single-shot planning and the mentioned idea can be seen as an exploration and selection process. The main reasons we use the single-shot manner are: 1) it’s faster and computationally efficient; and 2) selecting from multiple generated plans relies on having a reliable verifier, which is another open problem to explore.

8. Q6: For external planning methods, since our planning LLM is only an 8B model and the generation output is just a plan and not long (the number of tokens ranging from 100~200), the additional computing is very limited compared to the computing cost of the solving step with larger task-solving LLMs (such as 70B models or GPT-4).

We thank the reviewer for the thoughtful comments and hope to provide clear explanations during the discussion.

1. W1: We need to point out that our method is not complex in both the training and inference phases. During the training phase, we: (1) search for the optimal trajectory; (2) ask GPT-4o to analyze why this trajectory is optimal and add explanations; and (3) train the LLM to predict the optimal reasoning trajectory before solving the problem. In the inference phase, we only need to sample the plan onceone time. We have included an ablation study in section 3.6 to highlight the importance of the search process and the inclusion of explanations. It shows the effectiveness of each step designed in our method. The consistently strong performance of DOTS on various datasets and tasks also justified our overall approach design.

2. W2: We kindly note that our design includes decomposition (shown in line 177, Analysis Layer). Using a Python Interpreter as a tool is a common strategy when solving reasoning problems, and our exploration also incorporates this aspect (see line 181, Solution Layer). The main idea of the current version is to show the effectiveness of our idea. Adding more tools is naturally easy in our design as we could just add more possible actions in the Solution layer.

3. W3: We would like to clarify that the term "atomic trajectory" was not explicitly used in our paper. If you are referring to "atomic actions", these are based on prior work, as detailed in the first paragraph of Section 2.2. Alternatively, if your question refers to the "action trajectory after searching", the corresponding process is outlined in Section 2.3 and Algorithm 1. Specifically, this involves first enumerating all possible action trajectories, followed by testing each trajectory and collecting those that achieve higher accuracy. We hope this explanation addresses your query. If there is any additional information you would like us to elaborate on this question, we would be more than happy to engage in a deeper discussion.

4. Q1: Our few-shot setting demonstrates that our method can achieve improvements even with a limited number of instances and trajectories. However, we argue that few-shot examples help the LLM grasp high-level action planning ideas, such as "if the question involves math, use PoT." In contrast, more extensive training instances are necessary to teach the LLM fine-grained question-specific strategies, such as "use PoT when dealing with counting", "use CoT for geometric reasoning.", and "use verification for solving equations". Therefore, we argue that more training instances and trajectories would make planning ability better.

5. Q2: We appreciate the reviewer for raising this important future exploration. To extend this work with a process reward, we would like to mention a possible future work for searching for appropriate action in each process step: Assuming we have access to a process rewards model, our method could be applied stepwise by searching for the optimal action in each step. For example, the search might reveal that using a Python built-in function and verifying the result in step 1 achieves a higher process reward while applying query decomposition and CoT reasoning in step 2 yields better results. We could then construct training data with this trajectory, enabling it to learn the most appropriate action for each step.

Dear Reviewer in1A,

As we are approaching the end of the discussion period, we want to kindly remind you to read our response and share any further concerns you have. We are happy to address them before the end of the discussion. Thank you!

Thank you for your valuable comments. We have tried addressing each of the concerns below:

1. W1: For the concern about the OOD performance, please refer to our general reply. Besides, regarding additional computational costs compared with the prompting engineering method, we need to emphasize that our method does not introduce much additional computational cost during the inference stage. For external planning methods, since our planning LLM is an 8B model, the additional cost is very limited compared to larger task-solving LLMs (such as 70B models or GPT-4). For internal planning methods, in section 3.8, we compared our approach with prompt engineering methods, demonstrating that our method won’t cause additional computational overhead.

2. W2: Thanks for the suggestion. For this discussion, please refer to our general reply for the baselines of the hybrid training.

3. W3: For showing the combining of all non-empty actions, we added an extra baseline for using all three layers with non-empty action selection. Please refer to our general reply. Our method still outperforms this baseline.

4. W4: We sincerely appreciate the reviewer's thoughts; it is very helpful. However, it is out of the initial scope of our paper. This is because (1) We focus on the general and various reasoning tasks, not a specific task. Therefore, a code-based model, which may not be suitable for answering text-heavily reasoning tasks, is unlikely to be a candidate for the task-solving LLM compared with the general textual generation LLM. (2) If we consider a scenario where the task is only for code generation and the task-solving LLM is a code-based LLM, it will result in more successful PoT paths in searching. Consequently, our search result will always be the PoT and the tuned planner would tend to select PoT in all situations, which gives the exact same performance of using PoT.

2. More baselines

We added two sets of baselines in Appendix E. Our method still achieves significant improvement compared with these two baselines. The first set of baselines is about using all three layers in a static manner with the non-empty action selection. We reported the results when pairing different choices across the three layers (i.e., {Rewrite, Decompose} x {CoT, PoT} x {Verification}, a total of 4 baselines for each solver LLM setting). The result shows that our method still outperforms using all three layers. For detailed results, please refer to Appendix E in pdf.

The second baseline is about hybrid training, i.e., using a mixture of training data with both CoT and PoT. We use the same training instance and each instance is provided with both the PoT and CoT reasoning paths. We reported the results in Appendix E. The result shows that our method still outperforms the hybrid training. From the results, it can also be observed that for datasets where CoT outperforms PoT, training solely with CoT yields better performance. We analyzed several cases and found that although mixed training allows the LLM to choose between CoT and PoT dynamically, the results tend to perform random selection rather than actively analyzing the problem and selecting the most suitable path for the given question carefully because we didn’t guide the LLM learn this. This highlights the importance of searching for the optimal reasoning path before tuning the LLM.