SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In Section 3.1, we saw that the process of generating  $(G,\theta)$  follows two phases: first we construct  $G$  one edge at a time, until we sample a specific "stop" action, at which point we sample  $\theta$ , conditioned on  $G$ . All these actions are sampled using the forward transition probabilities  $P_{\phi}(G' \mid G)$  during the first phase, and  $P_{\phi}(\theta \mid G)$  during the second one. Following Deleu et al. (2022), we parametrize these forward transition probabilities using a hierarchical model: we first decide whether we want to stop the first phase or not, with probability  $P_{\phi}(\text{stop} \mid G)$ ; then, conditioned on this first decision, we either continue adding an edge to  $G$  to reach  $G'$  with probability  $P_{\phi}(G' \mid G, \neg \text{stop})$  (phase 1), or sample  $\theta$  with probability  $P_{\phi}(\theta \mid G, \text{stop})$  (phase 2). This hierarchical model can be written as
$$
P _ {\phi} \left(G ^ {\prime} \mid G\right) = \left(1 - P _ {\phi} (\operatorname {s t o p} \mid G)\right) P _ {\phi} \left(G ^ {\prime} \mid G, \neg \operatorname {s t o p}\right) \tag {9}
$$
$$
P _ {\phi} (\theta \mid G) = P _ {\phi} (\operatorname {s t o p} \mid G) P _ {\phi} (\theta \mid G, \operatorname {s t o p}). \tag {10}
$$
We use neural networks to parametrize each of the three components necessary to define the forward transition probabilities. Unlike in DAG-GFlowNet though, which uses a linear Transformer to define  $P_{\phi}(\mathrm{stop} \mid G)$  and  $P_{\phi}(G' \mid G, \neg \mathrm{stop})$ , we use a combination of graph network (Battaglia et al., 2018) and self-attention blocks (Vaswani et al., 2017) to encode information about the graph  $G$ , which appears in the conditioning of all the quantities of interest. This common backbone returns a graph embedding  $g$  of  $G$ , as well as 3 embeddings  $\mathbf{u}_i, \mathbf{v}_i, \mathbf{w}_i$  for each node  $X_i$  in  $G$
$$
\boldsymbol {g}, \left\{\boldsymbol {u} _ {i}, \boldsymbol {v} _ {i}, \boldsymbol {w} _ {i} \right\} _ {i = 1} ^ {d} = \text {S e l f A t t e n t i o n} _ {\phi} \left(\operatorname {G r a p h N e t} _ {\phi} (G)\right).
$$
We can parametrize the probability of selecting the "stop" action using  $\pmb{g}$  with  $P_{\phi}(\mathrm{stop} \mid G) = f_{\phi}(\pmb{g})$ , where  $f_{\phi}$  is a neural network with a sigmoid output; note that if we can't add any edge to  $G$  without creating a cycle, we force the end of the first phase by setting  $P_{\phi}(\mathrm{stop} \mid G) = 1$ . Inspired by Lorch et al. (2021), the probability of moving from  $G$  to  $G'$  by adding the edge  $X_i \rightarrow X_j$  is parametrized by
$$
P _ {\phi} \left(G ^ {\prime} \mid G, \neg \text {s t o p}\right) \propto \boldsymbol {m} _ {i j} \exp \left(\boldsymbol {u} _ {i} ^ {\top} \boldsymbol {v} _ {j}\right), \tag {11}
$$
where  $m_{ij}$  is a binary mask indicating whether adding  $X_i \rightarrow X_j$  is a valid action (i.e., if it is not already present in  $G$ , and if it doesn't introduce a cycle; Deleu et al., 2022). Finally, the probability of selecting the parameters  $\theta_i$  of the CPD for the variable  $X_i$  is parametrized with a multivariate Normal distribution with diagonal covariance (unless specified otherwise)
$$
P _ {\phi} \left(\theta_ {i} \mid G, \text {s t o p}\right) = \mathcal {N} \left(\theta_ {i} \mid \boldsymbol {\mu} _ {\phi} (\boldsymbol {w} _ {i}), \sigma_ {\phi} ^ {2} (\boldsymbol {w} _ {i})\right), \tag {12}
$$
where  $\pmb{\mu}_{\phi}$  and  $\sigma_{\phi}^{2}$  are two neural networks, with appropriate non-linearities to guarantee that  $\sigma_{\phi}^{2}(\pmb{w}_{i})$  is a well-defined diagonal covariance matrix. Note that  $P_{\phi}(\theta_i\mid G,\mathrm{stop})$  effectively approximates the posterior distribution  $P(\theta_i\mid G,\mathcal{D})$  once fully trained. Moreover, in addition to being an approximation of the joint posterior  $P(G,\theta \mid \mathcal{D})$ , the GFlowNet also provides an approximation of the marginal posterior  $P(G\mid \mathcal{D})$ , by only following the first phase of the generation process (to generate  $G$ ) until the "stop" action is selected, and not continuing into the generation of  $\theta$ .
3.5 Mini-batch training
Throughout the paper, we have assumed that we had access to the full dataset of observations  $\mathcal{D}$  in order to compute the reward  $R(G,\theta)$  in (5). However, beyond the capacity to have an arbitrary likelihood model  $P(\mathcal{D}|\theta ,G)$  (e.g., non-linear), another advantage of approximating the joint posterior
$P(G,\theta \mid \mathcal{D})$  is that we can train the GFlowNet using mini-batches of observations. Concretely, for a mini-batch  $\mathcal{B}$  of  $M$  observations sampled uniformly at random from the dataset  $\mathcal{D}$ , we can define
$$
\log \widehat {R} _ {\mathcal {B}} (G, \theta) = \log P (\theta \mid G) + \log P (G) + \frac {N}{M} \sum_ {\boldsymbol {x} ^ {(m)} \in \mathcal {B}} \log P (\boldsymbol {x} ^ {(m)} \mid G, \theta), \tag {13}
$$
which is an unbiased estimate of the log-reward. The following proposition shows that minimizing the estimated loss based on (13) wrt. the parameters  $\phi$  of the GFlowNet also minimizes the original objective in Section 3.3.
Proposition 3.2. Suppose that  $\mathcal{B}$  is a mini-batch of  $M$  observations sampled uniformly at random from the dataset  $\mathcal{D}$ , and let  $\widehat{\mathcal{L}}_{\mathcal{B}}(\phi)$  be the learning objective defined in Section 3.3, where the reward has been replaced by the estimate  $\widehat{R}_B(G,\theta)$  in (13). Then we have  $\mathcal{L}(\phi) \leq \mathbb{E}_{\mathcal{B}}[\widehat{\mathcal{L}}_{\mathcal{B}}(\phi)]$ .
The proof is available in Appendix C.3.4, only relies on the convexity of the square function to conclude, but no other property of this function; in practice, we use the Huber loss instead of the square loss for stability, which is also a convex function. In the case of the square loss, Proposition C.1 gives a stronger result in terms of unbiasedness of the gradient estimator.
Bayesian Structure Learning. There is a vast literature applying Markov chain Monte Carlo (MCMC) methods to approximate the marginal posterior  $P(G \mid \mathcal{D})$  over the graphical structures of Bayesian Networks (Madigan et al., 1995; Friedman and Koller, 2003; Giudici and Castelo, 2003; Viunikka et al., 2020). However, since the parameter space in which  $\theta$  lives depends on the graph structure  $G$ , approximating the joint posterior  $P(G, \theta \mid \mathcal{D})$  using MCMC requires additional trans-dimensional updates (Fronk, 2002), and has therefore received less attention than the marginal case.
Variational methods have been proposed to approximate the marginal posterior too (Annadani et al., 2021; Charpentier et al., 2022). Similar to MCMC though, approximating the joint posterior has also been less studied than its marginal counterpart, with the notable exceptions of DiBS (Lorch et al., 2021) and BCD Nets (Cundy et al., 2021). We provide an extensive qualitative comparison between our method JSP-GFN and prior variational inference and GFlowNet methods in Appendix A.
Generative Flow Networks. While they were initially developed to encourage the discovery of diverse molecules (Bengio et al., 2021), GFlowNets proved to be a more general framework to describe distributions over composite objects that can be constructed sequentially (Bengio et al., 2023). The objective of the GFlowNet is to enforce a conservation law such as the flow-matching conditions in (2), indicating that the total amount of flow going into any state is equal to the total outgoing flow, with some residual given by the reward. Alternative conditions, sometimes bypassing the need to work with flows altogether, have been proposed in order to learn these models more efficiently (Malkin et al., 2022; Madan et al., 2022; Pan et al., 2023). By amortizing inference, and thus treating it as an optimization problem, GFlowNets find themselves deeply rooted in the variational inference literature (Malkin et al., 2023; Zimmermann et al., 2022), and are connected to other classes of generative models (Zhang et al., 2022). Beyond Bayesian structure learning (Deleu et al., 2022), GFlowNets have also applications in modeling Bayesian posteriors for variational EM (Hu et al., 2023), combinatorial optimization (Zhang et al., 2023), biological sequence design (Jain et al., 2022), as well as scientific discovery at large (Jain et al., 2023).
Closely related to our work, Nishikawa-Toomey et al. (2023) proposed to learn the joint posterior  $P(G,\theta \mid \mathcal{D})$  over structures and parameters with a GFlowNet, combined with Variational Bayes to circumvent the challenge of learning a distribution over continuous quantities  $\theta$  with a GFlowNet. Atanackovic et al. (2023) also used a GFlowNet called DynGFN to approximate the posterior of a dynamical system. Similar to (Nishikawa-Toomey et al., 2023) though, they used the GFlowNet only to approximate the distribution over graphs  $G$ , making the parameters  $\theta$  a deterministic function of  $G$  (i.e.,  $P(\theta \mid G,\mathcal{D})\approx \delta (\theta \mid G;\phi)$ ). Here, we leverage the recent advances extending these models to general sample spaces (Lahlou et al., 2023), including continuous spaces (Li et al., 2023), in order to model the joint posterior within a single GFlowNet.
(a) Edge features
Figure 2: Comparison with the exact posterior distribution, on small graphs with  $d = 5$  nodes. (a) Comparison of the edge features computed with the exact posterior (x-axis) and the approximation given by JSP-GFN (y-axis); each point corresponds to an edge  $X_{i} \rightarrow X_{j}$  for each of the 20 datasets. (b) Quantitative evaluation of different methods for joint posterior approximation, both in terms of edge features and cross-entropy of sampling distribution and true posterior  $P(\theta \mid G, \mathcal{D})$ ; all values correspond to the mean and  $95\%$  confidence interval across the 20 experiments.
Edge features	EG,θ[-logP(θ|G,D)]
RMSE	Pearson's r
MH-MC3	0.357 ± 0.022	0.067 ± 0.143	5.39 ± 1.41 × 10^2
Gibbs-MC3	0.357 ± 0.022	0.028 ± 0.127	9.02 ± 1.54 × 10^5
B-GES*	0.263 ± 0.070	0.635 ± 0.180	1.56 ± 0.97 × 10^2
B-PC*	0.305 ± 0.057	0.570 ± 0.138	1.57 ± 0.87 × 10^2
DiBS	0.312 ± 0.038	0.737 ± 0.071	9.49 ± 7.34 × 10^3
BCD Nets	0.215 ± 0.055	0.819 ± 0.097	7.04 ± 3.21 × 10^1
VBG	0.237 ± 0.037	0.816 ± 0.064	1.24 ± 0.49 × 10^2
JSP-GFN (diag)	0.018 ± 0.005	0.998 ± 0.001	-4.91 ± 0.51 × 10^0
JSP-GFN (full)	0.019 ± 0.007	0.998 ± 0.001	-5.00 ± 0.52 × 10^0
(b) Quantitative comparison with the exact posterior
5.1 Joint posterior over small graphs
We can evaluate the accuracy of the approximation returned by JSP-GFN by comparing it with the exact joint posterior distribution  $P(G, \theta \mid \mathcal{D})$ . Computing this exact posterior can only be done in limited cases only, where (1)  $P(\theta \mid G, \mathcal{D})$  can be computed analytically, and (2) for a small enough  $d$  such that all the DAGs can be enumerated in order to compute  $P(G \mid \mathcal{D})$ . We consider here models over  $d = 5$  variables, with linear Gaussian CPDs. We generate 20 different datasets of  $N = 100$  observations from randomly generated Bayesian Networks. Details about data generation and modeling are available in Appendix D.2.
The quality of the joint posterior approximations is evaluated separately for  $G$  and  $\theta$ . For the graphs, we compare the approximation and the exact posterior on different margins of interest, also called features (Friedman and Koller, 2003); e.g., the edge feature corresponds to the marginal probability of a specific edge being in the graph. Figure 2 (a) shows a comparison between the edge features computed with the exact posterior and with JSP-GFN, proving that it can accurately approximate the edge features of the exact posterior, despite the modeling bias discussed in Appendix D.2.2. Compared to other methods in Figure 2 (b), JSP-GFN offers significantly more accurate approximations of the posterior, at least relative to the edge features. This observation still holds on the path and Markov features (Deleu et al., 2022); see Appendix D.2.2.
To evaluate the performance of the different methods as an approximation of the posterior over  $\theta$ , we also estimate the cross-entropy between the sampling distribution of  $\theta$  given  $G$  and the exact posterior  $P(\theta \mid G, \mathcal{D})$ . This measure will be minimized if the model correctly samples parameters from the true  $P(\theta \mid G, \mathcal{D})$ ; details about this metric are given in Appendix D.2.3. In Figure 2 (b), we observe that again both versions of JSP-GFN sample parameters  $\theta$  that are significantly more probable under the exact posterior compared to other methods.
5.2 Gaussian Bayesian Networks from simulated data
To evaluate whether our observations hold on larger graphs, we also evaluated the performance of JSP-GFN on data simulated from larger Gaussian Bayesian Networks, with  $d = 20$  variables. In addition to linear CPDs, as in Section 5.1, we experimented with non-linear Gaussian Bayesian Networks, where the CPDs are parametrized by neural networks. Following Lorch et al. (2021), we parametrized the CPDs of each variable with a 2-layer MLP, for a total of  $|\theta| = 2$ , 220 parameters. For both experimental settings, we used datasets of  $N = 100$  observations simulated from (randomly generated) Bayesian Networks; additional experimental details are provided in Appendix D.3.
(a) Linear Gaussian
(b) Non-linear Gaussian
(c) Non-linear Gaussian
Figure 3: Evaluation of JSP-GFN on Gaussian Bayesian Networks. (a-b) Comparison of the negative log-likelihood (NLL) on  $N' = 100$  held-out observations for different Bayesian structure learning methods, aggregated across 20 experiments on different datasets  $\mathcal{D}$ . (c) Linear correlation between the log-reward (x-axis) and the terminating state log-probability (y-axis) for 1,000 samples  $(G, \theta)$  from JSP-GFN; the color of each point indicates the number of edges in the corresponding graph.  $\Delta$  represents the slope of a linear function fitted using RANSAC (Fischler and Bolles, 1981).
We compared JSP-GFN against two methods based on MCMC (MH-MC $^3$  & Gibbs-MC $^3$ ; Madigan et al., 1995) and DiBS (Lorch et al., 2021) on both experiments, as well as two bootstrapping algorithms (B-GES $^*$  & B-PC $^*$ ; Friedman et al., 1999), BCD Nets (Cundy et al., 2021) and VBG (Nishikawa-Toomey et al., 2023) for the experiment with linear Gaussian CPDs, as they are not applicable for non-linear CPDs. In Figure 3 (a-b), we report the performance of these joint posterior approximations in terms of the (expected) negative log-likelihood (NLL) on held-out observations. We observe that JSP-GFN achieves a lower NLL than any other method on linear Gaussian models and is competitive on non-linear Gaussian models.
We chose the NLL over some other metrics, typically comparing with the ground-truth graphs used for data generation, since it is more representative of the performance of these methods on downstream tasks (i.e., predictions on unseen data), and measures the quality of the joint posterior instead of only the marginal over graphs. This choice is aligned with the shortcomings of these other metrics highlighted by Lorch et al. (2022), and it is further justified in Appendix D.3.3. Nevertheless, we also report the expected Structural Hamming Distance (SHD), as well as the area under the ROC curve (AUROC) in Appendix D.3.3 for completeness. To complement these metrics, and in order to assess the quality of the approximation of the posterior in the absence of reference  $P(G,\theta \mid \mathcal{D})$ , we also show in Figure 3 (c) how the terminating state log-probability  $\log P_{\phi}^{\top}(G,\theta)$  of JSP-GFN correlates with the log-reward for a non-linear Gaussian model. Indeed, as stated in Theorem 3.1, we should ideally have  $\log P_{\phi}^{\top}(G,\theta)$  perfectly correlated with  $\log R(G,\theta)$  with slope 1, as
$$
\log P _ {\phi} ^ {\top} (G, \theta) \approx \log P (G, \theta \mid \mathcal {D}) = \log R (G, \theta) - \log P (\mathcal {D}). \tag {14}
$$
We can see that there is indeed a strong linear correlation across multiple samples  $(G,\theta)$  from JSP-GFN, with a slope  $\Delta$  close to 1, suggesting that the GFlowNet is again an accurate approximation of the joint posterior, at least around the modes it captures. Details about how  $\log P_{\phi}^{\top}(G,\theta)$  is estimated are available in Appendix D.3.2.
5.3 Learning biological structures from real data
We finally evaluated JSP-GFN on real-world biological data for two separate tasks: the discovery of protein signaling networks from flow cytometry data (Sachs et al., 2005), as well as the discovery of a small gene regulatory network from gene expression data. The flow cytometry dataset consists of  $N = 4,200$  measurements of  $d = 11$  phosphoproteins from 7 different experiments, meaning that this dataset contains a mixture of both observational and interventional data. Furthermore, this dataset has been discretized into 3 states, representing the level of activity (Eaton and Murphy, 2007). For the gene expression dataset, we used a subset of  $N = 2,628$  observations of  $d = 61$  genes from (Sethuraman et al., 2023). Details about the experimental setups are available in Appendix D.4.
(a) Flow cytometry data (subsample,  $N = 100$
Figure 4: Performance of JSP-GFN on real-world biological data. (a) Comparison of the terminating state log-probability  $\log P_{\phi}^{\top}(G,\theta)$  returned by JSP-GFN with the log-reward  $\log R(G,\theta)$  on a subsample of  $N = 100$  datapoints of the flow cytometry dataset  $\mathcal{D}$ . (b) Same comparison with the full dataset  $\mathcal{D}$  of size  $N = 4,200$ . (c) Comparison of JSP-GFN with methods based on MCMC on both flow cytometry data and gene expression data, in terms of negative (interventional) log-likelihood on held-out data.
(b) Flow cytometry data
Flow cytometry
d=11,N=4,200	Gene expression
d=61,N=2,628
MH-MC3	2.578 × 104	1.073 × 106
Gibbs-MC3	3.490 × 105	4.882 × 106
JSP-GFN	2.679 × 104	2.651 × 105
(c) Negative log-likelihood
At this scale, using the whole dataset  $\mathcal{D}$  to evaluate the reward becomes impractical, especially for non-linear models. Fortunately, we showed in Section 3.5 that we can use an (unbiased) estimate of the reward, based on mini-batches of data, in place of  $R(G,\theta)$  in the loss function (8). In both experiments, we used non-linear models, where all the CPDs are parametrized with a 2-layer MLP. Figure 4 shows similar correlation plots as Figure 3 (c), along with an evaluation of the NLL on unseen observations and interventions. Beyond the ability to work with real data, these experiments allow us to highlight some other capacities of JSP-GFN: (1) handling discrete and (2) interventional data (flow cytometry), as well as (3) learning a distribution over larger graphs (gene expression).
To measure the quality of the posterior approximation returned by JSP-GFN, we compare in Figure 4 the terminating state log-probability  $\log P_{\phi}^{\top}(G,\theta)$  with the log-reward log  $R(G,\theta)$ , similar to Figure 3. We observe that there is correlation between these two quantities; unlike in Figure 3 though, we observe that the slope is not close to 1, suggesting that JSP-GFN underestimates the probability of  $(G,\theta)$ . We also observe that the graphs are "clustered" together; this can be explained by the fact that the posterior approximation is concentrated at only a few graphs, since the size of the dataset  $\mathcal{D}$  is larger. To confirm this observation, we show in Figure 4 (a) a similar plot on a subsample of  $N = 100$  datapoints randomly sampled from  $\mathcal{D}$ , matching the experimental setting of Section 5.2. In this case, we observe a much closer linear fit, with a slope closer to 1.
In addition to the comparison to the log-reward, we also compare in Figure 4 (c) JSP-GFN with 2 methods based on MCMC in terms of their negative log-likelihood on held-out data. We can see that JSP-GFN is competitive, and even out-performs MCMC on the more challenging problem of the discovery of gene regulatory networks from gene expression data, where the dimensionality of the problem is much larger ( $d = 61$ ). Note that the values reported for the discovery of protein signaling networks from flow cytometry data correspond to the negative interventional log-likelihood, on interventions unseen in  $\mathcal{D}$ .
We have presented JSP-GFN, an approach to approximate the joint posterior distribution over the structure of a Bayesian Network along with its parameters using a single GFlowNet. We have shown that our method faithfully approximates the joint posterior on both simulated and real data, and compares favorably against existing Bayesian structure learning methods. In line with Appendix B, future work should consider using more expressive distributions, such as those parametrized by normalizing flows or diffusion processes, to approximate the posteriors over continuous parameters, which would enable Bayesian inference over parameters in more complex generative models.
We would like to thank Antonio Góis, Salem Lahlou, Romain Lopez, Cristian Dragos Manta, and Mansi Rankawat for the useful discussions about the project and their feedback on the paper. We would like to also thank Lars Lorch for releasing the code for DiBS (Lorch et al., 2021), along with useful baselines used in this paper, as well as Chris Cundy for releasing the code for BCD Nets and for his help with the code. This research was enabled by compute resources provided by Mila (mila.quebec).