SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- initialize(theorem): Given the theorem to prove, LeanDojo returns the initial state. A valid state is a string representing current proof goals and local contexts (see the nodes in Fig. 1 Top left). When there are multiple goals, their strings are concatenated.
- run_tac(state, tactic): Run a tactic on a given state and return the next state. The returned state will be an error state if the tactic execution is not successful, e.g., due to timeout or inapplicable tactic. If the input state is an error, the result can only be an error.
Building this environment is technically challenging, as Lean is designed for human users, not machines. LeanDojo is the first tool that can interact with Lean reliably. Existing tool [19] is limited:  $21.1\%$  of the ground truth proofs are misjudged as incorrect, due to issues with how they construct the proof environment, which distorts the reported performance and produces unreliable feedback when used in reinforcement learning. In contrast, LeanDojo reduces the number of misjudgments to  $1.4\%$ . Details are in Appendix A.2.
5 ReProver: Retrieval-Augmented Theorem Prover
We develop the ReProver model that uses retrieval to select premises explicitly. At its core is a retrieval-augmented tactic generator (Fig. 1 Bottom). Given the current proof state, it retrieves a handful of potentially useful premises and generates a tactic conditioning on the concatenation of the state and retrieved premises. When proving theorems, the model generates multiple tactic candidates at each step, which are used in a standard best-first search algorithm to find proofs [16, 18, 19, 28].
Premise Retrieval. Our retriever is based on Dense Passage Retriever [26]. Given a state  $s$  as the query and a library of candidate premises  $\mathcal{P} = \{p_i\}_{i=1}^N$ , it retrieves a ranked list of  $m$  premises  $\{p_i'\}_{i=1}^m$  from  $\mathcal{P}$ . In DPR,  $s$  and  $p_i$  are both raw texts but are embedded in a vector space, and we retrieve the top  $m$  premises maximizing the cosine similarity between the state and the premise.
More formally, we have a function  $f$  parameterized by  $\theta$  for embedding both the state and the premises into a  $h$ -dimensional vector space:  $f(s,\theta), f(p_i,\theta) \in \mathbb{R}^h$ . We retrieve premises maximizing  $f(s,\theta)^T f(p_i,\theta) / (\|f(s,\theta)\|_2 \|f(p_i,\theta)\|_2)$ . We choose  $f$  to be a Transformer encoder [2] followed by average pooling:  $f(\cdot,\theta) = \mathrm{AvgPool}(\mathrm{Enc}(\cdot,\theta))$ .
The retrieval is efficient. The premise embeddings  $f(p_i, \theta)$  can be pre-computed, and we only need one forward pass to compute  $f(s, \theta)$ . We do not rerank the retrieved premises as in Magnushammer [49], which is more costly since it requires a separate forward pass for each retrieved premise.
Similar to DPR, we train the retriever by minimizing a contrastive loss between positive premises and in-batch negative premises. Specifically, suppose we have a batch of  $b$  states. For each state, we sample a positive premise from the ground truth and  $n$  negative premises from  $\mathcal{P}$ . They are called "in-batch" negatives because they are shared by all states in the batch—Every state is associated with all  $b \cdot (n + 1)$  premises; at least 1 of them is positive. Let  $l_{ij} \in \{0,1\}$  denote whether a state-premise pair  $(s_i, p_j)$  is positive. We minimize the mean squared loss:
$$
\mathcal {L} (\theta) = \sum_ {i = 1} ^ {b} \sum_ {j = 1} ^ {b \cdot (n + 1)} \left| l _ {i j} - \frac {f \left(s _ {i} , \theta\right) ^ {T} f \left(p _ {j} , \theta\right)}{\left\| f \left(s _ {i} , \theta\right) \right\| _ {2} \left\| f \left(p _ {j} , \theta\right) \right\| _ {2}} \right| ^ {2}. \tag {1}
$$
Another common type of proofs is "term-style proofs". Any term-style proof "X" can always be converted into an equivalent tactic-style proof "exact X", though such conversion may lead to unidiomatic proofs.
7When training the retriever, we ignore proof states followed by tactics without using any premise.
Retrieving from Accessible Premises. We incorporate into DPR two insights tailored to premise selection. First, instead of retrieving from all premises in the math library, we restrict to premises accessible to the current theorem. They include premises defined in the same file before the theorem, as well as those imported from other files. We compute accessible premises for each theorem, relying on LeanDojo's capability in program analysis (Sec. 4). Focusing on accessible premises makes  $\mathcal{P}$  much smaller. LeanDojo Benchmark contains 130,262 premises in total, but the average number of accessible premises is only 33,160.
In-file Negative Examples. DPR's performance depends critically on the quality of negative examples [91, 92]. In early experiments, we sampled all  $n$  negative premises randomly, and the model often mistakenly retrieved other premises from the same file as the positive one. Therefore, we propose a scheme that samples  $k$  in-file negatives and  $n - k$  random negatives for training.
Tactic Generation. As in Fig. 1 (Bottom), retrieved premises are concatenated with the state. $^{8}$  Then an encoder-decoder Transformer, ByT5 [44], takes them as input and generates the tactic. The model is trained to minimize the cross entropy loss w.r.t. human-written tactics.
Training ReProver takes substantially less compute than prior methods (120 GPU hours vs. more than 1000 hours [16, 17]). All existing LLM-based provers pretrain on datasets specific to math and coding [14-20]. The pretraining is computationally expensive, and the datasets are kept private. In contrast, we choose to avoid domain-specific pretraining and build upon google/byt5-small—a model checkpoint that is generic, publicly available, and relatively small (299M parameters vs. 837M [16] or 600M [17]). We could see further benefits from domain-specific pretraining, as in Minerva [57], or stronger LLMs like LLaMA [93] or StarCoder [94], but that is beyond our scope. In addition, our model is finetuned on human-written tactics only, without auxiliary data [16] or data collected through online interaction with Lean [17, 19]. These orthogonal directions are valuable but will significantly increase the method's complexity and compute requirements.
We evaluate ReProver on LeanDojo Benchmark. It outperforms baselines on premise selection and theorem proving, demonstrating the promise of theorem proving with retrieval-augmented language models. Experimental details and hyperparameters are in Appendix C.1.
Premise Selection. For premise selection, we only use tactics in LeanDojo Benchmark that have at least one premise. The model, based on a ByT5 encoder, uses the state before a tactic as the query to retrieve 100 premises. Then, we calculate standard metrics in information retrieval:  $\mathbb{R}@\mathbb{k}$  (recall for the top  $k$  retrieved premises) and MRR (mean reciprocal rank).
Our first baseline is a classical BM25 retriever [95] without machine learning. Results in Table 1 show that our method outperforms BM25 significantly across the board. However, it exhibits a large performance degradation on the challenging data split (comparing novel premises to random). This is consistent with the general observation that machine learning can be brittle in the presence of distribution shifts. In addition, we compare with two ablations: one retrieving from all premises (instead of accessible premises only) and the other without in-file negatives. They perform worse than our method, demonstrating the effectiveness of our two improvements upon DPR.
Theorem Proving Experimental Setup. Then we evaluate ReProver on theorem proving. The training has two stages: First, we train the retriever and use it to retrieve 100 premises for all proof states in LeanDojo Benchmark. Second, we train the tactic generator, taking as input the concatenation of the state and retrieved premises (truncated to a length limit). During evaluation, the tactic generator is combined with best-first search to prove theorems. We evaluate the  $Pass@1$  metric: The prover is given only one attempt and must find the proof within a wall time limit of 10 minutes. Training takes five days on a single NVIDIA A100 GPU with 80GB memory, and evaluation takes two days on eight V100 GPUs. Please see Appendix C.1 for details.
Baselines. Following prior work [16, 28], we include tidy as a baseline. It is a tactic in mathlib that tries to complete the proof using heuristics (without machine learning). We apply tidy directly
<sup>8</sup>We retrieve 100 premises, concatenate them with the state, and truncate the concatenation to a fixed length.
Table 1: Premise selection testing performance. For each method, we train and evaluate two models independently using different data splits (random and novel PREmises; see Sec. 4). R@k is the recall for the top  $k$  retrieved premises, and MRR is the mean reciprocal rank metric (higher is better). Our retriever outperforms BM25 and ablations. Results for Lean 4 are in Appendix D.
Method	random	novel_premises
R@1	R@10	MRR	R@1	R@10	MRR
BM25	6.7	17.2	0.15	5.9	15.5	0.14
w/ all premises	1.9	11.9	0.08	2.1	12.4	0.08
Ours	13.5	38.4	0.31	9.1	27.6	0.24
w/ all premises	11.7	36.2	0.27	7.1	23.1	0.20
w/o in-file negatives	10.8	33.1	0.25	7.9	25.7	0.22
to the original theorem and see if it can succeed within the wall time limit. Another baseline uses GPT-4 as the tactic generator. Given a state, it queries GPT-4 to generate 35 tactics in zero-shot. After removing invalid ones, the remaining tactics are combined with best-first search to find proofs. Data contamination is possible: Many proofs had been publicly available on GitHub before GPT-4's data cutoff date (September 2021). See Appendix C.2 for details.
Unfortunately, it is not feasible to compare with existing LLM-based provers in Lean [16, 17, 19]. None of them are open-source or can be reproduced with reasonable effort. Furthermore, we cannot compare directly with the numbers reported in their papers, due to differences in data, infrastructure, and training procedures (details in Appendix C.3). Many difficulties are due to the private nature of existing methods. By releasing our code and models, we hope to create accessible baselines for future work to build upon.
Table 2: Theorem proving Pass@1 (%) on the testing data of LeanDojo Benchmark. Our ReProver model outperforms tidy, GPT-4, and a baseline that generates tactics directly without retrieval. Results for Lean 4 are in Appendix D.
Method	random	novel_premises
tidy	23.8	5.3
GPT-4	29.0	7.4
ReProver (ours)	51.2	26.3
w/o retrieval	47.6	23.2
Results. Table 2 shows the results on the testing data of LeanDojo Benchmark. ReProver outperforms all baselines on two different data splits, demonstrating the effectiveness of retrieval-augmented theorem proving. GPT-4 performs substantially worse than our method, even though it may have seen the ground truth proofs due to data contamination. The task cannot be solved out of the box by state-of-the-art LLMs, calling for algorithmic innovations to make further progress.
Testing theorems in novel premises are indeed much more challenging. All methods in Table 2 perform substantially worse on novel premises than the random split. We argue that performance on challenging splits is more indicative of the prover's capability and should be emphasized in the future development of theorem proving.
Evaluation on MiniF2F and ProofNet. We run ReProver to prove theorems in MiniF2F [28] and ProofNet [29]. These two datasets are for testing only and do not have training theorems, which makes them challenging since the distribution of theorems is quite different from mathlib used to train ReProver. MiniF2F focuses on math olympiads, and ProofNet focuses on exercises in undergraduate math textbooks. On MiniF2F's test set in Lean, ReProver achieves a Pass@1 of  $26.5\%$ , which is competitive with state-of-the-art methods without RL (25.9% in Polu et al. [19]). On ProofNet, our Pass@1 is  $13.8\%$ , which is the first reported theorem proving result on this dataset. Further, many theorems do not have ground truth proofs in Lean. Our prover discovers 33 proofs in MiniF2F and 39 proofs in ProofNet that currently do not have Lean proofs. Please see Appendix C.4 for details, examples, and caveats.
We have introduced LeanDojo: an open-source playground for learning-based theorem proving in Lean, consisting of toolkits, models, and benchmarks. It extracts data from Lean and enables the model to interact with Lean programmatically. We have developed ReProver, the first retrieval-augmented LLM for theorem proving. Limitations and future work are discussed in Appendix F.
We have released our code, data, models, and documentation to facilitate future research:
- LeanDojo's codebase for data extraction and interaction with Lean: https://github.com/lean-dojo/LeanDojo
- LeanDojo's documentation: https://leandojo.readthedocs.io
- Datasets: (1) LeanDojo Benchmark: https://doi.org/10.5281/zenodo.8016385 with DOI 10.5281/zenodo.8016385. (2) LeanDojo Benchmark 4 (Appendix D): https://doi.org/10.5281/zenodo.8040109 with DOI 10.5281/zenodo.8040109.
- ReProver's code and models: https://github.com/lean-dojo/ReProver
- ChatGPT plugin (Appendix E): https://github.com/lean-dojo/LeanDojoChatGPT
- LeanDojo Website: https://leandojo.org
This work is partially supported by Caltech's Center for Autonomous Systems and Technologies. Kaiyu Yang is supported by the Computing, Data, and Society Postdoctoral Fellowship at Caltech. Alex Gu is supported by the National Science Foundation (NSF) Graduate Research Fellowship. Rahul Chalamala and Peiyang Song are supported by the Summer Undergraduate Research Fellowships (SURF) program at Caltech. Anima Anandkumar is partially supported by the Bren endowed chair. We appreciate the valuable feedback from Logan Murphy and members of the Anima AI+Science Lab on an initial version of this paper. We thank Junyan Xu for manually inspecting the proofs generated by our model on ProofNet. We also thank Jeremy Avigad and Mario Carneiro for insightful discussions on supporting Lean 4 in LeanDojo.
[97] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-5551, 2020. 23
[98] Jeff Rasley, Samyam Rajbhandari, Olatunjri Ruwase, and Yuxiong He. DeepSpeed: System optimizations enable training deep learning models with over 100 billion parameters. In International Conference on Knowledge Discovery and Data Mining (KDD), 2020. 24
[99] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019. 24
[100] Mathlib Community. Mathport: A tool for porting Lean 3 projects to Lean 4. URL https://github.com/leanprover-community/mathport.27
[101] OpenAI. ChatGPT plugins. https://openai.com/blog/chatgpt-plugins, 2023. URL https://openai.com/blog/chatgpt-plugins.29
[102] Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. arXiv preprint arXiv:2305.10601, 2023. 31, 33
[103] Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, and Caiming Xiong. CodeGen: An open large language model for code with multi-turn program synthesis. arXiv preprint arXiv:2203.13474, 2022. 32
[104] Qinkai Zheng, Xiao Xia, Xu Zou, Yuxiao Dong, Shan Wang, Yufei Xue, Zihan Wang, Lei Shen, Andi Wang, Yang Li, et al. CodeGeeX: A pre-trained model for code generation with multilingual evaluations on HumanEval-X. arXiv preprint arXiv:2303.17568, 2023. 32
[105] Lili Yu, Daniel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. MegaByte: Predicting million-byte sequences with multiscale transformers. arXiv preprint arXiv:2305.07185, 2023. 32
[106] Gautier Izacard and Édouard Grave. Leveraging passage retrieval with generative models for open domain question answering. In European Chapter of the Association for Computational Linguistics (EACL), 2021. 33
[107] Yi Tay, Vinh Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, et al. Transformer memory as a differentiable search index. In Neural Information Processing Systems (NeurIPS), 2022. 34
[108] Michele Bevilacqua, Giuseppe Ottaviano, Patrick Lewis, Scott Yih, Sebastian Riedel, and Fabio Petroni. Autoregressive search engines: Generating substrings as document identifiers. 2022.
[109] Ronak Pradeep, Kai Hui, Jai Gupta, Adam D Lelkes, Honglei Zhuang, Jimmy Lin, Donald Metzler, and Vinh Q Tran. How does generative retrieval scale to millions of passages? arXiv preprint arXiv:2305.11841, 2023. 34
[110] Norman Megill and David A Wheeler. *Metamath: a computer language for mathematical proofs*. Lulu. com, 2019. 36
A LeanDojo Technical Details 18
A.1 Extracting Premise Information from Lean's Elaborator 18
A.2 Reliable Interaction with Lean 18
A.3 Comparison with Existing Tools for Learning-Based Theorem Proving in Lean 19
B LeanDojo Benchmark 20
B.1 Dataset Format 20
B.2 Datasheet 22
B.3 Data Hosting, Licensing, and Maintenance 23
C Experiments 23
C.1 Details and Hyperparameters 23
C.2 The GPT-4 Baseline 24
C.3 Justifications for Not Comparing with Existing LLM-Based Provers 26
C.4 Evaluation on MiniF2F and ProofNet 26
D LeanDojo for Lean 4 27
E ChatGPT Plugin for Theorem Proving 29
F Limitations and Future Work 32
A LeanDojo Technical Details
We provide more information on how LeanDojo extracts data from and interacts with Lean. For further details, please check our open-source implementation.
A.1 Extracting Premise Information from Lean's Elaborator
"Premises" in this paper belong to a category of Lean expressions called "constants." In Lean, definitions of constants are grouped into nested, hierarchical namespaces. Therefore, each premise has a unique fully-qualified name. For example, mod_self in Fig. 2 is defined in the namespace nat; therefore, its fully qualified name is nat.mod_self. However, it would be too verbose if premises had to be referred to using full names. In practice, tactics often refer to premises using short names such as mod_self. In case multiple premises share the same short name, Lean automatically infers the correct one from the context through a process called "name resolution". LeanDojo is able to trace the input/output of Lean's name resolution and thereby extract accurate premise information for training the retriever.
Name resolution in Lean is implemented in a process called "elaboration," which happens after parsing but before the parsed expressions are checked by Lean's trusted kernel. Elaboration takes as input user-entered expressions (called "pre-expressions") that are concise, partially specified, and potentially ambiguous. It turns them into complete expressions ready to be checked by the kernel. This is realized by inferring not only full names but also missing types, implicit arguments, overloading, type coercion, etc. Please refer to de Moura et al. [96] for details on Lean's elaboration process. In LeanDojo, we modify Lean's internal implementation, intercepting the elaborator to record its input/output:
- Pre-expression: The input to Lean's elaborator, including where premises are used in proofs.
- Expression: The output of the elaborator, including the premise's full name and where it is defined.
Locations are spans in the source code, specified by the file name and the row/column numbers of its start/end. Our modification takes the form of a Git patch that LeanDojo can automatically apply to any version of Lean 3 after March 24, 2022.
A.2 Reliable Interaction with Lean
Polu et al. [19] introduced lean-gym. To our knowledge, it is the only mature, open-source tool before LeanDojo for interacting with Lean programmatically. However, we found severe issues with lean-gym: About  $21.1\%$  of the correct, human-written proofs are misjudged as incorrect, leading to two problems: First, it underestimates the prover's evaluation performance. Second, the results are too noisy as feedback signals for reinforcement learning.
After carefully analyzing lean-gym's implementation, we identified the root cause of the problem. When proving a theorem, the environment used by lean-gym is subtly different from the original environment used by humans. Specifically, lean-gym fails to handle namespaces correctly (illustrated in Fig. A). As a result, name resolution fails unexpectedly when checking correct proofs.
For example, Fig. A compares the correct environment and the environment constructed by lean-gym. The theorem should be inside the namespace "buffer". However, in lean-gym, it merely opens the namespace. These two scenarios are different when it comes to name resolution. Being inside a namespace instructs Lean to favor constants defined in that namespace, whereas opening a namespace does not have such an effect. In this example, the short name "read" is ambiguous: We have "monad_reader.read" defined in "init/control/reader lean" and "buffer.read" defined in "data/buffer lean". In the correct environment, the "read" in "unfold read" resolves to "buffer.read". Whereas in lean-gym's environment, it incorrectly resolved to "monad_reader.read". Lean complains that "read" is not an equational lemma, because it is referring to a wrong "read". LeanDojo does not suffer from this kind of error since it uses a different mechanism for constructing the environment. Specifically, it wraps the interaction code as a Lean tactic, which is inserted into the proof. Therefore, the environment is guaranteed to be correct.
"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.
We quantitatively compare lean-gym and LeanDojo on the number of proof checking errors. In this study, we use Lean v3.42.1 paired with mathlib version 6e5ca7d0097313e59f7533a42e3ea5197484c775 since they are supported by both tools. We use LeanDojo to extract all tactic-style proofs and enter them into both tools. These proofs are all correct, but lean-gym failed on  $21.1\%$  of them. In contrast, LeanDojo only failed on  $1.4\%$ , and its failures are a subset of lean-gym's. We include this study in our open-source repo and document example proofs from the remaining  $1.4\%$  to provide transparency on LeanDojo's limitations.[10]
Figure A: An example of correct proofs misjudged as incorrect by lean-gym, adapted from the theorem read_eq_read' in "data/buffer.lean" of Lean's standard library. The error message is because lean-gym failed to resolve the short name "read" to the correct fully-qualified name. The Lean code in this figure is only for illustrative purposes. It does not reflect the implementation technique used by lean-gym to construct the environment. Instead of generating actual Lean code, lean-gym uses Lean's metaprogramming APIs to construct the environment.
A.3 Comparison with Existing Tools for Learning-Based Theorem Proving in Lean
To our knowledge, LeanStep [16] and lean-gym [19] are the only published tools for learning-based theorem proving in Lean. There are a few unpublished prototypes, such as repl, lean-client-python, and lean-gym for Lean 4, none of which is mature enough or is under active development. Therefore, we only compare LeanDojo with LeanStep and lean-gym (summarized in Table A).
Functionality. LeanDojo supports both data extraction and interacting with Lean programmatically. In contrast, LeanStep is only for data extraction, and lean-gym is only for interacting with Lean. They are not actively maintained, so they do not support recent versions of mathlib (tested on August 11, 2023, using mathlib commit 19c869efa56bbb8b500f2724c0b77261edbfa28c). Also, neither of them support Lean 4 (Appendix D). LeanDojo fully supports recent mathlib and Lean 4. Furthermore, LeanStep cannot extract premise information and is not applicable to repos other than mathlib. Last, LeanDojo comes with comprehensive documentation and unit tests, whereas other tools barely have any.