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

Tristan Deleu1 Mizu Nishikawa-Toomey1 Jithendaraa Subramanian2
Nikolay Malkin1 Laurent Charlin3 Yoshua Bengio1,4
Mila – Quebec AI Institute 1Université de Montréal, 2McGill University, 3HEC
Montréal, 4CIFAR Senior Fellow.1Correspondence: Tristan Deleu
<EMAIL_ADDRESS>
Code available at: https://github.com/tristandeleu/jax-jsp-gfn

###### 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 ${\mathcal{D}}$.

In addition to being a challenging problem due to the super-exponentially
large search space, learning a single DAG structure from data may also lead to
confident but incorrect predictions (Madigan et al., 1994), especially in
cases where the evidence is limited. In order to avoid misspecifying the
model, it is therefore essential to quantify the epistemic uncertainty about
the structure of the Bayesian Network. This can be addressed by taking a
Bayesian perspective on structure learning and inferring the posterior
distribution $P(G\mid{\mathcal{D}})$ over graphs given our observations. This
(marginal) posterior can be approximated using methods based on Markov chain
Monte Carlo (MCMC; Madigan et al., 1995) or variational inference (Cundy et
al., 2021; Lorch et al., 2021). However, all of these methods rely on the
computation of the marginal likelihood $P({\mathcal{D}}\mid G)$, which can
only be done efficiently in closed form for limited classes of models, such as
linear Gaussian (Geiger and Heckerman, 1994), discrete models with Dirichlet
prior (Heckerman et al., 1995), or non-linear models parametrized with a
Gaussian Process (von Kügelgen et al., 2019).

While there exists a vast literature on Bayesian structure learning to
approximate the marginal posterior distribution, inferring the _joint
posterior_ $P(G,\theta\mid{\mathcal{D}})$ over both the DAG structure $G$ of
the Bayesian Network and the parameters $\theta$ of its conditional
probability distributions—the probability of each variable given its
parents—has received comparatively little attention. The main difficulty
arises from the mixed sample space of the joint posterior distribution, with
both discrete components (the graph $G$) and continuous components (the
parameters $\theta$), where the dimensionality of the latter may even depend
on $G$. However, modeling the posterior distribution over $\theta$ has the
notable advantage that the conditional probability distributions can be more
flexible (e.g., parametrized by neural networks): in general, computing
$P({\mathcal{D}}\mid G,\theta)$ is easier than computing the marginal
$P({\mathcal{D}}\mid G)$, lifting the need to perform intractable
marginalizations.

Since they provide a framework for generative modeling of discrete and
composite objects, Generative Flow Networks (GFlowNets; Bengio et al., 2021,
2023) proved to be an effective method for Bayesian structure learning. In
Deleu et al. (2022), the problem of generating a sample DAG from the marginal
posterior $P(G\mid{\mathcal{D}})$ was treated as a sequential decision
process, where edges are added one at a time, starting from the empty graph
over $d$ variables, following a learned transition probability. Nishikawa-
Toomey et al. (2023) have also proposed to use a GFlowNet to infer the joint
posterior $P(G,\theta\mid{\mathcal{D}})$; however, they used it in conjunction
with Variational Bayes to update the distribution over $\theta$, getting
around the difficulty of modeling a continuous distribution with a GFlowNet.

In this paper, we propose to infer the joint posterior over graphical
structures $G$ and parameters of the conditional probability distributions
$\theta$ of a Bayesian Network using _a single_ GFlowNet called JSP-GFN (for
_Joint Structure and Parameters GFlowNet_), leveraging recent advances
extending GFlowNets to continuous sample spaces (Lahlou et al., 2023), and
expands the scope of applications for Bayesian structure learning with
GFlowNets. The generation of a sample $(G,\theta)$ from the approximate
posterior now follows a two-phase process, where the DAG $G$ is first
constructed by inserting one edge at a time, and then the corresponding
parameters $\theta$ are chosen once the structure is completely known. To
enable efficient learning of the sampling distribution, we introduce new
conditions closely related to the ones derived in Deleu et al. (2022), based
on the subtrajectory balance conditions (Malkin et al., 2022), and show that
they guarantee that the GFlowNet does represent $P(G,\theta\mid{\mathcal{D}})$
once they are completely satisfied. We validate empirically that JSP-GFN
provides an accurate approximation of the posterior when those conditions are
approximately satisfied by a learned sampling model, and compares favorably
against existing methods on simulated and real data.

##  2 Background

##### Notations.

Throughout this paper, we will work with directed graphs $G=(V,E)$, where $V$
is a set of nodes, and $E\subseteq V\times V$ is a set of (directed) edges.
For a node $X\in V$, we denote by $\mathrm{Pa}_{G}(X)$ the set of parents of
$X$ in $G$, and $\mathrm{Ch}_{G}(X)$ the set of its children. For two nodes
$X,Y\in V$, $X\rightarrow Y$ represents a directed edge $(X,Y)\in E$ (denoted
$X\rightarrow Y\in G$), and $X\rightsquigarrow Y$ represents a directed path
from $X$ to $Y$, following the edges in $E$ (denoted $X\rightsquigarrow Y\in
G$).

In the context of GFlowNets (see Section 2.2), an _undirected path_ in a
directed graph ${\mathcal{G}}=({\mathcal{V}},{\mathcal{E}})$ between two
states $s_{0},s_{n}\in{\mathcal{V}}$ is a sequence of vertices
$(s_{0},s_{1},\ldots,s_{n})$ where either $(s_{i},s_{i+1})\in{\mathcal{E}}$ or
$(s_{i+1},s_{i})\in{\mathcal{E}}$ (i.e., following the edges of the graph,
regardless of their orientations).

###  2.1 Bayesian structure learning

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

$P(X_{1},\ldots,X_{d}\mid\theta,G)=\prod_{i=1}^{d}P\big{(}X_{i}\mid\mathrm{Pa}_{G}(X_{i});\theta_{i}\big{)},$
(1)

where $\theta=\\{\theta_{1},\ldots,\theta_{d}\\}$ represents the parameters of
the conditional probability distributions (CPDs) involved in this
factorization. When the structure $G$ is known, a Bayesian Network offers a
representation of the joint distribution that may be convenient for
probabilistic inference (Koller and Friedman, 2009), as well as learning its
parameters $\theta$ using a dataset
${\mathcal{D}}=\\{{\bm{x}}^{(1)},\ldots,{\bm{x}}^{(N)}\\}$, where each
${\bm{x}}^{(j)}$ represents an observation of the $d$ random variables
$\\{X_{1},\ldots,X_{d}\\}$.

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

###  2.2 Generative Flow Networks

A _Generative Flow Network_ (GFlowNet; Bengio et al., 2021, 2023) is a
generative model over a structured sample space ${\mathcal{X}}$. The structure
of the GFlowNet is described by a DAG ${\mathcal{G}}$ whose vertex set is the
_state space_ ${\mathcal{S}}$, where ${\mathcal{X}}\subseteq{\mathcal{S}}$.
This should not be confused with the DAG in a Bayesian Network: in fact, each
state in ${\mathcal{G}}$ could itself represent the structure of a Bayesian
Network (Deleu et al., 2022). A sample $x\in{\mathcal{X}}$ is constructed
sequentially by following the edges of ${\mathcal{G}}$, starting at a special
initial state $s_{0}$, until we reach the state $x$. We define a special
terminal state $s_{f}$, a transition to which indicates the end of the
sequential process. The states in ${\mathcal{X}}$ are those for which there is
a directed edge $x\rightarrow s_{f}$; they are called _complete states_
111Bengio et al. (2023) also call these states _terminating_ ; we follow the
naming conventions of Deleu et al. (2022) here to avoid ambiguity, as it is
closely related to our work. and correspond to valid samples of the
distribution induced by the GFlowNet. A path $s_{0}\rightsquigarrow s_{f}$ in
${\mathcal{G}}$ is called a complete trajectory. An example of the structure
of a GFlowNet is given in Section 3.1 and Figure 1.

Figure 1: Structure of the Generative Flow Network to approximate the joint
posterior distribution $P(G,\theta\mid{\mathcal{D}})$. A graph $G$ and
parameters $\theta$ are constructed as follows: starting from the empty graph
$G_{0}$, (1) the graph $G$ is first generated one edge at a time (blue), as in
(Deleu et al., 2022). Then once we select the action indicating that we stop
adding edges to the graph, (2) we generate the parameters $\theta$ (red),
conditioned on the graph $G$. Finally, given $G$ and $\theta$, (3) we receive
a reward $R(G,\theta)$ (green).

Every complete state $x\in{\mathcal{X}}$ is associated with a _reward_
$R(x)\geq 0$, indicating the unnormalized probability of $x$. We use the
convention $R(s)=0$ for any state $s\in{\mathcal{S}}\backslash{\mathcal{X}}$,
since they do not correspond to valid samples of the distribution. Bengio et
al. (2021) showed that if there exists a function $F_{\phi}(s\rightarrow
s^{\prime})\geq 0$ defined over the edges of ${\mathcal{G}}$, called a _flow_
, that satisfies the following _flow-matching conditions_

$\sum_{\mathclap{s\in\mathrm{Pa}_{{\mathcal{G}}}(s^{\prime})}}\
F_{\phi}(s\rightarrow
s^{\prime})-\sum_{\mathclap{s^{\prime\prime}\in\mathrm{Ch}_{{\mathcal{G}}}(s^{\prime})}}\
F_{\phi}(s^{\prime}\rightarrow s^{\prime\prime})=R(s^{\prime})$ (2)

for all the non-terminal states $s^{\prime}\in{\mathcal{S}}$, then the
GFlowNet induces a distribution over complete states proportional to the
reward. More precisely, starting from the initial state $s_{0}$, if we sample
a complete trajectory $(s_{0},s_{1},\ldots,s_{T-1},x,s_{f})$ following the
forward transition probability, defined as

$P(s_{t+1}\mid s_{t})\propto F_{\phi}(s_{t}\rightarrow s_{t+1}),$ (3)

with the conventions $s_{T}=x$ and $s_{T+1}=s_{f}$, then the marginal
probability that a trajectory sampled following $P$ terminates in
$x\in{\mathcal{X}}$ is proportional to $R(x)$. The flow function
$F_{\phi}(s\rightarrow s^{\prime})$ may be parametrized by a neural network
and optimized to minimize the error in (2), yielding a transition model that
can be used to approximately sample from the distribution on ${\mathcal{X}}$
proportional to $R$.

###  2.3 Structure learning with GFlowNets

Since GFlowNets are particularly well-suited to specifying distributions over
composite objects, Deleu et al. (2022) used this framework in the context of
Bayesian structure learning to approximate the (marginal) posterior
distribution over DAGs $P(G\mid{\mathcal{D}})$. Their model, called _DAG-
GFlowNet_ , operates on the state-space of DAGs, where each graph is
constructed sequentially by adding one edge at a time, starting from the empty
graph with $d$ nodes, while enforcing the acyclicity constraint at every step
of the generation (i.e., an edge is not added if it would introduce a cycle).
Its structure is illustrated at the top of Figure 1, where it forms the part
of the graph shown in blue.

Instead of working with flows $F_{\phi}(G\rightarrow G^{\prime})$, as in
Section 2.2, DAG-GFlowNet directly learns the forward transition probability
$P_{\phi}(G^{\prime}\mid G)$ that satisfies the following alternative
_detailed balance_ conditions for any transition $G\rightarrow G^{\prime}$
(i.e., $G^{\prime}$ is the result of adding a single edge to $G$):

$R(G^{\prime})P_{B}(G\mid G^{\prime})P_{\phi}(s_{f}\mid
G)=R(G)P_{\phi}(G^{\prime}\mid G)P_{\phi}(s_{f}\mid G^{\prime}),$ (4)

where $P_{B}(G\mid G^{\prime})$ is a fixed distribution over the parent states
of $G^{\prime}$ (e.g., uniform distribution over parents). Deleu et al. (2022)
showed that since all the states are complete here, satisfying the conditions
Equation 4 for all $G\rightarrow G^{\prime}$ still induces a distribution over
DAGs $\propto R(G)$. Therefore, to approximate the posterior distribution
$P(G\mid{\mathcal{D}})$, they used $R(G)=P({\mathcal{D}}\mid G)P(G)$ as the
reward of $G$. In particular, this requires evaluating the marginal likelihood
$P({\mathcal{D}}\mid G)$ efficiently, which is feasible only for limited
classes of models (e.g., linear Gaussian; Geiger and Heckerman, 1994).

##  3 Joint Bayesian inference of structure and parameters

Although Generative Flow Networks have been primarily applied to model
distributions over discrete objects such as DAGs, Lahlou et al. (2023) showed
that similar ideas could also be applied to continuous objects, and discrete-
continuous hybrids. Building on top of DAG-GFlowNet, we propose here to
approximate the joint posterior $P(G,\theta\mid{\mathcal{D}})$ over both the
structure of the Bayesian Network $G$, but also the parameters of its
conditional probability distributions $\theta$. Unlike in VBG though
(Nishikawa-Toomey et al., 2023), we use a single GFlowNet to approximate this
joint posterior. We call this model _JSP-GFN_ , for Joint Structure and
Parameters Bayesian inference with a GFlowNet.

###  3.1 Structure of the GFlowNet

Unlike in DAG-GFlowNet, where we model a distribution only over DAGs, here we
need to define a GFlowNet whose complete states are pairs $(G,\theta)$, where
$G$ is a DAG and $\theta$ is a set of (continuous-valued) parameters whose
dimension may depend on $G$. Complete states are obtained through two phases
(Figure 1): the DAG $G$ is first constructed one edge at a time, following
Deleu et al. (2022), and then the corresponding parameters $\theta$ are
generated, conditioned on $G$. We denote by $(G,\cdot)$ states where the DAG
$G$ has no parameters $\theta$ associated to it (states in blue in Figure 1);
they are intermediate states during the first phase of the GFlowNet, and do
not correspond to valid samples of the induced distribution. Using the
notations of Section 2.2, $(G,\theta)\in{\mathcal{X}}$, whereas
$(G,\cdot)\in{\mathcal{S}}\backslash{\mathcal{X}}$.

Starting at the empty graph $(G_{0},\cdot)$, the DAG is constructed one edge
at a time during the first phase, following the forward transition
probabilities $P_{\phi}(G^{\prime}\mid G)$. This first phase ends when a
special “stop” action is selected with $P_{\phi}$, indicating that we stop
adding edges to the graph; the role of this “stop” action is detailed in
Section 3.4. Then during the second phase, we generate $\theta$ conditioned on
$G$, following the forward transition probabilities $P_{\phi}(\theta\mid
G)$.222Since the states of the GFlowNet here are pairs of objects,
$P_{\phi}(\theta\mid G)$ (resp. $P_{\phi}(G^{\prime}\mid G)$) is an abuse of
notation, and represents $P_{\phi}((G,\theta)\mid(G,\cdot))$ (resp.
$P_{\phi}((G^{\prime},\cdot)\mid(G,\cdot))$). All the complete states
$(G,\theta)$, for a fixed graph $G$ and any set of parameters $\theta$, can be
seen as forming an (infinitely wide) tree rooted at $(G,\cdot)$.

Since we want this GFlowNet to approximate the joint posterior
$P(G,\theta\mid{\mathcal{D}})\propto P({\mathcal{D}},\theta,G)$, it is natural
to define the reward function of a complete state $(G,\theta)$ as

$R(G,\theta)=P({\mathcal{D}}\mid\theta,G)P(\theta\mid G)P(G),$ (5)

where the likelihood model $P({\mathcal{D}}\mid\theta,G)$ may be arbitrary
(e.g., a neural network), and decomposes according to Equation 1,
$P(\theta\mid G)$ is the prior over parameters, and $P(G)$ the prior over
graphs. When the dataset ${\mathcal{D}}$ is large, we can use a mini-batch
approximation to the reward (see Section 3.5).

###  3.2 Subtrajectory Balance conditions

To obtain a generative process that samples pairs of $(G,\theta)$
proportionally to the reward, the GFlowNet needs to satisfy some conditions
such as Equation 4. However, we saw in Section 2.3 that satisfying this
particular formulation of the detailed balance conditions in Equation 4 yields
a distribution $\propto R(\cdot)$ only if all the states are complete;
unfortunately, this is not the case here since there exists states of the form
$(G,\cdot)$ corresponding to graphs without their associated parameters.
Instead, we use a generalization of detailed balance to undirected paths of
arbitrary length, called the _Subtrajectory Balance_ conditions (SubTB; Malkin
et al., 2022); we give a brief overview of SubTB in Section C.2.

More precisely, we consider SubTB for any undirected path of the form
$(G,\theta)\leftarrow(G,\cdot)\rightarrow(G^{\prime},\cdot)\rightarrow(G^{\prime},\theta^{\prime})$
in the GFlowNet (see Figure C.2), where $G^{\prime}$ is therefore the result
of adding a single edge to $G$. Since both ends of these undirected paths of
length 3 are complete states, we show in Section C.3.1 that the SubTB
conditions corresponding to undirected paths of this form can be written as

$R(G^{\prime},\theta^{\prime})P_{B}(G\mid G^{\prime})P_{\phi}(\theta\mid
G)=R(G,\theta)P_{\phi}(G^{\prime}\mid G)P_{\phi}(\theta^{\prime}\mid
G^{\prime}).$ (6)

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

Although there is no guarantee in general that satisfying the SubTB conditions
would yield a distribution proportional to the reward, unlike with the
detailed balance conditions (Bengio et al., 2023), the following theorem shows
that the GFlowNet does induce a distribution $\propto R(G,\theta)$ if the
SubTB conditions in Equation 6 are satisfied for all pairs $(G,\theta)$ and
$(G^{\prime},\theta^{\prime})$.

######  Theorem 3.1.

If the SubTB conditions in Equation 6 are satisfied for all undirected paths
of length 3 between any $(G,\theta)$ and $(G^{\prime},\theta^{\prime})$ of the
form
$(G,\theta)\leftarrow(G,\cdot)\rightarrow(G^{\prime},\cdot)\rightarrow(G^{\prime},\theta^{\prime})$,
then we have

$P_{\phi}^{\top}(G,\theta)\triangleq P_{\phi}(G\mid G_{0})P_{\phi}(\theta\mid
G)\propto R(G,\theta),$

where $P_{\phi}(G\mid G_{0})$ is the marginal probability of reaching $G$ from
the initial state $G_{0}$ with any (complete) trajectory
$\tau=(G_{0},G_{1},\ldots,G_{T-1},G)$:

$P_{\phi}(G\mid G_{0})\triangleq\sum_{\tau:G_{0}\rightsquigarrow
G}\prod_{t=0}^{T-1}P_{\phi}(G_{t+1}\mid G_{t}),$

using the conventions $G_{T}=G$, and $P_{\phi}(G_{0}\mid G_{0})=1$.

The proof of this theorem is available in Section C.3.3. The marginal
distribution $P_{\phi}^{\top}(G,\theta)$ is also called the _terminating state
probability_ in Bengio et al. (2023).

###  3.3 Learning objective

One way to find the parameters $\phi$ of the forward transition probabilities
that enforce the SubTB conditions in Equation 6 for all $(G,\theta)$ and
$(G^{\prime},\theta^{\prime})$ is to transform this condition into a learning
objective. For example, we could minimize a non-linear least squares objective
(Bengio et al., 2021, 2023) of the form
${\mathcal{L}}(\phi)=\mathbb{E}_{\pi}[\Delta^{2}(\phi)]$, where the residuals
$\Delta(\phi)$ depend on the conditions in Equation 6, and $\pi$ is an
arbitrary sampling distribution of complete states $(G,\theta)$ and
$(G^{\prime},\theta^{\prime})$, with full support; see Malkin et al. (2023)
for a discussion of the effect of $\pi$ on training GFlowNets, and Section D.1
for further details in the context of JSP-GFN.

In addition to the SubTB conditions for undirected paths of length 3 given in
Section 3.2, we can also derive similar SubTB conditions for other undirected
paths, for example those of the form
$(G,\theta)\leftarrow(G,\cdot)\rightarrow(G,\tilde{\theta})$, where $\theta$
and $\tilde{\theta}$ are two possible sets of parameters associated with the
same graph $G$ (see also Figure C.2). When the reward function $R(G,\theta)$
is differentiable wrt. $\theta$, which is the case here, we show in Section
C.3.2 that we can write the SubTB conditions for these undirected paths of
length 2 in differential form as

$\nabla_{\theta}\log P_{\phi}(\theta\mid G)=\nabla_{\theta}\log R(G,\theta).$
(7)

This condition is equivalent to the notion of score matching (Hyvärinen, 2005)
to model unnormalized distributions, since $R(G,\theta)$ here corresponds to
the unnormalized posterior distribution $P(\theta\mid G,{\mathcal{D}})$, where
the normalization is over $\theta$. We also show in Section C.3.2 that
incorporating this information about undirected paths of length 2, via the
identity in Equation 7, amounts to _preventing_ backpropagation through
$\theta$ and $\theta^{\prime}$ (e.g., backpropagation with the
reparametrization trick) in the objective

${\mathcal{L}}(\phi)=\mathbb{E}_{\pi}\\!\left[\left(\log\frac{R\big{(}G^{\prime},\bot(\theta^{\prime})\big{)}P_{B}(G\mid
G^{\prime})P_{\phi}\big{(}\bot(\theta)\mid
G\big{)}}{R\big{(}G,\bot(\theta)\big{)}P_{\phi}(G^{\prime}\mid
G)P_{\phi}\big{(}\bot(\theta^{\prime})\mid
G^{\prime}\big{)}}\right)^{2}\right],$ (8)

where $\bot$ 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,\theta)$ follows two
phases: first we construct $G$ one edge at a time, until we sample a specific
“stop” action, at which point we sample $\theta$, conditioned on $G$. All
these actions are sampled using the forward transition probabilities
$P_{\phi}(G^{\prime}\mid G)$ during the first phase, and $P_{\phi}(\theta\mid
G)$ during the second one. Following Deleu et al. (2022), we parametrize these
forward transition probabilities using a hierarchical model: we first decide
whether we want to stop the first phase or not, with probability
$P_{\phi}(\mathrm{stop}\mid G)$; then, conditioned on this first decision, we
either continue adding an edge to $G$ to reach $G^{\prime}$ with probability
$P_{\phi}(G^{\prime}\mid G,\neg\mathrm{stop})$ (phase 1), or sample $\theta$
with probability $P_{\phi}(\theta\mid G,\mathrm{stop})$ (phase 2). This
hierarchical model can be written as

$\displaystyle P_{\phi}(G^{\prime}\mid G)$
$\displaystyle=\big{(}1-P_{\phi}(\mathrm{stop}\mid
G)\big{)}P_{\phi}(G^{\prime}\mid G,\neg\mathrm{stop})$ (9) $\displaystyle
P_{\phi}(\theta\mid G)$ $\displaystyle=P_{\phi}(\mathrm{stop}\mid
G)P_{\phi}(\theta\mid G,\mathrm{stop}).$ (10)

We use neural networks to parametrize each of the three components necessary
to define the forward transition probabilities. Unlike in DAG-GFlowNet though,
which uses a linear Transformer to define $P_{\phi}(\mathrm{stop}\mid G)$ and
$P_{\phi}(G^{\prime}\mid G,\neg\mathrm{stop})$, we use a combination of graph
network (Battaglia et al., 2018) and self-attention blocks (Vaswani et al.,
2017) to encode information about the graph $G$, which appears in the
conditioning of all the quantities of interest. This common backbone returns a
graph embedding ${\bm{g}}$ of $G$, as well as 3 embeddings
${\bm{u}}_{i},{\bm{v}}_{i},{\bm{w}}_{i}$ for each node $X_{i}$ in $G$

${\bm{g}},\\{{\bm{u}}_{i},{\bm{v}}_{i},{\bm{w}}_{i}\\}_{i=1}^{d}=\mathrm{SelfAttention}_{\phi}\big{(}\mathrm{GraphNet}_{\phi}(G)\big{)}.$

We can parametrize the probability of selecting the “stop” action using
${\bm{g}}$ with $P_{\phi}(\mathrm{stop}\mid G)=f_{\phi}({\bm{g}})$, where
$f_{\phi}$ is a neural network with a sigmoid output; note that if we can’t
add any edge to $G$ without creating a cycle, we force the end of the first
phase by setting $P_{\phi}(\mathrm{stop}\mid G)=1$. Inspired by Lorch et al.
(2021), the probability of moving from $G$ to $G^{\prime}$ by adding the edge
$X_{i}\rightarrow X_{j}$ is parametrized by

$P_{\phi}(G^{\prime}\mid
G,\neg\mathrm{stop})\propto{\bm{m}}_{ij}\exp\\!\big{(}{\bm{u}}_{i}^{\top}{\bm{v}}_{j}\big{)},$
(11)

where ${\bm{m}}_{ij}$ is a binary mask indicating whether adding
$X_{i}\rightarrow X_{j}$ is a valid action (i.e., if it is not already present
in $G$, and if it doesn’t introduce a cycle; Deleu et al., 2022). Finally, the
probability of selecting the parameters $\theta_{i}$ of the CPD for the
variable $X_{i}$ is parametrized with a multivariate Normal distribution with
diagonal covariance (unless specified otherwise)

$P_{\phi}(\theta_{i}\mid
G,\mathrm{stop})={\mathcal{N}}\big{(}\theta_{i}\mid\bm{\mu}_{\phi}({\bm{w}}_{i}),\bm{\sigma}_{\phi}^{2}({\bm{w}}_{i})\big{)},$
(12)

where $\bm{\mu}_{\phi}$ and $\bm{\sigma}_{\phi}^{2}$ are two neural networks,
with appropriate non-linearities to guarantee that
$\bm{\sigma}_{\phi}^{2}({\bm{w}}_{i})$ is a well-defined diagonal covariance
matrix. Note that $P_{\phi}(\theta_{i}\mid G,\mathrm{stop})$ effectively
approximates the posterior distribution $P(\theta_{i}\mid G,{\mathcal{D}})$
once fully trained. Moreover, in addition to being an approximation of the
joint posterior $P(G,\theta\mid{\mathcal{D}})$, the GFlowNet also provides an
approximation of the marginal posterior $P(G\mid{\mathcal{D}})$, by only
following the first phase of the generation process (to generate $G$) until
the “stop” action is selected, and not continuing into the generation of
$\theta$.

###  3.5 Mini-batch training

Throughout the paper, we have assumed that we had access to the full dataset
of observations ${\mathcal{D}}$ in order to compute the reward $R(G,\theta)$
in Equation 5. However, beyond the capacity to have an arbitrary likelihood
model $P({\mathcal{D}}\mid\theta,G)$ (e.g., non-linear), another advantage of
approximating the joint posterior $P(G,\theta\mid{\mathcal{D}})$ is that we
can train the GFlowNet using mini-batches of observations. Concretely, for a
mini-batch ${\mathcal{B}}$ of $M$ observations sampled uniformly at random
from the dataset ${\mathcal{D}}$, we can define

$\log\widehat{R}_{{\mathcal{B}}}(G,\theta)=\log P(\theta\mid G)+\log
P(G)+\frac{N}{M}\ \sum_{\mathclap{{\bm{x}}^{(m)}\in{\mathcal{B}}}}\ \log
P({\bm{x}}^{(m)}\mid G,\theta),$ (13)

which is an unbiased estimate of the log-reward. The following proposition
shows that minimizing the estimated loss based on Equation 13 wrt. the
parameters $\phi$ of the GFlowNet also minimizes the original objective in
Section 3.3.

######  Proposition 3.2.

Suppose that ${\mathcal{B}}$ is a mini-batch of $M$ observations sampled
uniformly at random from the dataset ${\mathcal{D}}$, and let
$\widehat{{\mathcal{L}}}_{{\mathcal{B}}}(\phi)$ be the learning objective
defined in Section 3.3, where the reward has been replaced by the estimate
$\widehat{R}_{B}(G,\theta)$ in Equation 13. Then we have
${\mathcal{L}}(\phi)\leq\mathbb{E}_{{\mathcal{B}}}\big{[}\widehat{{\mathcal{L}}}_{{\mathcal{B}}}(\phi)\big{]}$.

The proof is available in Section 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\mid{\mathcal{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 $\theta$ lives depends on the graph structure $G$,
approximating the joint posterior $P(G,\theta\mid{\mathcal{D}})$ using MCMC
requires additional trans-dimensional updates (Fronk, 2002), and has therefore
received less attention than the marginal case.

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

##### Generative Flow Networks.

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

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

##  5 Experimental results

###  5.1 Joint posterior over small graphs

Figure 2: Comparison with the exact posterior distribution, on small graphs
with $d=5$ nodes. (a) Comparison of the edge features computed with the exact
posterior (x-axis) and the approximation given by JSP-GFN (y-axis); each point
corresponds to an edge $X_{i}\rightarrow X_{j}$ for each of the 20 datasets.
(b) Quantitative evaluation of different methods for joint posterior
approximation, both in terms of edge features and cross-entropy of sampling
distribution and true posterior $P(\theta\mid G,{\mathcal{D}})$; all values
correspond to the mean and $95\%$ confidence interval across the 20
experiments.

We can evaluate the accuracy of the approximation returned by JSP-GFN by
comparing it with the exact joint posterior distribution
$P(G,\theta\mid{\mathcal{D}})$. Computing this exact posterior can only be
done in limited cases only, where (1) $P(\theta\mid G,{\mathcal{D}})$ can be
computed analytically, and (2) for a small enough $d$ such that all the DAGs
can be enumerated in order to compute $P(G\mid{\mathcal{D}})$. We consider
here models over $d=5$ variables, with linear Gaussian CPDs. We generate 20
different datasets of $N=100$ observations from randomly generated Bayesian
Networks. Details about data generation and modeling are available in Section
D.2.

The quality of the joint posterior approximations is evaluated separately for
$G$ and $\theta$. For the graphs, we compare the approximation and the exact
posterior on different 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 Section 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 Section D.2.2.

To evaluate the performance of the different methods as an approximation of
the posterior over $\theta$, we also estimate the cross-entropy between the
sampling distribution of $\theta$ given $G$ and the exact posterior
$P(\theta\mid G,{\mathcal{D}})$. This measure will be minimized if the model
correctly samples parameters from the true $P(\theta\mid G,{\mathcal{D}})$;
details about this metric are given in Section D.2.3. In Figure 2 (b), we
observe that again both versions of JSP-GFN sample parameters $\theta$ that
are significantly more probable under the exact posterior compared to other
methods.

###  5.2 Gaussian Bayesian Networks from simulated data

To evaluate whether our observations hold on larger graphs, we also evaluated
the performance of JSP-GFN on data simulated from larger Gaussian Bayesian
Networks, with $d=20$ variables. In addition to linear CPDs, as in Section
5.1, we experimented with non-linear Gaussian Bayesian Networks, where the
CPDs are parametrized by neural networks. Following Lorch et al. (2021), we
parametrized the CPDs of each variable with a 2-layer MLP, for a total of
$|\theta|=2,220$ parameters. For both experimental settings, we used datasets
of $N=100$ observations simulated from (randomly generated) Bayesian Networks;
additional experimental details are provided in Section D.3.

Figure 3: Evaluation of JSP-GFN on Gaussian Bayesian Networks. (a-b)
Comparison of the negative log-likelihood (NLL) on $N^{\prime}=100$ held-out
observations for different Bayesian structure learning methods, aggregated
across 20 experiments on different datasets ${\mathcal{D}}$. (c) Linear
correlation between the log-reward (x-axis) and the terminating state log-
probability (y-axis) for $1,000$ samples $(G,\theta)$ from JSP-GFN; the color
of each point indicates the number of edges in the correponding graph.
$\Delta$ represents the slope of a linear function fitted using RANSAC
(Fischler and Bolles, 1981).

We compared JSP-GFN against two methods based on MCMC (MH-MC3 & Gibbs-MC3;
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 Section D.3.3. Nevertheless, we also report the expected Structural Hamming
Distance (SHD), as well as the area under the ROC curve (AUROC) in Section
D.3.3 for completeness. To complement these metrics, and in order to assess
the quality of the approximation of the posterior in the absence of reference
$P(G,\theta\mid{\mathcal{D}})$, we also show in Figure 3 (c) how the
terminating state log-probability $\log P_{\phi}^{\top}(G,\theta)$ of JSP-GFN
correlates with the log-reward for a non-linear Gaussian model. Indeed, as
stated in Theorem 3.1, we should ideally have $\log P_{\phi}^{\top}(G,\theta)$
perfectly correlated with $\log R(G,\theta)$ with slope 1, as

$\log P_{\phi}^{\top}(G,\theta)\approx\log P(G,\theta\mid{\mathcal{D}})=\log
R(G,\theta)-\log P({\mathcal{D}}).$ (14)

We can see that there is indeed a strong linear correlation across multiple
samples $(G,\theta)$ from JSP-GFN, with a slope $\Delta$ close to 1,
suggesting that the GFlowNet is again an accurate approximation of the joint
posterior, _at least_ around the modes it captures. Details about how $\log
P_{\phi}^{\top}(G,\theta)$ is estimated are available in Section 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 Section D.4.

At this scale, using the whole dataset ${\mathcal{D}}$ to evaluate the reward
becomes impractical, especially for non-linear models. Fortunately, we showed
in Section 3.5 that we can use an (unbiased) estimate of the reward, based on
mini-batches of data, in place of $R(G,\theta)$ in the loss function Equation
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).

Figure 4: Performance of JSP-GFN on real-world biological data. (a) Comparison
of the terminating state log-probability $\log P_{\phi}^{\top}(G,\theta)$
returned by JSP-GFN with the log-reward $\log R(G,\theta)$ on a subsample of
$N=100$ datapoints of the flow cytometry dataset ${\mathcal{D}}$. (b) Same
comparison with the full dataset ${\mathcal{D}}$ of size $N=4,200$. (c)
Comparison of JSP-GFN with methods based on MCMC on both flow cytometry data
and gene expression data, in terms of negative (interventional) log-likelihood
on held-out data.

To measure the quality of the posterior approximation returned by JSP-GFN, we
compare in Figure 4 the terminating state log-probability $\log
P_{\phi}^{\top}(G,\theta)$ with the log-reward $\log R(G,\theta)$, similar to
Figure 3. We observe that there is correlation between these two quantities;
unlike in Figure 3 though, we observe that the slope is not close to $1$,
suggesting that JSP-GFN _underestimates_ the probability of $(G,\theta)$. We
also observe that the graphs are “clustered” together; this can be explained
by the fact that the posterior approximation is concentrated at only a few
graphs, since the size of the dataset ${\mathcal{D}}$ is larger. To confirm
this observation, we show in Figure 4 (a) a similar plot on a subsample of
$N=100$ datapoints randomly sampled from ${\mathcal{D}}$, matching the
experimental setting of Section 5.2. In this case, we observe a much closer
linear fit, with a slope closer to $1$.

In addition to the comparison to the log-reward, we also compare in Figure 4
(c) JSP-GFN with 2 methods based on MCMC in terms of their negative log-
likelihood on held-out data. We can see that JSP-GFN is competitive, and even
out-performs MCMC on the more challenging problem of the discovery of gene
regulatory networks from gene expression data, where the dimensionality of the
problem is much larger ($d=61$). Note that the values reported for the
discovery of protein signaling networks from flow cytometry data correspond to
the negative _interventional_ log-likelihood, on interventions unseen in
${\mathcal{D}}$.

##  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.

## Acknowledgements

We would like to thank António Góis, Salem Lahlou, Romain Lopez, Cristian
Dragos Manta, and Mansi Rankawat for the useful discussions about the project
and their feedback on the paper. We would like to also thank Lars Lorch for
releasing the code for DiBS (Lorch et al., 2021), along with useful baselines
used in this paper, as well as Chris Cundy for releasing the code for BCD Nets
and for his help with the code. This research was enabled by compute resources
provided by Mila (mila.quebec).

## References

  * Annadani et al. (2021) 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.
  * Atanackovic et al. (2023) 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.
  * Battaglia et al. (2018) 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.
  * Bengio et al. (2021) 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.
  * Bengio et al. (2023) Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J Hu, Mo Tiwari, and Emmanuel Bengio.  GFlowNet Foundations.  _Journal of Machine Learning Research (JMLR)_ , 2023.
  * Charpentier et al. (2022) Bertrand Charpentier, Simon Kibler, and Stephan Günnemann.  Differentiable DAG Sampling.  _International Conference on Learning Representations_ , 2022.
  * Chickering (2002) David Maxwell Chickering.  Optimal Structure Identification With Greedy Search.  _Journal of Machine Learning Research_ , 2002.
  * Cundy et al. (2021) Chris Cundy, Aditya Grover, and Stefano Ermon.  BCD Nets: Scalable Variational Approaches for Bayesian Causal Discovery.  _Advances in Neural Information Processing Systems_ , 2021.
  * Deleu et al. (2022) Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio.  Bayesian Structure Learning with Generative Flow Networks.  _Uncertainty in Artificial Intelligence_ , 2022.
  * Eaton and Murphy (2007) Daniel Eaton and Kevin Murphy.  Bayesian structure learning using dynamic programming and MCMC.  _Conference on Uncertainty in Artificial Intelligence_ , 2007.
  * Erdős and Rényi (1960) Paul Erdős and Alfréd Rényi.  On the evolution of random graphs.  _Publications of the Mathematical Institute of the Hungarian Academy of Sciences_ , 1960.
  * Fischler and Bolles (1981) 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.
  * Friedman and Koller (2003) Nir Friedman and Daphne Koller.  Being Bayesian About Network Structure. A Bayesian Approach to Structure Discovery in Bayesian Networks.  _Machine Learning_ , 2003.
  * Friedman et al. (1999) 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.
  * Friedman et al. (2000) 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.
  * Fronk (2002) Eva-Maria Fronk.  Model Selection for Dags via RJMCMC for the Discrete and Mixed Case.  Technical report, Ludwig-Maximilians-Universität München, 2002\.
  * Geiger and Heckerman (1994) Dan Geiger and David Heckerman.  Learning Gaussian Networks.  In _Uncertainty in Artificial Intelligence_. 1994.
  * Giudici and Castelo (2003) Paolo Giudici and Robert Castelo.  Improving Markov chain Monte Carlo model search for data mining.  _Machine learning_ , 2003.
  * Heckerman et al. (1995) David Heckerman, Dan Geiger, and David M Chickering.  Learning bayesian networks: The combination of knowledge and statistical data.  _Machine learning_ , 1995.
  * Hu et al. (2023) 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.
  * Hyvärinen (2005) Aapo Hyvärinen.  Estimation of Non-Normalized Statistical Models by Score Matching.  _Journal of Machine Learning Research_ , 2005.
  * Jain et al. (2022) 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.
  * Jain et al. (2023) Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez-Garcia, and Yoshua Bengio.  GFlowNets for AI-Driven Scientific Discovery.  _Digital Discovery_ , 2023.
  * Koller and Friedman (2009) Daphne Koller and Nir Friedman.  _Probabilistic Graphical Models: Principles and Techniques_.  MIT press, 2009.
  * Lahlou et al. (2023) Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin.  A Theory of Continuous Generative Flow Networks.  _International Conference on Machine Learning_ , 2023.
  * Lauritzen and Spiegelhalter (1988) 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.
  * Li et al. (2023) Yinchuan Li, Shuang Luo, Haozhi Wang, and Jianye Hao.  CFlowNets: Continuous control with Generative Flow Networks.  _International Conference on Learning Representations_ , 2023.
  * Liu and Wang (2016) Qiang Liu and Dilin Wang.  Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm.  _Advances in Neural Information Processing Systems_ , 2016.
  * Lorch et al. (2021) Lars Lorch, Jonas Rothfuss, Bernhard Schölkopf, and Andreas Krause.  DiBS: Differentiable Bayesian Structure Learning.  _Advances in Neural Information Processing Systems_ , 2021.
  * Lorch et al. (2022) Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf.  Amortized Inference for Causal Structure Learning.  _Advances in Neural Information Processing Systems_ , 2022.
  * Madan et al. (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.
  * Madigan et al. (1994) David Madigan, Jonathan Gavrin, and Adrian E Raftery.  Enhancing the Predictive Performance of Bayesian Graphical Models.  1994\.
  * Madigan et al. (1995) David Madigan, Jeremy York, and Denis Allard.  Bayesian Graphical Models for Discrete Data.  _International Statistical Review_ , 1995.
  * Malkin et al. (2022) Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio.  Trajectory balance: Improved credit assignment in GFlowNets.  _Neural Information Processing Systems_ , 2022.
  * Malkin et al. (2023) 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.
  * Mooij et al. (2020) Joris M Mooij, Sara Magliacane, and Tom Claassen.  Joint Causal Inference from Multiple Contexts.  _Journal of Machine Learning Research_ , 2020.
  * Nishikawa-Toomey et al. (2023) 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.
  * Pan et al. (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.
  * Sachs et al. (2005) 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.
  * Sethuraman et al. (2023) Muralikrishnna G Sethuraman, Romain Lopez, Rahul Mohan, Faramarz Fekri, Tommaso Biancalani, and Jan-Christian Hütter.  NODAGS-Flow: Nonlinear Cyclic Causal Structure Learning.  _International Conference on Artificial Intelligence and Statistics (AISTATS)_ , 2023.
  * Spirtes et al. (2000) Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman.  _Causation, Prediction, and Search_.  MIT press, 2000.
  * Vaswani et al. (2017) 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.
  * Viinikka et al. (2020) Jussi Viinikka, Antti Hyttinen, Johan Pensar, and Mikko Koivisto.  Towards Scalable Bayesian Learning of Causal DAGs.  _Advances in Neural Information Processing Systems_ , 2020.
  * von Kügelgen et al. (2019) Julius von Kügelgen, Paul K Rubenstein, Bernhard Schölkopf, and Adrian Weller.  Optimal experimental design via Bayesian optimization: active causal structure learning for Gaussian process networks.  _NeurIPS Workshop “Do the right thing”: machine learning and causal inference for improved decision making_ , 2019.
  * Wang et al. (2022) Benjie Wang, Matthew R Wicker, and Marta Kwiatkowska.  Tractable Uncertainty for Structure Learning.  _International Conference on Machine Learning_ , 2022.
  * Zhang et al. (2023) David Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan.  Robust Scheduling with GFlowNets .  _International Conference on Learning Representations_ , 2023.
  * Zhang et al. (2022) 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.
  * Zheng et al. (2018) 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.
  * Zimmermann et al. (2022) Heiko Zimmermann, Fredrik Lindsten, Jan-Willem van de Meent, and Christian A. Naesseth.  A Variational Perspective on Generative Flow Networks.  _arXiv preprint_ , 2022.

Appendix

##  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).

Table A.1: Comparison of different methods based on variational inference and
GFlowNets for Bayesian structure learning. See the text for a detailed
description of each category.

Joint Non DAG Discrete Max. Sampler Mini $G$ & $\theta$ Linear Support Obs.
Parents Batch VCN (Annadani et al., 2021) ✗ ✗ ✗ ✗ ✓ ✗ BCD Nets (Cundy et al.,
2021) ✓ ✗ ✓ ✗ ✗ ✓ ✗ DiBS (Lorch et al., 2021) ✓ ✓ ✗ ✓ ✗ ✗ ✓ Trust (Wang et
al., 2022) ✗ ✗ ✓ ✗ ✗ VI-DP-DAG (Charpentier et al., 2022) ✗ ✓ ✓ ✗ ✗ ✓ ✗ AVICI
(Lorch et al., 2022) ✗ ✓ ✗ ✗ ✓ ✗ DAG-GFlowNet (Deleu et al., 2022) ✗ ✗ ✓ ✓ ✓ ✓
✗ VBG (Nishikawa-Toomey et al., 2023) ✓ ✓ ✓ ✓ ✗ JSP-GFN (Ours) ✓ ✓ ✓ ✓ ✓ ✓ ✓

##### Joint $\bm{G}$ & $\bm{\theta}$.

This category indicates whether the model can approximate the joint posterior
distribution $P(G,\theta\mid{\mathcal{D}})$ over both graphical structures $G$
and parameters of the CPDs $\theta$, or if they are limited to approximating
the marginal posterior $P(G\mid{\mathcal{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ügelgen et al., 2019), we
only consider here methods that explicitly handle non-linearity (e.g., this
eliminates DAG-GFlowNet (Deleu et al., 2022), since the authors only
considered a linear Gaussian and discrete settings). Annadani et al. (2021)
mentioned the extension of VCN to non-linear causal models 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. AVICI
(Lorch et al., 2022) assumes access to a generative model $P({\mathcal{D}}\mid
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.

##### Maximum Parents.

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 ${\mathcal{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 $\theta$ makes all observations in ${\mathcal{D}}$
mutually _dependent_ (conditioned on $G$). See Section C.3.4 for details on
how to use mini-batch training with JSP-GFN.

##  Appendix 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_{\phi}(\theta\mid
G,\mathrm{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(\theta\mid
G,{\mathcal{D}})$ may be highly multimodal, especially when the model is non-
linear (the posterior is Normal when the model is linear Gaussian, see Section
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 $\theta$ leading to the same likelihood function, and under mild
assumptions on the priors $P(\theta\mid G)$ and $P(G)$, they would have the
same posterior probability $P(\theta\mid G,{\mathcal{D}})$.

To address this issue of unimodality, we can use more expressive posterior
approximations $P_{\phi}(\theta\mid G,\mathrm{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 $\theta$.

##### Biological plausibility of the acyclicity assumption.

One of the strengths of JSP-GFN, and DAG-GFlowNet before it (Deleu et al.,
2022), is the capacity to obtain a distribution over the DAG structure of a
Bayesian Network (and its parameters). The acyclicity assumption is
particularly important in order to properly define the likelihood model in
Equation 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 Section 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 ${\bm{m}}$ in Sec. 3.4. Alternatively,
we can view the generation of a cyclic graph by unrolling it, as in
Atanackovic et al. (2023).

##  Appendix C Details about Generative Flow Networks

Throughout this section, we will use both $P_{F}$ and $P_{\phi}$ (to emphasize
the parametrization on $\phi$, 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).

Figure C.1: Illustration of the different GFlowNet objectives. (a) The
detailed balance condition operates at the level of transitions $s\rightarrow
s^{\prime}$, whereas (b) the trajectory balance condition operates on complete
trajectories $s_{0}\rightsquigarrow s_{f}$. (c) The subtrajectory balance
condition operates on partial trajectories $s_{m}\rightsquigarrow s_{n}$, and
can be (d) generalized to undirected paths with a common ancestor $s_{k}$. We
use $P_{F}(s_{m}\rightsquigarrow s_{n})$ to denote the product of $P_{F}$
along the path $s_{m}\rightsquigarrow s_{n}$ (and similarly for $P_{B}$).

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\rightarrow s^{\prime}$ in the
GFlowNet as

$F(s)P_{F}(s^{\prime}\mid s)=F(s^{\prime})P_{B}(s\mid s^{\prime})$ (C.1)

where $F(s)$ is a flow function, that may also be parametrized by a neural
network. The detailed balance condition is illustrated in Figure C.1 (a).
Bengio et al. (2023) showed that if the detailed balance conditions are
satisfied for all the transitions $s\rightarrow s^{\prime}$ 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
$\tau=(s_{0},s_{1},\ldots,s_{T},s_{f})$, the trajectory balance condition is
given by

$Z\prod_{t=1}^{T}P_{F}(s_{t+1}\mid
s_{t})=R(s_{T})\prod_{t=1}^{T-1}P_{B}(s_{t}\mid s_{t+1}),$ (C.2)

with the convention $s_{T+1}=s_{f}$, and where $Z$ is the partition function
of the distribution (i.e., $Z=\sum_{x\in{\mathcal{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 $\tau=(s_{m},s_{m+1},\ldots,s_{n})$, the subtrajectory balance
condition is given by

$F(s_{m})\prod_{t=m}^{n-1}P_{F}(s_{t+1}\mid
s_{t})=F(s_{n})\prod_{t=m}^{n-1}P_{B}(s_{t}\mid s_{t+1}),$ (C.3)

where again $F(s)$ is a flow function (as in Equation C.1). This condition
encompasses both conditions in Section C.1, since we can recover the detailed
balance condition in Equation C.1 with partial trajectories of length 1 (i.e.,
transitions), and also the trajectory balance condition in Equation 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
$\mathrm{SubTB}(\lambda)$, inspired by $\mathrm{TD}(\lambda)$ in the
reinforcement learning literature.

This subtrajectory balance condition in Equation 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

$F(s_{m})\prod_{t=k}^{m-1}P_{B}(s_{t}\mid
s_{t+1})\prod_{t=k}^{n-1}P_{F}(s_{t+1}\mid
s_{t})=F(s_{n})\prod_{t=k}^{n-1}P_{B}(s_{t}\mid
s_{t+1})\prod_{t=k}^{m-1}P_{F}(s_{t+1}\mid s_{t}),$ (C.4)

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 Section C.3.3 that the GFlowNet does induce a
distribution $\propto R(s)$ in our case.

###  C.3 Proofs

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

Figure C.2: Illustration of the undirected paths of length 3 (red) and of
length 2 (blue) considered in this paper, and their corresponding
subtrajectory balance conditions.

Consider an undirected path of length 3 of the form
$(G,\theta)\leftarrow(G,\cdot)\rightarrow(G^{\prime},\cdot)\rightarrow(G^{\prime},\theta^{\prime})$,
where $G^{\prime}$ is the result of adding a new edge to the DAG $G$ (see
Figure C.2). Since the state $(G,\cdot)$ is a common ancestor of both complete
states $(G,\theta)$ and $(G^{\prime},\theta^{\prime})$, we can write the
subtrajectory balance conditions Equation C.4 as

$F(G,\theta)P_{B}(G\mid\theta)P_{F}(G^{\prime}\mid G)P_{F}(\theta^{\prime}\mid
G^{\prime})=F(G^{\prime},\theta^{\prime})P_{B}(G^{\prime}\mid\theta^{\prime})P_{B}(G\mid
G^{\prime})P_{F}(\theta\mid G),$ (C.5)

where we abuse the notation $P_{B}(G\mid\theta)$ again to denote
$P_{B}\big{(}(G,\cdot)\mid(G,\theta)\big{)}$. In fact, since the complete
state $(G,\theta)\in{\mathcal{X}}$ has only a single parent state $(G,\cdot)$,
we necessarily have $P_{B}(G\mid\theta)=1$ (and similarly for
$(G^{\prime},\theta^{\prime})$). Furthermore, we can use the observation from
(Deleu et al., 2022) to write the flow $F(G,\theta)$ of a complete state
$(G,\theta)$ as a function of its reward

$F(G,\theta)=\frac{R(G,\theta)}{P_{F}\big{(}s_{f}\mid(G,\theta)\big{)}}.$
(C.6)

We can simplify Equation C.6 even further by observing that in the GFlowNet
used here, $s_{f}$ is the _only_ child of the complete state
$(G,\theta)\in{\mathcal{X}}$. In other words, a complete state $(G,\theta)$ is
not directly connected to any other $(G,\tilde{\theta})$; this is the
(infinitely wide) tree structure rooted at $(G,\cdot)$ mentioned in Section
3.1. Since $s_{f}$ is the only child of $(G,\theta)$, we then necessarily have
$P_{F}\big{(}s_{f}\mid(G,\theta)\big{)}=1$, and therefore
$F(G,\theta)=R(G,\theta)$. With these simplifications, Equation C.5 becomes

$R(G,\theta)P_{F}(G^{\prime}\mid G)P_{F}(\theta^{\prime}\mid
G^{\prime})=R(G^{\prime},\theta^{\prime})P_{B}(G\mid
G^{\prime})P_{F}(\theta\mid G),$ (C.7)

which is the subtrajectory balance condition in Equation 6.

####  C.3.2 Integrating undirected paths of length 2

Similar to Section C.3.1, we consider here an undirected of length 2 of the
form $(G,\theta)\leftarrow(G,\cdot)\rightarrow(G,\tilde{\theta})$ (see Figure
C.2). Since $(G,\cdot)$ is a common ancestor (a common parent in this case) of
both complete states $(G,\theta)$ and $(G,\tilde{\theta})$, we can write the
subtrajectory balance conditions Equation C.4 as

$F(G,\theta)P_{B}(G\mid\theta)P_{F}(\tilde{\theta}\mid
G)=F(G,\tilde{\theta})P_{B}(G\mid\tilde{\theta})P_{F}(\theta\mid G).$ (C.8)

Using the same simplifications as in Section C.3.1
($P_{B}(G\mid\theta)=P_{B}(G\mid\tilde{\theta})=1$), we get the following
subtrajectory balance conditions for the undirected paths of length 2

$R(G,\theta)P_{F}(\tilde{\theta}\mid G)=R(G,\tilde{\theta})P_{F}(\theta\mid
G).$ (C.9)

Note that these conditions are effectively redundant if the SubTB conditions
over undirected paths of length 3 Equation 6 are satisfied for all possible
pairs of complete states $(G,\theta)$ and $(G^{\prime},\theta^{\prime})$.
Indeed, if we write these conditions between $(G,\theta)$ and
$(G^{\prime},\theta^{\prime})$ on the one hand, and between
$(G,\tilde{\theta})$ and $(G^{\prime},\theta^{\prime})$ on the other hand
(with a fixed $G^{\prime}$ and $\theta^{\prime}$)

$\displaystyle R(G^{\prime},\theta^{\prime})P_{B}(G\mid
G^{\prime})P_{F}(\theta\mid G)$ $\displaystyle=R(G,\theta)P_{F}(G^{\prime}\mid
G)P_{F}(\theta^{\prime}\mid G)$ (C.10) $\displaystyle
R(G^{\prime},\theta^{\prime})P_{B}(G\mid G^{\prime})P_{F}(\tilde{\theta}\mid
G)$ $\displaystyle=R(G,\tilde{\theta})P_{F}(G^{\prime}\mid
G)P_{F}(\theta^{\prime}\mid G),$ (C.11)

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

$\frac{R(G,\theta)}{P_{F}(\theta\mid
G)}=\frac{R(G^{\prime},\theta^{\prime})P_{B}(G\mid
G^{\prime})}{P_{F}(G^{\prime}\mid G)P_{F}(\theta^{\prime}\mid
G^{\prime})}=\frac{R(G,\tilde{\theta})}{P_{F}(\tilde{\theta}\mid G)}.$ (C.12)

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

$f_{G}(\theta)\triangleq\log R(G,\theta)-\log P_{F}(\theta\mid G)$ (C.13)

is constant, albeit with a constant that depends on the graph $G$. Since this
function is differentiable, this is equivalent to
$\nabla_{\theta}f_{G}(\theta)=0$, and therefore we get the differential form
of the subtrajectory balance conditions in Equation 7

$\nabla_{\theta}\log P_{F}(\theta\mid G)=\nabla_{\theta}\log R(G,\theta).$
(C.14)

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
${\mathcal{L}}(\phi)=\mathbb{E}_{\pi}[\tilde{\Delta}^{2}(\phi)]$, where
$\tilde{\Delta}(\phi)$ is a non-linear residual term

$\tilde{\Delta}(\phi)=\log\frac{R(G^{\prime},\theta^{\prime})P_{B}(G\mid
G^{\prime})P_{\phi}(\theta\mid G)}{R(G,\theta)P_{\phi}(G^{\prime}\mid
G)P_{\phi}(\theta^{\prime}\mid G^{\prime})}.$ (C.15)

Suppose that the parameters $\phi$ of the GFlowNet are such that the
subtrajectory balance conditions in Equation C.14 are satisfied for any
$(G,\theta)$. 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 Equation
C.9 and Equation C.7. Since $\theta$ and $\theta^{\prime}$ depend on $\phi$
(via the reparametrization trick since they are sampled on-policy, see Section
5), taking the derivative of $\tilde{\Delta}^{2}(\phi)$, we get

$\displaystyle\frac{d}{d\phi}\tilde{\Delta}^{2}(\phi)=\tilde{\Delta}(\phi)\cdot\frac{d}{d\phi}\big{[}\log
R(G^{\prime},\theta^{\prime})$ $\displaystyle+\log P_{\phi}(\theta\mid G)$
(C.16) $\displaystyle-\log R(G,\theta)-\log P_{\phi}(G^{\prime}\mid G)-\log
P_{\phi}(\theta^{\prime}\mid G^{\prime})\big{]}.$

Using the law of total derivatives, we have

$\displaystyle\frac{d}{d\phi}\big{[}\log P_{\phi}(\theta\mid G)-\log
R(G,\theta)\big{]}$
$\displaystyle=\underbrace{\left[\frac{\partial}{\partial\theta}\log
P_{\phi}(\theta\mid G)-\frac{\partial}{\partial\theta}\log
R(G,\theta)\right]}_{=\,0}\frac{d\theta}{d\phi}+\frac{\partial}{\partial\phi}\log
P_{\phi}(\theta\mid G)$ (C.17)
$\displaystyle=\frac{\partial}{\partial\phi}\log P_{\phi}(\theta\mid G),$
(C.18)

and similarly for the terms in $(G^{\prime},\theta^{\prime})$. The derivative
of the objective then becomes

$\frac{d}{d\phi}\tilde{\Delta}^{2}(\phi)=\tilde{\Delta}(\phi)\left[\frac{\partial}{\partial\phi}\log
P_{\phi}(\theta\mid G)-\frac{\partial}{\partial\phi}\log
P_{\phi}(\theta^{\prime}\mid G^{\prime})-\frac{d}{d\phi}\log
P_{\phi}(G^{\prime}\mid G)\right].$ (C.19)

An alternative way to obtain the same derivative in Equation C.17 is to take
$d\theta/d\phi=0$ instead, meaning that we would not differentiate through
$\theta$ (and $\theta^{\prime}$). Using the stop-gradient operation $\bot$,
this shows that the following objective

${\mathcal{L}}(\phi)\triangleq\mathbb{E}_{\pi}\big{[}\Delta(\phi)^{2}\big{]}=\mathbb{E}_{\pi}\left[\left(\log\frac{R\big{(}G^{\prime},\bot(\theta^{\prime})\big{)}P_{B}(G\mid
G^{\prime})P_{\phi}\big{(}\bot(\theta)\mid
G\big{)}}{R\big{(}G,\bot(\theta)\big{)}P_{\phi}(G^{\prime}\mid
G)P_{\phi}\big{(}\bot(\theta^{\prime})\mid
G^{\prime}\big{)}}\right)^{2}\right]$ (C.20)

takes the same value and has the same gradient Equation C.19 as the objective
in Equation C.15 when the subtrajectory balance conditions (in differential
form) over undirected paths of length 2 are satisfied.

While optimizing Equation 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

$\displaystyle\tilde{{\mathcal{L}}}(\phi)={\mathcal{L}}(\phi)+\frac{\lambda}{2}\mathbb{E}_{\pi}\Big{[}\big{\|}\nabla_{\theta}\log
P_{\phi}(\theta\mid G)$ $\displaystyle-\nabla_{\theta}\log
R(G,\theta)\big{\|}^{2}$ (C.21)
$\displaystyle+\big{\|}\nabla_{\theta^{\prime}}\log
P_{\phi}(\theta^{\prime}\mid G^{\prime})-\nabla_{\theta^{\prime}}\log
R(G^{\prime},\theta^{\prime})\big{\|}^{2}\Big{]}$

####  C.3.3 Marginal distribution over complete states

See 3.1

###### Proof.

We assume that the SubTB conditions are satisfied for all undirected paths of
length 3 between any $(G,\theta)$ and $(G^{\prime},\theta^{\prime})$, that is

$R(G^{\prime},\theta^{\prime})P_{B}(G\mid G^{\prime})P_{\phi}(\theta\mid
G)=R(G,\theta)P_{\phi}(G^{\prime}\mid G)P_{\phi}(\theta^{\prime}\mid
G^{\prime}).$ (C.22)

Let $G\neq G_{0}$ be a fixed DAG different from the initial state, and
$\theta$ a set of corresponding parameters. Let
$\tau=(G_{0},\ldots,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 $\theta_{t}$ is
a fixed set of parameters associated with $G_{t}$, then the SubTB conditions
above can written for every timestep as

$R(G_{t+1},\theta_{t+1})P_{B}(G_{t}\mid G_{t+1})P_{\phi}(\theta_{t}\mid
G_{t})=R(G_{t},\theta_{t})P_{\phi}(G_{t+1}\mid G_{t})P_{\phi}(\theta_{t+1}\mid
G_{t+1}),$ (C.23)

again, using the convention $\theta_{T}=\theta$. Taking the product of the
ratio between $P_{\phi}$ and $P_{B}$ over the trajectory $\tau$, we get

$\displaystyle\prod_{t=0}^{T-1}\frac{P_{\phi}(G_{t+1}\mid
G_{t})}{P_{B}(G_{t}\mid G_{t+1})}$
$\displaystyle=\prod_{t=0}^{T-1}\frac{P_{\phi}(\theta_{t}\mid
G_{t})R(G_{t+1},\theta_{t+1})}{R(G_{t},\theta_{t})P_{\phi}(\theta_{t+1}\mid
G_{t+1})}$ (C.24) $\displaystyle=\frac{P_{\phi}(\theta_{0}\mid
G_{0})R(G,\theta)}{R(G_{0},\theta_{0})P_{\phi}(\theta\mid G)}$ (C.25)

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

$\sum_{\tau:G_{0}\rightsquigarrow G}\prod_{t=0}^{T-1}P_{B}(G_{t}\mid
G_{t+1})=1.$ (C.26)

Therefore, we have

$\displaystyle P_{\phi}(G\mid G_{0})P_{\phi}(\theta\mid G)$
$\displaystyle=P_{\phi}(\theta\mid G)\left(\sum_{\tau:G_{0}\rightsquigarrow
G}\prod_{t=0}^{T-1}P_{\phi}(G_{t+1}\mid G_{t})\right)$ (C.27)
$\displaystyle=\frac{P_{\phi}(\theta_{0}\mid
G_{0})}{R(G_{0},\theta_{0})}R(G,\theta)\left(\sum_{\tau:G_{0}\rightsquigarrow
G}\prod_{t=0}^{T-1}P_{B}(G_{t}\mid G_{t+1})\right)$ (C.28)
$\displaystyle=\frac{P_{\phi}(\theta_{0}\mid
G_{0})}{R(G_{0},\theta_{0})}R(G,\theta)$ (C.29)

We saw in Section C.3.2 that $P_{\phi}(\theta_{0}\mid
G_{0})/R(G_{0},\theta_{0})$ is independent of the value of $\theta_{0}$ if the
SubTB conditions are satisfied for all undirected paths of length 3 (see
Equation C.12). This concludes the proof: $P_{\phi}(G\mid
G_{0})P_{\phi}(\theta\mid G)\propto R(G,\theta)$. ∎

####  C.3.4 Mini-batch training

See 3.2

###### Proof.

We will first show that $\log\widehat{R}_{{\mathcal{B}}}(G,\theta)$ defined in
Equation 13 is an unbiased estimate of the log-reward $\log R(G,\theta)$ under
a uniform distribution of the mini-batches ${\mathcal{B}}$. We can observe
that by conditional independence of the observations ${\bm{x}}^{(n)}$ given
$G$ and $\theta$, we have

$\log P({\mathcal{D}}\mid\theta,G)=\sum_{n=1}^{N}\log
P({\bm{x}}^{(n)}\mid\theta,G)=N\mathbb{E}_{{\bm{x}}}\big{[}\log
P({\bm{x}}\mid\theta,G)\big{]},$ (C.30)

where the expectation is taken wrt. the uniform distribution over the
observations in ${\mathcal{D}}$. It is important to note that we can decompose
the likelihood term as in Equation C.30 because the observations are mutually
independent given $G$ _and_ $\theta$; 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

$\mathbb{E}_{{\mathcal{B}}}\left[\sum_{{\bm{x}}^{(m)}\in{\mathcal{B}}}\log
P({\bm{x}}^{(m)}\mid\theta,G)\right]=M\mathbb{E}_{{\bm{x}}}\big{[}\log
P({\bm{x}}\mid\theta,G)\big{]}.$ (C.31)

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

$\displaystyle\mathbb{E}_{{\mathcal{B}}}\big{[}\log\widehat{R}_{{\mathcal{B}}}(G,\theta)\big{]}$
$\displaystyle=\frac{N}{M}\mathbb{E}_{{\mathcal{B}}}\left[\sum_{{\bm{x}}^{(m)}\in{\mathcal{B}}}\log
P({\bm{x}}^{(m)}\mid\theta,G)\right]+\log P(\theta\mid G)+\log P(G)$ (C.32)
$\displaystyle=\log P({\mathcal{D}}\mid\theta,G)+\log P(\theta\mid G)+\log
P(G)=\log R(G,\theta).$ (C.33)

Recall that the estimate of the loss is defined as
$\widehat{{\mathcal{L}}}_{{\mathcal{B}}}(\phi)=\mathbb{E}_{\pi}\big{[}\widehat{\Delta}_{{\mathcal{B}}}^{2}(\phi)\big{]}$,
where the residual is defined by

$\widehat{\Delta}_{{\mathcal{B}}}(\phi)=\log\frac{\widehat{R}_{{\mathcal{B}}}\big{(}G,\bot(\theta)\big{)}P_{B}(G\mid
G^{\prime})P_{\phi}\big{(}\bot(\theta)\mid
G\big{)}}{\widehat{R}_{{\mathcal{B}}}\big{(}G^{\prime},\bot(\theta^{\prime})\big{)}P_{\phi}(G^{\prime}\mid
G)P_{\phi}\big{(}\bot(\theta^{\prime})\mid G^{\prime}\big{)}}.$ (C.34)

Taking the expectation of this estimated loss wrt. a random mini-batch
${\mathcal{B}}$, we get

$\displaystyle\mathbb{E}_{{\mathcal{B}}}\big{[}\widehat{{\mathcal{L}}}_{{\mathcal{B}}}(\phi)\big{]}$
$\displaystyle=\mathbb{E}_{\pi}\big{[}\mathbb{E}_{{\mathcal{B}}}\big{[}\widehat{\Delta}_{{\mathcal{B}}}^{2}(\phi)\big{]}\big{]}$
(C.35)
$\displaystyle\geq\mathbb{E}_{\pi}\big{[}\mathbb{E}_{{\mathcal{B}}}\big{[}\widehat{\Delta}_{{\mathcal{B}}}(\phi)\big{]}^{2}\big{]}$
(C.36) $\displaystyle=\mathbb{E}_{\pi}\big{[}\Delta^{2}(\phi)\big{]}$ (C.37)
$\displaystyle={\mathcal{L}}(\phi),$ (C.38)

where we used the convexity of the square function and Jensen’s inequality in
Equation C.36, and the unbiasedness of
$\log\widehat{R}_{{\mathcal{B}}}(G,\theta)$ (as well as
$\log\widehat{R}_{{\mathcal{B}}}(G^{\prime},\theta^{\prime})$) in Equation
C.37; recall that $\Delta(\phi)$ is given in Equation C.20 (see also Equation
8). ∎

######  Proposition C.1.

The mini-batch gradient estimator is unbiased, i.e.,
$\nabla_{\phi}{\mathcal{L}}(\phi)=\mathbb{E}_{{\mathcal{B}}}\big{[}\nabla_{\phi}\widehat{{\mathcal{L}}}_{{\mathcal{B}}}(\phi)\big{]}$.
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

$\nabla_{\phi}\Delta(\phi)=\nabla_{\phi}\widehat{\Delta}_{{\mathcal{B}}}(\phi)$
(C.39)

since only the terms corresponding to the rewards differ between
$\Delta(\phi)$ and $\widehat{\Delta}_{{\mathcal{B}}}(\phi)$, and they do not
depend on $\phi$. Therefore,

$\displaystyle\mathbb{E}_{{\mathcal{B}}}\big{[}\nabla_{\phi}\widehat{{\mathcal{L}}}_{{\mathcal{B}}}(\phi)\big{]}$
$\displaystyle=\mathbb{E}_{{\mathcal{B}}}\big{[}\nabla_{\phi}[\widehat{\Delta}_{{\mathcal{B}}}(\phi)^{2}]\big{]}$
$\displaystyle=2\cdot\mathbb{E}_{{\mathcal{B}}}\big{[}\widehat{\Delta}_{{\mathcal{B}}}(\phi)\nabla_{\phi}\widehat{\Delta}_{{\mathcal{B}}}(\phi)\big{]}$
$\displaystyle=2\cdot\mathbb{E}_{{\mathcal{B}}}\big{[}\widehat{\Delta}_{{\mathcal{B}}}(\phi)\big{]}\nabla_{\phi}\Delta(\phi)$
$\displaystyle=2\Delta(\phi)\nabla_{\phi}\Delta(\phi)$
$\displaystyle=\nabla_{\phi}\big{[}\Delta(\phi)^{2}]$
$\displaystyle=\nabla_{\phi}{\mathcal{L}}(\phi),$

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

$\mathrm{Var}_{{\mathcal{B}}}\left[\log\frac{\widehat{R}_{{\mathcal{B}}}(G,\theta)}{\widehat{R}_{{\mathcal{B}}}(G^{\prime},\theta^{\prime})}\right],$

and showing the difference equals this constant yields an alternative proof. ∎

##  Appendix 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 MC3 (MH-MC3) samples both graphical
structures and parameters jointly at each move, whereas Metropolis-within-
Gibbs MC3 (Gibbs-MC3) alternates between updates of the structure, and updates
of the parameters; note that MC3 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\mid{\mathcal{D}})$ only, the
parameter sample $\theta$ corresponding to a DAG $G$ correspond to the
parameter inferred by $P(\theta\mid G,{\mathcal{D}})$ (see Section D.5.1).

Algorithm 1 JSP-GFN training procedure

1:Initialize the trajectory at $G_{0}$ the empty graph Interactions with the
environment $K$ steps

2:Sample whether the trajectory continues or not: $a\sim
P_{\phi}(\mathrm{stop}\mid G_{t})$ $a$ is the “$\mathrm{stop}$” action

3:Reset the trajectory: $G_{t+1}=G_{0}$ is the empty graph

4:Sample a new graph $G_{t+1}\sim P_{\phi}(G_{t+1}\mid
G_{t},\neg\mathrm{stop})$

5:Store the transition $G_{t}\rightarrow G_{t+1}$ in the replay buffer

6:A blank line Update of the parameters $\phi$

7:Sample a mini-batch ${\mathcal{B}}$ of transitions $G\rightarrow G^{\prime}$
from the replay buffer

8:Sample $\theta\sim P_{\phi}(\theta\mid G,\mathrm{stop})$ (and similarly for
$\theta^{\prime}$) for each transition in ${\mathcal{B}}$

9:Evaluate the rewards $R(G,\theta)$ and
$R(G^{\prime},\theta^{\prime})$$\triangleright$ Equation 5

10:Evaluate the loss ${\mathcal{L}}(\phi)$ based on
${\mathcal{B}}$$\triangleright$ Equation 8

11:$\phi\leftarrow\phi-\alpha\nabla_{\phi}{\mathcal{L}}(\phi)$$\triangleright$
Update the parameters with stochastic gradient methods convergence criterion

###  D.1 Sampling distribution

In Section 3.3, we saw that the learning objective of JSP-GFN can be written
as ${\mathcal{L}}(\phi)=\mathbb{E}_{\pi}[\Delta^{2}(\phi)]$, where $\pi$ is a
sampling distribution over $(G,\theta)$ and $(G^{\prime},\theta^{\prime})$
with full support. We use a combination of on-policy ($\pi=P_{\phi}$) and off-
policy ($\pi$ is different from $P_{\phi}$) in order to train the GFlowNet:
taking inspiration from (Deleu et al., 2022), transitions $G\rightarrow
G^{\prime}$ are sampled off-policy from a replay buffer, whereas their
corresponding parameters $\theta$ and $\theta^{\prime}$ are sampled on-policy
using our current $P_{\phi}(\theta\mid G)$. Therefore, a key difference with
Deleu et al. (2022) is that the reward $R(G,\theta)$ in Equation 5 is
calculated “lazily” when the loss is evaluated (i.e., only once $\theta$ and
$\theta^{\prime}$ 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ös-Rényi model
(Erdős and Rényi, 1960) over $d=5$ nodes, with $d$ edges on average (a setting
typically referred to as ER1). Once the structure of the graph $G^{\star}$ is
known, we sample the parameters $\theta^{\star}$ of the linear Gaussian model
randomly from a standard Normal distribution ${\mathcal{N}}(0,1)$. The linear
Gaussian model is defined as

$X_{i}=\sum_{X_{j}\in\mathrm{Pa}_{G}(X_{i})}\theta_{ij}^{\star}X_{j}+\varepsilon_{i},$
(D.1)

where $\theta_{ij}^{\star}\sim{\mathcal{N}}(0,1)$, and
$\varepsilon_{i}\sim{\mathcal{N}}(0,0.01)$; this defines all the CPDs
necessary for Equation 1. Finally, we use ancestral sampling to generate
$N=100$ observations to create the dataset ${\mathcal{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\big{(}X_{i}\mid\mathrm{Pa}_{G}(X_{i});\theta_{i}\big{)}={\mathcal{N}}(\mu_{i},\sigma^{2})$,
where

$\mu_{i}=\sum_{j=1}^{d}\mathbbm{1}\big{(}X_{j}\in\mathrm{Pa}_{G}(X_{i})\big{)}\theta_{ij}X_{j},$
(D.2)

and where $\sigma^{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 $\theta_{ij}$ of the model: $P(\theta_{ij}\mid
G)={\mathcal{N}}(0,1)$. This model differs from the widely used BGe score
(Geiger and Heckerman, 1994) in that $\sigma^{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(\theta\mid
G,{\mathcal{D}})$ can be computed analytically, and the proof is available in
Section D.5.1.

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

Figure D.1: Comparison of the marginals over graphs in terms of features
computed with the exact posterior (x-axis) and the approximation given by JSP-
GFN (y-axis). (a-c) Comparison with JSP-GFN (diag); (d-f) comparison with JSP-
GFN (full). Each point corresponds to a pair of variables $(X_{i},X_{j})$ for
each of the 20 datasets.

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}\rightarrow X_{j}$ is present in the graph:

$P(X_{i}\rightarrow
X_{j}\mid{\mathcal{D}})=\mathbb{E}_{P(G\mid{\mathcal{D}})}\big{[}\mathbbm{1}(X_{i}\rightarrow
X_{j}\in G)\big{]},$ (D.3)

where the expectation in Equation D.3 is either over the true (marginal)
posterior $P(G\mid{\mathcal{D}})$ (for the x-axis of Figure 2), or over the
distribution $P_{\phi}^{\top}$ induced by the GFlowNet (discarding $\theta$,
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}\rightsquigarrow 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

$\displaystyle P(X_{i}\rightsquigarrow X_{j}\mid{\mathcal{D}})$
$\displaystyle=\mathbb{E}_{P(G\mid{\mathcal{D}})}\big{[}\mathbbm{1}(X_{i}\rightsquigarrow
X_{j}\in G)\big{]}$ (D.4) $\displaystyle
P(X_{i}\in\mathrm{MB}(X_{j})\mid{\mathcal{D}})$
$\displaystyle=\mathbb{E}_{P(G\mid{\mathcal{D}})}\big{[}\mathbbm{1}(X_{i}\in\mathrm{MB}_{G}(X_{j}))\big{]},$
(D.5)

where $\mathrm{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 Equation D.4 & Equation 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_{\phi}(\theta_{i}\mid G,\mathrm{stop})$ in
Equation 12 is parametrized as a Normal distribution with a diagonal
covariance matrix; this adds a modeling bias since here the exact posterior
$P(\theta_{i}\mid G,{\mathcal{D}})$ is a Normal distribution with full
covariance (see Section D.5.1). To control for this bias, the second model
called _JSP-GFN (full)_ assumes that Equation 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 JSP-GFN (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 $\phi$ 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.

Table D.1: Quantitative comparison between different Bayesian structure
learning algorithms and the exact posterior on small graphs with $d=5$ nodes.
For each feature, the root meas-square error (RMSE) and Pearson’s correlation
coefficient between the features computed with the posterior approximation and
the exact posterior are reported. Values are reported as the mean and $95\%$
confidence interval across 20 different datasets.

Edge features Path features Markov features RMSE Pearson’s $r$ RMSE Pearson’s
$r$ RMSE Pearson’s $r$ MH-MC3 $0.357\pm 0.022$ $0.067\pm 0.143$ $0.368\pm
0.027$ $0.045\pm 0.179$ $0.341\pm 0.017$ $0.064\pm 0.217$ Gibbs-MC3 $0.357\pm
0.022$ $0.028\pm 0.127$ $0.367\pm 0.026$ $0.150\pm 0.162$ $0.341\pm 0.018$
$0.062\pm 0.159$ B-GES* $0.263\pm 0.070$ $0.635\pm 0.180$ $0.302\pm 0.080$
$0.544\pm 0.230$ $0.129\pm 0.022$ $0.955\pm 0.026$ B-PC* $0.305\pm 0.057$
$0.570\pm 0.138$ $0.349\pm 0.058$ $0.471\pm 0.154$ $0.354\pm 0.072$ $0.821\pm
0.087$ DiBS $0.312\pm 0.038$ $0.737\pm 0.071$ $0.357\pm 0.041$ $0.710\pm
0.079$ $0.504\pm 0.052$ $0.643\pm 0.093$ BCD Nets $0.215\pm 0.055$ $0.819\pm
0.097$ $0.266\pm 0.057$ $0.774\pm 0.109$ $0.327\pm 0.040$ $0.850\pm 0.067$ VBG
$0.237\pm 0.037$ $0.816\pm 0.064$ $0.284\pm 0.027$ $0.799\pm 0.050$ $0.434\pm
0.058$ $0.738\pm 0.091$ JSP-GFN (diag) $\mathbf{0.018\pm 0.005}$
$\mathbf{0.998\pm 0.001}$ $\mathbf{0.022\pm 0.005}$ $\mathbf{0.998\pm 0.001}$
$\mathbf{0.019\pm 0.006}$ $\mathbf{0.999\pm 0.001}$ JSP-GFN (full)
$\mathbf{0.019\pm 0.007}$ $\mathbf{0.998\pm 0.001}$ $\mathbf{0.021\pm 0.007}$
$\mathbf{0.998\pm 0.002}$ $\mathbf{0.020\pm 0.008}$ $\mathbf{0.999\pm 0.001}$

####  D.2.3 Evaluation of the posterior approximations over parameters

In order to evaluate the quality of the posterior approximation over $\theta$,
we measure how likely sample parameters from the approximation are under the
exact posterior distribution $P(\theta\mid G,{\mathcal{D}})$. More precisely,
we compute the cross-entropy between the posterior approximation
$P_{\phi}(G,\theta)$ and the exact joint posterior
$P(G,\theta\mid{\mathcal{D}})$: given a distribution $P_{\phi}(G,\theta)$
approximating the joint posterior $P(G,\theta\mid{\mathcal{D}})$, we estimate

$-\mathbb{E}_{P_{\phi}(G,\theta)}\big{[}\log P(\theta\mid
G,{\mathcal{D}})\big{]}\approx-\frac{1}{K}\sum_{k=1}^{K}\log
P(\theta^{(k)}\mid G^{(k)},{\mathcal{D}}),$ (D.6)

where $\\{(G^{(k)},\theta^{(k)})\\}_{k=1}^{K}$ are $K$ samples of the
posterior approximation $P_{\phi}(G,\theta)$. 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 Section D.2.1, except that we sample
ground truth graphs $G^{\star}$ from an Erdös-Rényi model with $2d$ 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
${\mathcal{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}\mid\mathrm{Pa}_{G}(X_{i});\theta_{i})={\mathcal{N}}(\mu_{i},\sigma^{2})$,
where

$\mu_{i}=\mathrm{MLP}({\bm{M}}_{i}{\bm{X}};\theta_{i}),$ (D.7)

where
${\bm{M}}_{i}=\mathrm{diag}\big{(}\mathbbm{1}(X_{1}\in\mathrm{Pa}_{G}(X_{i})),\ldots,\mathbbm{1}(X_{d}\in\mathrm{Pa}_{G}(X_{i}))\big{)}$
and ${\bm{X}}=(X_{1},\ldots,X_{d})$. The variance $\sigma^{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
$|\theta|=2,220$ parameters. The priors over parameters $P(\theta\mid G)$ and
over graphs $P(G)$ follow the ones described in Section 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,\theta\mid{\mathcal{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_{\phi}^{\top}(G,\theta)$ of JSP-GFN should ideally be equal to the joint
posterior (see Theorem 3.1), we have

$\log P_{\phi}^{\top}(G,\theta)\approx\log P(G,\theta\mid{\mathcal{D}})=\log
R(G,\theta)-\log P({\mathcal{D}}),$ (D.8)

where $\log P({\mathcal{D}})$ is a constant corresponding to the log-partition
function. Therefore, we can compare the log-terminating state probability
$\log P_{\phi}^{\top}(G,\theta)$ with the log-reward $\log R(G,\theta)$ (which
we can compute analytically) for different samples $(G,\theta)$, 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

$P_{\phi}^{\top}(G,\theta)=P_{\phi}(G\mid G_{0})P_{\phi}(\theta\mid
G)=P_{\phi}(\theta\mid G)\sum_{\tau:G_{0}\rightsquigarrow
G}\prod_{t=0}^{T-1}P_{\phi}(G_{t+1}\mid G_{t}),$ (D.9)

where the summation is over all the possible trajectories
$\tau=(G_{0},G_{1},\ldots,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}\mid G_{t+1})$ induces a distribution over the trajectories
$G_{0}\rightsquigarrow G$ (Bengio et al., 2023) to write $P_{\phi}(G\mid
G_{0})$ as

$\displaystyle P_{\phi}(G\mid G_{0})$
$\displaystyle=\sum_{\tau:G_{0}\rightsquigarrow G}P_{\phi}(\tau)$ (D.10)
$\displaystyle=\sum_{\tau:G_{0}\rightsquigarrow
G}\frac{P_{B}(\tau)}{P_{B}(\tau)}P_{\phi}(\tau)$ (D.11)
$\displaystyle=K!\cdot\mathbb{E}_{\tau\sim
P_{B}}\big{[}P_{\phi}(\tau)\big{]},$ (D.12)

where $P_{B}(\tau)=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_{\phi}(G\mid G_{0})$, hence of $P_{\phi}^{\tau}(G,\theta)$,
based on Monte-Carlo estimation:

$P_{\phi}(G\mid G_{0})\approx\frac{K!}{M}\sum_{m=1}^{M}P_{\phi}(\tau^{(m)}),$
(D.13)

where $\\{\tau^{(m)}\\}_{m=1}^{M}$ 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 Equation D.13 provides an unbiased estimate of $P_{\phi}(G\mid 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_{\phi}(\tau)$ may take. In order to reduce the variance, we
can first identify some trajectories that would contribute the most to the sum
in Equation D.10, and complement them with some randomly sampled trajectories
as in Equation D.13. In other words, if we have access to a subset
${\mathcal{T}}_{\mathrm{top}}$ of $B$ trajectories $\tau:G_{0}\rightsquigarrow
G$ that have a large $P_{\phi}(\tau)$, then

$P_{\phi}(G\mid
G_{0})=\sum_{\tau\in{\mathcal{T}}_{\mathrm{top}}}P_{\phi}(\tau)+\sum_{\tau\notin{\mathcal{T}}_{\mathrm{top}}}P_{\phi}(\tau)\approx\sum_{\tau\in{\mathcal{T}}_{\mathrm{top}}}P_{\phi}(\tau)+\frac{K!-B}{M}\sum_{m=1}^{M}P_{\phi}(\tau^{(m)}),$
(D.14)

where the sample trajectories $\\{\tau^{(m)}\\}_{m=1}^{M}$ can be obtained
using rejection sampling, with a uniform proposal as above. The estimate in
Equation D.14 is still unbiased, but with a lower variance. We can use beam-
search to find the “top-scoring” trajectories in
${\mathcal{T}}_{\mathrm{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}\cup\\{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:

$\widetilde{P}_{\phi}(G_{t+1}\mid G_{t})=\mathbbm{1}(e\in
G)P_{\phi}(G_{t+1}\mid G_{t}).$ (D.15)

In other words, we only keep transitions corresponding to adding edges that
are in $G$. Note that even though $\widetilde{P}_{\phi}$ 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

Figure D.2: (a-b) Comparison of JSP-GFN with other Bayesian structure learning
methods in terms of the expected-SHD to the ground truth graphs $G^{\star}$
used for data generation. (c-d) Comparison in terms of Area Under the ROC
curve (AUROC) to the ground truth graphs $G^{\star}$.

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^{\star}$ used for data generation. For example, the _expected
SHD_ , that is estimated from graphs $\\{G_{1},\ldots,G_{k}\\}$ sampled from
the posterior approximation as

$\mathbb{E}\mathrm{-SHD}\approx\frac{1}{n}\sum_{k=1}^{n}\mathrm{SHD}(G_{k},G^{\star}),$
(D.16)

where $\mathrm{SHD}(G,G^{\star})$ counts the number of edges changes (adding,
removing, reversing an edge) necessary to move from $G$ to $G^{\star}$. 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 Section D.2.2) and the target $G^{\star}$. 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,\theta\mid{\mathcal{D}})$), but merely how close the sampled graphs are
from $G^{\star}$. 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^{\star}$. Moreover, the expected SHD tends to favour overly sparse
graphs on the one hand, or posterior approximations that collapse completely
at $G^{\star}$ on the other hand, both situations indicating a poor
approximation of the true $P(G,\theta\mid{\mathcal{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}\mid\mathrm{Pa}_{G}(X_{i})\sim\mathrm{Categorical}(\pi_{i})$, where

$\pi_{i}=\mathrm{MLP}\big{(}{\bm{M}}_{i}{\bm{X}};\theta_{i}),$ (D.17)

where ${\bm{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 $|\theta|=6,545$ parameters. The priors over parameters
$P(\theta\mid G)$ and over graphs $P(G)$ follow the ones described in Section
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:

$P(X_{i}\mid\mathrm{Pa}_{G}(X_{i});\theta_{i})=\alpha_{i}\delta_{0}(X_{i})+(1-\alpha_{i}){\mathcal{N}}(\mu_{i},\sigma_{i}^{2})$
(D.18)

where $\mu_{i}$ is the result of a 2-layer MLP with 16 hidden units, as in
Equation D.17. The parameters of the CPDs contain the parameters of the MLP,
as well as the mixture parameter $\alpha_{i}$ and the variance of the
observation noise $\sigma_{i}^{2}$, for a total of $|\theta|=61,671$
parameters. The priors over parameters $P(\theta\mid G)$ and over graphs
$P(G)$ follow the ones described in Section 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

$\displaystyle P\big{(}X_{i}\mid\mathrm{Pa}_{G}(X_{i});\theta_{i}\big{)}$
$\displaystyle={\mathcal{N}}(\mu_{i},\sigma^{2})$
$\displaystyle\mathrm{where}$ $\displaystyle\mu_{i}$
$\displaystyle=\sum_{j=1}^{d}\mathbbm{1}\big{(}X_{j}\in\mathrm{Pa}_{G}(X_{i})\big{)}\theta_{ij}X_{j},$
(D.19)

and where $\sigma^{2}$ is a fixed hyperparameter. Moreover, we assume that the
parameters have a unit Normal prior associated to them, meaning that

$P(\theta_{ij}\mid
G)=\left\\{\begin{array}[]{ll}{\mathcal{N}}(\mu_{0},\sigma_{0}^{2})&\textrm{if
$X_{j}\rightarrow X_{i}\in G$}\\\
\delta_{0}&\textrm{otherwise,}\end{array}\right.$ (D.20)

where $\mu_{0}=0$, $\sigma_{0}^{2}=1$, and $\delta_{0}$ is the Dirac measure
at $0$, indicating that this parameter is always inactive.

We want to compute the posterior distribution $P(\theta_{i}\mid
G,{\mathcal{D}})$; this is sufficient, since we know that the parameters of
the different CPDs are mutually conditionally independent given $G$ and
${\mathcal{D}}$. Let $X\in{\mathbb{R}}^{N\times d}$ be the design matrix of
the dataset ${\mathcal{D}}$ (i.e., the observations ${\bm{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

$D_{i}=\mathrm{diag}\big{(}\mathbbm{1}\big{(}X_{1}\in\mathrm{Pa}_{G}(X_{i})\big{)},\ldots,\mathbbm{1}\big{(}X_{d}\in\mathrm{Pa}_{G}(X_{i})\big{)}\big{)}.$
(D.21)

We can rewrite the complete model above as

$\displaystyle P(\theta_{i}\mid G)$
$\displaystyle={\mathcal{N}}(D_{i}\mu_{0},\sigma_{0}^{2}D_{i})$ (D.22)
$\displaystyle P\big{(}X_{i}\mid\mathrm{Pa}_{G}(X_{i});\theta_{i}\big{)}$
$\displaystyle={\mathcal{N}}(XD_{i}\theta_{i},\sigma^{2}I_{N}).$ (D.23)

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 $\theta_{i}$ is a Normal
distribution

$\displaystyle P(\theta_{i}\mid G,{\mathcal{D}})$
$\displaystyle={\mathcal{N}}(\bar{\mu}_{i},\bar{\Sigma}_{i})$
$\displaystyle\mathrm{where}$ $\displaystyle\begin{aligned}
\bar{\mu}_{i}&=\bar{\Sigma}_{i}\left[\frac{1}{\sigma^{2}_{0}}D_{i}\mu_{0}+\frac{1}{\sigma^{2}}D_{i}X^{\top}X_{i}\right]\\\
\bar{\Sigma}_{i}^{-1}&=\frac{1}{\sigma_{0}^{2}}D_{i}^{-1}+\frac{1}{\sigma^{2}}D_{i}X^{\top}XD_{i}\end{aligned}$
(D.24)

where we used the conventions $1/0=\infty$ and $0\times\infty=0$. The masked
entries of $D_{i}$ will correspond to zeroed-out entries in $\bar{\mu}_{i}$,
and to rows and columns of $\bar{\Sigma}_{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(\theta_{i}\mid G,{\mathcal{D}})=\delta_{0}$.