Paper deep dive
Theorem Proving with Retrieval-Augmented Language Models — Yang, Swope, Gu, Chalamala, Song, Yu, Godil, Prenger, Anandkumar. NeurIPS 2023 (Datasets & Benchmarks Track, Oral).
TL;DR
LeanDojo is two things: a toolkit that turns Lean's mathlib library into 98,734 machine-readable theorems and proofs, and ReProver — a language model that proves new theorems by proposing one proof step at a time and searching over the results with best-first search, the same family of algorithm as Greedy in Lab 02.
The problem it targets: prior LLM theorem provers were mostly closed (private code, data, compute) and bad at picking the right lemma out of a huge math library. LeanDojo open-sources the whole pipeline and adds retrieval to fix that second problem.
Same three ingredients
Every search algorithm in Lab 01 & 02 needs the same three things: a node, an action, and something ranking the frontier. LeanDojo swaps out what fills each slot — the algorithm doesn't change.
| Search concept | Grid / graph search (Lab 01–02) | LeanDojo / ReProver |
|---|---|---|
| Search node | A state — a square, a puzzle configuration | A proof state: the current goal(s) + hypotheses in Lean, e.g. n : ℕ ⊢ gcd n n = n |
| Action | A move — go left, slide a tile | A tactic: a Lean command that transforms the goal, e.g. simp [gcd] or rw [gcd_comm] |
| Who proposes actions | The successor function (fixed, mechanical) | A language model, generating several candidate tactics per state |
| Frontier priority | h(n), or f(n)=g(n)+h(n) | The model's cumulative confidence (log-probability) in the tactic sequence so far |
| What narrows the branching factor | — | Premise retrieval: instead of guessing a lemma name from scratch, retrieve likely-relevant ones from mathlib first |
| Goal test | Reached the goal state | No open goals remain — the proof is complete |
This destroys a common misconception: search isn't just for route planning. Any problem you can break into states + actions + a way to rank candidates can be searched this way.
Interactive · illustrative
Goal (from LeanDojo's own documentation): n : ℕ ⊢ gcd n n = n — "for any natural number n, gcd(n,n) = n."
The tree below is a simplified example we built to show the mechanism — the real candidate tactics, scores, and branches a trained ReProver model would generate for this goal aren't published in the paper. The goal statement and the two example tactics (simp [gcd], rw [gcd_comm]) are real, taken from ReProver's own README.
Notice: the search briefly explores induction n (a reasonable-looking candidate) before backtracking to the shorter, higher-confidence simp [gcd] proof — because best-first search compares the whole frontier every step, not just the most recent branch. That's the same property that makes A* (not Hill Climbing) resistant to misleading first impressions in Lab 02.
Why retrieval matters
mathlib has 130,262 premises (definitions, lemmas, theorems) available to cite. Most aren't relevant, or aren't even in scope, for a given goal — on average, 33,160 are "accessible" per theorem. Asking a language model to guess the right one from scratch is like Greedy search with an enormous branching factor b. ReProver's retriever narrows that down to the top 100 likely-relevant premises before the tactic generator even runs.
Does it work?
Two test splits: random (ordinary held-out theorems) and novel_premises (theorems needing at least one premise never seen in that combination during training — the harder generalization test).
Retrieval's contribution is the gap between the two ReProver bars: +3.6pp on the easy split, +3.1pp on the hard one — small in absolute terms, but note how much further ahead ReProver already is of GPT-4 used zero-shot, especially on novel_premises (26.3% vs 7.4%). GPT-4 has broad math knowledge but no way to search Lean's specific proof state space; ReProver does.
Back to the lecture
Nothing about the search algorithm changed to make this work on mathematics instead of a grid. What changed is what fills in state, action, and frontier priority — and, because the branching factor here is enormous (thousands of possible lemmas per step), retrieval had to be added just to make search tractable at all. That's the same lesson as choosing a good heuristic in Lab 02: a better priority function means less wasted search.
← Back to Lab 02: Research Connections