Formal Theorem Proving by Rewarding LLMs to Decompose Proofs Hierarchically

Authors: kefan  dong
arvind  mahankali
tengyu  ma
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-26 04:08:56

Abstract: Mathematical theorem proving is an important testbed for large language models’ deep and abstract reasoning capability. This paper focuses on improving LLMs’ ability to write proofs in formal languages that permit automated proof verification/evaluation. Most previous results provide human-written lemmas to the theorem prover, which is an arguably oversimplified setting that does not sufficiently test the provers' planning and decomposition capabilities. Instead, we work in a more natural setup where the lemmas that are directly relevant to the theorem are not given to the theorem prover at test time. We design an RL-based training algorithm that encourages the model to decompose a theorem into lemmas, prove the lemmas, and then prove the theorem by using the lemmas. Our reward mechanism is inspired by how mathematicians train themselves: even if a theorem is too challenging to be proved by the current model, a positive reward is still given to the model for any correct and novel lemmas that are proposed and proved in this process. During training, our model proposes and proves lemmas that are not in the training dataset. In fact, these newly proposed correct lemmas consist of 37.7\% of the training replay buffer when we train on the dataset extracted from Archive of Formal Proofs (AFP). The model trained by our RL algorithm outperforms that trained by supervised finetuning, improving the pass rate from 40.8\% to 45.5\% on AFP test set, and from 36.5\% to 39.5\% on an out-of-distribution test set.

Main paper: 
 introduction: 
 The reasoning abilities of large language models (LLMs) are a significant marker of artificial intelligence and critical for complex and safety-sensitive applications. Yet recent studies highlight the limited performance of LLMs on reasoning tasks (e.g., Mündler et al. [2023] , Valmeekam et al. [2023] and references therein). Automated theorem proving by LLMs is an excellent reasoning task that abstracts away the need for numerical manipulation or tool use (e.g., using a calculator) and allows for precise correctness evaluation with an automatic verifier (such as Isabelle [Nipkow et al., 2002] and Lean [De Moura et al., 2015] ), even without ground truth. Thanks to tools such as Sledgehammer [Paulsson and Blanchette, 2012] that can automatically complete low-level details, the granularity of formal proofs is similar to natural language proofs (see Fig. 1 Left for an illustrative example). Note that verifying a proof is fundamentally much easier than generating the proof.1 Thus, learning to prove theorems from verifiers' supervision is reminiscent of weak-to-strong generalization [Burns et al., 2023] .Previous results in this area largely focus on setting where the theorem prover can use all the lemmas in the formal proof library, including those particularly written to decompose a specific theorem's Figure 1 : Left: in the dashed callout block, we show an example of an Isabelle proof and its explanation in natural language. Right: an example of a proof tree. The two child nodes correspond to the two new lemmas proposed in the proof of the root node.proof [Jiang et al., 2021, Polu and Sutskever, 2020] This setting arguably oversimplifies the problem and doesn't sufficiently test the models' planning and decomposition capabilities, and it is unclear whether the resulting models can be used to prove new theorems from scratch when such lemmas are not available at test time. Instead, we work in a more natural setup where the theorem prover needs to propose and prove lemmas to decompose the proof hierarchically itself (see Section 2 for more details). In Section C.2, we demonstrate that this task is indeed much more challenging.Figure 2 : Illustration of our algorithm Proof Tree Generator (ProD)-RL. In step 2b, locally correct means that the statement is proved correctly using the proposed lemmas, and globally correct means that all the proposed lemmas are also proved correctly. As an important feature of our algorithm, even if Theorem 1 is not proved by the model because the proof to Lemmas 1 is incorrect, we still train on the correct lemma (Lemma 2) by setting its reward r = 1.In addition, most existing proof-generation algorithms leverage the formal verifier by (a) providing the verifier's current proof state to the LLMs step-by-step, and (b) using best-first search algorithms such as A ⋆ to build a multi-step proof from many LLM-generated steps [Jiang et al., 2022a , Han et al., 2021] . The major challenge of these method is the high computation cost incurred in both (a) and (b) because (a) requires re-running LLMs on a different context that contains a long verifier's proof state at every step, and (b) the search is expensive. Consequently, the best method along this line of research requires more than 1k GPU days with A100s to train a model with 600M parameters [Lample et al., 2022] , whereas our method only takes less than 36 GPU days to train a 7B model.To address these issues, we design a method, called Proof Decomposer (ProD), that uses LLMs to propose and prove new lemmas hierarchically and generate whole proofs directly without searching. We augment the vanilla formal proof syntax so that the model can propose new lemmas by stating their statements during the proof, and then prove the lemmas separately. Hence, a complete proof of a theorem will form a tree structure where the child nodes are the lemmas proposed in the proof of the parent node (Fig. 1 Right), and the theorem is proved only if all the proofs in the tree are correct.We train our LLMs with reinforcement learning (RL) in a way that somewhat imitates the mathematician's process: we reward correct partial proofs (i.e., proof sub-trees) even if the original theorem (i.e., the root node) is not proved entirely. Since our model can generate and prove novel lemmas during training, it could still make progress even if the model is given a very challenging theorem. This is similar to mathematicians publishing papers on interesting lemmas that do not solve the original goal. We illustrate our algorithm in Fig. 2 , and defer the details to Section B.2. We test our model ProD-RL by generating proof trees on holdout theorems that the model is never trained on, and we show that our model ProD-RL outperforms several other baselines. Compared with the supervised fine-tuned (SFT) model on the same training set, our model improves the pass rate from 40.8% to 47.5% on the holdout test set, whereas vanilla reinforcement learning without lemma proposals during training does not improve the corresponding SFT model (see Section 3.2). This is partly because our method encourages the model to propose and prove additional lemmasin fact, 37.7% of the proved lemmas during training are not in the dataset. As a result, the model still improves even if it's already fine-tuned on the same dataset with human-written ground-truth proofs. 
setup: 
 Conditional proofs. We use the term conditional proof to denote a proof that, in addition to the standard formal proof syntax, can propose new lemmas by enclosing their statements by <invoke> and </invoke> tokens (examples shown in the blue boxes of Fig. 1 ). A conditional proof has the following format:t 1 <invoke> l 1 </invoke> t 2 <invoke> l 2 </invoke> • • • t k <invoke> l k </invoke> t k+1 where t 1 , • • • , t k+1 denoteproof segments in the original formal proof syntax (see e.g., Fig. 1 , proof texts in black), and l 1 , • • • , l t denote proposed lemma statements (see e.g. Fig. 1 , proof texts in red).2 Proof tree nodes. With the proposed lemmas, a complete proof forms a tree structure (as shown in Fig. 1 ). A node in a proof tree is a tuple of premises, context, a theorem statement, and a conditional proof. Premises represent the lemmas that are treated as common knowledge, which are typically not directly relevant to the proof. We allow the model to directly use them in the proof so that it does not have to repetitively prove all the fundamental facts, such as properties of continuous functions and natural numbers. Context represents the necessary contents to prepare the theorem statement, such as the definition of specific objects/functions. We use the context as part of the prompt for the LLMs to generate proofs, and to prepare the proof verifier to check the generated proofs.Correctness of conditional proofs and proof trees. A proof tree node n with conditional proof t 1 <invoke> l 1 </invoke> t 2 <invoke> l 2 </invoke> • • • t k <invoke> l k </invoke> t k+1 is locally correct if, after adding l 1 , . . . , l k to the set of premises, t 1 . . . t k+1 , is a proof to the statement of n acceptable by the formal verifier under the context of n.We consider a proof tree valid if, for every node, each of its child nodes corresponds to one proposed lemma and shares the same premises and context with its parent node. A tree node n is globally correct with respect to a given set of tree nodes N if we can construct a valid proof tree with root n using the locally correct tree nodes in N . Intuitively, global correctness corresponds to the standard notion of correctness (i.e., whether the theorem is proved), and local correctness is a weaker concept, referring to the correctness of conditional proofs assuming the proposed lemmas.Dataset construction. We construct the datasets by parsing raw proof-library files into examples of the form (premises, context, statement, conditional proof). For any theorem s, we construct an example where the premises are all the theorems from predecessor files. To compute the context, we iteratively remove all the theorem statement and its proof from the proof-library files if the theorem is not referred to in the remaining file contents (in other words, in the first iteration we peel off the root nodes of the proof trees from the file, and then the nodes in the next level, etc.). Finally, to augment the original proof of theorem s to a conditional proof with lemma proposal, for every proof step t j that calls lemmas l j,1 , . . . , l j,nj (and these lemmas were removed from the context), we insert the statements of these lemmas enclosed by the <invoke> and </invoke> tokens into the proof right before t j . We defer additional details of our dataset construction process to Appendix A.We split the training and test set (AFP test) based on the dependency of the files in the proof library so that the examples in the training set never refer to any files in the test set. We also construct an additional test set AFP 2023 by parsing the AFP files submitted after the knowledge cutoff date of the Llemma model (April 2023) to eliminate potential data leakage issues.Compared with prior works [Jiang et al., 2021 , First et al., 2023] , the two major differences in our setup are the availability of lemmas from the same file and the training/test split. In Section C.2, we discuss and test their effects in detail. 
reinforcement learning algorithm: 
 Due to limited space, we present a sketch of our RL algorithm here and defer the details to Appendix B.Proof tree generation. To generate proof trees using an autoregressive model π θ , we need to first fine-tune the model to follow a specific format:(a) the input x to the model π θ is the concatenation of a context and a theorem statement, and (b) the expected output y of the model is a special token t 0 followed by a conditional proof, where t 0 is either <use_invoke> or <no_invoke>, denoting whether the following conditional proof should propose new lemmas.We let the model generate the special token t 0 before a conditional proof so that we can upscale the probability of the <use_invoke> token during RL to propose more lemmas for better exploration. Using the fine-tuned model, we can generate conditional proofs and complete the proof tree recursively, as shown in Alg. 1.Reinforcement learning. Our method is illustrated in Fig. 2 . We start with a supervised fine-tuned model so that it can generate conditional proofs in the desired format. At every round, we sample a batch of examples from the training dataset and generate proof trees using the current model. Then we call the proof verifier to test the generated proofs, and assign rewards based on the global correctness of the conditional proofs. Finally, we train the model using REINFORCE with a replay buffer.We update the model using partial proofs (i.e., proof sub-trees) even if the original theorem from the dataset (i.e, the root of the proof tree) is not proved. Hence, our method can also be viewed as an instantiation of hindsight experience replay [Andrychowicz et al., 2017] , where the hindsight trajectories are correct proof sub-trees.This section reports the models' pass@k performance on AFP test and AFP 2023 datasets. For a given dataset, pass@k measures the percentage of the theorems proved by at least one of k proofs generated by the model. Recall that the theorems in the test set are selected based on the dependencies of the AFP files and are not used in any proofs from the training set (see Section C.2 for more details).In other words, we do not train the model on test datasets using reinforcement learning. Instead, we test whether ProD-RL is a fundamentally better model when tested on new theorems. As a baseline method, we train a model on a variant of the SFT dataset where all lemmas are kept in the context. It can be seen as a reproduction of First et al. [2023] with a slightly different way to obtain the context - First et al. [2023] includes all the file content before the statement of the theorem, whereas we only keep the statement of previous lemmas. We also run reinforcement learning on the same RL dataset as our method (see Section E.2 for more details). Table 1 shows the performance of our model on the AFP test sets. For a fair comparison, the baseline models are tested in our new setup without human-written lemmas. Note that the SFT model without lemma proposal outperforms the SFT model with lemma proposal. We hypothesize that it is because proposing correct lemmas itself is challenging, which distracts the model from learning to generate direct proofs. However, RL with lemma proposal improves the SFT model and outperforms others because the model proposes and proves additional lemmas that are not in the training dataset, whereas RL without lemma proposal yields no improvement.In Fig. 3 (Left), we plot the pass rates with different numbers of samples per theorem on both AFP test and AFP 2023. Fig. 3 shows that on AFP test, the ProD-RL model significantly improves upon baseline methods as well as the ProD-SFT. However, on AFP 2023, the improvement is minor over SFT w/o lemma proposal, while ProD-RL still outperforms ProD-SFT. The results suggest that the baseline methods are more robust to heavier distribution shifts, while our method has a larger improvement when the test distribution is closer to the training distribution.In Fig. 3 (Right), we decompose the proved theorems by the depth of their ground-truth proofs (shown on the x-axis) and the depth of generated proof trees (indicated by color). When there are multiple correct proof trees, we plot the one with the maximum depth. As a comparison, we also plot the success rates of the proofs generated by the RL model trained w/o lemma proposal. Fig. 3 (Right) shows that the improvement of ProD-RL mostly comes from proving theorems with low-to-medium difficulty where the depth of the ground-truth proof is at most 2. For more complex theorems, both models' pass rates are low and the improvement of our method is not significant, meaning that they are currently beyond the models' capability.Due to limited space, we defer the case study for the generated lemmas to Appendix C.4. 
conclusion: 
 In this paper, we design a reinforcement learning algorithm that encourages LLMs to write formal proofs by decomposing them hierarchically. We also design a more natural testing setup by removing the directly relevant lemmas from the context. We show that, by proposing and proving new lemmas that are not present in the training dataset, the resulting model ProD-RL outperforms or achieves comparable performance to baseline methods trained on the same dataset. 
 
 Appendix: In this section, we present the dataset construction process in detail.Recall that we construct the datasets by parsing raw proof-library files into examples of the form (premises, context, statement, conditional proof). Hence, we first segment each of the files into blocks c 1 s 1 p 1 • • • c l s l p l where the s i are theorem statements, the p i are the corresponding proofs, and the c i are the file contents between proofs, such as object definitions and local assumptions. Next, we build proof trees from the segmented file by removing (s i , p i ) pairs from the file if the theorem s i is not referred to in the remaining file contents (in other words, in the first iteration we peel off the root nodes of the proof trees from the file, and then the nodes in the next level, etc.). Note that some theorems cannot be peeled off by this process because they are referred to in some file content c j (e.g., lemmas used to instantiate local objects). We use T tree to denote the subset of theorems peeled off during the process.For every theorem s i , we construct an example where the context is the concatenation of {c j : j < i} and {s j : j < i, s j ̸ ∈ T tree } in the order they appear in the file. That is, we exclude the lemmas peeled off after we peel off s i -these are the lemmas that directly contributes to the proof of s i .To construct the conditional proof of theorem s i , we add the proposed lemma statements to the original proof p i . In particular, we split the proof p i into steps t 1 , • • • , t k using the formal language parser. Then for every step t j that uses lemmas l j,1 , . . . , l j,nj from T tree , we insert the statements of these lemmas enclosed by the <invoke> and </invoke> tokens, denoted by ζ j = <invoke> l j,1 </invoke> • • • <invoke> l j,nj </invoke>, into the proof right before t j . In other words, the conditional proof is the concatenation Jiang et al. [2022a] , we use Sledgehammer, a premise selection tool that automatically searches for proofs to the current goal, to replace proof steps that are originally generated by it (see Section E.3 for more details) so that the mode can focus less on the tedious low-level details.ζ 1 t 1 • • • ζ k t k . Similar toThe premises are all the theorems from predecessor files, which are typically not directly relevant to the theorem (otherwise they will be stated in the same file). Theorems in the premises set can be used directly in the proof, or they can be selected by Sledgehammer to search for proof steps. In our implementation, the premises are implicitly defined by the dependency graphs of the files.For every example in the training dataset, if the conditional proof proposes at least one lemma, we also add one augmented example by moving the proposed lemmas from the conditional proof to the context -this augmented example does not propose new lemmas and is always locally correct.In this section, we first describe how to use LLMs to generate proof trees, and then introduce our reinforcement learning method (ProD-RL) that rewards the model to decompose proofs hierarchically.To generate proof trees using an autoregressive model π θ , we need to first fine-tune the model to follow a specific format:(a) the input x to the model π θ is the concatenation of a context and a theorem statement, and (b) the expected output y of the model is a special token t 0 followed by a conditional proof, where t 0 is either <use_invoke> or <no_invoke>, denoting whether the following conditional proof should propose new lemmas.We let the model generate the special token t 0 before a conditional proof so that we can upscale the probability of the <use_invoke> token during reinforcement learning to force the model to propose more lemmas for better exploration.We summarize our proof-tree generation algorithm in Alg. 1. Given a theorem statement s and the corresponding context c, we first sample π θ autoregressively starting with the prompt x = c s, and ideally the model will output a special token t 0 followed by a conditional proof ρ (Line 3). Next, we parse the conditional proof ρ and collect the proposed lemmas l 1 , • • • , l k for the next round of generation (Line 5). We will force the model to generate conditional proofs without proposing new lemmas at a certain depth so that the proof tree doesn't grow indefinitely, which can be implemented easily by replacing the prompt with x ′ = c s <no_invoke> (Line 6).Algorithm 1 Generate proof trees (test time)1: Inputs: Model π θ , a set of contexts and statementsG 0 = {(c i , s i )} i , maximum depth d. 2: for ι ← 0, 1, • • • , d -1 do 3: Sample proofs (t 0,i ρ i ) ∼ π θ (• | c i s i ) for (c i , s i ) in G ι ,where t 0,i is the token represents whether the proof should use invoke and ρ i is the conditional proof.4:P ι ← {(c i , s i , ρ i ) | ∀(c i , s i ) ∈ G ι }. 5:Collect proposed lemmas (a condition proof ρ i might propose more than one lemma l j ):G ι+1 ← {(c i , l j ) | (c i , s i , ρ i ) ∈ P ι and l j is proposed in ρ i }. 6: Sample proofs ρ i ∼ π θ (• | c i s i <no_invoke>) for (c i , s i ) in G d . (▷) Truncate at depth d. 7: P d ← {(c i , s i , ρ i ) | ∀(c i , s i ) ∈ G d }. 8: Return ∪ d ι=0 P ι .Our reinforcement learning method is illustrated in Fig. 2 . We start with a supervised fine-tuned model so that it can generate conditional proofs in the desired format. Then at every round, we randomly sample a batch of examples from the training dataset and perform the following steps.Step 1: Generate proofs. To help exploration, we generate proof trees using a modified version of Alg. 1 (shown in Alg. 2 of Appendix E.1) with the following differences.(a) For the theorems where the probability π θ (<use_invoke> | x) is among the top 25% in the batch, we will force the model to generate conditional proofs with t 0 = <use_invoke>. Otherwise, we sample t 0 according to the probability of π θ (• | x). (b) For every theorem where the model generates a conditional proof with new lemmas, we also let the model generate another conditional proof without proposing lemmas. If any of these two conditional proofs is globally correct, it can be used to construct proof trees for other theorems.We also mix the ground-truth conditional proofs with proofs generated during RL by generating additional proof trees starting with the ground-truth conditional proof using the current model π θ . In other words, the root node of a proof tree constructed in this manner is a ground-truth conditional proof, but its descendants are generated by the model.Step 2: Determine the reward of an example. In this step, we first check the local correctness of each conditional proof using the formal verifier (Step 2a in Fig. 2 ).In addition to the verifiers' output, we apply two filters to help train the model: (a) we filter out trivial lemma proposals -if a proposed lemma can directly imply the theorem (e.g., if the proposed lemma has exactly the same statement as the theorem), we will simply discard this example, and (b) we remove unnecessary lemma proposals -if the conditional proof is still correct after removing all the references to a proposed lemma, we will delete this lemma from the conditional proof.Then we compute the global correctness of the generated proofs using its definition. Finally, we assign a binary reward r(c, s, ρ) to each tree node with context c, statement s, and conditional proof ρ based on its global correctness (Step 2b in Fig. 2 ).Step 3: Update the model by REINFORCE. In this step, we will first construct a training dataset consisting of examples with format (prompt, target, weight) from the conditional proofs collected inStep 1 and then update the model π θ using the weighted cross-entropy loss.For each generated conditional proof, we add one example to the training dataset where the prompt is the context concatenated with the theorem statement, and the target is the conditional proof prepended by the <use_invoke> or <no_invoke> token. Note that the reward of a conditional proof depends not only on the correctness of the conditional proof, but also on the correctness of the proposed lemmas.To reduce the variance of our gradient updates, we train a value function V ϕ that predicts the expected reward of the current policy on a given proof tree node (i.e., V ϕ ≈ E ρ∼π(•|c,s) [r(c, s, ρ)]). The weight of an example is the product of the value function outputs on invoked lemmas multiplied with a length penalization to prefer shorter proofs -for proof tree node with conditional proof length h and proposed lemmas l 1 , • • • , l k , the weight w of this example is w = γ h k i=1 V ϕ (l i ) with discount factor γ ∈ (0, 1), or w = 0 if the proof tree node is not locally correct.For easier implementation, we train the value function to predict two special tokens, <true> and <false>, conditioned on the context and theorem statement. Let p t be the probability of the <true> token conditioned the context and theorem statement, and p f the probability of the <false> token. The output of the value function is then p t /(p t + p f ).Similar to the SFT dataset, we add one augmented example by moving the proposed lemma from the conditional proof to the context for any locally correct conditional proof with new lemmas. In addition, we will use a replay buffer to stabilize the training.Remarks. We will update the model using partial proofs (i.e., proof sub-trees) even if the original theorem from the dataset (i.e, the root of the proof tree) is not proved. Hence, our method can also be viewed as an instantiation of hindsight experience replay [Andrychowicz et al., 2017] , where the hindsight trajectories are correct proof sub-trees.This section presents our experimental results. We will first list additional experiment details (Section C.1) and then compare our setup with prior works (Section C.2.Proof verification software. We use Isabelle [Nipkow et al., 2002] as our proof verification software since the proofs are declarative and human-readable without knowing the verifier's proof state, and we use PISA (Portal to ISAbelle, Jiang et al. [2021] ) to interact with Isabelle. To check whether a proof tree node is locally correct, we import all the theorems from its premises, move each of the proposed lemmas from the conditional proof to the context, and then add a fake proof indicated by the keyword 'sorry' to every lemma statement in the context (In Isabelle, 'sorry' will register the statement as a fact even without any actual proof.) The remaining proof steps will follow the original Isabelle syntax, and we can check its correctness directly. We set a 10s timeout for each proof step.Datasets. Our SFT dataset is the combination of theorems from Archive of Formal Proof3 (AFP, retrieved on 2022-12-06) and Isabelle built-in files (such as HOL which contains the theorems that define natural numbers, etc.). The resulting dataset contains 312k examples.To construct the test set, we parse the AFP files submitted after the knowledge cutoff date of our pretrained model (April 2023) to eliminate possible data leakages. 4 The test dataset is handled similarly, except that we only keep the theorems in T tree . This AFP 2023 test set contains 2k theorems.Testing setup. To measure the performance of the models, we sample k proof trees per theorem independently on the test set and report the pass@k performance (that is, a theorem is proved if at least one of the conditional proofs is globally correct with respect to all the generated tree nodes). When generating proofs, we use temperature 0.7 and truncate the context to only include the last 1k tokens. The proof trees are truncated at depth 2.Supervised fine-tuning. We start from the Llemma 7b model [Azerbayev et al., 2023] and fine-tune the model for 2 epochs with the standard cross entropy loss. 5 On theorems from AFP, we compute the loss only on the special token and the proof, but not on the context and statement. On Isabelle built-in theorems, we compute the loss on the statement to help the model internalize basic facts. We use the AdamW optimizer [Loshchilov and Hutter, 2018] with linear warmup, constant learning rate 1e-5, macro batch size 128, and context window 2048.Reinforcement learning. The dataset we use for the reinforcement learning stage is T tree , the set of theorems that are iteratively peeled off when parsing AFP files, which contains 109k examples. We initialize the model by supervised fine-tuning on a subset of AFP files, and then run RL for 20 rounds with a batch of 5k random examples per round. We truncate the proof tree at depth 3 and sample with temperature 0.7 during training. We use the same hyperparameters as the SFT stage to update the policy π θ , We initialize the value function V ϕ by a Llemma 7b model fine-tuned on our SFT dataset.In this section, we concretely compare our new setup with prior works [Jiang et al., 2021 , First et al., 2023] . Recall that there are two main differences in how we process our dataset:(a) we split the train/test set based on file dependencies so that no theorems in the test set are referred to in the training set, whereas PISA split theorems randomly, and (b) when testing a proof, we remove certain lemmas from the context.To show that our setup is indeed more challenging, we first construct datasets formatted similarly to those in [First et al., 2023] . Specifically, we parse the AFP files into examples using the method described in Section 2, with the only exception that all human-written lemmas are kept in the context. Then we select a subset of theorems as the test set D w/ l test based on the dependency of the AFP files, so that the examples in the test set are never used by the remaining theorems (see Section E.3 for more details). Then we split the remaining examples randomly into training and validation datasets, denoted by D w/ l train and D w/ l val respectively. We use D w/o l test to denote the test dataset of our setup where the lemmas are removed from the context. The validation dataset D w/ l val mimics prior works' setup [Jiang et al., 2021 , First et al., 2023] , and D w/ l test is an interpolation between prior works' setup and our setup. Table 2 shows the performance of the model supervised fine-tuned on D w/ l train plus all Isabelle built-in theorems. The results suggest that both features of our setup, removing the lemmas and splitting the training/test set by file dependency, increase the difficulty of the task. The effect of sampling temperature. In our preliminary experiments, We tune the sampling temperature using the models trained on the AFP training sets D w/ l train and D w/o l train (that is, training sets constructed with and without helper lemmas, respectively). We test the model on the AFP test set D w/o l test with different temperatures to decide the best choice for testing our models. Fig. 4 shows the performance of the SFT model without lemma proposal using different sampling temperatures. We conclude that the temperature 0.7 is best for testing both models.Results on miniF2F. We observe that the improvement of ProD-RL over SFT w/o lemma proposal is significant only when the test distribution is close to the training distribution. On miniF2F [Zheng et al., 2021] where the theorems are very different from theorems in the training dataset, ProD-RL performs worse than SFT w/o lemma proposal, as shown in Table 3 . We also observe that when tested on miniF2F theorems, our model failed to propose meaningful lemmas. This may be because proving to miniF2F-level mathematics questions typically does not require hierarchical decomposition. Therefore we leave it as future work to extend our methods to other domains such as miniF2F and PutnamBench [Tsoukalas et al., 2024] .In this section, we manually examine the new lemmas proposed during RL and list the typical cases where new lemmas are proposed. Note that many AFP files focus on complex concepts and results in mathematics or computer science, making manual examination challenging. Therefore, examples in this section are biased toward easier theorems. Case 2: The proposed lemma is a rephrase of an existing lemma. We also find that some proposed lemmas are rephrases of existing lemmas in the training dataset. Although in this case, the proposed lemma is not fundamentally useful for proving new theorems, they can be viewed as data augmentation to enhance the models' performance. In the following example, the model produces a lemma equivalent to one in an AFP file. Line 1 shows the original form of the lemma stated in the AFP file, while Lines 2-4 show an equivalent lemma proposed by our model during RL.1lemma icard_mono: "A ⊆ B =⇒ icard A ≤ icard B" 2 lemma icard_mono: 3 assumes "A ⊆ B" 4 shows "icard A ≤ icard B"Case 3: The proposed lemma is novel but not useful to the original proof. We also observe cases where the proposed lemma is novel, but the conditional proof of the theorem is incorrect. In the following example, the proposed lemma states that the shortest path between vertices u, v is a lower bound for the length of any path that connects u, v (in an unweighted and undirected graph):1 lemma shortest_path_lower_bound: This lemma is proposed to prove that the shortest path between the vertex u and itself has length 0 (which is a theorem in the AFP file). However, the conditional proof of the theorem contains a few mistakes while the proposed lemma is proved separately. In this case, we still train on the correct lemma even though it might not be directly useful to the theorem in the training set.Remarks.We observe that the lemmas proposed by the model typically do not involve complex ideas. We attribute this to two main factors: (a) the limited size of our model and formal proof dataset, and (b) the fact that many human-written lemmas in the AFP file are indeed about basic facts and basic properties (which are often used to prove more complex theorems later). Nevertheless, our model still proposes and proves reasonable lemmas that are not present in the training dataset, and our experiments demonstrate that with these proposed lemmas, ProD-RL outperforms ProD-SFT on holdout test sets. We leave it to future work to scale up our method and force the model to focus on more challenging theorems.Most existing methods use search-based algorithms to generate formal proofs with language models [Polu and Sutskever, 2020 , Jiang et al., 2021 , Han et al., 2021 , Polu et al., 2022 , Jiang et al., 2022a] .Prior works also show that RL can improve the search-based models' performance [Polu et al., 2022 , Wu et al., 2022 , Lample et al., 2022] . The major drawback of these methods is their high computation cost at test time. A recent work [First et al., 2023] trains LLMs to generate a whole proof directly without knowing the verifier's state. Our baseline model, SFT without lemma proposal, can be viewed as a reproduction of their method with a slightly different format. Orthogonally, Mikuła et al. [2023] improve premise selection tools (such as sledgehammer) using transformer-based retrieval models. Magnushammer could potentially be combined with our methods, which we leave as future works.Another line of research aims to translate natural language proofs into formal proofs [Jiang et al., 2022b , Zheng et al., 2023] . Xin et al. [2023] build a library of useful lemmas by decomposing natural language proofs into lemmas with an LLM and then formalizing the decomposed proofs. In contrast, we propose new lemmas entirely in formal language.In general, mathematical question-answering tasks (such as GSM8K [Cobbe et al., 2021] and MATH [Hendrycks et al., 2021] ) and theorem-proving tasks (such as Welleck et al. [2021] ) are well-accepted benchmarks for the reasoning capability of large language models. Prior works show that instruction tuning or RL can significantly improve the models' performance [Shao et al., 2024] . However, evaluation on these tasks is either performed by another language model (which is prone to errors) [Lightman et al., 2023] , or requires ground-truth answers that are hard to acquire at scale.In Alg. 2, we present the algorithm for generating proof trees during RL training. Recall that, compared with Alg. 1, there are two major differences:(a) for the theorems where the probability π θ (<use_invoke> | x) is among the top 25% in the batch, we will force the model to generate conditional proofs with t 0 = <use_invoke> (Line 4-6), and (b) for every theorem where the model generates a conditional proof with new lemmas, we also let the model generate another conditional proof without proposing new lemmas (Line 11).In this section, we describe the additional details for training the baseline models using reinforcement learning.Our RL training pipeline for the baseline models is similar to that of ProD-RL, except that the models only generate proofs without lemma proposal. For RL baselines, we use the same dataset and the same hyperparameters as our method. To mix the ground-truth conditional proofs with generated proofs, we convert the conditional proofs to proofs without lemma proposal by moving all the proposed Algorithm 2 Generate proof trees (train) Compute the invoke probability ∀i, p i = π θ (<use_invoke>|ci,si) π θ (<use_invoke>|ci,si)+π θ (<no_invoke>|ci,si) .Let κ be the 75% quantile of {p i } i . Sample proofs ρ i ∼ π θ (• | c i s i <no_invoke>) for lemmas (c i , s i , ρ ⋆ i ) in G.Sample proofs ρ i ∼ π θ (• | c i s i <use_invoke>) for lemmas (c i , s i , ρ ⋆ i ) in G.11:P ι ← {(c i , s i , ρ i ) | ∀i s.t. (c i , s i , ρ ⋆ i ) ∈ G} ∪ {(c i , s i , ρ i ) | ∀i s.t. (c i , s i , ρ ⋆ i ) ∈ G}.12:if ι = 0 then 13:P ι ← P ι ∪ G. (▷)In training, we also complete proof trees for ground-truth proofs.14:Extract proposed lemmas (note that a condition proof ρ might propose more than one lemma l): G ← {(c, l, Null) | (c, s, ρ) ∈ P ι and l is proposed in ρ}. 15: Return ∪ d ι=0 P ι .lemma in the conditional proof to the context, and the resulting proofs are always acceptable by the verifier.Details of using sledgehammer in the proof. Sledgehammer is a premise selection tool that can automatically generate proofs to solve the current goal. Although sledgehammer is not always applicable, Jiang et al. [2022a] shows that letting the model to call sledgehammer whenever it is applicable greatly improves the model's performance.To let the model use sledgehammer, we replace the actual proof steps in the training dataset by a call to sledgehammer if the proof step either (a) contains the proof tactics 'meson, metis, and smt' (these tactics are typically generated by sledgehammer), or (b) belongs to a predefined simple set of proof tactics that can be easily generated. In particular, they are[by auto, by simp, by blast, by fastforce, by force, by eval, by presburger, by sos, by arith, by linarith, by (auto simp: field_simps)]When testing a generated proof with calls to sledgehammer, we follow the pipeline of [Jiang et al., 2022b] -first, we try to replace the 'sledgehammer' command by one of the predefined tactics. If all the attempts fail, we call the actual premises selection tool in Isabelle with a 10s timeout. If the tool does not return a valid proof, we consider this step incorrect.Note that Jiang et al. [2022a] decide when to replace the actual proof step by a call to sledgehammer more aggressively. They attempt to call sledgehammer at every proof step, and replace the actual proof step by sledgehammer if the attempt is successful. In contrast, our decision is made without interacting with the formal verifier. This is because applying sledgehammer to every proof step requires a lot of compute, which would significantly slow down the reinforcement learning process.Dataset split. Here we describe how to split the training and test data based on the dependency of the AFP files. We first compute the dependency graph by crawling the AFP website https: //www.isa-afp.org/entries/, which lists the dependency of the AFP entries. Then we find the set of AFP entries that all other entries do not depend on using the dependency graph, in which we randomly sample 10% of the entries as the holdout test set. The resulting holdout entries are:[Verified_SAT_Based_AI_Planning, SIFPL, Khovanskii_Theorem, Bondy, Rewriting_Z, Decreasing-Diagrams-II, Registers, LocalLexing, FeatherweightJava, FFT, Knot_Theory, Eval_FO,The former is in P whereas the latter is undecidable in the worst caseTuring et al. [1936],Church [1936].38th Conference on Neural Information Processing Systems (NeurIPS2024).In this paper, we use the terms 'lemma' and 'theorem' relatively -theorem refers to the statement that we are currently focusing on, and lemma refers to the statement proposed during the proof. In other words, there is no fundamental difference between a lemma and a theorem.https://www.isa-afp.org/Here we use the archive of AFP retrieved on2023-11-22.In our preliminary experiments, we observe that the model overfits after 2 epoch 


Decision: The work has interesting ideas and a sound methodology with good experimental results. However, as Reviewer dXkd points out, the paper is quite difficult to read and the main text lacks many important details. We suggest that the authors take the time to revise the work to adhere to the submission format before the camera-ready version.

Reviews: review: title: title: Review of Formal Theorem Proving by Rewarding LLMs to Decompose Proofs Hierarchically
main_review: main_review: Summary:
This paper introduces an RL-based method that rewards intermediate correct lemmas generated within the proof tree, improving the performance of a fine-tuned 7B language model (llemma) in formal theorem proving.

Comments:
- I suggest the authors cite the Hindsight Experience Replay paper earlier (e.g., in the introduction) rather than later (i.e., in the experiment section). There are many prior works in RL addressing sparse reward issues, and it may be useful to briefly review them somewhere.
- In Figure 2 / Setup section, the authors may provide more illustrations on classical proof tree generators without lemma proposals (e.g., the flat tree structure, the classical proof format). This would enhance the paper's readability for common readers.
- Lines 44-45 ("the best method along this line of research requires more than 1K GPU days ... [Lample et al., 2022] ..., whereas ours only takes less than 36 GPU days") might be slightly misleading, as (i) the result comparison is not A2A; (ii) there are many recent works that achieve competitive results using a relatively small computation budget (e.g., ReProver [Yang et al., 2023] requires only 7 GPU days of training with a small 0.3B model). I suggest the authors rephrase this line and use a more fair comparison or a weaker tone. Optionally, they could provide a table summarizing existing works on training costs. Additionally, does the benefit of the proposed method include inference efficiency? It would be nice to discuss this as well.
- I am slightly confused by Line 47, which states, "generate whole proofs directly without searching," as this seems orthogonal to the paper's contribution. The conditional proof generation should work both with and without searching. It would be nice if the authors could provide more justification for this.
- I recommend the authors provide generated proof examples across the methods used in the appendix. It would help to analyze the additional proofs proved by the proposed method and what the method still struggles with. How does the model perform as theorem complexity increases (e.g., deeper proof trees)? More data on the number of tree depths and their correlation with success rates would strengthen the argument for hierarchical decomposition.
- The RL method does not bring much gain on AFP 2023. It might be useful to conduct ablation studies on exploration versus exploitation adjustments for RL methods. The authors should try out other base model as well, e.g., gemma, deepseek.

Overall, I believe the contribution of this work on hierarchical proof generation methods and RL training pipelines is sufficient as an workshop paper, and it might be enhanced by further improvements as mentioned in the comments.
title: title: Formal Theorem Proving by Rewarding LLMs to Decompose Proofs Hierarchically
main_review: main_review: This paper addresses a significant challenge in formal theorem proving with large language models (LLMs), aiming to improve their ability to write proofs autonomously without the crutch of pre-provided lemmas. It proposes a reinforcement learning (RL) algorithm called Proof Decomposer (ProD-RL) that encourages models to decompose theorems hierarchically into lemmas, aiming for better generalization and decomposition capabilities.

---

#### Strengths and Weaknesses

Originality:  
The paper introduces an interesting reinforcement learning-based approach to theorem proving, where the model proposes and proves new lemmas instead of relying on human-written lemmas. This is a notable shift from previous approaches that assumed all relevant lemmas were given, simplifying the problem. The hierarchical decomposition strategy is novel and potentially impactful, making the model more aligned with how mathematicians tackle proofs incrementally.

Quality:  
The submission is technically sound. The model's architecture and RL methodology are clearly explained, including the reward mechanism that encourages lemma generation even when a full proof cannot be achieved. The experimental results demonstrate improvement over baselines, with a 47.5% pass rate on holdout theorems. However, the presentation could benefit from more clarity on the experimental setup, particularly regarding the comparison between supervised fine-tuning (SFT) and RL-based methods. Some additional details on the computational cost breakdown would strengthen the claims about efficiency.

Clarity:  
The paper is fairly well-written but has room for improvement in explaining the transition from generating lemmas to complete proofs. More visual examples of proof trees or step-by-step walk-throughs of generated proofs would help readers understand the model's outputs better. The experimental section, while robust, could benefit from more detailed explanations of Table 1 and other results.

Significance:  
This work has the potential to impact future research on automated theorem proving and formal reasoning with LLMs. It provides a scalable solution for generating novel lemmas and proves that LLMs can propose useful intermediate steps when tackling difficult proofs. This advances the field of mathematical reasoning with LLMs and lays the groundwork for models capable of solving theorems more autonomously. The results are promising, but further validation is necessary to evaluate the method's scalability to more complex theorems.

---

#### Questions:
- Could you provide more detailed analysis on the types of theorems where lemma proposals worked better than direct proofs? This would clarify the strengths of hierarchical decomposition.
- How does the method handle situations where lemma generation fails? Do fallback strategies exist?

---

#### Limitations:
The paper addresses its limitations well, including the difficulty in proving theorems with fewer training lemmas. However, the potential for increased complexity in very deep proof trees is not fully discussed. The scalability of this method to very complex theorems with deep proof hierarchies is a concern that could be explored more.
title: title: Official Review
main_review: main_review: Summary: The authors recognize the problem statement of trying to prove theorems without relying on calls to existing relevant lemmas, instead additionally requiring proof search to generate relevant lemmas and their proofs is a better task to measure proof planning ability and the ability to decompose proofs into high-level steps. The authors propose an RL-based training algorithm and a way to produce dataset to begin the training process for this task. The results generally indicate that the approach works.

Originality: The work, to be the best of my knowledge, is original in this setting. There are other papers for subgoal search, like (Czechowski et al, 2021) but operate in a more constrained setting. Similarly producing relevant lemmas also bears resemblance to LEGOProver, but I think the exact problem statement, domain of application, and methodology are fairly different.

Significance: Learning to decompose proofs is a valuable task and I think this work would be considered valuable at this workshop, especially as I can see more people becoming interested in this problem in the future.

Weakness: The results of the model on the more-out-of-distribution AFP 2023 set are barely positive. I would judge the performance improvement by ProD-RL over SFT w/o Lemma Proposal to be statistically insignificant. There may need to be stronger justification for this approach's ability to generalize to unseen problems.

Q: In the very back of the submission's supplementary materials a mention about miniF2F is included regarding the license. I didn't see any mention anywhere, else but miniF2F could be a good set of problems to test your method on (using the model trained on AFP). It is out-of-distribution and can provide more evidence about the efficacy of the model.
title: title: Interesting methodology, very poorly written up, blatant disregard of required submission format
main_review: main_review: This paper proposes to use RL (specifically, using the REINFORCE estimator) to train a (previously pre-trained) LLM to become better at theorem proving. Specifically, the paper argues that current evaluations are biased by giving the necessary lemmas in-context; instead, they require the LLM to identify (and then prove) the required lemmas by itself.

Unfortunately, this paper is very hard to read, especially since it abruptly ends in the middle of discussing results, with all further discussion deferred to the appendices. There are multiple details that remain unclear to me; some glaring issues include there being (as far as I can tell) not a single reference to the origin of the AFP dataset anywhere in the main text of the paper, as well as no explanation of what the ProD-SFT method they compare against does. Reading between the lines, I assume ProD-SFT is the same as the SFT baseline but allowing the model to make <invoke> statements to call out to sub-theorems (lemmas). If that is the case, it is not at all clear what data exactly it is being trained on, nor why this would ever perform worse than the baseline SFT model (unless proof decomposition is not actually particularly beneficial on this dataset, contrary to the authors' intuition). Another example issue is that after having read the main part of this paper (i.e. the 4 pages of content which reviewers were instructed to read), I was led to believe that the reward being used was binary correctness (see e.g. Figure 2). However, only upon reading appendix C.1 (as a courtesy to the authors) was I made aware that they in fact train a value network to produce a soft reward instead. This is the kind of thing that should not be omitted from the main text, since it makes the discussion very misleading.

Pros:
- The paper is well-written in the sense that the language being used is clear
- The topic is relevant to the workshop
- The results look like they might be interesting

Cons:
- The paper does not adhere to the required 4 page format in any meaningful sense; note that it does not even contain a conclusion in the main text. This would not be hard to remedy (why is there a 2-page introduction to a 4-page paper?) so I can only presume that the authors decided to blatantly disregard the submission instructions.
- The methodology is not clear, in particular with respect to how the baselines differ (see above).
- Minor: I was surprised that there was no expert iteration baseline; if the claim is that the model learns to generate and then internalize novel lemmas that are useful at test time then I would expect this to happen with expert iteration too, and that might give more stable updates than REINFORCE.
- All the other issues mentioned above.

Overall, I think the core methodology and results in this paper would be of interest to the attendees of this workshop, and that it might make for an interesting poster. However, on the grounds of the blatant disregard of the submission format as well as the omission of several key experimental details from the main paper I cannot recommend its acceptance.