💻 computer science

Diffusion-Proof: Recipe for Formal Theorem Proving Beyond Auto-Regressive Generation

The paper introduces **Diffusion-Proof**, the first framework to apply diffusion-based Large Language Models for formal theorem proving, which outperforms traditional auto-regressive baselines by leveraging iterative denoising for long-range coherence and local correction, achieving significant benchmark improvements and solving an IMO problem that state-of-the-art models could not.

Original authors: Ruida Wang, Rui Pan, Pengcheng Wang, Shizhe Diao, Tong Zhang

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ruida Wang, Rui Pan, Pengcheng Wang, Shizhe Diao, Tong Zhang

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 Big Idea: Fixing the "One-Way Street" of AI Math

Imagine you are trying to solve a very complex math puzzle, like a proof in a formal language called Lean. This language is like a strict computer code where every single step must be logically perfect. If you make one tiny mistake, the whole proof collapses.

For years, the best AI models for this job have been Auto-Regressive (AR) models. Think of these models like a person writing a story one word at a time, strictly from left to right. They can't see what they are about to write, and they can't easily go back and change a word they wrote five minutes ago without rewriting the whole sentence.

The Problem:
In long math proofs, this "one-way street" approach causes two big issues:

The Domino Effect: If the AI makes a small mistake early on, it keeps building on that mistake, making the error worse and worse until the proof is garbage.
No "Looking Back": If the AI realizes halfway through that the first step was wrong, it can't easily fix just that first step. It has to start over or awkwardly try to patch the end of the sentence to match the beginning.

The Solution: The "Diffusion" Approach

The authors of this paper propose a new way to build AI for math called Diffusion-Proof. Instead of writing word-by-word, they use a Diffusion Large Language Model (dLLM).

The Analogy: The Sculptor vs. The Scribe

The Old Way (AR Model): Imagine a scribe writing a letter. Once the ink is dry, they can't change it. If they misspell a word, they have to cross it out and write a messy note next to it.
The New Way (Diffusion Model): Imagine a sculptor working with a block of clay. They don't just add clay; they start with a rough, noisy lump and iteratively refine it. They can look at the whole shape at once, smooth out a bump on the left side while looking at the curve on the right side, and fix mistakes by "denoising" the whole picture until it's perfect.

In technical terms, the Diffusion model generates text in blocks (chunks of words) rather than single words. It can look at the beginning and the end of a sentence simultaneously to figure out what fits in the middle.

How Diffusion-Proof Works: The Duo Team

The paper introduces a two-part team to solve these math problems:

1. The Architect (dLLM-Prover-7B)

This is the main builder. Because it uses the "sculptor" method (block diffusion), it is much better at long-range planning.

Why it matters: When building a long proof, the Architect can see the "big picture." It knows that the final step requires a specific condition, so it sets up the beginning of the proof to match that condition perfectly. It doesn't get lost in the middle because it can "see" the whole block of text at once.

2. The Inspector (dLLM-Corrector-7B)

Even the best Architect makes mistakes. Sometimes the proof looks good, but a specific step is wrong.

The Old Way: If an AR model makes a mistake, it often tries to fix it by writing more text after the error, which usually makes things worse.
The New Way: The Inspector is a special model trained to fill in the blanks. If the proof fails, the system takes the broken part, covers it with a "mask" (like a blank space), and asks the Inspector to rewrite just that specific chunk.
The Superpower: Because the Inspector uses the diffusion method, it can look at the text before the error and the text after the error to figure out exactly what the missing piece should be. It's like a editor who reads the sentence before and after a typo to guess the right word, rather than just guessing based on the previous word.

The Results: Beating the Giants

The researchers tested this new system on two famous math benchmarks:

MiniF2F: A collection of high-school and competition math problems.
ProofNet: A collection of college-level math problems.

The Scoreboard:

The new Diffusion-Proof system beat the best traditional (Auto-Regressive) models trained on the exact same data.
It solved 6.14% more problems on the MiniF2F test.
It solved 1.61% more problems on the ProofNet test.

The "IMO" Moment:
The most exciting result was a specific International Math Olympiad (IMO) problem. A very advanced, famous AI model called DeepSeek-Prover-V2 (which uses "Chain of Thought" reasoning) failed to solve it. Diffusion-Proof solved it.

Why? The DeepSeek model got stuck in a loop of trying to fix its own plan, but it couldn't see the long-range connection between the start and end of the proof. The Diffusion model, with its ability to look at the whole "block" of the proof at once, found the correct path.

Summary

The paper claims that by switching from "writing one word at a time" to "sculpting blocks of text," AI can become much better at formal math. It creates a system that plans better over long distances and can fix its own mistakes by looking at the whole context, not just the immediate past. This allows it to solve hard math problems that previous AI models simply couldn't crack.

Technical Summary: Diffusion-Proof

Problem Statement

Formal theorem proving requires maintaining long-range coherence and logical consistency across complex sequences of reasoning steps. While Large Language Models (LLMs) have advanced significantly in this domain, current state-of-the-art approaches rely heavily on Auto-Regressive (AR) models. These models generate text sequentially from left to right, predicting the next token based solely on previous tokens. This architecture suffers from inherent limitations, including:

Error Accumulation: Small errors in early tokens can compound exponentially over long sequences, leading to logical inconsistencies.
Lack of Long-Range Coherence: AR models struggle to maintain a coherent global plan for a proof, often failing to align intermediate tactics with the final goal.
Inability to Perform In-filling: AR models cannot naturally perform "in-filling" corrections (modifying a middle section of a proof while preserving the surrounding context), which is crucial for fixing local errors without rewriting the entire proof.

Although Diffusion LLMs (dLLMs) have shown promise in coding and general text generation by utilizing iterative denoising and bidirectional context, their application to formal mathematical reasoning remains largely unexplored.

Methodology

The authors propose Diffusion-Proof, the first framework designed to train and apply dLLMs specifically for formal theorem proving. The framework operates within the Lean4 environment and utilizes a dual-model architecture: dLLM-Prover-7B and dLLM-Corrector-7B.

1. Data Collection and Preprocessing

Dataset: The authors collected 5.5 million theorem proving records from previous works.
Cleaning: They standardized the data using rule-based methods to separate Natural Language (NL) statements, Formal Language (FL) statements, and FL proofs.
SFT Dataset: 300,000 records were selected to create a Supervised Fine-Tuning (SFT) dataset, maintaining a 1:2 ratio of pure Lean proofs to those with NL annotations.
Corrector Data: A subset of data featuring subgoal decomposition (marked by have keywords) was processed. Proof blocks were filled with placeholder tokens to create training data for "in-filling" corrections.

2. Model Architecture and Training

The framework is built upon Fast-dLLM-V2-7B, a block-based diffusion model. Unlike traditional dLLMs that mask tokens independently, this model divides sequences into blocks of size $B$ with bidirectional attention restricted to individual blocks while preserving causal relationships between blocks.

dLLM-Prover-7B:
- Goal: Perform whole-proof generation with long-range coherent tactic usage.
- Training: Fine-tuned on the 300k SFT dataset using a standard text-to-text framework.
- Strategy: Employs curriculum learning, organizing training data by proof complexity (increasing proof length) to stabilize the loss curve. The diffusion block size is set to 32.
dLLM-Corrector-7B:
- Goal: Perform local proof correction using bidirectional information.
- Training: Fine-tuned on the 128k subgoal decomposition dataset.
- Innovation: Introduces a Large Block Correction method. The diffusion block size is extended from 32 to 512. The model is trained to replace erroneous subgoal proofs with a sequence of mask tokens and then regenerate the block using the extended context (both prefix and suffix) to perform in-filling corrections.

3. Inference Pipeline

Generation: The dLLM-Prover-7B generates a complete proof for a given theorem statement using block-diffusion generation.
Verification: The proof is verified by Lean.
Correction: If the proof fails verification but the top-level skeleton (subgoal statements) is correct, the dLLM-Corrector-7B is invoked. It replaces the erroneous subgoal proofs with mask tokens and regenerates them using a 512-token diffusion block, leveraging bidirectional context to fix local errors while maintaining global coherence.

Key Contributions

First dLLM Framework for Theorem Proving: Diffusion-Proof is the first framework to adapt diffusion models for formal theorem proving, addressing the specific challenges of long-range coherence and error correction in mathematical reasoning.
Large Block Correction Training: The authors introduce a novel training method that leverages large diffusion blocks (512 tokens) and bidirectional information to enable effective in-filling corrections, a capability native to diffusion models but absent in AR models.
Dual-Model Synergy: The framework demonstrates a synergistic approach where a prover handles global planning and a corrector handles local refinement, outperforming single-model baselines.

Experimental Results

The framework was evaluated on MiniF2F-Test (244 problems) and ProofNet-Test (186 problems) using the pass@32 metric.

Performance Gains: Compared to a baseline AR model (Qwen-2.5-Lean-SFT-7B) trained on the exact same dataset:
- MiniF2F-Test: Diffusion-Proof achieved 50.00% accuracy, a 6.14% absolute improvement over the baseline (43.85%).
- ProofNet-Test: Diffusion-Proof achieved 7.53% accuracy, a 1.61% absolute improvement over the baseline (5.91%).
Problem Type Analysis: The model showed significant gains in challenging categories, particularly IMO (International Mathematical Olympiad) problems, where it improved from 5.00% to 15.00%.
Case Study: Notably, Diffusion-Proof successfully solved the IMO problem imo_1962_p2, which the more advanced DeepSeek-Prover-V2-7B (equipped with Long Chain-of-Thought reasoning) failed to solve under pass@32. The analysis indicated that the AR model struggled with the long-range dependencies and coherent tactic planning required, whereas the diffusion corrector successfully utilized bidirectional context to resolve the proof.
Ablation Studies:
- Removing the corrector reduced performance by 1.64% on MiniF2F, confirming the value of the correction step.
- Even without the corrector, the dLLM-Prover alone outperformed the AR baseline by 4.51%, validating the inherent advantages of the diffusion architecture for formal reasoning.

Significance and Claims

The paper claims that Diffusion-Proof highlights the unique advantages of diffusion models in formal theorem proving, specifically their ability to maintain long-range coherence and perform bidirectional in-filling corrections.

Architectural Advantage: The results suggest that the diffusion generation paradigm offers a structural advantage over AR models for tasks requiring rigorous logical consistency over long sequences, as evidenced by the lower validation loss distributions and higher success rates on complex proofs.
Future Potential: The authors position this work as a foundational step, demonstrating that dLLMs can surpass AR baselines even without Reinforcement Learning (RL) or Long CoT reasoning. They note that while current limitations include computational scale and the lack of RL integration, the framework provides a promising direction for future research into models with enhanced reasoning capabilities.

The authors emphasize that their primary goal was to investigate the architectural potential of dLLMs rather than to achieve state-of-the-art performance through massive scaling or RL, making the observed improvements a strong indicator of the diffusion paradigm's efficacy in formal reasoning.