Learning Formal Mathematics From Intrinsic Motivation — part 2/5
Cleaned-text code points [15778, 31597). Bibliography omitted here; source.md includes it.
To that end, our main insight is to leverage constrained decoding to obtain valid conjectures by construction. Our method will turn any language model — including a randomly initialized one, which is what we start with — into a probability distribution over strings that represent well-formed conjectures in dependent type theory over a given set of axioms. Ultimately, we will also train the LM to generate conjectures given a target difficulty. We use this ability to attempt to generate increasingly difficult problems for training, according to the agent’s current ability to prove theorems.
To reason about constraining the LM’s outputs, we leverage the abstraction of a completion engine , first introduced in the context of code generation with language models [ 29 ] . Assuming 𝒞 \mathcal{C} is the set of all valid conjectures, a completion engine will allow us to use any LM to sample strings from 𝒞 \mathcal{C} in an autoregressive fashion. Mathematically, 𝒞 \mathcal{C} is a function f 𝒞 : Σ ∗ → 𝒫 ( Σ ∗ ) f_{\mathcal{C}}{}:\Sigma^{*}\to\mathcal{P}(\Sigma^{*}) , taking a string s ∈ Σ ∗ s\in\Sigma^{*} and returning a set of strings f 𝒞 ( s ) f_{\mathcal{C}}{}(s) 1 1 1 This set can be implicitly defined by a regular expression, so it might be infinite. This allows a completion engine to represent constraints such as “what follows can be any valid identifier”. . Intuitively, we will sample conjectures from our LM by constraining it to strings that can be generated with iterated calls to f 𝒞 f_{\mathcal{C}} . Concretely, we will start with s = "" s=\texttt{""} , and query f 𝒞 ( s ) f_{\mathcal{C}}{}(s) to obtain the set of valid ways to begin to state a conjecture. After we choose one of those s 1 ∈ f 𝒞 ( s ) s_{1}\in f_{\mathcal{C}}{}(s) , we can then query f 𝒞 ( s 1 ) f_{\mathcal{C}}{}(s_{1}) to obtain the valid ways to proceed, and repeat until we have a complete conjecture. Our main challenge here is to construct a suitable f 𝒞 f_{\mathcal{C}} that it is sound (all conjectures obtained by this procedure are valid) and complete (all valid conjectures can be obtained by this procedure). After we define f 𝒞 f_{\mathcal{C}} , we can sample from any LM while guaranteeing that the output will belong to 𝒞 \mathcal{C} by using the Constrained Semantic Decoding (CSD) algorithm [ 29 ] .
To construct f 𝒞 f_{\mathcal{C}} , we analyze the minimal formulation of dependent type theory backing Peano [ 28 ] – essentially the classical Calculus of Constructions (CoC) [ 12 ] . In the CoC, mathematical propositions are constructions of type prop . Since generating a proposition might involve generating objects of arbitrary other types, depending on the axioms of the given domain, we will define a more general family of completion engines, f t f_{t} , which will constrain the LM to sample an object of an arbitrary type t t . At the end, we obtain f 𝒞 = f p r o p f_{\mathcal{C}}{}=f_{prop} .
To sample an object of type t t , we leverage the simple grammar of terms in Peano [ 28 ] to guide a recursive type-directed synthesis algorithm. Syntactically, terms in Peano are defined by the grammar e : := type | prop | x | ( e e ) | ( lambda x : e , e ) | [ ( x : e ) → e ] e::=\mbox{type}\ |\ \mbox{prop}\ |\ x\ |\ (e\ e)\ |\ (\mbox{lambda }x:e,e)\ |\ [(x:e)\rightarrow e] . The first two production rules give the names of two built-in base types, type and prop. We then have variables, function application, lambda functions, and function types (denoted in square brackets). As conventional in type theory, let Γ \Gamma be our typing context : a sequence of pairs of names and their types. For example, we might have Γ = [ n a t : t y p e , z : n a t , s u c c : [ n a t → n a t ] ] \Gamma=\left[nat:type,z:nat,succ:\left[nat\rightarrow nat\right]\right] , a context with three objects: a type nat , an object z having that type, and a function succ from nat to nat . Given a context, to obtain an object of type t t , it suffices to consider the formation rules in CoC to guide generation:
• If our target type is t = type t=\texttt{type} , we can generate either one of the names in Γ \Gamma having type t = type t=\texttt{type} (e.g., nat , for the example context above), or a function type.
• If our target type is t = prop t=\texttt{prop} , we can generate either one of the objects in Γ \Gamma having type prop , or a function type where the output type has type prop .
• If our target type is a function type, we must start by generating a square bracket; then, we (recursively) iteratively generate a type for the next parameter, or, if we already have at least one parameter, we can also generate the final output type.
• An object of any type t t can be formed by function application of a function f f chosen from Γ \Gamma , provided that the output type of f f can be unified with t t .
These rules allow us to determine the possible valid tokens at any point during generation. Besides the base types type and prop , objects of all other types are either in Γ \Gamma or the result of a function application. We use a recursive descent parser to parse the (incomplete) term we have so far (as originally done in [ 29 ] ), and compute the target type t t at the current point in generation. Then, we enumerate the possible next tokens for the LM by using the rules above, return the union of the sets of choices allowed by each rules as the output of f t f_{t} .
Note that this is a general procedure for searching for objects of any given type t t (i.e., inhabitants of that type). This is undecidable in general (theorem proving is a special case of type inhabitation), so this procedure might not terminate. Therefore, we set a maximum number of tokens K K (150, in our experiments), and ignore samples where the LM fails to generate a conjecture after K K tokens. In practice, we find the rejection rate for generating propositions to be low, < < 10% of the generations.
The above procedure forces the LM to generate a valid conjecture, but the LM still assigns a distribution over those. During training, we also aim to generate conjectures that are provable, but hard to prove. The signal on both success and difficulty is generated by running proof search (as we describe next) and, in cases where a proof is found, measuring it’s log-likelihood under the current policy, which correlates with how many iterations MCTS takes to find the proof (see Appendix).
### 3.2 Proof Search
Having a conjecture represented by a target type t t , we then perform proof search using Monte Carlo Tree Search (MCTS; [ 21 ] , [ 42 ] ), guided by a learned policy π θ \pi_{\theta} and value function V θ V_{\theta} . We represent both π θ \pi_{\theta} and V θ V_{\theta} using the same underlying language model that we use for conjecturing. We use Peano [ 28 ] as the environment for proof search. Peano exposes a finite action space, so we don’t need to generate actions using π θ \pi_{\theta} — it suffices to be able to evaluate them. More precisely, at a given state s s where we have actions a i ( s ) a^{(s)}_{i} available, we compute the distribution π θ ( a i ( s ) | s ) \pi_{\theta}(a^{(s)}_{i}|s) by evaluating the likelihood of each a i ( s ) a^{(s)}_{i} as the completion to a string of the form "STATE: <<s>>; POLICY:" . We read out the value of a state in a similar way — by considering the likelihood of 1 1 or 0 0 as being the next token following "STATE: <<s>>; VALUE:" . In both cases, the probability of the next token is normalized over only the choices that can lead to a valid completion.
States and actions in Peano are similar to several other interactive theorem provers, such as Lean and Coq. The state consists of a set of typed objects, along with a set of open proof goals (which are types). Objects in the state whose type is a proposition type are interpreted as evidence for that proposition (either a proof or an assumption). Actions might add new objects to the state (e.g., take the successor of one existing natural number, or use an equality to rewrite an existing proposition into a new one), change the goal (by backward chaining), or both (e.g., if the goal is to prove A → B A\rightarrow B , which is used to represent both logical implication and universal quantification, the intro action will add object of type A A to the state and change the goal to B B ).
When proof search succeeds, we can extract examples to train both the policy and the value function. For the policy, we simply extract the actions that lead to the proof as language modeling examples, using the same format we use to query π θ \pi_{\theta} . For the value function, we take the states in the subtree that complete the proof as examples with value 1 1 , and a random set of other states as examples with value 0 0 for training. When proof search fails, however, this naïve procedure does not extract any training examples. This happens often at the beginning, since our model initially generates a large number of conjectures it cannot prove (either because they are false, or because naïve proof search times out). But forward actions in the proof search tree often construct proofs for other propositions, even if they are irrelevant for proving the original goal. We levarage this fact for generating training data by hindsight relabeling, as we describe next.
### 3.3 Hindsight Relabeling
Even a conjecture that fails to be proven can be highly informative about the domain. During proof search, forward actions that apply functions whose result type is a proposition type (e.g., concluding A from (and A B) ) produce proofs, even when those proofs might not be useful for proving the original conjecture. In Reinforcement Learning, the well-known method of Hindsight Experience Replay [ 1 ] extracts training data for the policy from such failed trajectories by relabeling the trajectories with goals that were in fact achieved, as opposed to the original goal. For those alternative goals, the trajectory then represents a successful sequence of actions. We apply this idea to extract training examples for both the policy and value functions from proof search trees, by picking nodes after forward actions that produced a proof, and walking upwards in the tree until we find a backward action (since those change the goal). That path then becomes a successful trajectory after we relabel the goal. Two important steps to improve data quality are (1) we clean up the solutions by eliminating steps irrelevant for the proof of the new goal, and (2) we only add proofs of goals never seen before, to avoid oversampling trivial facts that are rediscovered extremely often (such as 0 = 0 0=0 ).
We go one step further and observe that hindsight relabeling can also be useful for training the conjecturer . Concretely, the procedure we described above produces a set of proofs p i p_{i} for relabeled statements g i g_{i} . All of these statements are therefore true in the mathematical domain, and we use them as examples of true conjectures. As a concrete example, in arithmetic, the agent will often conjecture simple expressions such as 2 + 1 = 0 2+1=0 . Most equalities generated at random will be false. However, by applying the Peano Axioms and facts about equality in succession, the agent eventually finds a proof of 2 + 1 = 3 2+1=3 . Evaluating the likelihood of the proof under π θ \pi_{\theta} gives a measure of the difficulty of this alternative statement for the current policy. This insight allows our agent to learn about hundreds of new unique, true statements in each proof search. As our experiments show, we find hindsight relabeling to be crucial for enabling the agent to steadily conjecture harder statements that it is able to prove.
### 3.4 Training loop
Finally, we tie all components together by alternating between (a) generating a batch of conjectures, (b) running proof search on them, (c) extracting training examples from the search trees, and finally (d) training the underlying LM using the standard cross-entropy objective. When a proof is found, either directly or by hindsight relabeling, we first compute a continuous difficulty metric of the statement by taking the log-likelihood of the proof under the current policy. To discretize difficulties, we then consider difficulty percentiles of the last batch of conjectures: we take the 10% least likely proofs to be associated with “hard” conjectures, the 50% most likely to be considered “trivial”, and the remaining are taken as “easy”. When put together, these components form a self-improving loop that starts only from the axioms of the given mathematical domain, as our experiments show.
## 4 Experiments
We now evaluate Minimo 2 2 2 Our code is available at https://github.com/gpoesia/minimo . on three mathematical domains. We experiment with axiomatic systems for (a) propositional logic, (b) arithmetic, and (c) abstract groups. All the axioms are given in the Appendix. We then train agents over 5 iterations of conjecturing and theorem proving, generating 200 conjectures in each batch, running MCTS for proof search with 1000 expansions, and evaluate our agents guided by the following research questions:
RQ1: Can our conjecturing method generate increasingly harder conjectures as training progresses?
Do agents improve at theorem proving as they train on their own generated problems?
Is hindsight relabeling effective at improving conjecturing and theorem proving?
Do our agents improve at proving an externally given set of human-selected theorems, even if these are not given during training?
The first three questions require intrinsic evaluations — they ask whether the learning dynamics of agents trained with Minimo matches the desiderata of self-improving at both conjecturing and theorem proving while given only the axioms. The last question is an extrinsic assessment of what our agents learn — we evaluate whether the learning progresses in a way that aligns with an external criteria — namely, the proficiency of the agent at proving theorems from human sources (a textbook and a Lean game).
### 4.1 Learning dynamics
We first investigate RQ1 and RQ2. Figure 2 shows the progression of difficulty across 5 iterations of the Minimo learning loop (first conjecturing, then running proof search, and training on collected examples). Here, we evaluate the average log-likelihood (y-axis) of conjectures proven at each conjecturing iteration (x-axis) under the policy after each iteration (line color). Policy 0 is the initial (randomly initialized) policy, while subsequent policies were trained on the examples obtained during proof search, with hindsight relabeling, up to the previous iteration. Lower log-likelihood generally means harder conjectures (i.e., they tend to take more search iterations, see Appendix).
##### Difficulty increases as training progresses (RQ1).
Fixing an early policy, the log-likelihood of proofs under that policy steadily decreases across training iterations. This is reflected in the negative slope of difficulty across iterations when the policy is fixed. In particular, the policy at iteration 0 serves as a naïve search baseline, since it is essentially uniform. We observe that, as training progresses, this policy struggles increasingly more with each new batch of conjectures. The same pattern repeats for subsequent policies when we consider conjectures generated in future iterations, giving a generally downward trend in log-likelihood of the solutions for any given policy, showing that conjectures get increasingly more challenging. This provides positive evidence for answering RQ1: in all 3 domains, the conjecturer is able to increasingly propose harder provable conjectures as training progresses.
##### Proof policy improves as training progresses (RQ2).