Learning Axioms to Compute Verifiable Symbolic Expression Equivalence Proofs Using Graph-to-Sequence Networks

We target the problem of proving semantic equivalence between two complex expressions represented as typed trees , and demonstrate our system on expressions from a rich multi-type symbolic language for linear algebra . We propose the first graph-to-sequence deep learning system to generate axiomatic proofs of equivalence between program pairs . We generate expressions which include scalars , vectors and matrices and 16 distinct operators combining them , with 147 distinct axioms of equivalence . We study the robustness of the system to generate proofs of increasing length , demonstrating how incremental graph-to-sequence networks can learn to represent complex and verifiable symbolic reasoning . It achieves 93 % average true positive coverage on 10,000 test cases while ensuring zero false positives by design . 1 INTRODUCTION . Deep neural networks have excelled at a variety of classification and reinforcement learning tasks ( Goodfellow et al. , 2016 ) . However , their stochastic nature tends to hinder their ability to learn representations useful for manipulating symbolic information . For applications requiring a guarantee of correctness , such as occur in relation to determining the semantic equivalence between two programs ( or symbolic expressions ) , a system that produces provably correct output must be developed . In this work we target the problem of automatically computing whether two input programs are semantically equivalent ( Kaplan , 1969 ) , under a well-defined axiomatic system for equivalence using semantics-preserving rewrite rules ( Dershowitz , 1985 ) . Program equivalence is summarized as determining whether two programs would always produce the same outputs for all possible inputs , and is a central problem in computing ( Kaplan , 1969 ; Godlin & Strichman , 2008 ; Verdoolaege et al. , 2009 ) . The problem ranges from undecidable , e.g . Goldblatt & Jackson ( 2012 ) , to trivial in cases of testing the equivalence of a program with itself . Our work directly studies the subset of programs represented by symbolic linear algebra expressions which include scalar , vector , and matrix types for both constants and variables , and 16 different operators with 147 distinct axioms of equivalence , creating a foundation for future work with broader program semantics . For example , the expression using matrices , scalars , and a vector : ( A+B ) I ( ( a+ ( b− b ) ) /a ) ~v−A~v can be proven equivalent to B~v by applying 10 axioms in sequence ; our work generates the proof steps between these expressions . While prior work has shown promise for deep networks to compute some forms of program equivalence ( Xu et al. , 2017 ; Alon et al. , 2019 ; Rabin et al. , 2020 ) , the system typically outputs only a probability of equivalence , without any reasoning or insight that can be verified easily : false positives can be produced . Programs can be represented as a tree ( or graph ) of symbols , and deep networks for symbolic reasoning have been studied , e.g . to compute the derivative of a symbolic expression ( Lample & Charton , 2020 ) , but such systems can produce output which is not directly explained by the system itself . In this work , we take a fundamentally different approach to the problem of symbolic program reasoning with deep networks : we make the system produce the sequence of steps that lead to rewriting one program into another , that is the reasoning for ( or proof of ) equivalence between the two programs , instead of producing directly the result of this reasoning ( e.g. , the transformed expression ( Lample & Charton , 2020 ) ) . We propose a method for generating training samples using probabilistic applications of production rules within a formal grammar , and then develop a graph-to-sequence ( Li et al. , 2016 ; Beck et al. , 2018 ) neural network system for program equivalence , trained to learn and combine rewrite rules to rewrite one program into another . It can deterministically prove equivalence , entirely avoids false positives , and quickly invalidates incorrect answers produced by the network ( no deterministic answer is provided in this case , only a probability of non-equivalence ) . In a nutshell , we develop the first graph-to-sequence neural network system to accelerate the search in the space of possible combinations of transformation rules ( i.e. , axioms of equivalence in the input language ) to make two graphs representing programs structurally identical without violating their original semantics . We make the following contributions : ( i ) We propose a machine learning system for program equivalence which ensures correctness for all non-equivalent programs input , and a deterministically checkable output for equivalent programs ( no false positives , specificity = 100 % ) . ( ii ) We introduce pe-graph2axiom , the first incremental graph-to-sequence neural network system targeting program equivalence to the best of our knowledge . We provide the first implementation of such graph-to-sequence systems in the popular OpenNMT-py framework ( Klein et al. , 2017 ) . ( iii ) We present a complete implementation of our system operating on a rich language for multitype linear algebra expressions . Our system provides a correct rewrite rule sequence between two equivalent programs for 93 % of the 10,000 test cases . The correctness of the rewrite rule is deterministically checkable in all cases in negligible time . 2 PROGRAM EQUIVALENCE USING DEEP LEARNING . We outline below the key intuitions and concepts for our program equivalence system . Extensive details , including a full formal framework for program equivalence reasoning , are presented in supplementary material , Secs . A-B . We illustrate below four very simple computations , represented as graphs , that are all equivalent under various axioms of natural arithmetic . ∗ ∗ 1 a + b c P1 : ( 1 ∗ a ) ∗ ( b+ c ) ≡ ∗ a + b c P2 : a ∗ ( b+ c ) ≡ + ∗ a b ∗ a c P3 : a ∗ b+ a ∗ c ≡ + ∗ a c ∗ a b P4 : a ∗ c+ a ∗ b Rewrite rules as axioms of equivalence In this work we represent programs with symbolic expressions made of variables ( e.g. , a , b , c ) , operators ( e.g. , + , * ) and neutral/absorbing elements ( e.g. , 1 ) . We consider a rich linear algebra expression language , supporting three variable types ( scalars as shown in P1-P4 , vectors , and matrices ) and 5 different variables per type ; 16 operators including operators mixing different variable types such as vector-matrix product . We represent these programs as dataflow graphs Buck & Lee ( 1993 ) with a root node that is to compute a single value . P1 is equivalent to P2 if we consider the axiom A1 : 1N ∗ x = x , ∀x ∈ N. This axiom is also a clear rewrite rule : the LHS expression 1N ∗ x ( with x ∈ N ) can be matched and replaced by the RHS expression x anywhere in the program without altering its semantics . An axiom , or equivalently here a graph rewrite rule , may be applied repeatedly to different subtrees . When applying A1 on a specific location , the node a of P1 , we obtain an equivalent and yet syntactically different program , we note P1 ≡ A1 ( a , P1 ) . To assess the validity of a transformation sequence S where P2 = S ( P1 ) , one simply needs to check for S , in sequence , that each axiom is applicable at that program point , apply it to obtain a new temporary program , and repeat the process for each axiom in the complete sequence . If the sequence is verified to be valid , and S ( P1 ) is structurally equivalent to P2 , then we have proved P1 ≡ P2 , and S forms the complete proof of equivalence between the two programs . These equivalences can be composed , incrementally , to form a complex transformation . Using A2 : x ∗ ( y + z ) = x ∗ y + x ∗ z , ∀x , y , z ∈ N and A3 : x + y = y + x , ∀x , y ∈ N , we have P1 ≡ P4 ≡ A3 ( + , A2 ( ∗ , A1 ( a , P1 ) ) ) , a verifiable proof of equivalence under our axioms between the programs 1a ( b + c ) and ac + ab , which involved structural changes including node deletion , creation and edge modification . Note the non-unicity of a sequence : by possibly many ways a program can be rewritten into another one , for example the sequence P4 ≡ A2 ( ∗ , A3 ( + , A1 ( a , P1 ) ) ) also correctly rewrites P1 into P4 . Conversely , a sequence may not exist : for example no sequence of the 3 above axioms allows rewriting a+ b into a ∗ b . We call these non-equivalent in our system , that is precisely if there is no sequence of axioms that can be applied to rewrite one program into the other . Our approach aims to compute some S for a pair of programs P1 , P2 , so that S is verified correct when P1 ≡ P2 . Consequently , if P1 6≡ P2 , no sequence S produced can be verified correct : true negatives are trivially detected . Pathfinding program equivalence proofs Intuitively , we can view the solution space as a graph , where every possible syntactically different program in the language is represented by its own vertex vi . And ∃ e ( Ak , x ) : vi → vj iff ∃Ak an axiom and x ∈ vi such that vj = Ak ( x , vi ) . Any two programs connected by a path in this graph are therefore semantically equivalent . Building S for P1 ≡ S ( P2 ) amounts to exposing one path between P1 and P2 in this graph when it exists , the path forming the proof of equivalence . We build a deep learning graph-to-sequence system to learn a stochastic approximation of an iterative algorithm to construct such feasible path when possible , trained only by randomly sampling pairs of programs and one carefully labeled path between them . This avoids the need to craft smart exploration heuristics to make this path-finding problem practical . Graph-to-sequence network for pathfinding The system in Fig . 1 takes as input two programs represented as symbolic trees , and produces a sequence of axioms along with their position of application ( or node ) that can be used to rewrite sequentially one input program into the other input program . Node initialization initializes the graph neural network , converting the input programs text ( e.g. , ( a + ( b + c ) ) into nodes and edges in the Graph Neural Network Scarselli et al . ( 2009 ) ; Xu et al . ( 2017 ) . The details of the network are covered in Sec . 5 . In a nutshell , the key principle is to combine a memory-based neural network approach , e.g. , using Long-Short Term Memory ( LSTM ) Hochreiter & Schmidhuber ( 1997 ) neurons and a graph neural network design ( which uses Gated Recurrent Units ( GRUs ) internally ) Beck et al . ( 2018 ) that matches our program graph representation . Token embedding is a neural network layer in which tokens are assigned a learnable multidimensional embedding vector Mikolov et al . ( 2013 ) . Each layer in LSTM 2 layers has 256 neurons , which support sequence generation . Token generator is the final output portion of the network . It learns to output the tokens based on the current LSTM hidden states and the Global Attention from the graph neural network . As each token is output , it feeds back into the LSTM layer through the embedding layer to affect its next state . We use a sequence generation principle , using a global attention mechanism Luong et al . ( 2015 ) to allow observation of program graph node information while generating the axiom and location on which it is applied . We specifically study the robustness of our approach to generate proofs of increasingly complex length , contrasting models to output the entire path at once with pe-graph2axiom which incrementally builds the sequence one step at a time in Sec . 6 .

Summary: The authors introduce a new synthetic dataset of equational proofs over the basic axioms of linear algebra.  The dataset consists of triples `(t1, t2, rewrites)` where `rewrites` is a sequence of rewrite instructions that can transform `t1` into `t2`, though some of their models emit one rewrite at a time and observe the result of applying the rewrite instruction to `t1`.  They develop a GNN for the task called pe-graph2axiom and show that it beats two baselines.  Finally, they show that their trained system can solve 15 problems from two Khan Academy modules that can be expressed in their fragment.  The appendix refers to supplementary material including code and data, but no supplementary material was submitted.