Do Transformers Understand Polynomial Simplification? 

Recently researchers have demonstrated that Transformers can be trained to learn symbolic tasks such as solving integration and differential equations in an end-toend fashion . In these setups , for an input symbolic expression , the Transformer predicts the final solution in a single step . Since such tasks may consist of a sequence of logical steps , question remains whether such networks have understood and learnt individual steps to reach the solution . To take a deeper look , we consider the task of polynomial simplification . Polynomials can be written in a simple normal form as a sum of monomials which are ordered in a lexicographic order . For a polynomial which is not necessarily in this normal form , a sequence of simplification steps is applied to reach the fully simplified ( i.e. , in the normal form ) polynomial . For this task , we describe a synthetic Polynomial dataset generation algorithm which generates polynomials with unique proof steps . Then , we conduct an extensive analysis of the Transformer ’ s abilities to learn the polynomial simplification task along different dimensions . 1 INTRODUCTION . With the state-of-the-art performance of Deep Neural Nets ( DNNs ) in perceptual tasks , researchers have started to explore their logical reasoning capabilities , in particular within the domain of Automated Theorem Proving ( ATP ) . In these domains ( LEAN ( de Moura et al. , 2015 ) , HOL Light and Mizar ( miz , 2020 ) ) , many recent works ( Paliwal et al. , 2020 ; Aygün et al. , 2020 ; Hahn et al. , 2020 ) have shown that Graph Neural Networks ( Gori et al. , 2005 ; Veličković et al. , 2018 ) and Transformers ( Vaswani et al. , 2017 ) can be trained to perform impressively on the theorem-proving task as part of a neuro-symbolic system . In a related but different development , recently Lample & Charton ( 2019 ) showed that for symbolic integration and differential equations , a large amount of synthetic end-to-end examples can be generated using symbolic systems . In these tasks , the authors show that Transformer networks can be trained to produce the final solution from an input integral ( or differential equation ) in a single step . This points to the exciting possibility of using deep neural nets to learn end-to-end theorem provers , and can be beneficial for formal mathematics ( Szegedy , 2020 ) . However , the setup combines multiple reasoning steps in a single shot . Additionally , integration ( or differential equation solving ) is a complex task requiring understanding of the integral symbols , functions , variables , and the basic concepts of arithmetic . As the system in Lample & Charton ( 2019 ) is simply trained to output the top solution ( s ) and a corresponding confidence score ( s ) , it is unclear what internal mechanisms enable these models to solve these problems . This lack of transparency has been noted in this context ( Davis , 2019 ) . An earlier work by Piotrowski et al . ( 2019 ) showed similar results for certain symbolic manipulation tasks and their work shares the same limitation . In this paper we ask if instead of only producing the end-result of symbolic manipulation or integral , can we have the model produce a human-readable proof as well . While we do not know if these models reason in the way humans do , one way to produce proofs would be to “ extract ” a proof from the models of the above type by “ probing ” them in some mannner . The problem of unraveling the inner workings of Transformers by probing is an active area of research ; however , at present our understanding is still evolving ( Rogers et al. , 2020 ) . Hence taking a detour , we instead train the model to produce the full proof . Inspired by Piotrowski et al . ( 2019 ) , we explore a novel but simpler setting of polynomial simplification . We illustrate the task with an example . We begin with a polynomial which is a sum of product of factors , where each factor is again a sum of monomials ( including constants ) , as shown below : P0 = ( 2 ∗ x22 ) ∗ factor︷ ︸︸ ︷ ( 3 ∗ x12︸ ︷︷ ︸ term +4 ) + product︷ ︸︸ ︷ ( 5 ∗ x21 + x11 ∗ x12 ) ∗ ( 3 ∗ x11 ) ∗ ( 2 ) , / * Initial * / To construct unique simplification steps , first each term in a factor is simplified . Once all factors are simplified ( facstep ) ; then within a product , all factors are multiplied ( mulstep ) . Lastly simplified products are summed ( sumstep ) . P0 = ( 2 ∗ x22 ) ∗ ( 3 ∗ x12 + 4 ) + ( 5 ∗ x21 + x11 ∗ x12 ) ∗ ( 3 ∗ x11 ) ∗ ( 2 ) , / * FACSTEP * / = ( 2 ∗ x22 ) ∗ ( 3 ∗ x2 + 4 ) + ( 5 ∗ x21 + x11 ∗ x12 ) ∗ ( 3 ∗ x11 ) ∗ ( 2 ) , ( P1 ) , / * FACSTEP * / = ( 2 ∗ x22 ) ∗ ( 3 ∗ x2 + 4 ) + ( 5 ∗ x21 + x1 ∗ x2 ) ∗ ( 3 ∗ x1 ) ∗ ( 2 ) , ( P2 ) , / * MULSTEP * / = ( 6 ∗ x32 + 8 ∗ x22 ) + ( 5 ∗ x21 + x1 ∗ x2 ) ∗ ( 3 ∗ x1 ) ∗ ( 2 ) , ( P3 ) , / * MULSTEP * / = ( 6 ∗ x32 + 8 ∗ x22 ) + ( 30 ∗ x31 + 6 ∗ x21 ∗ x2 ) , ( P4 ) , / * SUMSTEP * / = 30 ∗ x31 + 6 ∗ x32 + 6 ∗ x21 ∗ x2 + 8 ∗ x22 . ( P5 ) , / * ENDPOINT * / . Piotrowski et al . ( 2019 ) explores the task of learning symbolic re-write of an entire expression . In contrast , in our setting , for step-wise prediction , at each step the system needs to find the candidate sub-expression and a relevant simplification type to perform the simplification . This setup resembles the traditional ATP setup where a system needs to learn and execute symbolic steps to reach a final solution . But it is simpler as for each step only one type of simplification is applicable . By proof for an initial polynomial ( P0 ) we mean the sequence of simplification steps ( P1 to P5 ) . A model trained on step-wise prediction task , can be used to generate a full proof . Essentially , we start with an initial polynomial , and recursively feed the model output to itself , till it generates the final simplified polynomial ( in normal form ) . A proof is correct when all steps are correct . In the above setting ( termed COARSE ) , all terms in a factor are simplified at once in a facstep , and similarly all factors in a product are simplified at once in a mulstep . Additionally , we define another setting FINER , where a facstep involves simplification of a single term , and a mulstep involves multiplications of only two factors at once , illustrated below with an example ( for facstep ) : P0 = ( 5 ∗ x21 + x11 ∗ x12 ) ∗ ( 3 ∗ x11 ) ∗ ( 2 ) , / * FACSTEP * / = ( 5 ∗ x21 + x1 ∗ x2 ) ∗ ( 3 ∗ x11 ) ∗ ( 2 ) , / * FACSTEP * / = ( 5 ∗ x21 + x1 ∗ x2 ) ∗ ( 3 ∗ x1 ) ∗ ( 2 ) . As a state-of-the-art model , we explore Transformers . While both Graph Neural Networks and Transformers have been used for single-step representation learning of symbolic theorems and single step goal-theorem scoring , Transformer-based sequence-to-sequence networks have shown superiority in end-to-end tasks in integration , differential equations ( Lample & Charton , 2019 ) and temporal logic ( Hahn et al. , 2020 ) domains . Hence for the aforementioned tasks of step-wise polynomial simplification , we explore the Transformer ’ s ability along several dimensions . Our contributions are the following : 1 ) we propose polynomial simplification tasks requiring multiple steps of symbolic manipulation , 2 ) we show how datasets of different configurations can be generated synthetically for the task , 3 ) we propose an array of metrics to dissect the performance of Transformers , and 4 ) lastly through extensive experiments we show the performance of the Transformer on this task , establishing a strong baseline for future endeavors . Results Summary By varying over coefficient size , proof granularity and input representation ( in Tables 1 , 2 , Appendix Table 6 ) we observe that 1 ) full proof accuracy is only slightly lower than single-shot endpoint prediction accuracy in many 1-variable configurations , 2 ) coarse granular proofs help learn somewhat more accurate proofs , 3 ) prefix representation helps in most cases but infix sometimes provides higher accuracy . More than 80 % errors ( Tab . 7 and 8 in Appendix ) occur in multiplication steps , and we observe ( through independent experiments ) Transformer ’ s struggle to learn how to multiply numeric coefficients . By letting the system annotate the candidate subexpression , we observe that the system can understand candidate sub-expressions and which next step to perform explicitly ( Tables 3 , and Appendix Tables 9 , 10 , 11 ) . Also , through visualization we observe similar effects ( Figures 1 , 2 Appendix ) . We see systems trained for 2-variable outperform corresponding 1-variable systems on 1-variable test sets . For 1 variable , we observe steady and significant higher gains ( till 10 % for full proof ) using curriculum learning ( Table 17 Appendix ) . 2 RELATED WORK AND DISCUSSION . Unlike the problems dealt with by the aforementioned automatic theorem provers and related neuralbased systems , polynomial simplification does not involve any search . Our problem is simpler and tests a specific ability , namely certain kinds of symbol manipulations . This simplicity affords certain advantages ( shared by Piotrowski et al . ( 2019 ) and Lample & Charton ( 2019 ) ) : ( 1 ) We can generate artificial data to train models without limitations on the size . ( 2 ) It is easier to test the abilities of the models more thoroughly along multiple axes . ( 3 ) Accuracy achieved is much higher than for harder tasks , suggesting that fully solving such tasks may be possible in the near future . To compare with symbolic manipulation systems we note that in more detail the ability tested by our task is the following : the model must be able to identify smallest parts of the polynomial that can be simplified : ( 1 ) simplification of a factor , ( 2 ) multiplication of two factors , ( 3 ) addition of two sub-polynomials . Having identified what simplification to apply the model must produce a new polynomial with just that simplification . This ability is not tested by previous neural-based symbolic manipulation systems such as Piotrowski et al . ( 2019 ) and Lample & Charton ( 2019 ) and related works such as Saxton et al . ( 2019 ) and Hahn et al . ( 2020 ) . Several recent works have produced synthetic datasets for theorem proving tasks ( Aygün et al. , 2020 ; Wu et al. , 2020 ; Polu & Sutskever , 2020 ) , however , their focus remains more on search-based proofs .

Summary: The authors analyze the performance of Transformer models on simplifying polynomials and - importantly - generating proofs at the same time. This is a very nice idea that allows to study the performance of Transformers in depth and at the same time in an important setting where verification is performed as part of running the model. And the authors show a strong baseline, with models performing very well in a number of settings. A few areas seem to have been neglected though. For one, the authors only train a 4-layer 4-head model, which is quite small as far as Transformers go. Maybe it's irrelevant for this problem - but having at least one bigger model as a point of comparison would be good. Next, the out-of-distribution question warrants more experiments. Can the Transformers simplify polynomials with way more factors than trained on? With a higher number of variables? Higher degrees? The authors also show that one main problem for Transformers is learning to multiply the coefficients. But - assuming this reviewer understood correctly - the authors do not apply the proof requirement to multiplication. E.g., for "12*3" the model has to immediately output "36" rather than "10*3 + 2*3 = 30 + 6 = 36". Maybe this could help the Transformer learn and be more resilient to coefficient size? So while the current version of the paper is ok, there are a few areas for improvement which prevent it from being a clear accept.