EPFL-LARA/lean-feedback-qwen3-4b

Feedback (Qwen3-4B)

Lean 4 compiler feedback on GPUs. Given a Lean program, this Qwen3-4B model generates a feedback trace: the source annotated with predicted goal states and compiler diagnostics. Use the trace to guide proof development and derive a validity prediction with the included parser.

Paper results: 85.1% validity accuracy with greedy decoding; 87.4% with five-sample self-consistency.

Quickstart

Generate the annotated Lean program, then derive a validity prediction from its feedback trace:

pip install torch 'transformers==4.57.3' accelerate
hf download EPFL-LARA/lean-feedback-qwen3-4b example_transformers.py predict.py feedback_validity.py Example.lean --local-dir lean-feedback
cd lean-feedback
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 AutoModelForCausalLM, AutoTokenizer

from predict import format_source, VERDICT_STOPS
from feedback_validity import predict_validity

MODEL_ID = 'EPFL-LARA/lean-feedback-qwen3-4b'
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 = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, dtype=dtype, attn_implementation="sdpa",
).to(device).eval()

prompt = format_source(source_code)  # source_code.rstrip() + "\n\n"
inputs = tokenizer(
    prompt, add_special_tokens=False, return_tensors="pt",
).to(device)
prompt_length = inputs["input_ids"].shape[1]
budget = min(16384, 32768 - prompt_length - 8)

if budget < 1:
    result = {"is_valid": False, "finish_reason": "prompt_overflow", "prediction": None}
else:
    with torch.inference_mode():
        output = model.generate(
            **inputs, max_new_tokens=budget, do_sample=False,
            stop_strings=VERDICT_STOPS, tokenizer=tokenizer,
            pad_token_id=tokenizer.pad_token_id,
        )
    generated_ids = output[0, prompt_length:]
    prediction = tokenizer.decode(generated_ids, skip_special_tokens=True)
    eos_ids = model.generation_config.eos_token_id or tokenizer.eos_token_id
    eos_ids = [eos_ids] if isinstance(eos_ids, int) else eos_ids
    stopped = int(generated_ids[-1]) in eos_ids or any(
        marker in prediction for marker in VERDICT_STOPS
    )
    finish_reason = "stop" if stopped else "length"
    result = {
        "is_valid": predict_validity(prediction, finish_reason=finish_reason),
        "finish_reason": finish_reason,
        "prediction": prediction,
    }

print(json.dumps(result, ensure_ascii=False, indent=2))

The JSON output contains the generated trace in prediction, its finish_reason, and the parsed is_valid decision. predict_validity applies the feedback parser and the generation-limit policy. The Boolean is computed from the generated trace.

Command-line inference

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

pip install 'vllm==0.12.0'
python predict.py Example.lean --include-generated-feedback

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

Input and validity parsing

Supply the complete, unannotated Lean source, including imports and the theorem statement. The prompt is exactly source_code.rstrip() + "\n\n", tokenized as plain text with add_special_tokens=False.

The model generates annotated source; the Boolean verdict is computed from that feedback. The parser checks the error-marker regex --\s*type\s*:\s*error\b, terminal proof-state blocks, and the generation's stopping reason. Intermediate proof states are part of a normal trace.

The runner follows the paper's rule: a completed trace is valid when it has neither an error marker nor an unfinished terminal proof state. Length-limited generations and overlength inputs receive an invalid verdict. Total context is 32,768 tokens, with up to 16,384 output tokens within that budget.

The reusable parser is included in feedback_validity.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:

Feedback decoding	Accuracy	Macro F1	MCC
Greedy (quickstart default)	85.1%	0.851	0.733
Five-sample self-consistency	87.4%	0.873	0.762

Self-consistency samples five traces at temperature 0.8 and requires four valid votes, a rule selected on separate in-house validation data. Its validity vote score achieves 0.919 AUROC.

On the paper's 2,004-program feedback audit, overall exact block F1 is 0.796 and token F1 is 0.852. The paper measures 3.27 candidates/s on one H100 with speculative decoding as batch throughput; the quickstart uses standard greedy generation.

Pair it with the One-token model for the cascade: 87.8% accuracy with greedy feedback and 88.8% with five-sample feedback.

Training

Qwen3-4B receives Lean continued pretraining, followed by supervised training on 4,836,380 programs from the compiler-labelled training corpus. It learns source_code → annotated validation, with the training loss applied to the generated response.

Explore the project

Part of Learning Lean Feedback for Efficient Proof-Candidate Evaluation on GPUs.

Paper collection · Encoder model · One-token model