Mathematical Reasoning in Latent Space
Authors: Dennis Lee
Christian Szegedy
Markus Rabe
Sarah Loos
Kshitij Bansal
Venue:  (2020)
Keywords: machine learning
formal reasoning

Abstract: We design and conduct a simple experiment to study whether neural networks can perform several steps of approximate reasoning in a fixed dimensional latent space. The set of rewrites (i.e. transformations) that can be successfully performed on a statement represents essential semantic features of the statement. We can compress this information by embedding the formula in a vector space, such that the vector associated with a statement can be used to predict whether a statement can be rewritten by other theorems. Predicting the embedding of a formula generated by some rewrite rule is naturally viewed as approximate reasoning in the latent space. In order to measure the effectiveness of this reasoning, we perform approximate deduction sequences in the latent space and use the resulting embedding to inform the semantic features of the corresponding formal statement (which is obtained by performing the corresponding rewrite sequence using real formulas). Our experiments show that graph neural networks can make non-trivial predictions about the rewrite-success of statements, even when they propagate predicted latent representations for several steps. Since our corpus of mathematical formulas includes a wide variety of mathematical disciplines, this experiment is a strong indicator for the feasibility of deduction in latent space in general.

Published as a conference paper at ICLR 2020
MATHEMATICAL REASONING IN LATENT SPACE
Dennis Lee, Christian Szegedy, Markus N. Rabe, Kshitij Bansal, and Sarah M. Loos Google Research Mountain View, CA, USA
{Idennis, szegedy, mrabe, kbk,
com
ABSTRACT
We design and conduct a simple experiment to study whether neural networks can perform several steps of approximate reasoning in a fixed dimensional latent space. The set of rewrites (i.e. transformations) that can be successfully per- formed on a statement represents essential semantic features of the statement. We can compress this information by embedding the formula in vector space, such that the vector associated with statement can be used to predict whether a state- ment can be rewritten by other theorems. Predicting the embedding of a formula generated by some rewrite rule is naturally viewed as approximate reasoning in the latent space. In order to measure the effectiveness of this reasoning, perform sequences of rewrite steps both in formula space and in latent space, and compare the quality of embeddings of the resulting formulas to their predicted embeddings. Our experiments show that graph neural networks can make non-trivial predictions about the rewrite-success of statements, even when they propagate predicted la- tent representations for several steps. Since our corpus of mathematical formulas includes wide variety of mathematical disciplines, this experiment is a strong indicator for the feasibility of deduction in latent space in general.
INTRODUCTION
Automated reasoning has long been considered require development of logics and "hard" algo- rithms, such as backtracking search. Recently, approaches that employ deep learning have also been applied, but these have focused on predicting the next step of proof, which is again executed with a hard algorithm (Loos al.) 2017 Gauthier et al.) 2017 Lederman et al. 2018 Bansal et al. 2019b).
We raise the question of whether hard algorithms could be omitted from this process and mathe- matical reasoning performed entirely in the latent space. To this end, we investigate whether we can predict useful latent representations of the mathematical formulas that result from proof steps. Ideally, we could rely entirely on predicted latent representations to sketch out proofs and only go back to the concrete mathematical formulas to check if our intuitive reasoning was correct. This would allow for more flexible and robust system designs for automated reasoning.
In this work, we present a first experiment indicating directly that theorem proving in the latent space might be possible. We build on HOList, an environment and benchmark for automated the- orem provers based on deep learning (Bansal et al. (019b)) which makes use of the interactive theorem prover HOL Light (Harrison 1996), an interactive proof assistant. The HOList theorem database comprises over 19 thousand theorems and lemmas from variety of mathematical do- mains, including topology. multivariate calculus, real and complex analysis, geometric algebra, and measure theory. Concrete examples include basic properties of real and complex numbers such as (--1= = y) < (x: = y yー ), and also well-known theorems, such as Pythagoras' theorem, the funda- mental theorem of calculus, Skolem's 's theorem, Abel's theorem for complex power series, and that the eigenvalues of a complex matrix are the roots of its characteristic polynomial.
We focus rewrite rules (or rewrites in short). Rewrites are only one of several proof tactics in HOL Light, but they enable powerful transformations on mathematical formulas, as they can be given arbitrary theorems parameters. For example. the formula 32 can be rewritten to 3-3 by performing a rewrite with the parameter x² x x. Alternatively, a rewrite may diverge (as operates recursively) or it may return the same formula in both these cases consider the rewrite

---

Published as a conference paper at ICLR 2020
Formula
Formula
Formula 3
Formula
Param.
Param
Param. 3
R
1024
1024
R 1024
R 1024
Figure 1: Overview of latent space prediction propagated over multiple steps
to fail. For instance, in the example above, the rewrite would fail if we used equation x + y y as a rewrite parameter instead, since the expression 3 does not contain any operators to match with.
In our experiments, we first train a neural network map mathematical formulas into a latent space of fixed dimension. This network is trained by predicting - based on the latent representation being trained whether a given rewrite going to succeed (i.e. returns with a new formula). For successful rewrites we also predict the latent representation of the resulting formula. To evaluate the feasibility of reasoning in latent space over two steps, we first predict the latent representation of the result of rewrite, then we evaluate whether the predicted latent representation still allows for accurate predictions of the rewrite success of the resulting formula. For multi-step reasoning beyond two steps, predict the future latent representations based the previous latent representation only without seeing the intermediate formula (cf. Figure[]. Our experiments suggest that even after nine steps of reasoning purely in latent space, neural networks show non-trivial reasoning capabilities, despite not being trained on this task directly.
2 RELATED WORK
Our work is motivated by deep learning based automated theorem proving, but is also closely con- nected to model based reinforcement learning and approaches that learn to predict the future as part of reinforcement learning.
Model based reinforcement learning is concerned with creating a model of the environment while maximizing the expected reward (e.g. Ha & Schmidhuber (2018)). Early works have already shown that predicting the latent representations of reinforcement learning environments with deep learning sometimes feasible even over many steps (Oh et al.) 20155] Chiappa al.] [017 Burda et al., 2019). This can enable faster training, since it can preempt the need for performing expensive simulations the environment. Predicting latent representation was also proposed in Brunner al. 2018) as a regularization method for reinforcement learning.
One recent successful example of model based reinforcement learning Kaiser al. (2019) where the system learns to predict the pixel-wise output of the Atari machine. However this approach is based on actually simulating the environment directly in the "pixel space" as opposed to performing predictions in low dimensional semantic embedding space. More related to our work is Piotrowski cct al. (2019), which attempts to learn to rewrite simple formulas. The goal is there again is to predict the actual outcome of the rewrite rather than a latent representation of it. In Dosovitskiy & Koltun (017)), they predict "expected measurements" as an extra supervision in addition to the reward signal.
Graph neural networks have been used for premise selection in higher order logic (Wang et al. 2017) and more recently by PPaliwll et al. (2020) the component of DeepHOL., neural theorem prover for higher order logic. this work, we build upon their neural network architecture, but utilize it for different task.
HOL LIGHT
HOL Light (Harrison, 1996) is an interactive proof assistant (or interactive theorem prover) for higher-order logic reasoning. Traditionally, proof assistants have been used by human users for creating formalized proofs of mathematical statements manually. While proof assistants have some limited forms of automation, it is still cumbersome process formalize proofs, even when they
2

---

Published as a conference paper at ICLR 2020
already available in natural language. Some large scale formalization efforts were conducted successfully in HOL Light, Coq (Coq) and Isabelle/HOL and other theorem provers, for example the formal proof of the Kepler conjecture (Hales et al. 2017), the four color theorem (Gonthier, 2008), and the correctness of microkernel (Klein et al. 2009). They required significant meticulous manual work and expert knowledge of interactive proof assistants.
Lately, there have been several attempts to improve the automation of the proof assistants signifi- cantly by so called "hammers" (Kaliszyk& & Urban 2015). Still, traditional proof automation lacks the mathematical intuition of human mathematicians who can perform complicated intuitive argu- ments. The quest for modelling and automating fuzzy, "human style" reasoning is one of the main motivations for this work.
3.1 REWRITE TACTIC IN HOL LIGHT
The HOL Light system allows the user to specify a goal to prove, and then offers a number of tactics to apply to the goal. A tactic application consumes the goal and returns a list of subgoals. Proving all of the subgoals is equivalent to proving the goal itself. Accordingly, if a tactic application returns the empty list of subgoals, the parent goal is proven.
In this work, we focus on the rewrite tactic (REWRITE TAC) of HOL Light, which is a particularly common and versatile tactic. It takes list of theorems as parameters (though in this work only consider applications of rewrite with a single parameter). Parameters must be an equation or a conjunction of equations, possibly guarded by condition. Given a goal statement T and parameter P., the rewrite tactic searches for subexpressions in T that match the left side of one of the equations in P. and replaces it with the right side of the equation. The matching process takes care of variable names and types, such that minor differences can be bridged. The rewrite procedure recursive, and hence tries to rewrite the formula until no opportunities for rewrites are left. The rewrite tactic also has set of built-in rewrite rules, representing trivial simplifications, such as FST y) = .r. Note that REWRITE TAC uses "big step semantics", meaning that the application of each individual operation can perform multiple elementary rewrite steps recursively. For more details on REWRITE.TAC, refer to the manual (HOL Rewrite Tactic Reference).
REASONING IN LATENT SPACE
In this section we describe our training data, model architecture, and an evaluation setup that allows us to demonstrate mathematical reasoning in latent space. We use the embedding model by Paliwal jtt al. (2020) to map formulas into a fixed dimensional latent space, and train the model on two prediction tasks: (I) the outcome (success or failure) of a large number of possible formula rewrite operations, and (2) on the embedding of the outcome of the rewrite. For evaluation, we use the embedding prediction not only for predicting the embedding of the next formula, but also the formulas that result from several successive rewriting steps.
The challenge in the design of the model is preventing the embeddings from collapsing. Trained naively to minimize the 1 distance between the predicted and the true embedding of the result of the rewrite, the embedding model quickly learns to map most formulas to the same embedding destroying the information needed for performing multiple rewrite steps in latent space. To address this problem, we stop the gradient from updating the embedding of the outcome formula directly (but we share the embedding network for the source expression).
We have also found that adding extra noise the embedding during training makes the formulas more robust to prediction errors for multi-step step reasoning. We hypothesize that training with noisy embeddings acts as fuzzy groo-correction mechanism by preventing errors from accumu- lating as quickly during multiple steps of reasoning in the latent space. Both ideas are described in detail in Subsection
4.1 TRAINING DATA
We generate the training data from the theorem database of the HOList environment (Bansal et al.) 2019b), which contains 19591 theorems and definitions ordered such that later theorems use only earlier statements in their human proof. The theorems split into I 16655 training, 3668 validation,
3

---

Published as a conference paper at ICLR 2020
and 3620 testing theorems. To generate our training data, generate all pairs (T, P) of theorems from the training set, where P must occur before T in the database (to avoid circular dependencies). We then interpret theorem T goal prove and try to rewrite T with P using the REWRITE TAC of HOL Light. This can result in three different outcomes:
1. 7 is rewritten by theorem P successfully, and the result differs from T.
2. The rewrite operation terminates, but fails to change the input theorem T.
3. The rewrite operation times or diverges (becomes big).
In our experiments, we consider only the first outcome successful, i.e. when the application finishes within the specified time limit and changes the target, as a successful rewrite attempt. Each training example, therefore, consists of the pair (T, P). the success/fail-bit of the rewrite (1 for successful rewrites, O for failed rewrites), and, for successful tactic applications, the new formula that results from the rewrite, which we denote with R(T, P). Of the 115 million possible pairs of (T, P) where T falls in the training split, there are 1.6 million successful rewrites that are used as training examples.
4.2 MODEL ARCHITECTURE AND TRAINING METHODOLOGY
Let S denote the set of syntactically correct higher-order logic formu- las in HOL Light. Additionally, we define two latent spaces, I, = R and M R (k 1024), corresponding embeddings of theorems used as goals, and the statements used as parameters, respectively.
L2
We use two embedding networks 7 S one for each of the two formulas T and P. For both embedding net- works, follow the GNN approach by PPaliall al. (2020): We in- terpret the formulas' abstract syntax trees as graphs, feed the graphs into graph neural networks, and let the networks exchange messages between neighboring nodes for a fixed number of hops. The final em- beddings of the nodes are max-poot max-pooled form representation of the entire formula in fixed dimensional embedding space R
L and S - M,
E
π
Y
P
R(T. P
The embeddings are then concatenated and processed by a three-layer perceptron L x M - C, which is followed by C - R. a single output linear function producing a logit predicting whether theorem P rewrites T. The perceptron is trained with logistic re- gression on the success/fail-bii of the rewrite. Formally, this output is (((((((T, which we also refer as p(T, P) p(7(T), (P)).
Figure 2: Depiction of the
network architecture
The output of is also used as input to another linear function e: : C - L, which predicts the latent representation of the result (R(T, P)) by minimizing the 12 distance. Formally, this output e(c(T(T), n (P))), but we also use e(T, P) or ((>(T),
4.3 STOP GRADIENTS AND NOISE
Note that stop gradient backpropagation between the 12 loss and the second embedding tower (R,(T, P)). The purpose of this is to prevent the second tower from making the magnitude of the embedding of the rewrite outcome very small or "collapse". We discuss alternative solutions to this noise has the effect of more accurate downstream predictions in multiple step reasoning.
problem in S
In the case of the second goal embedding tower y, we add Gaussian noise with 0 mean and standard deviation 3 10-3 during training only. We can see the experimental section that training with
4.4 REASONING
After we have trained our model on the training set theorems (and theorem pairs) of the HOList benchmark, use it to perform rewrites in the latent space alone.
4

---

Published as a conference paper at ICLR 2020
We p: S S - R as a quality metric for the propagated embedding vector. Given approx- imation Ey of the latent representation y(T), we can evaluate p( Ey,, (P)) for a large number of tactic parameters P. This is compared with true rewrite successes of by P to assess the quality of the approximation.
To evaluate multiple steps of reasoning in the latent space, start with the embedding of formula Ti and predict the embeddings that result from rewriting by theorems PI, Pa. in that order. To assess the quality of the overall reasoning, the same sequence of rewrites is performed with the actual formulas and the final predicted embedding is evaluated with respect to the formula resulting from the sequence of formal rewrites.
In latent space L = R we start from some initial theorem Ti E S. The following schema indicates the sequence of predictions performed in latent space L:
Ti 3 11 E L ((,))( l2 E L c())m(Pb))
la E ((.,5((3))
This is compared with the following formal sequence of rewrites:
Ti R( T2 R(o.Pi)
T3 R(..Pa)
That is, le approximates the latent vector of T2, that is ((R(T, P1)) and approximates the latent vector of T3, that is (R)( P2). by construction.
ALTERNATIVE MODEL AADELLARCHITER ARCHITECTURES
Learning high quality embedding space while simultaneously learning to make predictions that space proved to be some challenge. As we described above, the embeddings quickly approach trained naïvely, minimizing the least squares distance between the e(T, P) and (R((T, P)) without improving the quality of Through the course of this work we evaluated several alternative approaches to tackling this problem. The empirical comparison between these architectures shown in Section 6.4
5.1 FIXED LATENT SPACE
Our first, straightforward approach to this problem was to fix the embedding space L and learn predictions from a separate space I.' The advantage of this approach is that it does not suffer from embedding collapse during the prediction of the embedding of the rewrite outcome. On the other hand, this approach incurs training of three different models, which is awkward and cumbersome. The three models are shown in Figure
Y
Y'
1. Rewrite success prediction a 5 x S
2. Rewrite outcome prediction 8 S - R x
T
P
3. Embedding alignment prediction
L'.
Figure 3: Depiction of the network ar- chitecture
The rewrite success prediction model (T, P) is identical to function p described above. Training fixes the embedding space L, from which we predict rewrite success.
In addition to a, we train separate rewrite outcome prediction (T, P). This model has an identical but with extra prediction layer e' to predict the embed- ding vector of the outcome of the rewritten formulas L. Here the embedding towers are denoted by n and n'. the combiner network is e' and the linear prediction layers p' and e'. That is: ('d)))) x' (P))). We denote the embedding space learned by
two-tower embedding architecture
(T, P) = (p' (d'( w, L'.
Since a and w produce latent vectors 7(T) E L and (T) E L' in different spaces, we need to align those spaces enable deduction purely in the embedding space. In particular, (R(((T, Po), Pi))
5

---

Published as a conference paper at ICLR 2020
is approximated by (R Po)).=((P)) (P)),, and need (R(,T, Po)). For that, train translation model a L - L', which is a three layer fully connected network with 1024 units in each layer and rectified linear activations, which predicts (T) given embedding of xTT). Note that a does not see 7 E S as an input, it makes its prediction based on the latent space representation
((T) of T alone.
Reasoning with this model can be summarized by:
T Th L' ( le e L " L L' (2)
E LL - Po)
5.2 CENTERING AND NORMALIZATION
A different approach to learning good embedding prediction is normalizing the embeddings fixed length. The na've implementation of this leads both and e to predict an almost-constant embedding vector independent T' with the relevant information encoded in the noise. Centering the embeddings before normalization mitigates this issuc. In our experiments, we found that using a variation of batch normalization without learned mean and scale parameters, followed by l2 nor- malization, is sufficient to force the model to predict meaningful embeddings from However, this model sensitive hyperparameters, in particular the scaling losses between p and e, and even in the best configuration we found performance was worse than our final model.
5.3 STOCHASTIC EMBEDDINGS
Another natural approach is to replace the mean squared error objective of the embedding prediction output with KL divergence objective, which is scale invariant. To do this, we attempted to substitute the deterministic embedding tower 7y with a stochastic variant, where the final layer of the tower is used to parameterize a diagonal Gaussian distribution. The final layer of the embedding network is similarly used to paramaterize a diagonal Gaussian. K KT(e(T, (R(,(R(T, P))) is the objective to minimize.
This model proved difficult to optimize: the naïve solution fails to converge. A carefully tuned information bottleneck mitigates this problem, but the performance of the model remains poor.
EXPERIMENTS
This section provides an experimental evaluation that aims to answer the following question: Is model capable of predicting useful embedding vectors multiple steps ahead? We explore the prediction quality of the embedding vectors (for rewrite success) and see how the quality of predicted embedding vectors degrades after is used for predicting multiple steps in the latent space alone.
6.1 NEURAL NETWORK ARCHITECTURE DETAILS
Our network has two independent towers, which each 16-hop graph neural networks with internal node representations in 122 The output of each of the two towers fed into layer that expands the dimension of the node representation IR 1024 with fully connected layer with shared weights for each node. This is followed by maximum pooling over all nodes of the network. We add Gaussian noise with 3 10-3 3 the embedding of the goal theorem. The two resulting embedding vectors are concatenated along with their element- wise multiplication, and are processed by three layer perceptron with 1024 hidden units and rectified linear activation between the layers.
Our model is trained with g 16 choices of T in each batch. For each Ti there are 1+1 = 16 choices one successful example in each group and random negatives. However all choices of P corresponding to other Teyt are used as negative instances as well, regardless of whether they would rewrite it - this is justified by the fact that only few theorems rewrite any given T; so this introduces only small amount of uncorrelated label noise. This training methodology is motivated by the fact that evaluating the combiner network is much cheaper than computing the embedding using the graph neural network. Based on Alemi et al. (2016), we expect that hard negative mining, or iteratively replacing random negative examples with examples that p(T, P) classifies incorrectly, would improve our results significantly, but it is left for future work.

---

Published as a conference paper at ICLR 2020
6.2 EVALUATION DATASET
In order to measure the performance of our models after multiple deduction steps are performed, generate datasets Do. Di...., D, successively by applying rewrites to randomly selected set of examples. We start with all theorems from the validation set HOList, denoted by To. For each theorem Tıj E Tú, perform all possible rewrites R(Ti,, Pt). k and include them in D, We then randomly select a new goal to include Ti accurately represent the full range of complexity of theorems present in the database, while the distribution of parameters P and the ratio of successful rewrites by each parameter matches that of the space of all possible rewrites.
in Ti+1- This method ensures that the goals in
6.3 EVALUATION OF REWRITE PREDICTION MODEL
To evaluate a model m S
x S - IR on dataset D,, we compute mTT, P) for each pair (T, P) in the dataset. We expect pairs where R(T, P) succeeds to score highly relative to pairs where R(T, P) fails. For concrete metric, under the curve (AUC) as our main metric. Higher curves and higher AUC values represent more accurate predictions of rewrite success. We measure how ROC curves change as different approximations of ((").
compute the ROC curve of the predictions and use the area
As we have trained our model only on pairs of theorems from the original database, the model exhibits increasing generalization error evaluate it on formulas with increasing num- ber of rewrites. The accuracy of the the model p evaluated directly on theorems after multi- ple rewrites is referred to as the "True" mea- surement. This gives an upper bound the rewrite success prediction since noisier embed- ding vectors generated by will produce worse results on average.
10
Occurrence
10*
10a
10
10*
5000
10000
15000
20000
Additionally, want to measure the rate at which the latent vectors degrade as we prop- agate them in embedding space as described in Subsection [4.4 Embedding predictions generated directly from theorem statements, ((n))") (P)), are referred to as "One Step", while predictions produced through propagat- ing are referred to as "Multi Step". A compar- ison these generation methods is shown in Figure
Parameter Sorted By Usefulness
Figure 5: The distribution of successful rewrites for cach parameter P, computed over all (T, P) pairs in the database. Note that small portion of parameters can be used for many theorems. This allows the Random and Usage baselines to achieve above-chance performance by predicting success based on parameter alone.
In order to evaluate the performance of our model in isolation we need to compare it with carefully selected baselines:
RCT:: Pg)
R(Ti, Pt)
RCTg Pa)
To
Theorem
Ep
Ej
E;
E
True
<Fiy s(P.))
aF miPs))
E2
One Step
odls atP
<LE>
Es
Step
Figure 4: Detailed view of embodding generation scheme for the True, Pred (One Step), and Pred (Multi Step) curves.
7

---

Published as a conference paper at ICLR 2020
1. As Figure shows, few theorems much more frequently applicable than others. We want to see how the prediction of rewrite success performs based on the rewrite parameter alone if we ignore the theorem to be rewritten. One way to establish such baseline is to feed a randomly selected theorem 7' to p instead of to predict its rewrite success. This is referred
the "Random" baseline. 2. A stronger "baseline" is achieved by utilizing the ground-truth rewrite statistics to rank each parameter P by the number theorems it successfully rewrites. This the best achievable prediction that does not depend on T'. We refer to this as the "Usage" baseline.
0.2
pos
neg
0.1
Rabio of Scares
True Positive Rate
0.2
pos
0.4
0,0
0.2
neg
Pred Rewrite
0.1
Random
0.0
1.0
-35
False Negative Rate
Score
Logit
Figure 6: Left: Histograms of scores first step rewrite predictions (top), ninth step of rewrite predictions (middle), and random baseline (bottom). Right: Corresponding ROC curves.
Figure 6 shows the distribution of the theorem pair prediction score logits of p(c()(T), x (P))). for the those "positive" pairs that rewrite and the "negative" pairs that do not rewrite. Note that the ratio is normalized separately for the positive and negative pairs as negative pairs occur much more frequently than positive pairs.
One can see that the quality of the rewrite success prediction degrades somewhat after nine steps of reasoning purely in the latent space, but it is still much better than the random baseline. This gives clear evidence that the embedding prediction manages to propagate much useful information over multiple steps in the latent space alone.
1.0
Area under RoC
Positive
0.4
0.90
0.88
0.86
Rewrite Steps
False Negative Rate
True
[One step!
(Muli step)
Usage Baseline
Baseiine
Figure 7: Left: Comparison of Area under ROC curves for each embedding method. Right: Comparison of ROC curves for each embedding method after the ninth rewrite.
In Figure 7 we make further measurements and comparisons on the quality of reasoning the em- bedding space. One see that model could perform reasoning for nine steps in the embedding space and still retain a lot the predictive power of the original model.
In order to appreciate the above results one should keep in mind that model is not trained on statements that were already rewritten. All the training was done only on the theorems present in the initial database. The reduction of prediction performance is apparent from downward trajectory of the "Pred (Multi step)" curve, which isolates this effect from that the error accumulated by the embedding predictions, the effect of which measured indirectly by the "True" curve.
8

---

Published as a conference paper at ICLR 2020
Ratio of Distances Ratio of Distances
2.5
0.5
0.4
0.3
2.0
0.2
12 Distance
1.5
0.0
0.5
0.4
0.3
0.5
0.2
0.1
0.0
0.0
0.5
1.0
2.5 12 Distance
1.5
2.0
3.0
3.5
2
9
Rewrite Steps
Pred (One step)
Pred (Muld step)
Random Basel ine
Figure 8: Left: Histogram of distances between each embedding method and the true embedding for rewrite
steps (top) and 9 (bottom). Right: Plot showing the mean 12 distances across rewrite steps.
In Figure 8 we measure 12 distance from the predicted embedding vectors to the true embedding vectors ((7) of formulas after multiple rewrite steps in the latent space. These results are consistent with our earlier findings on success of rewrite predication after rewrite steps in the latent space: while there is some divergence of the predicted embedding vectors from the true embedding vectors (as computed from the rewritten statements directly), the predicted embedding vectors are significantly closer to the true embedding vectors than randomly selected embeddings.
6.4 COMPARISON OF ALTERNATIVES
1.00
1.0
0.99
0.98
0.8
Area under RoC
0.97
True Positive Rate
0.6
0.96
0.4
0.94
0.93
0.2
0.92
0.0
0.91
2
3
0.2
0.4
0.6
0.8
1.0
Rewrite Steps
False Negative Rate
Space
Gracients
Figure 9: Left: Comparison of Area under RoC curves for each architecture, using predictions fully in the latent space (Pred (Multi Step)). Right: Corresponding ROC curves after the ninth rewrite step.
Our final model adds noise to the embeddings produced by 7/ during training time only. This forces the network produce embeddings sufficiently far apart distinguish between them. We found that this method provided the best experimental results compared to the models discussed in Section S as shown in Figure 9 We hypothesize that the noise makes the network robust future prediction mistakes as the prediction of the next rewrite is trained on noisy versions of the theorem embedding. This may act as a natural error correction mechanism. Figure Ø supports this hypothesis, since the difference between the noisy and noiseless versions of the model increases with the number of rewrite steps performed in the latent space. However, our experiments indicate that adding noise alone without stop gradients or other approaches is not sufficient to learn diverse latent space.
6.5 EVALUATION OF REWRITEABILITY
where R(T, P) fails and R( (TR, Po), P) succeeds, or vice versa. Particularly interesting rewrites are those that change the rewritability of a theorem, or rewrites

---

Published as a conference paper at ICLR 2020
First Rewrite Succeeds
6.25
Sccond Rewrits Succeeds
1.8
Secord Rumrita Falls
0.10
0.05
True Poskive Rate
0.5
First Rewrite Falls
0.20
Second Rewrita Succeeds
0.4
6.15
Sccond Rewritz Falls
0.10
0.2
0.05
First Rewiite Succoets
First Rowrite Falls
Negalive
Figure 10: Left: Histogram of scores for ((R(T, where R) succeeds (top) and where R(T, P) fails (bottom). Right: Corresponding ROC curves.
In Figure 10 we evaluate our model on these cases. We created a database of statements that are the result of RRR(T, Po), P) and split into groups, depending on whether R(T, P) succeeds or fails. Within each group, we score each rewrite and plot the resulting ROC curves. From this figure, we can see that the model is capable detecting when the rewritcability not change (First Rewrite Succeeds and Second Rewrite Succeeds, or First Rewrite Fails and Second Rewrite Fails).
theorem does
Most interestingly, cases where the first rewrite fails and the second rewrite succeeds are also scored correctly. This demonstrates that the predicted embeddings are precise enough to encode the ap- pearance of new possibilities. However, it is noteworthy that the model often fails to detect when rewrites are no longer possible. This may be due to label noise caused by randomly sampling neg- ative examples, some "negative" examples actually positives, leading generalization error. deeper investigation these failures an interesting subject for further work.
7 CONCLUSION
In this paper we studied the feasibility of performing complex reasoning for mathematical formulas fixed (1024) dimensional embedding space. We proposed new evaluation metric that measures the preservation semantic information under multiple reasoning steps in the embedding space. Al- though our models were not trained for performing rewrites on rewritten statements, nor were they trained for being able to deduce multiple steps in the embedding space, our approximate rewrite prediction model p has demonstrated significant prediction power on as many as nine rewrite steps in latent space alone. Our methods showcase simple but general methodology for reasoning in latent space.
This work can also be representing mathematical statements by neural networks. We hope that this will lead to improved neural network architectures for mathematies.
easy use, fast to train, and crisp evaluation methodology for
is likely that such representations prove helpful for faster learning to prove without imitating human proofs like that in DeepHOL-Zero Bansal al. (2019a), given that premise selection is closely related task to predicting the rewrite success of statements. Self-supervised pre-training or even co-training such models with premise selection could prove useful way of learning more semantic feature representations of mathematical formulas.
REFERENCES
Alexander A Alemi, François Chollet, Geoffrey Irving, Niklas Eén, Christian Szegedy, and Josef Ur- ban. Deepmath-deep sequence models for premise selection. In Advances in Neural Information Processing Systems, 2016.
Kshitij Bansal, Sarah M Loos, Markus N Rabe, and Christian Szegedy. Learning to reason in large theories without imitation. arXiv preprint arXiv: I9:e.:905/00. 10501, 2019a.
10

---

Published as a conference paper at ICLR 2020
Kshitij Bansal, Sarah M Loos, Markus N Rabe, Christian Szegedy, and Stewart Wilcox. HOList: An environment for machine learning of higher-order theorem proving. ICML 2019. International Conference on Machine Learning, 2019b.
Gino Brunner, Manuel Fritsche, Oliver Richter, and Roger Wattenhofer. Using state predictions for
value regularization in curiosity driven deep reinforcement learning. In 2018 IEEE 30th Interna- tional Conference on Tools with Artificial Intelligence (ICTAI), pp. 25-29. IEEE, 2018.
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A. Efros. Large-scale study of curiosity-driven learning. In ICLR, 2019.
Silvia Chiappa, Sébastien Racanière, Daan Wierstra, and Shakir Mohamed. Recurreetteenvrooomne simulators. CoRR, 011.1744..22252 2017. URL http::////aaaaaorr/r///////0.///
Coq. The Coq Proof Assistant, http::/ccoqinniaee fr UtLhttt::/ccq..ine
Alexey Dosovitskiy and Vladlen Koltun. Learning to act by predicting the future. ICLR 2017, 2017.
Thibault Gauthier, Cezary Kaliszyk, and Josef Urban. Tactictoe: Learning to reason with HOL4 tac- tics. In LPAR-221 21. 21st International Conference on Logic for Programming, Artificial Intelligence and Reasoning, volume 46, pp. 125-143, 2017.
Georges Gonthier. Formal proof-the four-color theorem. Notices of the AMS, 55(11):1382-139, 2008.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems 31. Curran Associates, Inc., 2018.
Thomas Hales, Mark Adams, Gertrud Bauer, Tat Dat Dang, John Harrison, Hoang Le Truong, Cezary Kaliszyk, Victor Magron, Sean McLaughlin, Tat Thang Nguyen, et al. A formal proof of the Kepler conjecture. In Forum of Mathematies, Pi, volume 5. Cambridge Univ. Press, 2017.
John Harrison. HOL. Light: A tutorial introduction. In FMCAD, pp. 265-269, 1996.
HOL Light Rewrite Tactic Reference. URL https:/ www. cl. cam. wu/.jr/i//.r////
TAC. html Accessed: 2019/09/23.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for Atari. arXiv preprint arXiv:/903.0m 000/99930007777 2019.
Cezary Kaliszyk and Josef Urban. Hol (y) hammer: Online atp service for hol light. Mathematics in Computer Science, 9():5-22, 2015.
Gerwin Klein, Kevin Elphinstone, Gernot Heiser, June Andronick, David Cock, Philip Derrin, Dhammika Elkaduwe, Kai Engelhardt, Rafal Kolanski, Michael Norrish, Thomas Sewell, Har- vey Tuch, and Simon Winwood. sel4: Formal verification of an os kernel. SOSP. ACM, 2009.
Gil Lederman, Markus N. Rabe, and Sanjit A. Seshia. Learning heuristics for automated reasoning through deep reinforcement learning. CoRR, abs/1807.0800 2018 URL http:://axxiv
org/ abs/ 1807. 08058
Sarah Loos, Geoffrey Irving, Christian Szegedy, and Cezary Kaliszyk. Deep network guided proof search. LPAR-21. s International Conference on Logic for Programming, Artificial Intelligence and Reasoning, 2017.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard I. Lewis, and Satinder Singh. Action-conditional video prediction using deep networks in Atari games. In Advances in neural information process- ing systems, pp. 2863-2871, 2015.
Aditya Paliwal, Sarah Loos, Markus Rabe, Kshitij Bansal, and Christian Szegedy. Graph represen- tations for higher-order logic and theorem proving. In AAAI, 2020.
Bartosz Piotrowski, Chad Brown, Josef Urban, and Cezary Kaliszyk. Can neural networks learn symbolic rewriting? 2019.
Mingzhe Wang, Yihe Tang, Jian Wang, and Jia Deng. Premise selection for theorem proving by deep graph embedding. In Advances in Neural Information Processing Systems, 2017.
11

---

Published as a conference paper at ICLR 2020
A VISUALIZATION OF LATENT SPACES
paireti
arth
25
20
25
Toor an
maters
15
Imagratton.m
Berooum
vacconum
10
reala Mumi
degres.nt mn
(b) Visualization of the embedding space >(2) evolving over rewrite steps. Theorems are colored by area mathematics they originate from, sourced from the the- orem database. The brightness of the embeddings corre- sponds to the rewrite step in which it was generated, with more recent embeddings being darker.
(a) Visualization of the embedding spaces produced by embedding prediction in the latent space after 2, and rewrite steps. The points are colored by their (i-distance to the true embedding (T').
Figure 11: iSNE visualization of embedding spaces.
In addition to numerically evaluating the performance of our latent embedding models, it is in- teresting to visually examine the embedding spaces produced. Figure Tla depicts the predicted embeddings the datasets D1 ... D4 by using PCA to reduce the dimension of each point in the joint dataset to 50, then using tSNE produce 2D visualization which is filtered by original step. The points are colored by their 1 !2-distance to the true embedding. From this graph it can be seen that the distribution of embeddings shifts meaningfully over multiple rewrites, further demonstrating that the rewrite task produces non-trivial changes.
We also wished compare the embeddings produced by our model to human understanding of mathematics, In the HOList database, each theorem comes with tags describing the area of mathe- matics originates from, such as topology set theory. In Figure [ we separate the embeddings of theorems from these categories by color, and examine the their evolution across multiple rewrites. can be seen that theorems originating from similar human categories correspond to similar em- beddings in the latent space.
12


Reviews:
[{"review_id": "rJxGlK6Ttr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a technique to perform reasoning on mathematical formulas in a latent space. The model is trained to predict whether a rewrite rule can be applied to a formula given its latent representation. When the rewrite is possible, the model also predicts the embedding of the resulting formula. Experiments show that the network can be applied multiple steps in a row, while operating only in the embedding space.\n\n1. As mentioned in the paragraph before Section 4.1, it would be much simpler to consider a single latent embedding space L. In that case, \\sigma and \\alpha become unnecessary and we only need to train \\omega. Did you try to have a single network? This seems a much more natural approach to me, and I'm surprised that you did not start with that. From my experience, aligning embedding spaces is something that usually does not work very well, especially in high dimension. The role of \\sigma seems very redundant given \\omega.\n\n2. If you consider \\sigma, why do you also predict the rewrite success with \\omega? Couldn't it be simply a function from S x S -> L ?\n\n3. The graph neural networks used in the model are not described in the paper, only a reference to Paliwal et al (2019) is given. It would be helpful to have a brief paragraph describing this architecture, for readers not familiar with the referenced paper.\n\n4. How large is the training set of (T, P) pairs? I don't think this is mentioned in the paper.\n\n5. To train \\sigma and \\omega, the negative instances are selected randomly. You mention that negative mining should improve over this strategy. What does negative mining correspond to in this context? Are there bad rewrites better than others?\n\n6. Did you consider using an inverse function (say G), that maps an embedding in L / L' back to S (i.e. the inverse function of gamma / gamma'). I would imagine that even if an embedding X is a bit noisy, because not exactly equal to gamma(P) where P is the expression it represents, you could consider doing the propagation with gamma(G(X)). This could be a possibility to remove the noise you have when doing multi-step operations (and potentially go way beyond 4 steps). Also, G could be used to check whether you obtain the expected formula after 4 steps, which would be a more informative information than the L2 distance between the resulting embedding and the embedding of the final formula.\n\nOverall, the model is a bit complicated (e.g. question 1.), but the results are promising, the paper is well written, and the ability to manipulate formula embeddings is probably going to be useful in the context of theorem proving.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HJl5cVvoYS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "= Summary\nEmbeddings of mathematical theorems and rewrite rules are presented. An in-depth analysis of the resulting embeddings is presented, showing that a network can learn to \"apply\" embedded rewrite rules to embedded theorems, yielding results that are similar to the embedding of the rewritten theorem. [i.e., app'(emb(thm), emb(rule)) is near to emb(app(thm, rule))] This is an interesting property for the application of deep learning to automated theorem proving, though not directly a breakthrough result.\n\n= Strong/Weak Points\n+ Simply a cute result, showing that proof search can remain in embedding space for a limited time horizon without having to switch back into the theorem prover environment.\n+ Nicely designed experiments testing this (somewhat surprising) property empirically\n- Missed opportunity of better analysis of which theorem/rewrite rule properties are more likely to fail\n- Writing sometimes a bit overcomplicated (e.g., Sect. 4.5 could just be a figure of a commuting diagram and two sentences...)\n- Architecture choice unclear: Why are $\\sigma$ and $\\omega$ separate networks. This is discussed on p4, but it's unclear to me how keeping $\\sigma$ separate is benefitial for the analysis, and this is not picked up again explicitly again?\n\n= Recommendation\nOverall, this is a nice, somewhat surprising result. The writing and experiments could use some improvement, but I believe that the majority of the ICLR audience would enjoy seeing this result (even though it would have no impact on most people's research)\n\n= Detailed Comments\n- page 4, Sect. 4.4: Architecture of $\\alpha$ would be nice (more than a linear layer?)\n- page 5, paragraph 3: \"we from some\" -> \"we start from some\"\n- p6par1: \"much cheaper then computing\" -> than\n- p6par6: \"on formulas that with\" -> no that\n- p6par7: \"measure how rate\" -> \"measure the rate\"\n- p8par1: \"approximate embedding $\\alpha(e(\\gamma'(...)))$ - $e$ is undefined and should probably be $e'$ (this is also the case in the caption of Fig. 5), and $c'$ should probably be included as well. However, I don't understand the use of $\\alpha$ here. If Fig. 4 is following Fig. 3 in considering $p(c(\\gamma(T), \\pi(P)))$, then Fig. 4 should plot the performance of, e.g., $p(c(e'(c'(\\gamma'(T_{i-1}), \\pi'(P_{i-1}))), \\pi(P_i)))$ (i.e., $p$ applied to approximate embedding of $T_i$ and (\"true\") embedding of $P_i$). I believe that's what \"Pred (One Step)\" expresses, but it would maybe be generally helpful to be more precise about the notation in Sect. 6.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "Bkx-GQrhtS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a method to do math reasoning purely using formula embeddings. The proposed method employs a graph neural network to embed math formulas to a latent space. The formula embeddings are then combined with theorem embeddings (also formulas, computed in the same way as formula embeddings) to predict whether one can do one step of math reasoning using the corresponding theorem, and also to predict the embeddings of the resulting formula. Empirically the authors demonstrate that the method can be chained end-to-end to do multiple steps of reasoning purely in the latent space.\n\nI tend to accept this paper, (but also OK if it gets rejected), for the following reasons: (1) the idea is novel and interesting; (2) the writing of the paper is below conference standard and very hard to read, especially the method and the experiment sections.\n\n===========================================================================\n\nNovelty and significance\n\nI really like the idea of doing math reasoning in latent space. The idea is definitely novel and interesting. It is related to existing works such as neural logic induction[1] and planning in latent space[2]. It is amazing that one can do multiple steps of math reasoning after only training the model using data from one single step. It would be interesting to see how it can improve existing learning-based theorem provers.\n\nMy question is if we want to integrate the proposed method into theorem provers, after multiple steps of math reasoning, how would us know the goal has been proved? Is it possible that we can train a decoder that maps back from the latent space to the formula space? Also can it work with theorems that decompose the current goal into several sub-goals? I know these are not the concerns of this paper, but I would be really grateful if you could provide some intuitive answers!\n\n===========================================================================\n\nWriting\n\nThe paper is not well-organized and not written in a consistent way. For the method and the experiment sections, I need to jump back and forth several times in order to understand what the authors are trying to say.\n\n1. Typo: Third paragraph in section 1, \"...which is makes use of ...\".\n2. It's very confusing when the authors introduce \\sigma and \\omega in the beginning of section 4: why would you need two networks predict the same thing?\n3. Mentioning \"merging \\sigma and \\omega, is left for future work\" is confusing before formally introducing \\sigma and \\omega.\n4. Even when the authors formally introduce \\sigma and \\omega in 4.2, it is still not clear that why both of them are used for modelling the success probability.\n5. In fact, I don't know why \\omega needs to output p. It's never mentioned in the experiment section.\n6. The rationale of the two tower design (why not combine two) is not clearly explained.\n7. Typo: Page 5 last paragraph, \"... negative instances for for each ...\".\n8. The itemized part in 5.3, \"...carefully selected baselines: 1.xxx, 2.xxx, 3. xxx, 4. xxx\". However, both 3 and 4 are not baselines!\n9. It is not clear that baseline 1 and 2 correspond to which baselines in later experiments.\n10. Reading the baselines before the experiments is very confusing. For example, for baseline 1, it is very hard to understand why would we want to use such an unusual baseline, and why it is called a \"random baseline\".\n11. Baseline 2 is actually referred to as \"usage baseline\" but this name is not introduced in the itemized part.\n\n\n\n[1] Rocktäschel, Tim, and Sebastian Riedel. \"End-to-end differentiable proving.\" Advances in Neural Information Processing Systems. 2017.\n[2] Srinivas, Aravind, et al. \"Universal planning networks.\" arXiv preprint arXiv:1804.00645 (2018).\n\n\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper was very well received by the reviewers with solid Accept ratings across the board.
The subject matter is quite interesting -  mathematical reasoning in latent space, and it was suggested by a reviewer that this could be a good candidate for an oral. The AC agrees and recommends acceptance as an oral. Some of the intuitions of what is being done in this paper could be better visualized and presented and I encourage the authors to think carefully about how to present this work if an oral presentation is granted by the PCs.
Decision: Accept (Talk)