DiBS: Differentiable Bayesian Structure Learning

Abstract:
Bayesian structure learning allows inferring Bayesian network structure from data while reasoning about the epistemic uncertainty—a key element towards enabling active causal discovery and designing interventions in real world systems. In this work, we propose a general, fully differentiable framework for Bayesian structure learning (DiBS) that operates in the continuous space of a latent probabilistic graph representation. Contrary to existing work, DiBS is agnostic to the form of the local conditional distributions and allows for joint posterior inference of both the graph structure and the conditional distribution parameters. This makes our formulation directly applicable to posterior inference of complex Bayesian network models, e.g., with nonlinear dependencies encoded by neural networks. Using DiBS, we devise an efficient, general purpose variational inference method for approximating distributions over structural models. In evaluations on simulated and real-world data, our method significantly outperforms related approaches to joint posterior inference.1

heading: 1 Introduction
text: 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.

heading: 2 Background
text: 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.

heading: 3 Related Work
text: 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.

heading: 4.2 Representing DAGs in a Continuous Latent Space
text: Generative model of directed graphs We define the latent variable Z as consisting of two embedding matrices U,V ∈ Rk×d, i.e., Z = [U,V]. Building on Kipf and Welling [47], we propose to use a bilinear generative model for the adjacency matrix G ∈ {0, 1}d×d of directed graphs using the inner product between the latent variables in Z:
pα(G |Z) = d∏ i=1 d∏ j 6=i pα(gij |ui,vj) with pα(gij = 1 |ui,vj) = σα(u>i vj) (6)
where σα(x) = 1/(1 + exp(−αx)) denotes the sigmoid function with inverse temperature α. We denote the corresponding matrix of edge probabilities in G given Z by Gα(Z) ∈ [0, 1]d×d with
Gα(Z)ij := pα(gij = 1 |ui,vj) . (7) Since we model acyclic graphs, which do not contain self-loops, we set pα(gii = 1 |Z) := 0. The latent dimensionality k trades off the complexity of the variable interactions with tractability during inference. For k ≥ d, the matrix of edge probabilities Gα(Z) is not constrained in rank, and the generative model in (6) can represent any adjacency matrix without self-loops. That said, the size of the latent representation Z only grows O(d ·k) and, in principle, k can be chosen independently of d. Note that the formulation in (6) models directed graphs since σα(u>i vj) 6= σα(u>j vi). We can even interpret ui and vi as node embeddings that may encode more information than mere edge probabilities, e.g., for graph neural networks. The fact that pα(G |Z) is invariant to orthogonal transformations of U and V is not an issue in practice.
Acyclicity via the latent prior distribution A major constraint in learning BNs is the acyclicity of G. With a latent graph model p(G |Z) in place, we design the prior p(Z) to act as a soft constraint enforcing that only DAGs are modeled. Specifically, we define the prior of Z as the product of independent Gaussians with a Gibbs distribution that penalizes the expected cyclicity of G given Z:
pβ(Z) ∝ exp ( −β Ep(G |Z) [ h(G) ]) ∏ ij N (zij ; 0, σ2z) (8)
Here, h is the DAG constraint function given in (4) and β is an inverse temperature parameter controlling how strongly the acyclicity is enforced. As β →∞, the support of pβ(Z) reduces to all Z that only model valid DAGs. The Gaussian component ensures that the norm of Z is well-behaved. Traditional graph priors of the form p(G) ∝ f(G) that induce, e.g., sparsity of G, can be flexibly incorporated into pβ(Z) by means of an additional factor involving f(Gα(Z)) or Ep(G |Z)[f(G)]. Unless highlighting a specific point, we omit writing the hyperparameters α and β to simplify notation.

heading: 4.3 Estimators for Gradient-Based Bayesian Inference
text: The extended generative model in (5) not only allows us to incorporate the notoriously difficult acyclicity constraint into the Bayesian framework. By rephrasing the posteriors p(G | D) and p(G,Θ | D) in terms of p(Z |D) and p(Z,Θ |D), respectively, it also makes Bayesian structure learning amenable to variational inference techniques that operate in continuous space and rely on the gradient of the unnormalized log posterior, also known as the score. The following result provides us with the score of both latent posteriors. Detailed derivations can be found in Appendix A.2. Proposition 2 (Latent posterior score). Under the generative graph model defined in (5), the gradient of the log posterior density of Z is given by
(a) ∇Z log p(Z | D) = ∇Z log p(Z) + ∇Z Ep(G |Z)
[ p(D |G) ] Ep(G |Z) [ p(D |G)
] (9) which is relevant when the marginal likelihood p(D |G) can be computed efficiently. In the general case, when inferring the joint posterior of Z and Θ, the gradients of the log posterior are given by
(b) ∇Z log p(Z,Θ | D) = ∇Z log p(Z) + ∇Z Ep(G |Z)
[ p(Θ,D |G) ] Ep(G |Z) [ p(Θ,D |G) ] (10)
(c) ∇Θ log p(Z,Θ | D) = Ep(G |Z)
[ ∇Θp(Θ,D |G) ] Ep(G |Z) [ p(Θ,D |G)
] (11) where p(Θ,D |G) = p(Θ |G)p(D |G,Θ), which is efficient to compute by construction.
The expectations have tractable Monte Carlo approximations because sampling from p(G |Z) in (6) is simple and parallelizable. The gradient terms of the form∇ZEp(G |Z)[ · ], which also appear inside ∇Z log p(Z), can be estimated using two different techniques, depending on the BN model inferred. Differentiable (marginal) likelihood Using the Gumbel-softmax trick [48, 49], we can separate the randomness from Z when sampling from p(G |Z) and obtain the following estimator:
∇ZEp(G |Z) [ p(D |G) ] ≈ Ep(L) [ ∇G p(D |G) ∣∣ G=Gτ (L,Z) · ∇Z Gτ (L,Z) ]
(12)
where L ∼ Logistic(0, 1)d×d i.i.d.. The matrix-valued function Gτ (·) is defined elementwise as
Gτ (L,Z)ij :=
{ στ ( lij + αu > i vj ) if i 6= j
0 if i = j (13)
The estimator applies equally to p(Θ,D |G) in place of p(D |G). For the estimator to be welldefined, p(D |G) or p(Θ,D |G), respectively, needs to be differentiable with respect to G. More specifically, the gradient ∇Gp(D |G) or ∇Gp(D,Θ |G) needs to be defined when G lies on the interior of [0, 1]d×d and not at its discrete endpoints. This depends on the parameterization of the BN model we want to infer. In case p(D |G) or p(Θ,D |G) is only defined for discrete G, it is possible to evaluate∇Gp(D |G) or∇Gp(Θ,D |G) using hard Gumbel-max samples of G (i.e., with τ =∞) and a straight-through gradient estimator. Since the DAG constraint h is differentiable, the Gumbel-softmax trick can always be applied inside∇Z log p(Z). In practice, we always use τ = 1. Non-differentiable (marginal) likelihood In general,∇Gp(D |G) or∇Gp(D,Θ |G) depending on the inference task might be not available or ill-defined. In this setting, the score function estimator provides us with a way to estimate the gradient we need [50]:
∇ZEp(G |Z) [ p(D |G) ] = Ep(G |Z) [( p(D |G)− b ) ∇Z log p(G |Z) ] (14)
The estimator likewise applies for p(Θ,D |G) in place of p(D |G). Here, b is a constant with respect to G that can be used for variance reduction [51], and ∇Z log p(G |Z) is trivial to compute. The derivations of both (12) and (14), alongside a more detailed discussion, can be found in Appendix B.

heading: 5 Particle Variational Inference for Structure Learning
text: In the previous section, we have proposed a differentiable formulation for Bayesian structure learning (DiBS) that is agnostic to the form of the local BN conditionals and, more importantly, translates learning discrete graph structures into an inference problem over the continuous variable Z. In the following, we overcome the remaining challenge of inferring the intractable DiBS posteriors p(Z | D) and p(Z,Θ | D) by employing Stein variational gradient descent (SVGD) [19], a gradient-based and general purpose variational inference method. The resulting algorithm infers a particle approximation of the marginal or joint posterior density over BNs given observational data.
SVGD for posterior inference Since Proposition 2 provides us with the gradient of the latent posterior score functions, we can apply SVGD off-the-shelf. SVGD minimizes the KL divergence to a target distribution by iteratively transporting a set of particles using a sequence of kernel-based transformation steps. We provide a more detailed overview of SVGD in Appendix C. Following this paradigm for DiBS, we iteratively update a fixed set {Z(m)}Mm=1 or {Z(m),Θ(m)}Mm=1 to approximate p(Z | D) or p(Z,Θ | D), respectively. If the BN model we aim to infer has a properly-defined likelihood gradient with respect to G, we use the Gumbel-softmax estimator in (12) to approximate the posterior score. Otherwise, we resort to the score function estimator in (14). We use a simple kernel for SVGD:
k ( (Z,Θ), (Z′,Θ′) ) := exp ( − 1 γz ‖Z− Z′‖22 ) + exp ( − 1 γθ ‖Θ−Θ′‖22 ) (15)
with bandwidths γz, γθ. For inference of p(Z | D), we leave out the second term involving Θ,Θ′. While Z is invariant to orthogonal transformations, more elaborate kernels that are, e.g., invariant to such transforms empirically perform worse in our experiments.
Algorithm 1 DiBS with SVGD [19] for inference of p(G | D) Input: Initial set of latent particles {Z(m)0 }Mm=1, kernel k, schedules for αt, βt, and stepsizes ηt Output: Set of discrete graph particles {G(m)}Mm=1 approximating p(G | D)
1: Incorporate prior belief of p(G) into p(Z) . See Section 4.2 2: for iteration t = 0 to T − 1 do 3: Estimate score∇Z log p(Z | D) given in (9) for each Z(m)t . See (12) and (14) 4: for particle m = 1 to M do 5: Z(m)t+1 ← Z (m) t + ηt φt(Z (m) t ) . SVGD step
where φt(·) := 1
M M∑ k=1 [ k(Z (k) t , · )∇Z(k)t log p(Z (k) t | D) +∇Z(k)t k(Z (k) t , · ) ] 6: return {G∞(Z(m)T )} M m=1 . See (16) and (17)
Annealing α and β The latent variable Z not only probabilistically models the graph G, but can also be viewed as a continuous relaxation of G, with α trading off smoothness with accuracy. As α→∞, the sigmoid σα(·) converges to the unit step function. Hence, as α→∞ in the graph model pα(G |Z) in (6), the expectations in Proposition 1 simplify to:
Ep(G | D) [ f(G) ] → Ep(Z | D) [ f ( G∞(Z) )] Ep(G,Θ | D) [ f(G,Θ) ] → Ep(Z,Θ | D) [ f ( G∞(Z),Θ
)] (16) where G∞(Z) denotes the single limiting graph implied by Z = [U,V] and is defined as
G∞(Z)ij := { 1 if u>i vj > 0 and i 6= j 0 otherwise
(17)
See Appendix A.3. In this case, pα(G |Z) converges to representing only a single graph. To be able to invoke this simplification, we anneal α→∞ over the iterations of SVGD and, upon termination, convert the latent variables Z to the single discrete G∞(Z). Furthermore, we similarly let β →∞ in the latent prior pβ(Z) over the iterations to enforce that the latent representation of G only models DAGs. By Equation (16), the resulting DAGs form a consistent particle approximation of p(G | D) or p(G,Θ | D), respectively. Algorithm 1 summarizes DiBS instantiated with SVGD for inference of p(G | D). The general case of inferring p(G,Θ | D) is given in Algorithm 2 of Appendix D. Single-particle approximation SVGD reduces to regular gradient ascent for the maximum a posteriori estimate when transporting only a single particle [19]. In this special case, DiBS with SVGD recovers some of the existing continuous structure learning methods: gradient ascent on a linear Gaussian likelihood solves an optimization problem similar to NOTEARS [9]. The cyclicity penalizer acts analogously. However, not only does DiBS automatically turn into a full Bayesian approach when using more particles, it is also not limited to settings such as linear Gaussian conditionals, where the adjacency matrix G and the parameters Θ can be modeled together by a weighted adjacency matrix.
Weighted particle mixture In high dimensional settings, it may be beneficial to move beyond a uniform weighting of the inferred particles of BN models to approximate p(G | D) or p(G,Θ | D). We consider a particle mixture that weights each particle by its unnormalized posterior probability p(G,D) or p(G,Θ,D), respectively, under the BN model. While we do not have a strong theoretical justification for the weighting, our motivation is that most of the particles in the empirical distribution will be unique as a consequence of the super-exponentially large space of DAGs, which may result in a crude approximation of their posterior probability mass function. In our experiments, DiBS and its instantiation with SVGD are used interchangeably, and DiBS+ denotes the weighted particle mixture.

heading: 6 Evaluation on Synthetic Data

heading: 6.1 Experimental Setup
text: Synthetic data We compare DiBS to a set of related methods in marginal and joint posterior inference of synthetic linear and nonlinear Gaussian BNs. Our setup follows [9, 24, 27, 52], who consider inferring BNs with Erdős-Rényi and scale-free random structures [53, 54]. For each graph,
here with d ∈ {20, 50} nodes and 2d edges in expectation, we sample a set of ground truth parameters and then generate training, held-out, and interventional data sets. In all settings, we use N = 100 observations for inference, emulating the use case of Bayesian structure learning where the uncertainty about the graph structure is significant.
Graph priors For Erdős-Rényi graphs, all methods use the prior p(G) ∝ q‖G‖1(1−q)( d 2)−‖G‖1 , capturing that each edge exists independently w.p. q [53]. For scale-free graphs, we define the prior p(G)∝ ∏d i=1(1+‖G>i ‖1)−3, analogous to their power law degree distribution p(deg)∼ deg
−3 [54]. Here, G>i is the i-th column of the adjacency matrix. DiBS implements either prior by using the corresponding term above as an additional factor in p(Z) with G := Gα(Z) (see Section 4.2).
Metrics Since neither density nor samples of the ground truth posteriors are available for BNs of d ∈ {20, 50} variables, we follow the evaluation metrics used by previous work. We define the expected structural Hamming distance to the ground truth graph G∗ under the inferred posterior as
E-SHD(p,G∗) := ∑ G p(G | D) · SHD(G,G∗) (18)
where SHD(G,G∗) counts the edge changes that separate the essential graphs representing the MECs of G and G∗ [8, 55]. In addition, we follow Friedman and Koller [20] and Ellis and Wong [41] and compute the area under the receiver operating characteristic curve (AUROC) for pairwise edge predictions when varying the confidence threshold under the inferred marginal p(gij = 1 | D). Finally, following Murphy [10], we also evaluate the ability to predict future observations by computing the average negative (marginal) log likelihood on 100 held-out observations Dtest:
neg. LL(p,Dtest) := − ∑ G,Θ p(G,Θ | D) · log p(Dtest |G,Θ) (19)
When inferring p(G | D), the corresponding neg. MLL metric instead uses p(Dtest |G). Analogously, we also compute the interventional log likelihoods I-LL and I-MLL, a relevant metric in causal inference [10, 12]. Here, an interventional data set (Dint, I) is instead used to compute p(Dint |G,Θ, I) and p(Dint |G, I) in (19), respectively. Scores are the average of 10 interventional data sets. All reported metrics in this section are aggregated for inference of 30 random synthetic BNs.
In the remainder, DiBS is always run for 3,000 iterations and with k = d for inference of d-variable BNs, which leaves the matrix of edge probabilities unconstrained in rank. We discard a DiBS particle in the rare case that a returned graph is cyclic. Complete details on Gaussian BNs, the evaluation metrics, hyperparameters, and all baselines can be found in Appendix E.

heading: 6.2 Linear Gaussian Bayesian Networks
text: Marginal posterior inference For linear Gaussian BNs, we first evaluate the classical setting of inferring the marginal posterior p(G | D) since p(D |G) can be computed in closed form. To this end, we employ the commonly used Bayesian Gaussian Equivalent (BGe) marginal likelihood, which
scores Markov equivalent structures equally [16, 17]. The form of the BGe score requires DiBS to use the score function estimator in (14).
We compare DiBS with the nonparametric DAG bootstrap [42] using the constraint-based PC [7] and the score-based GES [6] algorithms (BPC and BGES). For MCMC, we only consider structure MCMC (MC3) [37] as a comparison. Order MCMC or hybrid DP approaches bound the number of parents and thus often exclude the ground truth graph a priori, especially for scale-free BN structures. Burn-in and thinning for MC3 are chosen to make the wall time comparable with DiBS run on CPUs. In the remainder of the paper, each method uses 30 samples to approximate the posterior over BNs.
Figure 2 summarizes the results for 30 randomly generated BNs with d = 20 nodes. We find that DiBS+ performs well compared to the other methods, all of which were specifically developed for the marginal inference scenario evaluated here. DiBS+ appears to be preferable to DiBS.
Joint posterior inference When inferring the joint posterior p(G,Θ | D), we can employ a more explicit representation of linear Gaussian BNs, where the conditional distribution parameters are standard Gaussian. Here, DiBS can leverage the Gumbel-softmax estimator in (12) because p(G,Θ | D) is well-defined when G lies on the interior of [0, 1]d×d (see Appendix E.1). To provide a comparison with DiBS in the absence of an applicable MCMC method, we propose two variants of MC3 as baselines. Metropolis-Hastings MC3 (M-MC3) jointly samples parameters and structures, and Metropolis-within-Gibbs MC3 (G-MC3) alternates in proposing structure and parameters [56]. Moreover, we extend the bootstrap methods by taking the closed-form maximum likelihood estimate [57] as the posterior parameter sample for a given graph inferred using the BGe score (BPC∗ and BGES∗), an approach taken in, e.g., causal BN learning [14].
Figure 3 shows the results for d = 20 nodes, where E-SHD and AUROC are computed by empirically marginalizing out the parameters. We find that DiBS+ is the only considered method that performs well across all of the metrics, often outperforming the baselines by a significant margin. As for marginal posterior inference of linear Gaussian BNs, DiBS+ performs slightly better than DiBS.

heading: 6.3 Nonlinear Gaussian Bayesian Networks
text: We also consider joint inference of nonlinear Gaussian BNs where the mean of each local conditional Gaussian is parameterized by a 2-layer dense neural network with five hidden nodes and ReLU activation functions (see Appendix E.1). Since the marginal likelihood does not have a closed form, we are unable to use BPC∗ and BGES∗ as a means of comparison. Figure 4 displays the results for d = 20 variables, where a given BN model has |Θ| = 2,220 weights and biases. Analogous to joint inference of linear Gaussian BNs, DiBS and DiBS+ outperform the MCMC baselines across the considered metrics. To the best of our knowledge, this is the first time that such nonlinear Gaussian BN models have been inferred under the Bayesian paradigm, which opens up exciting avenues in the active learning of more sophisticated causal structures.
20
40
60
80
100
E rd
ő s-
R én
y i
E-SHD
0.45
0.60
0.75
0.90
AUROC
0
4×104 8×104
1.2×105 neg. LL
0
4×104 8×104
1.2×105 neg. I-LL
M -M
C 3
G -M
C 3
D iB
S
D iB
S+
20
30
40
50
S ca
le -f
re e
M -M
C 3
G -M
C 3
D iB
S
D iB
S+ 0.45
0.60
0.75
0.90
M -M
C 3
G -M
C 3
D iB
S
D iB
S+
0
3×104 6×104 9×104
1.2×105
M -M
C 3
G -M
C 3
D iB
S
D iB
S+
0
3×104 6×104 9×104
1.2×105
20
40
60
80
100
E rd
ő s-
R én
y i
E-SHD
0.45
0.60
0.75
0.90
AUROC
0
4×104 8×104
1.2×105 neg. LL
0
4×104 8×104
1.2×105 neg. I-LL
M -M
C 3
G -M
C 3
D iB
S
D iB
S+
20
30
40
50
S ca
le -f
re e
M -M
C 3
G -M
C 3
D iB
S
D iB
S+ 0.45
0.60
0.75
0.90
M -M
C 3
G -M
C 3
D iB
S
D iB
S+
0
3×104 6×104 9×104
1.2×105
M -M
C 3
G -M
C 3
D iB
S
D iB
S+
0
3×104 6×104 9×104
1.2×105
Figure 4: Joint posterior inference of nonlinear Gaussian BNs with d = 20 nodes. Here, the mean of the local conditional distribution of each node is parameterized by a 2-layer neural network with five hidden nodes. DiBS and DiBS+ perform favorably across the board, particularly in the graph metrics.
Appendix F complements Sections 6.2 and 6.3 with details on computing time and efficient implementation of DiBS with SVGD, showing that GPU wall times of the above inference experiments lie on the order of seconds or only a few minutes. In addition, Appendix G provides results for d = 50 variables, where DiBS+ likewise performs favorably when jointly inferring p(G,Θ | D). For marginal posterior inference of p(G | D) under the BGe marginal likelihood, DiBS appears to require more Monte Carlo samples to compensate for the high variance of the score function estimator in this high-dimensional setting.

heading: 6.4 DiBS with SVGD: Additional Analyses and Ablation Studies
text: Having compared DiBS and its instantiation with SVGD to existing approaches, we finally devote Appendix H to empirically analyzing some properties of the algorithm. One of our key results is that, all other things held equal, substituting our inner product model in (6) with pα(gij = 1 |Z) = σα(zij), where single scalars encode the edge probabilities, results in significantly worse evaluation metrics.
We additionally study the uncertainty quantification in (non)identifiable edge structures and show the effects of reducing the latent dimensionality k or the number of iterations T . Our findings suggest that reducing either hyperparameter still allows for competitive posterior approximations and enables trading off posterior inference quality with computational efficiency, e.g., in large-scale applications.

heading: 7 Application: Inferring Protein Signaling Networks From Cell Data
text: Table 1: Inference of protein signaling pathways with Gaussian BNs. Metrics are the mean± SD of 30 random restarts.
E-SHD AUROC
† MC3 34.0 ± 0.7 0.616 ± 0.027 BPC 25.5 ± 2.3 0.566 ± 0.020 BGES 33.7 ± 1.7 0.641 ± 0.030 DiBS 37.4 ± 0.5 0.647 ± 0.047 DiBS+ 34.7 ± 1.5 0.629 ± 0.045 § M-MC3 37.3 ± 3.5 0.551 ± 0.078 G-MC3 30.5 ± 3.2 0.527 ± 0.067 DiBS 23.4 ± 0.5 0.598 ± 0.052 DiBS+ 22.9 ± 2.7 0.557 ± 0.052 ¶ M-MC3 25.2 ± 3.0 0.526 ± 0.084 G-MC3 35.1 ± 3.2 0.540 ± 0.080 DiBS 22.6 ± 0.5 0.577 ± 0.039 DiBS+ 22.8 ± 1.9 0.535 ± 0.041
† Linear Gaussian BN; graph only via BGe marginal lik. § Linear Gaussian BN; graph and parameters jointly ¶Nonlinear Gaussian BN; graph and parameters jointly
A widely used benchmark in structure learning is the proteomics data set by Sachs et al. [3]. The data contain N = 7, 466 continuous measurements of d = 11 proteins involved in human immune system cells as well as an established causal network of their signaling interactions.
We infer both linear and nonlinear Gaussian BNs with Erdős-Rényi graph priors exactly as in Section 6. The AUROC results in Table 1 indicate that the posterior by DiBS under the BGe model provides the most calibrated edge confidence scores. Not penalizing model complexity as much, the marginal BGe posterior of DiBS (DiBS+) averages a high expected number of 39.0 (35.0) edges, compared to 12.7 (14.2) and 12.6 (14.2) for its joint posteriors over linear and nonlinear BNs, respectively. Appendix I provides further details and analyses on this matter. The E-SHD scores show that, among all the methods, DiBS is closest in structure to the consensus network when performing joint inference with nonlinear Gaussian BNs. This further highlights the need for nonlinear conditionals and joint inference of the graph and parameters in complex real-world settings.

heading: 8 Conclusion
text: We have presented a general, fully differentiable approach to inference of posterior distributions over BNs. Our framework is based on a continuous latent representation of DAGs, whose posterior can be equivalently inferred—without loss of generality and using existing black box inference methods. While we have used SVGD [19] for this purpose, our general approach could also be instantiated with, e.g., gradient-based sampling methods that rely on the score of the target density [58, 59]. This may improve upon the asymptotic runtime of DiBS with SVGD, which scales quadratically in the number of sampled particles. We expect that our end-to-end approach can be extended to handle missing and interventional data as well as to amortized contexts, where rich unstructured data is available.
Broader impact Our work is relevant to any scientific discipline that aims at inferring the (causal) structure of a system or reasoning about the effects of interventions. If algorithms and decisions are grounded in the structural understanding of a data-generating system and take into account the epistemic uncertainty, we expect them to be more robust and have fewer unforeseen side-effects. However, the assumptions allowing for a causal interpretation of DAGs, e.g., the absence of unmeasured confounders, are often untestable and to be taken with care [60], particularly in safety-critical and societally-sensitive applications. Hence, while potential misuse can never be ruled out, our presented method predominantly promises positive societal and scientific impact.
Acknowledgments and Disclosure of Funding
This project received funding from the Swiss National Science Foundation under NCCR Automation under grant agreement 51NF40 180545, the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation program grant agreement no. 815943, and was supported with compute resources by Oracle Cloud Services. This work was also supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039B, and by the Machine Learning Cluster of Excellence, EXC number 2064/1 – Project number 390727645. We thank Nicolo Ruggeri and Guillaume Wang for their valuable feedback.

Review summaries:
This work proposes a new approach to posterior approximation in the context of structure learning based on the Stein Variational Gradient Descent (SVGD) procedure and the differentiable acyclicity constraint inspired from Zheng et al 2018. A continuous latent variable Z (parent of G only) is introduced in the model to allow the computation of the gradient of its posterior (with respect to Z) which in turn enables one to use SVGD to sample from P(G|D). This gradient can be estimated using the Gumbel-softmax trick or score function estimator. The framework is validated on synthetic data as well as the protein signaling dataset from Sachs et al..

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.

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.