Gradient-Based Neural DAG Learning

1 INTRODUCTION . Structure learning and causal inference have many important applications in different areas of science such as genetics ( Koller & Friedman , 2009 ; Peters et al. , 2017 ) , biology ( Sachs et al. , 2005 ) and economics ( Pearl , 2009 ) . Bayesian networks ( BN ) , which encode conditional independencies using directed acyclic graphs ( DAG ) , are powerful models which are both interpretable and computationally tractable . Causal graphical models ( CGM ) ( Peters et al. , 2017 ) are BNs which support interventional queries like : What will happen if someone external to the system intervenes on variable X ? Recent work suggests that causality could partially solve challenges faced by current machine learning systems such as robustness to out-of-distribution samples , adaptability and explainability ( Pearl , 2019 ; Magliacane et al. , 2018 ) . However , structure and causal learning are daunting tasks due to both the combinatorial nature of the space of structures ( the number of DAGs grows super exponentially with the number of nodes ) and the question of structure identifiability ( see Section 2.2 ) . Nevertheless , these graphical models known qualities and promises of improvement for machine intelligence renders the quest for structure/causal learning appealing . The typical motivation for learning a causal graphical model is to predict the effect of various interventions . A CGM can be best estimated when given interventional data , but interventions are often costly or impossible to obtained . As an alternative , one can use exclusively observational data and rely on different assumptions which make the graph identifiable from the distribution ( see Section 2.2 ) . This is the approach employed in this paper . We propose a score-based method ( Koller & Friedman , 2009 ) for structure learning named GraNDAG which makes use of a recent reformulation of the original combinatorial problem of finding an optimal DAG into a continuous constrained optimization problem . In the original method named NOTEARS ( Zheng et al. , 2018 ) , the directed graph is encoded as a weighted adjacency matrix which represents coefficients in a linear structural equation model ( SEM ) ( Pearl , 2009 ) ( see Section 2.3 ) and enforces acyclicity using a constraint which is both efficiently computable and easily differentiable , thus allowing the use of numerical solvers . This continuous approach improved upon popular methods while avoiding the design of greedy algorithms based on heuristics . Our first contribution is to extend the framework of Zheng et al . ( 2018 ) to deal with nonlinear relationships between variables using neural networks ( NN ) ( Goodfellow et al. , 2016 ) . To adapt the acyclicity constraint to our nonlinear model , we use an argument similar to what is used in †Canada CIFAR AI Chair Correspondence to : sebastien.lachapelle @ umontreal.ca Zheng et al . ( 2018 ) and apply it first at the level of neural network paths and then at the level of graph paths . Although GraN-DAG is general enough to deal with a large variety of parametric families of conditional probability distributions , our experiments focus on the special case of nonlinear Gaussian additive noise models since , under specific assumptions , it provides appealing theoretical guarantees easing the comparison to other graph search procedures ( see Section 2.2 & 3.3 ) . On both synthetic and real-world tasks , we show GraN-DAG often outperforms other approaches which leverage the continuous paradigm , including DAG-GNN ( Yu et al. , 2019 ) , a recent nonlinear extension of Zheng et al . ( 2018 ) which uses an evidence lower bound as score . Our second contribution is to provide a missing empirical comparison to existing methods that support nonlinear relationships but tackle the optimization problem in its discrete form using greedy search procedures , namely CAM ( Bühlmann et al. , 2014 ) and GSF ( Huang et al. , 2018 ) . We show that GraN-DAG is competitive on the wide range of tasks we considered , while using pre- and post-processing steps similar to CAM . We provide an implementation of GraN-DAG here . 2 BACKGROUND . Before presenting GraN-DAG , we review concepts relevant to structure and causal learning . 2.1 CAUSAL GRAPHICAL MODELS . We suppose the natural phenomenon of interest can be described by a random vector X ∈ Rd entailed by an underlying CGM ( PX , G ) where PX is a probability distribution over X and G = ( V , E ) is a DAG ( Peters et al. , 2017 ) . Each node j ∈ V corresponds to exactly one variable in the system . Let πGj denote the set of parents of node j in G and let XπGj denote the random vector containing the variables corresponding to the parents of j in G. Throughout the paper , we assume there are no hidden variables . In a CGM , the distribution PX is said to be Markov to G , i.e . we can write the probability density function ( pdf ) of PX as p ( x ) = ∏d j=1 pj ( xj |xπGj ) where pj ( xj |xπGj ) is the conditional pdf of variable Xj given XπGj . A CGM can be thought of as a BN in which directed edges are given a causal meaning , allowing it to answer queries regarding interventional distributions ( Koller & Friedman , 2009 ) . 2.2 STRUCTURE IDENTIFIABILITY . In general , it is impossible to recover G given only samples from PX , i.e . without interventional data . It is , however , customary to rely on a set of assumptions to render the structure fully or partially identifiable . Definition 1 Given a set of assumptions A on a CGM M = ( PX , G ) , its graph G is said to be identifiable from PX if there exists no other CGM M̃ = ( P̃X , G̃ ) satisfying all assumptions in A such that G̃ 6= G and P̃X = PX . There are many examples of graph identifiability results for continuous variables ( Peters et al. , 2014 ; Peters & Bühlman , 2014 ; Shimizu et al. , 2006 ; Zhang & Hyvärinen , 2009 ) as well as for discrete variables ( Peters et al. , 2011 ) . These results are obtained by assuming that the conditional densities belong to a specific parametric family . For example , if one assumes that the distribution PX is entailed by a structural equation model of the form Xj : = fj ( XπGj ) +Nj with Nj ∼ N ( 0 , σ2j ) ∀j ∈ V ( 1 ) where fj is a nonlinear function satisfying some mild regularity conditions and the noises Nj are mutually independent , then G is identifiable from PX ( see Peters et al . ( 2014 ) for the complete theorem and its proof ) . This is a particular instance of additive noise models ( ANM ) . We will make use of this result in our experiments in Section 4 . One can consider weaker assumptions such as faithfulness ( Peters et al. , 2017 ) . This assumption allows one to identify , not G itself , but the Markov equivalence class to which it belongs ( Spirtes et al. , 2000 ) . A Markov equivalence class is a set of DAGs which encode exactly the same set of conditional independence statements and can be characterized by a graphical object named a completed partially directed acyclic graph ( CPDAG ) ( Koller & Friedman , 2009 ; Peters et al. , 2017 ) . Some algorithms we use as baselines in Section 4 output only a CPDAG . 2.3 NOTEARS : CONTINUOUS OPTIMIZATION FOR STRUCTURE LEARNING . Structure learning is the problem of learning G using a data set of n samples { x ( 1 ) , ... , x ( n ) } from PX . Score-based approaches cast this problem as an optimization problem , i.e . Ĝ = arg maxG∈DAG S ( G ) where S ( G ) is a regularized maximum likelihood under graph G. Since the number of DAGs is super exponential in the number of nodes , most methods rely on various heuristic greedy search procedures to approximately solve the problem ( see Section 5 for a review ) . We now present the work of Zheng et al . ( 2018 ) which proposes to cast this combinatorial optimization problem into a continuous constrained one . To do so , the authors propose to encode the graph G on d nodes as a weighted adjacency matrix U = [ u1| . . . |ud ] ∈ Rd×d which represents ( possibly negative ) coefficients in a linear SEM of the form Xj : = u > j X + Ni ∀j where Nj is a noise variable . Let GU be the directed graph associated with the SEM and let AU be the ( binary ) adjacency matrix associated with GU . One can see that the following equivalence holds : ( AU ) ij = 0 ⇐⇒ Uij = 0 ( 2 ) To make sure GU is acyclic , the authors propose the following constraint on U : Tr eU U − d = 0 ( 3 ) where eM , ∑∞ k=0 Mk k ! is the matrix exponential and is the Hadamard product . To see why this constraint characterizes acyclicity , first note that ( AUk ) jj is the number of cycles of length k passing through node j in graph GU . Clearly , for GU to be acyclic , we must have TrAU k = 0 for k = 1 , 2 , ... , ∞ . By equivalence ( 2 ) , this is true when Tr ( U U ) k = 0 for k = 1 , 2 , ... , ∞ . From there , one can simply apply the definition of the matrix exponential to see why constraint ( 3 ) characterizes acyclicity ( see Zheng et al . ( 2018 ) for the full development ) . The authors propose to use a regularized negative least square score ( maximum likelihood for a Gaussian noise model ) . The resulting continuous constrained problem is max U S ( U , X ) , − 1 2n ‖X−XU‖2F − λ‖U‖1 s.t . Tr eU U − d = 0 ( 4 ) where X ∈ Rn×d is the design matrix containing all n samples . The nature of the problem has been drastically changed : we went from a combinatorial to a continuous problem . The difficulties of combinatorial optimization have been replaced by those of non-convex optimization , since the feasible set is non-convex . Nevertheless , a standard numerical solver for constrained optimization such has an augmented Lagrangian method ( Bertsekas , 1999 ) can be applied to get an approximate solution , hence there is no need to design a greedy search procedure . Moreover , this approach is more global than greedy methods in the sense that the whole matrix U is updated at each iteration . Continuous approaches to combinatorial optimization have sometimes demonstrated improved performance over discrete approaches in the literature ( see for example Alayrac et al . ( 2018 , §5.2 ) where they solve the multiple sequence alignment problem with a continuous optimization method ) . 3 GRAN-DAG : GRADIENT-BASED NEURAL DAG LEARNING . We propose a new nonlinear extension to the framework presented in Section 2.3 . For each variable Xj , we learn a fully connected neural network with L hidden layers parametrized by φ ( j ) , { W ( 1 ) ( j ) , . . . , W ( L+1 ) ( j ) } where W ( ` ) ( j ) is the ` th weight matrix of the jth NN ( biases are omitted for clarity ) . Each NN takes as input X−j ∈ Rd , i.e . the vector X with the jth component masked to zero , and outputs θ ( j ) ∈ Rm , the m-dimensional parameter vector of the desired distribution family for variable Xj .1 The fully connected NNs have the following form θ ( j ) , W ( L+1 ) ( j ) g ( . . . g ( W ( 2 ) ( j ) g ( W ( 1 ) ( j ) X−j ) ) . . . ) ∀j ( 5 ) where g is a nonlinearity applied element-wise . Note that the evaluation of all NNs can be parallelized on GPU . Distribution families need not be the same for each variable . Let φ , { φ ( 1 ) , . . . , φ ( d ) } represents all parameters of all d NNs . Without any constraint on its parameter φ ( j ) , neural network j models the conditional pdf pj ( xj |x−j ; φ ( j ) ) . Note that the product∏d j=1 pj ( xj |x−j ; φ ( j ) ) does not integrate to one ( i.e . it is not a joint pdf ) , since it does not decompose according to a DAG . We now show how one can constrain φ to make sure the product of all conditionals outputted by the NNs is a joint pdf . The idea is to define a new weighted adjacency matrixAφ similar to the one encountered in Section 2.3 , which can be directly used inside the constraint of Equation 3 to enforce acyclicity .

Summary: This work addresses the problem of learning the structure of directed acyclic graphs in the presence of nonlinearities. The proposed approach is an extension of the NOTEARS algorithm which uses a neural network for each node in the graph during structure learning. This adaptation allows for non-linear relationships to be easily modeled. In addition to the proposed adaptation, the authors employ a number of heuristics from the causal discovery literature to improve the efficacy of the search. Empirical results are provided which compare the proposed algorithm to prior art.