Towards Learning to Reason at Pre-Training Scale

Thank you for your review and comments. We’re glad that the reviewer finds our work insightful for designing rewards for language model training. Please see below for responses to your comments and questions.

Unlike the motivation of the paper, the proposed method, RA, is not effective for pre-training scale, questioning the scalability of the proposed method.

Thanks for raising this point, we agree with the reviewer that our work does not solve the full unstructured, pretraining setting. We think it may be helpful to provide additional background, which we hope clarifies the contributions of our work:

As it becomes increasingly challenging and prohibitively expensive to curate large-scale (question, CoT, answer) datasets, the LLM reasoning community has begun focusing on a grand challenge: self-improving CoT reasoning on unstructured text at pretraining scale.

In this work, we introduce MMLU-FREE-FORM as a small step towards the grand challenge of truly unstructured text. As mentioned by Reviewer KyrV, it acts as “an intermediate benchmark between structured QA and general language modeling”. However, we show that even taking this small step renders current reward functions unusable. That is, standard reward functions cannot self-improve CoT reasoning even in this simplified free-form QA setting. Our work introduces RA to address this problem—it’s the only reward function which facilitates generalization when self-improving on MMLU-FREE-FORM. We also perform a comprehensive analysis of reward functions and how they affect what and where reasoning is rewarded. To our knowledge, our work is the first to provide this type of analysis on reward functions for self-improving CoT reasoning on unstructured text.

To be clear, our work does not solve the grand challenge mentioned. Our work demonstrates that current reward functions fail for even the small step of MMLU-FREE-FORM, introduces a novel reward function that mitigates this problem, and performs a comprehensive analysis with insights to help facilitate future research in this direction. We believe these to be major contributions to the literature, and have dramatically updated our Introduction, Section 5, Section 6, and Conclusion with more targeted descriptions of our contributions. The key points are updated with red text.

The paper measures the performance by using 'expected accuracy' metric, which makes comparison with other methods difficult. What is the absolute accuracy performance for Figure 4?

Thanks for the question. Expected accuracy is commonly used for evaluating language models on QA benchmarks, with many works simply referring to it as accuracy. For example, one of the most important recent works on self-improving LLM reasoning [1] reports accuracy, but it's only clear when going through their code that this is in fact expected accuracy.

Given a question and a reasoning trace, we compute expected accuracy as the probability of the correct answer. Another way to compute accuracy is by greedy decoding a single answer, or by sampling many answers and averaging the performance, which approaches expected accuracy given enough samples. Usually, expected accuracy is preferred because it represents the raw distribution learnt by the model, without depending on a specific decoding procedure. We have added additional clarification regarding this point near the end of Section 5.2 (blue text).

[1] Zelikman et al, Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking, 2024

The paper only uses a single backbone model to show the effect of the proposed method.

We thank the reviewer for raising this point, and we agree that running experiments with another model would strengthen the paper. Therefore, we have repeated the what and where experiments from Section 5.1 with another model, LLama 3.1 8B. We present the results in Appendix B Table 5 and find them to be consistent with our original results using Mistral 7B in Table 2.

How effective is RA compared to another baseline model which is directly trained to predict the final answer without training to generate CoT?

Thank you for pointing out that this important baseline was missing. We have included an updated Figure 7 in the latest version of the manuscript which includes a “no reasoning” baseline of training to predict the answer without any CoT. This new baseline performs worse than RA and many of the other baselines (Figure 7, right-side). In particular, the difference in performance is largest for reasoning style questions.

How much additional overhead occurs for applying RA during pre-training (Section 6)?

Our offline RL method has three main steps: (1) generate a large batch of CoTs, (2) self-insert them into the pretraining dataset, and (3) filter the ones with the highest scoring rewards and finetune the model. Thus, the main computational overhead occurs before training: generating the CoTs and using RA to score them. This can be efficiently parallelized, and thankfully, getting the score from RA only requires 2 forward passes (one with the CoT, and one for the “Empty CoT” baseline). During supervised finetuning, we just fit the dataset with the self-inserted CoTs as normal, which adds minimal overhead. So, the additional overhead is mostly before training, it can be massively parallelized (RA is parallelizable since it is loss-based), and it is relatively insignificant compared to the unavoidable cost of actually generating the CoTs.

We thank the reviewer for engaging deeply with the work. We’re glad that the reviewer found our investigations into effective reward functions for CoT reasoning both valuable and insightful.

We have updated the paper to clarify our contributions and address the reviewer’s concerns, with the main updates in red text. We address specific questions and comments below.

The technical contributions of the paper are relatively weak. The proposed MMLU-FREE-FORM is merely a simple adaptation of the original MMLU, and the introduced RA is only a minor modification based on token loss.

We agree that MMLU-FREE-FORM is not a substantial technical change from MMLU. However, our purpose for creating MMLU-FREE-FORM was not to create a radically new dataset, but to make the smallest possible change to MMLU that reveals the limitations of existing reward functions. It acts as an important middle-ground between improving CoT reasoning using curated (question, CoT, answer) datasets and the challenging, unsolved task of self-improving CoT reasoning on unstructured text. This is because MMLU-FREE-FORM does not allow for using exact-match accuracy as a reward metric (similar to unstructured pretraining text) and yet offers a higher density of clear opportunities for CoT reasoning compared to typical pre-training corpora, making it an ideal stepping-stone towards the ultimate goal of self-improving CoT reasoning on unstructured text. As mentioned by Reviewer KyrV, "the creation of MMLU-FREE-FORM as an intermediate benchmark between structured QA and general language modeling is clever and useful for the research community. The empirical results showing successful transfer learning to GSM8K math problems provide concrete validation of their approach." We have updated the Introduction and Section 5.2 to clarify the contribution of MMLU-FREE-FORM, with the main updates in red text.

We also agree with the reviewer that RA is a modification of token loss: by using clipping, subtracting the "empty CoT" baseline, and normalizing. As above, our goal was to make the smallest, simplest modification to the existing paradigm (standard loss) that has the potential to work for this setting. We believe that the contribution of RA is quite significant. It performs substantially better than token loss on the what to reward experiments (distinguishing effective CoT) and the where to reward experiments (picking out useful locations for producing CoT). Moreover, and possibly our most important result, only RA is able to facilitate generalization to the MMLU test set and zero-shot transfer to GSM8K when self-improving CoT reasoning on MMLU-FREE-FORM. In addition, we strongly believe that RA being based on token loss is a key advantage of this function. It requires only two forward passes, does not require an external strong model, is not limited to exact-match heuristics like accuracy-based functions (which fail on unstructured text), and allows the model to place weight over a distribution of valid answers. We have updated Section 4 to clarify this point, with the main updates in red text.

We hope that our explanations helped to clarify the contributions of MMLU-FREE-FORM and the RA reward function.

The paper somewhat overstates its contributions. The authors primarily demonstrate the positive impact of RA on MMLU-FREE-FORM, yet MMLU-FREE-FORM is derived from the structured MMLU dataset and cannot be regarded as a typical pre-training dataset. In fact, experiments on OpenWebMath show minimal improvement.

We thank the reviewer for their feedback on this point. Here, we aim to provide additional background, which we hope clarifies the important contributions of our work.

Curating large, challenging, and diverse (question, CoT, answer) datasets for improving LLM reasoning has become exceptionally expensive (millions of dollars) and very challenging (requiring thousands of expert hours). With this as motivation, the LLM reasoning community has recently begun focusing on a grand challenge: self-improving CoT reasoning on unstructured text at pretraining scale.

We agree with the reviewer that MMLU-FREE-FORM is not a typical pre-training dataset. It is a small step towards the grand challenge of truly unstructured text. However, we show that even taking this small step renders current reward functions unusable. That is, standard reward functions cannot self-improve CoT reasoning even in this simplified free-form QA setting. Our work introduces RA to address this problem—it’s the only reward function which facilitates generalization when self-improving on MMLU-FREE-FORM. We also perform a comprehensive analysis of reward functions and how they affect what and where reasoning is rewarded. To our knowledge, our work is the first to provide this type of analysis on reward functions for self-improving CoT reasoning on unstructured text.

[Continued in 2nd Response]

Thank you for your review and comments. We’re glad that you think our work addresses a crucial challenge in LLM development and that you see the introduction of RA as an innovative solution to the challenge of designing reward functions for reasoning. Please see below for responses to your comments and questions.

Some aspects of the reinforcement learning formulation could benefit from additional clarity, specifically regarding the choice of reward clipping values and the normalization strategies within RA. Additional explanation of these parameters and their impact on performance would make the approach more accessible.

Thanks for the great suggestions. We have updated the parts of Section 4 discussing clipping and normalization to be more clear. We have also added a detailed discussion of their impact on performance to the end of Section 5.1 (the two big groups of blue text), including an additional ablation for different values of the clipping threshold. Moreover, Appendix B.1 contains tables which show full results for additional combinations of clipping, baseline, and normalization.

Briefly summarizing how each design choice relevant to RA can be interpreted:

• Clipping value: the minimum value at which suffix token log-probabilities are clamped. This prevents any given token from having an outsized loss contribution. We have run an additional ablation on a range of clipping values to demonstrate its impact, and include the results in Appendix B.1 (Figure 5).

• Baseline: in RA, we compute the token log-probabilities for the suffix given the prefix and CoT, but subtract the "empty CoT" baseline, which is the token log-probabilities without conditioning on any CoTs (only the prefix). This ensures we are optimizing for CoTs that improve the suffix loss relative to not producing a CoT.

• Normalization: normalizing by the "empty CoT" baseline re-scales the reward to ensure that we don't provide high reward values for CoTs with trivial-to-predict suffixes (ie, when an "empty CoT" predicts the suffix well).

The experiments focus primarily on a limited scope of problems (e.g., MMLU and OpenWebMath). The model’s performance on broader tasks, such as Tool learning or agent problem-solving scenarios, would offer stronger evidence of the approach’s generalizability.

We thank the reviewer for their suggestion, it would indeed be interesting to evaluate the model’s tool learning and agentic behaviour. However, in keeping with much of the LLM reasoning literature, we chose to specifically focus our evaluations on a range of QA-style reasoning problems. One of the most important recent works on self-improving LLM reasoning, Quiet-STaR [1], uses two QA-style reasoning benchmarks to evaluate their method: GSM8K and CSQA. Similarly, we evaluate our self-improvement method from Section 5.2 on two QA-style reasoning benchmarks: MMLU and GSM8K.

By relying on the log-likelihood to evaluate the quality of intermediate reasoning (Chain-of-Thought) solely based on the model's ability to predict the following tokens, there is a risk that the model may overly focus on matching specific token patterns in the training data rather than developing generalized reasoning capabilities.

We view the strong zero-shot transfer performance to GSM8K after optimising for RA on MMLU-FREE-FORM as compelling evidence that the model learns generalisable reasoning — beyond just matching specific token patterns in the data. Moreover, we strongly believe that RA being based on token loss is actually a key advantage. Since it is loss-based, it requires only one forward pass, does not require an external strong model, is not limited to exact-match heuristics like accuracy-based functions (which fail on unstructured text), and allows the model to place weight over a distribution of valid answers.

It might also be worth mentioning that since we are starting from a pretrained LLM, much of the gains from attempting to match specific token patterns have already been exhausted. That is, at the beginning of standard LLM pretraining, most of the loss reduction is achieved by fitting specific token patterns like spelling and grammar rules, but by further reducing loss, models begin to learn higher-order skills such as CoT reasoning. We initialize our weights to a pretrained LLM, so the risk of overly focusing on specific token patterns is limited.

Thanks again for your review and helpful suggestions. We are excited to see the insights from our work help progress the field towards the grand challenge of self-improving CoT reasoning on unstructured text at the pretraining scale. If you feel that we have adequately addressed your concerns, we would appreciate your consideration to increase our score. And if you have any additional questions or needs for clarification, please don’t hesitate to ask.

[1] Zelikman et al, Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking, 2024

We thank the reviewer for the thoughtful review, and for recognizing the value in using MMLU-FREE-FORM as an intermediate benchmark between curated QA data and the unstructured text setting. We address your questions and comments below.

The paper's primary limitation appears in the scaling to general pre-training data, where the offline reinforcement learning approach that worked well on MMLU-FREE-FORM struggles to escape local optima of conservative reasoning. While the authors acknowledge this limitation and suggest future research directions, the paper doesn't fully solve the challenge of self-improving reasoning at pre-training scale.

We agree with the reviewer that our work does not solve the grand challenge of self-improving CoT reasoning on unstructured, pretraining-scale text. Our primary contributions are showing that standard reward functions fail even in the intermediate MMLU-FREE-FORM setting, introducing a novel reward function to solve this issue, and performing a comprehensive analysis of reward functions and how they affect what and where reasoning is rewarded. To our knowledge, our work is the first to provide this type of analysis on reward functions for self-improving CoT reasoning on unstructured text.

We also perform an exploratory experiment on the full unstructured setting and provide key insights to help facilitate future research in this direction. We believe these to be major contributions to the literature, and have dramatically updated our Introduction, Section 5, Section 6, and Conclusion with more targeted descriptions of our contributions. The key points are updated with red text.

While the authors demonstrate improved performance on mathematical reasoning tasks, there could be more exploration of how well their approach generalizes to other types of reasoning beyond mathematics.

We thank the reviewer for this suggestion. While GSM8K focuses purely on mathematics, MMLU contains questions which span various fields, including mathematics, sciences, law, etc. On MMLU, we actually find that our method leads to large gains on questions across a wide span of subjects (biology, physics, accounting, law, computer science, etc.) that involve quantitative reasoning—going beyond just mathematics. We have updated Appendix B.2 to better explain this point (blue highlighted text). Thanks again for the suggestion.

Have you explored whether the effectiveness of Reasoning Advantage (RA) varies across different types of reasoning tasks beyond math and standard QA?

In keeping with much of the LLM reasoning literature, we chose to specifically focus our evaluations on a range of QA-style reasoning problems. However, we do show positive results that go beyond just mathematics. As mentioned above, Appendix B.2 Figure 6 shows that RA significantly improves reasoning performance on “reasoning style questions” in the MMLU test set. These questions span a wide range of subjects beyond math, including physics, biology, accounting, law, computer science, etc.

In Section 5.2, you show that optimizing for RA leads to a 7% improvement on GSM8K. Could you provide more analysis of what specifically improved in the model's reasoning capabilities? Are there particular types of math problems where the improvement was more pronounced?

Great question. We observe that the GSM8K accuracy goes up due to fewer logical/arithmetic errors in the generated CoTs, but we don’t observe a single predominant qualitative change. However, we do notice something interesting regarding performance on the MMLU test set. As mentioned above, MMLU sees an improvement in questions across a wide range of subjects that require quantitative reasoning (biology, physics, accounting, computer science, etc.). Moreover, the improvement is far larger for these questions than for those that require recall (see Figure 6 in Appendix B.2). Thinking about it, this result makes a lot of sense, since CoT reasoning probably doesn't help as much when trying to recall a fact. However, for quantitative reasoning and problem-solving tasks, additional reasoning can clearly be of benefit. We have added a discussion of this to Appendix B.2 (blue text). Thank you for the suggestion!

The paper mentions that only 0.01% of generated CoTs achieve a reward above 0.2 on OpenWebMath. Have you analyzed these high-scoring CoTs to understand what makes them successful? This analysis could inform better prompting strategies.

Thanks for asking this great question. Upon manual inspection, many of the CoTs that passed the filtering threshold exhibited the conservative strategy described in the paper: they simply summarize past information from the context. This explains why the model learned to be overly conservative. However, these overly conservative CoTs which made it past the RA threshold were still superior to those that did not pass the threshold (the ones that did not pass the threshold mainly contained incorrect reasoning that predicted the subsequent tokens incorrectly). This indicates that RA actually succeeded at its job of identifying the best reasoning from the generated batch of CoTs, and that the main issue indeed lies with the lack of diversity in the generated CoTs. We mention a few potential ways to generate more diverse CoTs in the paper, but we agree that it would also be worth exploring different prompting strategies (we used a single system prompt to generate these CoTs, and did not spend much time on prompt engineering).

We have added this detailed discussion to Section 6.1 (the first blue block text, and then the rest discusses ways to increase diversity), and we think it dramatically improves the section. Thanks again for the great suggestion.