DRIFT: Decompose, Retrieve, Illustrate, then Formalize Theorems

AI summary

Overview

Research area: Artificial Intelligence / automated mathematical reasoning — specifically retrieval-augmented autoformalization, the task of translating informal mathematical statements into machine-checkable formal statements (Lean, Rocq, Isabelle).

Technical level: Advanced. The paper assumes familiarity with theorem proving, formal libraries (Mathlib), dense retrieval, and metrics such as Typecheck and BEq+.

Scope: A single paper introducing the DRIFT framework (Decompose, Retrieve, Illustrate, then Formalize Theorems) and evaluating it on three autoformalization benchmarks — ProofNet, MiniF2F-test, and ConNF — across frontier LLMs (GPT-4.1, Claude-Opus-4, DeepSeek-V3.1) and one specialized model (Goedel-V2-8B).

What This Paper Is About

Large Language Models can reason about mathematics informally, but they struggle to write formal statements that a proof assistant like Lean will accept, largely because they cannot reliably find and correctly use the exact definitions a statement depends on. Existing retrieval-augmented methods query a library such as Mathlib with the entire informal statement at once, and then hand back definitions without showing how they are actually used in code. DRIFT addresses both problems by splitting the informal statement into smaller atomic sub-queries, retrieving one formal premise per sub-query, and adding illustrative theorems that demonstrate how those premises are applied.

Key Contributions

1. A four-stage decomposition-driven framework. DRIFT autonomously breaks an informal statement into atomic, concept-focused sub-queries, retrieves a formal premise for each, selects demonstrative theorems that use those premises, and finally generates the formal statement conditioned on all of this context.

2. New state-of-the-art results in dependency retrieval and autoformalization. The paper reports gains over prior retrieval-augmented methods on ProofNet and ConNF, using both frontier LLMs and specialized open-source formalization models, with especially large gains on the out-of-distribution ConNF benchmark.

3. A systematic analysis of when retrieval helps. The paper argues that the usefulness of retrieved dependencies is conditioned on the gap between a model's parametric knowledge and the complexity of the statement, and that different pipeline stages (decomposition versus formalization) demand different model capabilities.

4. Released code and models, available at the project's GitHub repository and a Hugging Face model page (DRIFT-dpr-mathlib).

Main Findings

• Decomposition sharply improves dependency retrieval. Averaged across decomposer models, DRIFT improves the F1 score by 13.34, 2.08, and 7.74 absolute points over a no-decomposition baseline on ProofNet, MiniF2F-test, and ConNF respectively. On ProofNet the paper describes this as nearly doubling the F1 score relative to the DPR baseline: the no-decomposition baseline reaches 13.77 F1, while DRIFT reaches 27.68 with Claude-Opus-4, 27.01 with DeepSeek-V3.1, and 26.64 with GPT-4.1.

• Frontier models are largely interchangeable as decomposers. The top-performing decomposers differ by at most 2.05% F1 on any benchmark. Claude-Opus-4 is best on ProofNet (27.68%) and MiniF2F-test (3.83%); DeepSeek-V3.1 is best on ConNF (36.88%).

• The Illustrate module covers most premises with very few examples. Using a greedy coverage algorithm with a budget of at most three theorems (m = 3), the module achieves an average premise coverage rate of 74.59 ± 4.80% across all decomposers and benchmarks.

• Large gains on the out-of-distribution ConNF benchmark. The abstract reports BEq+@10 improvements of 42.25% and 37.14% using GPT-4.1 and DeepSeek-V3.1 respectively. With GPT-4.1, DRIFT reaches 62.33 BEq+@10 versus 20.08 for DPR (RAuto) and 6.76 zero-shot — even surpassing the oracle* baseline of 58.90 by 3.43%.

• Retrieval helps where knowledge is missing, and can hurt where it is not. All models score below 10% zero-shot BEq+ on ConNF, confirming a severe knowledge gap that retrieval closes. On MiniF2F-test, which has an average of only 0.43 Mathlib dependencies, the zero-shot-to-oracle* pass@10 gap is just 2.24% for GPT-4.1 and 0.89% for DeepSeek-V3.1, and added context can act as a distractor.

• Specialized model behavior differs from frontier models. Goedel-V2-8B produces many syntactically valid but semantically incorrect statements on ConNF (low BEq+ despite high Typecheck at pass@10). DRIFT nearly doubles its zero-shot BEq+ scores and outperforms DPR (RAuto), but Goedel-V2-8B degrades severely under DPR (RAuto) on the low-dependency MiniF2F-test while DRIFT maintains its effectiveness.

• Ablation results identify each component's role. For GPT-4.1, removing Decompose drops ProofNet BEq+@1 from 17.38 to 13.64, and removing Illustrate drops ConNF BEq+@1 from 54.84 to 35.90 and BEq+@10 from 62.33 to 35.90-scale degradation (BEq+@10 falls by 18.94 points versus DRIFT). Removing retrieval entirely causes the largest losses, especially on ConNF (BEq+@1 falls by 50.37 points).

• Formalization reveals model differences that decomposition hides. DeepSeek-V3.1 generally beats GPT-4.1 zero-shot, but GPT-4.1 overtakes it once retrieval is introduced, particularly on ConNF and on ProofNet BEq+@10 with DRIFT (21.93% versus 20.59%). The paper attributes this to stronger in-context synthesis, and notes the two stages rely on distinct capabilities.

Methodology in Plain English

DRIFT is a prompting-based pipeline with four stages.

Decompose. An off-the-shelf LLM, given a few expert-verified examples from the Putnam benchmark as a few-shot prompt, splits each informal statement into sub-queries. Each sub-query pairs a natural-language description of one mathematical concept with the LLM's own guess at its formal name; that guessed name acts as a syntactic anchor while the natural-language part carries the meaning. The number of sub-queries is decided by the model, not fixed by the authors, and the paper reports that it ranges from 5.21 to 6.42 sub-queries per statement.

Retrieve. Each sub-query is encoded with a BGE-M3-based dense passage retriever that was finetuned on the dependency-retrieval task introduced with RAutoformalizer. The retriever finds the single closest formal object in the library by cosine similarity for each sub-query, and the results are pooled with duplicates removed, so the final set size is at most the number of sub-queries.

Illustrate. Because a bare definition does not show how to use it in code, a greedy algorithm picks up to three theorems from the library that collectively cover as many of the retrieved premises as possible, with ties broken by semantic similarity to the original informal statement.

Formalize. A formalizer model receives a concatenated prompt — instruction, retrieved premises (with full names, declarations and source code), illustrative theorems (as informal/formal pairs), and the original informal statement — and produces the formal statement. Ten generations are produced per problem to compute pass@k.

Evaluation uses two benchmark families of metrics: Precision, Recall and F1 against the oracle* dependencies for retrieval, and Typecheck (TC, syntactic validity) plus BEq+ (semantic equivalence, verified by deterministic proof tactics that prove each statement from the other) for formalization. The authors use BEq+ instead of LLM-as-judge to avoid stochasticity, noting its reported alignment with human judgments (Pearson 0.974, Kendall 0.872).

The three benchmarks test different regimes: ProofNet (374 undergraduate-textbook theorems integrated with Mathlib, averaging 3.39 dependent premises from over 243k formal objects including 139k theorems) for in-distribution retrieval; MiniF2F-test (224 Olympiad-style theorems, 20 removed for duplication or compile failure on Lean v4.18.0, averaging 0.43 dependencies) as a self-contained boundary condition; and ConNF (961 research-level theorems from the con-nf library, averaging 3.92 premises from 1,348 formal objects, formalizing a consistency proof for Quine's New Foundations) as an out-of-distribution test, since the retriever and formalizers were not exposed to its formal objects.

Why This Matters

Impact on research. Retrieval-augmented autoformalization has been limited by treating mathematical statements as single monolithic queries. DRIFT shows that aligning retrieval granularity with the dependency structure of formal mathematics is a practical lever, and it introduces a way to supply usage context (illustrative theorems) rather than bare definitions. Its finding that retrieved context can hurt on low-dependency problems, and that model-specific knowledge boundaries determine retrieval utility, is a design constraint for future retrieval-augmented systems.

Real-world applications:

• Formal proof assistants in mathematics. Autoformalization is described as a critical step toward automated theorem proving systems that ultimately assist mathematicians in new discoveries.
• Keeping pace with evolving libraries. Finetuned models hallucinate formal objects that have been renamed, reorganized, or deprecated in libraries like Mathlib; dynamic test-time retrieval, as in DRIFT, addresses this without retraining.
• Formalization by specialized models. The results on Goedel-V2-8B show that precise retrieval lets a small specialized model spend its capacity on formalization rather than on recalling library contents.
• Education and proof checking. Typecheck and BEq+ provide compiler-verified correctness signals that can be used to validate machine-written or student-written formal statements.

Industry relevance. The work was conducted in part at Huawei Noah's Ark Lab, London, alongside the University of Cambridge's Language Technology Lab, and the code and retriever model were released publicly. Formal verification pipelines in industry — where statements must be machine-checkable — could benefit from a retrieval layer that scales to large, changing libraries and that behaves predictably across different base models.

Future Directions

• Adaptive retrieval strategies. The paper concludes that retrieval effectiveness depends on model-specific knowledge boundaries and points toward strategies that can assess when external knowledge genuinely complements a model's capabilities, rather than retrieving unconditionally.
• Better handling of low-dependency problems. MiniF2F-test shows that injected context can be a distractor; the paper's error taxonomy names specific failure modes such as force-fitting and over-complication, which remain open problems.
• Avoiding oracle limitations.* Because DRIFT with GPT-4.1 surpassed the imperfect oracle* on ConNF by 3.43% BEq+@10, the paper suggests that illustrative theorems supply information that ground-truth necessary premises alone do not — an open question is how to define a better retrieval upper bound.
• Extending beyond Lean and Mathlib. The authors state the framework is language-agnostic and adaptable to other formal systems with structured libraries, but all reported experiments are conducted in Lean 4; generalization to Rocq or Isabelle is untested. They also note that contamination on ConNF could not be conclusively validated without access to the underlying training data.

Target Audience

Researchers and engineers working on autoformalization, retrieval-augmented generation, and LLM-based mathematical reasoning; practitioners building theorem-proving or formal-verification pipelines on top of Lean and Mathlib; and readers interested in how query decomposition and context augmentation transfer from multi-hop question answering to formal mathematics. Some familiarity with theorem provers and retrieval metrics is helpful, but the four-stage design is described in plain terms.

Authors’ abstract

Automating the formalization of mathematical statements for theorem proving remains a major challenge for Large Language Models (LLMs). LLMs struggle to identify and utilize the prerequisite mathematical knowledge and its corresponding formal representation in languages like Lean. Current retrieval-augmented autoformalization methods query external libraries using the informal statement directly, but overlook a fundamental limitation: informal statements lack direct mappings to mathematical theorems and lemmata, nor do those theorems translate trivially into the formal primitives of languages like Lean. To address this, we introduce DRIFT, a novel framework that enables LLMs to decompose informal mathematical statements into smaller, more tractable "sub-components". This facilitates targeted retrieval of premises from mathematical libraries such as Mathlib. Additionally, DRIFT retrieves illustrative theorems to help models use premises more effectively in formalization tasks. We evaluate DRIFT across diverse benchmarks (ProofNet, ConNF, and MiniF2F-test) and find that it consistently improves premise retrieval, nearly doubling the F1 score compared to the DPR baseline on ProofNet. Notably, DRIFT demonstrates strong performance on the out-of-distribution ConNF benchmark, with BEq+@10 improvements of 42.25% and 37.14% using GPT-4.1 and DeepSeek-V3.1, respectively. Our analysis shows that retrieval effectiveness in mathematical autoformalization depends heavily on model-specific knowledge boundaries, highlighting the need for adaptive retrieval strategies aligned with each model's capabilities.