| --- |
| license: apache-2.0 |
| library_name: pytorch |
| tags: |
| - lean-4 |
| - mathlib |
| - theorem-proving |
| - tactic-prediction |
| - gnn |
| - graph-neural-network |
| datasets: |
| - jajostrains/Mathlib-Normalized-Sexpr |
| language: |
| - en |
| --- |
| |
| # Mathlib-Sexpr-GNN — GNN tactic and argument prediction for Lean 4 |
|
|
| Graph neural networks that predict **which tactic to apply** to a Lean 4 proof |
| goal, and **which arguments** (local hypotheses or library lemmas) to give it. |
| Trained on 250,814 tactic steps from Mathlib, represented as S-expression |
| goal-state DAGs from [jajostrains/Mathlib-Normalized-Sexpr](https://huggingface.co/datasets/jajostrains/Mathlib-Normalized-Sexpr). |
|
|
| Two bundles are published here, each a self-contained directory verified by |
| content hashes (weights, node vocabulary, tactic vocabulary): |
|
|
| | Directory | Model | What it predicts | |
| |---|---|---| |
| | `tactic-baseline-gat/` | GATv2 encoder + classifier | tactic family only (e.g. `rw`, `apply`, `simp`) | |
| | `pointer-gat-gru/` | GATv2 encoder + GRU pointer decoder + stop head, with a premise scorer | tactic family **and** its ordered argument sequence, plus a premise scorer for library-lemma retrieval | |
|
|
| ## How the models work |
|
|
| - **Goal representation.** Each proof state (goal + local context) is parsed |
| into a directed acyclic graph: expression nodes, hypothesis nodes with |
| binder annotations, and a State root. A GATv2 encoder (4 layers, 8 heads, |
| hidden 128, state+mean-attention readout) produces node embeddings. |
| - **Tactic head.** A classifier over the state embedding predicts one of 228 |
| tactic families. |
| - **Argument decoder (pointer bundle).** A GRU recurrent decoder runs up to 3 |
| decode steps. At each step it scores every premise-selectable DAG node |
| (pointer attention over node embeddings), picks one, and feeds it back into |
| the recurrence. A **stop head** decides when the argument sequence ends — |
| the count is learned, not taken from a static arity table. Already-selected |
| nodes are masked so no argument repeats. |
| - **Premise scorer.** A dot-product scorer over a unified candidate pool |
| (local DAG nodes + FAISS-retrieved library lemmas), for tactics that cite |
| Mathlib lemmas (`exact`, `apply`, `rw`). |
|
|
| ## Metrics |
|
|
| Validation split (4,260 tactic steps), pointer bundle: |
|
|
| | Metric | Value | |
| |---|---| |
| | Tactic top-1 / top-5 | 0.320 / 0.711 | |
| | Argument top-1 / top-5 (position-level) | 0.433 / 0.779 | |
| | Exact argument-sequence match | 0.037 | |
| | Stop-decision accuracy | 0.789 | |
| | Premise scorer top-1 / top-5 / MRR | 0.481 / 0.806 / 0.618 | |
|
|
| Test split: tactic top-1 0.315, premise top-1 0.468, premise MRR 0.600. |
|
|
| Baseline bundle: validation tactic top-1 0.327 / top-5 0.701; test top-1 0.311. |
|
|
| Caveat on the argument metrics: they are computed over the positions a |
| DAG-pointer can physically select — local-context targets within the 3-step |
| decode budget. Library-lemma citations are the premise scorer's task and are |
| excluded from the pointer's accuracy; longer argument lists (the corpus |
| median exceeds 3 for `rw`-family chains) are truncated at the budget and |
| reported separately. |
|
|
| ## Usage |
|
|
| Each bundle directory holds `bundle.json` (manifest with hashes and |
| metadata), `model.safetensors`, `config.json`, and the scorer (pointer |
| bundle only). The shared `vocab/` directory at the repo root holds the |
| canonical node and tactic vocabularies both bundles were trained against; |
| the manifests hash-bind the weights to them. |
|
|
| Loading with the project's API: |
|
|
| ```python |
| from maths_ai.gnn_inference.atp_lean_gnn.bundle import load_pointer_bundle |
| |
| loaded = load_pointer_bundle("pointer-gat-gru", device="cpu") |
| # loaded.model — TacticWithArgsClassifier (GATv2 + GRU pointer decoder) |
| # loaded.scorer — PremiseScorer (dot mode) |
| ``` |
|
|
| Bundle invariants enforced on load: weight and vocabulary hashes are |
| re-verified against `bundle.json`; a `pointer_gru` bundle cannot be loaded |
| as a legacy pointer or vice versa; a vocabulary of the same size but |
| different content is rejected (IDs are position-assigned, so a renumbered |
| vocabulary would silently mislabel every prediction otherwise). |
|
|
| ## Training data and provenance |
|
|
| - 250,814 train / 4,260 val / 4,506 test tactic steps from the LeanDojo |
| Mathlib 4 benchmark split (row selection, splits, and row indices inherited |
| unchanged from the upstream benchmark). |
| - Goal states elaborated with a pinned Pantograph REPL against Mathlib at the |
| benchmark's commit; S-expressions and per-position argument traces extracted |
| per row. |
| - Trained from a baseline-initialized checkpoint (tactic head warm-started |
| from `tactic-baseline-gat`'s encoder), 30 epochs, bfloat16 AMP. |
|
|
| ## Limitations |
|
|
| - Argument sequences are capped at 3 steps; longer `rw` chains are truncated. |
| - The premise scorer's FAISS index (318,953 Mathlib declarations) is **not** |
| published with this model — it is not yet hash-bound to the encoder that |
| produced its vectors, so serving it risks confidently wrong retrieval. It |
| will follow once that binding is enforced at load time. |
| - Closed-loop proof search (running the model against Lean) is not yet |
| evaluated; these are per-state prediction metrics. |
|
|
| ## License |
|
|
| Apache-2.0, inherited from Mathlib. The dataset carries the LeanDojo MIT |
| notice for the benchmark's row selection and split structure; see the |
| dataset card for details. |
|
|
|
|
