Deep symbolic regression: Recovering mathematical expressions from data via risk-seeking policy gradients

Discovering the underlying mathematical expressions describing a dataset is a core challenge for artificial intelligence . This is the problem of symbolic regression . Despite recent advances in training neural networks to solve complex tasks , deep learning approaches to symbolic regression are underexplored . We propose a framework that leverages deep learning for symbolic regression via a simple idea : use a large model to search the space of small models . Specifically , we use a recurrent neural network to emit a distribution over tractable mathematical expressions and employ a novel risk-seeking policy gradient to train the network to generate better-fitting expressions . Our algorithm outperforms several baseline methods ( including Eureqa , the gold standard for symbolic regression ) in its ability to exactly recover symbolic expressions on a series of benchmark problems , both with and without added noise . More broadly , our contributions include a framework that can be applied to optimize hierarchical , variable-length objects under a blackbox performance metric , with the ability to incorporate constraints in situ , and a risk-seeking policy gradient formulation that optimizes for best-case performance instead of expected performance . 1 INTRODUCTION . Understanding the mathematical relationships among variables in a physical system is an integral component of the scientific process . Symbolic regression aims to identify these relationships by searching over the space of tractable ( i.e . concise , closed-form ) mathematical expressions to best fit a dataset . Specifically , given a dataset ( X , y ) , where each point Xi ∈ Rn and yi ∈ R , symbolic regression aims to identify a function f : Rn → R that best fits the dataset , where the functional form of f is a short mathematical expression . The resulting expression can be readily interpreted and/or provide useful scientific insights simply by inspection . In contrast , conventional regression imposes a single model structure that is fixed during training , often chosen to be expressive ( e.g . a neural network ) at the expense of being easily interpretable . ∗Corresponding author . Symbolic regression exhibits several unique features that make it an excellent test problem for benchmarking automated machine learning ( AutoML ) and program synthesis methods : ( 1 ) there exist well-established , challenging benchmark problems with stringent success criteria ( White et al. , 2013 ) ; ( 2 ) there exist well-established baseline methods ( most notably , the Eureqa algorithm ( Schmidt & Lipson , 2009 ) ) ; and ( 3 ) the reward function is computationally expedient , allowing sufficient experiment replicates to achieve statistical significance . Most other AutoML tasks , e.g . neural architecture search ( NAS ) , do not exhibit these features ; in fact , even simply evaluating the efficiency of the discrete search itself is a known challenge within NAS ( Yu et al. , 2019 ) . The space of mathematical expressions is discrete ( in model structure ) and continuous ( in model parameters ) , growing exponentially with the length of the expression , rendering symbolic regression a challenging machine learning problem—thought to be NP-hard ( Lu et al. , 2016 ) . Given this large , combinatorial search space , traditional approaches to symbolic regression typically utilize evolutionary algorithms , especially genetic programming ( GP ) ( Koza , 1992 ; Schmidt & Lipson , 2009 ; Bäck et al. , 2018 ) . In GP-based symbolic regression , a population of mathematical expressions is “ evolved ” using evolutionary operations like selection , crossover , and mutation to improve a fitness function . While GP can be effective , it is also known to scale poorly to larger problems and to exhibit high sensitivity to hyperparameters . Deep learning has permeated almost all areas of artificial intelligence , from computer vision ( Krizhevsky et al. , 2012 ) to optimal control ( Mnih et al. , 2015 ) . However , deep learning may seem incongruous with or even antithetical toward symbolic regression , given that neural networks are typically highly complex , difficult to interpret , and rely on gradient information . We propose a framework that resolves this incongruity by tying deep learning and symbolic regression together with a simple idea : use a large model ( i.e . neural network ) to search the space of small models ( i.e . symbolic expressions ) . This framework leverages the representational capacity of neural networks to generate interpretable expressions , while entirely bypassing the need to interpret the network itself . We present deep symbolic regression ( DSR ) , a gradient-based approach for symbolic regression based on reinforcement learning . In DSR , a recurrent neural network ( RNN ) emits a distribution over mathematical expressions . Expressions are sampled from the distribution , instantiated , and evaluated based on their fitness to the dataset . This fitness is used as the reward signal to train the RNN using a novel risk-seeking policy gradient algorithm . As training proceeds , the RNN adjusts the likelihood of an expression relative to its reward , assigning higher probabilities to better expressions . We demonstrate that DSR outperforms several baseline methods , including two commercial software algorithms . We summarize our contributions as follows : ( 1 ) a novel method for symbolic regression that outperforms several baselines on a set of benchmark problems , ( 2 ) an autoregressive generative modeling framework for optimizing hierarchical , variable-length objects that accommodates in situ constraints , and ( 3 ) a novel risk-seeking policy gradient objective and accompanying Monte Carlo estimation procedure that optimizes for best-case performance instead of average performance . 2 RELATED WORK . Deep learning for symbolic regression . Several recent approaches leverage deep learning for symbolic regression . AI Feynman ( Udrescu & Tegmark , 2020 ) propose a problem-simplification tool for symbolic regression . They use neural networks to identify simplifying properties in a dataset ( e.g . multiplicative separability , translational symmetry ) , which they exploit to recursively define simplified sub-problems that can then be tackled using any symbolic regression algorithm . In GrammarVAE , Kusner et al . ( 2017 ) develop a generative model for discrete objects that adhere to a pre-specified grammar , then optimize them in latent space . They demonstrate this can be used for symbolic regression ; however , the method struggles to exactly recover expressions , and the generated expressions are not always syntactically valid . Sahoo et al . ( 2018 ) develop a symbolic regression framework using neural networks whose activation functions are symbolic operators . While this approach enables an end-to-end differentiable system , backpropagation through activation functions like division or logarithm requires the authors to make several simplifications to the search space , ultimately precluding learning certain simple classes of expressions like √ x or sin ( x/y ) . We address and/or directly compare to these works in Appendices C and E. AutoML and program synthesis . Symbolic regression is related to both automated machine learning ( AutoML ) and program synthesis , in that they all involve a search for an executable program ( i.e . expression ) to solve a particular task ( i.e . to fit data ) ( Abolafia et al. , 2018 ; Devlin et al. , 2017 ; Riedel et al. , 2016 ) . More specifically , our framework has many parallels to a body of works within AutoML that use an autoregressive RNN to define a distribution over discrete objects and use reinforcement learning to optimize this distribution under a black-box performance metric ( Zoph & Le , 2017 ; Ramachandran et al. , 2017 ; Bello et al. , 2017 ; Abolafia et al. , 2018 ) . For example , in neural architecture search ( Zoph & Le , 2017 ) , an RNN searches the space of neural network architectures , encoded by a sequence of discrete “ tokens ” specifying architectural properties ( e.g . number of neurons ) of each layer . The length of the sequence is fixed or scheduled during training ; in contrast , our framework defines a search space that is both inherently hierarchical and variable length . Ramachandran et al . ( 2017 ) search the space of neural network activation functions . While this space is hierarchical in nature , the authors ( rightfully ) constrain it substantially by positing a functional unit that is repeated sequentially , thus restricting their search space back to a fixed-length sequence . However , a repeating-unit constraint is not practical for symbolic regression because the ground truth expression may have arbitrary structure . Autoregressive models . The RNN-based distribution over expressions used in DSR is autoregressive , meaning each token is conditioned on the previously sampled tokens . Autoregressive models have proven to be useful for audio and image data ( Oord et al. , 2016a ; b ) in addition to the AutoML works discussed above ; we further demonstrate their efficacy for hierarchical expressions . GraphRNN defines a distribution over graphs that generates an adjacency matrix one column at a time in autoregressive fashion ( You et al. , 2018 ) . In principle , GraphRNN could be constrained to define a distribution over expressions , since trees are a special case of graphs . However , GraphRNN constructs graphs breadth-first , whereas expressions are more naturally represented using depth-first traversals ( Li et al. , 2005 ) . Further , DSR exploits the hierarchical nature of trees by providing the parent and sibling as inputs to the RNN , and leverages the additional structure of expression trees that a node ’ s value determines its number of children ( e.g . cosine is a unary operator and thus has one child ) . Risk-aware reinforcement learning . Many of the AutoML methods discussed above suffer from what we call the “ expectation problem. ” That is , policy gradient methods are fundamentally suited for optimizing expectations ; however , domains like neural architecture search and symbolic regression are evaluated by the few or single best-performing samples . Thus , there is a mismatch between the training objective function and the true desired objective : to maximize best-case performance . Abolafia et al . ( 2018 ) address the expectation problem by maintaining a priority queue of the best seen samples and using supervised learning to increase the likelihood of those top samples . Similarly , Liang et al . ( 2018 ) use a memory buffer to augment a policy gradient with off-policy training . These methods , however , only apply in the context of reinforcement learning environments with both deterministic transition dynamics and deterministic rewards . In contrast , the risk-seeking policy gradient introduced here is general , applying to any reinforcement learning environment and any stochastic policy gradient algorithm trained using batches , e.g . on Atari using proximal policy optimization ( Schulman et al. , 2017 ) . Lastly , our risk-seeking policy gradient is closely related to the EPOpt-ε algorithm used for robust reinforcement learning ( Rajeswaran et al. , 2016 ) , which is based on a risk-averse policy gradient formulation ( Tamar et al. , 2014 ) . 3 METHODS . Our overall approach involves representing mathematical expressions as sequences , developing an autoregressive model to generate expressions under a pre-specified set of constraints , and developing a risk-seeking policy gradient to train the model to generate better-fitting expressions . 3.1 GENERATING EXPRESSIONS WITH A RECURRENT NEURAL NETWORK . We leverage the fact that mathematical expressions can be represented using symbolic expression trees . Expression trees are a type of binary tree in which internal nodes are mathematical operators and terminal nodes are input variables or constants . Operators may be unary ( i.e . one argument , such as sine ) or binary ( i.e . two arguments , such as multiply ) . Further , we can represent an expression tree as a sequence of node values or “ tokens ” by using its pre-order traversal ( i.e . by visiting each node depth-first , then left-to-right ) . This allows us to generate an expression tree sequentially while still Library + 𝐬𝐢𝐧 𝐜𝐨𝐧𝐬𝐭 − × ÷ 𝐜𝐨𝐬 𝐥𝐨𝐠 𝐞𝐱𝐩 𝒚𝒙 Sampled token : Pre-order traversal Parent : Sibling : RNN : Categorical distribution : ÷ 𝐬𝐢𝐧 × 𝐜𝐨𝐧𝐬𝐭 𝒙 𝐥𝐨𝐠 𝒚 ÷ 𝐬𝐢𝐧 × × ÷ 𝐥𝐨𝐠 𝐜𝐨𝐧𝐬𝐭 𝐬𝐢𝐧 ÷ 𝐬𝐢𝐧 × 𝐜𝐨𝐧𝐬𝐭 𝒙 𝐥𝐨𝐠 𝒚 2 3 4 5 7 6 1 Expression tree A B C 1 2 3 4 6 75 Figure 1 : A . Example of sampling an expression from the RNN . For each token , the RNN emits a categorical distribution over tokens , a token is sampled , and the parent and sibling of the next token are used as the next input to the RNN . Subsequent tokens are sampled autoregressively until the tree is complete ( i.e . all tree branches reach terminal nodes ) . The resulting sequence of tokens is the tree ’ s pre-order traversal , which can be used to reconstruct the tree and instantiate its corresponding expression . Colors correspond to the number of children for each token . White circles represent empty tokens . Numbers indicate the order in which tokens were sampled . B . The library of tokens . C. The expression tree sampled in A . In this example , the sampled expression is sin ( cx ) / log ( y ) , where the value of the constant c is optimized with respect to an input dataset . maintaining a one-to-one correspondence between the tree and its traversal.1 Thus , we represent an expression τ by the pre-order traversal of its corresponding expression tree.2 We denote the ith token of the traversal as τi and the length of the traversal as |τ | = T . Each token has a value within a given library L of possible tokens , e.g . { + , − , × , ÷ , sin , x } . We generate expressions one token at a time along the pre-order traversal ( from τ1 to τT ) . A categorical distribution with parameters ψ defines the probabilities of selecting each token from L. To capture the “ context ” of the expression as it is being generated , we condition this probability upon the selections of all previous tokens in that traversal . This conditional dependence can be achieved very generally using an RNN with parameters θ that emits a probability vector ψ in an autoregressive manner . Specifically , the ith output of the RNN passes through a softmax layer ( with shared weights across time steps ) to produce vector ψ ( i ) , which defines the probability distribution for selecting the ith token τi , conditioned on the previously selected tokens τ1 : ( i−1 ) . That is , p ( τi|τ1 : ( i−1 ) ; θ ) = ψ ( i ) L ( τi ) , where L ( τi ) is the index in L corresponding to τi . The likelihood of the entire sampled expression is simply the product of the likelihoods of its tokens : p ( τ |θ ) = ∏|τ | i=1 p ( τi|τ1 : ( i−1 ) ; θ ) = ∏|τ | i=1 ψ ( i ) L ( τi ) . An example of the sampling process is illustrated in Figure 1 ; pseudocode is provided in Algorithm 2 in Appendix A . Note that different samples from the distribution have different tree structures of different size ; thus , the search space is inherently both hierarchical and variable length . Providing hierarchical inputs to the RNN . Conventionally , the input to the RNN when sampling a token would be a representation of the previously sampled token . However , the search space for symbolic regression is inherently hierarchical , and the previously sampled token may actually be very distant from the next token to be sampled in the expression tree . For example , the fifth and 1In general , a pre-order traversal is insufficient to uniquely reconstruct the tree . However , here we know how many children each node has based on its value , e.g . “ multiply ” is a binary operator and thus has two children . A pre-order traversal plus the corresponding number of children for each node is sufficient to uniquely reconstruct the tree . 2Given an expression tree , the corresponding mathematical expression is unique ; however , given an expression , its expression tree is not unique . For example , x2 and x× x are equivalent expressions but yield different trees . For simplicity , we use τ somewhat abusively to refer to an expression where it technically refers to an expression tree ( or equivalently , its pre-order traversal ) . sixth tokens sampled in Figure 1 are adjacent nodes in the traversal but are four edges apart in the expression tree . To better capture hierarchical information , we provide as inputs to the RNN a representation of the parent and sibling nodes of the token being sampled . We introduce an empty token for cases in which a node does not have a parent or sibling . Pseudocode for identifying the parent and sibling nodes given a partial traversal is provided in Subroutine 1 in Appendix A. Constraining the search space . Under our framework , it is straightforward to apply a priori constraints to reduce the search space . To demonstrate , we impose several simple , domain-agnostic constraints : ( 1 ) Expressions are limited to a pre-specified minimum and maximum length . We selected minimum length of 4 to prevent trivial expressions and a maximum length of 30 to ensure expressions are tractable . ( 2 ) The children of an operator should not all be constants , as the result would simply be a different constant . ( 3 ) The child of a unary operator should not be the inverse of that operator , e.g . log ( exp ( x ) ) is not allowed . ( 4 ) Descendants of trigonometric operators should not be trigonometric operators , e.g . sin ( x + cos ( x ) ) is not allowed because cosine is a descendant of sine . While still semantically meaningful , such composed trigonometric operators do not appear in virtually any scientific domain . We apply these constraints in situ ( i.e . concurrently with autoregressive sampling ) by zeroing out the probabilities of selecting tokens that would violate a constraint . Pseudocode for this process is provided in Subroutine 2 in Appendix A . This process ensures that samples always adhere to all constraints , without rejecting samples post hoc . In contrast , imposing constraints with GP requires rejecting evolutionary operations post hoc ( Fortin et al. , 2012 ) , which can be problematic ( Craenen et al. , 2001 ) , and as we show in our experiments , can actually reduce performance . Reward function . Once a pre-order traversal is sampled , we instantiate the corresponding symbolic expression and evaluate it with a reward function . A standard fitness measure in GP-based symbolic regression is normalized root-mean-square error ( NRMSE ) , the root-mean-square error normalized by the standard deviation of the target values , σy . That is , given a dataset ( X , y ) of size n and candidate expression f , NRMSE = 1σy √ 1 n ∑n i=1 ( yi − f ( Xi ) ) 2 . To bound the reward function , we apply a squashing function : R ( τ ) = 1/ ( 1 + NRMSE ) . Constant optimization . If the library L includes the constant token , sampled expressions may include several constant placeholders . These can be viewed as parameters ξ of the symbolic expression , which we optimize by maximizing the reward function : ξ ? = argmaxξ R ( τ ; ξ ) , using a nonlinear optimization algorithm , e.g . BFGS ( Fletcher , 2013 ) . We perform this inner optimization loop for each sampled expression as part of the reward computation before performing each training step .

Summary: The authors introduce a novel method for inferring a simple algebraic expression for an output in terms of an input. The method outputs a sequence tokens of the algebraic expression as represented by the pre-order traversal of an expression tree. For each token of the sequence, a recurrent neural network outputs a probability distribution on possible tokens, from which the token is sampled. Some notable contributions: