BCD Nets: Scalable Variational Approaches for Bayesian Causal Discovery

1 Introduction . One of the key uses of statistical methods is learning causal relationships from observed data a.k.a . causal discovery [ 39 ] . Causal models allow us to forecast the effects of interventions and counterfactuals in several real-world domains , such as economic policy [ 57 ] and medicine [ 47 ] . Although early approaches to statistical inference emphasised that ‘ correlation is not causation ’ [ 15 ] , it has since been shown that for certain families of data-generating processes , it is indeed possible to infer causal relationships from purely observational data [ 37 , 29 ] . One such widely studied data-generating process is the linear-Gaussian structural equation model ( SEM ) [ 37 ] , where the causal relationships between the random variables in the model can be represented via a weighted directed acyclic graph ( DAG ) . The value of any variable in the DAG of a linear-Gaussian SEM is given by a linear combination of the values of its parent nodes and additive noise . For causal discovery , naive Monte-Carlo sampling or enumeration of the possible DAGs quickly becomes intractable , since the number of possible DAGs over a model grows superexponentially with the number of variables [ 18 ] . A variety of methods have been developed over the years to efficiently sample or optimize over DAGs [ 11 , 53 , 51 , 45 ] . For example , a recent line of work scales to high dimensions by maximizing the likelihood of the model ( MLE ) over a set of continuous relaxations of adjacency matrices using gradient-based methods and specialized DAG regularization terms . [ 62 , 61 , 35 , 59 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . However , the majority of the aforementioned works for scalable causal discovery in linear-Gaussian SEMs focus on recovering point estimates for the underlying DAG via MLE . In many practical scenarios , however , a point estimate fails to reflect the uncertainty in inferring the underlying DAG . This includes scenarios where the true DAG is non-identifiable given ( infinite ) observational data as well as practical limitations due to an imperfect optimization algorithm , model mismatch , or simply a finite dataset . In any of the above scenarios , it is desirable to obtain an explicit posterior distribution over the unobserved DAG instead of a single point estimate [ 19 ] . Causal inference is increasingly being applied in situations with important real-world consequences , where such a Bayesian estimation procedure could be useful to sample and reason over alternative generative mechanisms for observed data . To achieve this goal , we propose Bayesian Causal Discovery Nets ( BCD Nets ) , an algorithmic framework for Bayesian causal discovery in linear-Gaussian SEMs based on modern variational inference [ 6 ] . Classical approaches to Bayesian causal discovery struggle in higher dimensions [ 23 , 16 ] , with limited improvements via Monte Carlo approximations [ 22 , 58 ] . In order to scale our variational method to high dimensions , we address several design challenges . First , we describe an expressive variational family of factorized posterior distributions over the SEM parameters ( edge weights and noise variance ) using deep neural networks . The factorization exploits the decomposition of DAGs into triangular matrices and permutations for specifying the distribution over edge weights and node orderings respectively . Further , for low-variance stochastic optimization of the variational objective , we exploit recent advances in modeling and reparametrizing distributions over permutations via continuous relaxations [ 31 , 28 ] . Finally , we employ a horseshoe prior [ 7 ] on the edge weights , which promotes sparsity.1 We evaluate BCD Nets for causal discovery on a range of synthetic and real-world data benchmarks . Our experiments demonstrate that with finite datasets , there is considerable uncertainty in the inferred posterior over DAGs . Using BCD Nets , we are able to effectively quantify the uncertainty and significantly outperform competing estimators [ 35 , 58 ] on the standard Structured Hamming Distance metric , especially in low data regimes . 2 Preliminaries . 2.1 Linear-Gaussian Structural Equation Models . A structural equation model ( SEM ) is a collection of random variables x1 , . . . , xd associated with a directed acyclic graph ( DAG ) G with d nodes [ 37 ] . The SEM consists of a series of equations xj = fj ( Pa ( xj ) , j ) , where Pa ( xj ) gives the values of the parents of the jth node in G and j is a noise variable . For a linear-Gaussian SEM , the equations fj are linear and the noise j is additive Gaussian . Considering X = [ x1 , . . . , xd ] as a vector in Rd and W as the weighted adjacency matrix of G , we can see that X must satisfy X = W > X + , where ∼ N ( 0 , Σ ) is the additive noise vector and Σ = diag { σ21 , . . . , σ 2 d } is a diagonal noise covariance matrix . We will denote the setting when all noise variances are equal i.e. , σ1 = . . . = σd = σ as an “ equal variance ” setting and “ non-equal variance ” otherwise . Figure 1 shows an illustration . Linear SEMs have been used to model microarray data [ 38 ] and protein pathways [ 13 ] , among many other systems . 1Code is available at github.com/ermongroup/BCD-Nets 2.2 Causal Discovery . Given some dataset Xn1 = { X1 , X2 , . . . , Xn } drawn i.i.d . from a linear-Gaussian SEM , we are interested in inferring the adjacency matrix W and the diagonal entries of the covariance matrix Σ . Maximum Likelihood Estimation ( MLE ) This approach obtains a point estimate forW and Σ by maximizing the likelihood of the datasetXn1 . Using the fact thatW is constrained to be a DAG , it can be shown that I−W > is always invertible.2 . Hence , we can rearrange to obtainX = ( I −W > ) −1 , obtaining X as the product of a matrix and a Gaussian vector . It follows that X itself has a Gaussian distribution , X ∼ N ( 0 , Θ−1 ) , with precision matrix Θ = ( I −W ) Σ−1 ( I −W ) > . This gives a joint log-likelihood over the dataset Xn1 as log p ( Xn1 ; Σ , W ) = n 2 ( log det Θ− d log ( 2π ) ) − 1 2 n∑ i=1 X > i ΘXi , ( 1 ) where Xi runs over the n points in the dataset and Θ is given above . The space of DAG adjacency matrices W is characterised by the acyclicity constraint . When viewed as an subset of the space of all adjacency matrices ( Rd×d ) , this is a piecewise linear manifold with a number of facets that grows super-exponentially in the dimension ( roughly as d ! 2d 2/2 ) , rendering any approach based on enumeration over all DAGs intractable . Many approaches have been proposed to scale causal discovery via MLE to high-dimensional data . This includes recent approaches which relax the DAG constraint to a larger , continuous set of adjacency matrices , such as Rd×d in [ 62 ] or the Birkhoff polytope of doubly-stochastic matrices in [ 5 ] , which enables the use of gradient-based optimization . However , there are a number of challenges with using point estimators for causal discovery . Fundamentally , the true SEM parameters are identifiable from observational data only under specific conditions . In fact , the map between the the data distribution parameters Θ and the SEM parameters { W , Σ } may not be bijective even in the limit of infinite data and an oracle optimizer . Theoretical results on identifiability in linear-Gaussian SEMs currently hold only under restricted settings ; notably , this only includes the equal variance setting [ 38 ] . Further , even if the noise variances are equal in the ground-truth SEM , we are limited in practice by the finite size of the dataset and MLE could still converge to an incorrect solution . Finally , such point estimators are typically not robust to potential misspecifications in the likelihood model , e.g. , a non-Gaussian noise model , non-linear cause-effect relationships , etc . For the above scenarios , it is useful to characterize the uncertainty in estimating the SEM parameters to guide downstream analysis . For example , in medical applications giving a distribution of possible causal pathways , with quantified uncertainty , could be much more useful than a point estimate of the most likely pathway . Bayesian Estimation . In contrast to point estimators , Bayesian methods explicitly characterize the uncertainty in the estimated parameters [ 19 ] . That is , we treat the unknown SEM parameters { W , Σ } as random variables associated with a prior distribution p ( W , Σ ) . The likelihood model p ( Xn1 |Σ , W ) follows the same expression as the RHS in equation ( 1 ) . Given the prior and the likelihood , we obtain a posterior distribution p ( W , Σ|Xn1 ) over { W , Σ } via Bayes rule , which quantifies the uncertainty in estimating the SEM parameters . As long as the likelihood model is well-specified and the prior includes the ground-truth SEM parameters in its support , as the dataset size increases the posterior will concentrate around a set of SEM parameters . In the equal variance case these will be the true SEM parameters . In the non-equal variance case the posterior will concentrate on the set of DAG parameters quasi-equivalent to the ground truth as defined in [ 35 ] . These are the set of DAG parameters which generate data with the same covariance as the ground truth , and are hence indistinguishable based on data alone . We discuss this further in the appendix , Section F. The key challenge in Bayesian estimation is tractable computation of the posterior distribution in high-dimensional spaces . With the exception of specific prior and likelihood families e.g. , conjugate distributions , computing the posterior is typically intractable . In the next section , we present a variational framework for scalable approximation of the posterior for Bayesian causal discovery . 2Writing W = PLP > from Section 3 , I −W > = ( I − PL > P > ) = P ( I − L > ) P > . Now ( I − L > ) is upper-triangular with unit diagonal , so det ( I−L > ) = 1 . Thus ( I−W > ) has unit determinant , so is invertible . 3 Causal Discovery via Bayesian Causal Discovery Nets . As discussed previously , we are interested in learning the posterior distribution p ( W , Σ | Xn1 ) over the unknown SEM parameters { W , Σ } given an observed dataset Xn1 . Unlike point estimators , such a posterior distribution will allow us to quantify the uncertainty in estimation . Our framework , Bayesian Causal Discovery Nets ( BCD Nets ) allows us to tractably estimate this posterior . As a first step , our approach involves parametrizing the adjacency matrix W as the product of a permutation matrix P and a strictly lower-triangular matrix L , so thatW = PLP > . In graphical terms , L is a weight matrix for a canonical DAG with a fixed ordering , while pre- and post-multiplication by P and P > modifies the ordering of nodes . L is parameterised by a vector of weights l ∈ Rd ( d−1 ) /2 , and the constraint that P is a permutation ensures that W is the adjacency matrix of a DAG . Our goal is to obtain the posterior distribution p ( P , L , Σ|Xn1 ) . Due to the intractable partition function , we can not directly compute the posterior . We turn to variational inference to deliver a tractable approximation to the posterior [ 24 ] . The key idea here is to cast inference as an optimization problem , wherein we approximate the true posterior with a tractable family of distributions qφ ( P , L , Σ ) parameterized by φ and optimize these parameters φ to minimize the KL divergence between the approximate and true posterior distributions : DKL ( qφ ( P , L , Σ ) ‖p ( P , L , Σ | Xn1 ) ) = −E ( P , L , Σ ) ∼qφ [ log p ( Xn1 |P , L , Σ ) − log qφ ( P , L , Σ ) p ( P , L , Σ ) ] ︸ ︷︷ ︸ ELBO ( φ ) + log p ( Xn1 ) . ( 2 ) Hence , minimizing the KL divergence above corresponds to maximizing the evidence lower bound ( ELBO ) w.r.t . variational parameters φ . With a sufficiently expressive variational family from which to choose q , maximizing the ELBO recovers the true posterior as qφ ( P , L , Σ ) = p ( P , L , Σ|Xn1 ) . In practice , we face important modeling choices which have a substantial impact on the quality of the posterior obtained , as well as the difficulty of optimizing the ELBO . These include the variable ordering to use when factorizing qφ ( P , L , Σ ) using the chain rule , choice of variational family for the individual ( conditional ) factors , as well as the prior distribution p ( P , L , Σ ) . We discuss these algorithmic design choices next .

Summary: This submission presents an approach to quantify the uncertainty in estimating linear DAG structure using variational inference. In particular, the authors present a set of modeling design choices that (i) lead to tractable gradients of the ELBO, and (ii) are robust to certain classes of model misspecification. They go on to demonstrate that their approach outperforms maximum likelihood-based approaches on synthetic and semi-synthetic data.