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 a general flow-based approach to learn DAGs from data which provides a unified view of existing continuous optimization methods for structure learning. As a side benefit, the authors demonstrate that the proposed method could naturally be modified to handle missing data. The authors provide empirical studies to show that their proposed method outperforms the other baselines.