Encoder (ModernBERT-large)

Fast Lean 4 validity scoring with a compact 395M-parameter encoder. Rank and screen proof candidates using a validity probability from a single forward pass.

Paper throughput: 264 candidates/s on one H100.

Quickstart

Score a complete Lean program with the sequence-classification head:

pip install torch 'transformers==4.57.3' accelerate
hf download EPFL-LARA/lean-encoder-modernbert-large example_transformers.py predict.py Example.lean --local-dir lean-encoder
cd lean-encoder
python example_transformers.py

Hugging Face Transformers

The complete code in example_transformers.py is below. Replace source_code with your candidate program, including its imports, context, and proof:

import json

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from predict import encode_head_tail

MODEL_ID = 'EPFL-LARA/lean-encoder-modernbert-large'
source_code = """import Mathlib

theorem nat_add_zero_identity (n : ℕ) : n + 0 = n := by
  rfl
"""

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_ID, dtype=dtype, attn_implementation="sdpa",
).to(device).eval()

inputs = encode_head_tail(
    tokenizer, [source_code], max_length=8192, head_fraction=0.25,
)
inputs = {key: value.to(device) for key, value in inputs.items()}
with torch.inference_mode():
    probabilities = model(**inputs).logits[0].float().softmax(dim=-1)

probability_valid = float(probabilities[1])  # Class 0: INVALID; class 1: VALID.
print(json.dumps({
    "is_valid": probability_valid >= 0.5,
    "valid_probability": probability_valid,
    "invalid_probability": float(probabilities[0]),
}, indent=2))

The JSON output contains is_valid, valid_probability, and invalid_probability. Class 0 is INVALID; class 1 is VALID.

Command-line inference

The downloaded Example.lean contains the same input. To score a file instead:

python predict.py Example.lean

The command-line runner also accepts Lean source on standard input.

Input and scoring

Supply the complete, unannotated Lean source, including imports and context.
The runner uses AutoModelForSequenceClassification and the training preprocessing: add_special_tokens=False, an 8,192-token context, and 25% head / 75% tail selection for long programs.
Class 0 is INVALID; class 1 is VALID. The default decision is valid_probability >= 0.5.

The reusable encode_head_tail helper is included in predict.py.

Results

Paper results on the 68,932-program evaluation benchmark, using Lean v4.15.0 labels from public-prover candidates on miniF2F and ProofNet:

Accuracy	Macro F1	MCC	AUROC
58.6%	0.563	0.372	0.862

Throughput is measured as batch throughput under the paper's serving configuration.

Training

Based on answerdotai/ModernBERT-large, with mean pooling and a binary classification head. Lean continued pretraining on approximately 130M tokens is followed by one supervised epoch on source_code → isValid pairs from the compiler-labelled training corpus.