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 proposes a method called DiBS (differentiable Bayesian structure learning) for learning the structure of a Bayesian network (BN) from data based on a fully differentiable reparameterization of the relevant likelihood expressions, allowing for the use of continuous optimization-based Bayesian inference algorithms. This method allows for posterior inference of graph structure alone, or joint inference over structure and conditional distribution parameters. First, they augment the standard generative model of data from a Bayesian network with a continuous-valued latent variable of node embeddings such that the probability of each edge in the graph adjacency matrix can be modeled by the inner product similarity of node embeddings with a temperature parameter. The authors use a previously-introduced function of the adjacency matrix that calculates graph cyclicity to include a prior that encourages acyclic structure. They derive expressions that reformulate the necessary score functions in terms of gradients with respect to the new continuous latent graph representation, using the Gumbel-softmax trick to reparameterize the gradient where applicable and the score function gradient estimator otherwise. The method uses Stein variational gradient descent as the inference algorithm, optimizing a set of particles for the latent graph representation while annealing the temperature parameters to result in a final set of particles that correspond to adjacency matrices of valid directed acyclic graphs. A variant of the base method called DiBS+ that weights each particle by its unnormalized posterior probability is also explored. On synthetic data from linear Gaussian BNs, DiBS and DiBS+ show comparable performance to greedy search-based (BPC, BGES) and sampling-based ($\text{MC}^3$) algorithms across a variety of metrics when inferring structure alone, and comparable or better performance when jointly inferring structure and parameters. On synthetic data with nonlinear local conditionals parameterized by a neural network, DiBS/DiBS+ perform comparably to or better than sampling-based methods for joint inference across metrics. These results are also reflected on real data from a protein signaling network, where DiBS/DiBS+ achieve the best recovery of the true network and the most well-calibrated edge confidence scores in the joint inference setting.