FCause: Flow-based Causal Discovery

1 INTRODUCTION . Understanding causal relationships between variables is crucial in many applications , including biology ( Koller & Friedman , 2009 ; Sachs et al. , 2005 ) , economics ( Varian , 2016 ; Cunningham , 2020 ) , and healthcare ( Tu et al. , 2019 ) . In addition , such information can also be used to advance other machine learning domains , such as fairness ( Kusner et al. , 2017 ; Chiappa , 2019 ) , privacy ( Tople et al. , 2020 ; Chandrasekaran et al. , 2021 ) and robustness ( Arjovsky et al. , 2019 ; Zhang et al. , 2020 ; Zheng et al. , 2021 ) . In practice , however , we often do not have a priori knowledge of causal relationships . We can gain this knowledge experimentally , e.g. , through randomized control trials ( Hariton & Locascio , 2018 ) , but experiments are sometimes cost-prohibitive or unethical . In such cases , causal discovery , the task of learning the causal relationships from existing data is essential . While causal discovery methods have been studied for decades , classical approaches tend to have scaling issues and rely on domain-specific assumptions . Recent advances in causal discovery use continuous optimization methods to overcome the scalability issues in classical approaches ( Zheng et al. , 2018 ; 2020 ) . They frame the combinatorial optimization problem of learning a directed acyclic graph ( DAG ) as a smooth optimization problem under constraints , allowing the use of efficient algorithms for continuous optimization . In addition to scalability advantages , this represents a first step in bridging the gap between deep learning and causal discovery . However , continuous optimization-based causal discovery remains experimental . These methods can not handle datasets with missing values and suffer from simple distributional assumptions , making performance sensitive to changes in the dataset ’ s scale ( Reisach et al. , 2021 ; Kaiser & Sipos , 2021 ) . These limitations directly impact these methods ’ reliability and applicability . In contrast , when the causal graph—the structure but not functional forms of causal relationships—is already known , causal-aware ( Bayesian ) deep learning methods have many application advantages . One recent framework , Carefl ( Khemakhem et al. , 2021 ) , fits an autoregressive flow model using the variables ’ causal ordering for the autoregressive transformations . This yields a model that can be used to answer complex causal queries , such as counterfactuals . Its performance for end-toend causal discovery , however , is limited by its reliance on classical causal discovery methods for inferring causal structures , the cost of fitting several flow-based models , and an inability to handle datasets with missing values . In this work , we show how continuous optimization methods can be combined with flow-based models to address these deficiencies . We propose FCause , an efficient flow-based causal discovery method able to capture complex nonlinear relations between variables , robust to data scaling , and able to handle datasets with missing values ( missing at random ) . FCause uses powerful flow-based models to capture complex relationships between variables ( Khemakhem et al. , 2021 ) , performs efficient causal discovery using ideas from continuous optimization-based causal discovery methods ( Zheng et al. , 2018 ) , and handles partially-observed datasets using amortized approximate inference to estimate distributions over missing values ( Kingma & Welling , 2013 ) . In section 6 , we perform extensive empirical evaluations of FCause , and observe that it achieves state-of-the-art results on multiple causal discovery benchmarks , and that its performance is robust with respect to changes in the scaling of the data ( e.g. , standardized data or not ) . In addition , FCauseremains competitive when used on datasets with 30 % of missing values , even comparing against baselines used on the corresponding fully-observed dataset . Finally , we present a unified formulation of several continuous optimization-based causal discovery algorithms based on flow-based models ( Khemakhem et al. , 2021 ) . This unified perspective facilitates the development of general new techniques to improve these methods , and allows a simple comparison between methods , shedding light on the benefits and drawbacks of each one . 2 RELATED WORK . Approaches for causal discovery from observational data can be classified into three broad groups : constraint-based , score-based and functional causal models ( Glymour et al. , 2019 ) . Constraint-based methods employ conditional independence tests to discover the underlying causal structure ( Spirtes & Glymour , 1991 ; Spirtes et al. , 2000 ) . Score-based methods find a causal structure by optimizing a score function ( Chickering & Meek , 2015 ; Chickering , 2020 ) . Functional causal models represent each variable as a function of its direct causes and some noise term ( Shimizu et al. , 2006 ; Hoyer et al. , 2008 ) . Classical causal discovery methods in these groups often struggle to scale to high dimensions , as the combinatorial space of possible causal orderings among variables grows superexponentially ( Peters et al. , 2017 ) . Thus , scalable causal discovery algorithms tend to reduce the search space through reliance on domain-specific assumptions . Recently , Zheng et al . ( 2018 ) introduced Notears , a new family of score-based method based on continuous optimization . Notears uses a novel algebraic characterization of directed acyclic graphs ( DAG ) , which allows an equality-constrained optimization problem to jointly learn the model parameters and adjacency relationships between variables . Notears inspired the development of other methods , Notears-MLP and Notears-Sob ( Zheng et al. , 2020 ) , Grandag ( Lachapelle et al. , 2019 ) , and DAG-GNN ( Yu et al. , 2019 ) , which extend the original formulation to model nonlinear relationships between variables . The methods ’ main benefits are its scalability and simplicity , a consequence of the fact that standard numerical solvers can be used to solve the resulting optimization problem . However , they have often been observed to be sensitive to different data scalings ( Kaiser & Sipos , 2021 ) , and can not handle missing values . Additionally , normalizing flows ( Rezende & Mohamed , 2015 ) have been used to build causal aware models ( Cai et al. , 2018 ; Khemakhem et al. , 2021 ) . These are based on the fact that the variables ’ causal ordering can be used for the transformations used by autoregressive flows ( Kingma et al. , 2016 ; Huang et al. , 2018 ) .1 These methods are able to model complex nonlinear relationships between variables . However , they rely on prior traditional methods or local search algorithms for causal discovery . Specifically , Self ( Cai et al. , 2018 ) uses a hill climbing procedure , and Carefl ( Khemakhem et al. , 2021 ) proposes to use a constraint-based method ( e.g . PC ) to find the graphs skeleton with as many oriented edges as possible and to fit several flow models to determine the orientation of the remaining edges . Apart from aforementioned methods , causal aware deep learning models have , in general , shown better properties regarding generalization and robustness ( Arjovsky et al. , 2019 ; Zhang et al. , 2020 ; Kyono et al. , 2020 ; Tople et al. , 2020 ) . 3 PRELIMINARIES . In this section , we describe the representations of causal relationships ( SEM and DAGs ) and explain how causal discovery can be formulated as an optimization task . We bring these components together when we introduce FCause in section 4 . 1While the formal connection to a restricted type of autoregressive flows was proposed by ( Khemakhem et al. , 2021 ) , some previous methods ( Mooij et al. , 2011 ; Cai et al. , 2018 ) use closely related ideas . Causal graphical models ( Peters et al. , 2017 ) are widely used to represent causal relationships between variables . Each variable is assigned a node in a graph ( typically a DAG ) , and causal relationships are represented as directed edges between nodes . These models are encoded using the graph ’ s adjacency matrix A ∈ { 0 , 1 } d×d , where d is the number of variables , or by explicitly stating the set of parents for each variables . If each of the edges is additionally assigned a weight , a weighted adjacency matrix is used W ∈ Rd×d instead . Here , zero entries indicate the absence of an edge , and non-zero entries indicate the presence of an edge with the corresponding weight . Structural Equation Models ( SEM ) are commonly used to describe a causal system . They characterize the value of a variable as a function of its causal parents and some external noise . Let x = [ x1 , . . . , xd ] represent d causally related variables , A ∈ { 0 , 1 } d×d the binary adjacency matrix between variables , and z = [ z1 , . . . , zd ] pairwise independent random noise . SEMs model each variable as xi = fi ( xpa ( i ) , zi ) , where fi a scalar function and xpa ( i ) = { xj : Aj , i = 1 } represents the set of parents of xi according to A . Additive Noise Models ( ANM ) represent a specific type of SEM for which the noise is additive : xi = fi ( xpa ( i ) ) + zi , i = 1 , . . . , d. ( 1 ) Equivalently , given a fixed adjacency matrix A , these d equations can be expressed jointly as x = fA ( x ) + z , ( 2 ) where fA ( x ) outputs a d-dimensional vector whose i-th component is given by fi ( xpa ( i ) ) . This is one of the most common functional forms in causal discovery . A crucial question regarding ANMs involves the identifiability of the underlying causal graph under observational data . Can the true adjacency matrix be recovered in the limit of infinite data ? Hoyer et al . ( 2008 ) showed that it is possible except when the underlying system is a combination of linear functions with Gaussian noise . In this work , we assume that the underlying functions are nonlinear without any specific assumption on their functional form . In this setting , the causal graph can be identified . A new algebraic characterization of DAGs was recently proposed by Zheng et al . ( 2018 ) . They showed that , given a weighted adjacency matrix W ∈ Rd×d , the quantity h ( W ) = tr ( eW W ) − d ( 3 ) is non-negative , and zero if and only if W represents a DAG . Since its introduction , this algebraic characterization of DAGs has been widely used to frame causal discovery problems with ANMs as continuous optimization tasks ( Zheng et al. , 2018 ; 2020 ; Lachapelle et al. , 2019 ) . All these methods propose to train a model ’ s parameters θ by maximizing a score subject to the constraint h ( W ( θ ) ) = 0 , whereW ( θ ) represents the weighted adjacency matrix as a function of the parameters θ . Due to the non-convexity of the set { θ : h ( W ( θ ) ) = 0 } , these methods use h ( W ( θ ) ) as a penalty term added to the loss , whose weight is increased as optimization proceeds .

Summary: This paper proposes to combine the continuous optimization-based causal discovery approach from notears with flow-based function learning. An extension is given to data missing (completely) at random. Also, a unifying framework is presented facilitating the comparison and exchange of ideas between different continuous optimization-based causal discovery methods, viewing all in terms of flows.