Title: Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network
Venue: NeurIPS 2023
Accepted: yes
Keywords: bayesian network

bayesian

structure learning

causal discovery

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

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.

1 Introduction

As a compact representation for complex probabilistic models, Bayesian Networks are a framework of choice in many fields, such as computational biology (Friedman et al., 2000; Sachs et al., 2005) and medical diagnosis (Lauritzen and Spiegelhalter, 1988). When the directed acyclic graph (DAG) structure of the Bayesian Network-which specifies the possible conditional dependences among the observed variables-is known, it can be used to perform probabilistic inference for queries of interest with a variety of exact or approximate methods (Koller and Friedman, 2009). However, if this graphical structure is unknown, one may want to infer it based on a dataset of observations D .

In addition to being a challenging problem due to the super-exponentially large search space, learning a single DAG structure from data may also lead to confident but incorrect predictions (Madigan et al., 1994), especially in cases where the evidence is limited. In order to avoid misspecifying the model, it is therefore essential to quantify the epistemic uncertainty about the structure of the Bayesian Network. This can be addressed by taking a Bayesian perspective on structure learning and inferring the posterior distribution P ( G | 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),

discrete models with Dirichlet prior (Heckerman et al., 1995), or non-linear models parametrized with a Gaussian Process (von K¨ ugelgen et al., 2019).

While there exists a vast literature on Bayesian structure learning to approximate the marginal posterior distribution, inferring the joint posterior P ( G,θ | D ) over both the DAG structure G of the Bayesian Network and the parameters θ of its conditional probability distributions-the probability of each variable given its parents-has received comparatively little attention. The main difficulty arises from the mixed sample space of the joint posterior distribution, with both discrete components (the graph G ) and continuous components (the parameters θ ), where the dimensionality of the latter may even depend on G . However, modeling the posterior distribution over θ 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,θ ) 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 et al., 2021, 2023) proved to be an effective method for Bayesian structure learning. In Deleu et al. (2022), the problem of generating a sample DAG from the marginal posterior P ( G | D ) was treated as a sequential decision process, where edges are added one at a time, starting from the empty graph over d variables, following a learned transition probability. Nishikawa-Toomey et al. (2023) have also proposed to use a GFlowNet to infer the joint posterior P ( G,θ | D ) ; however, they used it in conjunction with Variational Bayes to update the distribution over θ , getting around the difficulty of modeling a continuous distribution with a GFlowNet.

In this paper, we propose to infer the joint posterior over graphical structures G and parameters of the conditional probability distributions θ of a Bayesian Network using a single GFlowNet called JSP-GFN (for Joint Structure and Parameters GFlowNet ), leveraging recent advances extending GFlowNets to continuous sample spaces (Lahlou et al., 2023), and expands the scope of applications for Bayesian structure learning with GFlowNets. The generation of a sample ( G,θ ) 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 is completely known. To enable efficient learning of the sampling distribution, we introduce new conditions closely related to the ones derived in Deleu et al. (2022), based on the subtrajectory balance conditions (Malkin et al., 2022), and show that they guarantee that the GFlowNet does represent P ( G,θ | 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 = ( V, E ) , where V is a set of nodes, and E ⊆ V × V is a set of (directed) edges. For a node X ∈ V , we denote by Pa G ( X ) the set of parents of X in G , and Ch G ( X ) the set of its children. For two nodes X,Y ∈ V , X → Y represents a directed edge ( X,Y ) ∈ E (denoted X → Y ∈ 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 GFlowNets (see Section 2.2), an undirected path in a directed graph G = ( V , E ) between two states s 0 , s n ∈ V is a sequence of vertices ( s 0 , s 1 , . . . , s n ) where either ( s i , s i +1 ) ∈ E or ( s i +1 , s i ) ∈ E (i.e., following the edges of the graph, regardless of their orientations).

2.1 Bayesian structure learning

A Bayesian Network is a probabilistic model over d random variables { X 1 , . . . , X d } , whose joint distribution factorizes according to a directed acyclic graph (DAG) G as

where θ = { θ 1 , . . . , θ d } represents the parameters of the conditional probability distributions (CPDs) involved in this factorization. When the structure G is known, a Bayesian Network offers a representation of the joint distribution that may be convenient for probabilistic inference (Koller and

P

φ

(

θ

4

|

G

4

)

θ

4

R

(

G

4

, θ

4

)

A

B

C

G

0

A

B

C

G

1

A

B

C

G

2

A

B

C

G

4

A

B

C

G

3

A

B

C

G

5

A

B

C

G

6

A

B

C

G

7

Add

B

→

A

P

φ

(

G

1

|

G

0

)

Add

B

→

C

P

φ

(

G

3

|

G

1

)

. . .

. . .

. . .

Friedman, 2009), as well as learning its parameters θ using a dataset D = { x (1) , . . . , x ( N ) } , where each x ( j ) represents an observation of the d random variables { X 1 , . . . , X d } .

However, when the structure of the Bayesian Network is unknown, we can also learn the DAG G from data (Spirtes et al., 2000; Chickering, 2002). Using a Bayesian perspective, we may want to either model the (marginal) posterior distribution P ( G | D ) over only the DAG structures, or the joint posterior P ( G,θ | D ) over both the structure G , as well as the parameters of the CPDs θ .

2.2 Generative Flow Networks

A Generative Flow Network (GFlowNet; Bengio et al., 2021, 2023) is a generative model over a structured sample space X . The structure of the GFlowNet is described by a DAG G whose vertex set is the state space S , where X ⊆ S . This should not be confused with the DAG in a Bayesian Network: in fact, each state in G could itself represent the structure of a Bayesian Network (Deleu et al., 2022). A sample x ∈ X is constructed sequentially by following the edges of G , starting at a special initial state s 0 , until we reach the state x . We define a special terminal state s f , a transition to which indicates the end of the sequential process. The states in X are those for which there is a directed edge x → s f ; they are called complete states 1 and correspond to valid samples of the distribution induced by the GFlowNet. A path s 0 ⇝ s f in G is called a complete trajectory. An example of the structure of a GFlowNet is given in Section 3.1 and Figure 1.

Every complete state x ∈ X is associated with a reward R ( x ) ≥ 0 , indicating the unnormalized probability of x . We use the convention R ( s ) = 0 for any state s ∈ S\X , since they do not correspond to valid samples of the distribution. Bengio et al. (2021) showed that if there exists a function F ϕ ( s → s ′ ) ≥ 0 defined over the edges of G , called a flow , that satisfies the following flow-matching conditions for all the non-terminal states s ′ ∈ S , then the GFlowNet induces a distribution over complete states proportional to the reward. More precisely, starting from the initial state s 0 , if we sample a complete trajectory ( s 0 , s 1 , . . . , s T -1 , x, s f ) following the forward transition probability, defined as

with the conventions s T = x and s T +1 = s f , then the marginal probability that a trajectory sampled following P terminates in x ∈ X is proportional to R ( x ) . The flow function F ϕ ( s → s ′ ) may be parametrized by a neural network and optimized to minimize the error in (2), yielding a transition model that can be used to approximately sample from the distribution on X proportional to R .

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-GFlowNet , operates on the state-space of DAGs, where each graph is constructed sequentially by adding one edge at a time, starting from the empty graph with d nodes, while enforcing the acyclicity constraint at every step of the generation (i.e., an edge is not added if it would introduce a cycle). Its structure is illustrated at the top of Figure 1, where it forms the part of the graph shown in blue.

Instead of working with flows F ϕ ( G → G ′ ) , as in Section 2.2, DAG-GFlowNet directly learns the forward transition probability P ϕ ( G ′ | G ) that satisfies the following alternative detailed balance conditions for any transition G → G ′ (i.e., G ′ is the result of adding a single edge to G ):

where P B ( G | 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 ∝ R ( G ) . Therefore, to approximate the posterior distribution P ( G | D ) , they used R ( G ) = P ( D | G ) P ( G ) as 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-continuous hybrids. Building on top of DAG-GFlowNet, we propose here to approximate the joint posterior P ( G,θ | D ) over both the structure of the Bayesian Network G , but also the parameters of its conditional probability distributions θ . Unlike in VBG though (Nishikawa-Toomey et al., 2023), we use a single GFlowNet to approximate this joint posterior. We call this model JSP-GFN , for Joint Structure and Parameters Bayesian inference with a GFlowNet.

3.1 Structure of the GFlowNet

Unlike in DAG-GFlowNet, where we model a distribution only over DAGs, here we need to define a GFlowNet whose complete states are pairs ( G,θ ) , where G is a DAG and θ is a set of (continuousvalued) parameters whose dimension may depend on G . Complete states are obtained through two phases (Figure 1): the DAG G is first constructed one edge at a time, following Deleu et al. (2022), and then the corresponding parameters θ are generated, conditioned on G . We denote by ( G, · ) states where the DAG G has no parameters θ 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,θ ) ∈ X , whereas ( G, · ) ∈ S\X .

Since we want this GFlowNet to approximate the joint posterior P ( G,θ | D ) ∝ P ( D , θ, G ) , it is natural to define the reward function of a complete state ( G,θ ) as

Starting at the empty graph ( G 0 , · ) , the DAG is constructed one edge at a time during the first phase, following the forward transition probabilities P ϕ ( G ′ | G ) . This first phase ends when a special 'stop' action is selected with P ϕ , 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 θ conditioned on G , following the forward transition probabilities P ϕ ( θ | G ) . 2 All the complete states ( G,θ ) , for a fixed graph G and any set of parameters θ , can be seen as forming an (infinitely wide) tree rooted at ( G, · ) .

where the likelihood model P ( D | θ, G ) may be arbitrary (e.g., a neural network), and decomposes according to (1), P ( θ | 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).

3.2 Subtrajectory Balance conditions

To obtain a generative process that samples pairs of ( G,θ ) proportionally to the reward, the GFlowNet needs to satisfy some conditions such as (4). However, we saw in Section 2.3 that satisfying this particular formulation of the detailed balance conditions in (4) yields a distribution ∝ R ( · ) only if 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,θ ) ← ( G, · ) → ( G ′ , · ) → ( G ′ , θ ′ ) in the GFlowNet (see Figure C.2), where G ′ is therefore the result of adding a single edge to G . Since both ends of these undirected paths of length 3 are complete states, we show in Appendix C.3.1 that the SubTB conditions corresponding to undirected paths of this form can be written as

Note that the SubTB condition above is very similar to the detailed balance condition in (4) used in DAG-GFlowNet, where the probability of terminating P ϕ ( s f | G ) has been replaced by P ϕ ( θ | G ) , in addition to the reward now depending on both G and θ . Moreover, while it does not seem to appear in (6), the terminal state s f of the GFlowNet is still present implicitly, since we are forced to terminate once we have reached a complete state ( G,θ ) , and therefore P ϕ ( s f | G,θ ) = 1 (see App. C.3.1).

Although there is no guarantee in general that satisfying the SubTB conditions would yield a distribution proportional to the reward, unlike with the detailed balance conditions (Bengio et al., 2023), the following theorem shows that the GFlowNet does induce a distribution ∝ R ( G,θ ) if the SubTB conditions in (6) are satisfied for all pairs ( G,θ ) and ( G ′ , θ ′ ) .

Theorem 3.1. If the SubTB conditions in (6) are satisfied for all undirected paths of length 3 between any ( G,θ ) and ( G ′ , θ ′ ) of the form ( G,θ ) ← ( G, · ) → ( G ′ , · ) → ( G ′ , θ ′ ) , then we have

where P ϕ ( G | G 0 ) is the marginal probability of reaching G from the initial state G 0 with any (complete) trajectory τ = ( G 0 , G 1 , . . . , G T -1 , G ) :

using the conventions G T = G , and P ϕ ( G 0 | G 0 ) = 1 .

The proof of this theorem is available in Appendix C.3.3. The marginal distribution P ⊤ ϕ ( G,θ ) is also called the terminating 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,θ ) and ( G ′ , θ ′ ) 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 π [∆ 2 ( ϕ )] , where the residuals ∆( ϕ ) depend on the conditions in (6), and π is an arbitrary sampling distribution of complete states ( G,θ ) and ( G ′ , θ ′ ) , with full support; see Malkin et al. (2023) for a discussion of the effect of π on training GFlowNets, and Appendix D.1 for further details in the context of JSP-GFN.

In addition to the SubTB conditions for undirected paths of length 3 given in Section 3.2, we can also derive similar SubTB conditions for other undirected paths, for example those of the form ( G,θ ) ← ( G, · ) → ( G, ˜ θ ) , where θ and ˜ θ are two possible sets of parameters associated with the same graph G (see also Figure C.2). When the reward function R ( G,θ ) is differentiable wrt. θ , 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

This condition is equivalent to the notion of score matching (Hyv¨ arinen, 2005) to model unnormalized distributions, since R ( G,θ ) here corresponds to the unnormalized posterior distribution P ( θ | G, D ) , where the normalization is over θ . Wealso 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 θ and θ ′ (e.g., backpropagation with the reparametrization trick) in the objective where ⊥ denotes the 'stop-gradient' operation. This is aligned with the recommendations of Lahlou et al. (2023) to avoid backpropagation through the reward in continuous GFlowNets. The pseudo-code for training JSP-GFN is available in Algorithm 1.

3.4 Parametrization of the forward transition probabilities

In Section 3.1, we saw that the process of generating ( G,θ ) follows two phases: first we construct G one edge at a time, until we sample a specific 'stop' action, at which point we sample θ , conditioned on G . All these actions are sampled using the forward transition probabilities P ϕ ( G ′ | G ) during the first phase, and P ϕ ( θ | G ) during the second one. Following Deleu et al. (2022), we parametrize these forward transition probabilities using a hierarchical model: we first decide whether we want to stop the first phase or not, with probability P ϕ (stop | G ) ; then, conditioned on this first decision, we either continue adding an edge to G to reach G ′ with probability P ϕ ( G ′ | G, ¬ stop) (phase 1), or sample θ with probability P ϕ ( θ | G, stop) (phase 2). This hierarchical model can be written as

We use neural networks to parametrize each of the three components necessary to define the forward transition probabilities. Unlike in DAG-GFlowNet though, which uses a linear Transformer to define P ϕ (stop | G ) and P ϕ ( 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 a graph embedding g of G , as well as 3 embeddings u i , v i , w i for each node X i in G

We can parametrize the probability of selecting the 'stop' action using g with P ϕ (stop | G ) = f ϕ ( g ) , where f ϕ is a neural network with a sigmoid output; note that if we can't add any edge to G without creating a cycle, we force the end of the first phase by setting P ϕ (stop | G ) = 1 . Inspired by Lorch et al. (2021), the probability of moving from G to G ′ by adding the edge X i → X j is parametrized by

P ϕ ( G ′ | G, ¬ stop) ∝ m ij exp ( u ⊤ i v j ) , (11) where m ij is a binary mask indicating whether adding X i → X j is a valid action (i.e., if it is not already present in G , and if it doesn't introduce a cycle; Deleu et al., 2022). Finally, the probability of selecting the parameters θ i of the CPD for the variable X i is parametrized with a multivariate Normal distribution with diagonal covariance (unless specified otherwise)

where µ ϕ and σ 2 ϕ are two neural networks, with appropriate non-linearities to guarantee that σ 2 ϕ ( w i ) is a well-defined diagonal covariance matrix. Note that P ϕ ( θ i | G, stop) effectively approximates the posterior distribution P ( θ i | G, D ) once fully trained. Moreover, in addition to being an approximation of the joint posterior P ( G,θ | 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 θ .

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,θ ) in (5). However, beyond the capacity to have an arbitrary likelihood model P ( D | θ, G ) (e.g., non-linear), another advantage of approximating the joint posterior

P ( G,θ | D ) is that we can train the GFlowNet using mini-batches of observations. Concretely, for a mini-batch B of M observations sampled uniformly at random from the dataset D , we can define

which is an 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 ̂ L B ( ϕ ) be the learning objective defined in Section 3.3, where the reward has been replaced by the estimate ̂ R B ( G,θ ) in (13). Then we have L ( ϕ ) ≤ E B [ ̂ L B ( ϕ ) ] .

The proof is available in Appendix C.3.4, only relies on the convexity of the square function to conclude, but no other property of this function; in practice, we use the Huber loss instead of the square loss for stability, which is also a convex function. In the case of the square loss, Proposition C.1 gives a stronger result in terms of unbiasedness of the gradient estimator.

4 Related work

Bayesian Structure Learning. There is a vast literature applying Markov chain Monte Carlo (MCMC) methods to 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,θ | D ) using MCMC requires additional transdimensional updates (Fronk, 2002), and has therefore received less attention than the marginal case.

Variational methods have been proposed to approximate the marginal posterior too (Annadani et al., 2021; Charpentier et al., 2022). Similar to MCMC though, approximating the joint posterior has also been less studied than its marginal counterpart, with the notable exceptions of DiBS (Lorch et al., 2021) and BCD Nets (Cundy et al., 2021). We provide an extensive qualitative comparison between our method JSP-GFN and prior variational inference and GFlowNet methods in Appendix A.

Generative Flow Networks. While they were initially developed to encourage the discovery of diverse molecules (Bengio et al., 2021), GFlowNets proved to be a more general framework to describe distributions over composite objects that can be constructed sequentially (Bengio et al., 2023). The objective of the GFlowNet is to enforce a conservation law such as the flow-matching conditions in (2), indicating that the total amount of flow going into any state is equal to the total outgoing flow, with some residual given by the reward. Alternative conditions, sometimes bypassing the need to work with flows altogether, have been proposed in order to learn these models more efficiently (Malkin et al., 2022; Madan et al., 2022; Pan et al., 2023). By amortizing inference, and thus treating it as an optimization problem, GFlowNets find themselves deeply rooted in the variational inference literature (Malkin et al., 2023; Zimmermann et al., 2022), and are connected to other classes of generative models (Zhang et al., 2022). Beyond Bayesian structure learning (Deleu et al., 2022), GFlowNets have also applications in modeling Bayesian posteriors for variational EM (Hu et al., 2023), combinatorial optimization (Zhang et al., 2023), biological sequence design (Jain et al., 2022), as well as scientific discovery at large (Jain et al., 2023).

Closely related to our work, Nishikawa-Toomey et al. (2023) proposed to learn the joint posterior P ( G,θ | D ) over structures and parameters with a GFlowNet, combined with Variational Bayes to circumvent the challenge of learning a distribution over continuous quantities θ with a GFlowNet. Atanackovic et al. (2023) also used a GFlowNet called DynGFN to approximate the posterior of a dynamical system. Similar to (Nishikawa-Toomey et al., 2023) though, they used the GFlowNet only to approximate the distribution over graphs G , making the parameters θ a deterministic function of G (i.e., P ( θ | G, D ) ≈ δ ( θ | G ; ϕ ) ). Here, we leverage the recent advances extending these models to general sample spaces (Lahlou et al., 2023), including continuous spaces (Li et al., 2023), in order to model the joint posterior within a single GFlowNet.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Exact posterior

JSP-GFN (diag)

(a) Edge features

r

= 0

.

9984

Edge features

E

G,θ

[

-

log

P

(

θ

|

G,

D

)

]

RMSE

Pearson's

r

MH-MC

3

0

.

357

±

0

.

022

0

.

067

±

0

.

143

5

.

39

±

1

.

41

×

10

2

Gibbs-MC

3

0

.

357

±

0

.

022

0

.

028

±

0

.

127

9

.

02

±

1

.

54

×

10

5

B-GES

*

0

.

263

±

0

.

070

0

.

635

±

0

.

180

1

.

56

±

0

.

97

×

10

2

B-PC

*

0

.

305

±

0

.

057

0

.

570

±

0

.

138

1

.

57

±

0

.

87

×

10

2

DiBS

0

.

312

±

0

.

038

0

.

737

±

0

.

071

9

.

49

±

7

.

34

×

10

3

BCD Nets

0

.

215

±

0

.

055

0

.

819

±

0

.

097

7

.

04

±

3

.

21

×

10

1

VBG

0

.

237

±

0

.

037

0

.

816

±

0

.

064

1

.

24

±

0

.

49

×

10

2

JSP-GFN (diag)

0

.

018

±

0

.

005

0

.

998

±

0

.

001

-

4

.

91

±

0

.

51

×

10

0

JSP-GFN (full)

0

.

019

±

0

.

007

0

.

998

±

0

.

001

-

5

.

00

±

0

.

52

×

10

0

(b) Quantitative comparison with the exact posterior

5 Experimental results

5.1 Joint posterior over small graphs

We can evaluate the accuracy of the approximation returned by JSP-GFN by comparing it with the exact joint posterior distribution P ( G,θ | D ) . Computing this exact posterior can only be done in limited cases only, where (1) P ( θ | G, 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 d = 5 variables, with linear Gaussian CPDs. We generate 20 different datasets of N = 100 observations from randomly generated Bayesian Networks. Details about data generation and modeling are available in Appendix D.2.

The quality of the joint posterior approximations is evaluated separately for G and θ . 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 edge feature corresponds to the marginal probability of a specific edge being in the graph. Figure 2 (a) shows a comparison between the edge features computed with the exact posterior and with JSP-GFN, proving that it can accurately approximate the edge features of the exact posterior, despite the modeling bias discussed in Appendix D.2.2. Compared to other methods in Figure 2 (b), JSP-GFN offers significantly more accurate approximations of the posterior, at least relative to the edge features. This observation still holds on the path and Markov features (Deleu et al., 2022); see Appendix D.2.2.

To evaluate the performance of the different methods as an approximation of the posterior over θ , we also estimate the cross-entropy between the sampling distribution of θ given G and the exact posterior P ( θ | G, D ) . This measure will be minimized if the model correctly samples parameters from the true P ( θ | 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 a 2-layer MLP, for a total of | θ | = 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.

M-MC

3

G-MC

3

B-PC

*

B-GES

*

DiBS

BCD Nets

VBG

JSP-GFN

1

,

000

2

,

000

3

,

000

4

,

000

Negative log-likelihood

M-MC

3

G-MC

3

DiBS

JSP-GFN

1

,

000

3

,

000

5

,

000

7

,

000

9

,

000

-

2

,

600

-

2

,

400

-

800

-

700

-

600

-

500

-

400

log

P

(

D

, G, θ

)

log

P

/latticetop

φ

(

G,θ

)

30

40

50

60

∆ = 0

.

9864

(a) Linear Gaussian

(b) Non-linear Gaussian

(c) Non-linear Gaussian

We compared JSP-GFN against two methods based on MCMC (MH-MC 3 &Gibbs-MC 3 ; Madigan et al., 1995) and DiBS (Lorch et al., 2021) on both experiments, as well as two bootstrapping algorithms (B-GES * & B-PC * ; Friedman et al., 1999), BCD Nets (Cundy et al., 2021) and VBG (Nishikawa-Toomey et al., 2023) for the experiment with linear Gaussian CPDs, as they are not applicable for non-linear CPDs. In Figure 3 (a-b), we report the performance of these joint posterior approximations in terms of the (expected) negative log-likelihood (NLL) on held-out observations. We observe that JSP-GFN achieves a lower NLL than any other method on linear Gaussian models and is competitive on non-linear Gaussian models.

We chose the NLL over some other metrics, typically comparing with the ground-truth graphs used for data generation, since it is more representative of the performance of these methods on downstream tasks (i.e., predictions on unseen data), and measures the quality of the joint posterior instead of only the marginal over graphs. This choice is aligned with the shortcomings of these other metrics highlighted by Lorch et al. (2022), and it is further justified in Appendix D.3.3. Nevertheless, we also report the expected Structural Hamming Distance (SHD), as well as the area under the ROC curve (AUROC) in Appendix D.3.3 for completeness. To complement these metrics, and in order to assess the quality of the approximation of the posterior in the absence of reference P ( G,θ | D ) , we also show in Figure 3 (c) how the terminating state log-probability log P ⊤ ϕ ( G,θ ) of JSP-GFN correlates with the log-reward for a non-linear Gaussian model. Indeed, as stated in Theorem 3.1, we should ideally have log P ⊤ ϕ ( G,θ ) perfectly correlated with log R ( G,θ ) with slope 1, as

We can see that there is indeed a strong linear correlation across multiple samples ( G,θ ) from JSPGFN, with a slope ∆ close to 1, suggesting that the GFlowNet is again an accurate approximation of the joint posterior, at least around the modes it captures. Details about how log P ⊤ ϕ ( G,θ ) is estimated are available in Appendix D.3.2.

5.3 Learning biological structures from real data

We finally evaluated JSP-GFN on real-world biological data for two separate tasks: the discovery of protein signaling networks from flow cytometry data (Sachs et al., 2005), as well as the discovery of a small gene regulatory network from gene expression data. The flow cytometry dataset consists of N = 4 , 200 measurements of d = 11 phosphoproteins from 7 different experiments, meaning that this dataset contains a mixture of both observational and interventional data. Furthermore, this dataset has been discretized into 3 states, representing the level of activity (Eaton and Murphy, 2007). For the gene expression dataset, we used a subset of N = 2 , 628 observations of d = 61 genes from (Sethuraman et al., 2023). Details about the experimental setups are available in Appendix D.4.

-

3

,

000

-

2

,

500

-

2

,

000

-

2

,

250

-

2

,

000

-

1

,

750

-

1

,

500

-

1

,

250

log

P

(

D

, G, θ

)

log

P

/latticetop

φ

(

G,θ

)

5

10

15

20

-

3

.

26

-

3

.

24

-

3

.

22

-

3

.

2

·

10

4

-

2

,

400

-

2

,

300

-

2

,

200

-

2

,

100

-

2

,

000

log

P

(

D

, G, θ

)

10

12

14

16

18

∆ = 1

.

0272

∆ = 0

.

7817

(a) Flow cytometry data

(subsample,

N

= 100)

(b) Flow cytometry data

Flow cytometry

Gene expression

d

= 11

, N

= 4

,

200

d

= 61

, N

= 2

,

628

MH-MC

3

2

.

578

×

10

4

1

.

073

×

10

6

Gibbs-MC

3

3

.

490

×

10

5

4

.

882

×

10

6

JSP-GFN

2

.

679

×

10

4

2

.

651

×

10

5

(c) Negative log-likelihood

At this scale, using the whole dataset D to evaluate the reward becomes impractical, especially for non-linear models. Fortunately, we showed in Section 3.5 that we can use an (unbiased) estimate of the reward, based on mini-batches of data, in place of R ( G,θ ) in the loss function (8). In both experiments, we used non-linear models, where all the CPDs are parametrized with a 2-layer MLP. Figure 4 shows similar correlation plots as Figure 3 (c), along with an evaluation of the NLL on unseen observations and interventions. Beyond the ability to work with real data, these experiments allow us to highlight some other capacities of JSP-GFN: (1) handling discrete and (2) interventional data (flow cytometry), as well as (3) learning a distribution over larger graphs (gene expression).

To measure the quality of the posterior approximation returned by JSP-GFN, we compare in Figure 4 the terminating state log-probability log P ⊤ ϕ ( G,θ ) with the log-reward log R ( G,θ ) , similar to Figure 3. We observe that there is correlation between these two quantities; unlike in Figure 3 though, we observe that the slope is not close to 1 , suggesting that JSP-GFN underestimates the probability of ( G,θ ) . We also observe that the graphs are 'clustered' together; this can be explained by the fact that the posterior approximation is concentrated at only a few graphs, since the size of the dataset 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-GFN with 2 methods based on MCMC in terms of their negative log-likelihood on held-out data. We can see that JSP-GFN is competitive, and even out-performs MCMC on the more challenging problem of the discovery of gene regulatory networks from gene expression data, where the dimensionality of the problem is much larger ( d = 61 ). Note that the values reported for the discovery of protein signaling networks from flow cytometry data correspond to the negative interventional log-likelihood, on interventions unseen in D .

6 Conclusion

We have presented JSP-GFN, an approach to approximate the joint posterior distribution over the structure of a Bayesian Network along with its parameters using a single GFlowNet. We have shown that our method faithfully approximates the joint posterior on both simulated and real data, and compares favorably against existing Bayesian structure learning methods. In line with Appendix B, future work should consider using more expressive distributions, such as those parametrized by normalizing flows or diffusion processes, to approximate the posteriors over continuous parameters, which would enable Bayesian inference over parameters in more complex generative models.

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-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint , 2018.

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 J Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet Foundations. Journal of Machine Learning Research (JMLR) , 2023.

Bertrand Charpentier, Simon Kibler, and Stephan G¨ unnemann. 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´ onio G´ ois, 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 structure learning using dynamic programming and MCMC. Conference on Uncertainty in Artificial Intelligence , 2007.

Paul Erd˝ os and Alfr´ ed R´ enyi. 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: a 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 on Uncertainty in Artificial Intelligence , 1999.

Nir Friedman, Michal Linial, Iftach Nachman, and Dana Pe'er. Using bayesian networks to analyze expression data. In Proceedings of the fourth annual international conference on Computational molecular biology , 2000.

Eva-Maria Fronk. Model Selection for Dags via RJMCMC for the Discrete and Mixed Case. Technical report, Ludwig-Maximilians-Universit¨ at M¨ unchen, 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 M Chickering. Learning bayesian networks: The combination 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¨ arinen. Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research , 2005.

Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure F.P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological Sequence Design with GFlowNets. International Conference on Machine Learning , 2022.

Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez-Garcia, and Yoshua Bengio. GFlowNets for AI-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´ andezGarc´ ıa, L´ ena N´ ehale 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 Representations , 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¨ olkopf, 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¨ olkopf. Amortized Inference for Causal Structure Learning. Advances in 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 E 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.

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-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 GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning , 2023.

Karen Sachs, Omar Perez, Dana Pe'er, Douglas A Lauffenburger, and Garry P Nolan. Causal protein-signaling networks derived from multiparameter single-cell data. Science , 2005.

Muralikrishnna G Sethuraman, Romain Lopez, Rahul Mohan, Faramarz Fekri, Tommaso Biancalani, and Jan-Christian H¨ utter. NODAGS-Flow: Nonlinear Cyclic Causal Structure Learning. International Conference on 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, Łukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. Advances in Neural Information Processing Systems , 2017.

Jussi Viinikka, Antti Hyttinen, Johan Pensar, and Mikko Koivisto. Towards Scalable Bayesian Learning of Causal DAGs. Advances in Neural Information Processing Systems , 2020.

Julius von K¨ ugelgen, Paul K Rubenstein, Bernhard Sch¨ olkopf, and Adrian Weller. Optimal experimental design via Bayesian optimization: active causal structure learning for Gaussian process networks. NeurIPS Workshop 'Do the right thing': machine learning and causal inference for improved decision making , 2019.

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 GFlowNets and Beyond. International Conference on Machine Learning - Beyond Bayes workshop , 2022.

Xun Zheng, Bryon Aragam, Pradeep Ravikumar, and Eric P. Xing. DAGs with NO TEARS: Continuous Optimization for Structure Learning. In Advances in Neural Information Processing Systems , 2018.

Heiko Zimmermann, Fredrik Lindsten, Jan-Willem van de Meent, and Christian A. Naesseth. A Variational Perspective on Generative Flow Networks. arXiv preprint , 2022.

Appendix

A Positioning JSP-GFN in the Bayesian structure learning literature

We give a comparison between our method JSP-GFN, and various methods based on variational inference in Table A.1. We also include methods based on GFlowNets, namely DAG-GFlowNet (Deleu et al., 2022) and VBG (Nishikawa-Toomey et al., 2023), as they are effectively variational methods (Malkin et al., 2023; Zimmermann et al., 2022).

Joint G & θ . This category indicates whether the model can approximate the joint posterior distribution P ( G,θ | D ) over both graphical structures G and parameters of the CPDs θ , or if they are limited to approximating the marginal posterior P ( G | D ) . As we have seen in 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 a non-linear function (e.g., a neural network). While most methods approximating the marginal distribution may be applied to non-linear CPDs parametrized by a Gaussian Process (von K¨ ugelgen 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 as future work. While VBG (Nishikawa-Toomey et al., 2023) has only been applied to linear Gaussian models, the framework may also be applicable to non-linear models.

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 a 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 et al. (2022) reports that 1 . 50% of the graphs returned by DiBS 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 not enforce acyclicity by design, to allow flexibility on other domains (e.g., for modeling gene regulatory networks; Sethuraman et al., 2023). TRUST (Wang et al., 2022) guarantees acyclicity via a distribution over variable orders, that can be 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 a discrete domain, this extension can be found in the official code released. A VICI (Lorch et al., 2022) assumes access to a generative model P ( D | G ) , making it possibly applicable to discrete domains as

well. Since it builds on DAG-GFlowNet (Deleu et al., 2022), VBG (Nishikawa-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 DiBS as its underlying routine for structure learning, it should also inherit the properties of DiBS.

MaximumParents. This category indicates whether a 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 a sparsity-inducing prior (Lorch et al., 2021; Cundy et al., 2021), or use 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-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, or 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 (Nishikawa-Toomey et al., 2023) may be updated using mini-batches, since the maginalization over θ 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 in 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 1), 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 a diagonal covariance, except for JSP-GFN (full) in Section 5.1) to parametrize the approximation of the posterior over parameters P ϕ ( θ | G, stop) (see Section 3.4). This limits its expressivity to unimodal distributions only, and is an assumption which is commonly used with Bayesian neural networks. However in general, the posterior distribution P ( θ | G, D ) may be highly multimodal, especially when the model is non-linear (the posterior is Normal when the model is linear Gaussian,

s

s

′

P

F

(

s

′

|

s

)

P

B

(

s

|

s

′

)

s

0

s

1

. . .

s

T

s

f

R

(

s

)

P

F

(

s

0

/squiggleright

s

f

)

P

B

(

s

0

/squiggleright

s

T

)

s

m

. . .

s

n

P

F

(

s

m

/squiggleright

s

n

)

P

B

(

s

m

/squiggleright

s

n

)

s

k

. . .

s

m

. . .

s

n

(a) Detailed Balance

(b) Trajectory Balance

(c) Subtrajectory Balance

(d) Generalized

Subtrajectory Balance

see Appendix D.5.1). To see this, consider a non-linear model whose CPDs are parametrized by a 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 θ leading to the same likelihood function, and under mild assumptions on the priors P ( θ | G ) and P ( G ) , they would have the same posterior probability P ( θ | G, D ) .

To address this issue of unimodality, we can use more expressive posterior approximations P ϕ ( θ | G, stop) , such as ones parametrized with diffusion-based models, or with normalizing flows; both of these models are drop-in replacements in JSP-GFN, since their likelihood can be explicitly computed. An alternative is also to consider multiple steps of a continuous GFlowNet (Lahlou et al., 2023), instead of a single one, to generate θ .

Biological plausibility of the acyclicity assumption. One of the strengths of JSP-GFN, and DAGGFlowNet before it (Deleu et al., 2022), is the capacity to obtain a distribution over the DAG structure of a 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 a general framework though, the GFlowNet used in JSP-GFN can be adapted to 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 a cyclic graph by unrolling it, as in Atanackovic et al. (2023).

C Details about Generative Flow Networks

Throughout this section, we will use both P F and P ϕ (to emphasize the parametrization on ϕ , as in the main text) to denote equally the forward transition probability-the notation P F being more commonly used in the literature on GFlowNet (Bengio et al., 2023).

C.1 Alternative conditions

GFlowNets were initially introduced using the flow-matching conditions (Bengio et al., 2021), 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 to the reward).

One of those alternative conditions are the detailed balance conditions (Bengio et al., 2023), inspired by the literature on Markov chains. These conditions are given for any transition s → s ′ in the GFlowNet as 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 all the states of the GFlowNet are complete, in order to avoid having to learn a separate flow function (see Section 2.3).

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 τ = ( s 0 , s 1 , . . . , s T , s f ) , the trajectory balance condition is given by

with the convention s T +1 = s f , and where Z is the partition function of the distribution (i.e., Z = ∑ x ∈X R ( x ) ); in practice, Z is a parameter of the model that is being learned alongside the forward and backward transition probabilities. The trajectory balance condition is illustrated in Figure C.1 (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 a generalization of both the detailed balance and trajectory balance conditions to partial trajectories of arbitrary length. For a partial trajectory τ = ( s m , s m +1 , . . . , s n ) , the subtrajectory balance condition is given by

where again F ( s ) is a flow function (as in (C.1)). This condition encompasses both conditions in Appendix C.1, since we can recover the detailed balance condition in (C.1) with partial trajectories of length 1 (i.e., transitions), and also the trajectory balance condition in (C.2) with complete trajectories (note that F ( s 0 ) = 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 to create a novel objective called SubTB( λ ) , inspired by TD( λ ) 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 s m and s n , this (generalized) subtrajectory balance condition can be written as

where s k is a common ancestor of both s m and s n . This condition is illustrated in Figure C.1 (d). While these subtrajectory balance conditions (generalized or not) offer more flexibility, they are guaranteed to yield a GFlowNet inducing a distribution proportional to R ( s ) only if these conditions are satisfied for all the partial trajectories of any length. In particular, they provide no guarantee 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 ∝ R ( s ) in our case.

C.3 Proofs

C.3.1 Subtrajectory balance conditions for undirected paths of length 3

Consider an undirected path of length 3 of the form ( G,θ ) ← ( G, · ) → ( G ′ , · ) → ( G ′ , θ ′ ) , where G ′ is the result of adding a new edge to the DAG G (see Figure C.2). Since the state ( G, · ) is a common ancestor of both complete states ( G,θ ) and ( G ′ , θ ′ ) , we can write the subtrajectory balance conditions (C.4) as

F ( G,θ ) P B ( G | θ ) P F ( G ′ | G ) P F ( θ ′ | G ′ ) = F ( G ′ , θ ′ ) P B ( G ′ | θ ′ ) P B ( G | G ′ ) P F ( θ | G ) , (C.5) where we abuse the notation P B ( G | θ ) again to denote P B ( ( G, · ) | ( G,θ ) ) . In fact, since the complete state ( G,θ ) ∈ X has only a single parent state ( G, · ) , we necessarily have P B ( G | θ ) = 1 (and similarly for ( G ′ , θ ′ ) ). Furthermore, we can use the observation from (Deleu et al., 2022) to write the flow F ( G,θ ) of a complete state ( G,θ ) as a function of its reward

G

·

G

′

·

G

˜

θ

G

θ

G

′

θ

′

R

(

G,θ

)

R

(

G

′

, θ

′

)

R

(

G,

˜

θ

)

Undirected paths of length 3 R ( G ′ , θ ′ ) P B ( G | G ′ ) P ϕ ( θ | G ) = R ( G,θ ) P ϕ ( G ′ | G ) P ϕ ( θ ′ | G ′ )

Undirected paths of length 2 (differential form) log P ( θ G ) = log R ( G,θ )

∇ θ ϕ | ∇ θ

We can simplify (C.6) even further by observing that in the GFlowNet used here, s f is the only child of the complete state ( G,θ ) ∈ X . In other words, a complete state ( G,θ ) is not directly connected to any other ( G, ˜ θ ) ; this is the (infinitely wide) tree structure rooted at ( G, · ) mentioned in Section 3.1. Since s f is the only child of ( G,θ ) , we then necessarily have P F ( s f | ( G,θ ) ) = 1 , and therefore F ( G,θ ) = R ( G,θ ) . With these simplifications, (C.5) becomes

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,θ ) ← ( G, · ) → ( G, ˜ θ ) (see Figure C.2). Since ( G, · ) is a common ancestor (a common parent in this case) of both complete states ( G,θ ) and ( G, ˜ θ ) , we can write the subtrajectory balance conditions (C.4) as

Using the same simplifications as in Appendix C.3.1 ( P B ( G | θ ) = P B ( G | ˜ θ ) = 1 ), we get the following subtrajectory balance conditions for the undirected paths of length 2

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 ( G,θ ) and ( G ′ , θ ′ ) . Indeed, if we write these conditions between ( G,θ ) and ( G ′ , θ ′ ) on the one hand, and between ( G, ˜ θ ) and ( G ′ , θ ′ ) on the other hand (with a fixed G ′ and θ ′ )

we get the same subtrajectory balance conditions over undirected paths of length 2 as in (C.9):

However, since the SubTB conditions (6) are only satisfied approximately in practice, it might be advantageous to also satisfy (C.9) in addition to those in (6). The equation above provides an alternative way to express (C.9). Indeed, (C.12) shows that the function

is constant, albeit with a constant that depends on the graph G . Since this function is differentiable, this is equivalent to ∇ θ f G ( θ ) = 0 , and therefore we get the differential form of the subtrajectory balance conditions in (7)

As we saw in Section 3.3, one way to enforce the SubTB conditions over undirected paths of length 3 is to create a learning objective that encourages these conditions to be satisfied, and optimizing it

using gradient methods. The learning objective has the form L ( ϕ ) = E π [ ˜ ∆ 2 ( ϕ )] , where ˜ ∆( ϕ ) is a non-linear residual term

Suppose that the parameters ϕ of the GFlowNet are such that the subtrajectory balance conditions in (C.14) are satisfied for any ( G,θ ) . 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 θ and θ ′ depend on ϕ (via the reparametrization trick since they are sampled on-policy, see Section 5), taking the derivative of ˜ ∆ 2 ( ϕ ) , we get

and similarly for the terms in ( G ′ , θ ′ ) . The derivative of the objective then becomes

An alternative way to obtain the same derivative in (C.17) is to take dθ/dϕ = 0 instead, meaning that we would not differentiate through θ (and θ ′ ). Using the stop-gradient operation ⊥ , this shows that the following objective 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 2 are satisfied.

While optimizing (C.20) alone leads to eventually satisfying the subtrajectory balance conditions over undirected paths of length 2, it may be advantageous to explicitly encourage this behavior, especially in cases where d is larger and/or for non-linear models. We can incorporate some penalty to the loss function, such as

Theorem 3.1. If the SubTB conditions in (6) are satisfied for all undirected paths of length 3 between any ( G,θ ) and ( G ′ , θ ′ ) of the form ( G,θ ) ← ( G, · ) → ( G ′ , · ) → ( G ′ , θ ′ ) , then we have

where P ϕ ( G | G 0 ) is the marginal probability of reaching G from the initial state G 0 with any (complete) trajectory τ = ( G 0 , G 1 , . . . , G T -1 , G ) :

using the conventions G T = G , and P ϕ ( G 0 | G 0 ) = 1 .

Proof. We assume that the SubTB conditions are satisfied for all undirected paths of length 3 between any ( G,θ ) and ( G ′ , θ ′ ) , that is

̸

Let G = G 0 be a fixed DAG different from the initial state, and θ a set of corresponding parameters. Let τ = ( G 0 , . . . , G T -1 , G ) be an arbitrary trajectory from G 0 to G , where we use the convention G T = G . For any t < T , if θ t is a fixed set of parameters associated with G t , then the SubTB conditions above can written for every timestep as

again, using the convention θ T = θ . Taking the product of the ratio between P ϕ and P B over the trajectory τ , we get

Moreover, the backward transition probability P B , defined only over the transitions of the GFlowNet, induces a distribution over the trajectories from G 0 to G (Bengio et al., 2023), meaning that

Therefore, we have

We saw in Appendix C.3.2 that P ϕ ( θ 0 | G 0 ) /R ( G 0 , θ 0 ) is independent of the value of θ 0 if the SubTB conditions are satisfied for all undirected paths of length 3 (see (C.12)). This concludes the proof: P ϕ ( G | G 0 ) P ϕ ( θ | G ) ∝ R ( G,θ ) .

C.3.4 Mini-batch training

Proposition 3.2. Suppose that B is a mini-batch of M observations sampled uniformly at random from the dataset D , and let ̂ L B ( ϕ ) be the learning objective defined in Section 3.3, where the reward has been replaced by the estimate ̂ R B ( G,θ ) in (13). Then we have L ( ϕ ) ≤ E B [ ̂ L B ( ϕ ) ] .

Proof. We will first show that log ̂ R B ( G,θ ) defined in (13) is an unbiased estimate of the log-reward log R ( G,θ ) under a uniform distribution of the mini-batches B . We can observe that by conditional independence of the observations x ( n ) given G and θ , we have

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 in (C.30) because the observations are mutually independent given G and θ ; if we were only conditioning on G (i.e., using the marginal likelihood,

as in (Deleu et al., 2022)), then those observations would not be conditionally independent in general. Similarly, we have

Therefore, it shows that the estimate of the log-reward is unbiased:

Recall that the estimate of the loss is defined as ̂ L B ( ϕ ) = E π [ ̂ ∆ 2 B ( ϕ ) ] , where the residual is defined by

Taking the expectation of this estimated loss wrt. a random mini-batch B , we get

where we used the convexity of the square function and Jensen's inequality in (C.36), and the unbiasedness of log ̂ R B ( G,θ ) (as well as log ̂ R B ( G ′ , θ ′ ) ) in (C.37); recall that ∆( ϕ ) is given in (C.20) (see also (8)).

Proposition C.1. The mini-batch gradient estimator is unbiased, i.e., ∇ ϕ L ( ϕ ) = E B [ ∇ ϕ ̂ L B ( ϕ ) ] . 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 is unbiased. We observe that

∇ ϕ ∆( ϕ ) = ∇ ϕ ̂ ∆ B ( ϕ ) (C.39) since only the terms corresponding to the rewards differ between ∆( ϕ ) and ̂ ∆ B ( ϕ ) , and they do not depend on ϕ . Therefore,

as desired. This implies the expected mini-batch loss and full-batch loss differ by a constant and have the same set of local and global minima. This constant happens to equal and showing the difference equals this constant yields an alternative proof.

Algorithm 1 JSP-GFN training procedure

D Additional experiments & experimental details

In addition to Bayesian structure learning methods based on variational inference (Lorch et al., 2021; Cundy et al., 2021) or GFlowNets (Nishikawa-Toomey et al., 2023), we also consider 2 baseline methods based on MCMC, and 2 methods based on bootstrapping (Friedman et al., 1999), as introduced in (Lorch et al., 2021). Metropolis-Hastings MC 3 (MH-MC 3 ) samples both graphical structures and parameters jointly at each move, whereas Metropolis-within-Gibbs MC 3 (Gibbs-MC 3 ) alternates between updates of the structure, and updates of the parameters; note that MC 3 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 of the marginal posterior P ( G | D ) only, the parameter sample θ corresponding to a DAG G correspond to the parameter inferred by P ( θ | G, D ) (see Appendix D.5.1).

D.1 Sampling distribution

In Section 3.3, we saw that the learning objective of JSP-GFN can be written as L ( ϕ ) = E π [∆ 2 ( ϕ )] , where π is a sampling distribution over ( G,θ ) and ( G ′ , θ ′ ) with full support. We use a combination of on-policy ( π = P ϕ ) 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 off-policy from a replay buffer, whereas their corresponding parameters θ and θ ′ are sampled on-policy using our current P ϕ ( θ | G ) . Therefore, a key difference with Deleu et al. (2022) is that the reward R ( G,θ ) in (5) is calculated 'lazily' when the loss is evaluated (i.e., only once θ and θ ′ are known), as opposed to being computed during the interaction with the state space and stored in the replay buffer alongside the transitions. Algorithm 1 gives a description of the training algorithm of JSP-GFN.

D.2 Joint posterior over 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 a graph from an Erd¨ os-R´ enyi model (Erd˝ os and R´ enyi, 1960) over d = 5 nodes, with d edges on average (a setting typically referred to as ER1). Once the structure of the graph G ⋆ is known, we sample the parameters θ ⋆ of the linear Gaussian model randomly from a standard Normal distribution N (0 , 1) . The linear Gaussian model is defined as

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Exact posterior

JSP-GFN (diag)

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Exact posterior

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Exact posterior

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Exact posterior

JSP-GFN (full)

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Exact posterior

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Exact posterior

(a) Edge features

(b) Path features

(d) Edge features

(e) Path features

r

= 0

.

9984

r

= 0

.

9982

r

= 0

.

9980

r

= 0

.

9980

r

= 0

.

9976

r

= 0

.

9972

where θ ⋆ ij ∼ N (0 , 1) , and ε i ∼ N (0 , 0 . 01) ; this defines all the CPDs necessary for (1). Finally, we use ancestral sampling to generate N = 100 observations to create the dataset D .

Modeling. We use a linear Gaussian Bayesian Network to model the data, where the CPD for the variable X i can be written as P ( X i | Pa G ( X i ); θ i ) = N ( µ i , σ 2 ) , where and where σ 2 = 0 . 01 is a fixed variance across variables, matching the variance used for data generation. We place a unit Normal prior over the parameters θ ij of the model: P ( θ ij | G ) = N (0 , 1) . This model differs from the widely used BGe score (Geiger and Heckerman, 1994) in that σ 2 is treated as a hyperparameter here, instead of a 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 over graphs. Under this model, the posterior distribution P ( θ | G, D ) can be computed analytically, and the proof is available in Appendix D.5.1.

D.2.2 Comparing JSP-GFN with the exact posterior against features

In Section 5.1, we evaluated the accuracy of the posterior approximation returned by JSP-GFN by comparing them on the edge features, i.e., the marginal distribution of an edge X i → X j is present in the graph:

✶ 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 P ⊤ ϕ induced by the GFlowNet (discarding θ , 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 as the path feature and the Markov feature . The path feature corresponds to

±

±

±

±

±

±

the marginal probability of a path X i ⇝ X j being present in the graph, and the Markov feature is the marginal probability of a node X i being in the Markov blanket of X j . In other words

✶ where MB G ( X j ) denotes the Markov blanket of X j 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-GFN. The first one, called JSP-GFN (diag) , where P ϕ ( θ i | G, 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 ( θ i | G, D ) 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-Toomey et al., 2023).

In Figure D.1, we show a similar plot as in Figure 2 (a) for all these features, for both models JSP-GFN (diag) and JSP-GFN (full). We observe that for all features, the approximation of the posterior given by JSP-GFN is very accurate (as confirmed by the Pearson's correlation coefficients). Interestingly, the more expressive model JSP-GFN (full) seems to perform slightly worse than JSPGFN (diag); this is also confirmed in part by the quantitative measures in Table D.1. This can be explained by the additional number of parameters of the neural network ϕ necessary to output the higher dimensional full covariance matrix (more precisely, a lower-triangular matrix corresponding to its Cholesky decomposition). In Table D.1, we show a quantitative comparison across the different Bayesian structure learning methods on the three features, in terms of RMSE and Pearson's correlation coefficient. Similar to Section 5.1, 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 θ , we measure how likely sample parameters from the approximation are under the exact posterior distribution P ( θ | G, D ) . More precisely, we compute the cross-entropy between the posterior approximation P ϕ ( G,θ ) and the exact joint posterior P ( G,θ | D ) : given a distribution P ϕ ( G,θ ) approximating the joint posterior P ( G,θ | D ) , we estimate where { ( G ( k ) , θ ( k ) ) } K k =1 are K samples of the posterior approximation P ϕ ( G,θ ) . We use this measure as it can be estimated from samples.

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 Erd¨ os-R´ enyi model with 2 d edges on average (a setting commonly referred to 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 et al., 2021) with randomly generated weights. We also sample N = 100 observations to create the dataset D .

Modeling. We use a Gaussian Bayesian Network to model the data, where the CPD for the variable X i can be written as P ( X i | Pa G ( X i ); θ i ) = N ( µ i , σ 2 ) , where where M i = diag ( ✶ ( X 1 ∈ Pa G ( X i )) , . . . , ✶ ( X d ∈ Pa G ( X i )) ) and X = ( X 1 , . . . , X d ) . The variance σ 2 = 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 a ReLU activation function, for a total of | θ | = 2 , 220 parameters. The priors over parameters P ( θ | 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 JSP-GFN with the exact joint posterior P ( G,θ | D ) directly, since the latter becomes intractable (even with a linear Gaussian model) due to the super-exponential size of the sample space. Alternatively, since the terminating state probability P ⊤ ϕ ( G,θ ) of JSP-GFN should ideally be equal to the joint posterior (see Theorem 3.1), we have

where log P ( D ) is a constant corresponding to the log-partition function. Therefore, we can compare the log-terminating state probability log P ⊤ ϕ ( G,θ ) with the log-reward log R ( G,θ ) (which we can compute analytically) for different samples ( G,θ ) , and find a linear relation. This evaluation strategy was introduced in (Bengio et al., 2021).

However, recall from Theorem 3.1 that the terminating state probability is defined as

where the summation is over all the possible trajectories τ = ( G 0 , G 1 , . . . , G T ) from G 0 to G T = G . If G is a 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 P B ( G t | G t +1 ) induces a distribution over the trajectories G 0 ⇝ G (Bengio et al., 2023) to write P ϕ ( G | G 0 ) as

= K ! · E τ ∼ P B [ P ϕ ( τ ) ] , (D.12) where P B ( τ ) = 1 /K ! , since the backward transition probability here is fixed to be uniform over parent states. This suggests a way to get an unbiased estimate of P ϕ ( G | G 0 ) , hence of P τ ϕ ( G,θ ) , based on Monte-Carlo estimation:

where { τ ( m ) } M m =1 are trajectories from G 0 to G , sampled by removing one edge at time uniformly at random, starting at G (i.e., following the backward transition probabilities P B ).

While (D.13) provides an unbiased estimate of P ϕ ( G | G 0 ) , 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 P ϕ ( τ ) 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 T top of B trajectories τ : G 0 ⇝ G that have a large P ϕ ( τ ) , then

where the sample trajectories { τ ( m ) } M m =1 can be obtained using rejection sampling, with a 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 T top , with a beam-size B . More precisely, we need to run beam-search, starting at G 0 , in such a 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 a graph G t to G t +1 = G t ∪ { e } (using this notation to denote that G t +1 is the result of adding the edge e to G t ), with the following score:

˜ P ϕ ( G t +1 | G t ) = ✶ ( e ∈ G ) P ϕ ( G t +1 | G t ) . (D.15) In other words, we only keep transitions corresponding to adding edges that are in G . Note that even though ˜ P ϕ 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-likelihood on held-out data. In addition to the negative log-likelihood 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 { G 1 , . . . , G k } sampled from the posterior approximation as where SHD( G,G ⋆ ) counts the number of edges changes (adding, removing, reversing an edge) necessary to 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 in 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,θ | 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 on the one hand, or 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 a non-linear model with Categorical observations. The CPDs are parametrized using a 2-layer MLP with 16 hidden units and a ReLU activation function, i.e., X i | Pa G ( X i ) ∼ Categorical( π i ) , where

M-MC

3

G-MC

3

B-PC

*

B-GES

*

DiBS

BCD Nets

VBG

JSP-GFN

20

40

60

Expected SHD

M-MC

3

G-MC

3

DiBS

JSP-GFN

40

50

60

M-MC

3

G-MC

3

B-PC

*

B-GES

*

DiBS

BCD Nets

VBG

JSP-GFN

0

.

7

0

.

8

0

.

9

1

AUROC

M-MC

3

G-MC

3

DiBS

JSP-GFN

0

.

5

0

.

6

0

.

7

(a) Linear Gaussian

(b) Non-linear Gaussian

(c) Linear Gaussian

(d) Non-linear Gaussian

where X encodes the discrete inputs as one-hot encoded values, and the MLP has a softmax activation function for the output layer. In total, the model has | θ | = 6 , 545 parameters. The priors over parameters P ( θ | 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 a gene is expressed) or (exactly) zero values (when the gene is inhibited). To capture this type of observations, we model CPDs of the Bayesian Network as zero-inflated Normal distributions:

where µ i is the result of a 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 α i and the variance of the observation noise σ 2 i , for a total of | θ | = 61 , 671 parameters. The priors over parameters P ( θ | G ) and over graphs P ( G ) follow the ones described in Appendix D.2.1.

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

and where σ 2 is a fixed hyperparameter. Moreover, we assume that the parameters have a unit Normal prior associated to them, meaning that

where µ 0 = 0 , σ 2 0 = 1 , and δ 0 is the Dirac measure at 0 , indicating that this parameter is always inactive.

We want to compute the posterior distribution P ( θ i | G, D ) ; this is sufficient, since we know that the parameters of the different CPDs are mutually conditionally independent given G and D . Let X ∈ R N × d be the design matrix of the dataset D (i.e., the observations x ( n ) concatenated row-wise), and by abuse of notation, we denote by X i the i th column of this design matrix. Let D i be a diagonal matrix, dependent on G , defined as

We can rewrite the complete model above as

We abuse the notation above by treating a Dirac distribution at 0 as the limiting case of a Normal distribution with variance 0 . Given this form, we can easily identify that the posterior over θ i is a Normal distribution where we used the conventions 1 / 0 = ∞ and 0 ×∞ = 0 . The masked entries of D i will correspond to zeroed-out entries in ¯ µ i , and to rows and columns of ¯ Σ i being equal to zero, effectively reducing the dimensionality of the distribution to the number of parents of X i in G (e.g., this has an impact on the normalization constant of this distribution). In the limit case where X i has no parent in G , we recover P ( θ i | G, D ) = δ 0 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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).\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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\nQUESTIONS:\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- 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\nLIMITATIONS:\nThe authors only touch on some limitations and a more detailed discussion would strengthen the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
The authors utilize a generative flow network to jointly model the structure of the DAQ and its parameters.

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.

QUESTIONS:
None

LIMITATIONS:
The authors did not discuss this which I think they should.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-19

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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).\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nN/a", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}

comments: SUMMARY:
The present paper extends upon prior literature in making use of GFlowNets for structure learning.
The 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.

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.'_

* GFlowNets are an interesting alternative approach to sampling for various tasks of interest
* Learning a complete Graphical Structure is the optimum for structure learning
* Wonderful presentation, both theory and visualizations/schematics
* 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

QUESTIONS:
* What are the key assumptions (or 'sacrifices') one has to take for the real world example?

LIMITATIONS:
---
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-26

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nAlthough 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n/", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}