Combining Textual and Structural Information for Premise Selection in Lean

AI summary

Overview

Research area: Machine learning for formal mathematics — specifically automated premise selection for the Lean 4 interactive theorem prover, combining language-model embeddings with graph neural networks.

Technical level: Advanced. The paper assumes familiarity with dense retrieval (dual encoders), relational graph convolutional networks, InfoNCE contrastive losses, and the tactic-mode proof structure of Lean 4.

Scope: The paper builds a heterogeneous dependency graph from Lean/Mathlib source files, refines ByT5 text embeddings of premises and proof states with an RGCN, and evaluates retrieval on the LeanDojo Benchmark against the ReProver baseline.

What This Paper Is About

When an AI system tries to prove a theorem in Lean 4, it must first guess which existing definitions and lemmas from the library will be useful — a task called premise selection. Existing language-based systems like ReProver look at each premise in isolation, embedding its text with a ByT5 dual encoder, and ignore the network of references connecting premises to each other and to proof states. This paper asks whether feeding that structural information into the model improves retrieval.

Key Contributions

1. An augmented data extraction pipeline. The authors extend LeanDojo's dataset extraction to query the Lean proof assistant for reference information, building a heterogeneous directed dependency graph over the Mathlib library that captures state–premise and premise–premise relations in Lean 4 tactic mode.

2. A typed relation taxonomy. References are categorized by where they occur: premise names, signature local hypotheses, signature goals, and proof dependencies, with a corresponding set of relation types for proof states (local hypotheses, goals, next-tactic premise labels).

3. A graph-augmented retrieval architecture. A relational graph convolutional network (RGCN) is used to propagate structural information and produce graph-aware embeddings for both premises and proof states, trained with an InfoNCE contrastive objective.

4. An empirical comparison and a candid ablation. The method is compared to the ReProver baseline on the LeanDojo Benchmark, and ablation experiments reveal that removing parts of the dependency graph does not degrade — and sometimes improves — performance.

Main Findings

• Strong headline gains over ReProver. On the LeanDojo test set, the GNN-augmented retriever reaches R@1 of 17.98% (a relative gain of +33.98%), R@10 of 50.04% (+26.36%), and MRR of 0.4095 (+24.73%), compared with the ReProver baseline at 13.42%, 39.60%, and 0.3283 respectively.

• EMA adds a further small improvement. Adding exponential model averaging lifts the results to R@1 of 18.31% (+36.44%), R@10 of 50.33% (+27.10%), and MRR of 0.4140 (+26.10%).

• The baseline numbers differ from the original ReProver paper. The authors note that their baseline results differ from those reported in Yang et al. (2023), and that the ReProver authors attribute this on GitHub to improvements made to the ReProver system after publication.

• Ablations complicate the story. With no ensembling used for fair comparison, the full model scores R@1 of 17.43%, R@10 of 48.52%, and MRR of 0.4010; with EMA, 17.74%, 48.70%, and 0.4048. Removing the context graph yields 17.30%, 49.99%, 0.4008 (and 17.58%, 50.13%, 0.4057 with EMA); removing the premise graph yields 17.45%, 49.30%, 0.4008 (and 18.18%, 49.98%, 0.4096 with EMA). The authors state that removing parts of the dependency graph even improves performance in places.

• Alternative explanation offered for the gains. Because of the ablation outcomes, the authors raise the possibility that the improvement stems not from graph structure but from their different loss function and sampling strategy, including the decision to use all other premises in the library as negatives rather than sampling them.

• Graph scale. The augmented dependency graph contains 440,487 nodes total: 180,907 premise nodes and 259,580 state nodes. Premise-to-premise edges total 1,769,158 (652,484 signature local hypotheses, 626,318 signature goal, 490,356 proof dependency), and premise-to-state edges total 4,210,203, plus 379,861 next-tactic premise labels.

Methodology in Plain English

The starting point is ReProver's dual encoder: premises and proof states are each turned into text and mapped into a shared vector space using ByT5, so that a relevant premise ends up close to the state that needs it. The authors keep that step but add a second stage.

They extend LeanDojo's extraction to record, for every premise and proof state, which other premises it refers to and where — in the signature's hypotheses, in the goal, or in the proof body. This produces a directed graph with typed edges over all premises and proof states, which they summarize as a text-attributed directed graph.

An RGCN then runs over this graph, updating each node's vector by aggregating messages from its neighbors, with separate weight matrices per relation type and normalization by relation-specific neighbor counts. The number of layers is L = 2. Proof states are handled as temporary query nodes connected to the premises they reference, using the same architecture with different weights. Edge directions prevent information from flowing from states into dependencies or from future premises into earlier ones.

Training uses an InfoNCE loss with a temperature hyperparameter, pulling each state close to its valid premises and away from all others. The GNN is trained transductively on the full premise graph, but all edges from proof relations are excluded to prevent trivial memorization. The final system averages six independently trained models and applies exponential model averaging. Hyperparameters were tuned in two stages with Optuna — first architecture, then regularization — using Recall@10 on a training batch and then on the validation set. Final configuration: hidden size 1024, ReLU activation, dropout 0.256, edge dropout 0.142, residual connections, InfoNCE temperature 0.0138, learning rate 0.00499, weight decay 2.359e-5, batch size 1024, gradient accumulation over 2 batches, 120 epochs. Experiments used a cluster with three NVIDIA A6000 GPUs (48 GB each), two Intel Xeon Silver 4410Y CPUs (24 cores, 48 threads), and 512 GB of RAM.

Why This Matters

Research impact. Premise selection is described as a key bottleneck for scaling theorem proving in large formal libraries. If structural dependencies among premises carry useful signal beyond their text, that would change how retrieval systems for proof assistants are designed. This paper presents encouraging evidence but also a caution: its own ablations suggest the benefit may come from the training paradigm rather than the graph, which is an important methodological flag for the subfield.

Real-world applications:

• AI-assisted theorem proving in Lean 4 and related interactive provers, where retrieval quality directly affects whether proofs can be finished automatically.
• Tools like Sledgehammer, which the paper cites as an automated reasoning tool underpinned by premise selection.
• Formal verification of software and mathematics, where large libraries such as Mathlib must be searched efficiently.
• General retrieval from large technical codebases and formal knowledge bases, where reference graphs are available in addition to text.

Industry relevance. The paper notes that GNNs use relatively less memory than LLMs, which removes the need to sample negative examples — a practical efficiency argument for graph-based retrieval components inside larger AI proving pipelines. The extraction tooling and learning pipeline are released at https://github.com/JobPetrovcic/GNNReProver/tree/lighweight.

Future Directions

• Determine the true source of improvement. The authors explicitly leave for future work the question of whether the gains come from graph structure, the InfoNCE loss, or the sampling strategy — namely using all other library premises as negatives.
• Try stronger graph architectures. They propose exploring advanced GNN designs such as graph attention networks to better incorporate structural information.
• Evaluate on harder splits. The current results use the "random" split; the authors plan to test on the LeanDojo "novel" split or a split based on the creation time of premises.
• Use the file and import graph during training. That graph is currently only used at evaluation time to restrict premises to those available at that point in the file (imports and premises defined before the current entry); its use in training is deferred to future work.

Target Audience

Researchers and graduate students working at the intersection of machine learning, automated theorem proving, and interactive proof assistants — particularly those familiar with Lean 4, Mathlib, and retrieval-based premise selection. It is also relevant to practitioners building neural provers who want a concrete, reproducible comparison to ReProver, and to methodologists interested in a paper that honestly reports ablations weakening its own central claim. Readers without a background in graph neural networks or formal mathematics will find it demanding, though the high-level problem statement is accessible.

Authors’ abstract

Premise selection is a key bottleneck for scaling theorem proving in large formal libraries. Yet existing language-based methods often treat premises in isolation, ignoring the web of dependencies that connects them. We present a graph-augmented approach that combines dense text embeddings of Lean formalizations with graph neural networks over a heterogeneous dependency graph capturing both state-premise and premise-premise relations. On the LeanDojo Benchmark, our method outperforms the ReProver language-based baseline by over 25\% across standard retrieval metrics. These results suggest that relational information is beneficial for premise selection.