Poster
in
Workshop: Knowledge and Logical Reasoning in the Era of Data-driven Learning
LeanDojo: Theorem Proving with Retrieval-Augmented Language Models
Kaiyu Yang · Aidan Swope · Alexander Gu · Rahul Chalamala · Shixing Yu · Saad Godil · Ryan Prenger · Animashree Anandkumar
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.