Lean Refactor: Multi-Objective Controllable Proof Optimization via Agentic Strategy Search

Jialin Lu; Soonho Kong; Rodrigo Stehling; Kaiyu Yang; Zhangyang Wang; Weiran Sun; Wuyang Chen

Source-linked AI summary

Lean Refactor: Multi-Objective Controllable Proof Optimization via Agentic Strategy Search

Jialin Lu, Soonho Kong, Rodrigo Stehling, Kaiyu Yang, Zhangyang Wang, Weiran Sun, Wuyang Chen

arXiv:2605.20244v2May 18, 2026cs.LOcs.AIcs.CLcs.LGcs.SE

TL;DR

Lean proof refactoring must address verbosity, compilation cost, competing objectives, scarce paired data, and incompatibility with rapidly changing Lean/Mathlib versions. Lean Refactor uses a retrieval-augmented agentic framework with a densely annotated strategy bank to steer frozen LLMs at inference time. It achieves strong compression across competition and research proofs, reduces compilation cost, and improves version robustness, while cross-version research evaluation remains limited.

Problem
Lean proofs generated by LLMs are often correct but verbose and costly, while refactoring must jointly handle length, compilation cost, version compatibility, and scarce paired data.
Method
Lean Refactor steers frozen agentic LLMs with dynamically filtered and reranked retrievals from a strategy bank annotated with application guidance, expected cost reduction, and Lean/Mathlib compatibility.
Results
Lean Refactor achieves over 70% compression on competition proofs, over 20% on research repositories, and over 30% compilation-cost reduction while outperforming prior work and Claude Code.
Takeaways & Limitations
Metadata-aware retrieval supports controllable, multi-objective refactoring across LLM backbones and improves proof portability across future Lean/Mathlib releases.
Takeaways & Limitations
Cross-version evaluation is limited to PutnamBench because research projects are tied to specific Lean/Mathlib toolchains.

Abstract

from arXiv · show

We present Lean Refactor, a plug-and-play retrieval-augmented agentic framework for multi-objective, controllable, and version-robust refactoring of Lean proofs. LLM-generated proofs are notoriously correct-but-verbose and brittle across library versions, yet existing refactoring works overlook three practical challenges: 1) Lean refactoring is natively multi-objective (proof length, compilation cost, and version compatibility are often in tension); 2) Lean repositories have fragile compatibility, whereas LLM releases are unaware of Lean/Mathlib versions; 3) Training-based pipelines require repeated fine-tuning with each new LLM release, scaling neither with model churn nor with Lean's release cycle. Lean Refactor steers a frozen agentic LLM with retrievals from a curated database of multi-objective refactoring strategies, each densely annotated with metadata such as supported Lean/Mathlib versions and expected compilation-cost reduction. Experiments show over $70\%$ token-level compression on competition benchmarks, over $20\%$ on research repositories, and up to $60\%$ compilation-time reduction, outperforming prior work and Claude Code. Version-filtered retrieval further improves compression on the target Lean version, and refactored miniF2F proofs exhibit stronger zero-shot version transfer to future Lean releases than their unrefactored counterparts.

1 Introduction

Lean Refactor addresses verbose, costly, and version-fragile Lean proofs by steering a frozen LLM with version-aware retrieval. It targets competing objectives without fine-tuning and reports strong compression across competition and research proofs.

Motivation: LLM-generated Lean proofs are often correct but unnecessarily long, difficult to read, and expensive to compile.Reinforcement-learning rewards emphasize correctness, overlooking redundant steps and heavy automation.
Challenges: Lean refactoring must balance proof length, compilation cost, and compatibility across Lean versions.These objectives can conflict, while Lean and Mathlib evolve rapidly beyond static LLM knowledge cutoffs.
Challenges: Existing Lean corpora lack refactoring trajectories and noisy-clean proof pairs, especially compared with mainstream programming-language datasets.This scarcity limits training-based approaches for Lean proof refactoring.
Approach: Lean Refactor retrieves densely annotated strategies containing application instructions, examples, expected compilation-cost reduction, and validated Lean/Mathlib compatibility.Dynamic filtering and reranking steer a frozen agentic LLM toward user-specified objectives without model fine-tuning.
Results: Over 70% compression is achieved on competition proofs and over 20% on research repositories, while compilation cost falls by over 30%.The framework remains compatible with multiple LLM backbones and improves cross-version type-checking for miniF2F proofs.

2 Lean Refactoring is Challenging for LLMs

Lean proof optimization is difficult because shorter proofs do not reliably compile faster, and competition benchmarks underrepresent research-level structural complexity. LLMs must therefore reason about multiple objectives and substantially richer dependency contexts.

Proof Optimization is Multi-Objective: Proof length and compilation cost are important optimization objectives that are not always aligned.Short proofs can invoke heavy tactics, so prior work often misses joint optimization of length and runtime.
Proof Optimization is Multi-Objective: Prior Lean refactoring work generally ignores compilation cost or substitutes it for length rather than optimizing both jointly.Continuous runtime signals remain underused in formal-mathematics post-training.
Proof Optimization is Multi-Objective: Token length fails to predict compilation cost because short, heavy tactics can dominate runtime.This weak correlation makes proof length an insufficient proxy for compilation efficiency.
Competitions Fail to Reflect Research-Level Proofs: Competition proofs do not reflect the structural complexity of research-level Lean code.Research proofs rely more heavily on Mathlib and intra-project dependencies, requiring reasoning over larger and less familiar contexts.

3 Methods

Lean Refactor externalizes refactoring knowledge into a version- and cost-aware strategy bank, then uses objective-conditioned retrieval and a frozen-LLM agent loop to plan, execute, and debug proof transformations.

3.1 A Densely Annotated, Version-Aware Strategy Bank: Lean Refactor stores reusable refactoring strategies with structural fields, version compatibility, and expected compilation-cost metadata.Strategies include application conditions, step-by-step guidance, before/after examples, and potential length-reduction tags.
3.1 A Densely Annotated, Version-Aware Strategy Bank: Compilation-cost metadata is estimated by profiling every long–short pair and retaining each strategy’s median relative speedup.Profiling isolates proof execution from import overhead.
3.1 A Densely Annotated, Version-Aware Strategy Bank: The bank is built from diverse Lean sources, with theorem decontamination and synthesized long–short proof pairs used to distill strategies.Sources include NuminaMath-1.5, FineLeanCorpus, Mathlib, and ATLAS.
3.2 Multi-Objective Retrieval for Controllable Refactoring: A retrieval model maps proof segments to strategies using 639090 query–strategy pairs, in-batch negatives, and boundary augmentation.The query is the code segment operated on by a strategy, while its when-to-apply clause is the positive target.
3.2 Multi-Objective Retrieval for Controllable Refactoring: At inference, cosine retrieval, compilation-cost reranking, and version filtering select strategies according to user objectives over the shared bank.These rules can be composed, such as stacking version filtering on compilation-cost reranking, and extended through new metadata fields.
3.3 Iterative Agent Workflow: A frozen LLM agent annotates dependencies, retrieves strategies at multiple granularities, plans mapped refactoring steps, and replans after successful compilation and shortening.Persistent failures are logged and skipped; the loop stops after B LLM calls or when proof length falls below T.

4 Experiments

Lean Refactor is evaluated across competition and research benchmarks, backbones, objectives, and Lean versions. It improves proof compression and compilation cost, with targeted retrieval supporting cross-version robustness.

Experimental Setup: The evaluation covers seven benchmarks across competition and research regimes, using shared inputs where available and retaining research proofs’ internal and external dependencies.Experiments use Gemini 3 Flash by default, with additional Claude Haiku 4.5 and Claude Code comparisons.
Proof Length Reduction: Lean Refactor surpasses ProofOptimizer on competition benchmarks and achieves substantial additional compression on human-written research proofs.The framework beats every ablation on research-level proofs, with the competition gap widening on PutnamBench.
Ablations: Targeted strategy retrieval and planning outperform random or absent retrieval, showing that accurate strategy selection matters more than adding context alone.Without planning, the base agent stalls in persistent debugging loops; retrieval is effective when the planner sequences strategies.
LLM-Agnosticism: The framework remains effective across LLM backbones, and Claude Code improves when given strategy retrieval but remains below the full framework under matched budget and backbone.The comparison uses Claude Haiku 4.5 for both systems and matches Claude Code’s dollar budget to the framework’s runs.
Multi-Objective Control: Objective-aligned reranking achieves the largest compilation-cost reduction on every competition and research-level benchmark.Cosine-only retrieval can compress tokens while increasing compilation cost; on miniF2F its relative reduction is −103.98%.
Cross-Version Compatibility: Filtered Retrieval attains the best length reduction on four of six version–model cells and beats both baselines for GPT-OSS-20B at v4.22.0 and v4.16.0.Full Retrieval shows no consistent gain over No Retrieval, indicating that incompatible retrieved artifacts can disrupt refactoring.
Cross-Version Compatibility: Lean Refactor’s refactored miniF2F proofs continue to type-check across a longer span of future Mathlib releases than ProofOptimizer and original proofs.The advantage begins at v4.22.0 and widens at later releases without further refactoring.

5 Related Work

Prior Lean refactoring methods optimize limited objectives and do not address version drift, while existing corpora lack refactoring trajectories and noisy-clean proof pairs.

Prior Lean Refactoring: Prior Lean refactoring targets a single axis: ImProver requires per-objective example corpora, while ProofOptimizer fine-tunes for length and treats compilation cost only as an alternative inference-time objective.Neither approach handles Lean/Mathlib version drift.
Data Resources: Existing Mathlib and community-aggregated competition corpora provide neither refactoring trajectories nor noisy-clean code pairs.Lean also has a much smaller codebase than mainstream programming languages such as C/C++, Java, and Python.

6 Limitations

The evaluation’s cross-version compatibility analysis is limited to PutnamBench, and strategy metadata uses conservative cluster-level aggregation rather than marginal attribution.

Scope and Attribution: Cross-version evaluation is limited to PutnamBench because research projects are pinned to toolchains whose dependencies make version changes break compilation independently of refactoring quality.The strategy metadata also uses conservative cluster-level aggregation rather than marginal attribution.

7 Conclusion

Lean Refactor separates Lean-proof refactoring logic from the underlying LLM through a densely annotated strategy bank, retrieval, and compiler-aware validation. The framework targets maintainable refactoring across objectives, model backbones, and evolving Lean/Mathlib toolchains.

Conclusion: Lean Refactor decouples proof-refactoring logic from the underlying model through a densely annotated strategy bank and compiler-in-the-loop architecture.The framework steers frozen agentic LLMs toward proof conciseness, compilation efficiency, and toolchain compatibility.
Conclusion: The framework addresses Lean’s gap with mainstream refactoring settings, where paired data is scarce, APIs change frequently, and validity is compiler-checked.Its hybrid sourcing strategy spans competition statements, Mathlib4, and synthetic Mathlib-derived statements to broaden coding styles and proof architectures.
Conclusion: The pipeline includes decontamination against held-out evaluation statements before synthesizing the long–short proof pairs.Exact-match and embedding-based near-duplicate filtering are applied to candidate theorem statements.
Conclusion: The strategy bank is built from synthesized, verified long–short proof pairs and distilled into reusable refactoring strategies.The pipeline uses multiple prover models, dual-direction pair construction, dependency context, deduplication, and quality-control stages.
Conclusion: 9,237 unique refactoring strategies remain after deduplicating 481,567 raw extractions.Embedding-based shortlisting, a high-similarity shortcut, and an LLM judge fallback resolve semantic duplicates while keeping curation tractable.
Conclusion: Each strategy’s compatibility set conservatively includes only toolchains where all constituent shortened proofs compile successfully.The resulting metadata supports filtering incompatible strategies and ranking surviving candidates by median compilation speedup.

C Limitations

The cross-version evaluation is limited to PutnamBench because research projects are tightly coupled to pinned Lean/Mathlib toolchains. The metadata’s joint effects are summarized conservatively rather than attributed to individual strategies.

Benchmark coverage: Cross-version compatibility experiments are confined to PutnamBench because research projects’ deep dependencies make toolchain changes break compilation independently of refactoring quality.Version-portable research-level benchmarks are identified as an important direction for the community.
Joint-effect attribution: Strategy-level compilation annotations use cluster medians of joint pair-level reductions rather than marginal per-strategy attribution.Shapley-style attribution or regression over co-occurrence indicators is left for future work.
Length–compilation tradeoff: The example demonstrates that proof length and compilation time can diverge sharply: a proof exceeding 130 lines compiles more than 20× faster than its shorter counterpart.The shorter proof triggers interval_cases, simp_all, and nlinarith across hundreds of branches.
Implementation: The retrieval-model training configuration uses one epoch on four NVIDIA 6000Ada GPUs with a contrastive objective.The supplied passage gives implementation settings rather than a stated limitation of the method.

F Proof Dependency Reduction

Lean Refactor reduces proof dependencies across competition and research-level benchmarks, supporting shorter, more self-contained proofs and potentially improving maintainability and version resilience.

Dependency results: 64.0% of PutnamBench and 51.0% of miniF2F proofs reduced their external or intra-project dependencies.Research-level reductions reached 66.7% on Analysis and 60.0% on FLT.
Interpretation: The dependency reductions indicate that brevity comes from direct logical simplification rather than merely invoking heavier specialized library calls.The resulting proofs are more self-contained and require fewer external lemma names.
Version transfer: Zero-shot version transfer was evaluated by compiling v4.19.0 proofs against subsequent Mathlib releases without further refactoring.The comparison included LeanRefactor, original human-written proofs, and ProofOptimizer; transfer outcomes depended on the benchmark.
Version transfer: On MiniF2F, LeanRefactor transferred furthest among baselines, while on PutnamBench it fell short of the original proofs and tracked ProofOptimizer closely.The paper concludes that transfer depends on corpus composition and the tactics introduced during shortening.

I.1 Benchmarks

The evaluation spans three competition benchmarks and four research-level repositories, using proof-length, compilation-cost, and cross-version metrics with multiple baselines and ablations.

Benchmarks: Three competition benchmarks are evaluated: miniF2F, PutnamBench, and Putnam2025.miniF2F contains 194 theorems, PutnamBench 75, and Putnam2025 66.
Benchmarks: Four research-level repositories are evaluated: Analysis, FLT, PFR, and PhysLean, with 45 sampled theorems from each.Samples are stratified by initial proof length and retain internal and external dependencies.
Metrics: Proof brevity is measured as the average relative percentage decrease in token count using ProofOptimizer’s syntax-aware tokenizer.The metric is reported across all seven benchmarks.
Metrics: Compilation cost uses isolated proof timing for standalone competition files and heartbeats for research repositories where per-proof wall-clock attribution is unstable.Measurements use consistent environments, identical hardware, and means over five runs where timing is available.
Metrics: Cross-version evaluation compares version-filtered and unfiltered retrieval across Lean/Mathlib v4.14.0, v4.16.0, and v4.22.0.Additional experiments test zero-shot transfer and compare against ProofOptimizer and original proofs.
Baselines and ablations: Internal ablations remove the Planner, retrieval, or metadata-aware retrieval to isolate their contributions.A random-retrieval control distinguishes the value of selecting appropriate strategies from merely providing strategies.
Baselines and ablations: The framework is also tested with Claude Haiku 4.5 and compared with Claude Code using Lean MCP, Lean skills, and strategy retrieval configurations.The Claude Code comparison uses PutnamBench and Analysis as representative competition and research-level benchmarks.

K Comparison with ProofOptimizer’s Heartbeat-Optimized Proofs

Lean Refactor’s objective-conditioned reranking achieves lower heartbeat counts than ProofOptimizer’s heartbeat-optimized proofs on both compared competition benchmarks.

Heartbeat comparison: Lean Refactor with reranking attains the lowest average heartbeat count among all evaluated methods on both benchmarks.The comparison uses Lean’s deterministic heartbeat proxy because ProofOptimizer’s heartbeat-optimized model weights are unavailable for direct re-execution.
Heartbeat comparison: 1.2K fewer heartbeats are achieved on miniF2F, improving from ProofOptimizer’s 10.4K to 9.2K.These values compare Lean Refactor’s reranked proofs with ProofOptimizer’s heartbeat-optimized proofs.
Heartbeat comparison: 29K fewer heartbeats are achieved on PutnamBench, improving from 111K to 82K.The objective is carried in retrieval reranking through per-strategy heartbeat metadata rather than model weights.

L Compilation Cost for Research-Level Proofs

For research-level proofs, heartbeat measurements accommodate the lack of stable per-proof wall-clock attribution, and reranking is strongest by both reported compilation-cost metrics on every dataset.

Measurement: Research repositories use heartbeats because proof-level wall-clock time is confounded by imports, neighboring declarations, and kernel cache state.Competition proofs instead permit isolated wall-clock measurement in self-contained files.
Metrics: Table 12 reports average per-proof relative and absolute heartbeat reductions across Analysis, FLT, PFR, and PhysLean, each with 45 proofs.Higher values are better for both metrics.
Results: Lean Refactor with reranking is strongest by both heartbeat metrics on every research-level dataset.Non-reranking variants can show negative relative reductions because outliers increase the average, even when absolute means and most proofs improve.

M Additional Proof Length Reduction Results

Lean Refactor reduces proof length across verifiable code-generation and PDE benchmarks while applying reusable strategies and verifier-guided iterative refactoring. The evaluated strategies replace numeric derivations or per-entry matrix lemmas with more compact structural proofs.

Benchmark setup: The Verina proof-generation setup closes 14 of 189 problems, or 7.4%, using up to 16 compiler-feedback attempts.Each candidate is checked by Lean, with compiler errors returned as additional context and integrity and axiom checks applied.
Results: 53.50% average relative proof-length reduction on Verina and 35.34% on PDE demonstrates transfer across verification and mathematical domains.Verina covers verifiable code generation, while PDE covers formalized partial differential equation statements.
Existential-witness refactoring: For shifted existential goals, replace numeric witness derivations with the translated witness a + k and specialize the hypothesis at y - k.This applies when the hypothesis has the form ∃x, ∀y, P x y and the goal uses an affine reparameterization such as P x (y - k).
Matrix-equality refactoring: For finite-indexed matrix equalities, ext followed by fin_cases and simp replaces separate per-entry have lemmas and a final Matrix.ext proof.The pattern depends on Matrix.of or !![...] construction and on per-entry goals being solvable by a unified finisher.
Agentic refactoring pipeline: The retrieval-augmented algorithm segments proofs, retrieves aligned strategies, plans edits, compiler-checks candidates, self-debugges failures, and keeps only shorter successful proofs.It iterates under a step budget and dynamically replans after each accepted refactoring.