SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kaiyu Yang $^{1}$ , Aidan M. Swope $^{2}$ , Alex Gu $^{3}$ , Rahul Chalamala $^{1}$ , Peiyang Song $^{4}$ , Shixing Yu $^{5}$ , Saad Godil $^{*}$ , Ryan Prenger $^{2}$ , Anima Anandkumar $^{1,2}$ $^{1}$ Caltech,  $^{2}$ NVIDIA,  $^{3}$ MIT,  $^{4}$ UC Santa Barbara,  $^{5}$ UT Austin https://leandojo.org
Large language models (LLMs) have shown promise in proving formal theorems using proof assistants such as Lean. However, existing methods are difficult to reproduce or build on, due to private code, data, and large compute requirements. This has created substantial barriers to research on machine learning methods for theorem proving. This paper removes these barriers by introducing LeanDojo: an open-source Lean playground consisting of toolkits, data, models, and benchmarks. LeanDojo extracts data from Lean and enables interaction with the proof environment programmatically. It contains fine-grained annotations of premises in proofs, providing valuable data for premise selection—a key bottleneck in theorem proving. Using this data, we develop ReProver (Retrieval-Augmented Prover): an LLM-based prover augmented with retrieval for selecting premises from a vast math library. It is inexpensive and needs only one GPU week of training. Our retriever leverages LeanDojo's program analysis capability to identify accessible premises and hard negative examples, which makes retrieval much more effective. Furthermore, we construct a new benchmark consisting of 98,734 theorems and proofs extracted from Lean's math library. It features challenging data split requiring the prover to generalize to theorems relying on novel premises that are never used in training. We use this benchmark for training and evaluation, and experimental results demonstrate the effectiveness of ReProver over non-retrieval baselines and GPT-4. We thus provide the first set of open-source LLM-based theorem provers without any proprietary datasets and release it under a permissive MIT license to facilitate further research.
Reasoning is a cornerstone of human intelligence and a fundamental goal of AI [3]. One prominent task is automated theorem proving (ATP): automatically generating proofs for theorems expressed in formal logic. ATP is useful for formal mathematics, producing mathematical proofs that can be checked rigorously [4]. Furthermore, it underpins formal verification, which is essential for proving the correctness and safety of high-stakes applications [5, 6].
ATP is challenging since the search space is prohibitively large. In many applications, it is impractical to generate proofs fully automatically. Therefore, interactive theorem proving (ITP) has emerged as an alternative paradigm. In ITP, proofs are constructed by human experts interacting with software tools called proof assistants, such as Coq [7], Isabelle [8], and Lean [1]. Machine learning can automate such interactive theorem proving, opening up a new avenue for theorem proving [9]. The model can learn to interact with proof assistants, given data containing human-written proofs.
*Research conducted while Saad Godil was at NVIDIA.
Figure 1: Top right: LeanDojo extracts proofs in Lean [1] into datasets for training machine learning models. It also enables the trained model to prove theorems by interacting with Lean's proof environment. Top left: The proof tree of a Lean theorem  $\forall n \in \mathbb{N}$ , gcd  $n \text{ n} = n$ , where gcd is the greatest common divisor (details in Sec. 3). When proving the theorem, we start from the original theorem as the initial state (the root) and repeatedly apply tactics (the edges) to decompose states into simpler sub-states, until all states are solved (the leaf nodes). Tactics may rely on premises such as mod_self and gcd_zero_left defined in a large math library. E.g., mod_self is an existing theorem  $\forall n \in \mathbb{N}$ ,  $n \% n = 0$  used in the proof to simplify the goal. Bottom: Our ReProver model (Sec. 5). Given a state, it retrieves premises from the math library, which are concatenated with the state and fed into an encoder-decoder Transformer [2] to generate the next tactic.
Formal theorem proving serves as an important challenge for machine learning. From a computer science perspective, formal proofs can be treated as programs [10]. But unlike conventional programs in  $\mathrm{C}++$  or Python, the correctness of proofs can be verified using proof assistants. Therefore, theorem proving may be considered a special form of code generation, with rigorous evaluation and no room for the model to hallucinate. This can be consequential to current large language models (LLMs), as they have demonstrated exceptional capability in code generation [11] but have flaws in factuality and hallucination [12]. In addition, augmenting LLMs with external tools, such as proof assistants, has shown promise in improving their various capabilities, including multi-step reasoning [13].
Current research on LLMs for theorem proving is facing many barriers. To our knowledge, none of the existing LLM-based provers are open-source [14-21]. They all use private pretraining data, and the compute requirements can reach thousands of GPU days [17]. Furthermore, some rely on tailored infrastructure for distributed training and interaction with the proof assistant—both are not possible to fully reproduce without open-source code [17, 19]. We change the status quo by introducing LeanDojo: open-source toolkits, models, and benchmarks that give researchers access to state-of-the-art LLM-based provers with modest computational costs.
Tools for Data Extraction and Interaction. We focus on Lean, a proof assistant popular among mathematicians. Our framework LeanDojo provides two essential functions for learning-based theorem proving (Fig. 1): extracting data and enabling models to interact with Lean programmatically.
For data extraction, LeanDojo extracts training data not directly visible in the raw Lean code (Fig. 2), e.g., proof trees consisting of intermediate states between proof steps (Fig. 1 Top left). In addition, LeanDojo is the first tool to locate premises in Lean proofs, enabling training machine learning models for premise selection. For interaction, LeanDojo turns Lean into a gym-like interactive environment [22]. Using LeanDojo, the model can observe proof states, change the state by executing
2"Lean" in our paper refers to Lean 3 by default. Lean 4 is not backward-compatible but is also supported by LeanDojo. Our Lean 4 results are in Appendix D.
proof steps (referred to as "tactics" in proof assistants), and receive feedback from Lean. LeanDojo is the first tool capable of interacting with Lean reliably, reducing proof-checking errors in existing tools [19] (correct proofs misjudged as incorrect) from  $21.1\%$  to  $1.4\%$ .
Retrieval-Augmented LLMs for Theorem Proving. LeanDojo addresses a key bottleneck in theorem proving: premise selection [23, 24]. Existing LLM-based provers generate the next proof step (tactic), taking only the current state as input. However, proving theorems depends critically on the premises, such as lemmas and definitions, from a math library.
For example, Fig. 1 (Top left) illustrates the proof of “ $\forall n \in \mathbb{N}$ , gcd  $\texttt{n} = \texttt{n}$ ”, where gcd stands for greatest common divisor. The proof starts from the original theorem as the initial state and repeatedly applies tactics to decompose states into simpler sub-states, until all states are solved. Tactics may rely on premises such as mod_self and gcd_zero_left defined in a large math library. E.g., mod_self is an existing theorem “ $\forall n \in \mathbb{N}$ ,  $\texttt{n} \% \texttt{n} = 0$ ” useful for simplifying the goal.
Incorporating all possible premises is too large to fit into LLMs' input, given the limited context window. Existing methods must learn to memorize the association between the proof state and the name mod_self. It works if the premise has been used in the training data to solve similar goals, but does not generalize to truly novel scenarios, e.g., theorems requiring lemmas unseen in training.
One potential solution is to complement memorization with explicit premise selection. LeanDojo extracts premise data from Lean, including where they are defined and used. It enables us to tackle premise selection by augmenting LLMs with retrieval. We introduce ReProver (Retrieval-Augmented Prover) (Fig. 1 Bottom): Given the current state, it generates a tactic conditioning on a small number of premises retrieved from Lean's math library, mathlib [25].
We need to limit retrieval to a small number of premises for it to be effective, and ideally, they should contain the ground truth premise. Our retriever builds upon Dense Passage Retriever (DPR) [26] but incorporates two algorithmic innovations: First, not all premises are accessible when proving a theorem (Sec. 3). LeanDojo can perform program analysis on Lean code to determine accessible premises. On our data, that reduces the average number of premises from 128K to 33K, significantly simplifying the retriever's task. Second, DPR needs negative examples in training and benefits from hard negatives, i.e., irrelevant premises that are hard to distinguish from ground truth ones. We propose in-file negatives: a simple mechanism to find hard negatives in premise selection, which samples negative premises defined in the same Lean source file as the ground truth premise.
LeanDojo Benchmark. Using LeanDojo, we construct a benchmark containing 98,734 theorems/proofs extracted from math1ib. Our benchmark is one of the largest math-focused theorem-proving datasets. We find that the common practice of splitting theorems randomly into training/testing has led to an overestimated performance in the previous papers. LLMs can prove seemingly difficult theorems simply by memorizing the proofs of similar theorems during training. In LeanDojo Benchmark, we mitigate this issue by designing challenging data split requiring the model to generalize to theorems relying on novel premises that are never used in training.
We use LeanDojo Benchmark to train and evaluate ReProver. Training takes only five days on a single GPU. In evaluation, ReProver can prove  $51.2\%$  theorems, outperforming a baseline that generates tactics directly without retrieval  $(47.6\%)$  and another baseline using GPT-4 [27] to generate tactics in a zero-shot manner  $(29.0\%)$ . We also test ReProver on two existing datasets, MiniF2F [28] and ProofNet [29]. It can prove  $26.5\%$  theorems in MiniF2F and  $13.8\%$  in ProofNet, which is competitive with state-of-the-art methods without reinforcement learning [19], even though trained using far fewer resources. Moreover, it can prove 65 theorems that currently do not have proofs in Lean. Thus, our tool can also serve as an effective tool for augmenting existing math libraries in Lean.
Contributions. In summary, we make four main contributions: First, we introduce tools for extracting data from and interacting with Lean. Second, we develop ReProver, the first retrieval-augmented language model for theorem proving. Third, we construct a challenging benchmark for learning-based theorem proving and use it to validate the effectiveness of ReProver. Finally, we facilitate open research on LLMs for theorem proving by releasing our data, model, and code. Our method does not rely on private datasets and can be trained on a single GPU within a week. We believe this will significantly lower the barriers to academic research in this area and establish the first accessible baselines for future work to build upon. Further, our method can be used to automatically generate new Lean proofs without requiring human effort.
Theorem Proving. Classical provers express theorems in first-order logic and search for proofs automatically in a large space [30, 31]. Even with data-driven search heuristics [32, 33], they fail to scale to large formalization projects. Therefore, recent work on learning-based theorem proving has focused on an alternative paradigm: automating the interaction with proof assistants.
The architecture of learning-based provers progressed from classical machine learning algorithms such as KNN [34], to graph neural networks explicitly encoding the syntax of formal expressions [9, 35], and now Transformer-based LLMs treating expressions as plain strings [14]. Besides the model architecture, researchers have explored several complementary dimensions: proof search algorithms for assembling model-generated steps into complete proofs [17, 21]; overcoming data scarcity through reinforcement learning (RL) [17, 19, 36, 37] or synthetic/auxiliary data [16, 38-40]; as well as outsourcing some proof goals to classical provers [18, 41-43]. Our base model without retrieval is a combination of straightforward design choices. It generates tactics by finetuning an encoder-decoder Transformer, ByT5 [44], via supervised learning without RL or auxiliary data. Then it searches for proofs using best-first search. Our model's algorithmic novelty lies in the retrieval.
Premise Selection. Selecting useful premises is recognized as a key challenge in theorem proving [23, 24, 45, 46]. Machine learning methods for premise selection have also progressed from classical models [41, 47, 48], recurrent neural networks [24], graph neural networks [38], to Transformers [49, 50]. However, existing methods either tackle premise selection in isolation without theorem proving [24, 38, 48] or feed the premises to a symbolic prover [41, 47, 49]. To our knowledge, we are the first to augment a learning-based formal theorem prover with retrieved premises so that the prover can learn how to use them effectively. For example, it can decide whether to use an explicitly retrieved premise or an implicitly memorized one.
Data and Tools for Theorem Proving. Tools for data extraction and interacting with proof assistants have been crucial drivers of learning-based theorem proving. Existing tools and datasets can be divided by proof assistants: Coq has GamePad [51], CoqGym [9], and PRISM [52]; Isabelle has IsarStep [53] and PISA [15]; HOL Light has HOLList [54] and HoLStep [55], and Lean has LeanStep [16] and lean-gym [19]. MiniF2F [28] is the only cross-system dataset, with 488 theorems for evaluation. However, it does not have training theorems and is restricted to the domain of math olympiads.
Among available tools extracting data from proof assistants, LeanDojo is the only one that can extract premises for retrieval-augmented theorem proving. A few existing datasets also have premises [49, 54], but their data extraction tools are not public, making it difficult to construct new datasets. In addition, LeanDojo is the only tool that can interact with Lean robustly (Sec. 4) and can extract data from Lean 4. See Appendix A.3 for a detailed comparison between LeanDojo and alternatives.
Mathematical Reasoning in Natural Language. We focus on proving theorems expressed in formal logic, whereas researchers have also produced a plethora of work on mathematical reasoning in natural language [56-63]. A particularly relevant task is autoformalization, translating natural language texts into formal theorems and proofs [29, 64-72].
Retrieval-Augmented Language Models. Our ReProver is the first retrieval-augmented language model for formal theorem proving, though similar architectures have been studied extensively in NLP [73-81]. In addition, there have been many retrieval-augmented methods for code generation [82-88]. Most of them retrieve from a corpus not directly related to the current file, e.g., GitHub or Stack Overflow. In contrast, our retrieval corpus consists of premises accessible to the current file, which is determined by program analysis using LeanDojo. This is similar to what CoCoMIC [88] does for Python. However, their retrieval is based on heuristics, whereas ours is learned.
3 Background: Theorem Proving in Lean
At a high level, Lean is a programming language that allows you to write not only conventional programs but also theorems and proofs. To that end, it provides two pieces of machinery: First, it provides a unified language for defining programs, mathematical objects, theorems, and proofs, based on functional programming with dependent types [89]. Second, it provides a tactic system for constructing machine-checkable proofs semi-automatically.
Figure 2: Definition of greatest common divisor (gcd) in Lean and two related theorems. The proof of gcd_self (between "begin" and "end") relies on a premise mod_self imported from another file in the math library. Lean can run this proof to produce the proof tree in Fig.1 (Top left).
We use a simple example in Fig. 2 to illustrate how theorems are formalized and proved in Lean.<sup>3</sup> Here we want to formalize the greatest common divisor (gcd) of two natural numbers. First, we define gcd as a recursive function, taking two natural numbers as parameters and returning their gcd via the Euclidean algorithm. Then, we state a lemma named gcd_zero_left that  $\forall x \in \mathbb{N}$ , gcd  $0x = x$ , which can be proved simply by the definition of gcd. Finally, we state our main theorem gcd_self that  $\forall n \in \mathbb{N}$ , gcd  $n = n$ , followed by its proof consisting of five tactics. In theorem proving, we are only concerned with generating the proof, i.e., the part between "begin" and "end"; everything before "begin" is known, including other files imported.
The syntax of tactics is quite expressive. They can take arguments and can be combined into compound tactics. You can think of tactics as programs in a domain-specific language (DSL). Users can extend the DSL by defining new tactics. This discrete, combinatorial, and unbounded action space makes theorem proving challenging for machine learning.
Another challenge is premise selection. Premises are existing lemmas or definitions useful for proving a theorem. They are used as arguments in tactics. For example, in Fig. 2 and Fig. 1 (Top left), the tactic "rewrite mod_self" rewrites the goal using the premise mod_self, which is defined in another file imported by the current file. Proofs cannot use premises that haven't been defined. For example, gcd_self cannot be used to prove gcd_zero_left. In addition, they cannot use premises not imported to the current file. Still, premises come from a large math library containing hundreds of thousands of existing definitions and theorems, making it hard, for humans and machines alike, to select the right premises when generating a tactic. This is a key bottleneck in theorem proving and is what we aim to address through retrieval-augmented LLMs.
4 LeanDojo: Toolkit and Benchmark
LeanDojo serves two essential needs of learning-based theorem proving in Lean. First, it extracts training data from Lean, and we use this capability to construct a challenging theorem proving benchmark. Second, it enables the model to interact with Lean programmatically.
Data Extraction. Lean repos (e.g., mathlib or lean-liquid) contain source code of human-written theorems/proofs. However, the raw code is unsuitable for training the prover. It lacks runtime information that humans can access when using Lean, such as intermediate states between proof steps. Therefore, LeanDojo extracts the following information not directly visible in the code:
3The process is similar in many other proof assistants, though they may have different logical foundations.
- File dependencies and abstract syntax trees (ASTs): LeanDojo processes the repo to produce a directed acyclic graph whose nodes are files and edges are import relations between files. In addition, LeanDojo produces the AST of each file. File dependencies and ASTs are useful for program analysis, e.g., collecting theorems defined in a file or premises accessible to a theorem.
- States and tactics: LeanDojo extracts all tactics in proofs. For each tactic, it also extracts the states before/after the tactic, which allows us to reconstruct the proof tree in Fig. 1 (Top left).
- Premises: For each premise, such as mod_self in Fig. 2, LeanDojo records where it is defined (location in data/nat/lemma.lean) and where it is used (locations across many files). In addition, premises have unique fully qualified names (e.g., nat.mod_self) but are often used by ambiguous short names (mod_self), relying on Lean to perform name resolution. LeanDojo is capable of recording their full names.
Lean has basic support for exporting dependencies, ASTs, states, and tactics. However, it cannot resolve the premises' full names and locate their definitions. Therefore, we modify Lean to record this information (details in Appendix A.1). The modified Lean is used only for data extraction but not for evaluation, so we do not risk accidentally breaking Lean's logical soundness.
LeanDojo Benchmark. We construct a benchmark for premise selection and theorem proving, named LeanDojo Benchmark. The data is extracted from mathlib, $^{4}$  Lean's centralized math library covering diverse topics such as analysis, algebra, and geometry. $^{5}$  LeanDojo Benchmark is one of the largest math-focused theorem proving datasets, consisting of 98,734 theorems from 3,384 Lean files. Unlike existing datasets in Lean [16], LeanDojo Benchmark also contains the definitions of 130,262 premises, including not only theorems but also other definitions that can be used as premises (e.g., gcd in Fig. 2. Furthermore, the dataset has 217,776 tactics, 129,243 of them with at least one premise. The average number of premises is 2.13 among tactics with premises. Appendix B contains additional information on data format, datasheet [90], hosting, and licensing.
src/algebra/quadernion lean   
lemma conj.mul:  $(a*b).conj = b.conj*a.conj := begin$  ext;simp;ring_exp end   
lemma conj_conj.mul:(a.conj\*b).conj  $=$  b.conj\*a:=begin rw [conj.mul,conj_conj]   
end   
lemma conj.mul_conj:(a\*b.conj).conj  $=$  b\*a.conj:=begin rw [conj.mul,conj_conj]   
end
Figure 3: Similar theorems/proofs are common. If splitting them randomly into training/testing, the model can prove testing theorems by memorization.
LeanDojo Benchmark has 94,734/2,000/2,000 theorems for training/validation/testing. It features a challenging data split for testing the prover's generalization in more realistic scenarios. Splitting theorems randomly can overestimate the prover's performance, by allowing it to prove many theorems through memorization. In human-written Lean code, a common idiom is to have a block of similar theorems/proofs for slightly different properties of the same math concept. For example, in Fig. 3, the last two theorems not only look similar but have identical proofs. If one of them is in training, the model can easily prove the other one by memorization. This shortcut enables the model to prove seemingly nontrivial theorems, including those requiring premises to prove.
To mitigate this issue, besides the random split, we create a challenging data split named novel_premises. It requires testing proofs to use at least one premise that has never been used in training. For example, the last two theorems in Fig. 3 both use the premise conj.mul. If one theorem is in the training set of the novel_premises split, the other one must also be in training.
<sup>4</sup>We use the commit 19c869efa56bbb8b500f2724c0b77261edbfa28c released on October 11, 2023.
5More details, statistics, and visualizations of mathlib can be found at https://leanprover-community. github.io/mathlib/stats.html.
Interacting with Lean. Another important function of LeanDojo is to interact with Lean programmatically. It turns Lean into a gym-like environment [22], in which the prover can observe the proof state, run tactics to change the state, and receive feedback on errors or on proof completion. This environment is indispensable for evaluating/deploying the prover or training it through RL.
Below is LeanDojo's main interface for interacting with Lean through tactics. Lean also supports other proof styles not based on tactics. Although we only support tactic-style proofs, they are sufficiently general since any proof can be converted to a tactic-style proof.[6]