DiBS: Differentiable Bayesian Structure Learning

1 Introduction . Discovering the statistical and causal dependencies that underlie the variables of a data-generating system is of central scientific interest . Bayesian networks ( BNs ) [ 1 ] and structural equation models are commonly used for this purpose [ 2–5 ] . Structure learning , the task of learning a BN from observations of its variables , is well-studied , but computationally very challenging due to the combinatorially large number of candidate graphs and the constraint of graph acyclicity . While structure learning methods arrive at a single plausible graph or its Markov equivalence class ( MEC ) , e.g. , [ 6–9 ] , Bayesian structure learning aims to infer a full posterior distribution over BNs given the observations . A distribution over structures allows quantifying the epistemic uncertainty and the degree of confidence in any given BN model , e.g. , when the amount of data is small . Most importantly , downstream tasks such as experimental design and active causal discovery rely on a posterior distribution over BNs to quantify the information gain from specific interventions and uncover the causal structure in a small number of experiments [ 10–15 ] . A key challenge in Bayesian structure learning is working with a posterior over BNs—a distribution over the joint space of ( discrete ) directed acyclic graphs and ( continuous ) conditional distribution parameters . Most of the practically viable approaches to Bayesian structure learning revolve around 1Our Python JAX implementation of DiBS is available at : https : //github.com/larslorch/dibs 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Markov chain Monte Carlo ( MCMC ) sampling in combinatorial spaces and bootstrapping of classical score and constraint-based structure learning methods , e.g. , in causal discovery [ 10–15 ] . However , these methods marginalize out the parameters and thus require a closed form for the marginal likelihood of the observations given the graph to remain tractable . This limits inference to simple and by now well-studied linear Gaussian and categorical BN models [ 16–18 ] and makes it difficult to infer more expressive BNs that , e.g. , model nonlinear relationships among the variables . Due to the discrete nature of these approaches , recent advances in approximate inference and gradient-based optimization could not yet be translated into similar performance improvements in Bayesian structure learning . In this work , we propose a novel , fully differentiable framework for Bayesian structure learning ( DiBS ) that operates in the continuous space of a latent probabilistic graph representation . Contrary to existing work , our formulation is agnostic to the distributional form of the BN and allows for inference of the joint posterior over both the conditional distribution parameters and the graph structure . This makes our approach directly applicable to more flexible BN models where neither the marginal likelihood nor the maximum likelihood parameter estimate have a closed form . We instantiate DiBS with the particle variational inference method of Liu and Wang [ 19 ] and present a general purpose method for approximate Bayesian structure learning . In our experiments on synthetic and real-world data , DiBS outperforms all alternative approaches to joint posterior inference of graphs and parameters and when modeling nonlinear interactions among the variables , often by a significant margin . This allows us to narrow down plausible causal graphs with greater precision and make better predictions under interventions—an important stepping stone towards active causal discovery . 2 Background . Bayesian networks A Bayesian network ( G , Θ ) models the joint density p ( x ) of a set of d variables x = x1 : d using ( 1 ) a directed acyclic graph ( DAG ) G encoding the conditional independencies of x and ( 2 ) parameters Θ defining the local conditional distributions of each variable given its parents in the DAG . When modeling p ( x ) using a BN , each variable is assumed to be independent of its non-descendants given its parents , thus allowing for a compact factorization of the joint p ( x |G , Θ ) into a product of local conditional distributions for each variable and its parents in G. Bayesian inference of BNs Given independent observations D = { x ( 1 ) , . . . , x ( N ) } , we consider the task of inferring a full posterior density over Bayesian networks that model the observations . Following Friedman and Koller [ 20 ] , given a prior distribution over DAGs p ( G ) and a prior over BN parameters p ( Θ |G ) , Bayes ’ Theorem yields the joint and marginal posterior distributions p ( G , Θ | D ) ∝ p ( G ) p ( Θ |G ) p ( D |G , Θ ) , ( 1 ) p ( G | D ) ∝ p ( G ) p ( D |G ) ( 2 ) where p ( D |G ) = ∫ p ( Θ |G ) p ( D |G , Θ ) dΘ is the marginal likelihood . Thus , p ( G | D ) in ( 2 ) is only tractable in special conjugate cases where the integral over Θ can be computed in closed form . The Bayesian formalism allows us to compute expectations of the form Ep ( G , Θ | D ) [ f ( G , Θ ) ] or Ep ( G | D ) [ f ( G ) ] ( 3 ) for any function f of interest . For instance , to perform Bayesian model averaging , we would use f ( G , Θ ) = p ( x |G , Θ ) or f ( G ) = p ( x |G ) , respectively [ 21 , 22 ] . In active learning of causal BN structures , a commonly used f is the expected decrease in entropy of G after an intervention [ 10–12 , 14 , 15 ] . Inferring either posterior is computationally challenging because there areO ( d ! 2 ( d 2 ) ) possible DAGs with d nodes [ 23 ] . Thus , computing the normalization constant p ( D ) is generally intractable . Continuous characterization of acyclic graphs Orthogonal to the work on Bayesian inference , Zheng et al . [ 9 ] have recently proposed a differentiable characterization of acyclic graphs for structure learning . In this work , we adopt the formulation of Yu et al . [ 24 ] , who show that a graph with adjacency matrix G ∈ { 0 , 1 } d×d does not have any cycles if and only if h ( G ) = 0 , where h ( G ) : = tr [ ( I + 1dG ) d ] − d . ( 4 ) If h ( G ) > 0 , the function can be interpreted as quantifying the cyclicity or non-DAG-ness of G. Follow-up work has leveraged this insight to model nonlinear relationships [ 24–27 ] , time-series data [ 28 ] , in the context of generative modeling [ 29 , 30 ] , for causal inference [ 31–33 ] , and contributed to its theoretical understanding [ 34 , 35 ] . So far , a connection to Bayesian structure learning has been missing . 3 Related Work . The existing literature on Bayesian Structure Learning predominantly focuses on inferring the marginal graph posterior p ( G | D ) . Since this requires p ( D |G ) to be tractable , inference is limited to BNs with linear Gaussian or Categorical conditional distributions [ 16–18 ] . By contrast , the formulation we introduce overcomes this fundamental restriction by allowing for joint inference of the graph and the parameters , thereby facilitating the active ( causal ) discovery of more expressive BNs . MCMC Sampling from the posterior over graphs is the most general approach to approximate Bayesian structure learning . Structure MCMC ( MC3 ) [ 36 , 37 ] performs Metropolis-Hastings in the space of DAGs by changing one edge at a time . Several works try to remedy its poor mixing behavior [ 38–40 ] . Alternatively , order MCMC draws samples in the smaller but still exponential space of node orders , which typically requires a hard limit on the maximum parent set size [ 20 ] . Attempts to correct for its unintended structural bias are themselves NP-hard to compute and/or limit the parent size [ 38 , 41 ] . By performing variational inference in a continuous latent space , the method we propose circumvents such mixing issues and parent size limitations . Bootstrapping The nonparametric DAG bootstrap [ 42 ] performs model averaging by bootstrapping D , where each resampled data set is used to learn a single graph , e.g. , using the GES or PC algorithms [ 6 , 7 ] . The obtained set of DAGs approximates the posterior by weighting each unique graph by its unnormalized posterior probability . In simple cases , a closed-form maximum likelihood parameter estimate may be used to approximate the joint posterior [ 14 ] , but only if p ( D |G ) is tractable in the first place . Exact methods A few notable exceptions use dynamic programming to achieve exact marginal inference in time O ( d2d ) , which is only feasible for d ≤ 20 nodes [ 43 , 44 ] . In special cases , e.g. , for tree structures or known node orderings , exact inference can be performed more efficiently [ 45 , 46 ] . 4 A Fully Differentiable Framework for Bayesian Structure Learning 4.1 General Approach With the goal of moving beyond the restrictive conjugate setups required by most discrete sampling methods for Bayesian structure learning , we propose to transfer the posterior inference task into the latent space of a probabilistic graph representation . Our resulting framework is consistent with the original Bayesian structure learning task in ( 3 ) , enforces the acyclicity of G via the latent space prior , and provides the score of the continuous latent posterior , thus making general purpose inference methods applicable off-the-shelf . Without loss of generality , we assume that there exists a latent variable Z that models the generative process of G. Specifically , the default generative model described in Section 2 is generalized into the following factorization : p ( Z , G , Θ , D ) = p ( Z ) p ( G |Z ) p ( Θ |G ) p ( D |G , Θ ) ( 5 ) Figure 1 displays the corresponding graphical model . The following insight provides us with an equivalence between the expectation we ultimately want to approximate and an expectation over the posterior of the continuous variable Z : Proposition 1 ( Latent posterior expectation ) . Under the generative model in ( 5 ) , it holds that ( a ) Ep ( G | D ) [ f ( G ) ] = Ep ( Z | D ) [ Ep ( G |Z ) [ f ( G ) p ( D |G ) ] Ep ( G |Z ) [ p ( D |G ) ] ] , and ( b ) Ep ( G , Θ | D ) [ f ( G , Θ ) ] = Ep ( Z , Θ | D ) [ Ep ( G |Z ) [ f ( G , Θ ) p ( Θ |G ) p ( D |G , Θ ) ] Ep ( G |Z ) [ p ( Θ |G ) p ( D |G , Θ ) ] ] . A proof is provided in Appendix A.1 . Rather than approximating p ( G | D ) or p ( G , Θ | D ) , our goal will be to infer p ( Z | D ) or p ( Z , Θ | D ) instead , which by Proposition 1 allows us to compute expectations of the form in ( 3 ) . In the following , we first discuss how to define the two factors p ( G |Z ) and p ( Z ) in a way that models only directed acyclic graphs . Then , we provide the details necessary to perform black box inference of the posteriors of the continuous latent variable Z .

Summary: This paper addresses the structure learning problem, in which the goal is to find a directed graphical model (structure and parameters) that accurately models an observed dataset. The authors work in the Bayesian structure learning paradigm, defining a prior distribution over graphical models and aiming to infer a posterior. Rather than defining a discrete prior over the unknown graph structure $G$, the authors define a prior over a continuous latent variable $Z$, and a conditional distribution on $G$ given $Z$, such that the data likelihood, $P(D \mid Z) = E_{G \mid Z}(P(D \mid G))$, is differentiable with respect to $Z$. This enables gradient-based inference methods to be applied to the problem of inferring $Z$, marginalizing $G$—if the necessary gradients can be efficiently estimated. In principle, any gradient-based inference method could be applied, but the authors spell out details of a Stein Variational Gradient Descent procedure. This proposed technique, and a variant in which an additional reweighting step is applied, are evaluated on synthetic data, against MCMC-based and non-parametric bootstrapping baselines.