Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network

Tristan DeleuMizu Nishikawa-ToomeyJithendaraa SubramanianNikolay MalkinLaurent CharlinYoshua Bengio

article2023NeurIPS56 citations

Presents JSP-GFN, a single Generative Flow Network that jointly infers both the graph structure and continuous parameters of Bayesian networks through a two-phase sampling process, enabling tractable posterior inference for expressive non-linear models without intractable marginalizations.

Listen

Understanding complex causal and statistical dependencies across variables is critical for decision-making in domains like medical diagnostics and computational biology. While Bayesian networks provide a principled representation for these relationships, discovering both network structure and continuous mechanism parameters from limited observational data is exceptionally difficult. Traditional methods typically force restrictive linear assumptions or infer only network structures, failing to account for parameter uncertainty in expressive, non-linear models.

The article introduces and evaluates JSP-GFN (Joint Structure and Parameters Generative Flow Network), a probabilistic machine learning approach that infers the joint posterior distribution over both directed network structures and continuous distribution parameters. The primary objective is to demonstrate that a single generative framework can accurately quantify uncertainty and scale to complex non-linear relationships, discrete data, and interventions.

The approach models structure and parameter discovery as a unified sequential generation process. The network first builds an acyclic graph step by step and then generates the corresponding mechanism parameters conditional on that graph. By establishing generalized balance conditions for the network's learning flow, the system is optimized without intractable calculations. The authors evaluated the method across simulated benchmarks—from 5-variable linear networks to 20-variable non-linear systems with over 2,200 parameters—and real-world biological datasets, including protein signaling and gene regulatory expression data.

The findings show that JSP-GFN achieves state-of-the-art accuracy in recovering posterior distributions. On small benchmark graphs, the method reduced edge feature approximation error by roughly a factor of ten compared to baseline variational and Markov chain Monte Carlo methods (achieving a Pearson correlation of 0.998). On larger simulated models, it delivered superior or competitive predictive accuracy on held-out test data. Additionally, on high-dimensional gene expression benchmarks across 61 variables, JSP-GFN substantially outperformed standard sampling algorithms in predictive likelihood while supporting efficient mini-batch training.

These results establish that organizations can deploy flexible, neural-network-parameterized causal models without sacrificing rigorous uncertainty quantification. Enabling mini-batch training reduces the computational barriers to running Bayesian structure learning on large datasets, significantly lowering deployment and inference overheads while guarding against overconfident, erroneous structural conclusions.

Organizations analyzing complex biological, diagnostic, or transactional dependency structures should consider piloting joint inference approaches like JSP-GFN in place of single-estimate graph discovery algorithms. Further development should focus on testing the methodology in domains with cyclic feedback processes and incorporating multimodal parameter distributions—such as normalizing flows or diffusion models—to capture multiple competing mechanism hypotheses.

The main limitation lies in assuming unimodal parameter distributions and strictly acyclic graphs, which may simplify complex multimodal parameters or biological feedback loops. Nevertheless, the theoretical consistency and empirical validation provide strong confidence that the framework reliably discovers joint structural and parametric uncertainties across varied domains.

HEC Montréal McGill University Mila – Québec Artificial Intelligence Institute Université de Montréal bayesian inference conditional probabilities directed acyclic graphs Generative Flow Network joint Bayesian inference Bayesian structure learning

arXiv: 2305.19366 tristandeleu/jax-jsp-gfn

Paper: GFlowNet Foundations, Yoshua Bengio et al. (2023). It establishes the theoretical foundations and balance conditions of Generative Flow Networks, which provide the core sequential generative sampling mechanism used by JSP-GFN.
Paper: Generative Flow Networks for Discrete Probabilistic Modeling, Dinghuai Zhang et al. (2022). It introduces energy-based GFlowNets for discrete probabilistic modeling and posterior approximation, serving as the immediate precursor for applying GFNs to structured probabilistic objects.
Paper: DAGs with NO TEARS: Continuous Optimization for Structure Learning, Xun Zheng et al. (2018). It establishes modern benchmarks and formulations for continuous structure learning of directed acyclic graphs against which Bayesian structure learning methods are compared.
Paper: Learning Bayesian networks: The combination of knowledge and statistical data, David Heckerman et al. (1994). It provides foundational principles and metrics for Bayesian network structure and parameter learning from combined prior knowledge and observational data.
Paper: A Bayesian method for the induction of probabilistic networks from data, Gregory F. Cooper et al. (1992). It introduces seminal Bayesian scoring techniques and heuristic search methods for inferring belief network structures directly from discrete data.
Paper: A Tutorial on Learning with Bayesian Networks, David Heckerman (1999). It delivers a comprehensive tutorial on parameter estimation, causal inference, and marginal likelihood scoring in Bayesian networks.
Paper: Optimal Structure Identification With Greedy Search, David Maxwell Chickering (2002). It provides the theoretical framework for score-based search over Markov equivalence classes of directed acyclic graphs.
Paper: Variational Inference with Normalizing Flows, Danilo Jimenez Rezende et al. (2015). It presents normalizing flows for expressive approximate posterior inference, establishing continuous generative techniques relevant to non-linear parameter modeling.

Paper: Local Search GFlowNets, Minsu Kim et al. (2024). It develops local search training techniques for GFlowNets to alleviate mode exploration issues and improve sampling efficiency in large structured state spaces.
Paper: Better Training of GFlowNets with Local Credit and Incomplete Trajectories, Ling Pan et al. (2023). It improves GFlowNet training and credit assignment using local credit and incomplete trajectories, offering enhancements applicable to complex multi-step generative flows.
Paper: Nonparametric Identifiability of Causal Representations from Unknown Interventions, Julius von Kügelgen et al. (2023). It investigates nonparametric identifiability of causal structures under interventional distributions, extending the causal modeling principles explored in joint Bayesian inference.
Paper: Multi-View Causal Representation Learning with Partial Observability, Dingling Yao et al. (2024). It generalizes causal representation learning to multi-view systems with partial observability, expanding upon joint structure and parameter discovery in complex environments.
Paper: Sequential Neural Score Estimation: Likelihood-Free Inference with Conditional Score Based Diffusion Models, Louis Sharrock et al. (2024). It introduces conditional score-based diffusion methods for likelihood-free posterior parameter inference in complex simulator models.
Paper: An Introduction to Bayesian and Frequentist Simulation-Based Inference with Machine Learning, Maximilian Dax et al. (2026). It provides an extensive treatment of amortized Bayesian simulation-based inference using deep generative neural networks.
Paper: Identifying Weight-Variant Latent Causal Models, Yuhang Liu et al. (2026). It formulates latent causal model discovery under varying structural mechanisms and auxiliary variables, building on non-linear parameterization challenges.

Cover for Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network

Abstract

Generative Flow Networks (GFlowNets), a class of generative models over discrete and structured sample spaces, have been previously applied to the problem of inferring the marginal posterior distribution over the directed acyclic graph (DAG) of a Bayesian Network, given a dataset of observations. Based on recent advances extending this framework to non-discrete sample spaces, we propose in this paper to approximate the joint posterior over not only the structure of a Bayesian Network, but also the parameters of its conditional probability distributions. We use a single GFlowNet whose sampling policy follows a two-phase process: the DAG is first generated sequentially one edge at a time, and then the corresponding parameters are picked once the full structure is known. Since the parameters are included in the posterior distribution, this leaves more flexibility for the local probability models of the Bayesian Network, making our approach applicable even to non-linear models parametrized by neural networks. We show that our method, called JSP-GFN, offers an accurate approximation of the joint posterior, while comparing favorably against existing methods on both simulated and real data.

Topics

bayesian inference conditional probabilities directed acyclic graphs generative model inference linear models marginals neural networks posterior distributions network architecture

Keywords

Generative Flow Network joint Bayesian inference Bayesian structure learning Bayesian graphical models joint posterior distribution causal discovery variational inference marginal likelihood

1 Introduction
2 Background
2.1 Bayesian structure learning
2.2 Generative Flow Networks
2.3 Structure learning with GFlowNets
3 Joint Bayesian inference of structure and parameters
3.1 Structure of the GFlowNet
3.2 Subtrajectory Balance conditions
3.3 Learning objective
3.4 Parametrization of the forward transition probabilities
3.5 Mini-batch training
4 Related work
5 Experimental results
5.1 Joint posterior over small graphs
5.2 Gaussian Bayesian Networks from simulated data
5.3 Learning biological structures from real data
6 Conclusion
Acknowledgements
References
Appendix
A Positioning JSP-GFN in the Bayesian structure learning literature
B Broader impact & limitations
B.1 Broader impact
B.2 Limitations
C Details about Generative Flow Networks
C.1 Alternative conditions
C.2 Subtrajectory balance conditions
C.3 Proofs
C.3.1 Subtrajectory balance conditions for undirected paths of length 3
C.3.2 Integrating undirected paths of length 2
C.3.4 Mini-batch training
D Additional experiments & experimental details
D.1 Sampling distribution
D.2 Joint posterior over small graphs
D.2.1 Data generation & modeling
D.2.2 Comparing JSP-GFN with the exact posterior against features
D.2.3 Evaluation of the posterior approximations over parameters
D.3 Gaussian Bayesian Networks from simulated data
D.3.1 Data generation & modeling
D.3.2 Estimation of the log-terminating state probability
D.3.3 Additional comparisons with the ground-truth graphs
D.4 Learning biological structures from real data
D.4.1 Modeling
D.5 Proofs
D.5.1 Posterior of the linear Gaussian model

Knowls

Knowl 1 — JSP-GFN Framework for Joint Structure and Parameter Inference
model/method
Joint Structure and Parameters Generative Flow Network (JSP-GFN) models the joint posterior distribution $P(G, \theta \mid \mathcal{D})$ over both the directed acyclic graph (DAG) structure $G = (V, E)$ of a Bayesian Network over $d$ random variables $X = \{X_1, \dots, X_d\}$ and the continuous parameters $\theta = \{\theta_1, \dots, \theta_d\}$ of its conditional probability distributions (CPDs).

The state space $\mathcal{S}$ consists of intermediate graph states $(G, \cdot) \in \mathcal{S} \setminus \mathcal{X}$ and complete states $(G, \theta) \in \mathcal{X}$ . Generation proceeds in a two-phase sequential process:
1. Structure Generation: Starting from the empty graph $(G_0, \cdot)$ , directed edges are sequentially added one at a time following forward transition probabilities $P_\phi(G' \mid G)$ while enforcing acyclicity, until a dedicated "stop" action is selected.
2. Parameter Generation: Once the "stop" action terminates structure generation at DAG $G$ , parameters $\theta$ are sampled in a single step conditioned on $G$ following $P_\phi(\theta \mid G)$ .
The unnormalized target reward for a complete state $(G, \theta)$ is defined as the unnormalized joint posterior: $R(G, \theta) = P(\mathcal{D} \mid \theta, G) P(\theta \mid G) P(G)$ where $P(\mathcal{D} \mid \theta, G) = \prod_{n=1}^N \prod_{i=1}^d P(x_i^{(n)} \mid \text{Pa}_G(X_i); \theta_i)$ for dataset $\mathcal{D} = \{x^{(1)}, \dots, x^{(N)}\}$ , $P(\theta \mid G)$ is the parameter prior, and $P(G)$ is the graph prior.
Knowl 2 — Subtrajectory Balance Conditions for Joint Posterior Approximation
theoretical result
In JSP-GFN, complete states $(G, \theta)$ have a single parent $(G, \cdot)$ and terminate directly into the terminal state $s_f$ with probability $P_\phi(s_f \mid G, \theta) = 1$ . For any length-3 undirected path of the form $(G, \theta) \leftarrow (G, \cdot) \rightarrow (G', \cdot) \rightarrow (G', \theta')$ , where $G'$ is formed by adding one directed edge to $G$ , the Subtrajectory Balance (SubTB) condition is given by: $R(G', \theta') P_B(G \mid G') P_\phi(\theta \mid G) = R(G, \theta) P_\phi(G' \mid G) P_\phi(\theta' \mid G')$ where $P_B(G \mid G')$ is a fixed backward transition probability over parent graphs of $G'$ , $P_\phi(G' \mid G)$ is the forward transition probability of adding the edge, $P_\phi(\theta \mid G)$ is the forward probability of generating $\theta$ given $G$ , and $R(G, \theta) = P(\mathcal{D} \mid \theta, G) P(\theta \mid G) P(G)$ is the unnormalized joint reward.

Theorem (Terminating State Distribution): If the length-3 SubTB conditions are satisfied for all pairs $(G, \theta)$ and $(G', \theta')$ connected by valid single-edge additions $G \to G'$ , then the terminating state probability of the GFlowNet satisfies: $P_\phi^\top(G, \theta) \triangleq P_\phi(G \mid G_0) P_\phi(\theta \mid G) \propto R(G, \theta)$ where $P_\phi(G \mid G_0) = \sum_{\tau: G_0 \rightsquigarrow G} \prod_{t=0}^{T-1} P_\phi(G_{t+1} \mid G_t)$ is the marginal probability of reaching DAG $G$ from the empty graph $G_0$ over all valid edge-addition trajectories $\tau = (G_0, G_1, \dots, G_T = G)$ .
Knowl 3 — JSP-GFN Learning Objective with Stop-Gradient and Score Matching
model/method
For length-2 undirected paths $(G, \theta) \leftarrow (G, \cdot) \rightarrow (G, \tilde{\theta})$ sharing the same graph $G$ , the Subtrajectory Balance condition implies that $\log R(G, \theta) - \log P_\phi(\theta \mid G)$ is constant with respect to $\theta$ . When $R(G, \theta)$ is differentiable with respect to $\theta$ , this yields the differential score-matching condition: $\nabla_\theta \log P_\phi(\theta \mid G) = \nabla_\theta \log R(G, \theta)$

To satisfy length-3 SubTB while aligning optimization gradients with this differential score-matching condition (preventing path-derivative issues through sampled parameters), the GFlowNet parameters $\phi$ are trained by minimizing the squared log-ratio loss with stop-gradient operators $\bot$ on continuous parameter inputs: $\mathcal{L}(\phi) = \mathbb{E}_\pi \left[ \left( \log \frac{R(G', \bot(\theta')) P_B(G \mid G') P_\phi(\bot(\theta) \mid G)}{R(G, \bot(\theta)) P_\phi(G' \mid G) P_\phi(\bot(\theta') \mid G')} \right)^2 \right]$ where $\pi$ is a sampling distribution over transitions and parameter samples with full support.

To explicitly encourage parameter score matching in larger or non-linear models, an auxiliary penalty can be added: $\tilde{\mathcal{L}}(\phi) = \mathcal{L}(\phi) + \frac{\lambda}{2} \mathbb{E}_\pi \left[ \|\nabla_\theta \log P_\phi(\theta \mid G) - \nabla_\theta \log R(G, \theta)\|^2 + \|\nabla_{\theta'} \log P_\phi(\theta' \mid G') - \nabla_{\theta'} \log R(G', \theta')\|^2 \right]$
Knowl 4 — Forward Transition Parameterization via Graph Neural Networks
model/method
The forward transition probabilities of JSP-GFN are decomposed hierarchically into structure and parameter generation steps: $P_\phi(G' \mid G) = (1 - P_\phi(\text{stop} \mid G)) P_\phi(G' \mid G, \neg \text{stop})$ $P_\phi(\theta \mid G) = P_\phi(\text{stop} \mid G) P_\phi(\theta \mid G, \text{stop})$

A neural network backbone comprising a Graph Network and self-attention blocks processes DAG $G$ into a global graph embedding $g$ and per-node embeddings $u_i, v_i, w_i$ for each node $X_i$ ( $i = 1, \dots, d$ ): $g, \{u_i, v_i, w_i\}_{i=1}^d = \text{SelfAttention}_\phi(\text{GraphNet}_\phi(G))$

The individual forward policies are defined as follows:
1. Stop Decision: $P_\phi(\text{stop} \mid G) = f_\phi(g) \in (0, 1)$ via a sigmoid MLP (forced to $1$ if adding any edge would introduce a cycle).
2. Edge Addition: Adding directed edge $X_i \to X_j$ has probability: $P_\phi(G' \mid G, \neg \text{stop}) \propto m_{ij} \exp(u_i^\top v_j)$ where $m_{ij} \in \{0, 1\}$ is a binary mask indicating whether $X_i \to X_j \notin G$ and whether adding $X_i \to X_j$ preserves acyclicity.
3. CPD Parameters Distribution: Factorizes independently across variables as $P_\phi(\theta \mid G, \text{stop}) = \prod_{i=1}^d P_\phi(\theta_i \mid G, \text{stop})$ , parameterized as a multivariate Gaussian: $P_\phi(\theta_i \mid G, \text{stop}) = \mathcal{N}(\theta_i \mid \mu_\phi(w_i), \sigma_\phi^2(w_i))$ where $\mu_\phi$ and $\sigma_\phi^2$ output mean vectors and positive diagonal variance matrices.
Knowl 5 — Mini-Batch Estimation and Unbiased Loss Gradients
theoretical result
For a dataset $\mathcal{D} = \{x^{(1)}, \dots, x^{(N)}\}$ , evaluating the full likelihood $P(\mathcal{D} \mid \theta, G)$ is computationally expensive for large $N$ . For a mini-batch $\mathcal{B} \subset \mathcal{D}$ of $M$ observations sampled uniformly at random, the estimated log-reward is: $\log \hat{R}_\mathcal{B}(G, \theta) = \log P(\theta \mid G) + \log P(G) + \frac{N}{M} \sum_{x^{(m)} \in \mathcal{B}} \log P(x^{(m)} \mid G, \theta)$

Properties:
1. Unbiased Log-Reward: Because observations are mutually independent conditioned on both $G$ and $\theta$ , $\mathbb{E}_\mathcal{B}[\log \hat{R}_\mathcal{B}(G, \theta)] = \log R(G, \theta)$ .
2. Loss Upper Bound: For the mini-batch objective $\hat{\mathcal{L}}_\mathcal{B}(\phi) = \mathbb{E}_\pi [\hat{\Delta}_\mathcal{B}^2(\phi)]$ , convexity of the square loss implies $\mathcal{L}(\phi) \le \mathbb{E}_\mathcal{B}[\hat{\mathcal{L}}_\mathcal{B}(\phi)]$ .
3. Unbiased Gradient Estimator: The gradient of the mini-batch loss is unbiased: $\nabla_\phi \mathcal{L}(\phi) = \mathbb{E}_\mathcal{B}[\nabla_\phi \hat{\mathcal{L}}_\mathcal{B}(\phi)]$ Consequently, the expected mini-batch objective and full-batch objective share identical sets of local and global minima.
Knowl 6 — JSP-GFN Training Procedure
algorithm
The training procedure of JSP-GFN uses off-policy exploration for graph transitions stored in a replay buffer $\mathcal{R}$ combined with on-policy generation of continuous parameters $\theta, \theta'$ and lazy reward evaluation.
Input: Dataset $\mathcal{D}$ , learning rate $\alpha$ , rollout steps $K$ , transition batch size $|\mathcal{B}|$
Initialize current graph $G_t \leftarrow G_0$ (empty graph)
Initialize replay buffer $\mathcal{R} \leftarrow \emptyset$
repeat
    for step = 1 to $K$ do
        Sample action decision $a \sim P_\phi(\text{stop} \mid G_t)$
        if $a$ is the "stop" action then
            Reset trajectory: $G_{t+1} \leftarrow G_0$
        else
            Sample next graph $G_{t+1} \sim P_\phi(G_{t+1} \mid G_t, \neg\text{stop})$
            Store transition $(G_t \to G_{t+1})$ in replay buffer $\mathcal{R}$
             $G_t \leftarrow G_{t+1}$
        end if
    end for
    Sample a mini-batch of graph transitions $\mathcal{B} = \{(G \to G')\}$ from $\mathcal{R}$
    for each transition $(G \to G')$ in $\mathcal{B}$ do
        Sample parameters $\theta \sim P_\phi(\theta \mid G, \text{stop})$ on-policy
        Sample parameters $\theta' \sim P_\phi(\theta' \mid G', \text{stop})$ on-policy
        Evaluate rewards $R(G, \theta)$ and $R(G', \theta')$ (or mini-batch estimates $\hat{R}_\mathcal{B}$ )
    end for
    Compute loss $\mathcal{L}(\phi)$ over batch $\mathcal{B}$ with stop-gradients on $\theta, \theta'$
    Update model parameters: $\phi \leftarrow \phi - \alpha \nabla_\phi \mathcal{L}(\phi)$
until convergence criterion is met
Knowl 7 — Exact Parameter Posterior for Linear Gaussian Bayesian Networks
theoretical result
In a linear Gaussian Bayesian network over variables $\{X_1, \dots, X_d\}$ with fixed noise variance $\sigma^2$ and independent prior $P(\theta_{ij} \mid G) = \mathcal{N}(0, \sigma_0^2)$ for $X_j \in \text{Pa}_G(X_i)$ (where $\sigma_0^2 = 1$ ), the conditional distribution of variable $X_i$ is $P(X_i \mid \text{Pa}_G(X_i); \theta_i) = \mathcal{N}(\sum_{j=1}^d \mathbb{I}(X_j \in \text{Pa}_G(X_i)) \theta_{ij} X_j, \sigma^2)$ .

Let $X \in \mathbb{R}^{N \times d}$ be the observation matrix, $X_i$ the $i$ -th column of $X$ , and $D_i = \text{diag}(\mathbb{I}(X_1 \in \text{Pa}_G(X_i)), \dots, \mathbb{I}(X_d \in \text{Pa}_G(X_i)))$ . The exact conditional posterior distribution over parameters $\theta_i$ given DAG $G$ and dataset $\mathcal{D}$ is Gaussian: $P(\theta_i \mid G, \mathcal{D}) = \mathcal{N}(\bar{\mu}_i, \bar{\Sigma}_i)$ where $\bar{\Sigma}_i^{-1} = \frac{1}{\sigma_0^2} D_i^{-1} + \frac{1}{\sigma^2} D_i X^\top X D_i$ $\bar{\mu}_i = \bar{\Sigma}_i \left( \frac{1}{\sigma^2} D_i X^\top X_i \right)$ with conventions $1/0 = \infty$ and $0 \times \infty = 0$ , reducing active dimensions to $|\text{Pa}_G(X_i)|$ .
Knowl 8 — Estimation of Log-Terminating Probability via Beam Search and Monte Carlo
model/method
To evaluate the quality of posterior approximation on large graphs where exact marginalization is intractable, the terminating state log-probability $\log P_\phi^\top(G, \theta) = \log P_\phi(\theta \mid G) + \log P_\phi(G \mid G_0)$ is estimated and compared against the log-reward $\log R(G, \theta)$ .

Because computing $P_\phi(G \mid G_0) = \sum_{\tau: G_0 \rightsquigarrow G} P_\phi(\tau)$ involves $K!$ possible edge-addition permutations for a graph with $K$ edges, a variance-reduced unbiased estimator combines a set $\mathcal{T}_{\text{top}}$ of $B$ high-scoring trajectories identified by beam search with $M$ trajectories $\{\tau^{(m)}\}_{m=1}^M$ sampled uniformly from the remaining trajectories: $P_\phi(G \mid G_0) \approx \sum_{\tau \in \mathcal{T}_{\text{top}}} P_\phi(\tau) + \frac{K! - B}{M} \sum_{m=1}^M P_\phi(\tau^{(m)})$ Beam search finds trajectories in $\mathcal{T}_{\text{top}}$ from $G_0$ to $G$ using the masked forward scoring function $\tilde{P}_\phi(G_{t+1} \mid G_t) = \mathbb{I}(e \in G) P_\phi(G_{t+1} \mid G_t)$ , which only allows edges that belong to the target graph $G$ .

Knowl 9 — Posterior Feature and Parameter Cross-Entropy Benchmark on Small Graphs

data/table

Across 20 synthetic datasets of $N = 100$ observations generated from linear Gaussian DAGs ( $d = 5$ , ER1 topology), JSP-GFN (with diagonal or full covariance parameter distributions) was compared against exact posterior marginals and baseline methods. Graph posterior marginals were evaluated via Root Mean Square Error (RMSE) and Pearson's correlation coefficient $r$ across edge, path, and Markov blanket features. Parameter posterior approximation quality was evaluated via expected negative cross-entropy $\mathbb{E}_{G, \theta}[-\log P(\theta \mid G, \mathcal{D})]$ .

Method	Edge features		Path features		Markov features		Parameter Cross-Entropy
	RMSE	Pearson's $r$	RMSE	Pearson's $r$	RMSE	Pearson's $r$	$\mathbb{E}_{G, \theta}[-\log P(\theta \mid G, \mathcal{D})]$
MH-MC $^3$	$0.357 \pm 0.022$	$0.067 \pm 0.143$	$0.368 \pm 0.027$	$0.045 \pm 0.179$	$0.341 \pm 0.017$	$0.064 \pm 0.217$	$5.39 \pm 1.41 \times 10^2$
Gibbs-MC $^3$	$0.357 \pm 0.022$	$0.028 \pm 0.127$	$0.367 \pm 0.026$	$0.150 \pm 0.162$	$0.341 \pm 0.018$	$0.062 \pm 0.159$	$9.02 \pm 1.54 \times 10^5$
B-GES*	$0.263 \pm 0.070$	$0.635 \pm 0.180$	$0.302 \pm 0.080$	$0.544 \pm 0.230$	$0.129 \pm 0.022$	$0.955 \pm 0.026$	$1.56 \pm 0.97 \times 10^2$
B-PC*	$0.305 \pm 0.057$	$0.570 \pm 0.138$	$0.349 \pm 0.058$	$0.471 \pm 0.154$	$0.354 \pm 0.072$	$0.821 \pm 0.087$	$1.57 \pm 0.87 \times 10^2$
DiBS	$0.312 \pm 0.038$	$0.737 \pm 0.071$	$0.357 \pm 0.041$	$0.710 \pm 0.079$	$0.504 \pm 0.052$	$0.643 \pm 0.093$	$9.49 \pm 7.34 \times 10^3$
BCD Nets	$0.215 \pm 0.055$	$0.819 \pm 0.097$	$0.266 \pm 0.057$	$0.774 \pm 0.109$	$0.327 \pm 0.040$	$0.850 \pm 0.067$	$7.04 \pm 3.21 \times 10^1$
VBG	$0.237 \pm 0.037$	$0.816 \pm 0.064$	$0.284 \pm 0.027$	$0.799 \pm 0.050$	$0.434 \pm 0.058$	$0.738 \pm 0.091$	$1.24 \pm 0.49 \times 10^2$
JSP-GFN (diag)	$0.018 \pm 0.005$	$0.998 \pm 0.001$	$0.022 \pm 0.005$	$0.998 \pm 0.001$	$0.019 \pm 0.006$	$0.999 \pm 0.001$	$-4.91 \pm 0.51 \times 10^0$
JSP-GFN (full)	$0.019 \pm 0.007$	$0.998 \pm 0.001$	$0.021 \pm 0.007$	$0.998 \pm 0.002$	$0.020 \pm 0.008$	$0.999 \pm 0.001$	$-5.00 \pm 0.52 \times 10^0$

JSP-GFN matches exact posterior graph feature marginals with Pearson correlation $r \ge 0.998$ and obtains substantially lower cross-entropy with the true conditional parameter distribution than all baselines.

Knowl 10 — Real-World Biological Network Inference and Generalization Performance
empirical result
JSP-GFN was evaluated on biological network inference tasks using non-linear 2-layer MLP CPDs trained via mini-batch updates:
1. Flow Cytometry Data ( $d = 11$ phosphoproteins, $N = 4,200$ samples from 7 experiments): Discretized into 3 expression states, containing observational and interventional data. On held-out unseen interventions, JSP-GFN achieved a negative interventional log-likelihood of $2.679 \times 10^4$ , comparable to MH-MC $^3$ ( $2.578 \times 10^4$ ) and substantially outperforming Gibbs-MC $^3$ ( $3.490 \times 10^5$ ).
2. Gene Regulatory Networks from Gene Expression Data ( $d = 61$ genes, $N = 2,628$ samples): Modeled with zero-inflated Normal CPDs ( $|\theta| = 61,671$ ). On held-out test data, JSP-GFN achieved a negative log-likelihood of $2.651 \times 10^5$ , outperforming MH-MC $^3$ ( $1.073 \times 10^6$ ) and Gibbs-MC $^3$ ( $4.882 \times 10^6$ ).
Knowl 11 — Unimodality Limitation of Gaussian Parameter Posteriors in Non-Linear Models
limitation
Parametrizing the parameter posterior distribution $P_\phi(\theta \mid G, \text{stop})$ as a Gaussian distribution restricts the learned parameter posterior to unimodal distributions. In non-linear Bayesian networks where CPDs are parametrized by neural networks (e.g., 2-layer MLPs), the true posterior $P(\theta \mid G, \mathcal{D})$ is multimodal due to network symmetries such as hidden neuron permutations that yield identical likelihoods and equal posterior densities. A single Gaussian output distribution cannot capture these multiple disconnected posterior modes.

Coverage note — No substantial contributed theoretical or empirical material was omitted; descriptions of standard baseline algorithms and neural network architecture hyperparameter tables were condensed into their corresponding empirical results.

References

1.Yashas Annadani, Jonas Rothfuss, Alexandre Lacoste, Nino Scherrer, Anirudh Goyal, Yoshua Bengio, and Stefan Bauer. Variational Causal Networks: Approximate Bayesian Inference over Causal Structures. arXiv preprint, 2021.
2.Lazar Atanackovic, Alexander Tong, Jason Hartford, Leo J. Lee, Bo Wang, and Yoshua Bengio. DynGFN: Bayesian Dynamic Causal Discovery using Generative Flow Networks. arXiv preprint, 2023.
3.Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint, 2018.
4.Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. Neural Information Processing Systems, 2021.
5.Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet Foundations. Journal of Machine Learning Research (JMLR), 2023.
6.Bertrand Charpentier, Simon Kibler, and Stephan Günnemann. Differentiable DAG Sampling. International Conference on Learning Representations, 2022.
7.David Maxwell Chickering. Optimal Structure Identification With Greedy Search. Journal of Machine Learning Research, 2002.
8.Chris Cundy, Aditya Grover, and Stefano Ermon. BCD Nets: Scalable Variational Approaches for Bayesian Causal Discovery. Advances in Neural Information Processing Systems, 2021.
9.Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian Structure Learning with Generative Flow Networks. Uncertainty in Artificial Intelligence, 2022.
10.Daniel Eaton and Kevin Murphy. Bayesian structure learning using dynamic programming and MCMC. Conference on Uncertainty in Artificial Intelligence, 2007.
11.Paul Erdős and Alfréd Rényi. On the evolution of random graphs. Publications of the Mathematical Institute of the Hungarian Academy of Sciences, 1960.
12.Martin A Fischler and Robert C Bolles. Random sample consensus: a paradigm for model fitting with applications to image analysis and automated cartography. Communications of the ACM, 1981.
13.Nir Friedman and Daphne Koller. Being Bayesian About Network Structure. A Bayesian Approach to Structure Discovery in Bayesian Networks. Machine Learning, 2003.
14.Nir Friedman, Moises Goldszmidt, and Abraham Wyner. Data Analysis with Bayesian Networks: A Bootstrap Approach. Proceedings of the Fifteenth conference on Uncertainty in Artificial Intelligence, 1999.
15.Nir Friedman, Michal Linial, Iftach Nachman, and Dana Pe’er. Using bayesian networks to analyze expression data. In Proceedings of the fourth annual international conference on Computational molecular biology, 2000.
16.Eva-Maria Fronk. Model Selection for Dags via RJMCMC for the Discrete and Mixed Case. Technical report, Ludwig-Maximilians-Universität München, 2002.
17.Dan Geiger and David Heckerman. Learning Gaussian Networks. In Uncertainty in Artificial Intelligence. 1994.
18.Paolo Giudici and Robert Castelo. Improving Markov chain Monte Carlo model search for data mining. Machine learning, 2003.
19.David Heckerman, Dan Geiger, and David M Chickering. Learning bayesian networks: The combination of knowledge and statistical data. Machine learning, 1995.
20.Edward J. Hu, Nikolay Malkin, Moksh Jain, Katie Everett, Alexandros Graikos, and Yoshua Bengio. GFlowNet-EM for learning compositional latent variable models. International Conference on Machine Learning, 2023.
21.Aapo Hyvärinen. Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research, 2005.
22.Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure F.P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological Sequence Design with GFlowNets. International Conference on Machine Learning, 2022.
23.Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez-Garcia, and Yoshua Bengio. GFlowNets for AI-Driven Scientific Discovery. Digital Discovery, 2023.
24.Daphne Koller and Nir Friedman. Probabilistic Graphical Models: Principles and Techniques. MIT press, 2009.
25.Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A Theory of Continuous Generative Flow Networks. International Conference on Machine Learning, 2023.
26.Steffen L Lauritzen and David J Spiegelhalter. Local Computations with Probabilities on Graphical Structures and their Application to Expert Systems. Journal of the Royal Statistical Society: Series B (Methodological), 1988.
27.Yinchuan Li, Shuang Luo, Haozhi Wang, and Jianye Hao. CFlowNets: Continuous control with Generative Flow Networks. International Conference on Learning Representations, 2023.
28.Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. Advances in Neural Information Processing Systems, 2016.
29.Lars Lorch, Jonas Rothfuss, Bernhard Schölkopf, and Andreas Krause. DiBS: Differentiable Bayesian Structure Learning. Advances in Neural Information Processing Systems, 2021.
30.Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf. Amortized Inference for Causal Structure Learning. Advances in Neural Information Processing Systems, 2022.
31.Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. arXiv preprint, 2022.
32.David Madigan, Jonathan Gavrin, and Adrian E Raftery. Enhancing the Predictive Performance of Bayesian Graphical Models. 1994.
33.David Madigan, Jeremy York, and Denis Allard. Bayesian Graphical Models for Discrete Data. International Statistical Review, 1995.
34.Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems, 2022.
35.Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. International Conference on Learning Representations, 2023.
36.Joris M Mooij, Sara Magliacane, and Tom Claassen. Joint Causal Inference from Multiple Contexts. Journal of Machine Learning Research, 2020.
37.Mizu Nishikawa-Toomey, Tristan Deleu, Jithendaraa Subramanian, Yoshua Bengio, and Laurent Charlin. Bayesian learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayes. AAAI Workshop Graphs and More Complex Structures for Learning and Reasoning, 2023.
38.Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning, 2023.
39.Karen Sachs, Omar Perez, Dana Pe’er, Douglas A Lauffenburger, and Garry P Nolan. Causal protein-signaling networks derived from multiparameter single-cell data. Science, 2005.
40.Muralikrishnna G Sethuraman, Romain Lopez, Rahul Mohan, Faramarz Fekri, Tommaso Biancalani, and Jan-Christian Hütter. NODAGS-Flow: Nonlinear Cyclic Causal Structure Learning. International Conference on Artificial Intelligence and Statistics (AISTATS), 2023.
41.Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, Prediction, and Search. MIT press, 2000.
42.Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. Advances in Neural Information Processing Systems, 2017.
43.Jussi Viinikka, Antti Hyttinen, Johan Pensar, and Mikko Koivisto. Towards Scalable Bayesian Learning of Causal DAGs. Advances in Neural Information Processing Systems, 2020.
44.Julius von Kügelgen, Paul K Rubenstein, Bernhard Schölkopf, and Adrian Weller. Optimal experimental design via Bayesian optimization: active causal structure learning for Gaussian process networks. NeurIPS Workshop “Do the right thing”: machine learning and causal inference for improved decision making, 2019.
45.Benjie Wang, Matthew R Wicker, and Marta Kwiatkowska. Tractable Uncertainty for Structure Learning. International Conference on Machine Learning, 2022.
46.David Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust Scheduling with GFlowNets . International Conference on Learning Representations, 2023.
47.Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying Generative Models with GFlowNets and Beyond. International Conference on Machine Learning – Beyond Bayes workshop, 2022.
48.Xun Zheng, Bryon Aragam, Pradeep Ravikumar, and Eric P. Xing. DAGs with NO TEARS: Continuous Optimization for Structure Learning. In Advances in Neural Information Processing Systems, 2018.
49.Heiko Zimmermann, Fredrik Lindsten, Jan-Willem van de Meent, and Christian A. Naesseth. A Variational Perspective on Generative Flow Networks. arXiv preprint, 2022.

Citation

MLA

Deleu, T., et al. “Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network”. Advances in Neural Information Processing Systems, vol. 36, 2023, pp. 31204–31, https://proceedings.neurips.cc/paper_files/paper/2023/file/639a9a172c044fbb64175b5fad42e9a5-Paper-Conference.pdf.

APA

Deleu, T., Nishikawa-Toomey, M., Subramanian, J., Whitammer, E. S., Charlin, L., & Bengio, Y. (2023). Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network. Advances in Neural Information Processing Systems, 36, 31204–31231. https://proceedings.neurips.cc/paper_files/paper/2023/file/639a9a172c044fbb64175b5fad42e9a5-Paper-Conference.pdf

Chicago

Deleu, T., M. Nishikawa-Toomey, J. Subramanian, E. S. Whitammer, L. Charlin, and Y. Bengio. 2023. “Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network”. Advances in Neural Information Processing Systems 36: 31204–31. https://proceedings.neurips.cc/paper_files/paper/2023/file/639a9a172c044fbb64175b5fad42e9a5-Paper-Conference.pdf.

Harvard

Deleu, T. et al. (2023) “Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network”, Advances in Neural Information Processing Systems. Curran Associates, Inc., pp. 31204–31231. Available at: https://proceedings.neurips.cc/paper_files/paper/2023/file/639a9a172c044fbb64175b5fad42e9a5-Paper-Conference.pdf.

Vancouver

1. Deleu T, Nishikawa-Toomey M, Subramanian J, Whitammer ES, Charlin L, Bengio Y (2023) Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network. In: Advances in Neural Information Processing Systems. Curran Associates, Inc., pp 31204–31231

BibTeX

@inproceedings{deleu2023joint,
  title = {Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network},
  author = {Deleu, Tristan and Nishikawa-Toomey, Mizu and Subramanian, Jithendaraa and Whitammer, Esmeralda S. and Charlin, Laurent and Bengio, Yoshua},
  year = {2023},
  booktitle = {Advances in Neural Information Processing Systems},
  publisher = {Curran Associates, Inc.},
  volume = {36},
  pages = {31204-31231},
  url = {https://proceedings.neurips.cc/paper_files/paper/2023/file/639a9a172c044fbb64175b5fad42e9a5-Paper-Conference.pdf}
}

Metadata:DOI registry

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Published with permission

Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — JSP-GFN Framework for Joint Structure and Parameter Inference

Knowl 2 — Subtrajectory Balance Conditions for Joint Posterior Approximation

Knowl 3 — JSP-GFN Learning Objective with Stop-Gradient and Score Matching

Knowl 4 — Forward Transition Parameterization via Graph Neural Networks

Knowl 5 — Mini-Batch Estimation and Unbiased Loss Gradients

Knowl 6 — JSP-GFN Training Procedure

Knowl 7 — Exact Parameter Posterior for Linear Gaussian Bayesian Networks

Knowl 8 — Estimation of Log-Terminating Probability via Beam Search and Monte Carlo