BRiTE: Bootstrapping Reinforced Thinking Process to Enhance Language Model Reasoning

Comments on "Experimental Designs Or Analyses": Authors test their approach on math and code reasoning tasks. It would be interesting to see how this approach will work on other types of reasoning tasks (for ex, some popular benchmarks include MMLU, GPQA, BIG-Bench Hard, etc).

Response: We perform BRiTE on the model Qwen/Qwen2.5-7B with a mixed dataset (the size is around 40K) of RUC-AIBOX/STILL-3-Preview-RL-Data, MATH, and historical AIME problems (excluding AIME 2024). We evaluate models trained using RS baseline and BRiTE on a diverse set of reasoning benchmarks, including Math-500, Minerva Math, Olympiad Bench, AIME24, AMC23, and GPQA Diamond. The pass@1 accuracy (64 sampling times for AIME24 and AMC23 and 8 sampling times for the others) for each benchmark is reported in the table below.

Evaluation results show that BRiTE can surpass the RS baseline with a stable margin on these bechmarks, which demonstrates the power of BRiTE in the face of large datasets (40k data) and boosts the performance of the state-of-art model such as Qwen/Qwen2.5-7B.

Weakness on Training Cost: Proposed approach for BRiTE-SFT involves two stages of training, while traditional SFT methods authors compare with involve 2 stage of training. Authors to not provide any theoretical or actual cost estimates required to train with BRiTE. Given this process is two-staged, I assume it would use more GPU days and - depending on how much more - might not be directly comparable with direct SFT. Please include theoretical or actual cost estimates required to train with BRiTE-SFT, and add comparison with traditional SFT you used for Table 1.

Response: Thank you for your question. We will discuss data efficiency and computational efficiency as follows:

• Data Efficiency: (i) We want to emphasize that our algorithm only requires data without human annotation, making it more efficient than supervised fine-tuning (SFT) that relies on annotated data, which can be harder to collect; (ii) compared to other methods, such as rejection sampling, which also do not require data for the thinking process, our algorithm demonstrates greater efficiency since we have controlled for the same amount of data in all our experiments.
• Computational Efficiency: As mentioned earlier, our algorithm does not need data with a thinking process, so it is not fair to compare our algorithm (or other methods like rejection sampling) with SFT, which relies on human-annotated, step-by-step thinking data. Moreover, both our algorithm and rejection sampling (RS) share the same goal: generating a high-quality thinking process (through RL or RS) and then fine-tuning LLMs based on that. The generation process is the main contributor to computational costs, a factor that is unavoidable for all these methods. Specifically, the first stage of our method (rollout + RL updating) and the RS filtering process (rollout + filtering) both require approximately 30 hours. This main cost is attributed to the rollout of the Qwen2.5-7B model, generating outputs for 40.6K data points with 8 samples per point. Since RL provides more concise/shorter responses, the overall computational time (rollout + RL updating) is comparable to the RS filtering process. Conversely, the training phase (SFT for RS or $\theta$-updating in BRiTE) is significantly less demanding, requiring only 3.3 hours for 9K data points from RS, or 9 hours for 24K data points from RL, for four epochs. our method achieves significantly better performance without substantially increasing computational costs.

Additional Question: How do you select the best and worse response based on the correctness of y? Given 30 CoT, k will have correct answer, and (30-k) incorrect. How do you select best of k and worst of (30-k)? If it is not mentioned in the paper, please add elaboration.

Response: We randomly select one correct answer $A^+$ from the $k$ available correct answers and one incorrect $A^-$ answer from the $(30-k)$ incorrect answers. The pairs $\{(A^+, A^-)\}$ are then used to create a pair for DPO training. We will add this clarification in the revision.

Comments on "Claims and Evidence": In reality, datasets like MATH have very careless annotations, so the performance of SFT may not be as effective as LLMs generating their own data. This is a point that has been widely discussed in many papers and is not necessarily an interesting finding.

Response: Thank you for bringing this to our attention. If a widely used dataset like MATH still struggles with careless annotations, it strengthens our motivation for research, as obtaining data that accurately reflects a high-quality thinking process is challenging. Additionally, our algorithms demonstrate their effectiveness by comparing with other algorithms, showcasing their ability to learn without relying on high-quality step-by-step data.

Comments on "Methods And Evaluation Criteria": It would be better to include a broader range of tasks to evaluate its general applicability.

Response: We perform BRiTE on the model Qwen/Qwen2.5-7B with a mixed dataset (the size is around 40K) of RUC-AIBOX/STILL-3-Preview-RL-Data, MATH, and historical AIME problems (excluding AIME 2024). We evaluate models trained using RS baseline and BRiTE on a diverse set of reasoning benchmarks, including Math-500, Minerva Math, Olympiad Bench, AIME24, AMC23, and GPQA Diamond. The pass@1 accuracy (64 sampling times for AIME24 and AMC23 and 8 sampling times for the others) for each benchmark is reported in the table below.

Evaluation results show that BRiTE can surpass the RS baseline with a stable margin on these bechmarks, which demonstrates the power of BRiTE in the face of large datasets (40k data) and boosts the performance of the state-of-art model such as Qwen/Qwen2.5-7B.

Weaknesses: I did not find enough novel approaches that significantly contribute to improving LLM reasoning abilities. Instead, the paper primarily re-expresses previous methods within a unified framework.

Response: Thank you for recognizing that our proposed framework is general, encompassing many previous methods as special cases. We also want to highligh our four contributions:

1. Novel Framework and Algorithm: The unified framework itself, which innovatively utilizes a graphical model to incorporate both thought $z$ and the optimality $o$ in the explanation of LLM reasoning, is an important contribution to theoretically characterize LLM reasoning. Furthermore, the two-stage optimization algorithm BRiTE, derived from this framework, includes various RL paradigms as special cases and introduces new learning algorithms, demonstrating the framework's utility.
2. Unified Theoretical Guarantee: We provide a unified theoretical guarantee for BRiTE, thereby offering a theoretical foundation for a broad range of algorithms (for the first time) within this framework.
3. Superior Empirical Performance: Our method significantly outperforms previous methods on bootstrapping thinking process (please see the results in our paper and the updated results in the table below).
4. RL-Guided Reasoning Paradigm: Our work demonstrates that RL-guided reasoning paths are effective for generating high-quality thinking processes. This aligns with the concurrent DeepSeek-R1 project, which also employs RL training to generate SFT data.

Additional Question: Has anyone evaluated the quality of the generated data? Is there any intuition that explains why the new data is of higher quality?

Response: We evaluated the perplexity (the mean negative log probability of the model on the responses) of the base model (Gemma-2-9b-it) using the responses from both the human-annotated data in the corresponding dataset and RL-guided data (generated by model after $\psi$-update). The results demonstrate that the RL-guided data exhibits significantly lower perplexity, which means models are easier to learn the reasoning path from the RL-guided data.

Due to the page limit of the rebuttal, we could not provide a concrete comparison between the RL-guided response and the human-annotated response, which shows that the RL-guided reasoning path can provide more detailed and clear steps for thinking and is easier for the model to learn and derive the correct answer. If you are interested, we are pleased to provide this comparison in the new round of rebuttal.

Weakness: (1) Reasoning task diversity is limited. (2) The performance of BRiTE is not significantly better than other efficient methods.

Response: We perform BRiTE on the model Qwen/Qwen2.5-7B with a mixed dataset (the size is around 40K) of RUC-AIBOX/STILL-3-Preview-RL-Data, MATH, and historical AIME problems (excluding AIME 2024). We evaluate models trained using RS baseline and BRiTE on a diverse set of reasoning benchmarks, including Math-500, Minerva Math, Olympiad Bench, AIME24, AMC23, and GPQA Diamond. The pass@1 accuracy (64 sampling times for AIME24 and AMC23 and 8 sampling times for the others) for each benchmark is reported in the table below.

Evaluation results show that BRiTE can surpass the RS baseline with a stable margin on these bechmarks, which demonstrates the power of BRiTE in the face of large datasets (40k data) and boosts the performance of the state-of-art model such as Qwen/Qwen2.5-7B. Question 1: Could the authors discuss the algorithm in terms of efficiency?

Response: Thank you for your question. We will discuss data efficiency and computational efficiency as follows:

• Data Efficiency: (i) We want to emphasize that our algorithm only requires data without human annotation, making it more efficient than supervised fine-tuning (SFT) that relies on annotated data, which can be harder to collect; (ii) compared to other methods, such as rejection sampling, which also do not require data for the thinking process, our algorithm demonstrates greater efficiency since we have controlled for the same amount of data in all our experiments.
• Computational Efficiency: As mentioned earlier, our algorithm does not need data with a thinking process, so it is not fair to compare our algorithm (or other methods like rejection sampling) with SFT, which relies on human-annotated, step-by-step thinking data. Moreover, both our algorithm and rejection sampling (RS) share the same goal: generating a high-quality thinking process (through RL or RS) and then fine-tuning LLMs based on that. The generation process is the main contributor to computational costs, a factor that is unavoidable for all these methods. Specifically, the first stage of our method (rollout + RL updating) and the RS filtering process (rollout + filtering) both require approximately 30 hours. This main cost is attributed to the rollout of the Qwen2.5-7B model, generating outputs for 40.6K data points with 8 samples per point. Since RL provides more concise/shorter responses, the overall computational time (rollout + RL updating) is comparable to the RS filtering process. Conversely, the training phase (SFT for RS or $\theta$-updating in BRiTE) is significantly less demanding, requiring only 3.3 hours for 9K data points from RS, or 9 hours for 24K data points from RL, for four epochs. our method achieves significantly better performance without substantially increasing computational costs.

Question 2: It would be very helpful if the authors could explain the algorithm in steps. In my current understanding. The main contribution is to provide a new loss to optimize.

Response: Regarding the implemented algorithm, the main contribution is providing a new approach to generating high-quality latent thinking processes through RL objective with a reward function (as detailed in Section 3.4). Specifically, inspired by our theory, the learning algorithm guarantees provable convergence by iteratively updating equations (3.4) and (3.5). Equation (3.5) is relatively easier to implement, while the key challenge lies in solving equation (3.4). Guided by Proposition 3.7, solving the intractable distribution in (3.4) is equivalent to addressing an entropy-regularized Markov Decision Process (MDP) with the reward function $\log P(z, y, o\mid x)$. This can be further specified using the graphical model; for instance, if we let $P(o = 1 | z, y, x) = \exp(\beta \cdot R(x, z, y))$, we can define the reward function as $\log P(z, y | x) + \beta \cdot R(x, z, y)$ with a tuning parameter $\beta$. This can be attributed to each token through the autoregressive policy property, allowing us to assign the reward to the last token.

Beyond the algorithmic contributions, our main contribution includes a new unified framework, theoretical guarantee, superior performance, and the demonstration of RL-guided reasoning's potential. Due to space constraints, we refer you to our detailed response to Reviewer afMW.