Proof Artifact Co-Training for Theorem Proving with Language Models

1 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. , 2020 ) , and joint text-image modeling ( Radford et al . ) , 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 . 2 BACKGROUND AND RELATED WORK . LEAN 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. , 2019 ) , 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 . simpa [ le_antisymm_iff , norm_nonneg ] using @ norm_eq_zero α _ 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. , 2018 ) 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. , 2019c ; 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 : R ) , 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 the appendix . ) 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. , 2020 ) , Boolean satisfiability ( Finkbeiner et al. , 2020 ) , 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. , 2020 ; Wu et al. , 2021b ) . Our work both builds on and departs from these ideas in several ways . Unlike skip-tree training ( Rabe et al. , 2020 ) , 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. , 2018 ; 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 ) . 3 THE LEANSTEP DATASETS AND MACHINE LEARNING ENVIRONMENT . 3.1 HUMAN TACTIC PROOF STEPS . 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 the appendix for examples . 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 .

Summary: This paper addresses the general setup of using transformer models for interactive theorem proving (ITP) tasks. The ITP engine considered here is Lean. The contribution of the paper is a data augmentation method. This is achieved by mining low level artifacts from a given dataset of Lean proofs. This includes data extracted from additional type information inferred by Lean while processing the given proofs. Several prediction tasks are formed for this additional data, to allow pre-training and co-training in combination with the existing WebMath dataset. Results show that the additional datasets extracted in this way aid substantively when used for pretraining and cotraining.