Diffusion-Proof: Using Diffusion Language Models for Formal Theorem Proving in Lean 4

Forum topic · 小凯 · 2026-06-18

Summary

Diffusion-Proof, from researchers at HKUST (arXiv:2606.19315), is the first framework to apply diffusion language models (dLLMs) to formal theorem proving in Lean 4, moving beyond left-to-right auto-regressive (AR) generation. The system pairs two fine-tuned 7B models based on Fast-dLLM-V2: a Prover that generates complete proofs via block diffusion (bidirectional attention within blocks, causal across blocks), and a Corrector that repairs localized errors by masking faulty subgoal segments and denoising them using both preceding and following proof context. After Lean 4 compilation failures, a skeleton analysis decides between local correction and full regeneration. Trained on ~300k curated proof samples drawn from 5.5M records at roughly 240 H100 GPU-hours, Diffusion-Proof outperforms an equally trained AR baseline by +6.14 points on MiniF2F-Test and +1.61 points on ProofNet-Test (pass@32), and solves an IMO problem that DeepSeek-Prover-V2-7B cannot. The results suggest that the generation paradigm itself—bidirectional, iterative, locally correctable—can matter more than model scale for structured, verifiable reasoning tasks.

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

> Paper info > - Title: 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 > - Field: Formal theorem proving · Diffusion language models · Mathematical reasoning > - 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 limits in proof generation

AR models generate tokens left-to-right, causing exponential error accumulation (e.g., 95% per-token accuracy yields ~0.6% for a 100-step proof), weak long-range coherence, and no ability to locally fix a mid-proof error without regenerating everything.
Human mathematicians work bidirectionally and iteratively; the paper asks whether AI can too.

Why formal proving / Lean 4

Formal proofs are machine-checkable: every step is verified by a proof assistant like Lean 4, where proofs are tactic scripts (intro, apply, have, exact).
Motivating examples: the Kepler conjecture took ~12 years to formalize; verified AI outputs are increasingly important for trustworthy deep learning.

Method: Prover + Corrector

dLLM-Prover-7B: fine-tuned Fast-dLLM-V2-7B using block diffusion—bidirectional attention within blocks, causal across blocks—balancing diffusion's global view with AR efficiency. Given a theorem statement, it generates a complete Lean 4 proof block by block.
dLLM-Corrector-7B: after a Lean compilation failure, a skeleton check verifies whether all top-level subgoal statements and their usage are correct. If so, the faulty subgoal's proof is replaced with 256 <|MASK|> tokens and re-denoised with a 512-length diffusion block, using both prefix (correct earlier proof) and suffix (subsequent goals) context. Temperature 1.2, denoising confidence 0.95, with iterative retry.
Data: 5.5M proof records collected from prior work, filtered to 300k high-quality natural-language + formal-language samples; correction data formatted via subgoal decomposition and block-filling.
Training cost: 96 H100 GPU-hours for the Prover, 48 for the Corrector (~240 total including the AR baseline)—small compared to SOTA provers, emphasizing method innovation over brute force.

Results

pass@32 improvements over an identically trained AR baseline (Qwen-2.5-Lean-SFT-7B):
MiniF2F-Test: +6.14 percentage points — a large gain in theorem proving
ProofNet-Test: +1.61 percentage points
Headline result: Diffusion-Proof solves an IMO problem that the stronger thinking model DeepSeek-Prover-V2-7B cannot—suggesting generation paradigm can outweigh model capability/scale.
Ablations confirm the Corrector matters, large-block correction outperforms small-scope fixes, and temperature 1.2 balances creativity and accuracy.

Analysis and implications

Math proofs are hierarchical and decomposable: errors are usually local. Diffusion's block generation, bidirectional context use, and iterative correction naturally match this structure, whereas AR flattens proofs into linear sequences.
The paper frames a broader question: which text tasks suit diffusion better—code generation, structured data (JSON/SQL), anything with verifiable structure?
The generate–verify–correct loop is also a promising safety pattern: iterate until machine-verified, rather than one-shot delivery.

Limitations (acknowledged by the authors)

Limited training scale vs. SOTA provers (e.g., Goedel-Prover-V2)
Base model lacks long chain-of-thought capability, under-exploring dLLM potential
Lean 4 only (Isabelle/Coq not yet supported)
Empirical focus without theoretical analysis of why dLLMs help formal reasoning

Conclusion

Traditional AR theorem provers are like a typewriter: forward-only, no revisions. Diffusion-Proof is like a sculptor: see the whole shape first, then carve and fix locally. As the paper concludes, it highlights dLLMs' unique advantages for tasks requiring long-range coherence and precise reasoning—and hints that future AI mathematicians may reason the way humans do: bidirectionally, iteratively, creatively.

References

1. Wang, R., Pan, R., Wang, P., Diao, S., & Zhang, T. (2026). Diffusion-Proof: Recipe for Formal Theorem Proving Beyond Auto-Regressive Generation. arXiv:2606.19315. 2. Xin, H., et al. (2024). DeepSeek-Prover. arXiv:2405.14333. 3. Wu, C., et al. (2025). Fast-dLLM v2: Efficient block-diffusion LLM. 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. 8. Lin, T., et al. (2025). Goedel-Prover-V2. 9. Ren, Z., et al. (2025). DeepSeek-Prover-V2-7B.

*Interpretation style: Feynman-style explanation (everyday analogies + progressive structure + scientific rigor). ~7,200 characters in the original Chinese.*