Efficient Neural Causal Discovery without Acyclicity Constraints

1 INTRODUCTION . Uncovering and understanding causal mechanisms is an important problem not only in machine learning ( Schölkopf et al. , 2021 ; Pearl , 2009 ) but also in various scientific disciplines such as computational biology ( Friedman et al. , 2000 ; Sachs et al. , 2005 ) , epidemiology ( Robins et al. , 2000 ; Vandenbroucke et al. , 2016 ) , and economics ( Pearl , 2009 ; Hicks et al. , 1980 ) . A common task of interest is causal structure learning ( Pearl , 2009 ; Peters et al. , 2017 ) , which aims at learning a directed acyclic graph ( DAG ) in which edges represent causal relations between variables . While observational data alone is in general not sufficient to identify the DAG ( Yang et al. , 2018 ; Hauser & Bühlmann , 2012 ) , interventional data can improve identifiability up to finding the exact graph ( Eberhardt et al. , 2005 ; Eberhardt , 2008 ) . Unfortunately , the solution space of DAGs grows super-exponentially with the variable count , requiring efficient methods for large graphs . Current methods are typically applied to a few dozens of variables and can not scale so well , which is imperative for modern applications like learning causal relations with gene editing interventions ( Dixit et al. , 2016 ; Macosko et al. , 2015 ) . A promising new direction for scaling up DAG discovery methods are continuous-optimization methods ( Zheng et al. , 2018 ; 2020 ; Zhu et al. , 2020 ; Ke et al. , 2019 ; Brouillard et al. , 2020 ; Yu et al. , 2019 ) . In contrast to score-based and constrained-based ( Peters et al. , 2017 ; Guo et al. , 2020 ) methods , continuous-optimization methods reinterpret the search over discrete graph topologies as a continuous problem with neural networks as function approximators , for which efficient solvers are amenable . To restrict the search space to acyclic graphs , Zheng et al . ( 2018 ) first proposed to view the search as a constrained optimization problem using an augmented Lagrangian procedure to solve it . While several improvements have been explored ( Zheng et al. , 2020 ; Brouillard et al. , 2020 ; Yu et al. , 2019 ; Lachapelle et al. , 2020 ) , constrained optimization methods remain slow and hard to train . Alternatively , it is possible to apply a regularizer ( Zhu et al. , 2020 ; Ke et al. , 2019 ) to penalize cyclic graphs . While simpler to optimize , methods relying on acyclicity regularizers commonly lack guarantees for finding the correct causal graph , often converging to suboptimal solutions . Despite the advances , beyond linear , continuous settings ( Ng et al. , 2020 ; Varando , 2020 ) continuous optimization methods still can not scale to more than 100 variables , often due to difficulties in enforcing acyclicity . ∗Qualcomm AI Research is an initiative of Qualcomm Technologies , Inc . In this work , we address both problems . By modeling the orientation of an edge as a separate parameter , we can define the score function without any acyclicity constraints or regularizers . This allows for unbiased low-variance gradient estimators that scale learning to much larger graphs . Yet , if we are able to intervene on all variables , the proposed optimization is guaranteed to converge to the correct , acyclic graph . Importantly , since such interventions might not always be available , we show that our algorithm performs robustly even when intervening on fewer variables and having small sample sizes . We call our method ENCO for Efficient Neural Causal Discovery . We make the following four contributions . Firstly , we propose ENCO , a causal structure learning method for observational and interventional data using continuous optimization . Different from recent methods , ENCO models the edge orientation as a separate parameter . Secondly , we derive unbiased , low-variance gradient estimators , which is crucial for scaling up the model to large numbers of variables . Thirdly , we show that ENCO is guaranteed to converge to the correct causal graph if interventions on all variables are available , despite not having any acyclicity constraints . Yet , we show in practice that the algorithm works on partial intervention sets as well . Fourthly , we extend ENCO to detecting latent confounders . In various experimental settings , ENCO recovers graphs accurately , making less than one error on graphs with 1,000 variables in less than nine hours of computation . 2 BACKGROUND AND RELATED WORK . 2.1 CAUSAL GRAPHICAL MODELS . A causal graphical model ( CGM ) is defined by a distribution P over a set of random variables X = { X1 , ... , XN } and a directed acyclic graph ( DAG ) G = ( V , E ) . Each node i ∈ V corresponds to the random variable Xi and an edge ( i , j ) ∈ E represents a direct causal relation from variable Xi to Xj : Xi → Xj . A joint distribution P is faithful to a graph G if all and only the conditional independencies present in P are entailed by G ( Pearl , 1988 ) . Vice versa , a distribution P is Markov to a graph G if the joint distribution can be factorized as p ( X ) = ∏N i=1 pi ( Xi|pa ( Xi ) ) where pa ( Xi ) is the set of parents of the node i in G. An important concept which distinguishes CGMs from standard Bayesian Networks are interventions . An intervention on a variable Xi describes the local replacement of its conditional distribution pi ( Xi|pa ( Xi ) ) by a new distribution p̃ ( Xi|pa ( Xi ) ) . Xi is thereby referred to as the intervention target . An intervention is “ perfect ” when the new distribution is independent of the original parents , i.e . p̃ ( Xi|pa ( Xi ) ) = p̃ ( Xi ) . 2.2 CAUSAL STRUCTURE LEARNING . Discovering the graph G from samples of a joint distribution P is called causal structure learning or causal discovery , a fundamental problem in causality ( Pearl , 2009 ; Peters et al. , 2017 ) . While often performed from observational data , i.e . samples from P ( see Glymour et al . ( 2019 ) for an overview ) , we focus in this paper on algorithms that recover graphs from joint observational and interventional data . Commonly , such methods are grouped into constraint-based and score-based approaches . Constraint-based methods use conditional independence tests to identify causal relations ( Monti et al. , 2019 ; Spirtes et al. , 2000 ; Kocaoglu et al. , 2019 ; Jaber et al. , 2020 ; Sun et al. , 2007 ; Hyttinen et al. , 2014 ) . For instance , the Invariant Causal Prediction ( ICP ) algorithm ( Peters et al. , 2016 ; Christina et al. , 2018 ) exploits that causal mechanisms remain unchanged under an intervention except the one intervened on ( Pearl , 2009 ; Schölkopf et al. , 2012 ) . We rely on a similar idea by testing for mechanisms that generalize from the observational to the interventional setting . Another line of work is to extend methods working on observations only to interventions by incorporating those as additional constraints in the structure learning process ( Mooij et al. , 2020 ; Jaber et al. , 2020 ) . Score-based methods , on the other hand , search through the space of all possible causal structures with the goal of optimizing a specified metric ( Tsamardinos et al. , 2006 ; Ke et al. , 2019 ; Goudet et al. , 2017 ; Zhu et al. , 2020 ) . This metric , also referred to as score function , is usually a combination of how well the structure fits the data , for instance in terms of log-likelihood , as well as regularizers for encouraging sparsity . Since the search space of DAGs is super-exponential in the number of nodes , many methods rely on a greedy search , yet returning graphs in the true equivalence class ( Meek , 1997 ; Hauser & Bühlmann , 2012 ; Wang et al. , 2017 ; Yang et al. , 2018 ) . For instance , GIES ( Hauser & Bühlmann , 2012 ) repeatedly adds , removes , and flips the directions of edges in a proposal graph Published as a conference paper at ICLR 2022 NN NN NN X1 X3 X2 Distribution fitting X Y Z NN NN NN X Z Y Alternate between both steps Distribution fitting Graph fitting 𝑝 ( 𝑋 ! |… ) 𝑝 ( 𝑋 '' |… ) 𝑝 ( 𝑋 # |… ) X1 X2 X3 N1 N2 N3 X1 X3 X2 Alternate between both steps Distribution fitting Graph fitting 𝛾 ! '' 𝛾 '' ! 𝜃 ! '' 𝜃 # '' 𝜃 ! # 𝛾 ! # 𝛾 # ! 𝛾 # '' 𝛾 '' # @ @ ij L = ↵ · EX , C ij ⇥ LXi ! Xj ( Xj ) LXi 6 ! Xj ( Xj ) + sparse ⇤ C ( 1 ) C ( 2 ) C ( 3 ) LX1 ! X2 ( X2 ) LX1 6 ! X2 ( X2 ) LX3 ! X2 ( X2 ) LX3 6 ! X2 ( X2 ) LX1 ! X3 ( X3 ) LX1 6 ! X3 ( X3 ) LX2 ! X3 ( X3 ) LX2 6 ! X3 ( X3 ) @ @ 12 L̃ = 0 ( 12 ) · ( ✓12 ) · [ LX1 ! X2 ( X2 ) LX1 6 ! X2 ( X2 ) + sparse ] @ @ 32 L̃ = 0 ( 32 ) · ( ✓32 ) · [ LX3 ! X2 ( X2 ) LX3 6 ! X2 ( X2 ) + sparse ] @ @ 13 L̃ = 0 ( 13 ) · ( ✓13 ) · [ LX1 ! X3 ( X3 ) LX1 6 ! X3 ( X3 ) + sparse ] @ @ 23 L̃ = 0 ( 23 ) · ( ✓23 ) · [ LX2 ! X3 ( X3 ) LX2 6 ! X3 ( X3 ) + sparse ] @ @ ✓12 L̃ = @ @ ✓21 L̃ = 0 ( ✓12 ) · ( 12 ) · [ LX1 ! X2 ( X2 ) LX1 6 ! X2 ( X2 ) ] @ @ ✓13 L̃ = @ @ ✓31 L̃ = 0 ( ✓13 ) · ( 13 ) · [ LX1 ! X3 ( X3 ) LX1 6 ! X3 ( X3 ) ] f 1 f 2 f 3 1 @ @ ij L = ↵ · EX , C ij ⇥ LXi ! Xj ( Xj ) LXi 6 ! Xj ( Xj ) + sparse ⇤ 1 2 C ( 3 ) 1 6 3 6 2 2 1 6 LX2 6 ! X3 ( X3 ) @ @ 12 L̃ = 0 ( 12 ) · ( ✓12 ) · [ LX1 ! X2 ( X2 ) LX1 6 ! X2 ( X2 ) + sparse ] @ @ 32 L̃ = 0 ( 32 ) · ( ✓32 ) · [ LX3 ! X2 ( X2 ) LX3 6 ! X2 ( X2 ) + sparse ] @ @ 13 L̃ = 0 ( 13 ) · ( ✓13 ) · [ LX1 ! X3 ( X3 ) LX1 6 ! X3 ( X3 ) + sparse ] @ @ 23 L̃ = 0 ( 23 ) · ( ✓23 ) · [ LX2 ! X3 ( X3 ) LX2 6 ! X3 ( X3 ) + sparse ] @ @ ✓12 L̃ = @ @ ✓21 L̃ = 0 ( ✓12 ) · ( 12 ) · [ LX1 ! X2 ( X2 ) LX1 6 ! X2 ( X2 ) ] @ @ ✓13 L̃ = @ @ ✓31 L̃ = 0 ( ✓13 ) · ( 13 ) · [ LX1 ! X3 ( X3 ) LX1 6 ! X3 ( X3 ) ] 1 2 f 3 1 @ @ ij L = ↵ · EX , C ij ⇥ LXi ! Xj ( Xj ) LXi 6 ! Xj ( Xj ) + sparse ⇤ ( 1 ) 2 C ( 3 ) 2 ( 2 ) 1 6 2 ( 2 ) ( ) 3 6 2 ( 2 ) 3 ( 3 ) 1 6 3 ( 3 ) LX2 6 ! X3 ( X3 ) @ @ 12 L̃ = 0 ( 12 ) · ( ✓12 ) · [ LX1 ! X2 ( X2 ) LX1 6 ! X2 ( X2 ) + sparse ] @ @ 32 L̃ = 0 ( 32 ) · ( ✓32 ) · [ LX3 ! X2 ( X2 ) LX3 6 ! X2 ( X2 ) + sparse ] @ @ 13 L̃ = 0 ( 13 ) · ( ✓13 ) · [ LX1 ! X3 ( X3 ) LX1 6 ! X3 ( X3 ) + sparse ] @ 23 L̃ = 0 ( 23 ) · ( ✓23 ) · [ LX2 ! X3 ( X3 ) LX2 6 ! X3 ( X3 ) + sparse ] @ @ ✓12 L̃ = @ @ ✓21 L̃ = 0 ( ✓12 ) · ( 12 ) · [ LX1 ! X2 ( X2 ) LX1 6 ! X2 ( X2 ) ] @ @ ✓13 L̃ = @ @ ✓31 L̃ = 0 ( ✓13 ) · ( 13 ) · [ LX1 ! X3 ( X3 ) LX1 6 ! X3 ( X3 ) ] 1 2 f 3 1 X1 X2 X3 NN1 NN2 NN3 X1 X3 X2 Alterna between both steps Distribution fitting Graph fitting Figure 1 : Visualization of the two training stages of ENCO , distribution fitting and graph fitting , on an example graph with 3 variables ( X1 , X2 , X3 ) . The graph on the right further shows how the parameters γ and θ correspond to edge probabilities . We learn those parameters by comparing multiple graph samples on how well they generalize from observational to interventional data . until no higher-scoring graph can be found . The Interventional Greedy SP ( IGSP ) algorithm ( Wang et al. , 2017 ) is a hybrid method using conditional independence tests in its score function . Continuous-optimization methods are score-based methods that avoid the combinatorial greedy search over DAGs by using gradient-based methods ( Zheng et al. , 2018 ; Ke et al. , 2019 ; Lachapelle et al. , 2020 ; Yu et al. , 2019 ; Zheng et al. , 2020 ; Zhu et al. , 2020 ; Brouillard et al. , 2020 ) . Thereby , the adjacency matrix is parameterized by weights that represent linear factors or probabilities of having an edge between a pair of nodes . The main challenge of such methods is how to limit the search space to acyclic graphs . One common approach is to view the search as a constrained optimization problem and deploy an augmented Lagrangian procedure to solve it ( Zheng et al. , 2018 ; 2020 ; Yu et al. , 2019 ; Brouillard et al. , 2020 ) , including NOTEARS ( Zheng et al. , 2018 ) and DCDI ( Brouillard et al. , 2020 ) . Alternatively , Ke et al . ( 2019 ) propose to use a regularization term penalizing cyclic graphs while allowing unconstrained optimization . However , the regularizer must be designed and weighted such that the correct , acyclic causal graph is the global optimum of the score function .

Summary: The paper works on causal discovery on both observational and interventional data. It proposes a method for enforcing the acyclicity of DAGs which avoids imposing constraints on the optimization problem. It proves that under some assumptions, when all the variables have interventional data, the resulting graph converges to the ground-truth one. The experimental results show the properties and performance of the proposed method, including the cases with latent confounders and interventions on fewer variables.