Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network
Authors: Tristan Deleu
Mizu Nishikawa-Toomey
Jithendaraa Subramanian
Nikolay Malkin
Laurent Charlin
Yoshua Bengio
Venue: NeurIPS 2023 poster (2023)
Keywords: bayesian network
bayesian
structure learning
causal discovery
gflownet
TL;DR: We propose to approximate the joint posterior distribution over the graphical structure of a Bayesian Network and the parameters of its conditionals using a GFlowNet.

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.

Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network
Tristan Deleu
Mizu Nishikawa-Toomey
Jithendaraa Subramanian'
Nikolay Malkin
Laurent Charlin'
Yoshua Bengio 1.4
Mila Quebec Al Institute
Abstract
Generative Flow Networks (GFlowNets)), class of generative models over discrete and structured sample spaces, have been previously applied the problem of inferring the marginal posterior distribution over the directed acyclic graph (DAG) of a Bayesian Network, given dataset of observations. Based recent advances extending this framework to non-discrete sample spaces, we propose in this paper approximate the joint posterior over not only the structure of Bayesian Network, but also the parameters its conditional probability distributions. We use a single GFlowNet whose sampling policy follows 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 non-line 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.
Introduction
As compact representation for complex probabilistic models, Bayesian Networks are a framework of choice many fields, such computational biology (Friedman al., 2000; Sachs 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-$e known, it can be used to perform probabilistic inference for queries of interest with variety of exact or approximate methods (Koller and Friedman, 2009). However, this graphical structure is unknown, one may want to infer it based dataset of observations D.
In addition to being challenging problem due 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 addressed by taking Bayesian perspective on structure learning and inferring the posterior distribution P(G '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(D G). which can only be done efficiently in closed form for limited classes of models, such as linear Gaussian (Geiger and Heckerman, 1994),
Unniversité de Montréal, MMeGill University, HHEC Montréal, 'CIFAR Senior Fellow.
Correspondence: Tristan Deleu (deleutri (uili..eeemmllaeeeeee Code available at: https//ghhhb..ommttsssddddeeeppp
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

discrete models with Dirichlet prior (Heckerman et al., 1995), or non non-linear models parametrized with Gaussian Process (von Kügelgen al., 2019).
While there exists a vast literature on Bayesian structure learning to approximate the marginal posterior distribution, inferring the joint posterior P(G,0 0 'D) over both the DAG structure G of the Bayesian Network and the parameters 0 of its conditional probability distributions-thh 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 8), where the dimensionality of the latter may even depend on G. However, modeling the posterior distribution over 0 has the notable advantage that the conditional probability distributions can be more flexible (e.g., parametrized by neural networks): in general, computing P(D| G, 8) is easier than computing the marginal P(D 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 al., 2021, 2023) proved 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 D) was treated a time, starting from the empty graph over d variables, following a learned transition probability. Nishikawa-Tee Toomey et al. (2023) have also proposed to use a GFlowNet to infer the joint posterior P(G.0 0 D); however, they used it in conjunction with Variational Bayes to update the distribution over 0. getting around the difficulty of modeling a continuous distribution with GFlowNet.
sequential decision process, where edges are added one at
In this paper, we propose to infer the joint posterior over graphical structures G and parameters of the conditional probability distributions 0 of a Bayesian Network using 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 GelowNets. The generation of a sample (G, 0) 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 @ are chosen once the structure 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, 0 '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.
2 Background
Notations. Throughout this paper, we will work with directed graphs G E), where V is set of nodes, and E C V x V is a set of (directed) edges. For a node X E V. we denote by Pao (X) the set of parents of X in G, and Cha(X) the set of its children. For two nodes X, Y E V, X Y represents a directed edge (X, E E (denoted X + Y E G), and X ~ Y represents a directed path from X to Y. following the edges in E (denoted X Y G).
In the context of GelowNets (see Section 2.2), an undirected path in a directed graph g = (V,E) between two states So, 8n E V is sequence of vertices (So, $1...., 8n) where either (si, 81+1) 8 or ($141, st) E & (i.e., following the edges of the graph, regardless of their orientations).
2.1 Bayesian structure learning
Bayesian Network is probabilistic model over d random variables {Xi,..., whose joint distribution factorizes according to a directed acyclic graph (DAG) G as
d
P( Xa (.G) = II p(x (X)):
where 0 = {0),... 02) represents the parameters of the conditional probability distributions (CPDs) involved in this factorization. When the structure is known, Bayesian Network offers a rep- resentation of the joint distribution that may be convenient for probabilistic inference (Koller and

---

And
GS
B()
Figure 1: Structure of the Generative Flow Network to approximate the joint posterior distribution P(G, 0 D). A graph G and parameters are constructed as follows: starting from the empty graph Go. (1) the graph G is first generated one edge at 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 0 (red), conditioned on the graph G. Finally, given G and 0, (3) we receive a reward R(G, 0) (green).
Friedman, 2009), as well as learning its parameters 9 using a dataset D = (x(1), x x(N), where each (1) represents an observation of the d random variables [aa}
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 Bayesian perspective, we may want to either model the (marginal) posterior distribution P(G (D) over only the DAG structures, or the joint posterior P(G,6 A D) over both the structure as well as the parameters of the CPDs 0.
2.2 Generative Flow Networks
Generative Flow Network (GFlowNet; Bengio et al., 2021, 2023) is a generative model over structured sample space X. The structure of the GFlowNet is described by DAG G whose vertex set is the state space S, where X' S S. This should not be confused with the DAG in a Bayesian Network: in fact, each state in could itself represent the structure of Bayesian Network (Deleu et al., 2022). A sample x E X' is constructed sequentially by following the edges of g, starting at a special initial state 8o. until to which indicates the end the sequential process. The states in X' are those for which there is directed edge - 8f: they called complete states! and correspond distribution induced by the GFlowNet. A path 8p ~ 85 in G is called a complete trajectory. An example of the structure of GelowNet given in Section 3.1 and Figure 1.
reach the state We define special terminal state 8,, transition
valid samples of the
Every complete state x E X' is associated with a reward R(x) 2 O, indicating the unnormalized probability of x. We use the convention R(s)) = for any state E S'X, since they do not correspond to valid samples of the distribution. Bengio et al. (2021) showed that if there exists a function Fa(s + s') > 0 defined over the edges of G, called a flow, that satisfies the following flow-matching conditions
(s8)++))- - 1 F = R(s') d'ECho(s')
sEPag(')
for all the non-terminal states s' E S, then the GFlowNet induces distribution over complete states proportional to the reward. More precisely, starting from the initial state sp, we sample a complete trajectory $1...., 25. following the forward transition probability, defined as
P(s St) x Fo(s; + s(+1).
(3)
with the conventions 8T = x and 8T+1 = 8y, then the marginal probability that a trajectory sampled following P terminates in X' is proportional to R(x). The flow function Fs + 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 X' proportional to R.
Bengioo et al. (2023) also call these states terminating; we follow the naming conventions of Deleu et al. (2022) here avoid ambiguity, is closely related to our work.
3

---

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 D). Their model, called DAG-GF DAG-GFlowNer. 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 nodes, while enforcing the acyelicity constraint at every step of the generation (i.e., an edge is not added if it would introduce cycle). Its structure is illustrated at the top of Figure I, where it forms the part of the graph shown in blue.
Instead of working with flows F& + G'), as in Section 2.2, DAG-GFlowNel directly learns the forward transition probability (a G) that satisfies the following alternative detailed balance conditions for any transition G + G' (i.e., G' is the result of adding single edge to G):
PGGG) G') G) = (& G))s G'),
where PB 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 - G' still induces a distribution over DAGs x R(G). Therefore, to approximate the posterior distribution P(G D), they used R(G) P(D PGG) the reward of G. In particular, this requires evaluating the marginal likelihood P(D 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-continuuus hybrids. Building on top of DAG-GFlowNet, we propose here to approximate the joint posterior P(G, B (D) over both the structure of the Bayesian Network G, but also the parameters of its conditional probability distributions 0. Unlike in VBG though (Nishikawa-Toom 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 GFlowNet.
3.1 Structure of the GFlowNet
Unlike in DAG-GFlowNet, where we model distribution only over DAGs, here we need to define a GFlowNet whose complete states are pairs (G, 0), where G' is DAG and 0 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 time, following Deleu et al. (2022), and then the corresponding parameters 0 are generated, conditioned on G. We denote by (G. states where the DAG G has no parameters 0 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, 0) E X, whereas (G,-) E S SiX
Starting at the empty graph (Go. the DAG is constructed one edge a time during the first phase, following the forward transition probabilities Po(G' G). This first phase ends when a special "stop" action is selected with Ps. 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 0 conditioned on G, following the forward transition probabilities Po(O G). All the complete states (G, 0), for fixed graph G' and any set of parameters 0, can be seen as forming an (infinitely wide) tree rooted at (G, )).
Since we want this GFlowNet to approximate the joint posterior P(G,0 D) x P(D,O, G), is natural to define the reward function of a complete state (G, 0) as
R(G, 8) = P(D 6, P( G))(
(5)
where the likelihood model P(D G, may be arbitrary (c.g., a neural network), and decomposes according to (1), P(0 G') is the prior over parameters, and P(G) the prior over graphs. When the dataset D is large, we can use a mini-batch approximation to the reward (see Section 3.5).
"Since the states of the GelowNet here are pairs of objects, Pa(0 G) (resp. Ps(G" | G)) is an abuse of notation, and represents Po 0) (G. .)) (resp. Po((G", ..) (G.-)).
4

---

3.2 Subtrajectory Balance conditions
To obtain generative process that samples pairs of (G, 0) proportionally to the reward, the GelowNet needs to satisfy some conditions such (4). However, we saw in Section 2.3 that satisfying this particular formulation the detailed balance conditions in (4) yields a distribution x R(.) only all the states are complete; unfortunately, this is not the case here since there exists states of the form (G, 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. t (G.) (G', 8') in the GFlowNet (see Figure C.2), where G' is therefore the result of adding 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(G', "'G',0')E PB Go G) = RRG, P& G)P GG)..
(6)
Note that the SubTB condition above very similar to the detailed balance condition (4) used in DAG-GFlowNet, where the probability of terminating (s G') has been replaced by (o G). in addition to the reward now depending on both G and 0. Moreover, while does not seem to appear in (6), the terminal state 5 of the GFlowNet is still present implicitly, since are forced to terminate once we have reached a complete state (G, 0). and therefore Po G 0) = 1 (see App. C.3.I).
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 xx R(G, 0) if the SubTB conditions in (6) are satisfied for all pairs (G, 0) and (G', ,').
Theorem 3.1. If the SubTB conditions in (6) are satisfied for all undirected paths of length between any () and (G', ,0) of the form (G, 0) +- (G.-) - (G', .-) - (G', 0'), then we have
PJ (G, 2 (G Po G) x R(G,0),
where (o Go) is the marginal probubility of reaching G from the initial state Go with any (complete) trajectory T = (Ga, Gr,...,,TTI.. G):
GT-1;
T-1
Pd[G Go) A
Σ
II
PolGeet Gt).
T:Go-G 1=0
using the conventions Gr = G, and P (G Go) = 1.
The proof of this theorem is available in Appendix C.3.3. The marginal distribution PJ (G, is also called the ferminating state probability in Bengio et al. (2023).
3.3 Learning objective
One way to find the parameters & of the forward transition probabilities that enforce the SubTB conditions in (6) for all (G,0) and (G', 0¹) 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 L(Ø)= E, [o)]. where the residuals A()) depend on the conditions in (6), and is an arbitrary sampling distribution of complete states (G, 0) and (G", 0'), with full support; see Malkin al. (2023) for a discussion of the effect of π on training GFlowNets, and Appendix D.I for further details the context of JSP-GEN.
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, 0) + (G.-) > (G. where 0 and 0 are two possible sets of parameters associated with the same graph G (see also Figure C.2). When the reward function R(G, 0) is differentiable wrt. 0, 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
V log (0 G) V logog R(G, 0)
(7)
5

---

This condition is equivalent to the notion of score matching (Hyvärinen, 2005) to model unnormalized distributions, since R(G, 8) here corresponds to the unnormalized posterior distribution P G. G. where the normalization is over O. 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 0 and 6' (e.g., backpropagation with the reparametrization trick) in the objective
Em (108 R(G', PB(G ()) G) R() Pa(G' G) (0) G*)
21
[ = E, EE,
(8)
where IL denotes the "stop-gradins poop-gradient" operation. This is aligned with the recommendations of Lahlou et al. (2023) avoid backpropagation through the reward in continuous GFlowNets. The pscudo-code for training JSP-GFN is available in Algorithm 1.
3.4 Parametrization of the forward transition probabilities
In Section 3.1, we saw that the process of generating (G. 8) follows two phases: first we construct G one edge at time, until we sample a specific "stop" action, at which point we sample 0 conditioned on G. All these actions are sampled using the forward transition probabilities Po G)) during the first phase, and Po G) during the second one. Following Deleu 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& (stop) G); then, conditioned on this first decision, we either continue adding an edge to G to reach G" with probability (G G, stop)) (phase 1), or sample 0 with probability Po(O) G,stop) (phase 2). This hierarchical model can be written as
Po(G' G) (1 (1- (stopp G)) -stop)
Po(O G) (stop G)P G,stop).
(10)
We use neural networks to parametrize each of the three components necessary to define the forward transition probabilities. Unlike in DAG-GFlowNe though, which uses a linear Transformer to define (stop G) and (G | G, -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 graph embedding of G, as well as 3 embeddings u,, Vi, w; for each node X; in G
g, {u, Us wi } = SelfAttention (GraphNet, (G)).
We can parametrize the probability of selecting the "stop" action using 9 with (stopp G)= fo(g).
where fo is a neural network with sigmoid output; note that if we can't 't add any edge to G without creating cycle, we force the end of the first phase by setting (stsoop G) = 1. Inspired by Lorch et al. (2021), I), the probability of moving from G to G' by adding the edge X, - X, is parametrized by
(11) where my is binary mask indicating whether adding X; + X, is a valid action (i.e., if it is not already present in G. and if doesn't 't introduce a cycle; Deleu et al., 2022). Finally, the probability of selecting the parameters 0 of the CPD for the variable X; is parametrized with a multivariate Normal distribution with diagonal covariance (unless specified otherwise)
(G' G,, stop) x m.j cxp(xi, (u vz),
(0 G,stop) = (0, (w:),, (w:)).
(12)
where Me and o² are two neural networks, with appropriate non-linearities to guarantee that (² is a well-defined diagonal covariance matrix. Note that (e, | Gtstop) effectively approximates the posterior distribution P(0; G, D) once fully trained. Moreover, in addition to being an approximation of the joint posterior P(G,0 D). the GFlowNet also provides an approximation of the marginal posterior P(G 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 0.
3.5 Mini-batch training
Throughout the paper, we have assumed that we had access to the full dataset of observations D in order to compute the reward R(G, 0) (5). However, beyond the capacity to have an arbitrary likeli- hood model P(D) 0, G) (e.g., non-linear), another advantage of approximating the joint posterior

---

P(G.0 8 (D) is that we can train the GFlowNet using mini batihbaches of observations. Concretely, for
mini-batch B of M observations sampled uniformly at random from the dataset D, we can define
N M
log Re(G, 0) log P(O G) + log P(G)
+
Σ EB
log ((x(m)) |,
(13)
which is unbiased estimate of the log-reward. The following proposition shows that minimizing the estimated loss based on (13) wrt. the parameters & of the GFlowNet also minimizes the original objective in Section 3.3.
Proposition 3.2. Suppose that B is a mini-batch of M observations sampled uniformly at random from the dataset D, and let Lb()) be the learning objective defined in Section 3.3, where the reward has been replaced by the estimate R (G,0) in (13). Then we have ( S Eg [[o())]
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 convex function. In the case of the square loss, Proposition C.1 gives stronger result in terms of unbiasedness of the gradient estimator.
4 Related work
Bayesian Structure Learning, There is a vast literature applying Markov chain Monte Carlo (MCMC) methods approximate the marginal posterior P(G D) over the graphical structures of Bayesian Networks (Madigan et al., 1995; Friedman and Koller, 2003; Giudici and Castelo, 2003; Viinikka et al., 2020). However, since the parameter space in which lives depends on the graph structure G. approximating the joint posterior P(G,0 0 '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 al., 2022). Similar 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 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 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- Toommyy et al. (2023) proposed to learn the joint posterior P(G, 0 (D) over structures and parameters with a GFlowNet, combined with Variational Bayes to circumvent the challenge of learning a distribution over continuous quantities 0 with a GFlowNet. Atanackovic et al. (2023) also used a GFlowNet called DynGFN to approximate the posterior of a dynamical system. Similar to (Nishikawa-Toome)
et al., 2023) though, they used the GFlowNet only to approximate the distribution over graphs G, making the parameters 0 deterministic function of G (i.e., P(O G, D) x 8(0 Φ)). Here, we leverage the recent advances extending these models to general sample spaces (Lahlou al., 2023), including continuous spaces (Li al., 2023), in order to model the joint posterior within single GFlowNet.
7

---

Edge festures
Ba - Iog P( G, P)]
JSP-GEN (diag)
0.8
RMSE
Pearson's 's
MH-MC*
0.357 + 0.022
0.067 0.143
5.39 = 1.41 x 10a
Gibbe-MCa
0.357 + 0.022
0.028 0.127
9.02 + 1.54 x 10
B-GES"
0.263 +: 0.070
0.635 +: 0.180
1.56 = 0.97 x 10a
B-PC"
0.305 +: 0.067
0.570 0.138
1.57 = 0.87 x 102
DIBS
0.312 0.038
0.737 0.071
9.49 + 7.34 10
0.2
BCD Nets
0.215 4: 0.055
0.819 0.097
7.04 = 3.21 x 10"
0.9981
VBG
0.237 0.037
0.816 0.064
1.24 = 0.49 10
0
0.2
0.4
0.6
0.8
JSP-GEN (ding)
0.018 0.005
0.988
0.001
-4.91
10"
Exnct posterior
JSP-GFN (full)
0.019 0.007
0.998 0.001
-5.00
10
(a) Edge features
(b) Quanlitative comparisoo with the exact posterior
Figure 2: Comparison with the exact posterior distribution, 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; X; for each of the 20 datasets. (b) Quantitative evaluation of different methods for joint posterior approximation, both in terms of edge features and cross- cross-entropy of sampling distribution and true posterior P(0 G, D); all values correspond to the mean and 95% confidence interval across the 20 experiments.
5
Experimental results
5.1 Joint posterior over small graphs
We can evaluate the accuracy of the approximation returned by JSP-GEN by comparing it with the exact joint posterior distribution P(G.0 0 D). Computing this exact posterior can only be done in limited cases only, where (1) P(O 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' D). We consider here models over = 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 0. For the graphs, we compare the approximation and the exact posterior on different marginals of interest, also called features (Friedman and Koller, 2003); e.g., the edgefeature feature corresponds to the marginal probability of a specific edge being in the graph. Figure (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 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 Markoy features (Deleu et al., 2022); see Appendix D.2.2.
To evaluate the performance the different methods as an approximation of the posterior over 0, we also estimate the cross-entropy between the sampling distribution of 0 given G and the exact posterior P(O G, D). This measure will be minimized the model correctly samples parameters from the true P(0 G, G)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 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 2-layer MLP, for total of [0] = 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.
8

---

Negative log- fng-likelhood
9,000
4,000
7,000
log P (G, ()
4,000
5,000
-000
3,000
-700
1,000
1.000
-80
0.9864
M-Me
G-MC
B.PC
B-GES
Dibs
BCD Nets
VBG
JSP-GFN
M-Mcs
G-MCS
Dibs
Jsp JP--GEN
-2.600
-2,000
log P(D, G,0)
Linear Gaussian
(b) Non-linear Gaussian
(c) Non-lincar 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 D. (c) Linear correlation between the log-reward (X-axis) and the terminating state log pogg-robbbility (y-axis) for 1, 000 samples (G, 0) from JSP-GFN; the color of each point indicates the number of edges in the correponding graph. A 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² & Gibbs-MC*; 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-Toome) Toomey et al., 2023) for the experiment with linear Gaussian CPDs, as they are not applicable for non-linear CPDs. In Figure (a-b), report the performance of these joint posterior approximations in terms of the (cxpected) negative log- log-likelihood (NLL) on held-out observations. We observe that JSP-GEN achieves 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 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. 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 metries, and in order assess the quality of the approximation of the posterior in the absence reference P(G, 0 'D), we also show in Figure (c) how the terminating state log-probability log P (G, 0) of JSP-GFN correlates with the log-reward for a non non-lineat Gaussian model. Indeed, as stated Theorem 3.1, we should ideally have log P.J. (G, 0) perfectly correlated with log R(G, 0) with slope 1, as
also
log P (G,0) log P(G,8 8 D) log R(G, 0) log P(D).
(14)
We see that there is indeed strong linear correlation across multiple samples (G,0) from JSP- GFN, with slope A close to 1, suggesting that the GFlowNet is again an accurate approximation of the joint posterior, least around the modes it captures. Details about how log P (G, 0) 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 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 (Faton and Murphy, 2007). For the gene expression dataset, we used a subset of = 2, 628 observations of d = genes from (Sethuraman et al., 2023). Details about the experimental setups available in Appendix D.4.

---

2,000
1,250
log P, (G.0) ()
-1,500
-2, -0,000
16
Flow eytometry
11, N 4,200
Gene expression
61, 628
-1,750
-2,000
MH-MC²
2.578 10*
1.073 10"
Gibbe-MC3
3.490 10"
4.882 10"
2,300
12
JSP-GFN
2.679 10"
2.651 10
1.0272
A 0.7817
10
-3,500
log PDD, G.
log P(D, G,8)
-00
(a) Flow cytometry data (subsample, N 100)
(b) Flow cytometry data
(c) Negative log-likelihood
Figure 4: Performance of ISP-GEN on real-world biological data. (a) Comparison of the terminating state log-probability log P (G,0) returned by JSP-GFN with the log-reward log R(G, 0) on a subsample N = 100 datapoints of the flow cytometry dataset D. (b) Same comparison with the full dataset 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- lig-liked on held-out data.
At this scale, using the whole dataset DD 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 mini-batches of data, in place of R(G, 0) 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 Figure (c), along with an evaluation of the NLL on unseen observations and interventions. Beyond the ability to work with real data, these experiments allow highlight some other capacities of JSP-GFN: (1) handling discrete and (2) interventional data (flow cytometry), as well as (3) learming a distribution over larger graphs (gene expression).
To measure the quality of the posterior approximation returned by JSP-GFN, compare in Figure 4 the terminating state log-proo log-probability log PJ (G.0) with the log-reward log R(G, (G, similar to Fig- ure We observe that there is correlation between these two quantities; unlike in Figure though, we observe that the slope is not close to 1, suggesting that JSP-GFN underestimates the probability of (G, 0). We also observe that the graphs are "clustered" together; this can explained by the fact that the posterior upproximation is concentrated only few graphs, since the size of the dataset 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 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-GEN with 2 methods based on MCMC terms of their negative ligeliilldood 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- likeeiidooodd on interventions unseen in D.
Conclusion
We have presented JSP-GFN, approach to approximate the joint posterior distribution over the structure of Bayesian Network along with its parameters using 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 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.
10

---

Acknowledgements
We would like thank António 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 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.quebecc.
References
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.
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.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzale Vinicius Zambaldi, Mateusz Malinowski. Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint, 2018.
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.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet Foundations. Journal of Machine Learning Research (JMLR), 2023.
Bertrand Charpentier, Simon Kibler, and Stephan Günnemann. Differentiable DAG Sampling. International Conference on Learning Representations, 2022.
David Maxwell Chickering. Optimal Structure Identification With Greedy Search. Journal of Machine Learning Research, 2002.
Chris Cundy, Aditya Grover, and Stefano Ermon. BCD Nets: Scalable Variational Approaches for Bayesian Causal Discovery. Advances in Neural Information Processing Systems, 2021.
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.
Daniel Eaton and Kevin Murphy, Bayesian strueture learning using dynamic programming and MCMC. Conference on Uncertainty in Artificial Intelligence, 2007.
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.
Martin A Fischler and Robert C Bolles. Random sample consensus: paradigm for model fitting with applications to image analysis and automated cartography. Communications of the ACM, 1981.
Nir Friedman and Daphne Koller. Being Bayesian About Network Structure. A Bayesian Approach to Structure Discovery in Bayesian Networks. Machine Learning, 2003.
Nir Friedman, Moises Goldszmidt, and Abraham Wyner. Data Analysis with Bayesian Networks: A Bootstrap Approach. Proceedings of the Fifteenth conference Uncertainty in Artificial Intelligence, 1999.
Nir Friedman, Michal Linial, Iftach Nachman, and Dana er. Using bayesian networks to analyze expression data. In Proceedings of the fourth annual international conference on Computational molecular biology, 2000.

---

Eva-Maria Fronk. Model Selection for Dags via RJMCMC for the Discrete and Mixed Case. Technical report, Mathiig-Maxmmmliaan Univerrsittt München, 2002.
Dan Geiger and David Heckerman. Learning Gaussian Networks. In Uncertainty in Artificial Intelligence. 1994.
Paolo Giudici and Robert Castelo. Improving Markov chain Monte Carlo model search for data mining. Machine learning, 2003.
David Heckerman, Dan Geiger, and David Chickering. Learning bayesian networks: The combi- nation of knowledge and statistical data. Machine learning, 1995.
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.
Aapo Hyvärinen.
Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research, 2005.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector- -Botos, 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.
Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez Hernandez-Garcia. and Yoshua Bengio. GFlowNets for Al-Driven Scientific Discovery. Digital Discovery, 2023.
Daphne Koller and Nir Friedman. Probabilistic Graphical Models: Principles and Techniques. MIT press, 2009.
Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández- Gareía, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A Theory of Continuous Generative Flow Networks. International Conference on Machine Learning, 2023.
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.
Yinchuan Li, Shuang Luo, Haozhi Wang, and Jianye Hao. CFlowNets: Continuous control with
Generative Flow Networks. International Conference on Learning Representationss 2023.
Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. Advances in Neural Information Processing Systems, 2016.
Lars Lorch, Jonas Rothfuss, Bernhard Schölkopf, and Andreas Krause. DiBS: Differentiable Bayesian Structure Learning. Advances in Neural Information Processing Systems, 2021.
Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf. Amortized Inference for Causal Structure Learning. Advances Neural Information Processing Systems, 2022.
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.
David Madigan, Jonathan Gavrin, and Adrian Raftery. Enhancing the Predictive Performance of Bayesian Graphical Models. 1994.
David Madigan, Jeremy York, and Denis Allard. Bayesian Graphical Models for Discrete Data. International Statistical Review, 1995.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems, 2022.
12

---

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.
Joris M Mooij, Sara Magliacane, and Tom Claassen. Joint Causal Inference from Multiple Contexts. Journal of Machine Learning Research, 2020.
Mizu Nishikawa- 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,
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GelowNets with local credit and incomplete trajectories. International Conference on Machine Learning, 2023.
Karen Sachs, Omar Perez, Dana Pe'er, er, Douglas Lauffenburger, and Garry Nolan. Causal protein-signaling networks derived from multiparameter single-cell data. Science, 2005.
Muralikrishnna G Sethuraman, Romain Lopez, Rahul Mohan, Faramarz Fekri, Tommaso Bian- calami, and Jan-Christian Hütter. NODAGS-FIEN Flow: Nonlinear Cyclic Causal Structure Learning. International Conference Artificial Intelligence and Statistics (AISTATS), 2023.
Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, Prediction, and Search. MIT press, 2000.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention All You Need. Advances Neural Information Processing Systems, 2017.
Jussi Viinikka, Antii Hyttinen, Johan Pensar, and Mikko Koivisto. Towards Scalable Bayesian Learning Causal DAGs. Advances Neural Information Processing Systems, 2020.
Julius von Kügelgen, Paul Rubenstein, Bernhard Schölkopf. and Adrian Weller. Optimal experi- mental design via Bayesian optimization: active causal structure learning for Gaussian process networks. NeurlPS Workshop "Do the right thing": machine learning and causal inference for improved decision making, 2019.
Benjie Wang, Matthew R Wicker, and Marta Kwiatkowska. Tractable Uncertainty for Structure Learning. International Conference on Machine Learning, 2022.
David Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust Scheduling with GFlowNets. International Conference on Learning Representations, 2023.
Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying Generative Models with GelowNets and Beyond. International Conference on Machine Learning Beyond Bayes workshop, 2022.
Xun Zheng, Bryon Aragam, Pradeep Ravikumar, and Eric Xing. DAGS with NO TEARS: Continuous Optimization for Structure Learning. In Advances in Neural Information Processing Systems, 2018
Heiko Zimmermann, Fredrik Lindsten, Wil-Willem van de Meent, and Christian A. Naesseth. A Variational Perspective on Generative Flow Networks, arXiv preprint, 2022.
13

---

Appendix
A Positioning JSP-GEN in the Bayesian structure learning literature
We give comparison between method JSP-GFN, and various methods based variational inference in Table A.I. We also include methods based on GFlowNets, namely DAG-GFlowNet (Deleu et al., 2022) and VBG (Nishikawa-Toome) et al., 2023), they are effectively variational methods (Malkin et al., 2023; Zimmermann et al., 2022).
Table A.I:
Comparison different methods based on variational inference and GFlowNets for Bayesian structure learning. See the text for detailed description of each category.
Joint &
Non Linear
DAG
Discrete Obs.
Max. Parents
Mini Batch
Sampler
Support
VCN (Annadani et al., 2021)
X
X
BCD Nets (Cundy et al., 2021)
X
X
X
DiBS (Lorch et al., 2021)
/
/
X
/
X
X
TRUST (Wang et al., 2022)
X
X
V
X
X
VI-DP-DAG (Charpentier et al., 2022)
X
V
J
X
X
V
x
AVICI (Lorch al., 2022)
X
J
X
X
DAG-GFlowNet (Deleu et al., 2022)
X
X
V
VBG Keshikawa-Toomee et al., 2023)
JSP-GEN (Ours)
Joint G & 0. This category indicates whether the model can approximate the joint posterior distribution P(G, 0 D) over both graphical structures G and parameters of the CPDs 0, or if they are limited to approximating the marginal posterior P(G D). As we have seen Section 1, approximations of the marginal posteriors limit the classes of models these methods can be applied to, namely those where the marginal likelihood can be computed analytically.
Non-Linear. This indicates whether the model can be applied to Bayesian Networks whose CPDs are parametrized by non-linear function (c.g., neural network). While most methods approximating the marginal distribution may be applied to non-linear CPDs parametrized by Gaussian Process (von Kügelgen et al., 2019), we only consider here methods that explicitly handle non-linearity (e.g.. this eliminates DAG-GFlowNet (Deleu et al., 2022), since the authors only considered a linear Gaussian and discrete settings). Annadani et al. (2021) mentioned the extension of VCN to non-linear causal models future work. While VBG (Nishikawa-Toome Gaussian models, the framework may also be applicable to non-linear models.
al., 2023) has only been applied to linear
DAG Support. This indicates whether the posterior approximation is guaranteed to have support over the space of DAGs. VCN (Annadani et al., 2021) and DiBS (Lorch et al., 2021) only encourage acyclicity via prior term, inspired by continuous relaxations of the acyclicity constraint (Zheng et al., 2018), meaning the those methods may return graphs containing cycles; for example in practice, Deleu al. (2022) reports that 1.50% of the graphs returned by DīBS contain cycles (for d - 11). AVICI (Lorch et al., 2022) uses a similar prior term when applied to Structural Causal Models (SCMs), although in general this framework does enforce acyclicity by design, allow flexibility on other domains (e.g.. for modeling gene regulatory networks; Sethuraman et al., 2023). TRUST (Wang et al., 2022) guarantees acyclicity via distribution over variable orders, that can learned using Sum-Product Networks.
Discrete Observations. This indicates whether the posterior approximation may be applied to Bayesian Networks with discrete random variables. Although VCN (Annadani et al., 2021) was only applied to linear Gaussian models, the authors mention that this approach is also applicable to discrete random variables. Similarly, while there is no experiment in (Lorch et al., 2021) applying DiBS to discrete domain, this extension can found in the official code released. AVICI (Lorch al., 2022) assumes access to generative model P(D) G), making it possibly applicable to discrete domains as
14

---

well. Since it builds on DAG-GFlowNet (Deleu et al., 2022), VBG (Nishikawa-Toow) Toomey et al., 2023) should also inherit its properties, and therefore may also be applicable to discrete random variables. Similarly, since TRUST (Wang et al., 2022) may use DíBS as its underlying routine for structure learning, should also inherit the properties of DiBS.
Maximum Parents. This category indicates whether maximum number of parents can be specified for each variable in the DAGs returned by each method. Although this is a very common constraint used in the structure learning literature to improve efficiency (Koller and Friedman, 2009), none of the variational methods for Bayesian structure learning allow for such a (hard) constraint. Some methods may introduce sparsity-inducing prior (Lorch et al., 2021; Cundy et al., 2021), use post post-processing of the sampled DAGs (Charpentier et al., 2022) to reduce the number of edges in the sampled graphs. This can be naturally added in a GFlowNet, by masking out the actions adding certain edges that would violate this constraint; in fact, in the official code released for both DAG-GFlowNet (Deleu et al., 2022) and VBG (Nishikawa- Wishikawa-Toomey et al., 2023) (both using the same environment), this option is available.
Sampler. This category indicates whether one can sample graphs and parameters from the model once fully trained. DiBS (Lorch et al., 2021) uses a particle-based approach (Liu and Wang. 2016) to approximate the posterior (marginal, joint), and therefore the number of particles is fixed ahead of time; once fully trained, it is impossible to sample new pairs of graphs and parameters from this model. TRUST (Wang et al., 2022) can also use Gadget (an MCMC approach to Bayesian structure learning; Viinikka et al., 2020) as its routine for structure learning, and therefore this would allow sampling from the trained model.
Mini-Batch. This indicates whether the model can be updated with mini-batch of observations from D, or if the full dataset must be used. DiBS (Lorch et al., 2021) uses mini-batch updates for their experiments on protein signaling networks, where the number of datapoints N = 7466 is large. Note that unlike JSP-GFN here, neither DAG-GFlowNet (Deleu et al., 2022) nor VBG Wishikawa-Toomee et al., 2023) may be updated using mini-batches, since the maginalization over 0 makes all observations in D mutually dependent (conditioned on G). See Appendix C.3.4 for details on how to use mini-batch training with JSP-GFN.
B Broader impact & limitations
B.1 Broader impact
While structure learning of Bayesian Networks constitutes one of the foundations of causal discovery (also known as causal structure learning), it is important to emphasize that shy of any assumptions, the relationships learned from observations in a Bayesian Network are in general not causal, but merely statistical associations. As such, care must be taken interpreting the graphs sampled with JSP-GFN (or any other Bayesian structure learning method considered in this paper) as being causal. This is especially true when applying structure learning methods to the problem of scientific discovery (Jain et al., 2023). Assumptions that would allow causal interpretation of the graphs include using interventional data (as in Section 5.3), or parametric assumptions.
Although the graphs returned by JSP-GFN are not guaranteed to be causal, treating structure learning from a Bayesian perspective allows us to view identification of the causal relationships a softer way. Indeed, instead of returning a single graph which could be harmful from a causal perspective (notably due to the lack of data, see also Section I), having a posterior distribution over Bayesian Networks allows us to average out any possible model that can explain the data.
B.2 Limitations
Expressivity of the posterior approximation. Throughout this paper, we use Normal distributions (with diagonal covariance, except for JSP-GFN (full) in Section 5.1) to parametrize the approxima- tion of the posterior over parameters (o slo) (see Section 3.4). This limits its expressivity to unimodal distributions only, and is assumption which is commonly used with Bayesian neural networks. However in general, the posterior distribution P(O G, D) may be highly multimodal, especially when the model is non-linear (the posterior is Normal when the model is linear Gaussian,
15

---

Po(#o my)
Pe(sm sn)
Pala
Pol
Pulem
(d) Generalized Subtrajectory Balance
(a) Detailed Balance
(b) Trajectory Dalance
(c) Subtrajectory Balance
Figure C.I: Illustration of the different GFlowNet objectives. (a) The detailed balance condition operates at the level of transitions - s', whereas (b) the trajectory balance condition operates on complete trajectories 50 ~* s,. (c) The subtrajectory balance condition operates on partial trajectories 8m Sn, and can be (d) generalized to undirected paths with common ancestor Sk. We use Pp ~+ sn) to denote the product of Pp along the path 8m ~y Sn (and similarly for PB).
see Appendix D.5. To see this, consider non-linear model whose CPDs are parametrized by 2-layer MLP (as in Sections 5.2 and 5.3). The weights and biases of both layers can be transformed in such a way that the hidden units get permuted, while preserving the outputs; in other words, there are many sets of parameters 9 leading to the same likelihood function, and under mild assumptions on the priors P(0) G) and P(G), they would have the same posterior probability P(0 G,D).
To address this issue of unimodality, we can use more expressive posterior approximations Po(0 G, stop), such as ones parametrized with diffusion-based models, with normalizing flows; both of these models are drop-in replacements in JSP-GEN, since their likelihood can be explicitly computed. An alternative is also consider multiple steps of a continuous GFlowNet (Lahlou et al., 2023), instead of single one, to generate 0.
Biological plausibility of the acyelicity assumption. One of the strengths of JSP-GFN, and DAG- GFlowNet before it (Deleu et al., 2022), is the capacity to obtain a distribution over the DAG structure of Bayesian Network (and its parameters). The acyclicity assumption is particularly important in order to properly define the likelihood model in (1). However in some domains, such as biological systems, there may exist some feedback processes that cannot be captured by acyclic graphs (Mooij et al., 2020). In particular, the DAGs found in Section 5.3 and Appendix D.4 must be carefully interpreted. As general framework though, the GFlowNet used in JSP-GFN can be adapted ignore the acyclic nature of the graphs sampled by ignoring parts of the mask m in Sec. 3.4. Alternatively, we can view the generation of cyclic graph by unrolling it, as in Atanackovic al. (2023).
C Details about Generative Flow Networks
Throughout this section, we will use both PF and Po (to emphasize the parametrization on Φ, as in the main text) to denote equally the forward transition probability-the notation Pr being more commonly used in the literature on GFlowNet (Bengio et al., 2023).
C.1 Alternative conditions
GFlowNets were initially introduced using the fow-matching conditions (Bengio et al., 2021), 1), as described in Section 2.2. However, there have been multiple alternative conditions that, once satisfied, also offer the same guarantees as the original flow-matching conditions (namely, a GFlowNet satisfying any of those conditions would sample complete states proportionally the reward).
One of those alternative conditions are the detailed balance conditions (Bengio et al., 2023), inspired by the literature Markov chains. These conditions are given for any transition s' in the GFlowNet as
(C.1) where F(s) is a flow function, that may also be parametrized by a neural network. The detailed balance condition is illustrated in Figure C.1 (a). Bengio et al. (2023) showed that if the detailed balance conditions are satisfied for all the transitions s + s' in the GFlowNet, then the distribution induced by the GFlowNet is also proportional to R(s). Deleu et al. (2022) adapted the detailed balance conditions in the case where the states of the GelowNet are complete, in order to avoid having to learn separate flow function (see Section 2.3).
Fr)) 8) ((s)) 8')
16

---

Another alternative condition called the detailed balance conditions (Malkin et al., 2022). operates not at the level of transitions, but at the level of complete trajectories. For a complete trajectory (80, .......... 8T, 8),, the trajectory balance condition is given by
T-1
z II (S st) R(ST) II PB(s: $1+1), 1=1
(C.2)
t=1
with the convention 81+1 8f. and where Z is the partition function of the distribution (i.e., Z =meex R(x)); in practice, Z is parameter of the model that is being learned alongside the forward and backward transition probabilities. The trajectory balance condition is illustrated in Figure C.I (b). Again, if the trajectory balance conditions are satisfied for all complete trajectories in the GFlowNet, then the induced distribution is proportional to R(s).
C.2 Subtrajectory balance conditions
Also introduced in (Malkin et al., 2022), the subtrajectory balance conditions are generalization of both the detailed balance and trajectory balance conditions to partial trajectories of arbitrary length. For partial trajectory T = (sm: Sm+1.-... Sn), the subtrajectory balance condition is given by
n-1
F(sm)
II Pr(s441 I (s) = F(sn) II PB(se 8t+1), t=
(C.3)
t=
where again F(s) is a flow function (as in (C.1)). This condition encompasses both conditions in Appendix C.1, since we recover the detailed balance condition in (C.1) with partial trajectories of length (i.e., transitions), and also the trajectory balance condition in (C.2) with complete trajectories (note that F(so) Z: Bengio et al., 2023). The subtrajectory balance condition is illustrated in Figure C.1 (c). Madan et al. (2022) also proposed to combine subtrajectory balance conditions for partial trajectories of different lengths create novel objective called SubTB(A), inspired by TD(A) in the reinforcement learning literature.
This subtrajectory balance condition in (C.3) can also be generalized to undirected paths going "back and forth" (Malkin et al., 2022). For an undirected path between 8m and Sn, this (generalized) subtrajectory balance condition can be written as
((mm)
II t=k
Pn(st | 8141) II Pr(s1+1 | 8t) F(S.) t=
II t=k
Pn(se 8141) II t=k
PF |s),
(C.4) where 8& is a common ancestor of both 8m and 8n. This condition is illustrated in Figure C.1 (d). While these subtrajectory balance conditions (generalized or not) offer more flexibility, they are guaranteed yield GFlowNet inducing distribution proportional R(s) only if these conditions are satisfied for all the partial trajectories of any length. In particular, they provide no guarantce in general if those conditions are satisfied for all partial trajectories of fixed length, which is the case in this paper (see Section 3.2). Although this result may be extended with weaker assumptions, we prove in Appendix C.3.3 that the GFlowNet does induce a distribution o R(s) in our case.
C.3 Proofs
C.3.1 Subtrajectory balance conditions for undirected paths of length
Consider an undirected path of length 3 of the form (G, 0) + (G.-) (G', (G', where G' is the result of adding new edge to the DAG G (see Figure C.2). Since the state (G.-) is common ancestor of both complete states (G,0) and (G', 0'), we can write the subtrajectory balance conditions (C.4)
F (GG,8)PP PB P)P (B G F(G))(G P G')P((O (C.5)
where we abuse the notation PB(G | 0) again to denote PB (G. (G,0). In fact, since the complete state (G,8) E X has only a single parent state (G.-). we necessarily have PB(G 0) = 1 (and similarly for (G', 8')) Furthermore, we can use the observation from (Deleu et al., 2022) to
write the flow F(G, 8) of complete state (G, f) as function of its reward
R(G,0)
F( PF(ss (G)
(C.6)
17

---

Undirected paths of length 3 R(G", Pp(G P G) R(G,0)P G) 6*)
Undirected paths of length 2 (differential form) Volog Py(0 G) = V log R)
H(G.99
Figure C.2:
Illustration of the undirected paths of length 3 (red) and of length 2 (blue) considered in this paper, and their corresponding subtrajectory balance conditions.
We can simplify (C.6) even further by observing that the GFlowNet used here, sy is the only child of the complete state (G, 0) X. In other words, complete state (G, 8) is not directly connected to any other (G, 0); this is the (infinitely wide) tree structure rooted at (G, mentioned in Section 3.1. Since sy is the only child of (G, 0), we then necessarily have Pr ($y (G,0) 1, and therefore F(G,0) = R(G, 0). With these simplifications, (C.5) becomes
R(G,0) 0) (G' P) G') = R( PD (r) G):
(C.7)
which is the subtrajectory balance condition in (6).
C.3.2 Integrating undirected paths of length 2
Similar to Appendix C.3.1, we consider here an undirected of length 2 of the form (G,0) + (G. +
(G, 0) (see Figure C.2). Since -G is a common ancestor (a common parent in this case) of both complete states (G, 0) and (G, (), we write the subtrajectory balance conditions (C.4)
PGGG (0 G) = F(G,0) 1
Pr G)-
(C.8)
Using the same simplifications in Appendix C.3.1 (Pn(G 0) Pn(G 0) 1), get the following subtrajectory balance conditions for the undirected paths of length
() G) = R(G, O, (0 G).
(C.9)
Note that these conditions are effectively redundant if the SubTB conditions over undirected paths of length 3 (6) are satisfied for all possible pairs of complete states 0) and (G",0'. Indeed, if we write these conditions between (G, 0) and (G", 0') on the one hand, and between (G, 0) and (G',0" 0')
on the other hand (with a fixed G' and 0')
R(G", O(G PA G') G) PFG G) (G)
(C.10)
R(G', P'(G G') G) R(G, (G' G)
(G).
(C.11)
we get the same subtrajectory balance conditions over undirected paths of length 2 as in (C.9):
R(G,0 0) (GG G')
R(G,0)
(C.12)
(r G) PF Pr G')
Pr G)
However, since the SubTB conditions (6) are only satisfied approximately in practice it might be advantageous to also satisfy (C.9) addition to those in (6). The equation above provides an alternative way to express (C.9). Indeed, (C.12) shows that the function
(a A log R(G, 0) - log Pr (r (G)
(C.13)
constant, albeit with a constant that depends on the graph G. Since this function is differentiable. this is equivalent to V V0f = 0, and therefore we get the differential form of the subtrajectory
balance conditions (7)
V log (p G) V log R(G,8).
(C.14)
As we saw in Section 3.3, one way to enforce the SubTB conditions over undirected paths of length is to create learning objective that encourages these conditions to be satisfied, and optimizing
18

---

using gradient methods. The learning objective has the form (()) = E,, (0)],, where (Ø) is non-linear residual term
RIG', (O))PP(GGG)PPOC G (b) log O) ( G')
(C.15)
Suppose that the parameters Φ of the GFlowNet are such that the subtrajectory balance conditions in (C.14) are satisfied for any (G,0). Although this assumption is unlikely to be satisfied in practice, they will eventually be approximately satisfied over the course of optimization, given the discussion above about the relation between (C.9) and (C.7). Since f and 6' depend on ø (via the reparametrization trick since they are sampled on on-policy, see Section 5), taking the derivative of (o), we get
d
dø (0) (())
log R(G", 0²) + Po ( dø
(C.16)
- log R(G, 0) log Po(G" G) - log Po G')].
Using the law of total derivatives, we have
d
0 00
a
d8
+ a
log (0 G) log R(G,
log (o G)
log R(G,0)
log (o |G)
do
0e
dó
(C.17)
0 log (6 G).
(C.18)
and similarly for the terms in (G', The derivative of the objective then becomes
dΦ -=
d -
log Po G) log Po 1 as
a
G')
d
log P G)
(C.19)
dø
dø
An alternative way to obtain the same derivative in (C.17) is take d0/dø 0 instead, meaning that we would not differentiate through 0 (and 0'). Using the stop-gradient operation 1, this shows that the following objective
( a E. [ =E Ex ((oB P() G)P G) 27
(C.20)
B(G ()) | G*)
takes the same value and has the same gradient (C.19) as the objective in (C.15) when the subtrajectory balance conditions (in differential form) over undirected paths of length are satisfied.
While optimizing (C.20) alone leads to eventually satisfying the subtrajectory balance conditions over undirected paths of length it may be advantageous explicitly encourage this behavior, especially in cases where d is larger and/or for non-line linear models. We can incorporate some penalty to the loss function, such
A Br | log Po(0 G) Vo RoG 0|
(()) (())
(C.21)
2
Vor log Po G') Vo log R(G 8 8)1l
C.3.3 Marginal distribution complete states
Theorem 3.1. If the SubTB conditions in (6) satisfied for all undirected paths of length between any (G,0) and (G', 0') of the form (G, 0) + (G,-) - (G', - (G',0'), then we have
P (G. e (G) Go G) R(G,0),
where Po(G Go) is the marginal probability of reaching G from the initial state Go with any (complete) trajectory T = (Go, GI....,Grr GT-1, G):
T-1 Po Go) A Σ II
Po G.),
T:Gem t=a
using the conventions GT G, and Po
Go) 1.
19

---

Proof. We assume that the SubTB conditions are satisfied for all undirected paths of length 3 between any (G, 8) and (G', ('). that is
R(G', (2',0') PB G)), - R(G, G)P,(0' G*).
(C.22)
Let G + Go be fixed DAG different from the initial state, and a set of corresponding parameters. Let (Go..... GT-1. G) be an arbitrary trajectory from Go G. where use the convention GT = G. For any t < T, if Or is fixed set of parameters associated with Gr, then the SubTB conditions above can written for every timestep as
Pa G,) R(G,, (.P,B.)))(G... (.) Ge+1); (C.23)
again, using the convention 8 = 0. Taking the product the ratio between Pa and PB over the trajectory T, get
T-1 Po (G Gi)
T-1 Po G.)((C+,,,0)))) II t=o
(C.24)
PB (G. G++1)
R(G, Po Ge+1)
0
Po(Os Go) 0) R(Go. Ps G)
(C.25)
Moreover, the backward transition probability PB. defined only over the transitions of the GFlowNet, induces a distribution over the trajectories from Go to (Bengio et al., 2023), meaning that
T-1 Il
Σ
PB(G, G++1) 1.
(C.26)
G. t=0
Therefore, we have
T-1 II Po G+)
Po(G
Σ
(C.27)
G:GG tmm
Po(Oo | Go) R(G,0)
T-1
Σ
II Pa(G, G(11)
(C.28)
R(Go, 0o)
G:GomG t=0
Po Gu) R(G,0)
(C.29)
R(Go, 0o)
We saw in Appendix C.3.2 that Po Go)/R(Gg, 0b) is independent of the value of 9 the SubTB conditions are satisfied for all undirected paths of length 3 (see (C.12)). This concludes the proof: (G Go) ( G) x R).
C.3.4 Mini-batch training
Proposition 3.2. Suppose that a mini-batch of M observations sampled uniformly at random from the datasel D, and let Lg( be the learning objective defined Section 3.3, where the reward has been replaced by the estimate R Rg,(G,) in (13). Then we have L(Ø) I Es [Cs()).
Proof. We will first show that log Rg () defined (13) an unbiased estimate of the log-reward log R f) under uniform distribution of the mini-batches B. We observe that by conditional independence of the observations (n) given G and 8, have
N
log P(D ()
= Σ log (m 10, G) = NE log P( 6, G)], n=1
(C.30)
where the expectation is taken wrt. the uniform distribution over the observations in D. It is important to note that we can decompose the likelihood term as (C.30) because the observations are mutually independent given G and 0; if we were only conditioning G (i.c., using the marginal likelihood,
20

---

as in (Deleu et al., 2022)), then those observations would not be conditionally independent in general. Similarly, we have
Eg
log ((x(mm)
(m) 10, = MM log P(æ 00 G)
(C.31)
(m(m) EB
Therefore, shows that the estimate of the log-reward is unbiased:
N
E5 log RB(G, =
EB
Σ
log ((x( (
+ log P(O G) + log P(G) (C.32)
M
EB
log P(D 0, G) log P(0 G) log P(G) log R(G, 0).
(C.33)
Recall that the estimate the loss is defined as [x(Φ) E, A [)] where the residual is defined
by
Am(Φ) log R((( (L() G*)
(C.34)
Taking the expectation of this estimated loss wrt. random mini-batch B, get
Es [b({)]
= Ex Ex [6)
(C.35)
2 Ex Es AB(D)] 3
(C.36)
E,, [+)
(C.37)
L(d),
(C.38)
where we used the convexity of the square function and Jensen's inequality in (C.36), and the unbiasedness of log Rg(G, ()) (as well as log R 0')) in (C.37); recall that ( is given in (C.20) (see also (8)).
Proposition C.1. The mini- batch gradient estimator is unbiased, iee, vol(o)= [B(V Therefore, the local and global minima of the expected mini-batch loss coincide with those of the full-batch loss.
Proof. We now show that the gradient estimator unbiased. We observe that
( (B)))
(C.39)
since only the terms corresponding to the rewards differ between A(d) and (s()), and they do not depend on c. Therefore,
Es V LB()
Es V AB (0)
2 Es [B(o))
2- Es
=2A(Ø) A())
-volA(0)
-vL(Ø),
desired. This implies the expected mini-batch loss and full-batch loss differ by constant and have the same set of local and global minima. This constant happens to equal
RE (G,0)
Varg
log RB(G' 00)
and showing the difference equals this constant yields an alternative proof.
21

---

Algorithm JSP-GEN training procedure
1: Initialize the trajectory at Go the empty graph
2: repeat
3: Interactions with the environment
4: for K steps do
5: Sample whether the trajectory continues not: Pa(stop G+)
6: if a is the "stop" action then
7:
Reset the trajectory: G Go is the empty graph
8:
else
9:
Sample new graph Gt+1 Po Gr, stopp Store the transition G2 + G++1 in the replay buffer
10:
12:
Update of the parameters ø
Sample a mini-batch B of transitions G G' from the replay buffer
14:
Sample 8 ~ Po G. stop) (and similarly for 6') for each transition in B
Evaluate the rewards R(G, () and R(G',
D (5)
16:
Evaluate the loss (()) based on B b Update the parameters with stochastic gradient methods
b (8)
Φ + Φ - (f(
18: until convergence criterion
D Additional experiments & experimental details
In addition Bayesian structure learning methods based on variational inference (Lorch al., 2021; Cundy et al., 2021) or GFlowNets Kushikawa-Toomer et al., 2023). we also consider 2 baseline methods based on MCMC, and 2 methods based on bootstrapping (Friedman al., 1999). as introduced in (Lorch et al., 2021). Metropolis-Hattnngs MC (MH-MC*) samples both graphical structures and parameters jointly at each move, whereas vetropolis-within- alternates between updates of the structure, and updates of the parameters; note that MC² here refers to the Structure MCMC algorithm (Madigan et al., 1995). In terms of bootstrapping methods, we consider a variant (called B-GES*) based on GES (Chickering, 2002), and another (called B-PC") based on PC (Spirtes et al., 2000). However, since this would yield an approximation the marginal posterior P(G D) only, the parameter sample / corresponding to DAG G correspond to the parameter inferred by P(0 G, D) (see Appendix D.5.
MC3 (Gibbs-MC*)
D.1 Sampling distribution
In Section 3.3, we saw that the learning objective of JSP-GFN can be written = where æ is sampling distribution over (G,0) 0) and (G", 0') with full support. We use a combination of on-policy (π = Pa)) and off-policy (π is different from P&) in order to train the GFlowNet: taking inspiration from (Deleu et al., 2022), transitions G + G" are sampled oof-po from a replay buffer, whereas their corresponding parameters 0 and 6' are sampled on on-policy using current ( G).
Therefore, a key difference with Deleu et al. (2022) is that the reward R(G, 0) in (5) is calculated "lazily" when the loss is evaluated (i.e., only once and 0' are known), opposed to being computed during the interaction with the state space and stored the replay buffer alongside the transitions. Algorithm gives a description of the training algorithm JSP-GEN.
D.2
Joint posterior small graphs
D.2.1 Data generation & modeling
Data generation. We follow the same data generation process as in (Deleu et al., 2022: Lorch et al., 2021). More precisely, we first sample graph from Erdös-Rényi model (Erdös and Rényi. 1960) over = nodes, with d edges on average (a setting typically referred to as ERI). Once the structure of the graph G* known, we sample the parameters 6* of the linear Gaussian model randomly from a standard Normal distribution A 1). The linear Gaussian model is defined as
X, Σ
02 X1 + 6i,
(D.1)
(XE)
22

---

JSP GFP (diag)
0.9984
0.9982
0.9980
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Exact posterior
Exact posterior
Exact posterior
(a) Edge features
(b) Path features
(c) Markov features
0.6
9.5
0.5
GFP (full)
0.6
0.5
0.6
0.4
0.4
0.4
0.2
0.2
0.9950
0.9976
0.9972
0
Exacl posterior
Exact posterior
Exact posterior
(d) Edge features
(e) Path features
(I) Markuv features
Figure D.I: Comparison the marginals over graphs in terms features computed with the exact posterior (X-axis) and the approximation given by JSP-GFN (y-axis). (a-c) Comparison with JSP-GFN (diag); (d-f) comparison with JSP-GEN (full). Each point corresponds to a pair of variables (X,, for each of the 20 datasets.
where 01 ~ (0, 1), and e; ~ (0, 0.01); this defines all the CPDs necessary for (1). Finally, we ancestral sampling generate N 100 observations create the dataset D.
Modeling, We use linear Gaussian Bayesian Network to model the data, where the CPD for the variable X; can be written P(x: Pac (X:((:);00 NN o²), where
d #4 = Σ 1x, E (a) j=1
(D.2)
and where o² 0.01 is fixed variance across variables, matching the variance used for data generation, We place unit Normal prior over the parameters 0.j of the model: P(O.g G) -V 1). This model differs from the widely used BGe score (Geiger and Heckerman, 1994) that o² treated hyperparameter here, instead of parameter of the model, and therefore the resulting log-reward is not score equivalent (i.e., placing the same reward for Markov equivalent DAGs; Koller and Friedman, 2009). We used a uniform prior graphs. Under this model. the posterior distribution P(0 G, D) can be computed analytically, and the proof is available Appendix D.5.1.
D.2.2 Comparing JSP-GEN with the exact posterior against features
In Section 5. evaluated the accuracy of the posterior approximation returned by JSP-GEN by comparing them on the edge features, i.e., the marginal distribution of an edge X; - X, is present in the graph:
P(X, X, D) = EPIapp)(( [1(X, X, G)],
(D.3)
where the expectation in (D.3) is either over the true (marginal) posterior P(G (D) (for the x-axis of Figure 2), or over the distribution PJ induced by the GFlowNet (discarding 6, see Section 3.4, for the y-axis). Besides edge features, there exists other marginals of interest (Friedman and Koller, 2003:
Deleu et al., 2022), such the path feature and the Markov feature. The path feature corresponds to
23

---

Table D.I: Quantitative comparison between different Bayesian structure learning algorithmss and the exact posterior on small graphs with 5 nodes. For each feature, the root meas- meas-square error (RMSE) and Pearson's 's correlation coefficient between the features computed with the posterior approximation and the exact posterior are reported. Values are reported as the mean and 95% confidence interval across 20 different datasets.
Edge features
Path features
Markov features
RMSE
Pearson's 's
RMSE
Pearson's r
RMSE
Pearson's r
MH-MC
0.357 0.022
0.067 + 0.143
0.368 + 0.027
0.045 + 0.179
0.341 + 0.017
0.064 = 0.217
Gibbs-MC
0.357 1: 0.022
0.028 + 0.127
0.367 1 0.026
0.150 1 0.162
0.341 + 0.018
0.062 - 0.159
B-GES'
0.263 + 0.070
0.635 + 0.180
0.302 + 0.080
0.544 + 0.230
0.129 + 0.022
0.955 = 0.026
B-PC
0.305 0.057
0.570 0.138
0.349 0.058
0.471 1 0.154
0.354 0.072
0.821 + 0.087
DiBS
0.312 1 0.038
0.737 0.071
0.357 1 0.041
0.710 1 0.079 0.079
0.504 0.052
0.643 - 0.093
BCD Nets
0.215 = 0.055
=819
0.266 + 0.057
0.774 0.109
0.327 0.040
0.850 = 0.067
VBG
0.237 + 0.037
0.816 0.064
0.284 4 0.027
0.799 + 0.050
0.434
0.738 + 0.091
JSP-GEN (diag)
0.018 = 0.005
0.998 + 0.001
0. 0.022 + 0.005
0.998 + 0.001
0.019 - 0.006
0.999 = 0.001
JSP-GEN (full)
0.019 + 0.007
0.998 + 0.001
0.021 + 0.007
0.998 + 0.002
0.020 + 0.008
0.999 + 0.001
the marginal probability of a path X; ~ X, being present in the graph, and the Markov feature is the marginal probability of node X; being in the Markov blanket of Xj. In other words
P(X; ~ X5 D) = E P(G [ ~ X, E G)) P(X, MB(X,) D) (M())) 1(X; MBc(X,))
(D.4)
(D.5)
where (X, denotes the Markov blanket of X, in G. For the posterior approximation returned by JSP-GFN (and other methods), the expectations appearing in (D.4) & (D.5) are computed under the posterior approximation, and can be estimated using a Monte Carlo estimate over sample DAGs from the model.
In this experiment, we considered two variants of JSP-GEN, The first one, called JSP-GFN (diag), where Ps (0, stop) in (12) is parametrized as a Normal distribution with a diagonal covariance matrix; this adds a modeling bias since here the exact posterior P(0; G, G) is a Normal distribution with full covariance (see Appendix D.5.1). To control for this bias, the second model called JSP-GFN (full) assumes that (12) has a full covariance matrix, as in VBG (Nishikawa- Nisoik et al., 2023).
In Figure D.I, we show similar plot as in Figure 2 (a) for all these features, for both models JSP-GEN (diag) and JSP-GEN (full). We observe that for features, the approximation of the posterior given by JSP-GEN is very accurate (as confirmed by the Pearson's 's correlation coefficients). Interestingly, the more expressive model JSP-GEN (full) seems to perform slightly worse than JSP- GFN (diag); this is also confirmed in part by the quantitative measures in Table D. This can be explained by the additional number of parameters of the neural network 6 necessary to output the higher dimensional full covariance matrix (more precisely, lower-triangulat matrix corresponding to its Cholesky decomposition). In Table D. we show a quantitative comparison across the different Bayesian structure learning methods on the three features. in terms of RMSE and Pearson's 's correlation coefficient. Similar to Section 5. we observe that JSP-GFN provides a more accurate posterior approximation (at least in terms of its marginal over G) than other methods.
D.2.3 Evaluation of the posterior approximations over parameters
In order to evaluate the quality of the posterior approximation over 6, we measure how likely sample parameters from the approximation are under the exact posterior distribution P(0 G, D). More precisely, we compute the cross- coossoyy between the posterior approximation P& (G, f) and the exact joint posterior P(G,0 0 D): given distribution Po approximating the joint posterior P(G,0 0 D), we estimate
K Σ
Ep log P(0 [, = log P(B(k) ((()) D),
(D.6)
where {(G(A), { are K samples of the posterior approximation (G We use this measure as it can be estimated from samples.
24

---

D.3 Gaussian Bayesian Networks from simulated data
D.3.1 Data generation & modeling
Data generation. For the linear Gaussian experiment, the data generation process follows the process described for small graphs in Appendix D.2.1, except that we sample ground truth graphs G* from an Rénös-Rééyí model with 2d edges on average (a setting commonly referred as ER2), for d = 20 variables.
For the non-linear Gaussian experiment, the data generation process is also similar, except that the CPDs are parametrized using a 2-layer MLP with 5 hidden units and a ReLU activation function (Lorch al., 2021) with randomly generated weights. We also sample N 100 observations to create the dataset D.
Modeling. We use Gaussian Bayesian Network to model the data, where the CPD for the variable X; can be written as P(X; (X,()X;; 0,) = (N d²), where
N4 MLP(M.X;00.
(D.7)
where M, diag (X1, E Pac(X,)), ..... L(Xa Paa():)) and X (X1,..., Xa). The variance o² 0.01 is fixed across variables, and matches the variance used for data generation. Following Lorch et al. (2021) and matching the data generation process. we use a 2-layer MLP with 5 hidden units and ReLU activation function, for total of 0 2, 220 parameters. The priors over parameters P(0 G) and over graphs P(G) follow the ones described in Appendix D.2.1.
D.3.2 Estimation of the log-terminating state probability
When the graphs are larger, it becomes impossible to compare the posterior approximation returned by ISP-GFN with the exact joint posterior P(G, 0 'D) directly, since the latter becomes intractable (even with linear Gaussian model) due to the super-exponential size of the sample space. Alternatively, since the terminating state probability P. (G, 0) of JSP-GFN should ideally be equal to the joint posterior (see Theorem 3.1), we have
log P (G,0) v log P(G,0 0 D) = log R(G,0) - log P(D),
(D.8)
where log P(D) is constant corresponding to the log-partition function. Therefore, we can compare the log-terminating state probability log P. (G, 6) with the log-reward log R(G, 0) (which we can compute analytically) for different samples (G, 0), and find a linear relation. This evaluation strategy was introduced in (Bengio al., 2021).
However, recall from Theorem 3.1 that the terminating state probability defined as T-1 P (G, = Po Go) G) = Po Σ II Po Gt), GG), T.Gom E=
(D.9)
where the summation is over all the possible trajectories (Go, GI....,Gi GT) from Go to GT G.
G is DAG with K edges, then there are K! such trajectories (i.e., the K edges could be added in any order), meaning that this sum is also intractable. We can leverage the fact that the backward transition probability (B G++1) induces a distribution over the trajectories Go ~y G (Bengio et al., 2023) to write (& Go) as
(G Go)
Σ Po(T)
(D.10)
T:GonG
PB(T) PB(T)
Po(T)
(D.11)
G:
K! E..EE [P)(T))
(D.12)
where Pn(T) = 1/K!, since the backward transition probability here is fixed to be uniform over parent states. This suggests way get an unbiased estimate of P Go), hence of Pj (G,0). based Monte-Carlo estimation:
K!
M
Po Go) &
Σ
(m())
(D.13)
M
m==
25

---

where {T( Im=1 are trajectories from Gu to G, sampled by removing one edge at time uniformly at random, starting G (i.e., following the backward transition probabilities PR).
While (D.13) provides an unbiased estimate of (a Gn). in practice the variance of this estimate will be large due to the combinatorially large space of trajectories, and therefore due to the wide range of values (T) may take. In order to reduce the variance, we can first identify some trajectories that would contribute the most to the sum in (D.10), and complement them with some randomly sampled trajectories as in (D.13). In other words, if we have access to a subset Toop of B trajectories T: Gu ~- G that have large (T)),, then
M K!-B M Po(((m)). (D.14)
PG Go) = Σ + +
Po(T) ~ P
B
TETUp
TTToop
TETTOp
where the sample trajectories (r(m) Jm=1 can obtained using rejection sampling, with uniform proposal as above. The estimate in (D.14) is still unbiased, but with a lower variance. We can use beam-search to find the "top-scoring" trajectories in Topp with a beam-size B. More precisely, we need to run beam-search, starting at Go. in such way that the trajectories are guaranteed to end at G. We can achieve this by constraining the set of actions one can take at each step of expansion to move from graph G2 to G(+1 G2 U {e} (using this notation to denote that G is the result of adding the edge e to Gt), with the following score:
(G G,) = 1(e E G) | Gt).
(D.15)
In other words, only keep transitions corresponding adding edges that in G. Note that even though Ps is not a properly defined probability distribution (it does not sum to 1), we can still use this scoring function to run beam-search in order to find "top-scoring" trajectories.
D.3.3 Additional comparisons with the ground-truth graphs
In Section 5.2, we compare JSP-GFN against other Bayesian structure learning in terms of their negative log-lik log-likelihood on held-out data. In addition to the negative log- ligliike though, there exists standard metrics in the structure learning literature that compare the posterior approximation with the ground truth graphs G* used for data generation. For example, the expected SHD, that is estimated from graphs {Gi,..., G sampled from the posterior approximation as
E SHH R Σ SHD(G&. G*),
(D.16)
where SHD(G, G*) counts the number of edges changes (adding, removing, reversing an edge) necessary move from G to G*. There is also the area under the ROC curve (AUROC) that compares the edge marginals estimated from the posterior approximation (i.e., the edge features, see Appendix D.2.2) and the target G*. We report these metrics Figure D.2.
Although these metrics are used in the Bayesian structure learning literature, they also suffer from a number of drawbacks (Lorch et al., 2022). Namely, these metrics do not properly assess the quality of the posterior approximation (i.e., how close the approximation is the the true P(G, A D)), but merely how close the sampled graphs are from G In general, and especially when the data is limited, the graphs sampled from the true posterior have no reason a priori to match exactly G*. Moreover, the expected SHD tends to favour overly sparse graphs the one hand, posterior approximations that collapse completely at G* on the other hand, both situations indicating a poor approximation of the true P(G, D).
D.4 Learning biological structures from real data
D.4.1 Modeling
Protein signaling networks from flow cytometry data. Since the flow cytometry data has been discretized, we use non-linear model with Categorical observations. The CPDs parametrized using a 2-layer MLP with 16 hidden units and a ReLU activation function, i.c., X, Paa((X)) ~ Categorical( where
T MLP MMP(MMXX.) M,X;0,)..
(D.17)
26

---

Expocted
20
M-Mos
G-Mes
B-PO" B-GRs
bitss
Bop Net
Vna
JSP-GPS
M-Mes
G-Mca
Dibs Jap 8PP-GPM
(a) Linear Gaussian
(b) Non-liness Gaussian
0.7
AUROC
0.88
0.5
M-3roe
G-MC
B.PC
B-GES
DRRS
BCD Nets
VBG
JSP JSP-GPN
M-MeC
G-Me
DBb
JSP-PPN
[c] Linear Gaussian
(d) Non-linear Gaussian
Figure D.2: (a-b) Comparison of JSP-GFN with other Bayesian structure learning methods in terms of the expected-SHD to the ground truth graphs G* used for data generation. (c-d) Comparison in terms of Area Under the ROC curve (AUROC) to the ground truth graphs G*.
where X encodes the discrete inputs one-hot encoded values, and the MI.P has a softmax activation function for the output layer. In total, the model has |0] = 6, 545 parameters. The priors over parameters P(0 G) and over graphs P(G) follow the ones described in Appendix D.2.1.
Gene regulatory networks from gene expression data. Gene expression data is composed of either non-zero continuous data (when gene is expressed) or (exactly) zero values (when the gene inhibited). To capture this type of observations, model CPDs of the Bayesian Network as zero-inflated Normal distributions:
(D.18) where is the result of 2-layer MLP with 16 hidden units, as in (D.17). The parameters of the CPDs contain the parameters of the MLP,, as well as the mixture parameter a; and the variance the observation noise o², for total of [0] = 61, 671 parameters. The priors over parameters P(0 G) and over graphs P(G) follow the ones described in Appendix D.2.1.
P(X, Pag (X:); 02) = (x,) (1- (1 of)
D.5 Proofs
D.5.1 Posterior of the linear Gaussian model
Recall that the CPD for the linear Gaussian model can be written as
d
P(X: Pac(X)):0.) NW o²) where
Σ 1(x Pac(X.))
(D.19)
j=1
and where o² is fixed hyperparameter. Moreover, we assume that the parameters have unit Normal prior associated to them, meaning that
P(O, G) =
V(Ho.o7)
+
(D.20)
do otherwise,
27

---

where μo = 0, o0 1, and ōg is the Dirac measure at 0, indicating that this parameter is always inactive.
We want to compute the posterior distribution P(o. D) this is sufficient, since we know that the parameters of the different CPDs mutually conditionally independent given G and D. Let X Nxd be the design matrix of the dataset D (i.e., the observations x!*) concatenated row-wise), and by abuse of notation, we denote by X; the ith column of this design matrix. Let D; be a diagonal matrix, dependent on G, defined
Dj diag (1 (X1 E Paa((x))...... 1 1Xx E Pac((:)))
(D.21)
We can rewrite the complete model above as
P(0: G) NDD, D.) P(X, Pag (X.); 0.) = -,(XX,,,, I
(D.22)
(D.23)
We abuse the notation above by treating Dirac distribution the limiting case of Normal distribution with variance O. Given this form, we can easily identify that the posterior over 8, is Normal distribution
14 = Egj
Diuo + D.XTR x,
0
0²
P(O: D, -N E,)
where
(D.24)
Σ
1 0 DT D.
DIXTTX
where we used the conventions 1/0 oc and 0o 0. The masked entries of D; will correspond to zeroed-out entries in fú, and to rows and columns of Σ; being equal to zero, effectively reducing the dimensionality of the distribution the number of parents of X; in (e.g.. this has impact on the normalization constant of this distribution). In the limit case where X; has no parent in G, we
recover P(O, G, D) = 60.
28


Reviews:
[{"review_id": "lCkjlBafR6", "reviewer": "Reviewer_MmTz", "summary": "This paper uses GFlowNets to learn a generative model that samples from the posterior distribution of both the graphical structure and corresponding parameterisation of Bayesian networks given some observation. The paper demonstrates that sub-trajectory balance conditions suffice to ensure that the GFlowNet induce a distribution that is well proportional to the reward function associated with a final sample. Experimental results demonstrate that this approach is able to recover that follows closely the true posterior distribution. The method outperform alternative strategies for linear gaussian CPD and is competitive with alternative approach for non-linear Gaussian.", "questions": "- Why did not you try more expressive parameterisation of the CPD, such as with normalizing flows?\n- Why did you only consider N=100 in 5.2?\n- If the edge feature is a probability why do you used MSE, rather than a better metric for comparing distributions such as the KL?", "limitations": "/", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The problem statement and extension of Gflownets proposed to solve it are very sound. Jointly learning discrete structures and continuous parameterisation has many promising applications in various domains\n- Although merging continuous and discrete GFlownets does not seem very original in itself, it requires to adapt the balancing conditions appropriately and making it work is a strong results in itself. \n- Experimental results are very promising.", "weaknesses": "Although I have found the idea presented in the paper of learning a generative model over the structure and continuous parameters very interesting, I also find the paper hard to read and understand technically for a reader non-familiar with GFlowNets. For instance, it is unclear exactly how the policy is parameterised and optimised. Although section 3.3 discusses the learning objective it does not provide a clear explanation on how exactly the training process happens, which makes it hard for me to argue for more than a weak accept. I consider important for a paper at a top venue to be clear by itself and to not require to check all supplementary materials and or related work to be understood by someone with the sufficient mathematical background. I would suggest authors to merge sections 2 and 3 in order to save space and give a more precise description of the parameterisation of the policy and training algorithms used. Moreover, I find strange that 5.3 is also so weakly developed in the main materials whereas it hints very important capacity of the proposed method, which might handle well interventional data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their review and their suggestions to improve the presentation of our method.\n\n> *I also find the paper hard to read and understand technically for a reader non-familiar with GFlowNets.*\n> \n\nWe are sorry to hear that the presentation was not sufficiently clear. Our objective was precisely to make the description of GFlowNets as self-contained as possible in the paper, in particular in Section 2. We introduced all necessary notions in Section 2.2 with a description of how objects are sampled from the GFlowNet, and described the prior work DAG-GFlowNet in Section 2.3 since our submission builds on top of it. We also included additional details about training objectives in Appendix C.1 & C.2.\n\n> *For instance, it is unclear exactly how the policy is parameterised and optimised. Although section 3.3 discusses the learning objective it does not provide a clear explanation on how exactly the training process happens, which makes it hard for me to argue for more than a weak accept.*\n> \n\nWe invite you to read the general rebuttal comment (**Practical implementation**) for a detailed response. We want to again emphasize that all the components were described in details in the main text.\n\n> *I consider important for a paper at a top venue to be clear by itself and to not require to check all supplementary materials and or related work to be understood by someone with the sufficient mathematical background.*\n> \n\nWe absolutely agree with the reviewer here, and that’s why we strived for making this submission as self-contained as possible. However with only 9 pages available, we made our best effort to provide a description of GFlowNets as extensive as possible in the main text, and we necessarily have to defer details to the supplementary material, to leave enough space for a complete description of our contributions (Section 3).\n\n> *I would suggest authors to merge sections 2 and 3 in order to save space and give a more precise description of the parameterisation of the policy and training algorithms used.*\n> \n\nSection 2 discusses the necessary background for our submission (not our contributions), whereas Section 3 corresponds to our contributions, and as such we find that it would be confusing to merge these two sections to save space. Moreover, as mentioned above, the parametrization of the forward transition probability is described in great details in Section 3.4 already. What additional precisions would the reviewer like to be included?\n\n> *Moreover, I find strange that 5.3 is also so weakly developed in the main materials whereas it hints very important capacity of the proposed method, which might handle well interventional data.*\n> \n\nDue to limited space, we had to sacrifice the analysis of our experiments on real biological data for legibility, and defer them to Appendix D.5.2 (referenced in the main text in line 350). We only provided a description of the experiments and their relevance, in particular for interventional data. Upon acceptance and with an additional page in the camera ready version, we propose to move the complete Appendix D.5.2 (including Figure D.3) back into the main text.\n\n> *Why did not you try more expressive parameterisation of the CPD, such as with normalizing flows?*\n> \n\nIndividually, each CPD may be simple (Gaussian, although with non-linear functions), but as a whole the Bayesian network can model quite a complex joint distribution. Using normalizing flows at the level of individual CPDs would be interesting, but it is not clear if this would significantly increase the expressiveness of the joint distribution. However, our method is not limited to Gaussian distributions (e.g., our experiments in Sec 5.3 on real data), and it could be applied to CPDs parametrized by normalizing flows too. We leave this as future work.\n\n> Why did you only consider N=100 in 5.2?\n> \n\nWe used $N=100$ in our experiments in Section 5.2 in order to match the experimental setups used in a vast majority of the literature on Bayesian structure learning (Cundy et al., 2021; Lorch et al., 2021; Deleu et al., 2022; Wang et al., 2022). The motivation for a small dataset size is to highlight the advantage of the Bayesian perspective to capture uncertainty about the graphical structures.\n\n> *If the edge feature is a probability why do you used MSE, rather than a better metric for comparing distributions such as the KL?*\n> \n\nThe edge features correspond to the marginal probability of a specific edge being in the graph, and is defined precisely in Equation D.4. This is not a probability *distribution* over edges, and therefore we can’t compare distributions using the KL divergence.\n\nWhat we could do though is compare the (marginal) distribution over graphs returned by the GFlowNet ($P_{\\phi}^{\\top}(G)$, using the notation of Theorem 3.1, discarding $\\theta$ as described in Appendix D.3.2 line 794) and the exact marginal posterior $P(G\\mid \\mathcal{D})$ using the KL divergence (or any other divergence, such as the Jensen-Shannon divergence for symmetry): $\\mathrm{KL}(P_{\\phi}^{\\top}||P(\\cdot\\mid \\mathcal{D}))$. However, we wouldn’t be able to compute this metric for most of the baseline methods, since they only provide samples of the posterior approximation, and not the full distribution (we could instead use the cross-entropy metric, as we do for the component in $\\theta$).\n\n---\n\n(Cundy et al., 2021) Chris Cundy, et al. BCD Nets: Scalable Variational Approaches for Bayesian Causal Discovery. NeurIPS 2021.\n\n(Lorch et al., 2021) Lars Lorch, et al. DiBS: Differentiable Bayesian Structure Learning. NeurIPS 2021.\n\n(Deleu et al., 2022) Tristan Deleu, et al. Bayesian Structure Learning with Generative Flow Networks. UAI 2022.\n\n(Wang et al., 2022) Benjie Wang, et al. Tractable Uncertainty for Structure Learning. ICML 2022."}, {"review_id": "L8lG8r2hs0", "reviewer": "Reviewer_rtyV", "summary": "The present paper extends upon prior literature in making use of GFlowNets for structure learning.\nThe extension is concerned with finding an adequate approach to learning a complete Bayesian network, that is, the graph's structure and its parameters. Theory to support the consistency of the representation and empirical results comparing to state of the art.", "questions": "* What are the key assumptions (or 'sacrifices') one has to take for the real world example?", "limitations": "---", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "_Note to the authors: I apologize in advance for being unable to provide an extensive, constructive feedback for your manuscript due to capacity issues on my end. Therefore, the confidence is being kept low. Also, there is an OpenReview bug regarding \"First Time Reviewer,\" which sates 'Yes' although it is 'No.'_\n\n* GFlowNets are an interesting alternative approach to sampling for various tasks of interest\n* Learning a complete Graphical Structure is the optimum for structure learning\n* Wonderful presentation, both theory and visualizations/schematics\n* Self-enclosed work showcasing most/all relevant aspects to answering this research question", "weaknesses": "* 'Just' learning the parameters brings a completion to the structure learning approach for GFlowNets but IMHO poses a more marginal/incrimental improvement", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their review and their enthusiasm regarding our submission.\n\n> *'Just' learning the parameters brings a completion to the structure learning approach for GFlowNets but IMHO poses a more marginal/incrimental improvement*\n> \n\nWe invite you to read the general comment (**Advantage of joint posterior**).\n\n> *What are the key assumptions (or 'sacrifices') one has to take for the real world example?*\n> \n\nOne major consideration when working with real data is to make sure that our modeling assumptions (how we parametrize the CPDs) match the statistics of the data as closely as possible. In particular, gene expression data initially consists of counts that can take a wide range of values, and a standard preprocessing technique consists in transforming this data through a $\\log(1+x)$ function to make it so that the data “looks Gaussian”.\n\nHowever, since the raw data corresponds to counts, and a majority of those entries are 0 due to a lack of detection of a particular gene, this is also reflected after the transformation through $\\log (1+x)$. Therefore, the processed data is not Gaussian, but closer to a mixture of a Gaussian and a Dirac at 0 (for these undetected genes): this is well modeled by a zero-inflated Normal distribution, as described in Equation D.29. Modeling this with only a Gaussian CPD (as is standard in the structure learning literature) would create a mismatch with the nature of the data.\n\nBut even disregarding the heavy bias towards the value 0, since counts are positives, the function $\\log(1+x)$ transforms the non-zero counts into “Gaussian-like” data with a positive mean. However when working with linear-Gaussian CPDs, it is often standard to work with normalized data, with mean 0, to avoid having to introduce a bias term in the linear function. Here, we can’t easily standardize data due to the Dirac at 0. Therefore, it can be advantageous to consider a more expressive function for the CPD, such as a MLP, to even offer the option to capture complex biological mechanisms. JSP-GFN has the advantage to be able to work with such non-linear models (parametrized by a MLP) with non-standard distributions (zero-inflated Normal)."}, {"review_id": "o5yEczjy43", "reviewer": "Reviewer_5gCu", "summary": "This paper presents a new method for Bayesian structure learning from observational data, based on the framework of GFlowNets. \n\nGFlowNets have been used for Bayesian structure learning, but previous approaches view the distributions over the parameters and structures of a model modular and may learn them with separate methods. The paper introduces a joint method that learns them with a single GFlowNet. To do so, the paper extends the previous approach (Deleu et al., 2022) by introducing a new reward function, deriving the subtrajectory balance for the new reward function, and introducing a few adaptations to the parametrization of the forward transition probabilities of GFlowNet.\n\nThe paper conducts experiments on small synthetic graphs to evaluate how well the posterior of the model of the graph and parameters matches the ground-truth. Moreover, experiments with data generated by linear/non-linear Gaussian Bayesian networks are also conducted, with a focus on evaluating different models' performance of predicting held-out data (likelihood).", "questions": "1. Is it possible to add the results of BCD Nets and VBG in the non-linear case of Figure 3.b?\n\n2. Is it possible to provide complexity analysis between the proposed method and VBG?", "limitations": "N/a", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Structure learning is an important area in machine learning. GFlowNets is a new and emerging framework for structure learning. There are only a few studies in this direction, which adds to the novelty of the paper.\n\n2. The paper studies the problem of how to jointly learn parameters and structures in Bayesian structure learning with a single GFlowNet. The proposed solution to the problem is conceptually intuitive and technically sound, including the introduction of subtrajectory balance conditions and the parametrization of the forward transition probabilities.\n\n3. The experiments of the paper are quite comprehensive. Benchmark datasets and recent advances are included. Multiple metrics are compared.\n\n4. The quality of writing of the paper is good. The detailed derivations and settings are given. The paper is easy to follow.", "weaknesses": "1. Jointly learn parameters and structures in Bayesian structure learning with a single GFlowNet seems to be an interesting idea. But the benefit of doing so is less convincing in the paper. We have not seen theoretical or analytical study on why jointly learning them with one GFlowNet is useful. Empirically from the experiments of the paper, it is a bit hard to see that the joint learning scheme can help the quality of both the parameters and the structure. Specifically, the proposed method shows advantages on small graphs in Section 5.1, while its performance gain over others in Section 5.2 is a bit marginal in terms of the likelihood of held-out data. More importantly, one of the main purpose of structure learning is to reveal the true graph of the data, where the proposed method has not shown advantage over VBG (the method learns graphs with GFlowNet and parameters with variational inference, respectively) in Figure D.2. Therefore, it is a bit unclear whether people should learn them jointly.\n\n2. The experiments in Figure 3 are about predicting held-out data. In the linear Gaussian case, BCD Nets shows very competitive performance. Given that BCD Nets and VBG are linear models, one can still use them to model data generated from non-linear Gaussian models and predict the held-out data, although it might be hard to use them to compare SHD or AUROC. These two methods may serve as informative baselines for the non-linear case (Figure 3.b).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed review and comments on our submission.\n\n> *But the benefit of doing so is less convincing in the paper. We have not seen theoretical or analytical study on why jointly learning them with one GFlowNet is useful.*\n\nWe invite you to read the general comment (**Advantage of joint posterior**).\n\n> *Empirically from the experiments of the paper, it is a bit hard to see that the joint learning scheme can help the quality of both the parameters and the structure.*\n\nJust to be clear, all methods considered in our paper (both baseline methods, and our own method JSP-GFN) are jointly learning the parameters and structure. Our objective is not to show that finding the joint posterior over structures and parameters $P(G, \\theta\\mid D)$ is advantageous compared to the marginal posterior $P(G\\mid D)$. Our comparison with the exact posterior in Sec 5.1 on small graphs show a clear advantage of JSP-GFN over existing methods, and this is confirmed on the downstream performance on larger graphs in Sec 5.2; the difference is statistically significant against most methods on linear Gaussian models. On non-linear Gaussian, the downstream performance alone (Fig 3b) is not enough to conclude, although JSP-GFN seems to offer some benefit. That’s why we complemented this analysis with a comparison with the log-joint probability in Fig 3c, to show the quality of the approximation. Note that this analysis with correlation plots in Fig 3c, which is a reliable metric for Bayesian structure learning as it compares with the exact posterior, cannot be done on other methods since they only provide samples of the posterior approximation, and the log-probability $\\log Q(G, \\theta)$ cannot be evaluated.\n\n> *More importantly, one of the main purpose of structure learning is to reveal the true graph of the data, where the proposed method has not shown advantage over VBG […] in Figure D.2. Therefore, it is a bit unclear whether people should learn them jointly.*\n\nWe agree that the objective of structure learning is to recover the true graph generating the data, under identifiability assumptions. However, we respectfully disagree that this is the objective of **Bayesian** structure learning: the objective here is to provide a faithful approximation of the posterior distribution over graphs, not to find a single graph that generated the data. We argue in our submission that metrics such as SHD and AUROC (which are standard metrics in the structure learning literature) are not representative of the quality of the posterior approximation in the main text (lines 321-322), and expand on it in App D.4.3 (lines 886-893).\n\nOne argument we give in the paper is that a method that would always sample the true graph would fare very well under these two metrics (with an expected SHD of 0 and an AUROC of 1), even though this would most likely be a very poor approximation of the posterior distribution $P(G\\mid D)$. This would be a very good structure learning method, just not a good Bayesian structure learning one. That’s why we focused our analysis on metrics that do not compare with the ground truth graph (Figs 3 & D.3). As mentioned in our paper (lines 322-324), we included the expected SHD and AUROC in the Appendix only for completeness, since prior works on Bayesian structure learning still use these metrics.\n\nFor the comparison with VBG specifically, VBG performs significantly worse than JSP-GFN in terms of negative log-likelihood on held-out data in Fig 3a. But beyond performance, JSP-GFN offers three advantages over VBG (some of which are summarized in Table A.1):\n- In its original formulation, VBG can only be applied to linear-Gaussian models. The authors mentioned that it is possible to apply VBG to non-linear CPDs, but this has not been theoretically or empirically validated yet (we mention this in our comparison with other methods in App A lines 503-504). On the other hand, JSP-GFN can naturally handle non-linear CPDs.\n- JSP-GFN is conceptually simpler than VBG, where training simply consists in minimizing the objective in Eq 8 using stochastic gradient methods. On the other hand, training VBG consists in an inner-loop/outer-loop procedure, where the GFlowNet needs to the re-trained on a new reward function during the inner-loop, and the reward function is updated during the outer-loop. This introduces a number of hyperparameters, which might be hard to tune (e.g., number of training iterations for the GFlowNet during the inner-loop).\n- It is impossible to use mini-batch training of the GFlowNet in VBG, due to the same limitations DAG-GFlowNet suffers from. This can limit the scope of application of VBG to cases where the size of the dataset is small.\n\n> *Is it possible to add the results of BCD Nets and VBG in the non-linear case of Figure 3.b?*\n\nThank you for the suggestion. We only used methods that could be applied to non-linear models, to match the data generation process, following (Lorch et al., 2021).\n\nHowever, we followed your suggestion and evaluated both BCD Nets and VBG (methods based on linear Gaussian models) on our non-linear datasets as well. The results in terms of negative log-likelihood on held-out data are as follows (we’re reporting median & 25th/75th percentile here):\n\n| Algorithm | NLL (held-out) |\n| --- | --- |\n| M-MC3 | 2911.99 [1920.47, 3000.35] |\n| G-MC3 | 4224.05 [2982.29, 5385.90] |\n| DiBS | 3181.80 [2376.39, 4018.54] |\n| BCD Nets* | 2714.62 [1832.27, 2998.93] |\n| VBG* | 3030.28 [2311.04, 4805.90] |\n| JSP-GFN (ours) | 3235.74 [2327.23, 3429.66] |\n\nThe rows with * denote the new results, and all other rows correspond to the results in Fig 3b. It is interesting to see that methods with linear Gaussian models perform well even on these datasets generated with non-linear functions, although the advantage is not statistically significant.\n\n---\n\n(Lorch et al., 2021) Lars Lorch, et al. DiBS: Differentiable Bayesian Structure Learning. NeurIPS 2021."}, {"review_id": "CPYMi7hpnL", "reviewer": "Reviewer_c3AM", "summary": "The authors utilize a generative flow network to jointly model the structure of the DAQ and its parameters.", "questions": "None", "limitations": "The authors did not discuss this which I think they should.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors utilize GFlowNets to jointly model the structure of the DAQ and its parameters. The experiments on simulated data and real-world data are promising.", "weaknesses": "I am not an expert on this area and I did not bid this paper. So I will listen to other reviewers' opinion on weakness. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for taking the time to read our submission, despite it not being in their area of expertise.\n\nThe limitations of our work, along with broader impacts, are discussed in Appendix B."}, {"review_id": "wysgfnyf80", "reviewer": "Reviewer_jHMM", "summary": "The authors present a method for learning a posterior distribution over Bayesian networks using an extended version of GFlowNets that can jointly construct the underlying DAG and associated parameters in a two-stage process. They use an objective based on a generalization of detailed balance to train their flow network and provide a theoretical justification that optimizing this generalized objective leads to the desired distribution.\nThe method is evaluated on synthetic problems and gene regulatory networks, outperforming recent methods based on MCMC.\n", "questions": "- Fig. D.3c: is this the negative LL (lower=better) or just LL? Depending on the answer the interpretation of the results would change. \n- What is the relationship between accuracy vs. number of training data?\n- How does the method scale with the number of nodes? Can the authors provide some intuition about the size of Bayesian networks one might be able to learn with this method, also in relation to the number of training data?\n- Why is the additional use of a graph neural network necessary compared to the architecture used in DAG-GFlowNet? \n- The authors discuss how $P_{\\phi}(\\theta_i|G,stop)$ can be parameterized by a Normal/diagonal-Normal distribution, or a 2-layer MLP for each $\\theta_i$. How do the authors ensure that correlations between $\\theta$s of different $G_i$ can be learned if different MLPs are used for each $\\theta_i$? \n", "limitations": "The authors only touch on some limitations and a more detailed discussion would strengthen the paper. \n\n\n", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The problem of learning Bayesian networks is relevant to the community. The authors provide a deep theoretical framework for the presented problem and derive it thoroughly. The strength of the method is that it provides not just a point estimate, but a full posterior distribution over Bayesian networks and parameters in an amortized fashion.", "weaknesses": "- While the theoretical framework is explained in much detail, the actual implementation is quite opaque to me. The authors could spend more time in explaining the architecture and implementation details. \n- “You are strongly encouraged to submit code that produces this result.” [NeurIPS guidelines].\n- A simple educational example that is dissected in detail could help to better understand the framework. \n- The experimental results are restricted to some performance metrics, but especially in the biological datasets, it would be interesting to dig deeper. Could JSP-GFN discover known gene regulatory networks? How do posterior predictives look compared to the real data? What role does interventional vs. observational data play?\n\n- The experiments provide little intuition about the regime on which this method can be expected to work in terms of the number of training data, the number of nodes, and the number of parameters per node.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed comments. Due to limited space, we focus our responses to the Questions, and defer the responses to other points raised in the review in a separate comment below.\n\n> *Fig. D.3c: is this the negative LL (lower=better) or just LL? Depending on the answer the interpretation of the results would change.*\n\nThis is indeed a typo, and we would like to thank you for pointing this out. The results reported in Figure D.3c are the log-likelihoods. We will update the values in the table to report the negative log-likelihoods for consistency with the rest of the paper.\n\n> *What is the relationship between accuracy vs. number of training data?*\n\nIn terms of accuracy of the posterior approximation, we expect JSP-GFN (or any other Bayesian structure learning method) to provide a more accurate approximation of the posterior when the dataset is small. When the dataset is larger, there are some challenges that arise in JSP-GFN, which we detail below.\n\n> *How does the method scale with the number of nodes? Can the authors provide some intuition about the size of Bayesian networks one might be able to learn with this method, also in relation to the number of training data?*\n\nOne limitation for scaling in terms of number of nodes is the size of the neural network that is used to parametrize the forward transition probability: larger graphs mean larger inputs and therefore it is more expensive from a computational perspective. That being said, using a graph network instead of a transformer as in DAG-GFlowNet offers an advantage here (see below).\n\nAnother limitation when it comes to scaling, which is beyond the scope of this paper, is the problem of effective exploration when the state space becomes so large (recall that it is super exponential in $d$). For the experiments we ran in this paper (up to $d=61$ nodes), we found that a naive exploration strategy as described in the general comment (**Practical implementation**) was sufficient. However, we envision that this could become a bottleneck for much larger graphs, where the state space becomes too large. Since this is a challenge that arises from GFlowNets themselves, this is an issue that is better studied in isolation and not in the narrow context of Bayesian structure learning (e.g., inspiration can be taken from the RL literature).\n\nBut from our experiments, we can confidently attest that JSP-GFN is an effective method to approximate the posterior distribution for up to $d=50$ nodes (which is typically the scale attained by prior works), and maybe pushing it to $d=100$ nodes.\n\nIn terms of the size of the dataset, JSP-GFN shares some limitations that were discussed in DAG-GFlowNet: when the dataset becomes large, the posterior distribution becomes more concentrated around a few graphs (and parameters). Therefore, the neural network that parametrizes the forward transition probability must be able to handle large variations in outputs, which can be challenging. This comes from the fact that Equation 8 corresponds (roughly speaking) to a regression problem where the target is the delta score $\\log R(G', \\theta') - \\log R(G, \\theta)$: if the posterior is very “peaky”, then the delta-score may take values with large magnitude (positive or negative, it scales roughly with the size of $\\mathcal{D}$). One standard practice in machine learning would be to normalize the target of the regression problem, but this is impossible here since this would change the nature of the distribution we are trying to approximate. One solution suggested by the authors of DAG-GFlowNet, which could be applied to JSP-GFN too, would be to take inspiration of simulated annealing and introduce a temperature parameter in the posterior distribution that will decrease over training.\n\nIn practice, JSP-GFN can handle datasets up to a few thousands examples (in our experiments, we had up to $N=4200$ in the flow cytometry data experiment).\n\n> *Why is the additional use of a graph neural network necessary compared to the architecture used in DAG-GFlowNet?*\n\nThe graph network is not necessary, and we could have used a similar transformer architecture as the authors of DAG-GFlowNet. Our choice was motivated by the fact that the graph network was a much simpler architecture, specifically created for graph inputs, without compromising in terms of performance. The transformer used in DAG-GFlowNet is bound to have a fixed input size of $d^2$, regardless of the number of edges in the input graph; in contrast, the input graphs in the graph networks will have at most $d(d-1)/2$ edges, and will often have much fewer edges in practice. This allows us to better scale in terms of the size of graphs, as opposed to using a transformer.\n\n> *The authors discuss how $P_{\\phi}(\\theta_{i} \\mid G, stop)$ can be parameterized by a Normal/diagonal-Normal distribution, or a 2-layer MLP for each $\\theta_{i}$. How do the authors ensure that correlations between $\\theta$s of different $G_{i}$ can be learned if different MLPs are used for each $\\theta_{i}$?*\n\nWe (implicitly) assume in this paper that the priors over graphs $P(G)$ and over parameters $P(\\theta\\mid G)$ are modular, which is a standard assumption in the structure learning literature, meaning that\n\n$$P(G) = \\prod_{i=1}^{d}P(G_{i}) \\qquad \\mathrm{and} \\qquad P(\\theta\\mid G) = \\prod_{i=1}^{d}P(\\theta_{i}\\mid G_{i})$$\n\nIn practice, we used a uniform prior over graphs, and a unit Gaussian prior over parameters (e.g., lines 783-787). Under modularity assumption, it’s easy to see that the posterior over parameters also decomposes as $P(\\theta\\mid G, \\mathcal{D}) = \\prod_{i}P(\\theta_{i}\\mid G_{i}, \\mathcal{D})$. This is in fact at the core of our proof in Appendix D.6.1. Since $P_{\\phi}(\\theta_{i}\\mid G, stop)$ effectively approximates the posterior over parameters (as mentioned in lines 231-232), there is no need to model correlations between $\\theta$s, since they do not exist in the true posterior."}, {"review_id": "hNNS0YDrf2", "reviewer": "Reviewer_9gh3", "summary": "In a recent line of research, Generative Flow Networks (GNFs) are used for structure learning by drawing DAGs (that represent BN structures) from an implicit DAG distribution that fits the data (Deleu et al 2022). The contribution of the present paper is extending this line of research by drawing joint DAG and the parameters (of the factors associated with the DAG's nodes).\nTo be more precise, firstly they construct a graph progressively, by adding new links between nodes (as in Deleu et al 2022 paper), and then the parameters are chosen to form a pair of (structure, parameter).", "questions": "1. I would appreciate it if the authors could comment on the mentioned measure theoretic problem. \n\n2. In lines 176-179, there is a confusing statement that I interpret as follows: \"Although there is no guarantee that equation (6) guarantees the convergence but if we assume that (6) holds for all pairs (which I suppose, means if we assume that the dimensionality is fixed which is not (?)) then the convergence is guaranteed.\" \nCould you please clarify these lines?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The problem that this paper addresses i.e. \"structure learning\" is an important task. \n2. The paper is overall well written and the literature review is thorough. \n3. The reported results are good. \n4. This is an interesting paper but unfortunately, I believe its theoretical flaws are serious and have to be addressed before publishing the work.", "weaknesses": "1. The core theory behind this work, i.e. equation (6), seems fundamentally incorrect and has measure theoretic problems:\nDetailed balance is sufficient for \"fixed point condition\" --which in this case means, the convergence of the probability of the terminal structure/parameters  (G, theta) to a posterior probability that is proportional to its reward R(G, theta)-- only if the dimensionality is fixed. When the sub-spaces have different dimensions, measure-theoretic problems rise due to mapping a single point (in a low-dimensional subspace) to infinite points (in a high-dimensional subspace). \nThe authors correctly mention (in line 242) that MCMC encounters trans-dimensional move problems. \nHowever, this (as discussed in Green's Reversible Jump MCMC paper) is due to the detailed balance condition which is the foundation of the present work as well. As such, as far as I can say, the same issue should be addressed here too otherwise the correct convergence is not guaranteed. \n \n2. The code is not provided. As such, the reported results are not reproducable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their comments; we believe that checking the soundness of theoretical results when available is an integral part of the reviewing process, and we appreciate it that you took the time to do so. **However, we want to draw the attention of all reviewers and the Area Chairs to the fact that there is no theoretical flaw in our submission, including in the context raised by Reviewer 9gh3.**\n\nThe complete measure-theoretical treatment of continuous GFlowNets can be found in (Lahlou et al., 2023), but we provide here some additional details in the specific context of this submission. For some DAG $G$, let $\\Theta_{G}$ denote the measurable space of parameters $\\theta$ associated with $G$ (with its Borel $\\sigma$-algebra, which we omit here for clarity). **Note that the spaces $\\Theta_{G}$ do not need to be the same for different graphs, and they may have different dimensionalities.** From this perspective, the reward $R(G, d\\theta)$ and the forward transition probability $P_{\\phi}(d\\theta\\mid G)$ are now **measures** over $\\Theta_{G}$ (using infinitesimal notations).\n\nFor any transition $G \\rightarrow G'$ in the GFlowNet (i.e., $G'$ is the result of adding a single edge to $G$), Eq 6 of the submission represents **an equality between measures over the product space $\\Theta_{G}\\times \\Theta_{G'}$**. Putting it differently, for any bounded measurable function $h: \\Theta_{G}\\times \\Theta_{G'}\\rightarrow \\mathbb{R}$, the condition reads:\n\n$$\n\\iint_{\\Theta_{G}\\times \\Theta_{G'}}h(\\theta, \\theta')R(G', d\\theta')P_{B}(G\\mid G')P_{\\phi}(d\\theta\\mid G) = \\iint_{\\Theta_{G}\\times \\Theta_{G'}}h(\\theta, \\theta')R(G, d\\theta)P_{\\phi}(G'\\mid G)P_{\\phi}(d\\theta'\\mid G')\n$$\n\nWe would like to emphasize once again that this condition is valid even when $\\Theta_{G}$ and $\\Theta_{G'}$ do not have the same dimensionality, since this is an equality between measures over the product space (not either individual space $\\Theta$). In (Lahlou et al., 2023), it was shown how similar conditions can be translated into constraints in terms of probability **densities**.\n\nOne mild (implicit) assumption we make in this submission is that for any DAG $G$, there exists a (common base) measure $\\mu_{G}$ over $\\Theta_{G}$ such that $R(G, d\\theta) \\ll \\mu_{G}(d\\theta)$ and $P_{\\phi}(d\\theta\\mid G) \\ll \\mu_{G}(d\\theta)$ are both absolutely continuous wrt. $\\mu_{G}$. This assumption is indeed mild because as mentioned in Section 3.4 of our submission, $P_{\\phi}(d\\theta\\mid G)$ approximates the posterior distribution $P(d\\theta\\mid G, \\mathcal{D})$ (up to a multiplicative constant), and $R(G, d\\theta) = P(\\mathcal{D}, G, d\\theta)$. Using the notations $R(G, \\theta)$ and $P_{\\phi}(\\theta\\mid G)$ to denote their Radon-Nikodym derivatives wrt. $\\mu_{G}$, Eq 6 corresponds to the equality between measures above, written as an equality in terms of their Radon-Nikodym derivatives:\n\n$$\nR(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')\n$$\n\nAll the results in the submission hold if we explicitly work with measures instead of their derivatives, regardless of the dimensionalities of $\\Theta$. In particular Thm 3.1 of our submission (satisfying these conditions implies that the GFlowNet induces a distribution $\\propto R$) still holds true. We provide a detailed proof from a measure theoretical perspective in a separate comment below for completeness.\n\nUnlike the detailed balance condition in MCMC, Eq 6 (which is *not* the detailed balance condition for Markov chains) does not define a fixed-point condition for an invariant distribution to match some target distribution (here, the joint posterior). The objective of the GFlowNet is to find a forward transition probability $P_{\\phi}$ such that its marginal distribution over complete states (called the terminating state probability in line 186) matches the joint posterior, via satisfying Eq 6. $P_{\\phi}$ itself is not the invariant distribution that would arise from a fixed-point condition like the detailed balance.\n\nThe key difference with RJ-MCMC is that in the GFlowNet presented here, we never have to make trans-dimensional moves: once the first phase ends (Sec 3.1), we commit to a specific graph $G$, and therefore we can sample from a distribution over $\\Theta_{G}$; then once this is done (end of the second phase), we can start over from the initial state $(G_{0}, \\cdot)$. In other words, we never make any move from one space $\\Theta_{G}$ (or more precisely $(\\{G\\}\\times \\Theta_{G})$) to another $\\Theta_{G'}$ with possibly a different dimensionality, unlike in RJ-MCMC. That’s why there is no need for an invertible transformation across dimensions in our Eq 6.\n\n> *In lines 176-179, there is a confusing statement […] Could you please clarify these lines?*\n> \n\nWe see where this confusion comes from, and we apologize for it. “The SubTB conditions” in line 176 refers to generic subtrajectory balance conditions, as described in App C.2, and not the ones specifically for our setting in Eq 6. What we meant was that satisfying the subtrajectory balance conditions of the form in Eq C.3 do not guarantee (in general) that the GFlowNet induces a distribution proportional to $R$. However, satisfying the SubTB conditions for the specific subtrajectories of the form $(G, \\theta) \\leftarrow (G, \\cdot) \\rightarrow (G', \\cdot) \\rightarrow (G', \\theta')$ considered in Sec 3.2, corresponding to the conditions in Eq 6, do guarantee this result (Thm 3.1). We propose to replace lines 176-179 by:\n\n*Although there is no guarantee in general that satisfying **the SubTB conditions in (C.3) for arbitrary subtrajectories** would yield…*\n\nWe hope that this clarified the misunderstanding about our theoretical results, and we strongly encourage the reviewer to update their score since this was the only major source of confusion.\n\n---\n\n(Lahlou et al., 2023) Salem Lahlou, et al. A Theory of Continuous Generative Flow Networks. ICML 2023."}]

Author rebuttal:
We would like to thank all the reviewers for taking the time to review our submission and for their valuable comments. While we are addressing specific questions in individual rebuttal comments, we would like to add a few points here which will be of interest for all reviewers and Area Chairs.

If any comment in the reviews has not been addressed in our individual rebuttal responses, this is due to the limited space available. We will provide these responses as soon as the discussion period starts.

### **Theory**

We would like to draw the attention of all reviewers and the Area Chairs here that **there is no theoretical flaw in our submission, including in the context raised by Reviewer 9gh3.** 

### **Advantage of joint posterior**

We mentioned in the introduction of our submission (lines 37-46) why it is advantageous to consider the joint posterior $P(G, \theta\mid \mathcal{D})$ as opposed to the marginal posterior $P(G\mid \mathcal{D})$. However, as pointed out by reviewers rtyV & 5gCu, this could have been better highlighted in the paper.

To paraphrase what we said in lines 37-46, the main advantage of looking at $P(G, \theta\mid \mathcal{D})$ is that this allows us to consider Bayesian networks with more complex CPDs (e.g., non-linear CPDs, parametrized by a neural network). This has a significant impact in practice, since marginal posteriors are generally limited to linear Gaussian models (some exceptions exist, as mentioned in the paper in lines 32-36), and therefore approximating the joint posterior significantly increases the range of applications of Bayesian structure learning methods.

For the specific advantages of JSP-GFN over existing methods, we provide a table of comparison in Table A.1, where we compared our method to other methods based on variational inference and GFlowNets, on 7 different points: (1) if the method was approximating the joint posterior $P(G, \theta\mid \mathcal{D})$, (2) if the method could be applied to Bayesian networks with non-linear CPDs, (3) if the support of the distribution was guaranteed to be limited to the space of DAGs, (4) if the method could be applied to discrete data, (5) if the maximum number of parents for each node in the graph can be explicitly specified (common assumption in non-Bayesian structure learning), (6) if the method provides a method to sample from $P(G, \theta\mid \mathcal{D})$, and (7) if the method could learn the variational approximation based on mini-batches of data.

To the best of our knowledge, JSP-GFN is the only variational method that can approximate the joint distribution $P(G, \theta\mid \mathcal{D})$ for Bayesian networks with non-linear CPDs (including discrete random variables), while guaranteeing that the support of the distribution is limited to DAGs, which is essential for Bayesian networks. As such, it constitutes a major development in Bayesian structure learning.

### **Practical implementation**

While we described extensively all the necessary components of our method JSP-GFN in our submission, we would like to thank reviewers MmTz & jHMM for pointing out that the description of the practical implementation could be improved.

The way the policy (more precisely, the forward transition probability $P_{\phi}$) is parametrized is precisely described in Section 3.4; only the exact architectures of the neural networks used are not detailed (number of layers, number of hidden units), but we believe that they do not contribute to the overall description. The learning objective is described in Section 3.3, and explicitly formulated in Equation 8. Moreover, the way the behavior policy $\pi$ (used in the objective in Equation 8) is again described in details in Appendix D.1 (referenced in line 193 in the main text).

Training the GFlowNet then simply consists of standard stochastic optimization:

1. Draw samples $(G, \theta, G', \theta') \sim \pi$ according to the behavior policy $\pi$ as described in Appendix D.1, to obtain a Monte Carlo estimate of Equation 8;
2. Compute the different forward transition probabilities $P_{\phi}(G'\mid G)$, $P_{\phi}(\theta\mid G)$ and $P_{\phi}(\theta'\mid G')$ with the parametrization detailed in Section 3.4;
3. Evaluate the objective in Equation 8 (Section 3.3), and optimize it using stochastic gradient methods.

We can add this as a pseudo-code in the Appendix of the paper (the space in the main text unfortunately does not allow us to put it there). For example in Appendix D, in a new section before Appendix D.3.

The last component that is not described in out paper is how the transitions $G\rightarrow G'$ are obtained before putting them in the replay buffer, as described in Appendix D.1 (for the policy $\pi$). This is using a simple exploration strategy using $\epsilon$-sampling: we transition in the GFlowNet following $P_{\phi}(G'\mid G)$ with probability $\epsilon$, and select a new edge to add (among the valid edges) uniformly at random with probability $1-\epsilon$. We will add this as part of the pseudo-code above.

### **Code release**

While the source-code was not included as part of the supplementary material, as mentioned by reviewers jHMM & 9gh3, we have since released the code publicly. A link to an anonymized version of the code is available [here](https://anonymous.4open.science/r/jsp-gfn-AD5D/), and a link to the repository will be added to the final version of the paper upon acceptance. Our code is largely based on the code released by (Deleu et al., 2022) available [here](https://github.com/tristandeleu/jax-dag-gflownet).

---

(Deleu et al., 2022) Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, Yoshua Bengio. Bayesian Structure Learning with Generative Flow Networks. UAI 2022.

Meta-review: This paper presents a method for Bayesian structure learning from observational data, based on the GFlowNets framework. It shows that sub-trajectory balance conditions suffice to ensure that the GFlowNets induce a distribution that is well proportional to the reward function associated with a final sample. 

All the reviewers recognize the importance of the problem this paper is addressing and that the proposed approach is sound. Although the method may seem incremental, it is important to acknowledge that adjusting the theory in GFlowNets (re: balancing conditions) and the promising results are substantial contributions. Several issues were raised regarding (i)  clarity of the presentation and the background necessary on GFlowNets; (ii)  details of the implementation; (iii) reproducibility and code and (iv) a potential theoretical flaw (with measure-theoretic and trans-dimensional issues). All these have been resolved by the authors’ rebuttal.

Because of the low confidence in the original set of assigned reviewers at the initial stages of the reviewing process, I recruited more reviewers and we ended up with 6 reviews. Overall, I am confident this paper is worthy of presentation at NeurIPS.
Decision: Accept (poster)