LeanDojo: Theorem Proving with Retrieval-Augmented Language Models

Large language models (LLMs) have shown promise in proving formal theorems in proof assistants such as Lean. However, existing methods are difficult to reproduce or build upon, due to private code, data, interactive environments, and distributed systems with hundreds of GPUs. These have created substantial barriers to research on machine learning for theorem proving. In this paper, we break the barriers through open toolkits, benchmarks, and models. We introduce LeanDojo: A tool for extracting data from Lean and interacting with the proof environment programmatically. LeanDojo features fine-grained annotations of premises in proofs, providing valuable data for premise selection—a key bottleneck in theorem proving. Using LeanDojo, we construct a challenging benchmark for theorem proving. Furthermore, we develop ReProver (Retrieval-Augmented Prover): the first LLM-based prover augmented with retrieval for selecting premises from a vast math library. And we design novel recipes for training the premise retriever. Experimental results demonstrate the effectiveness of our method over non-retrieval baselines and GPT-4. ReProver is the first LLM-based prover trainable on a single GPU and without proprietary datasets. We will release our data, model, and code to facilitate future research.