Lean State Search: Search Mathlib theorems with proof states

Stream: Machine Learning for Theorem Proving

[Yicheng Tao]

Announcing Lean State Search: A Proof State based Search Engine for Mathlib Theorems

We are thrilled to introduce Lean State Search, a semantic search engine powered by a pretrained model specifically designed for retrieving Mathlib theorems using formal proof states as queries. Developed by the AI4Math team at Renmin University of China, this tool aims to streamline formalization in Lean.

Resources

• arXiv Paper: Technical details and methodology.

• GitHub Repository: Open-source code for training and inference.

• Hugging Face Model: Pretrained model for direct integration.

Clarification
Unlike existing tools like LeanSearch.net and Moogle.ai, Lean State Search focuses exclusively on formal proof states as input—mirroring the premise retrieval functionality of LeanDojo.

Open-Source Commitment
We have open-sourced our model and training/inference code. Self-hosting support will be released shortly. The lightweight BERT-based architecture ensures seamless deployment even on CPU-only laptops, lowering barriers to adoption.
Community-Driven Improvement
User feedback is vital to refining our model. All anonymized interaction data will be publicly shared and periodically updated to foster transparency and collaborative progress.

We welcome your engagement and look forward to supporting the formalization community with this tool.

[Floris van Doorn]

This seems nice.

I did a single test with it, using a small lemma proven in the Carleson project:

import Mathlib.MeasureTheory.Integral.Average

open scoped ENNReal
lemma laverage_mono_ae {α : Type*} [MeasurableSpace α] (μ : Measure α)
    {f g : α → ℝ≥0∞} (h : ∀ᵐ a ∂μ, f a ≤ g a) :
    ⨍⁻ a, f a ∂μ ≤ ⨍⁻ a, g a ∂μ := by
  sorry

This is proven in the Carleson project using the term-mode proof

lintegral_mono_ae <| h.filter_mono <| Measure.ae_mono' Measure.smul_absolutelyContinuous

(implicitly unfolding the average)

The first proof state is:

α : Type u_1
inst✝ : MeasurableSpace α
μ : Measure α
f g : α → ℝ≥0∞
h : ∀ᵐ (a : α) ∂μ, f a ≤ g a
⊢ ⨍⁻ (a : α), f a ∂μ ≤ ⨍⁻ (a : α), g a ∂μ

This gives as second suggestion a useful lemma here: docs#MeasureTheory.setLaverage_eq

After simp_rw [MeasureTheory.laverage_eq] we get the proof state

α : Type u_1
inst✝ : MeasurableSpace α
μ : Measure α
f g : α → ℝ≥0∞
h : ∀ᵐ (a : α) ∂μ, f a ≤ g a
⊢ (∫⁻ (a : α), f a ∂μ) / μ Set.univ ≤ (∫⁻ (a : α), g a ∂μ) / μ Set.univ

here the tool has more trouble, the lemma I want is docs#ENNReal.div_le_div_right, which is the 30th result. Many other lemmas look similar to the goal, but don't apply.

After apply ENNReal.div_le_div_right (or gcongr ?_ / _) the remaining goal is

α : Type u_1
inst✝ : MeasurableSpace α
μ : Measure α
f g : α → ℝ≥0∞
h : ∀ᵐ (a : α) ∂μ, f a ≤ g a
⊢ ∫⁻ (a : α), f a ∂μ ≤ ∫⁻ (a : α), g a ∂μ

and the model finds docs#MeasureTheory.lintegral_mono_ae as its first hit.

[Floris van Doorn]

I think this is a useful tool to find related lemmas, especially when you are not experienced with the library yourself. You still have to heavily judge yourself which lemmas are useful, and which tactic to use to apply the lemmas.

The interface is very clean, and I am excited about the fact that the model already supports two versions of Lean. I hope that means that it will stay updated with a pretty recent (<= 1 month old) version of Mathlib for a while.

[Yicheng Tao]

Floris van Doorn said:

I think this is a useful tool to find related lemmas, especially when you are not experienced with the library yourself. You still have to heavily judge yourself which lemmas are useful, and which tactic to use to apply the lemmas.

The interface is very clean, and I am excited about the fact that the model already supports two versions of Lean. I hope that means that it will stay updated with a pretty recent (<= 1 month old) version of Mathlib for a while.

Thanks for the comment! We have expected that this tool could be more helpful for those who are new to Lean and Mathlib to find useful theorems or just relevant modules. Rather than directly find the exact theorem in your mind, it more likely to provide a bunch of premises that might be useful or instructive, serving as supplementary information.

Our paper introduces a re-ranking procedure to improve the accuracy of the model. But for the concern of computation cost, we didn't adopt re-ranking in this application. So the best-match may not be presented as first few results.

[Yicheng Tao]

We plan to keep our database up-to-date with Mathlib. But our model may not be updated frequently. We may retrain our model every half a year. For now, it was trained on data extracted from v4.10.0 of Mathlib. As to the next update, we can use the newest version along with user feedbacks collected from this website.

[Kim Morrison]

Can we make it callable directly from VSCode, like the #leansearch command? PRs to the https://github.com/leanprover-community/LeanSearchClient repository very welcome!

[Yicheng Tao]

Kim Morrison said:

Can we make it callable directly from VSCode, like the #leansearch command? PRs to the https://github.com/leanprover-community/LeanSearchClient repository very welcome!

We are now serving on a AWS EC2 server with 8 vcpus and 16G memory, which has low load-bearing capacity. So for now we haven't exposed our API. But we can have a try until any issue happens.

[Yicheng Tao]

I have made a PR to LeanSearchClient https://github.com/leanprover-community/LeanSearchClient/pull/13. cc @Siddhartha Gadgil

[Siddhartha Gadgil]

Thanks @Yicheng Tao
Some suggestions:

• Can you add a brief decription to the PR?
• Please move the examples to the LeanSearchClientTest directory like https://github.com/leanprover-community/LeanSearchClient/blob/main/LeanSearchClientTest/MoogleExamples.lean
• For configuration, I feel it is more idiomatic in Lean to use set_option and that is what is done for the other searches. Please add register_option commands to https://github.com/leanprover-community/LeanSearchClient/blob/main/LeanSearchClient/Basic.lean for the Lean version and number of responses like the ones there and use these. It is fine to also override with syntax (though as the Lean version is per project set_option does feel more natural).

[Yicheng Tao]

Siddhartha Gadgil said:

Thanks Yicheng Tao
Some suggestions:

• Can you add a brief decription to the PR?
• Please move the examples to the LeanSearchClientTest directory like https://github.com/leanprover-community/LeanSearchClient/blob/main/LeanSearchClientTest/MoogleExamples.lean
• For configuration, I feel it is more idiomatic in Lean to use set_option and that is what is done for the other searches. Please add register_option commands to https://github.com/leanprover-community/LeanSearchClient/blob/main/LeanSearchClient/Basic.lean for the Lean version and number of responses like the ones there and use these. It is fine to also override with syntax (though as the Lean version is per project set_option does feel more natural).

Thanks for the suggestions. I have committed my changes.

[Yicheng Tao]

BTW, the tactic I implement sends the current main goal to the search engine. The user can not manually modify the query. I did this because I think it would be inconvenient for people to enter the proof state by hand. But it will cost flexibility. If the proof state is some kind of "dirty", the search results will be less satisfying. Any suggestions? Or what do you prefer?

[Siddhartha Gadgil]

I think the current proof state is usually the best thing

[Siddhartha Gadgil]

It looks fine to me now. If @Floris van Doorn and @Kim Morrison think it is ready, one of us can merge.

[Isak Colboubrani]

Are there any plans to integrate this into the premise selection API (https://github.com/leanprover/lean4/pull/7061)?

[Siddhartha Gadgil]

My personal view is that if and when someone hosts a premise selection server, we can add support.

[Siddhartha Gadgil]

Personally, I would be very happy using such a server.

[Yicheng Tao]

For that API, I think the service provider should be an option decided by the user. Relying on any public server may be not a wise choice.