SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tristan Deleu<sup>1</sup> Mizu Nishikawa-Toomey<sup>1</sup> Jithendarraa Subramanian<sup>2</sup>
Nikolay Malkin<sup>1</sup> Laurent Charlin<sup>3</sup> Yoshua Bengio<sup>1,4</sup>
Mila - Quebec AI Institute
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.
As a compact representation for complex probabilistic models, Bayesian Networks are a framework of choice in many fields, such as computational biology (Friedman et al., 2000; Sachs et al., 2005) and medical diagnosis (Lauritzen and Spiegelhalter, 1988). When the directed acyclic graph (DAG) structure of the Bayesian Network—which specifies the possible conditional dependences among the observed variables—is known, it can be used to perform probabilistic inference for queries of interest with a variety of exact or approximate methods (Koller and Friedman, 2009). However, if this graphical structure is unknown, one may want to infer it based on a dataset of observations  $\mathcal{D}$ .
In addition to being a challenging problem due to the super-exponentially large search space, learning a single DAG structure from data may also lead to confident but incorrect predictions (Madigan et al., 1994), especially in cases where the evidence is limited. In order to avoid misspecifying the model, it is therefore essential to quantify the epistemic uncertainty about the structure of the Bayesian Network. This can be addressed by taking a Bayesian perspective on structure learning and inferring the posterior distribution  $P(G \mid \mathcal{D})$  over graphs given our observations. This (marginal) posterior can be approximated using methods based on Markov chain Monte Carlo (MCMC; Madigan et al., 1995) or variational inference (Cundy et al., 2021; Lorch et al., 2021). However, all of these methods rely on the computation of the marginal likelihood  $P(\mathcal{D} \mid G)$ , which can only be done efficiently in closed form for limited classes of models, such as linear Gaussian (Geiger and Heckerman, 1994),
<sup>1</sup>Université de Montréal, <sup>2</sup>McGill University, <sup>3</sup>HEC Montréal, <sup>4</sup>CIFAR Senior Fellow.
Correspondence: Tristan Deleu (deleutri@mila.quebec)
Code available at: https://github.com/tristandeleu/jax-jsp-gfn
discrete models with Dirichlet prior (Heckerman et al., 1995), or non-linear models parametrized with a Gaussian Process (von Kugelgen et al., 2019).
While there exists a vast literature on Bayesian structure learning to approximate the marginal posterior distribution, inferring the joint posterior  $P(G,\theta \mid \mathcal{D})$  over both the DAG structure  $G$  of the Bayesian Network and the parameters  $\theta$  of its conditional probability distributions—the probability of each variable given its parents—has received comparatively little attention. The main difficulty arises from the mixed sample space of the joint posterior distribution, with both discrete components (the graph  $G$ ) and continuous components (the parameters  $\theta$ ), where the dimensionality of the latter may even depend on  $G$ . However, modeling the posterior distribution over  $\theta$  has the notable advantage that the conditional probability distributions can be more flexible (e.g., parametrized by neural networks): in general, computing  $P(\mathcal{D}\mid G,\theta)$  is easier than computing the marginal  $P(\mathcal{D}\mid G)$ , lifting the need to perform intractable marginalizations.
Since they provide a framework for generative modeling of discrete and composite objects, Generative Flow Networks (GFlowNets; Bengio et al., 2021, 2023) proved to be an effective method for Bayesian structure learning. In Deleu et al. (2022), the problem of generating a sample DAG from the marginal posterior  $P(G \mid \mathcal{D})$  was treated as a sequential decision process, where edges are added one at a time, starting from the empty graph over  $d$  variables, following a learned transition probability. Nishikawa-Toomey et al. (2023) have also proposed to use a GFlowNet to infer the joint posterior  $P(G, \theta \mid \mathcal{D})$ ; however, they used it in conjunction with Variational Bayes to update the distribution over  $\theta$ , getting around the difficulty of modeling a continuous distribution with a GFlowNet.
In this paper, we propose to infer the joint posterior over graphical structures  $G$  and parameters of the conditional probability distributions  $\theta$  of a Bayesian Network using a single GFlowNet called JSP-GFN (for Joint Structure and Parameters GFlowNet), leveraging recent advances extending GFlowNets to continuous sample spaces (Lahlou et al., 2023), and expands the scope of applications for Bayesian structure learning with GFlowNets. The generation of a sample  $(G,\theta)$  from the approximate posterior now follows a two-phase process, where the DAG  $G$  is first constructed by inserting one edge at a time, and then the corresponding parameters  $\theta$  are chosen once the structure is completely known. To enable efficient learning of the sampling distribution, we introduce new conditions closely related to the ones derived in Deleu et al. (2022), based on the subtrajectory balance conditions (Malkin et al., 2022), and show that they guarantee that the GFlowNet does represent  $P(G,\theta \mid \mathcal{D})$  once they are completely satisfied. We validate empirically that JSP-GFN provides an accurate approximation of the posterior when those conditions are approximately satisfied by a learned sampling model, and compares favorably against existing methods on simulated and real data.
Notations. Throughout this paper, we will work with directed graphs  $G = (V, E)$ , where  $V$  is a set of nodes, and  $E \subseteq V \times V$  is a set of (directed) edges. For a node  $X \in V$ , we denote by  $\operatorname{Pa}_G(X)$  the set of parents of  $X$  in  $G$ , and  $\operatorname{Ch}_G(X)$  the set of its children. For two nodes  $X, Y \in V$ ,  $X \to Y$  represents a directed edge  $(X, Y) \in E$  (denoted  $X \to Y \in G$ ), and  $X \rightsquigarrow Y$  represents a directed path from  $X$  to  $Y$ , following the edges in  $E$  (denoted  $X \rightsquigarrow Y \in G$ ).
In the context of GFlowNets (see Section 2.2), an undirected path in a directed graph  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  between two states  $s_0,s_n\in \mathcal{V}$  is a sequence of vertices  $(s_0,s_1,\ldots ,s_n)$  where either  $(s_i,s_{i + 1})\in \mathcal{E}$  or  $(s_{i + 1},s_i)\in \mathcal{E}$  (i.e., following the edges of the graph, regardless of their orientations).
2.1 Bayesian structure learning
A Bayesian Network is a probabilistic model over  $d$  random variables  $\{X_1, \ldots, X_d\}$ , whose joint distribution factorizes according to a directed acyclic graph (DAG)  $G$  as
$$
P \left(X _ {1}, \dots , X _ {d} \mid \theta , G\right) = \prod_ {i = 1} ^ {d} P \left(X _ {i} \mid \operatorname {P a} _ {G} \left(X _ {i}\right); \theta_ {i}\right), \tag {1}
$$
where  $\theta = \{\theta_{1},\dots,\theta_{d}\}$  represents the parameters of the conditional probability distributions (CPDs) involved in this factorization. When the structure  $G$  is known, a Bayesian Network offers a representation of the joint distribution that may be convenient for probabilistic inference (Koller and
Figure 1: Structure of the Generative Flow Network to approximate the joint posterior distribution  $P(G, \theta \mid \mathcal{D})$ . A graph  $G$  and parameters  $\theta$  are constructed as follows: starting from the empty graph  $G_0$ , (1) the graph  $G$  is first generated one edge at a time (blue), as in (Deleu et al., 2022). Then once we select the action indicating that we stop adding edges to the graph, (2) we generate the parameters  $\theta$  (red), conditioned on the graph  $G$ . Finally, given  $G$  and  $\theta$ , (3) we receive a reward  $R(G, \theta)$  (green).
Friedman, 2009), as well as learning its parameters  $\theta$  using a dataset  $\mathcal{D} = \{\pmb{x}^{(1)},\dots,\pmb{x}^{(N)}\}$ , where each  $\pmb{x}^{(j)}$  represents an observation of the  $d$  random variables  $\{X_{1},\ldots ,X_{d}\}$ .
However, when the structure of the Bayesian Network is unknown, we can also learn the DAG  $G$  from data (Spirtes et al., 2000; Chickering, 2002). Using a Bayesian perspective, we may want to either model the (marginal) posterior distribution  $P(G \mid \mathcal{D})$  over only the DAG structures, or the joint posterior  $P(G, \theta \mid \mathcal{D})$  over both the structure  $G$ , as well as the parameters of the CPDs  $\theta$ .
2.2 Generative Flow Networks
A Generative Flow Network (GFlowNet; Bengio et al., 2021, 2023) is a generative model over a structured sample space  $\mathcal{X}$ . The structure of the GFlowNet is described by a DAG  $\mathcal{G}$  whose vertex set is the state space  $\mathcal{S}$ , where  $\mathcal{X} \subseteq \mathcal{S}$ . This should not be confused with the DAG in a Bayesian Network: in fact, each state in  $\mathcal{G}$  could itself represent the structure of a Bayesian Network (Deleu et al., 2022). A sample  $x \in \mathcal{X}$  is constructed sequentially by following the edges of  $\mathcal{G}$ , starting at a special initial state  $s_0$ , until we reach the state  $x$ . We define a special terminal state  $s_f$ , a transition to which indicates the end of the sequential process. The states in  $\mathcal{X}$  are those for which there is a directed edge  $x \rightarrow s_f$ ; they are called complete states<sup>1</sup> and correspond to valid samples of the distribution induced by the GFlowNet. A path  $s_0 \rightsquigarrow s_f$  in  $\mathcal{G}$  is called a complete trajectory. An example of the structure of a GFlowNet is given in Section 3.1 and Figure 1.
Every complete state  $x \in \mathcal{X}$  is associated with a reward  $R(x) \geq 0$ , indicating the unnormalized probability of  $x$ . We use the convention  $R(s) = 0$  for any state  $s \in S \backslash \mathcal{X}$ , since they do not correspond to valid samples of the distribution. Bengio et al. (2021) showed that if there exists a function  $F_{\phi}(s \to s') \geq 0$  defined over the edges of  $\mathcal{G}$ , called a flow, that satisfies the following flow-matching conditions
$$
\sum_ {s \in \operatorname {P a} _ {\mathcal {G}} \left(s ^ {\prime}\right)} F _ {\phi} \left(s \rightarrow s ^ {\prime}\right) - \sum_ {s ^ {\prime \prime} \in \operatorname {C h} _ {\mathcal {G}} \left(s ^ {\prime}\right)} F _ {\phi} \left(s ^ {\prime} \rightarrow s ^ {\prime \prime}\right) = R \left(s ^ {\prime}\right) \tag {2}
$$
for all the non-terminal states  $s' \in S$ , then the GFlowNet induces a distribution over complete states proportional to the reward. More precisely, starting from the initial state  $s_0$ , if we sample a complete trajectory  $(s_0, s_1, \ldots, s_{T-1}, x, s_f)$  following the forward transition probability, defined as
$$
P \left(s _ {t + 1} \mid s _ {t}\right) \propto F _ {\phi} \left(s _ {t} \rightarrow s _ {t + 1}\right), \tag {3}
$$
with the conventions  $s_T = x$  and  $s_{T + 1} = s_f$ , then the marginal probability that a trajectory sampled following  $P$  terminates in  $x \in \mathcal{X}$  is proportional to  $R(x)$ . The flow function  $F_{\phi}(s \to s')$  may be parametrized by a neural network and optimized to minimize the error in (2), yielding a transition model that can be used to approximately sample from the distribution on  $\mathcal{X}$  proportional to  $R$ .
<sup>1</sup>Bengio et al. (2023) also call these states terminating; we follow the naming conventions of Deleu et al. (2022) here to avoid ambiguity, as it is closely related to our work.
2.3 Structure learning with GFlowNets
Since GFlowNets are particularly well-suited to specifying distributions over composite objects, Deleu et al. (2022) used this framework in the context of Bayesian structure learning to approximate the (marginal) posterior distribution over DAGs  $P(G \mid \mathcal{D})$ . Their model, called DAG-GFlowNet, operates on the state-space of DAGs, where each graph is constructed sequentially by adding one edge at a time, starting from the empty graph with  $d$  nodes, while enforcing the acyclicity constraint at every step of the generation (i.e., an edge is not added if it would introduce a cycle). Its structure is illustrated at the top of Figure 1, where it forms the part of the graph shown in blue.
Instead of working with flows  $F_{\phi}(G \to G')$ , as in Section 2.2, DAG-GFlowNet directly learns the forward transition probability  $P_{\phi}(G' \mid G)$  that satisfies the following alternative detailed balance conditions for any transition  $G \to G'$  (i.e.,  $G'$  is the result of adding a single edge to  $G$ ):
$$
R \left(G ^ {\prime}\right) P _ {B} (G \mid G ^ {\prime}) P _ {\phi} \left(s _ {f} \mid G\right) = R (G) P _ {\phi} \left(G ^ {\prime} \mid G\right) P _ {\phi} \left(s _ {f} \mid G ^ {\prime}\right), \tag {4}
$$
where  $P_B(G \mid G')$  is a fixed distribution over the parent states of  $G'$  (e.g., uniform distribution over parents). Deleu et al. (2022) showed that since all the states are complete here, satisfying the conditions (4) for all  $G \to G'$  still induces a distribution over DAGs  $\propto R(G)$ . Therefore, to approximate the posterior distribution  $P(G \mid \mathcal{D})$ , they used  $R(G) = P(\mathcal{D} \mid G)P(G)$  as the reward of  $G$ . In particular, this requires evaluating the marginal likelihood  $P(\mathcal{D} \mid G)$  efficiently, which is feasible only for limited classes of models (e.g., linear Gaussian; Geiger and Heckerman, 1994).
3 Joint Bayesian inference of structure and parameters
Although Generative Flow Networks have been primarily applied to model distributions over discrete objects such as DAGs, Lahlou et al. (2023) showed that similar ideas could also be applied to continuous objects, and discrete-continuous hybrids. Building on top of DAG-GFlowNet, we propose here to approximate the joint posterior  $P(G,\theta \mid \mathcal{D})$  over both the structure of the Bayesian Network  $G$ , but also the parameters of its conditional probability distributions  $\theta$ . Unlike in VBG though (Nishikawa-Toomey et al., 2023), we use a single GFlowNet to approximate this joint posterior. We call this model JSP-GFN, for Joint Structure and Parameters Bayesian inference with a GFlowNet.
3.1 Structure of the GFlowNet
Unlike in DAG-GFlowNet, where we model a distribution only over DAGs, here we need to define a GFlowNet whose complete states are pairs  $(G,\theta)$ , where  $G$  is a DAG and  $\theta$  is a set of (continuous-valued) parameters whose dimension may depend on  $G$ . Complete states are obtained through two phases (Figure 1): the DAG  $G$  is first constructed one edge at a time, following Deleu et al. (2022), and then the corresponding parameters  $\theta$  are generated, conditioned on  $G$ . We denote by  $(G,\cdot)$  states where the DAG  $G$  has no parameters  $\theta$  associated to it (states in blue in Figure 1); they are intermediate states during the first phase of the GFlowNet, and do not correspond to valid samples of the induced distribution. Using the notations of Section 2.2,  $(G,\theta)\in \mathcal{X}$ , whereas  $(G,\cdot)\in S\backslash \mathcal{X}$ .
Starting at the empty graph  $(G_0, \cdot)$ , the DAG is constructed one edge at a time during the first phase, following the forward transition probabilities  $P_{\phi}(G' \mid G)$ . This first phase ends when a special "stop" action is selected with  $P_{\phi}$ , indicating that we stop adding edges to the graph; the role of this "stop" action is detailed in Section 3.4. Then during the second phase, we generate  $\theta$  conditioned on  $G$ , following the forward transition probabilities  $P_{\phi}(\theta \mid G)$ . All the complete states  $(G, \theta)$ , for a fixed graph  $G$  and any set of parameters  $\theta$ , can be seen as forming an (infinitely wide) tree rooted at  $(G, \cdot)$ .
Since we want this GFlowNet to approximate the joint posterior  $P(G, \theta \mid \mathcal{D}) \propto P(\mathcal{D}, \theta, G)$ , it is natural to define the reward function of a complete state  $(G, \theta)$  as
$$
R (G, \theta) = P (\mathcal {D} \mid \theta , G) P (\theta \mid G) P (G), \tag {5}
$$
where the likelihood model  $P(\mathcal{D} \mid \theta, G)$  may be arbitrary (e.g., a neural network), and decomposes according to (1),  $P(\theta \mid G)$  is the prior over parameters, and  $P(G)$  the prior over graphs. When the dataset  $\mathcal{D}$  is large, we can use a mini-batch approximation to the reward (see Section 3.5).
<sup>2</sup>Since the states of the GFlowNet here are pairs of objects,  $P_{\phi}(\theta \mid G)$  (resp.  $P_{\phi}(G' \mid G)$ ) is an abuse of notation, and represents  $P_{\phi}((G, \theta) \mid (G, \cdot))$  (resp.  $P_{\phi}((G', \cdot) \mid (G, \cdot))$ ).
3.2 Subtrajectory Balance conditions
To obtain a generative process that samples pairs of  $(G,\theta)$  proportionally to the reward, the GFlowNet needs to satisfy some conditions such as (4). However, we saw in Section 2.3 that satisfying this particular formulation of the detailed balance conditions in (4) yields a distribution  $\propto R(\cdot)$  only if all the states are complete; unfortunately, this is not the case here since there exists states of the form  $(G,\cdot)$  corresponding to graphs without their associated parameters. Instead, we use a generalization of detailed balance to undirected paths of arbitrary length, called the Subtrajectory Balance conditions (SubTB; Malkin et al., 2022); we give a brief overview of SubTB in Appendix C.2.
More precisely, we consider SubTB for any undirected path of the form  $(G,\theta)\gets (G,\cdot)\to (G^{\prime},\cdot)\to$ $(G^{\prime},\theta^{\prime})$  in the GFlowNet (see Figure C.2), where  $G^{\prime}$  is therefore the result of adding a single edge to  $G$ . Since both ends of these undirected paths of length 3 are complete states, we show in Appendix C.3.1 that the SubTB conditions corresponding to undirected paths of this form can be written as
$$
R \left(G ^ {\prime}, \theta^ {\prime}\right) P _ {B} (G \mid G ^ {\prime}) P _ {\phi} (\theta \mid G) = R (G, \theta) P _ {\phi} \left(G ^ {\prime} \mid G\right) P _ {\phi} \left(\theta^ {\prime} \mid G ^ {\prime}\right). \tag {6}
$$
Note that the SubTB condition above is very similar to the detailed balance condition in (4) used in DAG-GFlowNet, where the probability of terminating  $P_{\phi}(s_f \mid G)$  has been replaced by  $P_{\phi}(\theta \mid G)$ , in addition to the reward now depending on both  $G$  and  $\theta$ . Moreover, while it does not seem to appear in (6), the terminal state  $s_f$  of the GFlowNet is still present implicitly, since we are forced to terminate once we have reached a complete state  $(G, \theta)$ , and therefore  $P_{\phi}(s_f \mid G, \theta) = 1$  (see App. C.3.1).
Although there is no guarantee in general that satisfying the SubTB conditions would yield a distribution proportional to the reward, unlike with the detailed balance conditions (Bengio et al., 2023), the following theorem shows that the GFlowNet does induce a distribution  $\propto R(G,\theta)$  if the SubTB conditions in (6) are satisfied for all pairs  $(G,\theta)$  and  $(G^{\prime},\theta^{\prime})$
Theorem 3.1. If the SubTB conditions in (6) are satisfied for all undirected paths of length 3 between any  $(G,\theta)$  and  $(G^{\prime},\theta^{\prime})$  of the form  $(G,\theta)\gets (G,\cdot)\to (G^{\prime},\cdot)\to (G^{\prime},\theta^{\prime})$ , then we have
$$
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)$  is the marginal probability of reaching  $G$  from the initial state  $G_0$  with any (complete) trajectory  $\tau = (G_0, G_1, \ldots, G_{T-1}, G)$ :
$$
P _ {\phi} (G \mid G _ {0}) \triangleq \sum_ {\tau : G _ {0} \sim G} \prod_ {t = 0} ^ {T - 1} P _ {\phi} \left(G _ {t + 1} \mid G _ {t}\right),
$$
using the conventions  $G_{T} = G$ , and  $P_{\phi}(G_0\mid G_0) = 1$ .
The proof of this theorem is available in Appendix C.3.3. The marginal distribution  $P_{\phi}^{\top}(G,\theta)$  is also called the terminating state probability in Bengio et al. (2023).
3.3 Learning objective
One way to find the parameters  $\phi$  of the forward transition probabilities that enforce the SubTB conditions in (6) for all  $(G,\theta)$  and  $(G^{\prime},\theta^{\prime})$  is to transform this condition into a learning objective. For example, we could minimize a non-linear least squares objective (Bengio et al., 2021, 2023) of the form  $\mathcal{L}(\phi) = \mathbb{E}_{\pi}[\Delta^{2}(\phi)]$ , where the residuals  $\Delta (\phi)$  depend on the conditions in (6), and  $\pi$  is an arbitrary sampling distribution of complete states  $(G,\theta)$  and  $(G^{\prime},\theta^{\prime})$ , with full support; see Malkin et al. (2023) for a discussion of the effect of  $\pi$  on training GFlowNets, and Appendix D.1 for further details in the context of JSP-GFN.
In addition to the SubTB conditions for undirected paths of length 3 given in Section 3.2, we can also derive similar SubTB conditions for other undirected paths, for example those of the form  $(G,\theta)\gets (G,\cdot)\to (G,\tilde{\theta})$ , where  $\theta$  and  $\tilde{\theta}$  are two possible sets of parameters associated with the same graph  $G$  (see also Figure C.2). When the reward function  $R(G,\theta)$  is differentiable wrt.  $\theta$ , which is the case here, we show in Appendix C.3.2 that we can write the SubTB conditions for these undirected paths of length 2 in differential form as
$$
\nabla_ {\theta} \log P _ {\phi} (\theta \mid G) = \nabla_ {\theta} \log R (G, \theta). \tag {7}
$$
This condition is equivalent to the notion of score matching (Hyvärinen, 2005) to model unnormalized distributions, since  $R(G,\theta)$  here corresponds to the unnormalized posterior distribution  $P(\theta \mid G,\mathcal{D})$ , where the normalization is over  $\theta$ . We also show in Appendix C.3.2 that incorporating this information about undirected paths of length 2, via the identity in (7), amounts to preventing backpropagation through  $\theta$  and  $\theta'$  (e.g., backpropagation with the reparametrization trick) in the objective
$$
\mathcal {L} (\phi) = \mathbb {E} _ {\pi} \left[ \left(\log \frac {R \left(G ^ {\prime} , \bot \left(\theta^ {\prime}\right)\right) P _ {B} \left(G \mid G ^ {\prime}\right) P _ {\phi} \left(\bot \left(\theta\right) \mid G\right)}{R \left(G , \bot \left(\theta\right)\right) P _ {\phi} \left(G ^ {\prime} \mid G\right) P _ {\phi} \left(\bot \left(\theta^ {\prime}\right) \mid G ^ {\prime}\right)}\right) ^ {2} \right], \tag {8}
$$
where  $\perp$  denotes the "stop-gradient" operation. This is aligned with the recommendations of Lahlou et al. (2023) to avoid backpropagation through the reward in continuous GFlowNets. The pseudo-code for training JSP-GFN is available in Algorithm 1.
3.4 Parametrization of the forward transition probabilities