B-STaR: Monitoring and Balancing Exploration and Exploitation in Self-Taught Reasoners

Q4: Such evaluation would classify as correct a mathematical proof that follows wrong reasoning but arrives at some form of correct conclusion?

A4: That’s a good point, you are right that the response would be classified as correct if the final answer is correct. Evaluating correctness solely based on the final answer can indeed cause flawed reasoning paths to be classified as correct, and almost all the self-improvement methods admit this limitation as it is non-trivial to guarantee the correctness of the intermediate steps. One approach to mitigate this issue is to introduce finer-grained supervision signals, such as process reward models (PRMs), which evaluate reasoning step-by-step. This approach is included in our work. How to evaluate intermediate steps faithfully is still an challenging topic in this direction, which we leave as future work to explore.

Q5: The conclusions from Section 3.2 are based on modifying individual hyperparameters while keeping the others fixed..... Could the authors provide further motivation on why these exact fixed values should be chosen and maybe extended empirical results showing the trend for other fixed values?

A5: Section 3.2 aims to provide a preliminary analysis to motivate our methods. The purpose of Section 3.2 is mainly to show the optimal hyperparameters for the best query effect may not be constant over training, thus we only provide a case where they are not constant, and those fixed values are commonly used heuristic values without cherry-picking (temperature=1, reward threshold=0, sample size=32). While we analyze sample size, temperature, reward threshold in Section 3.2, we did not really use any conclusion on temperature and reward threshold in our B-STaR algorithm directly -- in the actual B-STaR algorithm, we perform grid search of temperature and threshold together to select the best hyperparameters, rather than fixing the other to select each one separately. We did use the conclusion on sample size in our final algorithm, where we always chose the sample size to be the max allowable value within our budget. We agree that this conclusion on sample size is not comprehensively validated, this is because our B-STaR is designed to mainly focus on adjusting temperature and reward threshold in the current version, and we intended to fix the sample size and not adjust it dynamically to save the adjustment cost. We will explore adjusting sample size together with the other hyperparameters in the future.

Q6: Whether any adaptation of the reward threshold parameter is necessary, since only in the beginning its value differs from other stages of the training？

A6: To answer this question, we conduct two new experiments to demonstrate the necessity of dynamic adjustment of reward threshold.

1. We perform an ablation study where we only adjust the temperature while fixing the reward threshold. Specifically, we fix the reward threshold as 0.0 that is the one used by the default online RFT setting. The results are shown below.

   	MATH	GSM 8K
   Online RFT	23.2	46.8
   B-STaR (only adjust temperature)	25.0	53.1
   B-STaR (adjust temperature and reward threshold)	27.8	53.8

   The results demonstrate that changes in the reward threshold are helpful, and fixing the reward threshold leads to a noticeable drop.

2. The finer-grained hyper-parameter adjustments in B-STaR (Appendix C): In the previous experiments, due to the coarse adjustment with the granularity of 0.1, the dynamic change in hyperparameter was not obvious. Therefore, we further conduct finer-grained hyper-parameters search with the granularity of 0.05 for temperature and 0.01 for reward threshold. In Table 3, we observe more dynamic change to temperature and reward thresholds throughout the training process.

Therefore, we conclude that the dynamic tuning is necessary for the sustained growth and higher performance of B-STaR.

Thanks for your time and review! We respond to your comments below.

Q1:  My main concern about the paper is its generalizability to a broader spectrum of complex reasoning tasks? Could the authors provide a summary of the scale of the experimental evaluation of the main papers that they compare with?

A1:  Thank you for your suggestion. Research on self-improvement techniques for enhancing reasoning capabilities has primarily concentrated on mathematical and coding tasks, as demonstrated by works such as RFT [1], IRPO [2], REST-EM [3] and V-Star [4] . Some of these studies might further extend to commonsense reasoning tasks like ARC-Challenge and CommonsenseQA.

Here, to validate the generalization of B-STaR, following IRPO [2], we add the evaluation results on the ARC Challenge dataset that consists of multiple-choice science questions designed to evaluate commonsense reasoning beyond mathematics and coding challenges. We have added these results to Section 4 of the paper (Table 1 right and Figure 1 (d)), and we report a simplified version below as well. It demonstrates that B-STaR is not only highly effective for mathematical reasoning and code generation tasks but is also capable of generalizing to other reasoning domains. We would like to clarify that the our current evaluation on math, coding, and commonsense reasoning (for ARC) tasks aligns with or is already broader than many previous works on self-improving [1, 2, 3].

[1] Yuan et al. Scaling Relationship on Learning Mathematical Reasoning with Large Language Models. 2023

[2] Pang et al. Iterative Reasoning Preference Optimization. NeurIPS 2024.

[3] Singh et al. Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models. TMLR, 2024

[4] Hosseini et al. V-STaR: Training Verifiers for Self-Taught Reasoners, COLM 2024.

Q2: Specific concepts related to self-improvement training ..... formulating the paper's concepts in a mathematical language would greatly improve clarity.

A2: Thank you for your suggestion! We have added an introduction to specific concepts related to self-improving training and some related mathematical formulations in the Appendix B of the paper, specifically from Line 734 to Line 791.

Q3: Some ambiguous statements are included throughout the paper?

A3: We apologize for any confusion caused. We have improved the clarity of the mentioned contents by the reviewer. Specifically, we clarify the respective contents below.

Lines 155- 164 (previously 155-160) :

This paragraph introduces two underlying conditions essential for effective self-improvement: diverse exploration to ensure the discovery of high-quality responses and an effective reward function capable of accurately distinguishing these responses. These two conditions correspond to the concepts of exploration and exploitation in our paper.

-- "the distribution of candidates impacts how the reward model differentiates them, which suggests the distribution somehow influencing the reward function"

The distribution of candidates does not influence the reward function, and you are right that our reward function is fixed. Here we meant when the input to the fixed reward function changes, the reward function performance (the ability to distinguish the inputs) may change. This is like the same reward function could have different classification accuracy for different data distributions. We revised this part to improve the clarify (Line 166-178)

-- "What is the difference between stating that (the number of) high-quality respones among the selected ones cannot be too small and the percentage of the high-quality responses among the selected ones must be large enough"

The key difference between these two lies in that the former emphasizes the quantity of selected high-quality responses, while the latter emphasizes the ratio (but not the absolute quantity) of the high-quality responses among all selected responses. For example, if we select only 2 responses and they are correct, then the ratio is 100% while the absolute number is only 2, which cannot give us enough training data; if we select all 64 candidates, suppose 16 of them are correct and others are wrong, then we have 16 high-quality responses, but the ratio is only 25% -- feeding many incorrect responses into training is undesirable. Therefore, these two capture different aspects of the data. We have made clarification of this to Line 319- 341.

Thank you for the response! For your additional comments, we address them below.

Q1：While the two important factors are motivated heuristically, they are not grounded theoretically. Developing at least some conceptualization of why these are important would further improve the paper.

A1: Thank you for this suggestion. The objective of self-improvement can be expressed in the framework of reinforcement learning as follows:

$$\pi_\theta^{t+1} = \arg\max_{\pi_\theta^t} E_{x,y^* \sim \mathcal{D}, \hat{y} \sim \pi_\theta^t[\cdot|x]} \left[ R( \hat{y}, y^*) \right]$$

where $\mathcal{D}$ represents the dataset, $x$ and $y^*$ denote the sampled input and its corresponding ground-truth answer, respectively, and $\hat{y}$ is the sampled response. Here, $\pi_\theta^t$ corresponds to the language model in the $t$-th iteration. $R$ is the reward function. According to the policy gradient algorithm,

$$\nabla_{\theta} E_{x,y^* \sim \mathcal{D}, \hat{y} \sim \pi_\theta^t[\cdot|x]} \left[ R( \hat{y}, y^*) \right] = E_{x,y^* \sim \mathcal{D}, \hat{y} \sim \pi_\theta^t[\cdot|x]}\nabla_{\theta}R(\hat{y}, y^*)\log \pi_\theta^t[\hat{y}|x]$$

When $R(\hat{y}, y^*)$ is binary, the above equation turns to be simple data selection and supervised training loss that is exactly what we are doing. Thus, self-improvement can be viewed as a form of reinforcement learning, where maintaining a balance between exploration and exploitation is crucial and has been studied for years [1,2,3,4]. Conceptually in classic RL, exploration involves exploring the environment (analogous to reasoning tasks in our paper) by trying random actions (corresponding to sampling multiple candidates in our work) to gather more information about the environment. Exploitation, on the other hand, involves utilizing the known information to maximize the reward (similar to how we use the reward function to select data samples). Insufficient exploration may cause the training process to stagnate, while insufficient exploitation can lead to instability and large variance during training. We appreciate the reviewer's suggestion, and we have added this discussion to Appendix F due to the limited time and space available during the rebuttal period. For the next formal revision, we will ensure to clarify the connection above in the main body of the paper as well.

[1] Wikipedia contributors. "Exploration-exploitation dilemma." Wikipedia, The Free Encyclopedia, 25 Sept. 2024, 07:07 UTC. Available at: https://en.wikipedia.org/w/index.php?title=Exploration-exploitation_dilemma&oldid=1247645791. Accessed 26 Nov. 2024

[2] Weng, L. (2018). The multi-armed bandit problem and its solutions. Retrieved from https://lilianweng.github.io

[3] Şimşek et al. An intrinsic reward mechanism for efficient exploration. ICML 2006

[4] Sutton et al. Reinforcement learning: An introduction. MIT press, 2018.

Q2: Why do we strictly regard wrong responses as those that are of lower quality?

A2: That’s a good point. In previous works, wrong responses were often regarded as those of lower quality when additional reward models were not used. In such cases, the correctness of the responses is typically the only signal available for selecting data, as practiced in most works [1, 2, 3]. In our paper, we follow this approach for the APPS and ARC-Challenge datasets, where we do not train a reward model. However, for mathematical reasoning, where we did train additional reward models, we combined the correctness of the response and the scores from the reward models to compute the final reward (as described in Eq. 1 of the paper). Therefore, in our case, we did not strictly regard wrong responses as lower quality. Instead, responses with incorrect answers but high scores from the reward model were still likely to be selected in our approach.

[1] Avi et al. Beyond human data: Scaling self-training for problem-solving with language models. TMLR, 2024

[2] Yuan et al. Scaling Relationship on Learning Mathematical Reasoning with Large Language Models. 2023

[3] Zelikman et al. STaR: Bootstrapping Reasoning With Reasoning. NeurIPS 2022

Thanks for your time and comments! We respond to your comments below.

Q1: Limited Applicability Beyond Specific Tasks: Evaluation on other reasoning tasks, the robust of the query effect metric across different tasks beyond mathematical reasoning and code generation？

A1: Thank you for your suggestion. We follow IRPO [3] and add the evaluation results on the ARC Challenge dataset, a benchmark consisting of multiple-choice science questions designed to evaluate commonsense reasoning beyond mathematics and coding challenges, as shown below.

The results clearly indicate that B-STaR outperforms other baseline methods, achieving a significant improvement in reasoning performance. We note that the our current evaluation on math, coding, and commonsense reasoning (for ARC) tasks aligns with or is already broader than many previous works on self-improving [1, 2, 3]. These added results and more details are also included in Section 4 of the paper.

For other tasks such as open-domain QA and language understanding suggested by the reviewer, CoT typically does not help these tasks [4]. Also, these tasks are typically not considered as "reasoning" tasks and rarely used for self-improving reasoning in the literature. As our paper specifically focuses on complex reasoning, these non-reasoning tasks fall outside the scope of our paper.

[1] Avi et al. Beyond human data: Scaling self-training for problem-solving with language models. TMLR, 2024

[2] Zhang et al. Small Language Models Need Strong Verifiers to Self-Correct Reasoning. COLM 2024.

[3] Pang et al. Iterative Reasoning Preference Optimization. NeurIPS 2024.

[4] Sprague et al. To CoT or not to CoT? Chain-of-thought helps mainly on math and symbolic reasoning. Preprint 2024.

Q2: Unknown advantage over hyperparameter shift? It seems that the final dynamics of the parameters do not fluctuate a lot. It should the excluded that most of the benefit is gain by adding an offset rather that dynamically change these parameters.

A2: This is a good point. To answer this question, we conduct two new experiments to demonstrate the necessity of dynamic adjustment of hyperparameters.

1. First, we report comprehensive grid search results of various temperature and reward threshold combinations, where these hyperparameters are decided in the beginning and remain static. We present the complete results of online RFT in Figure 7 and Table 4 of Appendix D. Below we briefly show four optimal configurations found in grid search and compare them with our B-STaR.

The results indicate that while fixed hyperparameter values derived from grid search can lead to some performance improvements in online RFT, they remain inferior to the adaptive parameter adjustments offered by B-STaR. This highlights the necessity of dynamic parameter optimization in achieving superior performance. These results and more details are added as Appendix D of the updated PDF.

2. Second, part of the reason that previous experiments did not show significant fluctuations might be that we adopted a coarse adjustment scheme with the granularity of 0.1. To verify this, we further conduct finer-grained hyper-parameters search with the granularity of 0.05 for temperature and 0.01 for reward threshold. In Table 3 of Appendix C, we observe more dynamic changes throughout the training process.

Q3: From figure 4, it seems that after 1000 iterations, parameters like temperature and threshold remain constant. So essentially the dynamics freeze. Does this imply that original hyperparameter is simply not set to optimal?

A3: Please refer to A2 to Q2, which verifies the effectiveness and the necessity of the proposed B-STaR method.

Thanks for your time and encouraging review! We respond to your comments below.

Q1: Reliance on fixed reward models

A1: We use fixed reward models following most literature in this direction [1, 2, 3]. We agree that it is possible to update the reward models during training to obtain further gains, we leave exploration of this for future work.

[1] Hosseini et al. V-STaR: Training Verifiers for Self-Taught Reasoners, COLM 2024

[2] Sun et al. Easy-to-Hard Generalization: Scalable Alignment Beyond Human Supervision, 2024

[3] Wang et al. Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations, ACL 2024

Q2: complexity of configuration tuning

A2: We clarify that our method does not increase the complexity of configuration tuning, instead, it simplifies and reduces the efforts to tune hyperparameters such as temperature and reward threshold. This is because the proposed B-STaR automatically decides the critical hyperparameters through the introduced Query Effect metric, with negligible cost during training (Line 420-422). In contrast, previous methods need to tune these hyperparameters manually after going through multiple expensive training processes in a trial-and-error fashion.

Q3: limited generalization across task types

A3: Thank you for your suggestion. We follow IRPO [3] and add the evaluation results on the ARC Challenge dataset, a benchmark consisting of multiple-choice science questions designed to evaluate commonsense reasoning beyond mathematics and coding challenges, as shown below.

The results clearly indicate that B-STaR outperforms other baseline methods, achieving a significant improvement on ARC Challenge as well. This demonstrates that B-STaR is not only highly effective for mathematical reasoning and code generation tasks but is also capable of generalizing to other complex reasoning domains. We note that the our current evaluation on math, coding, and commonsense reasoning (for ARC) tasks aligns with or is already broader than most previous works on self-improving [1, 2, 3]. These added results and more details are also included in Section 4 and Figure 1 (d) of the paper.

[1] Avi et al. Beyond human data: Scaling self-training for problem-solving with language models. TMLR, 2024

[2] Zhang et al. Small Language Models Need Strong Verifiers to Self-Correct Reasoning. COLM 2024.

[3] Pang et al. Iterative Reasoning Preference Optimization. NeurIPS 2024.

Q4: potentially saturate in long iterative training

A4: This is a good point. We agree that with extended iterative training, model performance will eventually saturate due to many factors such as the training dynamics, the model's intrinsic capabilities, the data coverage, and the quality of the reward function. While overcoming all these fundamental limitations remain challenging, our proposed approach, B-STAR, specifically focuses on the training dynamics aspect and achieves more sustained growth and higher performance compared to baseline methods. Achieving scalable and sustainable improvements over long iterative training requires to address many other challenges, which we leave as future work to explore.