Toward Self-Improvement of LLMs via Imagination, Searching, and Criticizing

We greatly appreciate the reviewer’s insightful feedback. Your recognition of our novelty, clear identification of challenges, and the effectiveness of our results is is highly encouraging to us.

---

[W1] It could be more clear in some parts and including some examples would be helpful. For example, in the part of 4.2 data synthesizing, I’m curious about how the synthesized data looks like.

We understand that some sections of the paper, such as the data synthesizing part in section 4.2, could benefit from further clarification and examples. To address your query about the synthesized data, we used the method described in [1] to generate synthesized questions. Here's an example synthesized based on GSM8K:

Question: Sandy's monthly phone bill expense is equal to ten times her age now. In two years, Sandy will be three times as old as Kim. If Kim is currently x years old, calculate Sandy's monthly phone bill expense.\nIf we know the answer to the above question is 340, what is the value of unknown variable x?\n

We hope this example provides a clearer understanding of how our data synthesizing process works. We are open to further discussions to improve the clarity of our work.

[1]. Yu, Longhui, et al. "Metamath: Bootstrap your own mathematical questions for large language models." arXiv preprint arXiv:2309.12284 (2023).

---

[W2] Authors identify that one of the challenges in working on LLMs’ self-improvement is the difficulty to get clear evaluations, and they proposed some method to get the evaluations, i.e. value function, PRM, ORM. But they do not discuss whether these methods are reliable to get perfect feedback.

We appreciate your insightful comment regarding the reliability of the evaluation methods such as the value function, PRM, and ORM for obtaining ideal feedback in the context of LLMs' self-improvement. In our study, we have utilized a combination of these evaluation methods to ensure better calibration results. The accuracy obtained on the GSM8K dataset for each critics is as follows:

• Value: 0.832
• PRM: 0.841
• ORM: 0.877

As a combination of them, we have achieved the following calibration scores:

• ECE: 0.1319
• AUROC: 0.958
• Brier Score: 0.0429

These results indicate that our proposed evaluation methods have provided reasonably good feedback for LLMs' self-improvement. We appreciate the suggestion and will continue to explore more reliable and robust evaluation techniques to enhance the self-improvement capabilities of LLMs in future research.

Thank you for your thoughtful feedback and insightful questions. We are encouraged by your approval of the novelty and effectiveness of ALPHALLM, as well as the clarity of this approach.

---

[W1] Given only final math answer annotations, ALPHALLM essentially performs the final-label (reward) classification problem and tree search based on the predicted reward to generate rationales for each final answer. It is still unclear how ALPHALLM could outperform other LLM-SFTs that utilize both rationale annotations and final answer annotations.

Our primary motivation stems from the observation that as LLMs continue to evolve, achieving human parity in numerous tasks, the necessity and quality of explicit rationale annotations in datasets such as GSM8K and MATH become points of consideration.

• The rationale annotations provided in datasets like GSM8K and MATH, while useful, are not necessarily optimal or superior to those that can be generated through advanced methods like MCTS guided approaches. The quality and suitability of these human-provided rationales can vary, and they might not always align with the nuanced reasoning capabilities of advanced LLMs. Our hypothesis is that the rationales generated by ALPHALLM, guided by MCTS, are potentially more robust and aligned with the model’s training process than those annotated in GSM8K and MATH.
• ALPHALLM utilizes final math answer annotations in conjunction with tree search based on predicted rewards to derive rationales. This method allows the model to explore a wider range of reasoning paths and select the most plausible ones based on calculated rewards, rather than being confined to potentially suboptimal human-provided annotations.
• As presented in our experimental results, ALPHALLM demonstrates superior performance compared to WizardMath, which utilizes over 96k data points for math-specific training, including both rationale and final answer annotations

---

[W2] ALPHALLM is a general framework. The authors should apply ALPHALLM to other tasks where learning signals are clear to demonstrate the overall effectiveness of this self-improvement framework.

We appreciate the reviewer’s suggestion regarding the application of ALPHALLM to a broader range of tasks. We are actively working on extending our framework to other STEM reasoning tasks to further validate its effectiveness. Preliminary results are promising, and we are committed to sharing a comprehensive report on these additional tasks in the near future!

---

[Q1] Have you ever tried using your value/reward function to perform RL fine-tuning with PPO?

We have not yet explored vanilla PPO but optimizing LLMs using trajectories from ALPHALLM with PPO/DPO is indeed on our roadmap for future work. To provide a comparative analysis, we included the results of Best-of-N ranked by our ORM in the current study. Best-of-N is often considered an alternative to PPO, particularly when N is sufficiently large. As demonstrated in Table 5 in the submission, ALPHALLM significantly outperforms Best-of-N, highlighting the effectiveness of our MCTS approach. This substantial margin of improvement underscores the potential of our method over traditional approaches like PPO, even before we integrate them.

Thank you for your response. Below is my feedback:

• To support the hypothesis that the rationales generated by ALPHALLM are more robust and aligned, the author should provide additional examples of these rationales and analyze the differences between human-annotated rationales and those generated by ALPHALLM. Without this, the theoretical strength of ALPHALLM remains unclear.
• After reading comments from Reviewers SX7Z and PPNK, I agree with their concerns about the implementation details. I also believe that the code should be provided to allow for verification of each component of your method.

Due to the main concerns raised above, I have decided to lower my score. However, given that the experimental results are indeed remarkable, I still support the acceptance of this paper.

[W1] It is not that clear how the option-level is implemented to separate sentences

Thank you for your feedback regarding the clarity of our writing. To clarify, the termination function for options operates differently depending on the dataset:

• For the GSM8K dataset, the termination condition occurs at the end of each line. This is based on the typical structure of this dataset where each line represents a distinct step or point.
• For the MATH dataset, due to its complexity and the base model's tendency to generate many '\n\n' line breaks with some less meaningful content between them, termination occurs at the end of a line if a formula pattern is detected. During inference, if '\n\n' is encountered, we perform a rule-based check for formula patterns. It terminates if a pattern is found or continues generating until the next '\n\n'.

We will include the clarifications in the revision.

---

[W2.1, W2.2] The ablation studies are not enough. 2.1 Other heuristic search methods; 2.2 Token consumption

In Table 5 of the original submission, a basic comparison of accuracy and the number of rollouts for self-consistency (majority-vote) and re-ranking methods for both GSM8K and MATH datasets was included. Following the suggestion, beam search (BFS) has also been incorporated into the experiments on GSM8K, as shown in the table below:

Token consumption is estimated by tracking MCTS rollouts and multiplying by the average tokens per rollout for Llama2-70b on GSM8K. The table indicates that \etaMCTS outperforms and is more efficient than other baselines.

---

[W2.3] 2.3 The ablation studies seem to be too simple

In addition, we have conducted ablation studies on GSM8K with base Llama2-70b to further demonstrate the influence of various components on the performance:

State Merge: The impact of the choice of heuristic functions (w/ hyperparameters) or model-based state merge is not very significant.

Fast-rollout: Number of rollouts: As the number of fast-rollouts increases, there is a corresponding improvement in performance. This is due to the reduction in the variance of the estimates. We used n=4 in our experiments for better trade-off between performance and efficiency.

Fast-rollout model: Using Llama-2-70b instead of Abel-7B-002 improves performance by reducing bias from a smaller model, but Abel-002-7B is faster with similar computational resources due to higher concurrency and quicker processing.

ORM: During training, each question is sampled multiple times, and the ORM predicts the correctness of each trajectory as True or False. The ORM's score (probability of outputting True) is used in the data generation process, along with the scores from value and PRM, eliminating the need for a threshold.

Adaptive branching: with adaptive branching, the performance score is 84.9, whereas without it, and using a fixed number of 4 branches, the score drops to 79.5.

---

[Q1] What is the difference between value function and PRM training? Also, please clarify your contributions and provide proper citations in the critic section.

Value function training and PRM training share some similarities, but they differ in task formulation, model architecture, and training loss.

Value function training focuses on estimating the expected FUTURE return of a given state or state-action pair. In contrast, PRM evaluates the quality of the current state or node. While PRM ideally requires quality labels for each state, due to the high cost and time involved in obtaining these, MC estimation is used as a proxy.

The Value Function model extends from a policy model by adding a value head and is trained using regression loss to predict expected returns, a continuous variable. PRM, trained using LM loss, evaluates the quality of states/nodes and benefits from the contextual understanding of language models. The differences in architecture and training result in distinct behaviors: the Value Function model has better precision (0.82 vs. 0.62) and calibration (ECE: 0.032 vs. 0.375), while PRM has superior recall (0.90 vs. 0.79).

Finally, our contribution lies in utilizing a trio of critics, each with strengths in different aspects, to provide reliable signals for guiding MCTS. We will restructure the critic section to clarify these points further, and will also ensure to appropriately cite the references [1,2,3,4] in that section to acknowledge their contributions to this field.

---

[Q2] Why choose to include the fast-rollout?

We use rollouts to the terminal state, because empirical evidence shows that the value network's accuracy improves with detailed information closer to the terminal state. For example, an early version of the value network had 89.5% accuracy at the terminal state and 83.2% accuracy three steps before, indicating more reliable predictions near the terminal state. Estimations at the end of rollout trajectories have lower bias but higher variance, so using the mean score from multiple rollouts offers low-bias, low-variance estimations. To mitigate bias from fast rollouts of a smaller model, we also keep value estimations at the actual progress point. Additionally, using the outcome reward model through fast rollouts provides extra signals that enhance learning, making our approach more robust and effective.

Due to the word limit, we included the references used in the rebuttal for reviewer PPNk in this comment.

We would also like to greatly thank the reviewer PPNk for the valuable feedback and insightful questions. We appreciate the recognition of our technical contributions and the potential of our algorithm. The reviewer's questions and suggestions have been crucial in improving the clarity of our work.

Reference

[1] Feng, Xidong, et al. "Alphazero-like tree-search can guide large language model decoding and training." arXiv preprint arXiv:2309.17179 (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] Lightman, Hunter, et al. "Let's verify step by step." arXiv preprint arXiv:2305.20050 (2023). [4] Uesato, Jonathan, et al. "Solving math word problems with process-and outcome-based feedback." arXiv preprint arXiv:2211.14275 (2022).

Your suggestions have helped us a lot to improve the work, though we do think there might exist misunderstandings about our method. We hope the following clarifications and additional evidences would be possible for a re-evaluation.

---

[W1] Lots of experimental details are missing. There is very little discussion of how AlphaLLM is trained or what base model it is (it's only mentioned in the introduction!). In particular, no information is given about the actual training of either the policy, the Critics, or the heuristic functions p_vM

Model and training details: As detailed in Appendix A.6 (line 609-624) in the original submission, the base models used are Llama-2-70b for GSM8K and WizardMath-70B-V1.0 for MATH, respectively. More details regarding the training (LR, warm-up, etc.), critics (data and model), and the fast-rollout are also provided in A.6.

Heuristic functions: As mentioned in Sec 4.3.3, the heuristic function p_vM can either be a faster rule-based measurement (e.g. edit distance) or a model-based method (e.g. prompting a LLM). For our experiments, we used edit distance. We also include additional ablation studies demonstrating that the performance is not very sensitive to the choice of p_vM:

---

[W2] The explanation of the options is lacking. Is Beta learnt or generated in some other way? Are the set of options I, learnt before a rollout or generated on the fly?

The termination function $\beta$ can be either be learnt or rule-based. In practice, for the GSM8K, we used a newline character (\n) as the termination condition, as shown in Table 4. For the MATH, termination condition occurs when a line with a formula pattern is found, with line breaks indicated by '\n\n'. During inference, if the generated output has '\n\n', we check for formula patterns using rules. It terminates if a pattern is detected or continues until another '\n\n' appears.

The set $I \subseteq S$ represents a set of initial states (not a set of options). It can be the state of a question or a previously terminated step.

---

[W3] Ablations are missing - in particular using different state methods, removing the Adaptive Branching, State Merge or Option-Level MCTS.

We have included ablations for different state methods, including PRM, fast-rollout with ORM, state merge, and large number of rollouts, tool-augmented ORM and option-level MCTS in Table 3, with a discussion in Sec 5.3. Additional ablations on adaptive branching are included here:

• With adaptive branching: 84.9
• Without adaptive branching: 79.5

---

[W4] No Code provided.

We plan to release the code after it undergoes a legal review by our entity.

---

[W5] The prompts in the appendix are reduced so we can’t replicate.

All the prompt templates will be included with the code release. These templates are detailed in Appendix A.5 , with specific rubrics for PRM and ORM are provided in Section A in the comment below.

---

[W6] Weird choice of baselines - surely to keep evaluations consistent with other LLMs, you should allow models to generate the same amount of output tokens with access to the PRM or ORM? I find it hard to believe COT on Claude-2 is a fair comparison, when the amount of optimisation pressure applied to multiple components of the method.

Our primary focus in this paper is to explore the self-improvement capabilities of LLMs. The experimental results demonstrate that our method, AlphaLLM, can achieve significant self-improvement. Moreover, we observed that the performance of Llama-2-70b, when employs our proposed \etaMCTS, is comparable to that of Claude-2 and GPT-4, highlighting the potential for self-improvement.

We included proprietary models mainly to illustrate the potential of our self-improvement method. The definition of a "fair comparison" can vary. Proprietary models, such as the latest versions of Claude-2 and GPT-4, are trained with unknown data quantities and types, often iteratively refined. By default, these models already incorporate chain-of-thought processes. It is plausible that proprietary models leverage System-2 results to enhance System-1 outputs, complicating direct comparisons.

In addition, we have demonstrated superior performance when compared with recent work [1] that also used MCTS and utilizes a similar token usage amount. For example, AlphaLLM achieved 51.0 on MATH while [1] achieved 34.3.

We hope this clarifies our approach and rationale behind the baseline choices. We welcome further suggestions to improve our evaluation methodology.

---

[W7] There is no example of a rollout provided.

Please refer to the example in the Section B in the comment.

---

[Q1] What policy was used in the end for the fast rollout?

As mentioned in Appendix A.6, the fast-rollout model is Abel-002-7B.

---

[Q2] The prompt templates are missing key parts - In particular when comparing to GPT-4 or other model perform was an equally informative prompt used? (Or at least a prompt of equal input tokens).

Please refer to the responses for [W5] and [W6].

---

[Q3] If the synthetic data generated is used few-shot by llama2, how does your model perform?

The synthetic data would still have the same question format as [2] but with answer steps from MCTS. Therefore, Llama2 with few-shot prompting would exhibit similar performance to that observed on the questions in [2].

---

[Q4] I think the related work is missing some key papers such as the STAR family of papers

We acknowledge that we missed the STaR papers and will include them in the related work section. Although we have cited some similar papers [3,4], we appreciate the importance of adding this relevant literature.

Due to the word limit, I am providing the prompts, examples, and references used in our rebuttal for reviewer SX7Z in this comment.

A. Rubrics in Prompt Templates

As mentioned in our response to [W5], the prompt templates for PRM and ORM are detailed in Appendix A.5 , with specific rubrics for PRM and ORM are

• PRM:

You are given a math problem, followed by a step-by-step reasoning process. Your task is to read the problem carefully, understand the solving steps, and check the correctness of the last reasoning step. Output 'True' if the last step is correct, and 'False' otherwise.

• ORM:

Assess a solution including final answer to a given math problem by following below steps.\n- Evaluate the method used for solving the problem.\n- Review each calculation step for accuracy. Check for computational errors, incorrect formula applications, or arithmetic mistakes.\n- The solution should use all the information provided in the question.\n- Examine the final answer for correctness, considering the calculations and method used.\n.

B. Rollout Example

Here's a detailed rollout example addressing [W7]:

Consider the following GSM-like question:

Question: Sandy's monthly phone bill expense is equal to ten times her age now. In two years, Sandy will be three times as old as Kim. If Kim is currently x years old, calculate Sandy's monthly phone bill expense.\nIf we know the answer to the above question is 340, what is the value of the unknown variable x?\n

A node in the second layer could have the following content:

Answer: We know that Sandy's monthly phone bill is 10 times her age. In two years, Sandy will be 3 times as old as Kim. The sum of Sandy's age now and 2 years is 3 times the sum of Kim's age now and two years.\nSandy's age now is 340/10 = <<340/10=34>>34. In two years, Sandy's age will be 34 + 2 = <<34+2=36>>36.\n

The parent of this node has the content:

Answer: We know that Sandy's monthly phone bill is 10 times her age. In two years, Sandy will be 3 times as old as Kim. The sum of Sandy's age now and 2 years is 3 times the sum of Kim's age now and two years.\n

And one of its fast-rollout paths could be:

The sum of Sandy's age now and 2 years is 36. The sum of Kim's age now and two years is x + 2.\n36 = 3(x + 2)\n36 = 3x + 6\n3x = 30\nx = 10\n#### 10

C. Reference

[1]. Zhang, Dan, et al. "ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search." arXiv preprint arXiv:2406.03816 (2024). [2]. Yu, Longhui, et al. "Metamath: Bootstrap your own mathematical questions for large language models." arXiv preprint arXiv:2309.12284 (2023). [3]. Li, Xian, et al. "Self-alignment with instruction backtranslation." arXiv preprint arXiv:2308.06259 (2023). [4]. Guo, Hongyi, et al. "Human-instruction-free llm self-alignment with limited samples." arXiv preprint arXiv:2401.06785 (2024).