Causal Discovery with Reinforcement Learning

1 INTRODUCTION . Discovering and understanding causal mechanisms underlying natural phenomena are important to many disciplines of sciences . An effective approach is to conduct controlled randomized experiments , which however is expensive or even impossible in certain fields such as social sciences ( Bollen , 1989 ) and bioinformatics ( Opgen-Rhein and Strimmer , 2007 ) . Causal discovery methods that infer causal relationships from passively observable data are hence attractive and have been an important research topic in the past decades ( Pearl , 2009 ; Spirtes et al. , 2000 ; Peters et al. , 2017 ) . A major class of such causal discovery methods are score-based , which assign a score S ( G ) , typically computed with the observed data , to each directed graph G and then search over the space of all Directed Acyclic Graphs ( DAGs ) for the best scoring : min G S ( G ) , subject to G ∈ DAGs . ( 1 ) While there have been well-defined score functions such as the Bayesian Information Criterion ( BIC ) or Minimum Description Length ( MDL ) score ( Schwarz , 1978 ; Chickering , 2002 ) and the Bayesian Gaussian equivalent ( BGe ) score ( Geiger and Heckerman , 1994 ) , Problem ( 1 ) is generally NP-hard to solve ( Chickering , 1996 ; Chickering et al. , 2004 ) , largely due to the combinatorial nature of its acyclicity constraint with the number of DAGs increasing super-exponentially in the number of graph nodes . To tackle this problem , most existing approaches rely on local heuristics to enforce the acyclicity . For example , Greedy Equivalence Search ( GES ) enforces acyclicity one edge at a time , explicitly checking for the acyclicity constraint when an edge is added . GES is known to find global minimizer with infinite samples under suitable assumptions ( Chickering , 2002 ; Nandy et al. , 2018 ) , but this is not guaranteed in the finite sample regime . There are hybrid methods , e.g. , the max-min hill climbing method ( Tsamardinos et al. , 2006 ) , which use constraint-based approaches to reduce ∗Work was done during an internship at Huawei Noah ’ s Ark Lab . the search space before applying score-based methods . However , this methodology generally lacks a principled way of choosing a problem-specific combination of score functions and search strategies . Recently , Zheng et al . ( 2018 ) introduced a smooth characterization for the acyclicity . With linear models , Problem ( 1 ) was then formulated as a continuous optimization problem w.r.t . the weighted graph adjacency matrix by picking a proper loss function , e.g. , the least squares loss . Subsequent works Yu et al . ( 2019 ) and Lachapelle et al . ( 2019 ) have also adopted the evidence lower bound and the negative log-likelihood as loss functions , respectively , and used Neural Networks ( NNs ) to model the causal relationships . Note that the loss functions in these methods must be carefully chosen in order to apply continuous optimization methods . Unfortunately , many effective score functions , e.g. , the generalized score function proposed by Huang et al . ( 2018 ) and the independence based score function from Peters et al . ( 2014 ) , either can not be represented in closed forms or have very complicated equivalent loss functions , and thus can not be easily combined with this approach . We propose to use Reinforcement Learning ( RL ) to search for the DAG with the best score according to a predefined score function , as outlined in Figure 1 . The insight is that an RL agent with stochastic policy can determine automatically where to search given the uncertainty information of the learned policy , which can be updated promptly by the stream of reward signals . To apply RL to causal discovery , we use an encoder-decoder NN model to generate directed graphs from the observed data , which are then used to compute rewards consisting of the predefined score function as well as two penalty terms to enforce acyclicity . We resort to policy gradient and stochastic optimization methods to train the weights of the NNs , and our output is the graph that achieves the best reward , among all graphs generated in the training process . Experiments on both synthetic and real datasets show that our approach has a much improved search ability without sacrificing any flexibility in choosing score functions . In particular , the proposed approach with BIC score outperforms GES with the same score function on Linear Non-Gaussian Acyclic Model ( LiNGAM ) and linear-Gaussian datasets , and also outperforms recent gradient based methods when the causal relationships are nonlinear . 2 RELATED WORK . Constraint-based causal discovery methods first use conditional independence tests to find causal skeleton and then determine the orientations of the edges up to the Markov equivalence class , which usually contains DAGs that can be structurally diverse and may still have many unoriented edges . Examples include Sun et al . ( 2007 ) ; Zhang et al . ( 2012 ) that use kernel-based conditional independence criteria and the well-known PC algorithm ( Spirtes et al. , 2000 ) . This class of methods involve a multiple testing problem where the tests are usually conducted independently . The testing results may have conflicts and handling them is not easy , though there are certain works , e.g. , Hyttinen et al . ( 2014 ) , attempting to tackle this problem . These methods are also not robust as small errors in building the graph skeleton can result in large errors in the inferred Markov equivalence class . Another class of causal discovery methods are based on properly defined functional causal models . Unlike constraint-based methods that assume faithfulness and identify only the Markov equivalence class , these methods are able to distinguish between different DAGs in the same equivalence class , thanks to the additional assumptions on data distribution and/or functional classes . Examples include LiNGAM ( Shimizu et al. , 2006 ; 2011 ) , the nonlinear additive noise model ( Hoyer et al. , 2009 ; Peters et al. , 2014 ; 2017 ) , and the post-nonlinear causal model ( Zhang and Hyvärinen , 2009 ) . Besides Yu et al . ( 2019 ) ; Lachapelle et al . ( 2019 ) , other recent NN based approaches to causal discovery include Goudet et al . ( 2018 ) that proposes causal generative NNs to functional causal modeling with a prior knowledge of initial skeleton of the causal graph and Kalainathan et al . ( 2018 ) that learns causal generative models in an adversarial way but does not guarantee acyclicity . Recent advances in sequence-to-sequence learning ( Sutskever et al. , 2014 ) have motivated the use of NNs for optimization in various domains ( Vinyals et al. , 2015 ; Zoph and Le , 2017 ; Chen et al. , 2017 ) . A particular example is the traveling salesman problem that was revisited in the work of pointer networks ( Vinyals et al. , 2015 ) . Authors proposed a recurrent NN with nonparametric softmaxes trained in a supervised manner to predict the sequence of visited cities . Bello et al . ( 2016 ) further proposed to use the RL paradigm to tackle the combinatorial problems due to their relatively simple reward mechanisms . It was shown that an RL agent can have a better generalization even when the optimal solutions are used as labeled data in the previous supervised approach . Alternatively , the RL based approach in Dai et al . ( 2017 ) considered combinatorial optimization problems on ( undirected ) graphs and achieved a promising performance by exploiting graph structures , in contrast with the general sequence-to-sequence modeling . There are many other successful RL applications in recent years , e.g. , AlphaGo ( Silver et al. , 2017 ) , where the goal is to learn a policy for a given task . As an exception , Zoph and Le ( 2017 ) applied RL to neural architecture search . While we use a similar idea as the RL paradigm can naturally include the search task , our work is different in the actor and reward designs : our actor is an encoder-decoder model that generates graph adjacency matrices ( cf . Section 4 ) and the reward is tailored for causal discovery by incorporating a score function and the acyclicity constraint ( cf . Section 5.1 ) . 3 MODEL DEFINITION . We assume the following model for data generating procedure , as in Hoyer et al . ( 2009 ) ; Peters et al . ( 2014 ) . Each variable xi is associated with a node i in a d-node DAG G , and the observed value of xi is obtained as a function of its parents in the graph plus an independent additive noise ni , i.e. , xi : = fi ( xpa ( i ) ) + ni , i = 1 , 2 , . . . , d , where xpa ( i ) denotes the set of variables xj so that there is an edge from xj to xi in the graph , and the noises ni are assumed to be jointly independent . We also assume causal minimality , which in this case reduces to that each function fi is not a constant in any of its arguments ( Peters et al. , 2014 ) . Without further assumption on the forms of functions and/or noises , the above model can be identified only up to Markov equivalence class under the usual Markov and faithful assumptions ( Spirtes et al. , 2000 ; Peters et al. , 2014 ) ; in our experiments we will consider synthetic datasets that are generated from fully identifiable models so that it is practically meaningful to evaluate the estimated graph w.r.t . the true DAG . If all the functions fi are linear and the noises ni are Gaussian distributed , the above model yields the class of standard linear-Gaussian model that has been studied in Bollen ( 1989 ) ; Geiger and Heckerman ( 1994 ) ; Spirtes et al . ( 2000 ) ; Peters et al . ( 2017 ) . When the functions are linear but the noises are non-Gaussian , one can obtain the LiNGAM described in Shimizu et al . ( 2006 ; 2011 ) and the true DAG can be uniquely identified under favorable conditions . In this paper , we consider that all the variables xi are scalars ; extending to more complex cases is straightforward , provided with a properly defined score function . The observed data X , consisting of a number of vectors x : = [ x1 , x2 , . . . , xd ] T ∈ Rd , are then sampled independently according to the above model on an unknown DAG , with fixed functions fi and fixed distributions for ni . The objective of causal discovery is to use the observed data X , which gives the empirical version of the joint distribution of x , to infer the underlying causal DAG G . 4 NEURAL NETWORK ARCHITECTURE FOR GRAPH GENERATION . Given a dataset X = { xk } mk=1 where xk denotes the k-th observed sample , we want to infer the causal graph that best describes the data generating procedure . We would like to use NNs to infer the causal graph from the observed data ; specifically , we aim to design an NN based graph generator whose input is the observed data and the output is a graph adjacency matrix . A naive choice would be using feed-forward NNs to output d2 scalars and then reshape them to an adjacency matrix in Rd×d . However , this NN structure failed to produce promising results , possibly because the feed-forward NNs could not provide sufficient interactions amongst variables to capture the causal relations . Motivated by recent advances in neural combinatorial optimization , particularly the pointer networks ( Bello et al. , 2016 ; Vinyals et al. , 2015 ) , we draw n random samples ( with replacement ) { xl } nl=1 from X and reshape them as s : = { x̃i } di=1 where x̃i ∈ Rn is the vector concatenating all the i-th entries of the vectors in { xl } nl=1 . In an analogy to the traveling salesman problem , this represents a sequence of d cities lying in an n-dim space . We are concerned with generating a binary adjacency matrix A ∈ { 0 , 1 } d×d so that the corresponding graph is acyclic and achieves the best score . In this work we consider encoder-decoder models for graph generation : Encoder We use the attention based encoder in the Transformer structure proposed by Vaswani et al . ( 2017 ) . We believe that the self-attention scheme , together with structural DAG constraint , is capable of finding the causal relations amongst variables . Other attention based models such as graph attention network ( Veličković et al. , 2018 ) may also be used , which will be considered in a future work . Denote the outputs of the encoder by enci , i = 1 , 2 , . . . , d , with dimension de . Decoder Our decoder generates the graph adjacency matrix in an element-wise manner , by building relationships between two encoder outputs enci and encj . We consider the single layer decoder gij ( W1 , W2 , u ) = u T tanh ( W1 enci +W2 encj ) , where W1 , W2 ∈ Rdh×de , u ∈ Rdh×1 are trainable parameters and dh is the hidden dimension associated with the decoder . To generate a binary adjacency matrix A , we pass each entry gij into a logistic sigmoid function σ ( · ) and then sample according to a Bernoulli distribution with probability σ ( gij ) that indicates the probability of existing an edge from xi to xj . To avoid self-loops , we simply mask the ( i , i ) -th entry in the adjacency matrix . Other decoder choices include the neural tensor network model ( Socher et al. , 2013 ) and the bilinear model that build the pairwise relationships between encoder outputs . Another choice is the Transformer decoder which can generate an adjacency matrix in a row-wise manner . Empirically , we find that the single layer decoder performs the best , possibly because it contains less parameters and is easier to train to find better DAGs while the self-attention based encoder has provided sufficient interactions amongst the variables for causal discovery . Appendix A provides more details regarding these decoders and their empirical results with linear-Gaussian data models .

Summary: This work addresses the task of causal discovery. The proposed contribution is to apply prior work which uses reinforcement learning for combinatorial optimization to structure learning. Specifically, the proposed optimization problem seeks to maximize a penalized score criterion subject to the acyclicity constraint proposed by Zheng, et al. Empirical results show the proposed method performing favorably in contrast to prior art.