💻 computer science

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

Lean Refactor is a retrieval-augmented agentic framework that optimizes Lean proofs for multiple objectives—including token compression, compilation speed, and version compatibility—by dynamically selecting curated refactoring strategies without requiring model retraining.

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

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Over-Engineered" Proof

Imagine a brilliant but overly enthusiastic architect (the AI) who has just built a house (a mathematical proof). The house is structurally sound—it won't fall down, and it passes all safety inspections. However, the architect used 500 bricks to build a wall that only needed 50. They used a sledgehammer to crack a nut.

In the world of Lean (a language used to write mathematical proofs that computers can verify), AI models often produce proofs that are correct but incredibly long, messy, and slow to compile.

Too Long: The proof is hard for humans to read.
Too Slow: Because the proof is bloated with unnecessary steps, it takes a long time for the computer to check if it's valid.
Fragile: The Lean language changes frequently (like software updates). A proof written for "Version 1.0" might break completely when the computer updates to "Version 2.0," even if the math is still right.

Existing tools tried to fix this by either retraining the AI (which is expensive and slow) or just asking it to "be shorter" (which often makes the proof slower to run).

The Solution: The "Smart Librarian" System

The authors created Lean Refactor. Instead of trying to retrain the AI, they built a "plug-and-play" system that acts like a super-smart librarian for the AI.

Here is how it works, using a metaphor:

1. The Strategy Bank (The Library)

Imagine a massive library filled with "Refactoring Cards." Each card describes a specific trick to shorten a proof.

The Trick: "Instead of listing every number one by one, use this one magic formula."
The Metadata: Crucially, every card has a label. It says:
- How much time does this save? (e.g., "Saves 30% compilation time").
- Which version of the language does this work on? (e.g., "Works on Lean v4.16 and v4.22").
- How much shorter does it make the proof?

The paper claims they built the largest library of these cards ever, containing over 9,000 unique strategies distilled from hundreds of thousands of proof examples.

2. The Agent (The Architect + The Librarian)

When the AI needs to fix a proof, it doesn't guess. It follows a loop:

The Planner: The AI looks at the messy proof and says, "This section looks like it needs a 'Magic Formula' trick."
The Librarian (Retrieval): The system goes to the library and pulls out the specific cards that match that section.
- If you want the shortest proof: It pulls cards that promise the biggest size reduction.
- If you want the fastest compilation: It pulls cards that promise the biggest speed boost.
- If you are using an old version of Lean: It filters out any cards that won't work on your version.
The Refactorer: The AI applies the trick from the card to rewrite the proof.
The Debugger: The computer checks the new proof. If it breaks, the AI tries to fix the error locally without undoing the whole improvement.

Why This is Special (The Paper's Claims)

1. It balances competing goals (The "Multi-Objective" Magic)
Usually, you have to choose: "Do I want the proof to be short, or do I want it to compile fast?"

Old Way: You pick one, and the other suffers.
Lean Refactor: You can tell the system, "I want it short, but also fast." The system looks at the library cards, sees which ones offer both, and picks the best balance.
Result: On competition math problems, they shrank proofs by over 70% and cut compilation time by over 60%.

2. It survives software updates (The "Version Robustness")
Because the library cards are labeled with the specific software versions they work on, the system can instantly filter out "outdated" tricks.

Result: If you ask for a proof for "Lean v4.16," the system only uses cards that work for v4.16. This prevents the AI from hallucinating (making up) tools that don't exist in that version.

3. It works with any AI (The "Model Agnostic" Feature)
You don't need to train a new AI for this. The system works with any "frozen" (pre-trained) AI model, whether it's from Google (Gemini), Anthropic (Claude), or OpenAI (GPT). The "intelligence" comes from the library of strategies, not from the AI's brain.

Result: They tested it on different AI models, and it improved the performance of all of them, even beating a specialized coding agent called "Claude Code."

The Bottom Line

Lean Refactor is like giving a construction crew a set of blueprints and a toolbox of pre-tested, labeled shortcuts. Instead of guessing how to build a house, they look up the best way to build a specific wall based on the tools they have and the version of the building code they are using.

The paper claims this approach:

Shrinks proofs by 70%+ on math competitions.
Makes them compile 60% faster.
Keeps them working even when the software language updates.
Does all this without needing to retrain the AI models.

Technical Summary: Lean Refactor

Problem Statement

Large Language Models (LLMs) trained with reinforcement learning (RL) for formal theorem proving in Lean often generate proofs that are correct but notoriously verbose, brittle, and computationally expensive. While existing refactoring efforts exist, they face three fundamental limitations when applied to the Lean ecosystem:

Multi-Objective Tension: Lean refactoring is inherently multi-objective. Optimizing for proof length often conflicts with minimizing compilation cost (e.g., shorter proofs may invoke heavy tactics that inflate compile time), and existing methods lack principled mechanisms to balance these trade-offs or adapt to unseen user objectives at inference time.
Version Fragility: Lean and its standard library, Mathlib, evolve on a near-weekly basis with frequent API changes and lemma renames. Frontier LLMs have static knowledge cutoffs, leading to hallucinations of stale tactics and incompatibility with current toolchains.
Data Scarcity: Unlike mainstream programming languages, Lean lacks large-scale paired datasets (noisy-clean code pairs) required for training-based refactoring pipelines. Furthermore, training-based approaches require repeated fine-tuning for every new LLM release, failing to scale with model churn or Lean's release cycle.

Methodology: Lean Refactor

The authors propose Lean Refactor, a plug-and-play, retrieval-augmented agentic framework designed to refactor Lean proofs without fine-tuning the underlying LLM. The system decouples refactoring logic from the model by externalizing knowledge into a structured, densely annotated database.

1. Densely Annotated Strategy Bank

The core of the framework is a curated database of 9,237 unique refactoring strategies distilled from over 200,000 long-short proof pairs.

Structure: Each strategy includes a title, description, precondition ("when to apply"), step-by-step application guide, and abstract before/after examples.
Metadata: Crucially, each strategy is annotated with metadata the LLM cannot infer from training data:
- Compilation-Time Reduction: The median relative speedup observed when applying the strategy.
- Version Compatibility: A set of validated Lean/Mathlib toolchains (e.g., v4.14.0, v4.16.0, v4.22.0) under which the refactored proof successfully compiles. This is determined by recompiling shortened proofs across multiple versions.

2. Multi-Objective Retrieval and Reranking

At inference, the framework employs a retrieval pipeline that adapts to user-specified objectives without model retraining:

Shortest Proof: Retrieves strategies based purely on semantic similarity (cosine similarity) to the proof segment.
Compilation-Time Aware: Retrieves a broad pool of semantically similar candidates and then reranks them based on the annotated compilation-cost reduction metadata.
Version-Specific: Filters the candidate pool to include only strategies whose compatibility metadata includes the user's target Lean version.

3. Agentic Loop

The framework wraps these components in a model-agnostic agent loop consisting of:

Planner: Decomposes the refactoring task into specific steps based on retrieved strategies.
Refactorer: Executes the plan to generate a shorter proof.
Debugger: Uses compiler feedback to repair errors in a closed loop (up to $D$ rounds).
Backbone Agnosticism: The planner, refactorer, and debugger can be instantiated by any frozen frontier LLM (e.g., Gemini, Claude, GPT series), allowing the system to transfer across models without retraining.

Key Contributions

Dataset Release: The authors release the largest Lean refactoring dataset to date, comprising over 200K long-short proof pairs and 9K refactoring strategies, each annotated with compilation cost and version compatibility.
Performance: Lean Refactor achieves >70% token-level compression on competition benchmarks (miniF2F, PutnamBench) and >20% on research repositories (Analysis, FLT, PhysLean). It also reduces compilation time by up to 60% on competition benchmarks.
Controllability and Robustness: The system supports multi-objective control (length vs. speed vs. compatibility) via retrieval rules rather than model weights. Refactored proofs exhibit stronger zero-shot version transfer to future Lean releases compared to unrefactored counterparts and prior optimization methods.
Model Agnosticism: The framework improves performance across different LLM backbones (Gemini 3 Flash, Claude Haiku 4.5) and outperforms general-purpose coding agents like Claude Code, even when the latter is given access to the same strategy retrieval tool.

Experimental Results

Length Reduction: On competition benchmarks, Lean Refactor outperforms the prior state-of-the-art, ProofOptimizer, achieving 88.40% reduction on miniF2F and 70.38% on PutnamBench (vs. 57.20% for ProofOptimizer). On research-level proofs, it achieves an average of 20–33% reduction, significantly outperforming ablation baselines.
Compilation Time: By reranking strategies based on compile-cost metadata, the framework achieves a 42.62% relative reduction in compilation time on PutnamBench, compared to 37.11% for the standard retrieval approach.
Version Compatibility: When targeting specific Lean versions (e.g., v4.16.0), version-filtered retrieval consistently yields better length reduction than unfiltered retrieval. Refactored proofs maintain type-checking capability across more future Lean releases (up to v4.29.0) than those produced by ProofOptimizer or original verbose proofs.
Dependency Reduction: The refactored proofs reduce external and intra-project dependencies in 42–67% of cases, indicating genuine logical simplification rather than mere offloading to heavy library lemmas.

Significance

The paper argues that Lean Refactor addresses the critical bottlenecks of maintaining formal mathematical libraries: the tension between proof conciseness and compilation efficiency, and the fragility of code across rapidly evolving toolchains. By externalizing refactoring knowledge into a metadata-rich strategy bank, the framework enables controllable, multi-objective optimization without the need for expensive model fine-tuning. This approach allows the system to remain robust against Lean/Mathlib version drift and adaptable to new LLM backbones immediately, laying the groundwork for more maintainable and version-resilient formal proofs.