---
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.