Timezone: »

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.

Author Information

Kaiyu Yang (Caltech)
Kaiyu Yang

Kaiyu Yang is a postdoctoral researcher at Caltech in the Computing + Mathematical Sciences (CMS) Department, working with Prof. Anima Anandkumar. His research aims to make machine learning capable of symbolic reasoning. It includes (1) applying machine learning to symbolic reasoning tasks, such as automated theorem proving; and (2) introducing symbolic components into machine learning models to make them more interpretable, verifiable, and data-efficient. In addition, he has also worked on constructing and analyzing machine learning datasets, especially focusing on fairness, privacy, and mitigating dataset bias. His research is recognized with a Siebel Scholar award. Before joining Caltech, he received his Ph.D. from the Department of Computer Science at Princeton University, advised by Prof. Jia Deng.

Aidan Swope (NVIDIA)
Alexander Gu (MIT)
Rahul Chalamala (Caltech)
Shixing Yu (University of Texas at Austin)
Saad Godil (Hippocratic AI)
Ryan Prenger (NVIDIA)
Animashree Anandkumar (Caltech and NVIDIA)

More from the Same Authors