Skip to content

mike1729/BurnQED

BranchesTags

Repository files navigation

 ██████╗ ██╗   ██╗██████╗ ███╗   ██╗       ██████╗ ███████╗██████╗
 ██╔══██╗██║   ██║██╔══██╗████╗  ██║      ██╔═══██╗██╔════╝██╔══██╗
 ██████╔╝██║   ██║██████╔╝██╔██╗ ██║█████╗██║   ██║█████╗  ██║  ██║
 ██╔══██╗██║   ██║██╔══██╗██║╚██╗██║╚════╝██║▄▄ ██║██╔══╝  ██║  ██║
 ██████╔╝╚██████╔╝██║  ██║██║ ╚████║      ╚██████╔╝███████╗██████╔╝
 ╚═════╝  ╚═════╝ ╚═╝  ╚═╝╚═╝  ╚═══╝       ╚══▀▀═╝ ╚══════╝╚═════╝

Lean 4 theorem prover combining LLM policy with Energy-Based Model value function, trained via expert iteration.

A single DeepSeek-Prover-V2-7B backbone serves both autoregressive tactic generation (policy) and mean-pooled state embeddings (value), AlphaZero-style. A GoalConditionedEnergyHead (~25M params, PyTorch) scores proof states. LLM fine-tuning runs in Python with LoRA/PEFT.

Results

Tree Search vs pass@N Token Efficiency (miniF2F v2s_test, DeepSeek-Prover-V2-7B base)

Method Proofs Rate Token Budget / Theorem Tokens / Proof
pass@32 132/243 54.3% 262K 483K
pass@128 151/244 61.9% 1,049K 1,695K
Hybrid BFS 161/242 66.5% ~524K ~788K

Token budget = samples x max_tokens (pass@N) or search_budget x tokens_per_round (BFS). Tokens/proof = budget / proof_rate.

Key findings:

  • Hybrid BFS beats pass@128 by +4.6pp while using 2x fewer tokens
  • pass@128 uses 4x more tokens than pass@32 for only +7.6pp gain
  • Tree search is 2.1x more token-efficient than pass@128 (788K vs 1,695K tokens per proof)
  • Median nodes to first proof: 8 — most proofs found early via directed search

Why tree search wins — branching factor collapse:

Depth Branching Factor
0 → 1 5.0
1 → 2 1.7
2 → 3 1.5
3 → 4 1.3
4 → 5 1.0
5+ < 1.0

Best-first search rapidly prunes the space — by depth 5, most branches are dead. Blind pass@N sampling cannot exploit this structure and wastes tokens exploring equally in all directions.

Benchmarks

Benchmark Source Theorems Description
miniF2F v1 yangky11/miniF2F 244 test + 244 valid Standard math competition (AMC, AIME, etc.)
miniF2F v2s roozbeh-yz/miniF2F_v2 244 test + 244 valid Harder variant, different problems
IMO-Steps lemmas roozbeh-yz/IMO-Steps 1,328 steps Incremental proof steps from 13 IMO problems
IMO-Steps theorems roozbeh-yz/IMO-Steps 21 theorems Full IMO problem proofs

Architecture

DeepSeek-Prover-V2-7B (custom inference server, sgl.Engine)
├── Policy head: autoregressive tactic generation (LM head)
└── Mean-pool hidden states → Vec<f32> via /encode endpoint
                                    │
                                    ▼
                    Energy Head (burn-rs, trainable)
                    SpectralNorm MLP: 4096 → 2048 → 1024 → 512 → 1
                    Output: scalar energy (lower = more provable)

Lean 4 REPL Pool (tokio, Pantograph JSON protocol)
└── Verifies tactics against proof states, returns new goals

Best-first search expands nodes by combined LLM log-probability and EBM energy score, verified against Lean 4 via Pantograph.

Project Structure

BurnQED/
├── crates/
│   ├── lean-repl/       # Async Pantograph client, worker pool, ProofHandle pattern
│   ├── policy/          # SGLang HTTP client for tactic generation + hidden-state extraction
│   ├── search/          # Best-first proof search engine, trait-based
│   ├── trajectory/      # Parquet I/O for search trajectories
│   ├── ebm/             # Energy-Based Model: SpectralNorm MLP, training, inference
│   ├── prover-core/     # CLI binary tying everything together
│   └── burn-contrib/    # Upstream burn-rs PR modules (stub)
├── python/
│   ├── encode_server.py     # Embedding extraction server (nf4)
│   ├── encode_embeddings.py # Direct PyTorch encoding (bypasses SGLang batch bug)
│   ├── training/            # LoRA fine-tuning (train_llm.py, export_llm.py)
│   ├── joint/               # v2 joint training (EBM head, dataset, losses, model, train)
│   └── data/                # Dataset downloads, benchmark conversion
├── scripts/             # Pipeline orchestration scripts (paths from _lib.sh)
├── configs/             # search.toml, models.toml
├── vendor/Pantograph/   # Git submodule (Lean 4 REPL, Mathlib v4.27.0)
├── docs/                # Architecture plan, experiment guide, known issues
│
└── data/                # ALL generated/downloaded artifacts (gitignored except benchmarks/)
    ├── lean/            # Raw HF dataset downloads
    ├── benchmarks/      # Evaluation benchmark JSONs (tracked in git)
    ├── sft/             # Formatted SFT training data
    ├── models/          # Model weights (base/ + merged/)
    ├── checkpoints/     # Training checkpoints (lora/ + ebm/)
    ├── trajectories/    # Search result parquets
    ├── embeddings/      # Extracted embeddings per iteration
    ├── evals/           # Evaluation results & reports
    └── logs/            # Training and search logs

Prerequisites

  • Rust (stable, edition 2021)
  • Lean 4 via elan
  • Python 3.10+ with sglang, torch, transformers, peft, accelerate
  • GPU required for LLM inference; EBM training supports both CUDA and CPU
  • DeepSeek-Prover-V2-7B weights (HuggingFace)

Quick Start

# Clone with submodules
git clone --recurse-submodules https://github.com/<you>/BurnQED.git
cd BurnQED

# Build Pantograph (Lean 4 REPL)
./scripts/setup_pantograph.sh

# Build the prover
cargo build --release -p prover-core

# Smoke test (~2-3 min on A100: LLM search → train EBM → search with EBM → compare)
./scripts/smoke_test.sh

For cloud GPU setup (installs Rust, elan, Python deps, builds everything):

bash scripts/setup_runpod.sh   # RunPod (RTX 4090) — recommended
bash scripts/setup_lambda.sh   # Lambda Labs (A100)

Expert Iteration Pipeline

Each iteration has two phases: train then search.

For iteration i = 0, 1, 2, ...:

  run_iteration_train.sh i
  ┌───────────────────────────────────────────┐
  │ Step 0: Pre-training eval (train subset)  │
  │ Step 1: LLM LoRA fine-tuning (Python)     │
  │ Step 1b: Export merged safetensors        │
  │ Step 2: Restart inference server          │
  │ Step 3: Post-training eval (train subset) │
  │ Step 4: EBM training (encode + train)     │  ← skip for iter 0
  │ Step 5: miniF2F evaluation + ablation     │
  └───────────────────────────────────────────┘

  run_iteration_search.sh i
  ┌───────────────────────────────────────────┐
  │ Step 2: Proof search (trajectory collect) │
  │ Step 3: Trajectory summary                │
  └───────────────────────────────────────────┘

Both scripts support START_STEP=N to skip earlier steps.

# Full pipeline
./scripts/run_iteration_train.sh 1
./scripts/run_iteration_search.sh 1

# Skip to EBM training + eval
START_STEP=4 ./scripts/run_iteration_train.sh 1

# Skip to miniF2F eval only
START_STEP=5 ./scripts/run_iteration_train.sh 1

EBM Training

EBM training uses pre-computed embeddings (direct PyTorch encoding, bypassing SGLang's broken batch hidden states). The run_ebm_train.sh script handles the full cycle: stop server, encode, train, restart server.

# Standalone EBM training
./scripts/run_ebm_train.sh 2

# Environment variables
EBM_STEPS=50000 EBM_LR=3e-5 LOSS_TYPE=info_nce ./scripts/run_ebm_train.sh 2

Inference Server

Custom server wrapping sgl.Engine with in-process mean-pooling for encoding:

./scripts/start_sglang.sh data/models/merged/iter_1
PORT=30000 TP=1 ./scripts/start_sglang.sh

CLI Reference

All commands are subcommands of cargo run --release -p prover-core --:

search — Run proof search

cargo run --release -p prover-core -- search \
    --config configs/search.toml \
    --server-url http://localhost:30000 \
    --theorems data/benchmarks/iter1_search_theorems.json \
    --output data/trajectories/iter_1.parquet \
    --ebm-path data/checkpoints/ebm/iter_1 \  # optional: EBM value guidance
    --concurrency 8 \
    --num-workers 8 \
    --imports Mathlib

eval — Multi-budget evaluation

cargo run --release -p prover-core -- eval \
    --config configs/search.toml \
    --server-url http://localhost:30000 \
    --theorems data/benchmarks/minif2f_test.json \
    --budgets 600 \
    --output data/evals/iter_1.json \
    --ebm-path data/checkpoints/ebm/iter_1 \
    --num-candidates 16 \
    --imports Mathlib

train-ebm — Train EBM from trajectories

cargo run --release -p prover-core -- train-ebm \
    --trajectories data/trajectories/iter_0.parquet data/trajectories/iter_1.parquet \
    --server-url http://localhost:30000 \
    --output-dir data/checkpoints/ebm/iter_2 \
    --steps 50000 \
    --embeddings-cache data/checkpoints/ebm/iter_2/embeddings.parquet \
    --loss-type info_nce

summary / compare

cargo run --release -p prover-core -- summary --input data/trajectories/iter_1.parquet
cargo run --release -p prover-core -- compare \
    --results data/evals/iter_0.json data/evals/iter_1.json

Configuration

configs/search.toml

[search]
max_nodes = 600               # Node budget per theorem
max_depth = 25                # Max proof depth
num_candidates = 16           # Tactics generated per node (T≥1.0 yields 6-8 unique)
alpha = 0.5                   # LLM log-prob weight
beta = 0.5                    # EBM score weight
timeout_per_theorem = 600     # seconds
batch_expansion_size = 1      # Nodes popped per batch
batch_encode_size = 4         # Max states per EBM encode batch
# Hybrid whole-proof search
hybrid_num_proofs = 16
hybrid_min_proofs = 4
hybrid_max_rounds = 60
hybrid_max_tokens = 1024
hybrid_budget = 512

[lean_pool]
num_workers = 8
max_requests_per_worker = 1000
max_lifetime_secs = 1800      # Worker recycling
tactic_timeout_secs = 60

Testing

# Unit tests (no external dependencies)
cargo test --workspace

# Integration tests (require Pantograph, run single-threaded)
cargo test -p lean-repl -- --ignored --test-threads=1    # ~60-90s
cargo test -p search -- --ignored --test-threads=1       # ~60s
cargo test -p prover-core -- --ignored --test-threads=1  # ~15s

# Policy integration tests (require running SGLang server)
cargo test -p policy -- --ignored --test-threads=1

Key Dependencies

Crate Version Purpose
burn 0.16 EBM training and inference
reqwest 0.12 HTTP client for inference server
tokio 1 Async runtime for Lean worker pool
arrow / parquet 53 Trajectory data I/O
clap 4 CLI argument parsing

Known Issues

  • SGLang batch hidden states: return_hidden_states=True is broken in batch mode (SGLang #8066). Workaround: python/encode_embeddings.py uses direct PyTorch encoding for training data. Search-time encoding uses sequential calls (one at a time), which is correct. See docs/encoding_bug.md.
  • Dataset version gaps: Training datasets (Lean Workbook v4.8, Goedel v4.9, NuminaMath v4.15) lag behind our Pantograph v4.27. Mathlib lemma renames may break some tactics. Pantograph validation (tasks 0.3d-f) measures actual compatibility. See docs/datasets.md.

License

Apache 2.0

For the full architecture plan, see docs/burn-qed_plan.md.

About

Lean 4 theorem prover combining LLM policy with Energy-Based Model value function, trained via expert iteration.

Resources

Readme

License

Apache-2.0 license
Activity

Stars

1 star

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages