Towards Learning to Reason at Pre-Training Scale

Authors: Thomas Foster, Eltayeb Ahmed, Jonathan Cook, Shalev Lifshitz, Tim Rocktäschel, Jakob Nicolaus Foerster
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, self-improvement, reasoning

Abstract:
Prompting a Large Language Model (LLM) to output Chain-of-Thought (CoT) reasoning improves performance on complex problem-solving tasks. Moreover, several popular approaches exist to "self-improve" the CoT reasoning abilities of LLMs on tasks where supervised (question, answer) datasets are already available. An emerging line of work explores whether self-improvement is possible without these supervised datasets, instead utilizing the same large, unstructured text corpora as used during pre-training. This would overcome the data availability bottleneck present in current self-improvement methods, and open the door towards compute-only scaling of language model reasoning ability. We investigate a fundamental question in this line of work: What constitutes a suitable reward function for learning to reason during general language model pretraining? We outline the desirable qualities of such a reward function and empirically demonstrate how different functions affect what reasoning is learnt and where reasoning is rewarded. Using these insights, we introduce a novel reward function called Reasoning Advantage (RA) that facilitates self-improving CoT reasoning on free-form question-answering (QA) data, where answers are unstructured and difficult to verify. We also perform an exploratory experiment optimizing RA on general unstructured text using offline RL, and our analysis indicates that future work should investigate methods for generating a more diverse set of CoTs.

Decision: reject
Reviewer scores: 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1729439660507

thread_type:
official_review

<thread_object>
creator: reviewer_mbau

<content>
summary:
The paper explores effective rewards that could be applied during LLM pretraining. Especially, the paper explores various reward functions based on what reasoning is learnt and where reasoning is rewarded. Based on the findings, the paper suggests, RA (Reasoning advantage) which facilitates self-improving CoT reasoning on free- form question-answering (QA) data.

strengths:
- The paper provides useful insights for designing rewards for language model training.
- The authors explores the effectiveness of RA on multiple experimental settings.

weaknesses:
- Unlike the motivation of the paper, the proposed method, RA, is not effective for pre-training scale, questioning the scalability of the proposed method. 
- 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?
- The paper only uses a single backbone model to show the effect of the proposed method.

questions:
- How effective is RA compared to another baseline model which is directly trained to predict the final answer without training to generate CoT?
- How much additional overhead occurs for applying RA during pre-training (Section 6)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730347771324

thread_type:
official_review

<thread_object>
creator: reviewer_zh4y

<content>
summary:
This paper explores a method for self-improving CoT reasoning in LLMs without relying on curated datasets. By leveraging reinforcement learning on general pre-training data, the authors aim to enhance models’ reasoning abilities across diverse tasks. They introduce a new reward function, Reasoning Advantage (RA), which better identifies effective reasoning, and demonstrate its impact on open-ended question-answering tasks. The paper highlights RA’s potential but also suggests that more advanced optimization methods are needed for scalable CoT improvements in broader, less structured contexts.

strengths:
1. This paper addresses an important issue: achieving self-improvement of CoT reasoning during the pre-training phase. This approach has significant potential to help overcome the data bottleneck in LLMs.

2. The paper explores several types of reward functions and establishes criteria for an effective reward function, which is valuable and insightful for future research in this area.

weaknesses:
1. 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.

2. 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. Typical pre-training datasets often include substantial noise, such as HTML elements, which is a key challenge in achieving self-improvement CoT during the pre-training phase.

3. The paper lacks discussion on relevant work in reasoning enhancement during the pre-training phase, such as https://arxiv.org/pdf/2404.07965.

4. The experiments are insufficiently comprehensive, as they are conducted on only one model and one dataset. Testing with models of different parameter sizes within the same series or different architectures could help demonstrate the generalizability of RA.

5. The presentation of the paper could be improved. Some key findings should be in the main body rather than the Appendix, such as Appendix D and the definition of RA in Appendix A. Essential parameters, like the type of LLM used and inference hyperparameters, should also be included in the main text.

    Minor:
    - Punctuation should be added at the end of each equation.
    - Some quotation marks are unmatched, such as in line 265 and line 349.
    - Figure 1 appears somewhat rudimentary.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730600687551

thread_type:
official_review

<thread_object>
creator: reviewer_qddb

<content>
summary:
The paper titled explores the potential for self-improvement in large language models' ability to perform CoT reasoning without the need for supervised datasets. The authors frame this as a reinforcement learning problem where an LLM generates a CoT to predict subsequent tokens in a text corpus, receiving a reward based on the effectiveness of the CoT in predicting the next tokens. Their approach explores generating CoTs for next-token prediction in unstructured data, aiming to improve general-purpose reasoning abilities.

strengths:
1. The paper presents a novel approach to improving CoT reasoning in LLMs, exploring reinforcement learning as a framework for unsupervised self-improvement. The introduction of RA offers an innovative solution to the reward function challenge.
2. This work addresses a crucial challenge in LLM development—achieving autonomous improvement in reasoning without reliance on human-generated data. If successful, this approach could significantly reduce reliance on expensive, curated datasets and enable more scalable reasoning improvement across diverse domains.

weaknesses:
1. 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.
2. 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.
3. 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.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715529040

thread_type:
official_review

<thread_object>
creator: reviewer_kyrv

<content>
summary:
This paper explores how to enable large language models (LLMs) to self-improve their Chain-of-Thought (CoT) reasoning abilities using general pre-training data rather than supervised datasets. The authors investigate what makes a good reward function for learning reasoning during language modeling, examining how different reward functions affect both what reasoning is rewarded and where reasoning is applied. They introduce a novel "Reasoning Advantage (RA)" reward function that combines clipping and normalization techniques, and demonstrate its effectiveness on a new free-form question-answering dataset called MMLU-FREE-FORM, showing improved transfer to math reasoning tasks.

strengths:
The systematic analysis of reward functions and their properties is thorough and well-motivated. The introduction of the RA reward function addresses key limitations of existing approaches, particularly in distinguishing good reasoning from random text and identifying appropriate contexts for reasoning. 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.

weaknesses:
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. Additionally, 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.

questions:
Have you explored whether the effectiveness of Reasoning Advantage (RA) varies across different types of reasoning tasks beyond math and standard QA?
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?
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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732056140344

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response From Authors to Reviewers

comment:
Thanks to all the reviewers for your time and effort during the review process. We appreciate that you found our work insightful, and we’re glad that there is excitement about our progress towards self-improving CoT reasoning. Your thoughtful reviews have helped us dramatically improve the clarity and rigour of our submission.

We have responded to each reviewer individually, and also include a general response to all reviewers here. If you find our answers responsive to your concerns, we would be grateful if you considered increasing your score, and if you have additional questions, we’re happy to engage further. 


>Additional experiments.

Based on the reviewers’ questions and comments, we have added further experiments to strengthen our results. In response to reviewers ZH4y and MbaU, we have added results going beyond a single model to also include Llama 3.1 8B (see Table 5 in the appendix). In response to reviewer QdDb, we have included an additional ablation for different values of the clipping threshold (Figure 5 in the appendix). And in response to reviewer KyrV, we have added a qualitative analysis of the highest- and lowest-scoring CoTs in Section 6.1.

>Clarification of contributions.

We have updated the submission pdf to clarify the contributions of our work (main updates in red text). A brief summary is provided below:

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 task: self-improving CoT reasoning on unstructured, pretraining-scale text. To be clear, our work does not solve this grand challenge. 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. There is still more work to be done in order to solve the full unstructured pretraining setting, and we present an exploratory experiment that provides 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 manuscript (especially the Introduction, Section 5, Section 6, and Conclusion) with more targeted and clear descriptions of our contributions. 

Note: The main updates to the manuscript which clarify our contributions are made with *red* text. The blue text refers to changes made in response to specific reviewer questions, and we reference them in the individual responses below.

>Updated appendix.

We have also updated the appendix to be more clear, and have reworked the previous Appendix D into the main paper as the newly added Section 6.2 (per the request from reviewer ZH4y). To make space for these changes we have:
- Converted the former Figure 3 bar plot into a table (Table 2). Notice that full results with confidence bounds can still be found in Appendix B.1.
- Moved the former Figure 1 to the Appendix (new Figure 7).
- Moved the former Figure 5 to the Appendix (new Figure 8).

>Clarified hyperparameters.

We have made the specific details about model architecture and inference hyperparameters (i.e., temperature) more clear in the manuscript. These changes are made in blue text.


We again thank the reviewers for their engagement and we appreciate all the suggestions that we believe will make the paper significantly stronger!
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces the Reasoning Advantage (RA) reward function to improve CoT reasoning in LLMs using unstructured, pretraining-scale text instead of curated datasets. RA demonstrates improved performance on MMLU-FREE-FORM and transferability to GSM8K math tasks, while also providing a detailed analysis of reward functions and their properties. The work addresses a critical challenge in LLM reasoning, introducing MMLU-FREE-FORM as an intermediate benchmark that bridges curated QA datasets and unstructured text. The experiments highlight RA’s ability to improve reasoning performance and generalization to diverse tasks. However, some reviewers felt the contributions were incremental, with RA being a slight modification of token loss and MMLU-FREE-FORM a minor adaptation of MMLU. They also noted the limited scope of models and datasets tested, along with minimal improvement on OpenWebMath, raising concerns about scalability to pretraining-scale data. While the approach shows promise in intermediate settings, the lack of comprehensive evaluations and broader generalizability testing would limit its impact. Future work might expand experiments to include larger and more diverse datasets (e.g., BBH, CSQA) and evaluate performance across multiple model sizes to validate scalability and robustness.