Deep Symbolic Superoptimization Without Human Knowledge

1 INTRODUCTION . Superoptimization refers to the task of simplifying and optimizing over a set of machine instructions , or code ( Massalin , 1987 ; Schkufza et al. , 2013 ) , which is a fundamental problem in computer science . As an important direction in superoptimization , symbolic expression simplification , or symbolic superoptimization , aims at transforming symbolic expression to a simpler form in an effective way , so as to avoid unnecessary computations . Symbolic superoptimization is an important component in compilers , e.g . LLVM and Halide , and it also has a wide application in mathematical engines including Wolfram2 , Matlab , and Sympy . Over recent years , applying deep learning methods to address symbolic superoptimization has attracted great attention . Despite their variety , existing algorithms can be roughly divided into two categories . The first category is supervised learning , i.e . to learn a mapping between the input expressions and the output simplified expressions from a large number of human-constructed expression pairs ( Arabshahi et al. , 2018 ; Zaremba & Sutskever , 2014 ) . Such methods rely heavily on a human-constructed dataset , which is time- and labor-consuming . What is worse , such systems are highly susceptible to bias , because it is generally very hard to define a minimum and comprehensive axiom set for training . It is highly possible that some forms of equivalence are not covered in the training set , and fail to be recognized by the model . In order to remove the dependence on human annotations , the second category of methods leverages reinforcement learning to autonomously discover simplifying equivalence ( Chen et al. , 2018 ) . However , to make the action space tractable , such systems still rely on a set of equivalent transformation actions defined by human beings , which again suffers from the labeling cost and learning bias . In short , the existing neural symbolic superoptimization algorithms all require human input to define equivalences . It would have benefited from improved efficiency and better simplification if there were algorithms independent of human knowledge . In fact , symbolic superoptimization should have been a task that naturally keeps human outside the loop , because it directly operates on machine code , whose consumers and evaluators are machines , not humans . ∗Authors contributed equally to this paper . 1The code is available at https : //github.com/shihui2010/symbolic_simplifier . 2https : //www.wolframalpha.com/ Therefore , we propose Human-Independent Symbolic Superoptimization ( HISS ) , a reinforcement learning framework for symbolic superoptimization that is completely independent of human knowledge . Instead of using human-defined equivalence , HISS adopts a set of unsupervised techniques to maintain the tractability of action space . First , HISS introduces a tree-LSTM encoder-decoder architecture with attention to ensure that its exploration is confined within the set syntactically correct expressions . Second , the process of generating a simplified expression is broken into two stages . The first stage selects a sub-expression that can be simplified and the second stage simplifies the sub-expression . We performed a set of evaluations on artificially generated expressions as well as a publicly available code dataset , called the Halide dataset ( Chen & Tian , 2018 ) , and show that HISS can achieve competitive performance . We also find out that HISS can automatically discover simplification rules that are not included in the human predefined rules in the Halide dataset . 2 RELATED WORK . Superoptimization origins from 1987 with the first design of Massalin ( 1987 ) . With the probabilistic testing to reduce the testing cost , the brute force searching is aided with a pruning strategy to avoid searching sub-spaces that contains pieces of code that have known shorter alternatives . Due to the explosive searching space for exhaustive searching , the capability of the first superoptimizer is limited to only very short programs . More than a decades later , Joshi et al . ( 2002 ) presented Denali , which splits the superoptimization problem into two phases to expand the capability to optimize longer programs . STOKE ( Schkufza et al. , 2013 ) follows the two phases but sacrifices the completeness for efficiency in the second phase . Recent attempts to improve superoptimization are categorized into two fields : exploring transformation rules and accelerating trajectory searching . Searching the rules are similar to the problem of superoptimization on limited size program , but targeting more on the comprehensiveness of the rules . Buchwald ( 2015 ) exhaustively enumerates all possible expressions given the syntax and checks the equivalence of pairs of expressions by SMT solver . A similar method with an adaption of the SMT solver to reuse the previous result is proposed by Jangda & Yorsh ( 2017 ) . On the other hand , deep neural networks are trained to guide the trajectory searching ( Cai et al. , 2018 ; Chen & Tian , 2018 ) . Considering transformation rule discovery as a limited space superoptimization , the large action space and sparse reward are the main challenges for using neural networks . Special neural generator structures are proposed for decoding valid symbolic programs , which leverage the syntax constraints to reduce the searching space as well as learn the reasoning of operations , and are gaining popularity in program synthesis ( Parisotto et al. , 2016 ; Zhong et al. , 2017 ; Bunel et al. , 2018 ) , program translation ( Chen et al. , 2018 ; Drissi et al. , 2018 ) , and other code generation tasks ( Ling et al. , 2016 ; Alvarez-Melis & Jaakkola , 2016 ) . Among the symbolic expression decoders , the family of tree structure RNNs ( Parisotto et al. , 2016 ; Drissi et al. , 2018 ; Alvarez-Melis & Jaakkola , 2016 ; Chen et al. , 2018 ) are more flexible than template-based predictors ( Ling et al. , 2016 ; Zhong et al. , 2017 ) . 3 THE HISS ARCHITECTURE . In this section , we will detail our proposed HISS architecture . We will first introduce a few notations . T denotes a tree ; a denotes a vector , and A denotes a matrix . We introduce an LSTM ( · ) function that summarizes standard one-step LSTM operation as [ ht , ct ] = LSTM ( xt , ht−1 , ct−1 ) , ( 1 ) where ht , ct and xt denote the output , cell and input at time t of a standard LSTM respectively . 3.1 FRAMEWORK OVERVIEW . Our problem can be formulated as follows . Given a symbolic expression TI , represented in the expression tree form , our goal is to find a simplified expression TO , such that 1 ) the two expressions are equivalent , and 2 ) TO contains a smaller number of nodes than TI . It is important to write the symbolic expressions in their expression tree form , rather than strings because HISS will be operating on tree structures . An expression tree assigns a node for each operation or variable . Each non-terminal node represents an operation , and each terminal node , or leaf node , LSTM1 LSTM2 LP LP Att LP LSTMout Softmax ⋮ ⋮ LSTM LP LSTM LSTM LP Sigmoid LP Softmax⋯ ⋯ ⋯ ⋯ ⋯ ⋯ ⋮ ⋮ ⋯ ⋮ ⋮ ⋯ # $ # % # & # ' ( % ( ' ( & ) % →+ , - . / % →+ , - . ) % →0 . / % →0 . ) % . / % . ) & . / & . ) ' . / ' . ) % → & . / % → & . ) ' /' ) & / & ) % / % { ) % } ) % →+ , - . Prob . of choosing sub-tree ) % D % Copy & Split Concatenate LSTM Output & Cell Input Output Previous Output & Cell LP Linear Projection ( Feed Forward ) Att Attention Module ( Eq . ( 6 ) ) ( a ) The tree encoder ( green ) and subtree selector ( orange ) . LSTM1 . LSTM2 LP LP Att LP LSTMout Softmax ⋮ ⋮ LSTM LP LSTM LSTM LP Sigmoid LP Softmax⋯ ⋯ ⋯ ⋯ ⋯ ⋯ ⋮ ⋮ ⋯ ⋮ ⋮ ⋯ # $ # % # & # ' ( % ( ' ( & ) % →+ , - . / % →+ , - . ) % →0 . / % →0 . ) % . / % . ) & . / & . ) ' . / ' . ) % → & . / % → & . ) ' /' ) & / & ) % / % { ) % } ) % →+ , - . Prob . of choosing sub-tree ) % D % Copy & S lit Concatenate LSTM Output & Cell Input Output Previous Output & Cell LP Linear Projection ( Feed Forward ) Att Attention Module ( Eq . ( 6 ) ) ( b ) The tree decoder . LSTM1 LSTM2 LP LP Att LP LSTMout Softmax ⋮ ⋮ LSTM LP LSTM LSTM LP Sigmoid LP Softmax⋯ ⋯ ⋯ ⋯ ⋯ ⋯ ⋮ ⋮ ⋯ ⋮ ⋮ ⋯ # $ # % # & # ' ( % ( ' ( & ) % →+ , - . / % →+ , - . ) % →0 . / % →0 . ) % . / % . ) & . / & . ) ' . / ' . ) % → & . / % → & . ) ' /' ) & / & ) % / % { ) % } ) % →+ , - . Prob . of choosing sub-tree ) % D % Copy & Split Concatenate LSTM Output & Cell Input Output Previous Output & Cell LP Linear Projection ( Feed Forward ) Att Attention Module ( Eq . ( 6 ) ) Figure 1 : The HISS architecture , illustrated on a three-node binary subtree , where i is the parent of j , k , and p is the parent of i. represents a variable or a constant . The arguments of operation are represented as the descendant subtrees of the corresponding node . Compared to string representation , tree representation naturally ensures any randomly generated expression in its form is syntactically correct . It also makes working with subexpressions easier : simply by working with subtrees . HISS approaches the problem using the reinforcement learning framework , where the action of generating simplified expressions are divided into two consecutive actions . The first action is to pick a subexpression ( or subtree ) that can be simplified , and the second action generates the simplified expression for the selected subexpression . Accordingly , HISS contains three modules . The first module is a tree encoder , which computes an embedding for each subtree ( including the entire tree ) of the input expression . The embeddings are useful for picking subtree for simplification as well as simplifying a subtree . The second module is a subtree selector , which selects a subtree for simplification . The third module is a tree decoder with an attention mechanism , which generates a simplified expression based on the input subtree embedding . The subsequent subsections will introduce each module respectively . 3.2 THE TREE ENCODER . The tree encoder generates embedding for every subtree of the input expression . We apply the N -ary Tree LSTM as proposed in Tai et al . ( 2015 ) , where N represents the maximum number of arguments that an operation has . It is important to note that although different operations have a different number of arguments , for structural uniformity , we assume that all operations have N arguments , with the excessive arguments being a NULL symbol . The tree encoder consists of two layers . The first layer is called the embedding layer , which is a fullyconnected layer that converts the one-hot representation of each input symbol to an embedding . The second layer is the tree LSTM layer , which is almost the same as the regular LSTM , except that the cell information now flows from the children nodes to their parent node . Formally , denote ci , hi , and xi as the cell , output and input of node i respectively . Then the tree LSTM encoder performs the following information [ hi , ci ] = LSTM xi , ⋃ j∈D ( i ) hj , ⋃ j∈D ( i ) cj , ( 2 ) where D ( i ) denotes the set of children of node i . Fig . 1 ( a ) plots the architecture of the tree LSTM encoder ( in green ) . Since each node fuses the information from its children , which again fuse the information from their own children , it is easy to see that the output hi summarizes the information of the entire subtree led by node i , and thus can be regarded as an embedding for this subtree .

Summary: This paper presents a method for symbolic superoptimization — the task of simplifying equations into equivalent expressions. The main goal is to design a method that does not rely on human input in defining equivalence classes, which should improve scalability of the simplification method to a larger set of expressions. The solution uses a reinforcement learning method for training a neural model that transforms an equation tree into a simpler but equivalent one. The model consists of (i) a tree encoder, a recursive LSTM that operates over the input equation tree, (ii) a sub-tree selector, a probability distribution over the nodes in the input equation tree, and (iii) a tree decoder, a two layer LSTM that includes a tree layer and a symbol generation layer. The RL reward uses an existing method for determining soft equivalence between the output tree and the input tree along with a positive score for compressing.