Lean-STaR: Learning to Interleave Thinking and Proving

Authors: Haohan Lin, Zhiqing Sun, Sean Welleck, Yiming Yang
Primary area: reinforcement learning
Keywords: Automated Theorem Proving, Chain-of-Thought, Reinforcement Learning, Reasoning

Abstract:
Traditional language model-based theorem proving assumes that by training on a sufficient amount of formal proof data, a model will learn to prove theorems. Our key observation is that a wealth of informal information that is not present in formal proofs can be useful for learning to prove theorems. For instance, humans think through steps of a proof, but this thought process is not visible in the resulting code. We present Lean-STaR, a framework for training language models to produce informal thoughts prior to each step of a proof, thereby boosting the model's theorem-proving capabilities. Lean-STaR uses retrospective ground-truth tactics to generate synthetic thoughts for training the language model. At inference time, the trained model directly generates the thoughts prior to the prediction of the tactics in each proof step. Building on the self-taught reasoner framework, we then apply expert iteration to further fine-tune the model on the correct proofs it samples and verifies using the Lean solver. Lean-STaR significantly outperform base models (43.4% → 46.3%, Pass@64). We also analyze the impact of the augmented thoughts on various aspects of the theorem proving process, providing insights into their effectiveness.

Decision: accept (spotlight)
Reviewer scores: 8, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1729214348396

thread_type:
official_review

<thread_object>
creator: reviewer_meco

<content>
summary:
Generally, this paper takes motivation from a hypothesis that informal thought process can be valuable for LLMs to learn formal theorem proving. This work presents the following main contributions:

1. The first thought (informal language) augmented for theorem proving in Lean (formal math).
2. Lean-STaR: a framework which combines tactic prediction, thought augmentation with synthetic data, and two iterations of expert iteration.

strengths:
1. Prior to this work, people have been arguing how informal thought process can potentially benefit formal theorem proving, yet no one has taken the effort to explore how to interleave those two. People either used a standard tactic prediction model that operates only on formal code, or use a general-purpose model that only chats in an informal manner. While the methods (e.g. STaR, expert iteration) are not new themselves, the idea of combining the two and curating a resulting dataset sounds to me as a nice step. It is nice to have this work verify that known methods apply well to this domain of AI for math, especially given the many differences between this and other application areas.

2. Theorem proving in Lean as a field suffers from a lack of data. Works like this with a large dataset size (50,000 from mathlib + 50,000 synthetic data) can greatly benefit the field.

3. It seems that the framework is quite generally usable, and multiple models can be plugged in. This means the method proposed has a great potential to generalize. Similarly, its potential to be used alongside other synthetic data generation methods such as autoformalization without any approach-wise conflict sounds promising too.

weaknesses:
1. The thought data generation in this paper is by few-shot prompting a general-purpose model (GPT-4). Data quality is usually a concern of synthetic data generation. People in the field either try to design data filtering mechanisms or recruit domain experts to sample and check. However, it seems that this work did not apply such techniques or analyses after synthetically generating thought data.

2. The framework is composed of multiple methods. This paper could benefit a lot from having more thorough ablations. For instance, the authors can do an ablation on the thought augmentation, which is connected to their main hypothesis and contribution. What would the performance if all other settings (expert iteration, tactic prediction) are kept as-is but no thought augmentation is added? The authors can perform similar analyses on other factors they feel important too.

3. Potential data contamination seems insufficiently discussed. Is it possible that InternLM2-Math-base-7b may have been exposed to parts of the test data? I think some analysis would be needed here.

questions:
1. Related to the first weakness, the heavy use of GPT-4 to generate thought data (which are usually not short) can add a lot of computational and financial burdens. I would be curious to see a direct analyses of such overheads.

2. The authors only performed two iterations of expert iteration. Both iterations lead to performance gains, and the first gains more than the second. Such findings all make sense and align with my expectation of applying expert iteration. However the author stopped right after only 2 iterations. I was wondering if a bottleneck would be hit very soon, or would the framework continue to benefit from more iterations?

3. Minor suggestions: A few works that appear related to this paper, which the authors may consider citing/discussing.

* [1] Multilingual Mathematical Autoformalization. Albert Q. Jiang, Wenda Li, Mateja Jamnik. arXiv 2311.03755. 2023.

[1] is a work in autoformalization, which the authors did discuss in their related works. One thing that especially connects [1] to this paper is that [1] constructs <informal, formal> statement pairs in a backward manner, leveraging the fact that auto-informalization is a much easier task than autoformalization for LLMs. I feel this aligns in spirit with the retrospective rationale generation in this paper, which generates informal thought data from formal data using GPT-4.

* [2] Towards Large Language Models as Copilots for Theorem Proving in Lean. Peiyang Song, Kaiyu Yang, Anima Anandkumar. arXiv 2404.12534.

[2] is an open-source tool work in using LLMs as Copilots in Lean. The authors did discuss one related work LLMStep as practical tools in the field and may want to include [2] for completion. To my knowledge, they are so far the only two works that try to set up a ML framework easily usable in Lean.

* Finally (and a very minor point), I think when discussing related works in this field, it may be beneficial to include some ML + ITP works in other proof assistants. Of course one does not need to expand in detail about them, but would be great to have a broader view.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730502121851

thread_type:
official_review

<thread_object>
creator: reviewer_qmzb

<content>
summary:
Lean-STaR uses the STaR (Self-Taught Reasoner) framework to mix step-by-step reasoning ("thoughts") with formal proof steps. It combines this with expert iteration, where the model learns from successful proofs to improve further. This approach allows Lean-STaR to reach a pass@32 rate of 45.4% on miniF2F, indicating that adding informal CoT reasoning can enhance formal automated proving capabilities.

strengths:
**1. Simple yet very effective approach:**
The authors use retrospective rationale generation to come up with the thoughts given the state and ground truth action. The idea to train another model Lean-CoT for generation of thought-augmented tactics, is simple but works very well in this setting.

**2. Use of expert iteration while informal thought generation:**
While expert iteration has been used in previous works to improve formal theorem proving, the authors demonstrate that the same holds true for models that generate thought-augmented tactics.

**3. Search based on Sampling:**
The authors used sampling-based search instead of BFS which is new and seems more effective for this setting. The idea that negative log-likelihood doesn't make proof search effective and a simple sampling technique can perform better is very elegant. The authors avoid the overkill of training a value function model which predicts the quality of generated tactics for choosing the search direction during the proof search. I'm surprised to see the effectiveness of a simple sampling-based search and would like to see a broader comparison with BFS in a formal proof-search setting.

weaknesses:
**1. Data Leakage:**
While the authors talk about the limitations of their approach, they don't discuss potential data leakage due to the use of GPT-4 for generating thought annotations. It can very well happen that miniF2F-related proof-steps are leaked in the thought generation process. Also, miniF2F has been used for a while for evaluating approaches for formal theorem proving, there is a chance that the base pre-trained model itself is trained on miniF2F repository. The worst part is that proofs are already checked in miniF2F repository, I would like to understand how much overlap exists between proofs generated by LeanStar and the proofs already checked in the miniF2F repository. Even the LeanDojo dataset, has existed in the public domain for quite some time, the pre-trained models like InternLM might as well be trained on this mix. I would like to see a broad study about how data leakage "impacts"/"doesn't impact" the overall results. I can also recommend some new benchmarks like Putnam Bench (https://arxiv.org/abs/2407.11214), which don't have proofs checked in online, for further evaluation.

It is a well-written paper, I will be more than happy to increase my score if the authors address some of my concerns.

questions:
My main concern is the impact of data leakage (as highlighted in the weakness section) which needs to be studied further. The paper is well-written, I don't have any other questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730505748027

thread_type:
official_review

<thread_object>
creator: reviewer_kzhz

<content>
summary:
This paper introduces a framework called Lean-STaR for training LLMs that incorporates both informal reasoning in the form of natural language (called thoughts) with formal reasoning in the form of Lean code to improve a LLMs theorem-proving capabilities. More concretely, the paper proposes to construct a tactic predictor as p(tactic | proof context) = Expectation_{thought ~ p(thought | proof context)} p(tactic | thought, proof context). This contrasts with direct tactic prediction which directly construct p(tactic | proof context) as is done in previous works. To train the tactic predictor, the authors propose the following pipeline. First, they introduce retrospective rationale generation to create a "thought-augmented" dataset (called CoT dataset) which consists of tuples of (tactic, proof context, thought) where each (tactic, proof context) are taken from an existing Lean corpus (Mathlib) and each thought is generated by GPT-4 with a specially-crafted prompt. Second, they take InternLM2-Math-base-7b, a LLM that is pre-trained on Mathlib, and perform supervised fine-tuning on the CoT dataset to create a model called Lean-CoT. Third, the authors use expert iteration to construct a Lean-STaR model by fine-tuning the Lean-CoT model on a dataset of proof trajectories that it generates. The authors show that Lean-STaR outperforms direct tactic prediction on minif2f and LeanDojo benchmarks.

strengths:
- The idea for constructing and training a tactic predictor p(tactic | proof context) = Expectation_{thought ~ p(thought | proof context)} p(tactic | thought, proof context) by introducing a latent variable in the form of thoughts is novel. 
- There are numerous experiments.

weaknesses:
The presentation of the paper could be improved and I have one major concern about soundness of the results (see questions for the authors).

- In general, I found the abbreviations and terminology of the paper to be confusing. For example, the text refers to the dataset as a thought-augmented dataset but it is abbreviated as CoT dataset. A model for direct tactic prediction is abbreviated as a SFT model. Lean-STaR is both the name of the framework and a model that is produced by the framework.
- Automatic theorem proving and interactive theorem proving are used interchangeably throughout the text even though they are not the same concept.
- The paper seems to argue that GPT-4 does not perform well in generating "correct rationale through few-shot prompting" (lines 196) and then proceeds to use GPT-4 to generate rationale for the CoT dataset. 
- Figure 3 is not very clear for me. For instance, I would expect the Base Dataset to be fed into GPT-4 as opposed to directly into the CoT dataset. The Expert Iteration loop is also not clear. I would expect Expert Iteration from Lean-CoT to feed into Lean-STaR, and a separate arrow kind/color to indicate the usage of models in inference mode to generate datasets (also for GPT-4).
- Table 7 -> Table 1.

questions:
- I believe that the learned distribution is p(tactic, thought | proof context) as opposed to  p(tactic | proof context) = Expectation_{thought ~ p(thought | proof context)} p(tactic | thought, proof context) as presented in the paper. In particular, I cannot find at any point in the paper how the expectation is handled (e.g., with a Monte Carlo estimate). Moreover, the dataset associates a single thought with a single (tactic, proof context) tuple so that we do not have multiple thoughts for a given proof context. Finally, my understanding of Figure 2 is that it is consistent with p(tactic, thought | proof context) since a single thought is generated and then used to generate a tactic as opposed to multiple thoughts that are marginalized. If so, it's not really a fair comparison with direct tactic prediction since you are basically comparing p(tactic | thought, proof context) vs. p(tactic | proof context) which means your approach gets strictly more information. In other words, the thought is not really treated as a latent variable and what is presented in the paper is not what is actually done. Looking forward to clarifying this.
- Is there an intuitive reason why sampling works better than best first-search for your method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730545860647

thread_type:
official_review

<thread_object>
creator: reviewer_twp5

<content>
summary:
This paper introduces a novel approach to neural theorem proving (NTP). One previous drawback of step-by-step approaches to NTP is the lack of Chain-of-Thought (CoT) or SketchPad techniques before generating each tactic. These techniques have proven to be very useful in the domain of mathematical word problems (MWP). To address this, the paper proposes using GPT-4 to generate an initial CoT before each tactic on the Mathlib dataset, and subsequently bootstrapping the model with expert iteration. The experimental results show that adding CoT before generating the tactic leads to performance gains.

strengths:
- The problem the paper aims to address is long-standing, and the paper provides strong empirical results demonstrating that the Chain-of-Thought (CoT) reasoning before each tactic application is useful.
- The newly proposed sampling method is novel and, based on its performance, appears to be a promising approach compared to Best-First Search.
- The proposed CoT-augmented dataset should be useful for future research in neural theorem proving.
- The paper is well-written and easy to follow.

weaknesses:
- One issue with the generated COT is the absence of a verification mechanism or alternative methods to ensure its quality. The paper does not provide a detailed evaluation of these generated COTs; it only presents the final pass rate information. I believe that having a robust evaluation is crucial for improving the performance of this method. Moreover, in prompting GPT-4 to generate COTs, the authors have only explored one approach in constructing the prompt. I hope for a more in-depth analysis on how to create these prompts and a detailed examination of the structure of COTs that would better facilitate formal theorem proving.
- Another concern is that while the method demonstrates effectiveness, the ablation study reveals only a modest improvement of 2.5% (Lean-STaR versus Expert iteration (SFT) as shown in Table 3 in InternLM-Plus and 3.3% in InternLM-Base. Given that COT is expected to significantly enhance performance in the MWP (Math Word Problem) domain, I had anticipated a more substantial improvement. Could you provide any insights into why the performance of COT in NTP is not as strong as it is in MWP.

questions:
- Are there any alternative methods employed in Lean-StaR to ensure the quality of the COT generated by GPT-4? Why was the state and ground truth tactic the sole approach used to generate these COTs? Would incorporating additional proving context improve the quality of the generated COTs?

- In Table 1, does the 'N' in 'Sampling' and 'Search' have different meanings? Is 'N' in 'Search' referring to the total number of expansions (nodes to explore), while 'N' in 'Sampling' refers to the number of failure retry attempts when the LLM generates a faulty tactic? I believe clarifying this in the caption would enhance understanding.

- In Lines 372 and 398, should 'Table 7' be corrected to 'Table 1'?

- It appears that the SFT finetune model does not show significant performance differences compared to InternLM2-7B. Is this due to the fact that these models have already been pre-trained on the Lean dataset?

- On Line 449, where does the 39.8% figure originate? Should it be 34.0% instead?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732443035049

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We are grateful for all the comments, questions and suggestions from the reviewers. Our responses are organized into two parts:
 1) the generic responses to the shared questions or suggestions from multiple reviewers, and
 2) the detailed responses to each point by individual reviewers.

Generic Responses

We are pleased to see the reviewers agree on the strengths of the paper, including:
 1) Focusing on an **important problem** (TwP5, MECo) in AI for math;
 2) Proposing a **novel approach** (MECo, kzhz, QMzb, TwP5) that combines the strengths of informal CoT reasoning, formal math proving with tactic generation and sampling-enhanced search;
 3) Providing **strong experimental results** (TwP5, QMzb, kzhz);
 4) Creating **a large new dataset with CoT-augmented math proving**, which should be useful for future research (TwP5, MECo).
 5) The paper is **well written** (TwP5, QMZb) but some terminology and wording need improvement (kzhz).

As for some shared questions or concerns, our Generic Responses are listed below:

### Is the proposed method lack of a verification or filtering for quality control of GPT-4 generated thoughts? (questioned by TwP5, MECo)

Our Clarification: The GPT-4 generated informal CoT thoughts are in natural language that is hard (or impossible) to verify directly. Instead, we filter out the lower-quality CoT in the STaR process by using LLM to generate tactics with CoT and checking whether the tactics can successfully prove the theorem. This process indeed can filter out the lower-quality thoughts and keep the ones that lead to correct tactics for further finetuning. In short, we do have an effective filtering process for quality control of the GPR-4 generated CoT. We will make this point more clear in our revised version of the paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a novel framework, Lean-STaR, which interleaves informal reasoning in natural language (or "thoughts") with formal reasoning in Lean (tactic prediction). The Lean-STaR framework combines several interesting techniques including LLM fine-tuning, self-taught reasoner, and expert iteration; although each idea is not new, the combination in theorem proving regarding leveraging informal reasoning is novel. The effectiveness of Lean-STaR has been shown convincingly through extensive experimental evaluations. Overall, Lean-STaR is an instance of a much general and promising framework with great potentials in theorem proving.