PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Structure Learning with Generative Flow Networks

Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, Yoshua Bengio

arXiv:2202.13903v2cs.LGstat.ML

TL;DR

Bayesian structure learning seeks posterior distributions over DAGs, but their combinatorially large space makes this difficult and often requires approximation. DAG-GFlowNet uses a GFlowNet to construct DAGs edge by edge and approximates the posterior, with evaluations showing accurate approximation on smaller graphs and simulated and real data.

  • Problem

    Inferring the posterior over DAG structures accounts for model uncertainty, but the combinatorially large DAG space and acyclicity constraint make this generally intractable.

  • Method

    DAG-GFlowNet uses a GFlowNet whose states are valid DAGs and constructs each graph sequentially by adding one edge at a time, with rewards proportional to the Bayesian posterior.

  • Results

    DAG-GFlowNet accurately approximates the exact posterior on smaller graphs and performs across simulated and real discrete and linear-Gaussian Bayesian networks using observational and interventional data.

  • Takeaways & Limitations

    DAG-GFlowNet offers an effective way to approximate posterior distributions over DAGs when data is limited and epistemic uncertainty matters.

  • Takeaways & Limitations

    The flow-matching condition can produce flows orders of magnitude larger near the initial state, making its parameterization challenging.

Abstract

from arXiv · show

In Bayesian structure learning, we are interested in inferring a distribution over the directed acyclic graph (DAG) structure of Bayesian networks, from data. Defining such a distribution is very challenging, due to the combinatorially large sample space, and approximations based on MCMC are often required. Recently, a novel class of probabilistic models, called Generative Flow Networks (GFlowNets), have been introduced as a general framework for generative modeling of discrete and composite objects, such as graphs. In this work, we propose to use a GFlowNet as an alternative to MCMC for approximating the posterior distribution over the structure of Bayesian networks, given a dataset of observations. Generating a sample DAG from this approximate distribution is viewed as a sequential decision problem, where the graph is constructed one edge at a time, based on learned transition probabilities. Through evaluation on both simulated and real data, we show that our approach, called DAG-GFlowNet, provides an accurate approximation of the posterior over DAGs, and it compares favorably against other methods based on MCMC or variational inference.

1 INTRODUCTION

Bayesian structure learning seeks posterior distributions over DAGs rather than single graphs, but the large constrained graph space makes exact inference difficult. DAG-GFlowNet addresses this with sequential graph generation and is evaluated against existing approximation methods.

  • Learning a DAG from data can reveal statistical or possibly causal relationships when expert knowledge does not specify the network structure.
  • Posterior inference over DAGs accounts for epistemic uncertainty, whereas single-graph methods can produce poorly calibrated predictions, especially with limited data.
  • The combinatorially large DAG space and acyclicity constraint make the full posterior generally intractable, motivating MCMC and variational approximations.
  • DAG-GFlowNet constructs graphs one edge at a time as a sequential decision process using learned transitions to approximate the posterior over DAGs.
  • DAG-GFlowNet samples iid. DAGs from scratch and restricts its support to valid DAGs, unlike some variational methods that may sample cyclic graphs.
  • The method is evaluated on simulated and real data, including discrete and linear-Gaussian networks and observational and interventional settings.

2 RELATED WORK

Related work approximates Bayesian network structure posteriors mainly through MCMC, variational inference, or sequential graph-generation procedures, each addressing the discrete DAG space differently.

  • Markov chain Monte Carlo: MCMC methods explore DAGs through local edge additions and removals, but direct DAG-space sampling can mix slowly; node-order sampling improves mixing at the cost of bias.
  • Variational Inference: Variational approaches face difficulties from discrete structure and encode DAG distributions using continuous relaxations or soft acyclicity priors.
  • Sequential decisions: Sequential graph-generation methods construct graphs through actions, while DAG-GFlowNet applies this idea without assuming a fixed node ordering.

3 BACKGROUND

GFlowNets generate structured objects by sequential transitions whose learned flows match rewards, producing complete samples with probability proportional to reward. Detailed balance offers a transition-based alternative to flow matching.

  • Generative Flow Networks: GFlowNets define distributions over discrete structured objects by constructing samples sequentially from an initial state to a terminal state.
  • Generative Flow Networks: Complete states receive nonnegative rewards, while incomplete states conventionally receive zero reward.
  • Flow matching: The flow-matching condition equates flow entering a state with flow leaving it plus the state's residual reward, and its parameters can be learned by regression.
  • Sampling: When the flow conditions hold, sampling transitions from the initial state yields complete states with probability proportional to their rewards.
  • Sampling: Unlike MCMC, GFlowNet samples are constructed from scratch, so irreducibility of the underlying Markov process is unnecessary when complete states remain reachable.
  • Detailed balance: Flow matching can create very large flows near the initial state, making its parameterization difficult.
  • Detailed balance: Detailed balance directly parameterizes forward and backward transition probabilities to enforce reversibility, with backward transitions distributed over parent states.
  • Detailed balance: Satisfying detailed balance preserves reward-proportional sampling; fixing backward probabilities can reduce the search space so only forward transitions are learned.

4 GFLOWNET OVER DIRECTED ACYCLIC GRAPHS

DAG-GFlowNet represents valid DAGs as GFlowNet states and constructs them sequentially by adding edges while filtering actions that would create cycles. Neural transition models separately learn whether to terminate or which valid edge to add.

  • 4.1 STRUCTURE OF THE GFLOWNET: The GFlowNet state space consists of valid DAGs, starting from the completely disconnected graph and adding one edge per transition.A graph with k edges can be reached through k! edge-addition orders.
  • 4.1 STRUCTURE OF THE GFLOWNET: Acyclicity is enforced by masking edges that are already present or would introduce a cycle.The mask is built from the adjacency matrix and the transitive closure of its transpose.
  • 4.2 FORWARD TRANSITION PROBABILITIES: Forward transitions use a hierarchical neural model with separate heads for terminating and adding a new graph edge.The two probabilities are Pθ(sf | G) and Pθ(G′ | G, ¬sf).
  • 4.2 FORWARD TRANSITION PROBABILITIES: The edge-addition head assigns probabilities over d2 possible directed edges, then masks and renormalizes invalid actions.Self-loops are included in the candidate action set but are always invalid.
  • 4.3 PARAMETRIZATION WITH LINEAR TRANSFORMERS: A Linear Transformer processes edge-set representations to provide permutation invariance, edge-level outputs, and reduced scaling with graph size.Each edge is represented using source, target, and presence embeddings; linearized attention uses Q, K, and V transformations.

5 APPLICATION TO BAYESIAN STRUCTURE LEARNING

DAG-GFlowNet approximates the Bayesian posterior over DAG structures by assigning each graph a prior-times-marginal-likelihood reward and training transitions with detailed balance. Its expected loss uses sampled transitions, including replay-buffer data and exploratory random transitions.

  • 5 APPLICATION TO BAYESIAN STRUCTURE LEARNING: The method targets the posterior P(G | D) over Bayesian-network DAGs rather than returning a single graph.This formulation accounts for epistemic uncertainty, which is especially relevant when data are limited.
  • 5 APPLICATION TO BAYESIAN STRUCTURE LEARNING: The reward R(G) = P(G)P(D | G) makes GFlowNet sampling proportional to the Bayesian posterior.P(G) is the structure prior and P(D | G) is the marginal likelihood.
  • 5.1 MODULARITY & COMPUTATIONAL EFFICIENCY: Modular structure and parameter priors make the reward modular, so log R(G) decomposes into local scores for variables and their parents.The resulting score corresponds to Bayesian scores such as BDe and BGe.
  • 5.1 MODULARITY & COMPUTATIONAL EFFICIENCY: Because the DAG state space is superexponential, the detailed-balance loss is minimized in expectation over a full-support transition distribution.The sampling distribution may be on-policy, provided every possible next state receives nonzero probability.
  • 5.1 MODULARITY & COMPUTATIONAL EFFICIENCY: Off-policy training stores transitions and their delta scores in a replay buffer, with uniformly random transitions added with probability ε for exploration.Mini-batches sampled from the buffer are used to estimate the loss and update θ.

6 EXPERIMENTAL RESULTS

Experiments evaluate DAG-GFlowNet against exact posteriors, MCMC, bootstrapping, and variational methods on simulated and real Bayesian-network data. The results show accurate posterior approximation, competitive structural performance, predictive likelihood near the ground-truth graph, and diversity across DAGs and Markov equivalence classes.

  • 6.1 COMPARISON WITH THE EXACT POSTERIOR: DAG-GFlowNet’s estimated structural-feature probabilities strongly correlate with exact posterior marginals across 20 randomly generated Bayesian networks.The evaluation considers edge, path, and Markov-blanket features for five-node linear-Gaussian networks.
  • 6.2 SIMULATED DATA: On 20-node graphs, DAG-GFlowNet is competitive with all evaluated methods on E-SHD and AUROC, including MCMC-based methods.E-SHD measures structural distance to the ground-truth graph, while AUROC evaluates edge features.
  • 6.2 SIMULATED DATA: On held-out data, DAG-GFlowNet produces predictive log-likelihood concentrated near the ground-truth DAG’s log-likelihood.The comparison uses log P(G, D′ | D) on a held-out dataset D′.
  • 6.3 APPLICATION: FLOW CYTOMETRY DATA: On flow cytometry data, DAG-GFlowNet offers a trade-off between E-SHD, AUROC, and assigning higher probability to DAGs with more edges.BCD Nets and Bootstrap PC obtain smaller E-SHD but tend to sample fewer-edge graphs; 1.50% of DiBS samples contained a cycle.
  • 6.3 APPLICATION: FLOW CYTOMETRY DATA: DAG-GFlowNet covers multiple high-scoring MECs and samples multiple equivalent DAGs within MECs rather than collapsing to one representative.Its maximum-a-posteriori MEC scores higher than Gadget’s and lower than MC3’s; the best GES MEC scores −10,716.12.
  • 6.4 APPLICATION: INTERVENTIONAL DATA: For interventional and observational flow-cytometry data, DAG-GFlowNet is compared with an exact-posterior AUROC upper bound and MC3 using a modified BDe score.The exact-posterior method provides edge marginals but not a distribution over DAGs.

7 CONCLUSION

The paper presents DAG-GFlowNet as a GFlowNet-based method for approximating posterior distributions over DAGs, with enhancements designed for Bayesian structure learning. It is intended for settings where data are limited and epistemic uncertainty matters, while its current performance can degrade as dataset size increases.

  • DAG-GFlowNet treats sample-graph generation as a sequential decision problem within a GFlowNet framework.
  • The method introduces enhancements to standard GFlowNets specifically for approximating posterior distributions over DAGs P(G | D).
  • DAG-GFlowNet offers an effective solution when data are limited and measuring epistemic uncertainty is critical.
  • The current method may suffer limitations as the size of dataset D increases.
  • Future work may adapt the GFlowNet to essential graphs and study causal discovery with latent variables.

(Supplementary material)

The supplementary material identifies the paper’s authors and their affiliations across Mila, Université de Montréal, the Technical University of Munich, and KTH Stockholm.

  • The paper is authored by Tristan Deleu, António Góis, Chris Emezue, and Mansi Rankawat.
  • The author list also includes Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio.
  • The affiliations include Mila, Université de Montréal, the Technical University of Munich, and KTH Stockholm.
  • The listed authors hold CIFAR roles including AI Chair, Azrieli Global Scholar, and Senior Fellow.

A LIMITATIONS OF DAG-GFLOWNET

DAG-GFlowNet approximates the posterior accurately for moderate dataset sizes, but training becomes more difficult as datasets grow because the posterior becomes sharply peaked and score changes widen.

  • DAG-GFlowNet accurately approximates P(G | D) when dataset size is moderate.
  • As dataset size increases, fitting the detailed-balance loss becomes more challenging.
  • Larger datasets make the posterior very peaky, so adding an edge can drastically increase or decrease its score.
  • These large delta-score fluctuations require the transition network to compensate, making training harder.
  • Input normalization cannot be applied because normalizing rewards would change the approximated distribution to P(G | D)τ.
  • Temperature schedules and transition reparameterization are proposed as future directions for handling large delta-score fluctuations.

B DETAILED-BALANCE CONDITION WITH ALL COMPLETE STATES

The supplementary material explains a detailed-balance condition for complete GFlowNet states and describes efficient legality masking for DAG construction. It also defines evaluation metrics and presents the d = 50-node experimental setting.

  • Detailed-balance condition: For complete GFlowNet states, the detailed-balance condition uses forward transitions, backward transitions, rewards, and terminal-state probabilities.
  • Detailed-balance condition: The alternative detailed-balance condition requires parameterizing only the forward transition probability, including transitions to the terminal state.
  • Online DAG mask updates: After adding C →A, the adjacency and transitive-closure components are updated separately and recombined.
  • Online DAG mask updates: The DAG mask excludes edges already present and edges whose addition would create a cycle, using adjacency and transitive-closure matrices.
  • Online DAG mask updates: The two mask-update operations can be performed in O(d^2), where d is the number of DAG nodes.
  • Evaluation and Figure 7: Evaluation uses expected SHD and AUROC, with SHD counting edge additions, removals, and reversals needed to reach the ground-truth graph.
  • Evaluation and Figure 7: For d = 50 linear-Gaussian networks, E-SHD and AUROC aggregate over 10 datasets, while log P(G, D′ | D) uses one dataset and compares against the ground-truth graph.

D.2 SIMULATED DATA

The simulated-data experiments evaluate DAG-GFlowNet on graphs of varying size using synthetic linear-Gaussian Bayesian networks generated from random DAGs and datasets of 100 observations.

  • For larger simulated graphs with d = 50 nodes, DAG-GFlowNet remained competitive with other algorithms on E-SHD, AUROC, and joint log-likelihood.The comparison used a held-out dataset D′ and the same experimental setup as the d = 20-node experiments.
  • Synthetic graphs were sampled from an Erdős-Rényi model with 2d edges on average.
  • The simulated-data generation procedure began by sampling a DAG, then its parameters, and finally observations from the resulting Bayesian network.
  • The sampled graph parameters followed a linear-Gaussian model with coefficients drawn from N(0, 1) and noise from N(0, 0.01).
  • Each completed Bayesian network generated a dataset D of N = 100 datapoints through ancestral sampling.

D.3 FLOW CYTOMETRY DATA

On real flow cytometry data, DAG-GFlowNet captured high-scoring posterior regions while retaining diversity across Markov equivalence classes and DAGs within classes. Comparisons showed distinct diversity and scoring weaknesses for variational methods.

  • DAG-GFlowNet sampled both high-scoring graphs and diverse graph structures across Markov equivalence classes and unique DAGs within a class.In Figure 5, the number of unique DAGs within each equivalence class is represented by point size.
  • The flow cytometry comparison included DAG-GFlowNet, MCMC methods, BCD Nets, and DiBS to assess posterior approximation quality.
  • BCD Nets returned only 2 unique DAGs among 1,000 samples, concentrated in two Markov equivalence classes.Both classes had BGe scores around −10,950, indicating low diversity in the learned posterior approximation.
  • DiBS sampled very low-scoring DAGs, reaching BGe scores as low as −12,600.The best MEC obtained with GES had a BGe score of −10,716.12.
  • DiBS covered many Markov equivalence classes but sampled only unique DAGs per class, unlike the equal within-class probabilities implied by the BGe posterior.
Loading 2202.13903v2…