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 authors propose a differentiable causal discovery method from interventional data. Each edge is learned separately using a score-based approach that relies on Monte-Carlo sampling. The main difference is that - since edges are learned separately - acyclicity need not be enforced.