Proof Artifact Co-training for Theorem Proving with Language Models

Introduction

Deep learning-driven automated theorem proving in large libraries of formalized mathematics (henceforth “neural theorem proving”) has been the focus of increased attention in recent years. Labeled data for imitation learning of theorem proving is scarce—formalization is notoriously labor-intensive, with an estimated cost of 2.5 man-years per megabyte of formalized mathematics (Wiedijk, 2000), and complex projects require years of labor from human specialists. Within a fixed corpus of (possibly unproven) theorem statements, it is possible to augment a seed dataset of human proofs with new successful trajectories using reinforcement learning or expert iteration. However, for some large models this can be quite computationally intensive, and without a way to expand the curriculum of theorems, the agent will inevitably saturate and suffer from data starvation.

Data scarcity is a particularly thorny obstruction for applying large language models (LLMs) to neural theorem proving. LLMs have achieved spectacular success in data-rich regimes such as plain text (Brown et al., 2020), images (Dosovitskiy et al., 2021), and joint text-image modeling (Radford et al., 2021), and the performance of decoder-only Transformers has been empirically shown to obey scaling power laws in model and data size (Henighan et al., 2020). However, existing datasets of human proof steps for neural theorem proving are extremely small and exist at scales at which overfitting occurs extremely rapidly, disrupting the scaling of performance with respect to model size (Kaplan et al., 2020).

We make two contributions towards addressing the problem of data scarcity in the context of formal mathematics. First, we introduce PACT (Proof Artifact Co-Training), a general methodology for extracting self-supervised auxiliary tasks for jointly training a language model alongside a tactic prediction objective for interactive theorem proving. Second, we present LeanStep, a collection of datasets and a machine learning environment for the Lean 3 theorem prover with support for PACT, supervised learning of tactic prediction, theorem proving evaluation, and reinforcement learning.

We train large language models on these data and demonstrate that PACT significantly improves theorem proving success rate on a held-out suite of test theorems, from 32% to 48%. We then embark on a careful study of the effects of pre-training vs. co-training and show that PACT combined with WebMath pre-training (Polu & Sutskever, 2020) achieves the best validation loss and theorem proving success rate. Finally, on an out-of-distribution collection of thousands of theorems (some involving novel definitions) added to Lean’s mathematical library after we extracted our train/test data, we achieve a theorem proving success rate of $37\%$, suggesting strong generalization and usefulness at the frontier of formalized mathematics.

Background and related work

Lean is an interactive theorem prover and functional programming language (de Moura et al., 2015). It has an extremely active community and is host to some of the most sophisticated formalized mathematics in the world, including scheme theory (Buzzard et al., 2021), forcing (Han & van Doorn, 2020), perfectoid spaces (Buzzard et al., 2020), and condensed mathematics (Scholze, 2020). Lean’s foundational logic is a dependent type theory called the calculus of inductive constructions (Pfenning & Paulin-Mohring, 1989). This design means that terms, types and proofs are all represented with a single datatype called an expression. A proof term is a Lean expression whose type is a proposition, i.e. a theorem. This proof term serves as a checkable artifact for verifying the proposition. Lean uses a small, trusted kernel to verify proof terms. The primary repository of formalized mathematics in Lean is mathlib (mathlib, 2020). At the time of writing, 140 contributors have added almost 500,000 lines of code; mathlib contains over 46,000 formalized lemmas backed by over 21,000 definitions, covering topics such as algebraic geometry, computability, measure theory, and category theory. The range of topics and the monolithic, unified organization of mathlib make it an excellent foundation for a neural theorem proving dataset.

Machine learning in interactive theorem proving

In a tactic-based interactive theorem prover (ITP) such as Lean, a proof is a list of tactics, i.e. small proof-term-generating programs. Tactics can be simple one-word commands, e.g. refl, or be composed of many nested parts, e.g.

Here the brackets enclose a list of simplifier rules (which often are just lemmas from the library), and @norm_eq_zero α _ g is a proof term applying the lemma norm_eq_zero to the local variables α and g.

Other ML and neural theorem provers for tactic-based ITPs take one of two approaches to tactic generation. TacticToe (Gauthier et al., 2021) for HOL4 and Tactician (Blaauwbroek et al., 2020) for Coq use k-NN to select similar tactics in the training set and apply modifications to the result, e.g. swapping the tactic variables with those found in the local context. HOList/DeepHOL (Bansal et al., 2019b; a; Paliwal et al., 2020) for HOL Light; TacticZero (Wu et al., 2021a) for HOL4; and CoqGym/ASTactic (Yang & Deng, 2019) and ProverBot9001 (Sanchez-Stern et al., 2020) for Coq hard-code the DSL for every tactic command. The model chooses a tactic command, and then fills in the tactic arguments using specialized argument selectors (such as a lemma selector, a local hypothesis selector, and/or a variable selector). None of these selectors currently synthesize arbitrary terms. This prevents the tactic synthesis from constructing tactics with proof terms, such as @norm_eq_zero α _ g, or directly proving an existential, e.g. ∃ (x : ℝ), x + 3 = 0, by supplying the witnessing term -3.

Directly applying generative language modeling to tactic generation allows this setup to be considerably simplified. Our tactic generator is able to synthesize tactics of any form found in mathlib including, for example, the simpa example above as a one line proof to a test theorem, even though the string @norm_eq_zero does not occur in our dataset. (See more examples in Appendix D.) We leave as future work the possibility of re-integrating specialized components, e.g. lemma selection, found in other works (possibly as, say, a source of additional prompts for the language model).

Language models have also been explored in the first-order ITP Mizar for conjecturing and proof synthesis (Urban & Jakubuv, 2020). While their work shows the promise of such approaches, is not intended as a complete end-to-end theorem prover. For Metamath, which does not use tactics, language modeling approaches have been quite successful. Holophrasm (Whalen, 2016), MetaGen (Wang & Deng, 2020), and GPT-f (Polu & Sutskever, 2020) all use RNNs or Transformers to generate proof steps. Indeed, our paper builds on the work of Metamath GPT-f (Polu & Sutskever, 2020) (MM GPT-f). Whereas MM GPT-f trained primarily on the Metamath proof step objective (i.e. guessing the next lemma to be applied to a goal, which is similar to our NEXTLEMMA task in Section 3.2), we co-train on a diverse suite of self-supervised tasks extracted from Lean proof terms and demonstrate significant improvements in theorem proving performance when doing so. This is our main result.

Reasoning with Transformers

Besides theorem proving, a number of recent papers have shown that language models, especially Transformers, are capable of something like mathematical and logical reasoning in integration (Lample & Charton, 2020), differential equations (Charton et al., 2021), Boolean satisfiability (Hahn et al., 2021), and inferring missing proof steps (Li et al., 2021).

A closely-related vein of work has shown that pre-training Transformers on data engineered to reflect inductive biases conducive to mathematical reasoning is beneficial for downstream mathematical reasoning tasks (Rabe et al., 2021; Wu et al., 2021b). Our work both builds on and departs from these ideas in several ways. Unlike skip-tree training (Rabe et al., 2021), which focuses solely on predicting masked subterms of theorem statements, PACT derives its self-supervised training data from far more complex proofs. Unlike LIME (Wu et al., 2021b), which uses purely synthetic data and is presented as a pre-training methodology, our self-supervised tasks are extracted from non-synthetic human proofs. Moreover, we show that not only are Transformers capable of performing well on auxiliary tasks gathered from low-level proof artifact data, but that we can directly leverage this low-level data by jointly training a language model to greatly improve its performance at high-level theorem proving.

Machine learning with proof artifacts

The idea of mining low-level proof artifacts was previously explored by Kaliszyk and Urban in the context of automated lemma extraction (Kaliszyk & Urban, 2015b; Kaliszyk et al., 2015). It has also been previously observed that training on fully elaborated Coq terms (Nie et al., 2020) helps with a downstream theorem naming task. However, similar to previous work on skip-tree training, their dataset focuses solely on theorem statements, i.e. types, does not cover the far more complex proof terms, and does not evaluate the effect of such training on theorem proving evaluations.

While there exist environments and datasets for other formal mathematics libraries (Kaliszyk et al., 2017; Li et al., 2021; Huang et al., 2019; Kaliszyk & Urban, 2015a), LeanStep is the first and only tactic proof dataset for the Lean theorem prover. This makes available a large set of formal mathematical data to researchers covering a diverse and deep spectrum of pure mathematics. Moreover, LeanStep is unique in that it contains both high-level human-written tactics as well as kernel-level proof terms, which enables the extraction of self-supervised tasks for PACT (Section 3.2).

The LeanStep datasets and machine learning environment

Tactics in Lean are metaprograms (Ebner et al., 2017), which can construct Lean expressions, such as proof terms. A tactic state which tracks the list of open goals and other metadata (like the partial proof term constructed so far) is threaded through each tactic invocation. Lean has special support for treating tactics as an extensible domain-specific language (DSL); this DSL is how Lean is typically used as an interactive theorem prover. The DSL amounts to a linear chain of comma-separated invocations. The Lean proof step task is to predict the next tactic given this goal state. We refer the reader to Appendix A for examples and further explanation.

Our human tactic proof step dataset consists of source-target pairs of strings, one for each tactic invocation in the Lean core library and in mathlib. The source string is the pretty-printed tactic state. The target string is the tactic invocation as entered by a human author of the source code. This data is gathered by hooking into the Lean parser and Lean’s compilation process. We refer to the task of predicting the next human tactic proof step given a tactic state as the proofstep objective.

2 Proof artifact co-training

In this section, we describe the PACT task suite and how data for these tasks are extracted.

For every proof term $\tau$, we record the type $\Gamma$ of $\tau$, its name nm, and a list ps of all premises (i.e. named references to other lemmas in the library) which are used in $\tau$. We then recurse through $\tau$, tracking a list bs of bound variables which we update whenever navigating into the body of a $\lambda$-expression. At every sub-term $\tau^{\prime}\subseteq\tau$ we record $\tau^{\prime}$, its type $\Gamma^{\prime}$, the current state of bs, and the following data:

A tactic state, where the goal is set to be $\Gamma^{\prime}$ and the list of hypotheses in the local context is set to be the list bs, i.e. those bound variables in scope at $\tau^{\prime}$.

A partial proof term, i.e. $\tau$ with $\tau^{\prime}$ masked out.

A premise selection bitmask, i.e. Boolean labels for every p in ps indicating whether p is used in $\tau^{\prime}$.

A local context bitmask, i.e. similar Boolean labels for every b in bs indicating whether b is used in $\tau^{\prime}$.

An optional next lemma: if the first step of $\tau^{\prime}$ is to apply a premise p in ps, we record p.

Whenever we record a term, we record both pretty-printed and far more explicit fully elaborated versions of it. The fully elaborated terms explicitly display enormous amounts of type information which are usually silently inferred by Lean. From these data, we assemble the following language modeling tasks:

Next lemma prediction. Given the tactic state, predict the next lemma to be applied.

Proof term prediction. Given the tactic state, predict the entire proof term $\tau^{\prime}$.

Skip-proof. Given the partial proof term, predict the masked subterm $\tau^{\prime}$.

Type prediction. Given the partial proof term, predict the type $\Gamma^{\prime}$ of the masked subterm $\tau^{\prime}$.

Tactic state elaboration. Given the tactic state, predict the fully elaborated tactic state.

Proof term elaboration. Given $\tau$, predict the fully elaborated version of $\tau$.

Premise classification. Given the tactic state and a premise p $\in$ ps, predict either or according to the premise selection bitmask.

Local context classification. Given the tactic state (which consists of a list of local assumptions bs and the goal $\Gamma^{\prime}$), predict the sublist of bs which is true on the local context bitmask.

Theorem naming. Given the type $\Gamma$ of the top-level proof term $\tau$, predict the name nm.

We remark that our next lemma prediction task is precisely the low-level PROOFSTEP objective studied in (Polu & Sutskever, 2020), and our skip-proof task superficially resembles, but is much more difficult than the skip-tree task studied in (Rabe et al., 2021), as proof terms tend to be far more complex than the syntax trees of theorem statements.

3 The LeanStep machine learning environment

We instrument Lean for automatic theorem proving with a language model, including utilities for (1) setting the runtime environment at a particular theorem (ensuring proofs are never circular), (2) serializing the tactic state as environment observations for a theorem-proving agent, (3) exposing Lean’s parser to re-parse strings emitted by a language model into tactic invocations, and (4) executing and capturing the results of the re-parsed tactics, enabling the recording of trajectories for expert iteration and reinforcement learning.

In addition to this general instrumentation, we implement a generic best-first search algorithm for theorem proving; it forms the basis for our evaluations and is written entirely in Lean itself. The algorithm is parametrized by an oracle (Ω : tactic_state → list (string × float)) that accepts a tactic state and returns a list of strings and heuristic scores. The search is controlled by a priority queue of search nodes, which consist of a tactic state (i.e. a partial proof) and search metadata. In the outer loop of the algorithm—which continues until either the theorem is completely proved (i.e. no goals are remaining on the current node), the priority queue is empty (i.e. the search has failed), or a pre-set timeout or budget of iterations is exceeded—we pop a node off the queue, serialize the associated tactic state and use it to query the oracle, producing a list of candidates cs : list (string × float). We then loop over the candidates cs to produce a list of new search nodes, by re-parsing each string into a tactic and adding a new node if the parsed tactic advances the proof without raising errors. These new search nodes are then re-inserted into the queue in order of decreasing priority and the search continues. We optionally constrain the search by enforcing maximum width and depth limits $w_{\operatorname{max}}$ and $d_{\operatorname{max}}$ that guard insertion into the queue. When considering nodes for insertion, any node whose depth exceeds $d_{\operatorname{max}}$ is ignored, and all nodes are ignored if the queue size is strictly larger than $w_{\operatorname{max}}$. Due to the flexibility in assigning heuristic scores and in choosing the maximum width and depth hyperparameters, our algorithm is quite general—for example, it reduces to (1) a greedy depth-first search when $w_{\operatorname{max}}=0$, and (2) a naïve breadth-first search when heuristic scores are identical and $w_{\operatorname{max}}=d_{\operatorname{max}}=\infty$.

Experiments

In all of our experiments, we use decoder-only Transformers similar to GPT-3 (Brown et al., 2020). Unless mentioned otherwise, all of our models have $24$ layers with $d_{\operatorname{model}}=1536$ and $24$ heads, accruing to $837$M trainable parameters. They are also pre-trained on WebMath (Polu & Sutskever, 2020) for $72$B tokens. We use the standard BPE encoding (Brown et al., 2020), a batch size of $512$ and a learning rate of $0.00025$ with a cosine schedule and a $100$-step ramp-up.

We use an 80-5-15 train-validation-test split. We split all datapoints deterministically by theorem name, by hashing each name to a float in $(0,1)$. This ensures, for example, that proof steps used to prove a test theorem never appear in the training data and vice-versa.

When fine-tuning a model we load its saved parameters but re-initialize the optimizer. We start each training for a fixed number of tokens (defining the cosine schedule) and record the number of tokens consumed as we reach a minimal validation loss. We use the minimum validation loss snapshot to evaluate each model on our held-out test set.

We partition our datasets into three groups:

tactic: the dataset described in Section 3.1.

mix1: the union of the PACT tasks next lemma prediction and proof term prediction (Section 3.2), selected because of their close relation to tactic.

mix2: all other datasets described in Section 3.2.

This grouping is motivated by the impossibility to ablate each dataset separately given our compute budget. They nonetheless enable us to study the effect of tasks that are very close to the tactic objective in comparison to others. Our choice of next lemma prediction and proof term prediction for mix1 is motivated by the observation that these tasks are closely related to the theorem proving objective: a proof can be given entirely in terms of a sequence of lemmas to apply (as in Metamath), or the proof can be finished in one step by supplying the entire proof term. Despite their logical similarity to the PROOFSTEP objective, we nevertheless use different keywords in the prompt to the model to disambiguate (NEXTLEMMA and PROOFTERM) from (PROOFSTEP) because the data is noisy and represents a significant distribution shift: during pretty-printing, subtrees of proof terms beyond a certain depth are dropped entirely, there is generally no guarantee that they can be re-parsed, and the data is much more verbose than what humans typically supply in source code.

Theorem proving evaluation

We run theorem-proving evaluations on our held-out test set, comprising $3071$ theorems. Since the split was conducted by theorem name, the proofs of these theorems never appear in the training data. For each theorem in the test set, we set the runtime environment to the location where the theorem is proved in the source code, preventing the use of theorems defined later in mathlib and ensuring that we never derive circular proofs. We compare against existing proof automation In Lean by also evaluating the tactics refl, which attempts to prove statements via definitional equality, and tidy, which conducts a greedy depth-first search using a fixed list of tactics at each step. We re-implement tidy as a special case of our best-first search algorithm using an oracle which always emits the same list of tactics, and so henceforth refer to it as tidy-bfs. In all of our experiments, we use a maximum width of $w_{\operatorname{max}}=16$, a maximum depth of $d_{\operatorname{max}}=128$, a maximum budget of $512$ iterations of the outer loop, a timeout of $5$ seconds per tactic execution, and a global timeout of $600$ seconds per theorem. Because sampling completions from our models is much slower ($\approx 1$ second) than querying the constant tidy-bfs oracle (instantaneous), the tidy-bfs search runs many more iterations than gptf before timeout.

We report the pass-rate (i.e. percentage of theorems proved) from the randomly-chosen held-out test set, following (Whalen, 2016), (Bansal et al., 2019c), and others. We provide an alternative pass-rate at the end of this section, using theorems added to mathlib after our dataset was collected. We average over three evaluation runs when reporting the pass rate.

Effect of co-training vs pre-training

We first study the effects of pre-training versus co-training with the mix1 and mix2 datasets. We pre-train using the methodology described above (potentially pre-training first on WebMath, and then on a PACT dataset in sequence). For co-training, we simply concatenate and shuffle the datasets together without applying any particular weight to a given dataset.

The main results are presented in Figure 2. Pre-training exhibits an effective transfer from mix-1 and/or mix-2 but the best result is achieved by co-training with both these datasets. With this setup, we are able to train for much longer (71B tokens vs 22B+18B for the best pre-training setup) before overfitting on the PROOFSTEP task. We hypothesize that PACT regularizes overfitting to the PROOFSTEP task while still imparting useful knowledge to the model due to large amounts of mutual information, and that this is the main driver of increased performance.

Ablating WebMath pre-training

Next, we ablate the effect of WebMath pre-training (instead starting with a model pre-trained on the same English language mix as GPT-3). As expected, co-trained models suffer from a performance drop without Webmath pretraining. but we were more interested in measuring the effect on pre-trained models on mix-1 and mix-2, as they may not benefit from WebMath as much due to the two successive pre-training steps.

We report the optimal validation losses in Figure 3. WebMath appears as substantially beneficial even in the sequential pre-training setup. This indicates that PACT is not a replacement for WebMath pre-training, but rather a complementary method for enhancing the performance of language models for theorem proving.

Ablating regularization

We rule out the possibility that the benefits from PACT come from simply regularizing our models on the scarce tactic data alone. We checked that a WebMath > tactic model trained with 15% residual dropout achieved a minimum validation loss of 1.01 and 33.6% pass rate, far below the 48.4% PACT pass rate.

Effect of model size

Finally, we study how performance scales with respect to model size. We use the best training setup reported in Figure 2, WebMath > mix1 + mix2 + tactic. The 837m model is our main model. The 163m and 121m models respectively have $12$ and $6$ layers, with $d_{\operatorname{model}}=768$. The learning rates are respectively adjusted to $0.0014$ and $0.0016$.

As demonstrated by Figure 4, performance is highly correlated with model size, with larger models generally achieving better generalization even in the overfitted regime. We leave as future work a careful study of how evaluation performance is affected when scaling to multi-billion parameter models, as well as the feasibility of deploying them for interactive use by Lean users.

Time-stratified evaluation

In the 5 week period that separated our last dataset extraction and the writing of this paper, mathlib grew by 30K lines of code, adding 2807 new theorems. Evaluating our models on these new theorem statements gives a unique way to assess their capability to assist humans in formalizing proofs and to test their generalization to completely unseen theorems and definitions. This evaluation set also addresses one of the weaknesses of using a random split of theorems from a formal mathematics library, namely that the split is non-chronological; e.g. test theorems can appear as lemmas in proofs of train theorems.

We call this temporally held-out test set future-mathlib and evaluate our best model as well as the refl and tidy-bfs baselines on it. In contrast to evaluation on our test split, the refl baseline (simply attempting a proof by the refl tactic) closes $328$ proofs ($11.6\%$), demonstrating an important skew towards trivial boilerplate lemmas generally defined to provide alternate interfaces to new definitions. The tidy-bfs baseline closed $611$ proofs ($21.8\%$), and our best model wm-tt-m1-m2 closed $1043$ proofs ($37.1\%$), proving $94\%$ of the refl lemmas. We attribute the weaker performance to heavy distribution shift: by the nature of the dataset, the future-mathlib theorems frequently involve new definitions and concepts which the model was never exposed to during training. Nevertheless, the success rate remains high enough to suggest strong generalization and usefulness at the frontier of formalized mathematics.

Discussion

In Lean, multiple tactic commands can be chained together using semicolons. Our data pipeline treats these tactic chains as a single sequence in our training data, and they are occasionally predicted by the model. Such chained tactic applications are difficult for human formalizers to synthesize on their own, as they require reasoning about the semantics of multiple tactics in sequence and their effects on the tactic state, and the examples present in the training data are usually optimized by hand from longer, less succinct proofs. We observed that PACT significantly boosts the capability of our models to successfully predict longer chained tactic applications. This occurs despite the fact that the tactic chaining idiom is specific to the tactic proofstep dataset and does not appear in the PACT data whatsoever. We supply more detail in Section C.1.

Theorem naming

We also evaluate our best PACT model (wm-to-tt-m1-m2) on the theorem naming task, using the theorem statements and human-supplied names from the future-mathlib evaluation set. It achieved 20% acc@1, 27% acc@10, and 30% acc@16. An inspection of its outputs reveals that even when its predictions diverge from the ground truth, they are often idiomatic and semantically correct alternatives. We supply more detail in Section C.2.

Impact on Lean community

Lean’s mathlib (mathlib, 2020) is a rapidly growing open source library of formal mathematics which has grown considerably in size each year for the past four years.See https://leanprover-community.github.io/mathlib_stats.html for up-to-date statistics on mathlib’s size and growth over time. Our work has been welcomed by members of this community, with Lean power users describing some of the new proofs found by GPT-f as “nontrivial” and “clever”. More than one-third of the proofs found by our models are shorter and produce smaller proof terms (sometimes by several orders of magnitude) than the ground truth. Manually inspecting a small, non-cherry picked sample of these shorter proofs has led to 19 GPT-f co-authored commits to mathlib, some of which reduce proof term sizes and theorem compilation times by an order of magnitude (see Appendix D).

Potential societal impact

Strong automated reasoning systems have enormous potential impact for mathematical research and scientific progress in other disciplines. The methods that we discuss in this paper could accelerate the development of strong automated reasoning systems. We have also observed that our language models absorb stylistic biases from their training data which could be amplified via reinforcement learning. However, since we focus on mathematics codified in proof assistants, we believe that there is little immediate negative societal impact from our work.

Future directions

There are many elaborations on the training data, training methodology, and tree search wrapping lean-gptf which can be reasonably expected to improve its performance at theorem proving. Our dataset can be synthetically augmented using similar methods as (Polu & Sutskever, 2020). Our dataset could be cleaned further, and proofs minimized. Merely making the decoded rewrites robust by only using the largest prefix of successful rewrites significantly boosts the success rate of suggested rewrites. In a similar vein, predicted lemmas generated as arguments to unsuccessful tactic applications could be cached and re-used as hints for an intermittently-queried hammer. The increased success rate of chained tactic predictions mentioned above shows the feasibility of having language models perform multiple reasoning steps in a single query, potentially improving the efficiency of the proof search. From the experiments described in Section 4, it is clear that the composition of the dataset used for co-training significantly affects performance on theorem proving. Although we uniformly sampled across all co-training tasks, it would be interesting to optimize a dynamic mixture schedule, perhaps annealing towards a desired task.

Conclusion

There is a sense in which PACT is merely an application of the well known principle that compute in the form of search should be exchanged for training signal whenever possible. In Lean, typeclass inference relies on a backtracking Prolog-style search; the elaborator performs search to disambiguate overloaded notation and infer types; Lean tactics have complex semantics precisely because they can perform search to find subproofs automatically. The work done by these subroutines is preserved in the proof artifacts, and PACT can be viewed as a way of extracting this information offline for more training signal.

We have presented PACT as a way of addressing the data scarcity issue for learning theorem proving from human tactic scripts in proof assistant libraries. Another well-studied solution for this is expert iteration and reinforcement learning. In the setting of HOL Light, and under the assumption of a hardcoded finite action space of tactics, Bansal et al. (2019a) in conjunction with supervised seed data was able to achieve up to 70% proof success rate on the HOList theorem proving task. Similarly, in a set-up much closer to ours, MM GPT-f demonstrated the feasibility of expert iteration when using generative language models for theorem proving.

Within a fixed corpus of theorems (and hence proof terms), however, both PACT and RL are fundamentally constrained by a lack of exploration—as the performance of the theorem proving agent improves, it will eventually saturate and become starved for data, and its curriculum will need to be expanded. Although self-supervised methods such as PACT represent a way to significantly improve the data-efficiency of reinforcement learning loops over existing theorem prover libraries, the development of continuously self-improving and infinitely scalable neural theorem provers remains contingent on sufficiently powerful exploration and automated curriculum generation; we consider these challenges to be of paramount importance.

Acknowledgments

We thank the members of the Lean community, in particular Kevin Buzzard, Simon Hudon, Johan Commelin, Mario Carneiro, Bhavik Mehta, and Gabriel Ebner for their valuable feedback on our work. We are indebted to Markus Rabe and Christian Szegedy for many hours of helpful discussion. We also thank Daniel Selsam, Tom Hales, and Josef Urban for feedback on earlier drafts of this paper.

Reproducibility statement

The source code used to generate the Lean datasets and run the evaluation is open source and made available in the following repositories:

: https://github.com/jesse-michael-han/lean-tpe-public

: https://github.com/jasonrute/lean_proof_recording

: https://github.com/jesse-michael-han/lean-step-public

Our Transformer model was pre-trained on two proprietary datasets. The first is the same mix used by GPT-3 (Brown et al., 2020) and the second is WebMath (Polu & Sutskever, 2020). More details can be found in Appendix B.

While our weights and the API through which we query our models are not currently public, techniques for training decoder-only transformers and efficiently performing inference with them are well-known. Our released theorem proving code is agnostic to these implementation details and will work with any language model exposed via an HTTP server. The provided code also supports querying a locally hosted Transformer from the open-source library fairseq via the Fairseq CLI (Ott et al., 2019).

We have released a simplified version of the proof search described in Section 3.3 as a tactic to the Lean community in a public beta, opening the way for our models to directly accelerate the development of formalized mathematics and for human experts to provide feedback and additional training signal in a virtuous cycle. The tactic and code are available at https://github.com/jesse-michael-han/lean-gptf, and users who sign up for the beta are granted access to our Transformer model through an API.

References

Appendix A Additional Background

Tactics

Tactics in Lean are metaprograms Ebner et al. (2017), which can construct Lean expressions, such as terms. A tactic state which tracks the list of open goals and other metadata is threaded through each tactic invocation. Lean has special support for treating tactics as an extensible domain-specific language (DSL); this DSL is how Lean is typically used as an interactive theorem prover. The DSL amounts to a linear chain of comma-separated invocations. The process of interactive proving is mediated through Lean’s language server, which will present the context and type for the current goal in the proof to the user, dependent on where their cursor is in the source text. The tactic prediction task is to predict the next tactic given this goal state. We extract supervised training data for this task by extracting all human-supplied proof steps from Lean’s mathlib.

An object called the tactic state is threaded through each invocation of a tactic. Among other things, the tactic state maintains a context of metavariables: placeholders in to which expressions will be substituted later. At each point in the proof, one or more of these metavariables are selected as the goal of the tactic state which is present As the proof progresses, there are multiple values to be found

Example

Consider this (modified) example of a tactic proof from the library.

Each tactic line modifies the proof state, which we explicitly annotate below with comments between each tactic.

Our proofstep objective is to predict the tactic applied to a given tactic state.

Lean stores this proof internally as a proof term:

Since this proof term is just stored internally as a tree, any branch of this term tree can be removed, to create a hole _, for example:

Lean will automatically provide a list of both the local context and the type of a term needed to fill that hole as shown below. Notice this is the same as a tactic state we saw from the term proof above.

Using this methodology of following proof term trees, we can mine low level proof data for every node of a term proof to produce the PACT dataset described in Section 3.2.

Appendix B Datasets

We pre-train on WebMath as described in (Polu & Sutskever, 2020). All models, including the WebMath pre-trained models, and the non-WebMath models used in ablations, were first pre-trained on the mix used by GPT-3 (Brown et al., 2020) which includes a filtered CommonCrawl, WebText2, Book1, Book2 and Wikipedia. WebMath includes Python-only GitHub data, as well as arXiv and Math StackExchange.

From these datasets, a potential risk for test-set contamination (presence of mathlib) exists for the crawled datasets, namely CommonCrawl, WebText2, and (in case of a filtering bug) Python-only GitHub. The other datasets (in particular arXiv and Math StackExchange) may contain short references of mathlib code but in shape and forms that would not lead to effective contamination.

To assess the contamination risk related with the crawled datasets, we searched CommonCrawl, WebText2, arXiv, Python-only GitHub, and Math StackExchange for test theorems. For example, given the test theorem nat.div_eq_sub_div we searched for any occurrences of the string div_eq_sub_div. Of over 3000 test theorem names, we found 595 which occurred in the datasets. Many instances were innocuous, but some were in Lean files, and in some cases there was a proof of a test theorem. There were also 160 additional test theorems with no underscore in their name, which we did not check, but whose name is likely to be found in the datasets. (There is no need to check for training theorems since they are already in the training data and it would not constitute contamination.) We re-calculated the pass-rates of the results in Figure 2 omitting these 755 test theorems. This decreases the reported pass-rates slightly, ranging from $0.6$ to $1.1$ percentage points. The adjusted pass-rate of our best model WebMath > mix1 + mix2 + tactic is $47.4\%,$ a decrease of $1$ percentage point. Our main results still hold even with the adjusted pass-rates.

Additionally we also look at the results for the 1,350 test theorems in our dataset that were added to Lean and mathlib after April 18, 2020, which is after CommonCrawl and WebText2 were gathered, and the 544 test theorems added after September 11, 2020, which is after WebMath was gathered. Unlike future-mathlib, these theorems were part of the originally extracted data. The pass-rates for the WebMath > mix1 + mix2 + tactic model on these restricted sets of test theorems are $45.6\%$ and $43.3\%$, respectively.

We also looked for the following Metamath specific and HOL specific strings in CommonCrawl, WebText2, and Python-only GitHub:

We found occurrence of the Metamath-related strings but interestingly found a non-negligible amount of HOL-related documents, which does not constitute a test-set contamination but potentially benefits the downstream tasks studied in this paper.

While our results show a significant benefit to pre-training on WebMath, it is unclear exactly how pre-training helps. Since Lean’s theorem names are made of coded mathematical phases, e.g. affine.simplex.dist_circumcenter_eq_circumradius, it is not unreasonable to suspect that important statistical connections are extracted from math sources. It is even possible that simple instances of auto-formalization or ITP translation are happening. There is prior work (Gauthier & Kaliszyk, 2015; Wang et al., 2018; 2020) suggesting that both of these are possible. From the point of view of a lean-gptf end-user, any such extraction of prior, publicly available data is useful and helpful. Nonetheless, our results are of a different nature than other AI for theorem proving research which do not use data outside of a given theorem proving library. This should be taken into account in any future comparisons and benchmarks.

B.2 Dataset sizes

B.3 Example datapoints

We present datapoints extracted from a toy example, namely the proof of the Peirce identity, viz.

From this, we can extract four tactic datapoints (i.e. human-generated tactic proof steps):

In contrast, we can extract dozens of raw PACT datapoints. Due to space constraints, we list a representative sample of four such datapoints, from each of which we can derive the nine self-supervised auxiliary PACT tasks studied in our present work. For example, proof term prediction is precisely predicting the "proof_term" given the concatenation of "hyps", "⊢", and the "goal", skip-proof is predicting the "proof_term" given "result", etc.

Appendix C Experiments

Individual Lean tactics are chained together with commas. However, the Lean interactive tactic DSL also includes a number of other tactic combinators for creating composite tactics. A frequently used combinator is the infix semicolon t; s which will perform the tactic t and then apply the tactic s to each of the resulting subgoals produced by t. Our data pipeline for human tactic proof steps treats these semicolon-chained tactics as a single string for the language modeling objective. Thus, our models learn to occasionally emit multiple-step tactic predictions using semicolons. For example, wm-to-tt-m1-m2 solved the following lemma in category theory with a single prediction chaining four tactics in a row:

One way of measuring the sophistication of predicted tactics is to consider the number of successful proofs on the evaluation set which have this composite form using semicolon-chaining. We display this analysis in Table 1, which shows that training with PACT in addition to the human-made tactics causes longer semicolon-chained tactics to be successfully predicted during theorem proving. This is remarkable because the semicolon idiom is specific to the tactic DSL which does not occur in the PACT data whatsoever, and yet the co-training causes longer and more frequent successful composite tactic predictions.

C.2 Theorem naming case study

We included theorem naming as part of the PACT task suite. By mathlib convention, theorem names are essentially snake-cased, natural language summaries of the type signature of a theorem, and so the theorem naming task is analogous to a formal-to-informal translation task. We evaluate the ability of our best model (in terms of theorem proving success rate) wm-to-tt-m1-m2 on its ability to guess theorem names on the completely unseen future-mathlib set of theorems. The distribution shift inherent in the future-mathlib dataset particularly impacts the theorem naming task, because many of the ground-truth names will involve names for concepts that were only defined in mathlib after we extracted our training data.

On the $\approx$$2.8$K future-mathlib theorems, we queried wm-to-tt-m1-m2 for up to $N=16$ candidates. We order these candidates into a list xs by decreasing cumulative log-probability and calculate the top-$K$ accuracy by checking if any of the first $K$ candidates of xs match the ground truth exactly. The model wm-to-tt-m1-m2 was able to achieve 20.1% top-1 accuracy, 21.1% top-3 accuracy, 26.7% top-10 accuracy, and 30.0% top-16 accuracy. We display a sample of correct top-1 guesses (Figure 5) and a sample of failed guesses in (Figure 6). We note that the failed guesses, while containing no syntactic matches, are both semantically reasonable and syntactically very similar to the ground truth.

C.3 Test set evaluation breakdown by module

Lean’s mathlib is organized into top-level modules, which roughly organize theorems into mathematical subject area. In Figure 7, we break down the evaluation results on our test set between our PACT-trained models wm-to-tt-m1-m2 and wm-to-tt-m1 and our baselines wm-to-tt and tidy. We see that full PACT mostly dominates over co-training on just the mix1 tasks over all subject areas, and that wm-to-tt-m1 dominates the model wm-to-tt trained on human tactic proof steps only.

C.4 Baseline description

The tidy backend is determined by a constant oracle

which always returns the same list of tactics, namely:

Unlike the gptf backend, which generates a list of candidates in parallel independently, tidy enjoys the advantage that the list of tactics it emits is carefully chosen and ordered in order to optimize the proof search—this is based on the “waterfall” technique of the human-style automated theorem prover described in (Ganesalingam & Gowers, 2017).

C.5 Computational resource estimates

For each evaluation loop over the test set, we distributed the theorems over a pool of $32$ CPU workers whose inference requests were load-balanced over $4$ V100 GPUs. Each evaluation required $\approx$$10$ hours with $\approx$$30\%$ GPU utilization. We observed that our evaluation was bottlenecked by inference and in practice, we hosted up to three evaluation loops at once on a VM with 80 logical cores without achieving full CPU utilization. In addition to the wall-clock timeout of 600s, we also limited the proof search to a logical timeout of 512 iterations, where one iteration corresponds to a single expansion of a node of the BFS search tree. In practice, so much time was spent either blocked on inference or performing the tactic executions in the inner loop of each iteration that we rarely exceeded the logical timeout, usually exceeding the wall-clock timeout instead.

Appendix D Example proofs

Lean’s mathlib is one of the most active open-source software projects in the world. More than one-third of the proofs found by our models are shorter and produce smaller proof terms than the ground truth, leading to dozens of GPT-f co-authored commits to mathlib. We examine some of the proofs found by our models in more detail.

This proof produces a proof term which is 4X smaller than the original:

The original, human-written proof is much longer, viz.

D.2 primrec.of_equiv

This proof produces a proof term which is 12X smaller than the original:

The author of the original proof and maintainer of that package commented:

encode_iff.1 primrec.encode is clever, it’s a way to translate primrec across an equivalence when the encode function is defined as encode x = encode (e x) where e is the isomorphism.

As far as they knew, this trick was never used before in the computability package.

D.3 real.tan_eq_sin_div_cos

This proof demonstrates our model’s library knowledge and ability at premise selection.

D.4 sym2.is_diag_iff_proj_eq

The proof of this lemma is longer than the ground truth and was not contributed to mathlib, but we describe it here because the proof is original and includes a nontrivial instantiation of an existential quantifier.

This goal state never appeared in mathlib.

D.5 norm_le_zero_iff

The following proof is remarkable because it uses fewer tactic steps and takes a different route to the proof than the ground truth, uses a complex idiom simpa [...] using @..., and was predicted in one shot.

The lemmas supplied between the square brackets are used to simplify the main goal. The lemma supplied after the keyword using can further simplify the lemmas supplied between the square brackets. The @ modifier makes all arguments explicit. The string @norm_eq_zero never appeared in our training data but the prediction includes the correct number of correctly typed arguments, and even replaces the second argument with a placeholder _, correctly guessing that it can be inferred by the elaborator. Finally, this again showcases the strength of our models as premise selectors: all three lemmas le_antisymm_iff, norm_nonneg, and norm_eq_zero were not used in the human-supplied proof but are necessary for this proof.

Moving forward, we hope that our neural theorem provers will continue to find ways to improve mathlib and assist in creating new proofs. More generally, we hope neural theorem proving will one day be become a routine part of the formalization workflow.