Exploring AXLE integration + lemma retrieval for the proving pipeline #4
Description
Hey — following up on my benchmarks question (#3). I've been digging into the AUTOLEAN pipeline and thinking about what could improve proving pass rates. A few ideas I want to explore:
1. AXLE integration
Right now the pipeline shells out to lake env lean for compilation. Axiom's AXLE API (https://axle.axiommath.ai) gives you that plus sorry2lemma(), repair_proofs(), and extract_theorems() for free over HTTP. The sub-goal decomposition alone (sorry2lemma) could help a lot — instead of proving everything in one shot, break hard proofs into steps and tackle each one.
2. Lemma retrieval
The current prompts ask the LLM to guess which Mathlib lemmas to use. With 230k+ theorems in Mathlib, that's a huge search space. Systems like Lean Finder and REAL-Prover show that injecting relevant lemmas into the prompt significantly improves pass rates. Even a basic embedding search over Mathlib names + docstrings could help.
3. Proof state feedback
When proving fails, AUTOLEAN feeds back the compiler error. But it doesn't show the LLM what goals remain or what hypotheses are available. Extracting the actual Lean proof state (via AXLE or LeanDojo-style) and feeding it into the repair prompt would make the retry loop much smarter.
I already put up a PR on the AUTOLEAN repo with a module refactor + response caching (T3S1AMAX/autolean#1), and I'm planning to run a benchmark on some classic Erdős theorems to get baseline numbers.
Would you be open to collaborating on any of these? I'm a student at UC Berkeley and happy to do the implementation work. Just want to make sure we're not duplicating effort if you already have plans in these areas.