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: This paper proposes a model for verifying semantic equivalence  between symbolic linear algebra expressions. Expressions are represented by trees and equivalence is proven by a sequence of axioms applied to the first expression. The proposed model encodes the expression/program trees as nodes on a graph connected by edges representing one of a set of axioms being applied to one of the elements in the first expression to yield a node in the second expression. The output of the model is a path, a sequence of edges, on this constructed graph that correspond to a sequence of axioms applied to the first expression to arrive at the second.