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: In this paper, the authors propose an RL-based structure searching method for causal discovery. The authors reformulate the score-based causal discovery problem into an RL-format, which includes the reward function re-design, hyper-parameter choose, and graph generation. To my knowledge, it’s the first time that the RL algorithm is applied to causal discovery area for structure searching.