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: The paper proposes a variational inference approach for approximate Bayesian causal discovery. In that, it focuses on modelling distribution over linear-Gaussian (LG) SEMs. The LG-SEMs are represented as tuples of an observation noise matrix  $\Sigma$ and a weighted adjacency matrix $W = PLP^\top$  and, which are modelled as combination of strictly lower-triangular real-valued matrix $L$ followed by a variable permutation represented by permutation matrix $P$. The proposed algorithm maximizes the ELBO w.r.t. to the parameters $\phi$ of the variational posterior $q_\phi(P|L, \Sigma) = q_\phi(P|L, \Sigma) q_\phi(L, \Sigma)$. Here, $q_\phi(P|L, \Sigma)$ is modelled as conditional Bolzmann distribution where re-parametrized sampling of permutation matrices $P$ is facilitated via the Gumble-Sinkhorn relaxation - this way pathwise gradient estimators of the ELBO can be obtained. In the experiment section, the proposed method is compared to three related approaches on simulated Erdos-Renyi graphs and on the well-known protein-signalling task.