This project focuses on enhancing Lean 4's automated tactics, specifically grind and aesop, through neural-guided proof search and efficiency-aware data collection.
The neural_grind tactic replaces grind's native split-candidate heuristic with a lightweight MLP model (~1M parameters) that ranks branching options by predicted proof efficiency.
Import:
import NeuralTacticUsage:
example (h : P ∨ Q) (hp : P → R) (hq : Q → R) : R := by
neural_grindStarting the inference server (required once per session):
Native binary (fastest):
GRIND_NATIVE_WEIGHTS=training/experiments/exp09_heuristics/model.native.bin \
./NeuralTactic/native_serveOr via Python:
python3 training/serve.py --model training/experiments/exp09_heuristics/model.ptKey environment variables (all optional):
| Variable | Effect |
|---|---|
GRIND_NATIVE_WEIGHTS |
Path to .native.bin weights for the C++ server |
GRIND_MODEL |
Path to PyTorch .pt checkpoint (Python server) |
GRIND_SERVE |
Path to serve.py or native executable |
GRIND_SERVE_NATIVE |
Set to 1 when GRIND_SERVE is the native binary |
GRIND_NO_MODEL |
Set to 1 to disable model and use stock grind heuristic |
GRIND_MARGIN_MILLI |
Minimum logit margin (milli-logits) required to trust model ranking |
GRIND_DECISION_LOG |
Path to a JSONL file for per-decision diagnostic output |
If the server is unavailable or GRIND_NO_MODEL=1, neural_grind falls back to grind's native heuristic. Proofs that close with stock grind will also close with neural_grind; the model only changes split ordering, not soundness. To use stock grind, simply call grind as usual — the two tactics are independent.
aesop_with_overrides runs standard aesop but reassigns the success probabilities of unsafe rules before the search begins. The probabilities come from a graph-based topological scoring of ~53,000 empirical rule-application traces: rules that tended to be chosen earlier in successful proofs receive higher scores, so aesop tries them first.
Import:
import AesopWithOverridesUsage:
example (s t : Finset α) (h : s ⊆ t) : s ∩ t = s := by
aesop_with_overridesSetting the overrides file (required; without it the tactic is identical to aesop):
export AESOP_OVERRIDES_JSON=/path/to/aesop_rule_ordering/aesop_overrides.jsonThe JSON file maps fully-qualified declaration names to floats in [0, 1]:
{ "Exists": 0.833, "Or": 0.833, "Aesop.BuiltinRules.ext": 0.5 }Only unsafe rules are affected — safe and norm rules use integer penalties and cannot have their success probabilities overridden.
Custom @[aesop ...] attributes work exactly as with stock aesop; the tactic accepts the same clauses (add, erase, rule_sets, config, simp_config). An aesop_with_overrides? variant (analogous to aesop?) is also available to emit a suggested proof script.
Our data collection process followed two primary strategies to capture a diverse range of mathematical contexts and theorem styles.
Mathlib Data Collection:
We collected data from Mathlib by systematically extracting theorems and their environmental contexts. This was achieved by parsing import statements, namespaces, and open statements to reconstruct the precise mathematical state. We utilized a custom tactic to execute grind and aesop across these theorems, extracting the goal states and proof traces to build a comprehensive map of Mathlib's successful automated proofs.
External Dataset Integration: The second strategy involved leveraging several large-scale Lean datasets, collectively providing a vast search space for theorem automation:
- NuminaMath Lean: 104,155 competition-style theorems.
- Lean Workbook: 57,231 formal statements.
- Herald: 580,000 NL-FL statement pairs.
- FineLeanCorpus: 509,358 formalization entries.
- LeanDojo-v2: 120,000 statements
These datasets provide a vast collection of theorem statements that typically do not yet have accompanying proofs. From this pool, we have successfully verified and extracted proof traces for approximately 53,000 theorems that work with either grind or aesop.
This collection forms the core of our training data, representing a diverse "Winning Spine" for both grind and aesop across different domains.
Verification and Filtering:
For grind, we tested each theorem to see if the tactic could autonomously close the goal. Successful cases were added to a JSONL dataset for training. This verification process was facilitated by AXLE by axiom.ai.
For aesop, we employed a similar verification protocol but introduced a refinement step: we first attempted to solve the goal using simp. If simp failed but aesop subsequently succeeded, the theorem was added to our collection. This filtering ensures that our dataset prioritizes non-trivial proofs where aesop's search capabilities provide unique value beyond standard simplification, reflecting its internal strategy of attempting simplification before proceeding with more complex search heuristics.
Aesop trace extraction was performed with our aesop_collect tactic, which functions identically to aesop but when it closes the goal it additionally emits a message documenting the final set of rule applications that lead aesop to the goal, including the state at that application, the rule itself, and the other options that aesop had that it did not choose.
We instrumented the grind tactic to capture the internal state and decision-making process at every "split" (branching point) during proof search. This instrumentation allows us to record the precisely timed "snapshots" of the prover's state.
The Trace Data: For every successful proof, we extract a JSON trace containing:
- Goal Features: Numerical summaries such as
assertedCount(total known facts),ematchRounds, andsplitDepth. - Candidate Pool: The full set of available branching options (e.g., case splits on an
if-then-elseor an implication), including their origins (ematch,ext,input, etc.). - Environmental Context: A "Pretty-Printed" version of the local context and the internal state of the E-graph (the robot's "brain").
- Winning Spine: We prune these traces to retain only the "Winning Spine"—the sequence of successful decisions that directly led to the proof, discarding any exploratory dead ends.
Data Preparation
- We collect all the messages emitted by
aesop_collectin the training data (we replace all instances ofaesopin the training data withaesop_collect). We split this into a train set and a test set. These collections can be found inaesop_data_split/train.json, andaesop_data_split/test.json. - We further split the messages into their individual rule applications, found in
aesop_data_split/train_pairs.json, andaesop_data_split/test_pairs.json.
Training
These steps correspond to the python files marked 01 through 05 in aesop_rule_ordering/:
- We build a directed graph representing our training data. Each node represents a rule and an edge from node a to node b with weight w indicates that there are w instances in the training data in which rule a was
chosenand rule b appeared inallowedUnsafeRules. - We remove two-cycles by doing the following: wherever there is an edge from a to b with weight n and there is an edge from b to a with weight k, such that n > k, we replace both edges with a new edge from a to b that has weight n/k.
- We remove the rest of the cycles from the graph by applying a custom ILP solver.
- We separate each connected component of the graph into layers. Nodes on layer 0 are nodes with no edges into them. Nodes on layer 1 have only edges from nodes on layer 0. Nodes on layer 2 have only edges from layers 0 and 1. And so on.
- We produce an override JSON based on the graph. If a rule's node is in layer n of a connected component with N layers, then we predict its override value as being (N - n) / (N + 1).
To leverage the extracted trace data, we trained a lightweight neural model designed to predict the optimal next step in a grind proof search.
Model Architecture:
- Type: A 3-layer Multi-Layer Perceptron (MLP).
- Size: ~74,500 parameters
- Input: A 32-dimensional feature vector encoding the goal state, the specific candidate branch, and its relationship to the rest of the candidate pool.
- Inference: Optimized for speed, the model runs in a native C++ server with an inference time of < 1ms, enabling its use directly within the Lean tactic loop.
Training Strategy:
- Objective: The model is trained using Cross-Entropy Loss to maximize the score of the "Winning Spine" candidate within its original pool of alternatives.
-
Efficiency Weighting: We apply a weighting formula (
$W = 1 / \sqrt{\text{total \_ splits}}$ ) to the training loss. This incentivizes the model to prioritize decisions from shorter, more efficient proofs, effectively learning "shortcuts" over the native heuristic's often exhaustive search. - Evaluation: We maintain a stable benchmark of 75 problems (balanced across Mathlib, Workbook, and Numina) that are strictly excluded from the training set to measure the model's generalization and split-reduction capabilities.
This project serves as a focused investigation into the scaling potential of neural-guided proof search. Our primary interest was in projecting what the results would look like in the "actual end"—leveraging the full scale of the 53,000+ theorem dataset to move toward a highly efficient, zero-exploration automated prover.