PaperlayerView exact revision on arXiv

Source-linked AI summary

DiBS: Differentiable Bayesian Structure Learning

Lars Lorch, Jonas Rothfuss, Bernhard Schölkopf, Andreas Krause

arXiv:2105.11839v3cs.LGstat.ML

TL;DR

Bayesian structure learning needs posterior distributions over networks to represent uncertainty relevant to intervention design, but existing methods face computational and modeling limits. DiBS uses differentiable inference in a continuous latent graph space to jointly infer structures and parameters, and outperforms alternatives in joint posterior inference.

  • Problem

    Bayesian structure learning seeks posterior uncertainty over graphs for tasks such as experimental design, while existing methods' marginal-likelihood requirements limit inference to simpler network conditionals.

    Source: §1 · pp. 1–2

  • Method

    DiBS performs differentiable variational inference in a continuous latent graph representation to jointly infer graph structure and conditional-distribution parameters.

    Source: §1 · p. 2, §4.2 · p. 4, §5 · p. 6

  • Results

    DiBS outperforms alternative approaches to joint graph-and-parameter posterior inference, often by a significant margin for nonlinear interactions.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The method narrows plausible causal graphs and improves predictions under interventions, as a step toward active causal discovery.

    Source: §1 · p. 2

    Causal interpretation of DAGs relies on assumptions such as the absence of unmeasured confounders, which are often untestable.

    Source: §8 · p. 10

Abstract

from arXiv · show

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 Introduction

Bayesian structure learning seeks a posterior over graphs, rather than one plausible graph, to represent uncertainty that supports experimental design and active causal discovery. DiBS addresses computational and modeling constraints with differentiable joint inference over graph structure and parameters.

1 Introduction

  • Structure learning is computationally challenging because it searches a combinatorially large set of candidate graphs subject to acyclicity.Bayesian structure learning targets a full posterior over networks rather than a single graph or Markov equivalence class.Source: §1 · p. 1
  • Existing Bayesian approaches often marginalize parameters, requiring a tractable closed-form marginal likelihood and limiting inference to linear-Gaussian or categorical networks.This makes inference of more expressive networks, such as those with nonlinear relationships, difficult.Source: §1 · p. 2
  • DiBS jointly infers graph structures and conditional-distribution parameters in a fully differentiable latent-graph framework, and outperforms alternatives in joint posterior inference.The reported advantage is especially notable for nonlinear interactions, often by a significant margin.Source: §1 · p. 2

2 Background

The background defines Bayesian networks and Bayesian inference over their graph and parameter structures. It highlights the computational challenge of posterior inference and the role of differentiable acyclicity constraints.

2 Background

  • A Bayesian network pairs a DAG encoding conditional independencies with parameters for local conditional distributions, yielding a compact factorization of the joint density.Each variable is assumed independent of its non-descendants given its parents.Source: §2 · p. 2
  • Bayesian inference places priors on DAGs and network parameters, then uses observations to infer posterior distributions over these structures.The marginal likelihood integrates the data likelihood over the parameter distribution and is tractable only in special conjugate cases.Source: §2 · p. 2
  • Posterior inference is generally intractable because the number of DAGs grows as O(d!2^(d choose 2)), making normalization difficult.Expected decreases in graph entropy after interventions are one use of posterior expectations in active causal learning.Source: §2 · p. 2
  • A graph is acyclic exactly when the differentiable constraint function h(G) equals zero; positive values quantify cyclicity.The cited differentiable characterization had not yet been connected to Bayesian structure learning.Source: §2 · p. 2

3 Related Work

Prior Bayesian structure-learning methods face constraints from marginal-likelihood tractability, graph-space sampling, and limits on exact inference. DiBS instead moves inference into a continuous latent space.

3 Related Work

  • Most existing Bayesian structure-learning work targets the marginal graph posterior, which requires tractable p(D | G) and restricts inference to linear-Gaussian or categorical conditionals.Joint graph-and-parameter inference can overcome this restriction for more expressive networks.Source: §3 · p. 3
  • Structure MCMC can mix poorly, while order MCMC remains exponential and typically imposes a maximum parent-set size.Correcting order MCMC's structural bias is NP-hard and/or also limits parent size; DiBS's continuous latent variational inference circumvents these issues.Source: §3 · p. 3
  • DAG bootstrapping approximates the posterior using graphs learned from resampled datasets, but joint-posterior approximations require tractable marginal likelihoods.Exact dynamic-programming inference is feasible only up to 20 nodes, except in special cases such as tree structures or known node orderings.Source: §3 · p. 3

4 A Fully Differentiable Framework for Bayesian Structure Learning

DiBS transfers Bayesian structure inference to a continuous latent representation of graphs, enabling gradient-based inference while preserving the target posterior. Its graph model and estimators address acyclicity and likelihood-gradient constraints.

4.1 General Approach

  • The latent generative model lets inference target p(Z | D) or p(Z, Θ | D) while recovering expectations for the original graph posteriors.This reframing makes general-purpose continuous-space inference applicable to Bayesian structure learning.Source: §4.1 · pp. 3–4

4.2 Representing DAGs in a Continuous Latent Space

  • The graph model represents Z with paired embedding matrices and uses their bilinear interactions to parameterize directed-edge probabilities.The latent dimension k trades off interaction complexity and inference tractability; k ≥ d permits any adjacency matrix without self-loops.Source: §4.2 · p. 4
  • A Gibbs prior penalizes expected graph cyclicity, softly enforcing acyclicity; as β → ∞, its support reduces to latent representations of valid DAGs.The Gaussian component controls the norm of Z, and additional factors can incorporate graph priors such as sparsity.Source: §4.2 · p. 4

4.3 Estimators for Gradient-Based Bayesian Inference

  • The framework supplies scores for continuous latent posteriors, allowing gradient-based inference for both marginal graph and joint graph-parameter posteriors.When the marginal likelihood is nondifferentiable or unavailable, a score-function estimator can estimate the required gradient.Source: §4.3 · pp. 4–5
  • Monte Carlo estimates of latent gradients are tractable because graph sampling is simple and parallelizable.For differentiable likelihoods, the Gumbel-softmax estimator requires gradients with respect to G on the interior of [0, 1]^(d×d); discrete-only likelihoods can use hard Gumbel-max samples with a straight-through estimator.Source: §4.3 · p. 5

5 Particle Variational Inference for Structure Learning

DiBS turns discrete graph-structure learning into inference over a continuous latent representation, then uses particle variational inference to approximate Bayesian network posteriors. Annealing maps the latent particles back to discrete DAGs, while extensions include joint graph–parameter inference and weighted particle mixtures.

5 Particle Variational Inference for Structure Learning

  • SVGD transports latent particles with kernel-based updates to approximate marginal or joint Bayesian network posteriors.DiBS uses Gumbel-softmax when the likelihood gradient with respect to the graph is defined, and otherwise uses a score-function estimator; more elaborate orthogonal-invariant kernels performed worse in experiments.Source: §5 · p. 5
  • As α and β increase, the graph model approaches a single limiting graph and the latent prior enforces DAGs, yielding discrete particles approximating the graph or joint posterior.At termination, DiBS converts latent variables Z to G∞(Z).Source: §5 · p. 6
  • With one particle, SVGD becomes gradient ascent for a maximum a posteriori estimate; DiBS can recover methods such as NOTEARS in linear-Gaussian settings.With more particles, DiBS becomes a Bayesian approach and is not restricted to linear-Gaussian conditionals.Source: §5 · p. 6
  • DiBS+ weights particles by their unnormalized posterior probabilities, motivated by the possibility that unique particles poorly represent posterior mass in high dimensions.The authors note that they lack a strong theoretical justification for this weighting.Source: §5 · p. 6

6 Evaluation on Synthetic Data

Synthetic evaluations cover marginal and joint posterior inference on linear and nonlinear Gaussian Bayesian networks, along with analyses of DiBS variants and hyperparameters. DiBS+ performs well across the reported linear joint-inference metrics, while DiBS and DiBS+ outperform MCMC baselines on nonlinear networks.

6.1 Experimental Setup

  • The synthetic evaluation compares methods on linear and nonlinear Gaussian BNs generated from Erdős–Rényi and scale-free random graph structures.Inference uses 100 observations in settings with 20 or 50 nodes and evaluates 30 random synthetic BNs.Source: §6.1 · pp. 6–7

6.2 Linear Gaussian Bayesian Networks

  • DiBS+ performs well in marginal posterior inference of 20-node linear Gaussian BNs and appears preferable to DiBS.The comparison uses the BGe marginal likelihood.Source: Figure 2 · p. 7, §6.2 · p. 8
  • DiBS+ is the only considered method performing well across all metrics in joint inference of 20-node linear Gaussian BNs, and it slightly outperforms DiBS.BGES∗, the next-best alternative, performs substantially worse in E-SHD.Source: §6.2 · p. 8, Figure 3 · p. 8

6.3 Nonlinear Gaussian Bayesian Networks

  • DiBS and DiBS+ outperform MCMC baselines across the considered metrics for nonlinear Gaussian BNs with neural-network-parameterized conditional means.The local conditional means use two-layer dense neural networks with five hidden nodes and ReLU activations.Source: §6.3 · p. 8, Figure 4 · p. 9

6.4 DiBS with SVGD: Additional Analyses and Ablation Studies

  • Replacing the inner-product edge-probability model with scalar edge probabilities significantly worsens evaluation metrics.Reducing latent dimensionality k or iteration count T can still yield competitive posterior approximations while trading inference quality for computational efficiency.Source: §6.4 · p. 9

7 Application: Inferring Protein Signaling Networks From Cell Data

The protein-signaling application evaluates linear and nonlinear Gaussian BNs on 7,466 measurements of 11 proteins. DiBS under the BGe model provides the most calibrated edge confidence scores, while joint nonlinear DiBS is closest to the consensus network in structure.

7 Application: Inferring Protein Signaling Networks From Cell Data

  • The protein-signaling data contain 7,466 continuous measurements of 11 proteins and an established causal network of their signaling interactions.The table reports metrics as mean ± SD over 30 random restarts.Source: §7 · p. 9, Table 1 · p. 9
  • DiBS is closest in structure to the consensus network among the methods when jointly inferring nonlinear Gaussian BNs.The BGe-model DiBS posterior provides the most calibrated edge confidence scores.Source: §7 · p. 9
  • Marginal BGe posteriors average 39.0 (35.0) edges, compared with 12.7 (14.2) for joint linear and 12.6 (14.2) for joint nonlinear posteriors.The passage attributes the higher marginal-posterior edge counts to less penalization of model complexity.Source: §7 · p. 9

8 Conclusion

DiBS provides a fully differentiable approach to Bayesian-network posterior inference through a continuous latent representation of DAGs. The conclusion notes possible extensions and a quadratic-in-particle-count runtime for its SVGD implementation.

8 Conclusion

  • DiBS infers Bayesian-network posteriors through a continuous latent representation of DAGs, with posterior inference possible using black-box methods.The authors used SVGD, but note that other gradient-based sampling methods relying on the target-density score could also instantiate the approach.Source: §8 · p. 10
  • SVGD’s runtime scales quadratically with the number of sampled particles, and alternative score-based gradient samplers may improve its asymptotic runtime.The paper also expects the approach could be extended to missing and interventional data and amortized settings with rich unstructured data.Source: §8 · p. 10

A.3 Equations (16) and (17)

  • As the temperature α tends to infinity, the relaxed edge probabilities converge to a binary graph, simplifying the probability mass function and expectation.The limiting graph is the single graph implied by the latent variable Z.Source: §A.3 · p. 16

B.1 Gumbel-Softmax Estimator for the Likelihood Gradient

  • For finite temperature τ, the Gumbel-softmax likelihood-gradient estimator is biased because it relaxes the true distribution.The reparameterization estimator generally has lower variance, but requires the likelihood gradient with respect to the graph.Source: §B.1 · p. 17

B.2 Score Function Estimator for the Likelihood Gradient

  • The score-function gradient estimator is unbiased when a constant baseline is added, and the baseline can reduce Monte Carlo variance.The experiments use b = 0.Source: §B.2 · p. 18

C Background: Stein Variational Gradient Descent

  • SVGD transports a fixed set of particles toward a target distribution; its update combines movement toward high-density regions with repulsion between particles.The repulsive term helps prevent particles from collapsing into the target distribution’s modes.Source: C Background: Stein Variational… · p. 18

E.1 Gaussian Bayesian Networks

  • The nonlinear Gaussian BN specification uses neural networks for local conditionals, with parameters comprising weights and biases for each node’s network.The mask representation ensures each variable is independent of its non-descendants given its parents.Source: §E.1 · pp. 19–20

E.2 Evaluation metrics

The evaluation uses structural, predictive, and interventional metrics, with hyperparameters selected on held-out instances. It compares DiBS with structure MCMC and bootstrap-based methods.

E.2 Evaluation metrics

  • Beyond five variables, the ground-truth posterior is inaccessible, so evaluation relies on metrics established in prior work.The metrics include expected structural Hamming distance, AUROC, held-out log likelihood, and interventional log likelihood.Source: §E.2 · p. 20

E.3 Hyperparameters

  • DiBS uses 128 Monte Carlo gradient samples, and its kernel bandwidths and temperature-schedule slope are selected on held-out instances for each setting.The stated hyperparameters are expected to apply to problem settings of comparable magnitude.Source: §E.3 · p. 21, Table 2 · p. 21

E.4 Baseline Methods

  • Structure MCMC samples DAGs through single-edge changes, while nonparametric DAG bootstrap approximates a posterior by weighting learned graphs by unnormalized posterior probability.The bootstrap methods use PC or GES; starred variants use closed-form maximum-likelihood parameter estimates for linear Gaussian networks with known graphs.Source: §E.4 · pp. 21–22

F Efficient Implementation and Computing Resources

DiBS uses continuous tensors, Monte Carlo estimation, and parallel gradient updates to support vectorized and hardware-accelerated implementation. GPU runtimes for medium-sized tasks are reported on the order of seconds to minutes.

F Efficient Implementation and Computing Resources

  • DiBS’s continuous-tensor operations and Monte Carlo gradient updates support vectorization, automatic differentiation, just-in-time compilation, and hardware acceleration.The implementation uses JAX; the authors describe distributed particle updates across multiple GPUs as a possibility for larger problems.Source: F Efficient Implementation and… · p. 22
  • GPU wall times for inference tasks of up to around 20 nodes and 30 particles are on the order of seconds or a few minutes.The reported wall times are summarized in Table 3; some larger runs exceeded GPU memory.Source: Table 3 · p. 22, F Efficient Implementation and… · p. 22

G Results for Gaussian Bayesian networks with d = 50 variables

At d = 50, DiBS+ outperforms alternatives across metrics for joint inference on linear Gaussian networks, while results differ between marginal linear and joint nonlinear settings. The nonlinear comparison uses fewer DiBS particles to keep CPU wall time comparable to MC3 baselines.

G Results for Gaussian Bayesian networks with d = 50 variables

  • For marginal inference in linear Gaussian BNs, DiBS and DiBS+ are competitive on E-SHD and AUROC, while marginal-inference baselines perform favorably on likelihood-based metrics.The authors hypothesize that high variance in DiBS’s score-function estimator contributes to the likelihood-metric results and suggest more than 128 Monte Carlo samples may be needed for comparable performance.Source: Figure 5 · p. 23
  • DiBS+ outperforms all alternatives across metrics for joint posterior inference of graph structures and conditional parameters in linear Gaussian BNs.Figure 6 aggregates results over 30 random BNs for each of two graph types.Source: Figure 6 · p. 23
  • For joint inference in nonlinear Gaussian BNs, DiBS and DiBS+ perform favorably against MC3 baselines with 13,050 conditional-distribution parameters.DiBS uses 10 particles in this setting to make CPU wall time comparable to M-MC3 and G-MC3.Source: Figure 7 · p. 24

H Additional Analyses and Ablation Studies

These supplementary analyses test graph representations, inference settings, and uncertainty estimates; they compare performance while varying individual design choices.

H.1 Graph Embedding Representation

  • The bilinear graph model significantly improves performance over the scalar-per-edge variant.The authors hypothesize that coupling edges smooths the densities and may make them less prone to local minima in gradient-based methods.Source: Figure 8 · p. 24, §H.1 · p. 25

H.2 Graph Embedding Dimensionality

  • With k=15 or even k=10, performance remains competitive across metrics, while E-SHD appears most affected by reduced complexity.The bandwidth parameters were tuned for k=d=20, so alternative settings could improve lower-rank performance; full O(d^2) representations also pose computational challenges.Source: Figure 9 · p. 25, §H.2 · pp. 25–26

H.3 Particle Transport Iterations

  • Fewer particle-transport iterations still achieve competitive results, and variance in neg. LL and neg. I-LL decreases as transforms increase.Variation in E-SHD remains roughly unchanged; Figure 10 reports performance after about 1,000 iterations as well as across transport-step counts.Source: Figure 10 · p. 25, §H.3 · p. 26

H.4 Uncertainty Quantification Within a Markov Equivalence Class

  • DiBS+ correctly captures uncertainty in the nonidentifiable edge directions and high confidence in the identifiable v-structure.Unweighted DiBS particles infer too much uncertainty; DiBS+ weights particles by their unnormalized posterior probabilities.Source: Table 4 · p. 27, §H.4 · p. 27

I Experimental Details for Application to Protein Signaling Networks

The protein-signaling experiments use standardized data and specified BN-model hyperparameters. Additional results examine how those specifications affect inferred network structure and evaluation metrics.

I Experimental Details for Application to Protein Signaling Networks

  • The Sachs data are standardized, DiBS uses minibatches of 100 observations, and the graph embedding uses k=d=11.For joint linear-Gaussian inference, BPC∗ and BGES∗ use the BGe marginal likelihood, but their scores are not reported because they are very similar to BPC and BGES.Source: I Experimental Details for Application… · p. 27
  • Changing the BGe effective sample size does not significantly change metrics; at αμ=10, DiBS and DiBS+ average 39.6 and 35.4 edges, respectively.The reported metrics are means ± SD over 30 random restarts.Source: Table 5 · p. 28
  • Higher fixed noise levels yield fewer inferred edges, lower E-SHD, and worse AUROC, indicating less calibrated edge-confidence scores.These hyperparameters specify the inferred BN models rather than the posterior-inference methods.Source: I Experimental Details for Application… · p. 27, Table 6 · p. 28
Loading 2105.11839v3…