🤖 AI

ImProver 2: Iteratively Self-Improving LMs for Neurosymbolic Proof Optimization

The paper introduces ImProver 2, a neurosymbolic framework that combines data-efficient expert iteration with a structural scaffolding system to enable small language models to effectively optimize complex formal proofs in Lean 4, outperforming significantly larger models while establishing proof optimization as a scalable, learnable task.

Original authors: Riyaz Ahuja, Tate Rowney, Jeremy Avigad, Sean Welleck

Published 2026-05-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Riyaz Ahuja, Tate Rowney, Jeremy Avigad, Sean Welleck

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

Imagine you have a massive, growing library of mathematical proofs. These aren't just stories; they are rigid, computer-checked blueprints that prove things are true. Recently, computers (specifically AI) have started helping humans write these proofs, causing the library to explode in size.

But there's a problem: The library is getting messy. Some proofs are correct but written in a confusing style, others are too long, and some rely on a huge list of other rules that make them hard to maintain. It's like having a house where the plumbing works, but the pipes are tangled, the walls are painted in clashing colors, and you need to carry a backpack full of extra tools just to turn on a light.

Enter ImProver 2.

Think of ImProver 2 as a super-organized, self-improving librarian who doesn't just shelve books but rewrites them to be better. Its job is to take a correct proof and rewrite it to be shorter, cleaner, or more modular, without breaking the math.

Here is how it works, using some simple analogies:

1. The "Practice Makes Perfect" Loop (Iterative Self-Improvement)

Usually, to teach a computer to do a task, you show it thousands of examples written by humans. But for proof optimization, good examples are rare.

ImProver 2 solves this by teaching itself.

The Draft: It takes a proof and tries to rewrite it in many different ways (like a writer brainstorming 10 different endings for a story).
The Grading: It checks which rewrites are still mathematically correct. Then, it scores them based on what we want: Is it shorter? Does it use fewer external references? Is it broken into clearer steps?
The Lesson: It compares its "winning" rewrites against its "losing" ones. It learns, "Oh, when I did this, the proof got shorter and stayed correct. When I did that, it broke."
The Replay Buffer: To prevent the AI from forgetting what it learned or getting stuck in a loop of bad ideas, it keeps a "memory bank" of old good examples mixed with new ones. This ensures it keeps getting better over time, rather than just repeating the same mistakes.

2. The "Neurosymbolic Scaffold" (The Training Wheels)

Imagine trying to fix a complex engine without a manual or a diagram. It's hard. Now imagine having a diagram that highlights exactly which part you are working on, explains what that part does in plain English, and lists the tools you need nearby.

ImProver 2 gives the AI this "diagram" for every proof. This is called Neurosymbolic Augmentation. It provides:

The Map: It shows the current state of the proof (what is proven so far, what is left to do).
The Context: It pulls up the specific definitions and rules relevant to that specific proof, so the AI doesn't have to guess.
The Translation: It gives a "plain English" summary of what the proof is trying to do, helping the AI understand the goal before it starts writing the code.

This "scaffold" helps even small, less powerful computers perform like much larger, smarter ones.

3. The Three Goals (Metrics)

The librarian doesn't just want the proof to be "correct." It wants to optimize for specific qualities, like a chef adjusting a recipe:

Length (The "Golf" Metric): Just like in golf, the goal is to get the ball in the hole with the fewest swings. ImProver 2 tries to shorten proofs by removing unnecessary steps.
Dependencies (The "Self-Contained" Metric): Imagine a recipe that says "add the secret sauce from the neighbor's house." That's a dependency. ImProver 2 tries to rewrite proofs so they don't rely on so many external "neighbor's sauces," making them easier to understand and use on their own.
Modularity (The "Lego" Metric): A messy proof is like a giant block of clay. A modular proof is like a set of Legos—small, distinct, reusable pieces. ImProver 2 tries to break big proofs into smaller, independent sub-proofs (like "have h1," "have h2") so the logic is clearer.

The Results: Small is Mighty

The most surprising finding is that ImProver 2 took a small AI model (7 billion parameters) and trained it to be better at these tasks than massive AI models (some with hundreds of billions of parameters) that were not given this special training.

It's like taking a compact, efficient sports car and tuning its engine so perfectly that it beats a massive, heavy truck in a specific race, even though the truck has a bigger engine. The small model, when given the right "scaffold" and the right "practice loop," learned to restructure complex mathematical arguments better than the giants.

In summary: ImProver 2 is a system that teaches small AI models to be expert proof editors. By giving them a clear map of the problem and letting them learn from their own best attempts, it can clean up messy mathematical libraries, making them shorter, cleaner, and easier to maintain, all without needing the most expensive, massive supercomputers.

Technical Summary: ImProver 2

Problem Statement

The rapid expansion of formal mathematics libraries (e.g., Lean's Mathlib) has created a dual challenge: the maintainability of these libraries is threatened by heterogeneous proof styles and low modularity, while the quality of training data for neural theorem provers is degraded by proofs that are correct but structurally suboptimal. Existing automated proof optimization methods face significant hurdles, including heterogeneous and heuristically defined objectives, scarce high-quality optimization data, and the prohibitive costs of training and inference for large-scale models. Furthermore, prior work often relied on general-purpose closed-source models or focused narrowly on token counts, failing to address broader structural metrics like modularity and dependency management.

Methodology

The authors introduce ImProver 2, a neurosymbolic framework designed to iteratively self-improve small language models (SLMs) for proof optimization in Lean 4. The system operates through a multi-component pipeline:

1. Neurosymbolic Augmentation (Scaffolding)

To bridge the gap between informal reasoning and formal verification, the framework exposes the model to rich context from the Lean environment. This "scaffold" ( $\Psi$ ) includes:

Context Slices: Relevant definitions, theorems, and lemmas directly referenced by the target proof, serialized with metadata.
Chain-of-States (CoS): Explicit traces of goal states before and after each tactic step, interleaved as comments to reveal the proof's logical flow.
Auto-informalization: Natural language summaries of the proof steps, generated by translating the formal proof into a sequence of high-level descriptions.

2. Optimization Metrics

The system targets three distinct structural objectives, defined as computable functions $\mu$ :

Length: Minimizing the number of tactics used (proof "golfing").
Dependencies: Minimizing the count of unique external lemmas explicitly named in the proof, encouraging self-contained proofs.
Modularity: Maximizing the number of "effective spawned goals" (independent subproofs introduced via have, calc, etc.), which serves as a proxy for proof decomposition and maintainability.

3. Iterative Preference Optimization (IRPO) Pipeline

ImProver 2 employs an iterative training loop to bootstrap SLMs:

Generation: The current model $G_t$ generates $n$ candidate rewrites for each theorem in the training set, utilizing the neurosymbolic scaffold.
Filtering & Replay: Candidates are scored based on correctness (Lean kernel verification) and metric improvement. A replay buffer is used to balance newly generated data with historical data, preventing model collapse and ensuring monotonic improvement. Trivial or low-quality problems are filtered out.
Preference Pair Construction: The dataset is partitioned into "winners" (compiling, high-improvement proofs) and "losers" (failures or low-improvement proofs). Two types of preference pairs are formed: winner–winner (ordered by improvement magnitude) and winner–loser.
Training: The model is updated using Iterative Reasoning Preference Optimization (IRPO), which combines a Direct Preference Optimization (DPO) loss with a negative log-likelihood (NLL) loss on the winner proofs. This process repeats until convergence or budget exhaustion.

Key Contributions

Self-Improving SLMs for Refactoring: The paper demonstrates that iterative preference optimization can effectively bootstrap small language models (specifically a 7B parameter model) to optimize proofs in specialized research libraries, making them competitive with significantly larger models.
Structural Optimization Metrics: Beyond simple length reduction, the authors introduce and evaluate modularity and dependency metrics. These metrics leverage proof structure and library context to target distinct refactoring goals relevant to library maintainability.
Neurosymbolic Augmentation: The framework extracts and integrates goal-state traces, relevant context, and auto-informalizations. The authors show this augmentation significantly boosts performance for both small and frontier models, facilitating the discovery of valid, high-value rewrites.

Experimental Results

The system was evaluated on miniCTX-v2, a benchmark of human-written research-level mathematics, using a 7B parameter base model (DeepSeek-R1-Distill-Qwen-7B).

Performance vs. Scale: ImProver 2 outperforms the much larger DeepSeek-R1 671B model on all three metrics (Length, Modularity, Dependency) when unscaffolded. For instance, on modularity, ImProver 2 achieved a mean improvement of 0.143 compared to 0.055 for the 671B model.
Comparison to Frontier Models: The 7B ImProver 2 model is competitive with mid-tier frontier models (e.g., GPT-5-mini) on length and dependency metrics. It leads all evaluated unscaffolded systems on modularity.
Impact of Scaffolding: Adding the neurosymbolic scaffold significantly improves performance across all model families. For example, the 7B model's length improvement score rose from 0.118 (base) to 0.236 (scaffolded). Even large models like GPT-5-high saw substantial gains (0.660 to 0.875) with the scaffold.
Iteration Dynamics: Most performance gains occur within the first 2–3 IRPO iterations. The system shows that small models can learn to restructure research-level proofs effectively, though performance eventually plateaus as the model absorbs common refactoring patterns.
Accuracy Trade-offs: The paper notes a trade-off between optimization and compilation accuracy. While optimized proofs often yield higher metric improvements, they may have slightly lower compilation rates than conservative rewrites, reflecting the riskier nature of structural transformations.

Significance and Claims

The paper posits that proof optimization is not solely a function of model scale but is significantly driven by task-specific specialization and iterative self-improvement. By combining a replay-buffered preference optimization pipeline with neurosymbolic scaffolding, ImProver 2 establishes that small, accessible models can effectively restructure complex, research-level proofs.

The authors claim that their work establishes proof optimization as a scalable, learnable task. They argue that their approach addresses the critical need for maintaining formal libraries and improving training data quality for future provers, offering a path to automate the refactoring of proofs for maintainability without relying exclusively on massive, expensive frontier models. The work highlights that with proper scaffolding and training, small models can match or exceed the performance of substantially larger systems on specific structural objectives.