[Paper Review] LeanDojo: Theorem Proving with Retrieval-Augmented Language Models

Kaiyu Yang, Aidan Swope+7 more|arXiv (Cornell University)|Jun 27, 2023

Mathematics, Computing, and Information Processing38 citations

TL;DR

LeanDojo introduces an open-source Lean playground with data, models, and benchmarks, and presents ReProver, a retrieval-augmented prover that uses premises retrieved from Lean's math library to improve theorem proving with LLMs.

ABSTRACT

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.

Motivation & Objective

• Provide open, reproducible tools for extracting Lean data and interacting with Lean programmatically.
• Develop a retrieval-augmented prover (ReProver) that selects premises from mathlib to generate tactics.
• Create a large, challenging Lean-based benchmark to evaluate premise selection and proving performance.
• Demonstrate that retrieval augmentation improves prover performance over non-retrieval baselines and GPT-4 on LeanMath benchmarks.

Proposed method

• LeanDojo extracts runtime proof data from Lean (states, tactics, premises) and augments Lean with full premise naming and accessibility information.
• ReProver uses a retrieval-augmented tactic generator that conditions on a small set of retrieved premises.
• Premise retrieval is based on Dense Passage Retriever with enhancements: restricting to accessible premises and using in-file negative examples.
• Training Finetunes ByT5 encoder-decoder on states+premises to generate tactics, followed by best-first search for proofs.
• The dataset LeanDojo Benchmark contains 98,734 theorems/proofs and 130,262 premises, with novel_premises data split to test generalization.

Experimental results

Research questions

• RQ1 Can retrieval-augmented prompting using retrieved premises improve interactive theorem proving in Lean compared to non-retrieval baselines?
• RQ2 How does restricting retrieval to accessible premises and using in-file negatives affect premise recall and proof success?
• RQ3 What is the impact of the LeanDojo Benchmark splits (random vs novel_premises) on generalization to novel premises?
• RQ4 How does ReProver perform on external datasets MiniF2F and ProofNet relative to existing methods without heavy RL?

Key findings

• ReProver achieves 51.2% Pass@1 on LeanDojo Benchmark random split, outperforming non-retrieval baselines (47.6%) and GPT-4 (29.0%).
• On novel_premises, ReProver attains 26.3% Pass@1, outpacing non-retrieval baselines (23.2%) and GPT-4 (7.4%).
• Premise retrieval with accessible-pr Premises and in-file negatives significantly improves recall metrics over baselines (e.g., in Table 1, Ours achieves R@1 of 13.5 on random vs 6.7 for BM25).
• ReProver proves 26.5% of MiniF2F and 13.8% of ProofNet test sets in Lean, competitive with state-of-the-art non-RL methods and discovering dozens of proofs lacking Lean proofs.
• Training requires only five GPU days on a single GPU and evaluation on eight GPUs, with open-source code, data, and models released under MIT license.
• LeanDojo Benchmark is one of the largest math-focused theorem proving datasets and emphasizes challenging generalization data splits.