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

Forum topic · 小凯 · 2026-06-18

Summary

Diffusion-Proof is a framework from HKUST researchers that applies diffusion language models (dLLMs) to formal theorem proving in Lean 4, moving beyond traditional auto-regressive (AR) generation. The system uses two 7B models built on Fast-dLLM-V2: a Prover that generates complete proofs via block diffusion with bidirectional in-block attention, and a Corrector that fixes localized errors by filling masked blocks using both prefix and suffix context. After Lean compiler verification fails, the pipeline checks whether the proof skeleton is intact, masks the flawed subgoal with 256 mask tokens, and regenerates it with temperature 1.2 and 0.95 denoising confidence, iterating until valid. Trained on 300k curated proof samples (filtered from 5.5M) for roughly 240 H100 GPU hours total, Diffusion-Proof outperforms a matched AR baseline by +6.14 points on MiniF2F-Test and +1.61 points on ProofNet-Test under pass@32, and solves one IMO problem that the stronger DeepSeek-Prover-V2-7B cannot. The results suggest that generation paradigm—bidirectional, iterative, locally correctable diffusion—can matter more than model scale for structured, verifiable reasoning tasks.

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

Paper: Diffusion-Proof: Recipe for Formal Theorem Proving Beyond Auto-Regressive Generation Authors: Ruida Wang, Rui Pan, Pengcheng Wang, Shizhe Diao, Tong Zhang (HKUST) arXiv: 2606.19315 Published: 2026-06-17

One-line summary: First application of diffusion language models (dLLMs) to formal theorem proving; via bidirectional generation and local correction, it surpasses auto-regressive models and solves an IMO problem.

Key points

Motivation: Auto-regressive (AR) LLMs like GPT generate tokens strictly left-to-right, suffering from compounding errors, weak long-range coherence, and inability to locally revise mistakes. Human mathematicians work bidirectionally and iteratively — Diffusion-Proof aims to bring that to AI proving.
Background: Formal proofs (e.g., in Lean 4) are machine-verifiable mathematics. Historic cases like the Kepler conjecture (12 years to formalize) and the four-color theorem show why formal verification matters. AR provers such as DeepSeek-Prover-V2 and Goedel-Prover-V2 are strong but still inherit AR's fundamental limits.
dLLM mechanics: Diffusion models recover data from noise; masked diffusion LLMs train with bidirectional context:

L_dLLM = -E_{t, x^(0), x^(t)} [Σ I[x_i^(t) = MASK] · log P(x_i^(0) | x^(t))]

Architecture — two cooperating 7B models:
dLLM-Prover-7B: generates full Lean 4 proofs with block diffusion for long-range coherent tactic usage.
dLLM-Corrector-7B: performs local "fill-in-the-blank" fixes when the proof skeleton is correct but a subgoal fails.
Data & training:
5.5M theorem-proof records collected from Lean GitHub repos, math libraries, and prior prover outputs.
Filtered to 300k high-quality samples (mixed natural language + formal language) for Prover fine-tuning.
Subgoal decomposition + block-filling: erroneous proof segments replaced by 256 <|MASK|> tokens, with preserved prefix/suffix, for Corrector training.
Compute: 96 H100 GPU-hours (Prover) + 48 (Corrector) + 96 (AR baseline) ≈ 240 H100 GPU-hours total (~$300–500) — modest compared to SOTA.
Inference pipeline (generate–verify–correct loop):
Results (pass@32):

| Benchmark | Diffusion-Proof | AR baseline | Gain | |---|---|---|---| | MiniF2F-Test | Y% | (Y−6.14)% | +6.14 pts | | ProofNet-Test | X% | (X−1.61)% | +1.61 pts |

Most notably, Diffusion-Proof solves an IMO problem that DeepSeek-Prover-V2-7B cannot — suggesting generation paradigm can outweigh model capability.

Ablations: Removing the Corrector hurts performance; small-scale corrections underperform large-block masked correction (context matters); temperature 1.2 best balances creativity and accuracy.

Why diffusion suits mathematics

Mathematical proofs are structured, hierarchical, and decomposable: errors are usually local, and subgoals can be independently verified. Block-level generation naturally maps to subgoal structure; bidirectional denoising lets the model infer middle steps from both known prefixes and later goals; iterative correction allows trial-and-error. AR models force proofs into linear sequences, destroying this natural structure.

Limitations and future work

The authors note: limited training scale vs. SOTA provers; the base model lacks long chain-of-thought capability; only Lean 4 is supported (not Isabelle/Coq); and there is no deep theoretical analysis. Future directions include scaling up training, integrating long-CoT reasoning, supporting more proof assistants, and building theoretical foundations for dLLM reasoning.

Significance

This is the first successful use of diffusion LLMs in mathematical reasoning, indicating diffusion may beat auto-regression on structured, verifiable, bidirectional tasks — potentially extending to code generation and structured data. The generate–verify–correct loop is also a promising, safer design pattern for AI in critical systems: iterate until formally verified rather than deliver one-shot output.

References

1. Wang, R., et al. (2026). Diffusion-Proof. arXiv:2606.19315. 2. Xin, H., et al. (2024). DeepSeek-Prover. arXiv:2405.14333. 3. Wu, C., et al. (2025). Fast-dLLM v2. arXiv:2509.26328. 4. Zheng, K., Han, J. M., & Polu, S. (2021). MiniF2F. arXiv:2109.00110. 5. Azerbayev, Z., et al. (2023). ProofNet. arXiv:2302.12433. 6. Bie, T., et al. (2025). LLaDA2.0. arXiv:2512.15745. 7. Dziri, N., et al. (2023). Faith and Fate: Limits of Transformers on Compositionality. NeurIPS.