Co-Training Proof Generation and Search Heuristics for Lean 4 Theorem Proving
Learning not only what tactic to try, but which proof state to explore next.
Trained models will be released on Hugging Face soon.
Neural theorem provers have improved dramatically β DeepSeek-Prover-V2 achieves 88.9% on MiniF2F, Goedel-Prover-V2 reaches 90.4%. But they all share one asymmetry: the tactic generator is trained, the search strategy is not.
CoProver fixes this. We co-train two models:
| Tactic Generator | Value Model | |
|---|---|---|
| Model | DeepSeek-Prover-V2-7B | Llama-3.2-1B |
| Task | Generate complete Lean 4 proofs | Score proof states by distance to QED |
| Training | GRPO (Lean verification reward) | SFT on proof trajectories ( |
| Output | Lean 4 code | Scalar |
The value model predicts
βββββββββββββββββββββββββββββββββββββββββββββββ
β CoProver System β
β β
β βββββββββββββββββββββ βββββββββββββββββ β
β β Tactic Generator β β Value Model β β
β β DeepSeek-V2 7B β β Llama 1B β β
β β β β β β
β β Input: theorem β β Input: state β β
β β Output: proof β β Output: Ξ³^d β β
β ββββββββββ¬ββββββββββββ βββββββββ¬ββββββββ β
β β β β
β βΌ βΌ β
β ββββββββββββββββββββββββββββββββββββββββββββ
β β Value-Guided Best-First Search ββ
β β ββ
β β priority = Ξ±Β·logprob + (1-Ξ±)Β·V(state) ββ
β β ββ
β β 1. Pop highest-value state ββ
β β 2. Generator proposes K tactics ββ
β β 3. Lean verifies each tactic ββ
β β 4. Value model scores new states ββ
β β 5. Push scored states to queue ββ
β β 6. Repeat until QED or budget ββ
β βββββββββββββββββββ¬βββββββββββββββββββββββββ
β β β
β βΌ β
β ββββββββββββββββββββββββββββββββββββββββββββ
β β Kimina Lean Server ββ
β β Parallel Lean 4 Proof Checking ββ
β β POST /api/check β valid / invalid ββ
β ββββββββββββββββββββββββββββββββββββββββββββ
βββββββββββββββββββββββββββββββββββββββββββββββ
The two models improve each other through iterative co-training:
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β CO-TRAINING ROUND N β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ£
β β
β Phase 1: PROOF SEARCH β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β Generator attempts proofs on 5,000 theorems β β
β β Value model guides search (if available) β β
β β Lean verifies each attempt β β
β β β Collect trajectories with Ξ³^d labels β β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β Phase 2: TRAIN VALUE MODEL β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β SFT on ALL accumulated trajectories (rounds 0..N) β β
β β Positive: Ξ³^d where d = steps remaining to QED β β
β β Negative: 0.0 (failed proof paths) β β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β Phase 3: TRAIN GENERATOR (GRPO) β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β Whole-proof GRPO: generate proof β Lean verifies β β
β β Reward = 1.0 (verified) or 0.0 (failed) β β
β β 64 prompts Γ 8 samples = 512 per rollout β β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β Phase 4: EVALUATE β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β pass@1, pass@8, pass@32 on MiniF2F / PutnamBench β β
β β Search efficiency: nodes expanded per proof β β
β ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β
β Better search β more proofs β better data β better models β
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
Training the models independently works, but co-training compounds their improvements:
Independent Training: Co-Training:
Generator v1 βββΊ eval Generator v1 βββΊ search βββΊ trajectories
Value v1 βββΊ eval β β
βΌ βΌ
Generator v2 βββΊ eval Value v1 trained βββββ from round 1 data
Value v2 βββΊ eval β
βΌ
(no interaction) Generator v2 βββΊ search with Value v1
β β
βΌ βΌ
Value v2 trained βββββ from round 1+2 data
β
βΌ
(models adapt to each other)
The value model must track the generator β a value model trained on Generator v1's trajectories gives poor scores for Generator v3's states. Co-training keeps them calibrated.
Unlike binary (proved/not proved) labels, we use discounted distance to give the value model a gradient of information:
Proof trace: Stateβ β Stateβ β Stateβ β Stateβ β QED
β β β β
Label: Ξ³β΄=0.81 Ξ³Β³=0.86 Ξ³Β²=0.90 Ξ³ΒΉ=0.95 (Ξ³=0.95)
Failed trace: Stateβ β Stateβ β Stateβ β ERROR
β β β
Label: 0.00 0.00 0.00
This tells the value model not just whether a state leads to a proof, but how close it is. A state two steps from QED (label 0.90) is more valuable than one five steps away (label 0.77).
coprover/
βββ README.md
βββ Makefile # make train-generator, make eval, make co-train
βββ requirements.txt
βββ .gitignore
β
βββ configs/
β βββ generator_grpo.yaml # GRPO hyperparameters
β βββ value_model_sft.yaml # Value model training config
β βββ co_training.yaml # Full co-training loop config
β
βββ models/
β βββ deepseek-prover-v2-7B.sh # Megatron model architecture args
β
βββ ββ Training ββββββββββββββββββββββββββββββββββββββββββββββ
β
βββ train_step_grpo.sh # SLIME GRPO for generator
βββ train_value_slime.sh # SLIME SFT for value model
βββ train_value_model.py # HF Trainer alternative
βββ lean_reward.py # Reward: kimina verification β 1.0/0.0
βββ value_reward.py # Combined: Lean + value model bonus
β
βββ ββ Search & Inference ββββββββββββββββββββββββββββββββββββ
β
βββ value_model.py # Llama-3.2-1B + MLP value head β sigmoid
βββ value_guided_search.py # Best-first search with V(s) scoring
βββ value_rollout.py # Custom SLIME rollout with value guidance
βββ lean_generate.py # Single-tactic generation rollout
β
βββ ββ Data Pipeline βββββββββββββββββββββββββββββββββββββββββ
β
βββ trajectory_collector.py # Generate β verify β record Ξ³^d labels
βββ prepare_minif2f.py # Download MiniF2F + Kimina promptset
βββ prepare_value_data.py # Trajectories β SFT format (with oversampling)
βββ prepare_all.sh # One-shot data preparation
βββ filter_dataset.py # Frontier difficulty filtering
β
βββ ββ Evaluation & Analysis βββββββββββββββββββββββββββββββββ
β
βββ evaluate.py # pass@k with unbiased estimator
βββ compare_checkpoints.py # A/B model comparison
βββ analyze_training.py # Parse SLIME logs, plot curves
β
βββ ββ Orchestration βββββββββββββββββββββββββββββββββββββββββ
β
βββ co_train.py # Multi-round co-training orchestrator
βββ co_train.sh # Docker wrapper with server lifecycle
βββ launch_servers.sh # Start kimina + SGLang with health checks
βββ convert_and_serve.sh # Megatron β HuggingFace β SGLang
βββ check_status.sh # GPU, server health, latest metrics
βββ run_docker.sh # Launch SLIME container
β
βββ ββ Infrastructure ββββββββββββββββββββββββββββββββββββββββ
β
βββ lean_server.py # LeanDojo verification HTTP server
βββ lean_server_pool.py # Persistent Lean REPL pool with LRU cache
If you want to exercise the prompt and verification path without downloading the
full MiniF2F split, use the checked-in fixture at examples/minif2f_smoke.jsonl.
It follows the same formal_statement / prompt schema produced by
prepare_minif2f.py, so it is safe to swap into evaluate.py or other
whole-proof scripts.
python3 evaluate.py \
--dataset examples/minif2f_smoke.jsonl \
--sglang-url http://localhost:30000 \
--kimina-url http://localhost:8000 \
--n-samples 4 \
--max-problems 12 \
--output results/minif2f_smoke.jsonThe fixture is intentionally tiny:
nat_add_comm_smallis a familiar arithmetic theorem that should verify cleanly when the model emits a standardsimpa [Nat.add_comm]proof.list_append_nil_smallchecks that the model can finish a simple theorem overList, which is a useful sanity check before moving to MiniF2F-scale evaluation.and_left_smallexercises proposition-style proof search with a named hypothesis, which catches prompt or verifier regressions that equality-only smoke tests can miss.and_intro_smallchecks conjunction construction from two named hypotheses, which is a compact way to exercise proposition-level constructor proofs.exists_witness_smallchecks witness construction for an existential goal, which is a lightweight way to catch regressions in constructor-style proofs before moving on to full MiniF2F preparation.or_inl_smallchecks disjunction introduction from a named hypothesis, which is a compact sanity check for proposition-level branching before running the full dataset.or_elim_smallchecks disjunction elimination with two continuations, which is a lightweight way to exercise case-split reasoning before scaling up to full MiniF2F prompts.imp_chain_smallchecks a tiny implication chain (P -> Q -> R), which is a good smoke test forintroplus hypothesis application before running larger theorem sets.eq_rewrite_succ_smallchecks direct rewriting through an equality hypothesis, which is a compact sanity check forrw/simpa-style proofs before moving on to larger arithmetic or algebra examples.false_elim_smallchecks explosion from a contradictory hypothesis, which is a lightweight sanity check forexfalso/False.elim-style proofs before scaling to larger proposition-heavy benchmarks.iff_forward_smallchecks that the prover can project the forward direction of an equivalence, which is a compact sanity check forIff.mp/h.mpstyle reasoning before moving on to larger proposition-heavy benchmarks.exists_elim_rewrite_smallchecks existential elimination followed by rewriting a hypothesis onto the target index, which is a useful smoke test forrcasesplusrw atstyle proof steps before scaling up to larger theorem sets.
For the full benchmark and Kimina promptset, run:
# 1. Pull images
docker pull slimerl/slime:latest
docker pull projectnumina/kimina-lean-server:2.0.0
# 2. Start Lean verification server
docker run -d --name kimina-lean-server \
--ulimit nofile=65536:65536 \
-p 8000:8000 --restart unless-stopped \
projectnumina/kimina-lean-server:2.0.0
# 3. Enter SLIME container
bash run_docker.sh
# This shell can reach host-side services like kimina/SGLang via host.docker.internal.
# or override defaults when your formal workspace lives elsewhere
WORKSPACE=/mnt/filesystem-m5/formal DOCKER_IMAGE=slimerl/slime:latest bash run_docker.sh
# 4. Convert model weights to Megatron format
source models/deepseek-prover-v2-7B.sh
PYTHONPATH=/root/Megatron-LM python /root/slime/tools/convert_hf_to_torch_dist.py \
${MODEL_ARGS[@]} \
--hf-checkpoint /workspace/models/DeepSeek-Prover-V2-7B \
--save /workspace/models/DeepSeek-Prover-V2-7B_torch_dist
# 5. Train the generator
bash train_step_grpo.sh
# 6. Run full co-training loop
bash co_train.sh --num-rounds 5| Parameter | Value |
|---|---|
| Base model | DeepSeek-Prover-V2-7B (30 layers, 4096 hidden) |
| Learning rate | 1Γ10β»βΆ (constant) |
| Batch size | 512 (64 prompts Γ 8 samples) |
| Reward | Binary Lean verification |
| Clipping | Ξ΅_low=0.2, Ξ΅_high=0.28 |
| Temperature | 1.0 |
| Max tokens | 4,096 |
| GPUs | 4 train (TP=2) + 4 rollout (SGLang) |
| Framework | SLIME (Megatron-LM + SGLang) |
| Parameter | Value |
|---|---|
| Base model | Llama-3.2-1B (16 layers, 2048 hidden) |
| Learning rate | 2Γ10β»β΅ (cosine with warmup) |
| Labels | Ξ³^d where d = steps to QED, Ξ³=0.95 |
| Positive oversampling | 3Γ |
| Epochs | 3 per round |
| GPUs | 1 |
Why whole-proof GRPO over step-level? DeepSeek-Prover-V2-7B generates complete proofs naturally. Step-level GRPO requires a separate Lean REPL session per tactic, which is 100Γ slower than whole-proof verification via kimina-lean-server.
Why no KL penalty? With only 3-8% reward rate on hard theorems, KL regularization overly constrains exploration. The model needs freedom to discover novel proof strategies. We tried KL coefficients of 0.001 and 0.0001 β both stagnated training.
Why Ξ³^d over binary labels? Binary labels (proved/not proved) throw away information about how close a failed attempt was. Ξ³^d preserves the gradient β the value model learns that "3 steps from QED" is better than "10 steps away", even when neither attempt completed.
Why Llama-3.2-1B for value model? Small enough for fast inference during search (must score every expanded state), large enough to understand proof state structure. Served on SGLang alongside the generator with negligible overhead.
# MiniF2F (244 problems)
python evaluate.py \
--dataset data/minif2f_test.jsonl \
--sglang-url http://localhost:30000 \
--kimina-url http://localhost:8000 \
--n-samples 32 --output eval_results/minif2f.json
# Compare two checkpoints
python compare_checkpoints.py \
--sglang-url-a http://localhost:30000 \
--sglang-url-b http://localhost:30001 \
--label-a "baseline" --label-b "co-trained" \
--dataset data/minif2f_test.jsonl \
--n-samples 32 --output eval_results/comparison.json| System | Generator Training | Search | Value/Critic | Co-Training |
|---|---|---|---|---|
| DeepSeek-Prover-V2 | GRPO (whole-proof) | Fixed tree search | No | No |
| InternLM2.5-StepProver | Expert iteration | Critic-guided BFS | Yes (separate) | No (one round) |
| Kimina-Prover | GRPO (whole-proof) | Fixed | No | No |
| BFS-Prover | Expert iteration | Fixed BFS | No | No |
| AlphaProof | RL (unpublished) | MCTS | Yes (hand-tuned) | Unclear |
| CoProver (ours) | GRPO | Learned value-guided | Yes (co-trained) | Yes (iterative) |
- 8Γ NVIDIA H200 (140 GB VRAM each)
- 4 GPUs for Megatron-LM training (tensor parallelism 2)
- 4 GPUs for SGLang rollout generation
- kimina-lean-server on host CPU (Docker)
- ~3 min per GRPO training step (512 samples)
@misc{coprover2026,
title={CoProver: Co-Training Proof Generation and Search Heuristics for Automated Theorem Proving},
year={2026},
howpublished={\url{https://github.com/coprover}},
}Proprietary. All rights reserved. Models will be released on Hugging Face.