Bayesian Structure Learning with Generative Flow Networks

Introduction

Bayesian networks (Pearl, 1988) are a popular framework of choice for representing uncertainty about the world in multiple scientific domains, including medical diagnosis (Lauritzen and Spiegelhalter, 1988; Heckerman and Nathwani, 1992), molecular biology (Friedman, 2004; Sebastiani et al., 2005), and ecological modeling (Varis and Kuikka, 1999; Marcot et al., 2006). For many applications, the structure of the Bayesian network, represented as a directed acyclic graph (DAG) and encoding the statistical dependencies between the variables of interest, is assumed to be known based on knowledge from domain experts. However, when this graph is unknown, we can learn the DAG structure of the Bayesian network from data alone in order to discover these statistical (or possibly causal) relationships. This may form the basis of novel scientific theories.

Given a dataset of observations, most of the existing algorithms for structure learning return a single DAG (or a single equivalence class; Chickering, 2002), and in practice those may lead to poorly calibrated predictions (Madigan et al., 1994), especially in cases where data is limited. Instead of learning a single graph candidate, we can view the problem of structure learning from a Bayesian perspective and infer the posterior over graphs $P(G\mid{\mathcal{D}})$, given a dataset of observations ${\mathcal{D}}$, to account for the epistemic uncertainty over models. Except in limited settings (Koivisto, 2006; Meilă and Jaakkola, 2006), characterizing a whole distribution over DAGs remains intractable because of its combinatorially large sample space and the complex acyclicity constraint. Therefore, we must often resort to approximations of this posterior distribution, e.g., based on MCMC or, more recently, variational inference.

In this paper, we propose to use a novel class of probabilistic models called Generative Flow Networks (GFlowNets; Bengio et al., 2021a, b) to approximate this posterior distribution over DAGs. A GFlowNet is a generative model over discrete and composite objects that treats the generation of a sample as a sequential decision problem. This makes it particularly appealing for modeling a distribution over graphs, where sample graphs are constructed sequentially, starting from the empty graph, by adding one edge at a time. In the context of Bayesian structure learning, we also introduce improvements over the original GFlowNet framework, including a novel flow-matching condition and corresponding loss function, a hierarchical probabilistic model for forward transitions, and using additional tools from the literature on Reinforcement Learning (RL). We call our method DAG-GFlowNet, to emphasize that the support of the distribution induced by the GFlowNet is exactly the space of DAGs, unlike some variational approaches that may sample cyclic graphs (Annadani et al., 2021; Lorch et al., 2021). Compared to MCMC, which operates through local moves in the sample space (here, adding or removing edges of a graph) and is therefore subject to slow mixing (Friedman and Koller, 2003), DAG-GFlowNet yields a sampling process that samples iid. DAGs, each of them constructed from scratch.

We evaluate DAG-GFlowNet on various problems with simulated and real data, on both discrete and linear-Gaussian Bayesian networks. Furthermore, we show that DAG-GFlowNet can be applied on both observational and interventional data, by modifying standard Bayesian scores (Cooper and Yoo, 1999). On smaller graphs, we also show that it is capable of learning an accurate approximation of the exact posterior distribution. The code is available online.https://github.com/tristandeleu/jax-dag-gflownet

Related Work

Methods based on MCMC have been particularly popular in Bayesian structure learning to approximate the posterior distribution. Structure MCMC (MC3; Madigan et al., 1995) simulates a Markov chain in the space of DAGs, through local moves (e.g. adding or removing an edge). Working directly with DAGs leads to slow mixing though; to improve mixing, Friedman and Koller (2003) proposed a sampler in the space of node orders, that introduced a bias (Ellis and Wong, 2008). This was further refined by either modifying the underlying space of the Markov chain (Kuipers and Moffa, 2017; Niinimäki et al., 2016), or its local moves (Mansinghka et al., 2006; Eaton and Murphy, 2007a; Kuipers et al., 2021). Recently, Viinikka et al. (2020) incorporated many of these advances into an efficient MCMC sampler called Gadget.

Variational Inference

In the context of structure learning, applying the recent advances in approximate inference based on gradient methods can be difficult due to the discrete nature of the problem (Lorch et al., 2021). Cundy et al. (2021) decomposed the adjacency matrix of a DAG into a triangular matrix and a permutation, and used a continuous relaxation to parametrize a distribution over permutations. Other methods (Annadani et al., 2021; Lorch et al., 2021) encode the acyclicity constraint into a soft prior $P(G)$, based on continuous characterizations of acyclicity (Zheng et al., 2018). While the effect of this prior can be made arbitrarily strong, this does not guarantee that the graphs sampled from the resulting distribution are acyclic. By contrast, our approach guarantees by construction that the support of the posterior approximation is exactly the space of DAGs.

Sequential decisions

In this work, we treat the construction of a sample graph from the posterior as a sequential decision problem, starting from the empty graph and adding one edge at a time. Li et al. (2018) use a similar process for creating a generative model over graphs with a fixed ordering over nodes. Similarly, although they do not consider a distribution over graphs, Buesing et al. (2020) use a variant of Monte Carlo Tree Search to approximate a distribution over a pre-specified ordering of discrete random variables. Our method, based on Generative Flow Networks, does not make any assumption on the order in which the edges are added, and multiple edge insertion sequences may lead to the same DAG. Zhu et al. (2020) learn a single high-scoring structure using RL; however, unlike our approach, the creation of this graph does not involve sequential decisions.

Background

A Bayesian network is a probabilistic model over $d$ random variables $\{X_{1},\ldots,X_{d}\}$, whose joint distribution factorizes according to a DAG $G$ as

Originally introduced to encourage the discovery of diverse modes of an unnormalized distribution (Bengio et al., 2021a), Generative Flow Networks (GFlowNets; Bengio et al., 2021b) are a class of generative models over a discrete and structured sample space ${\mathcal{X}}$. The structure of a GFlowNet is defined by a DAG over some states $s\in{\mathcal{S}}$; in general, the sample space over which we wish to define a distribution is only a subset of the overall state space of the GFlowNet: ${\mathcal{X}}\subseteq{\mathcal{S}}$. Samples $s\in{\mathcal{X}}$ are constructed sequentially by following the edges of the DAG, starting from a fixed initial state $s_{0}$. We also define a special absorbing state $s_{f}$, called the terminal state, indicating when the sequential construction terminates; some of the states $s\in{\mathcal{X}}$ are connected to $s_{f}$, and we call them complete states.“Complete” here means that the state is a valid sample from the distribution induced by the GFlowNet. This must not be confused with a “complete graph”, where all the nodes are connected to one another, when the states are DAGs (see Section 4). For example, Bengio et al. (2021a) use a GFlowNet to define a distribution over molecules, where ${\mathcal{X}}$ would correspond to the space of all (complete) molecules, which are constructed piece by piece by attaching a new block to an atom in a possibly partially constructed molecule (i.e. a state in ${\mathcal{S}}\backslash{\mathcal{X}}$). Another example of a GFlowNet structure is given in Fig. 1, illustrating the sequential process of constructing a DAG. A GFlowNet is structurally equivalent to a Markov Decision Process (MDP; Puterman, 1994) with deterministic dynamics, or a Markov Reward Process (Howard, 1971).

In addition to the DAG structure over states, every complete state $s\in{\mathcal{X}}$ is associated with a reward $R(s)\geq 0$, indicating a notion of “preference” for certain states. By convention, $R(s)=0$ for any incomplete state $s\in{\mathcal{S}}\backslash{\mathcal{X}}$. The goal of a GFlowNet is to find a flow that satisfies, for all states $s^{\prime}\in{\mathcal{S}}$, the following flow-matching condition:

where $F_{\theta}(s\rightarrow s^{\prime})\geq 0$ is a scalar representing the flow from state $s$ to $s^{\prime}$, typically parametrized by a neural network. Putting it in words, the overall flow going into $s^{\prime}$ is equal to the flow going out of $s^{\prime}$, plus some residual $R(s^{\prime})$. To learn the parameters $\theta$ of the flow with SGD, we can turn Equation 2 into a regression problem, e.g. using a least squares objective over sampled states.

If the conditions in Equation 2 are satisfied for all states $s^{\prime}$, a GFlowNet induces a generative process to sample complete states $s\in{\mathcal{X}}$ with probability proportional to $R(s)$. Starting from the initial state $s_{0}$, if we sample a complete trajectory $(s_{0},s_{1},\ldots,s_{T},s,s_{f})$ using the transition probability defined as the normalized outgoing flow

with the conventions $s_{T+1}=s$ and $s_{T+2}=s_{f}$, then $s$ is sampled with probability $P(s)\propto R(s)$. Note that the linear system in Equation 2 is in general underdetermined, and therefore it may admit many solutions $F_{\theta}(s\rightarrow s^{\prime})$ that all induce the same distribution $\propto R(s)$. Unlike MCMC, each sample $s\in{\mathcal{X}}$ is constructed from scratch, starting at the initial state $s_{0}$, instead of traversing ${\mathcal{X}}$ from sample to sample. Therefore, the underlying Markov process of the GFlowNet does not have to be irreducible, which is typically necessary in MCMC, but merely requires all the complete states to be reachable from the initial state. Finally, although GFlowNets borrow terminology from RL and control theory (e.g. MDPs, rewards, trajectories), their objective is different from the typical RL training objective: the latter seeks to maximize a function of the rewards, while the goal of GFlowNets is to model the whole distribution proportional to the rewards.

2 Detailed-balance condition

Since the flows are added together, one of the downsides of the flow-matching condition is that flows tend to be orders of magnitude larger the closer we are of the initial state (Bengio et al., 2021a), making it challenging to parametrize $F_{\theta}$. Bengio et al. (2021b) proposed an alternative characterization of GFlowNets inspired by the detailed-balance equations from the literature on Markov chains (Grimmett and Stirzaker, 2020). Instead of working with flows, this condition uses a parametrization of the forward transition probability $P_{\theta}(s_{t+1}\mid s_{t})$ directly, together with a backward transition probability $P_{B}(s_{t}\mid s_{t+1})$ to enforce reversibility. As opposed to $P_{\theta}(s_{t+1}\mid s_{t})$, which is a distribution over the children of $s_{t}$, $P_{B}(s_{t}\mid s_{t+1})$ is a distribution over the parents of $s_{t+1}$ in the structure of the GFlowNet. If all the states of the GFlowNet are complete (except the terminal state $s_{f}$), which will be the case here for generating DAGs, then we show in Appendix B that we can write the detailed-balance condition for all transitions $s\rightarrow s^{\prime}$ as follows:

Similar to Section 3.1, finding $P_{\theta}$ and $P_{B}$ that satisfy this condition for all the transitions $s\rightarrow s^{\prime}$ of the GFlowNet also yields a sampling process of complete states $s$ with probability proportional to $R(s)$, based on the forward transition probability $P_{\theta}(s_{t+1}\mid s_{t})$. Because this system of equations also admits many solutions, similar to Equation 2, we can set the backward transition probability $P_{B}$ to some fixed distribution (e.g. the uniform distribution over the parent states) to reduce the search space, making $P_{\theta}$ the only quantity to learn and, with enough capacity (to satisfy the constraints), there is a unique solution $P_{\theta}$ (Bengio et al., 2021b).

To fit the parameters $\theta$ of the forward transition probability, we can minimize the following non-linear least squares objective for all the transitions $s\rightarrow s^{\prime}$ of the GFlowNet, called the detailed-balance loss:

Alternatively, we can minimize this loss in expectation, using a distribution $\pi(s\rightarrow s^{\prime})$ with full support over transitions (i.e. for all transitions $s\rightarrow s^{\prime}$ in the GFlowNet, we have $\pi(s\rightarrow s^{\prime})>0$; see Section 5.2).

GFlowNet over Directed Acyclic Graphs

Our objective in this paper is to construct a distribution over DAGs. This is a challenging problem in general, as the space of DAGs is discrete and combinatorially large. We propose to use a GFlowNet to model such a distribution; this is particularly appropriate here since graphs are composite objects, and the acyclicity constraint can be obtained by constraining the allowed actions in each state (as in Figure 1). Note that the DAGs in this section and thereafter represent the states of the GFlowNet, and they must not be confused with the DAG structure of the GFlowNet itself.

We consider a GFlowNet where the states are DAGs over $d$ (labeled) nodes. Since the states of the GFlowNet are graphs, we will use the notation $G$ to denote a state, in favour of $s$ as in Section 3.1, except for the terminal state $s_{f}$. A transition $G\rightarrow G^{\prime}$ in this GFlowNet corresponds to adding an edge to $G$ to obtain the graph $G^{\prime}$; in other words, the graphs are constructed one edge at a time, starting from the initial state $G_{0}$, which is the fully disconnected graph over $d$ nodes. Since we assume that all the states $G$ of the GFlowNet are valid DAGs, they are all complete (i.e. connected to the terminal state $s_{f}$) with a corresponding reward $R(G)$. Figure 1 shows an illustration of the structure of such a GFlowNet, where the states are DAGs over $d=3$ nodes. This application to graphs highlights the importance of the DAG structure of the GFlowNet, since there can be multiple paths leading to the same state: for any graph $G$ with $k$ edges, there are $k!$ possible paths from $G_{0}$ leading to $G$, because the edges of $G$ may have been added in any order.

To guarantee the integrity of the GFlowNet, we have to ensure that adding a new edge to some state $G$ also yields a valid DAG, meaning that this edge (1) must not be already present in $G$, and (2) must not introduce a cycle. Fortunately, we can filter out invalid actions using some mask ${\bm{m}}$ associated to the graph, built from the adjacency matrix of $G$ and the transitive closure of its transpose, and that can be updated efficiently after the addition an edge (Giudici and Castelo, 2003). A description of this update is given in Appendix C for completeness.

2 Forward transition probabilities

Following Section 3.2, the GFlowNet may be parametrized only by the forward transition probabilities $P_{\theta}(G_{t+1}\mid G_{t})$; here, $G_{t+1}$ might be the terminal state $s_{f}$ by abuse of notation. To make sure that the detailed-balance conditions can be satisfied, we need to define these transition probabilities using a sufficiently expressive function, such as a neural network. We use a hierarchical model, where the forward transition probabilities are defined using two neural networks: (1) a network modeling the probability of terminating $P_{\theta}(s_{f}\mid G)$, and (2) another giving the probability $P_{\theta}(G^{\prime}\mid G,\neg s_{f})$ of transitioning to a new graph $G^{\prime}$, given that we do not terminate. The probability of taking a transition $G\rightarrow G^{\prime}$ is then given by

In practice, as $G^{\prime}$ is the result of adding an edge to the DAG $G$, we can model $P_{\theta}(G^{\prime}\mid G,\neg s_{f})$ as a probability distribution over the $d^{2}$ possible edges one could add to $G$—this includes self-loops, for simplicity, even though these actions are guaranteed to be invalid. We can use the mask ${\bm{m}}$ introduced in Section 4.1 to filter out actions that would not lead to a valid DAG $G^{\prime}$ and set $P_{\theta}(G^{\prime}\mid G,\neg s_{f})=0$ for any invalid action (as well as normalize $P_{\theta}$ accordingly).

3 Parametrization with Linear Transformers

Beyond having enough capacity to satisfy as well as possible the detailed-balance condition at all states, we choose to parametrize the forward transition probabilities with neural networks to benefit from their capacity to generalize to states not encountered during training. In practice, instead of defining two separate networks to parametrize $P_{\theta}(s_{f}\mid G)$ and $P_{\theta}(G^{\prime}\mid G,\neg s_{f})$, we use a single neural network with a common backbone and two separate heads, to benefit from parameter sharing. The full architecture is given in Figure 2.

Our choice of neural network architecture is motivated by multiple factors: we want an architecture (1) that is invariant to the order of the inputs, since $G$ is represented as a set of edges, (2) that transforms a set of input edges into a set of output probabilities for each edge to be added, in order to define $P_{\theta}(G^{\prime}\mid G,\neg s_{f})$, and (3) whose parameters $\theta$ do not scale too much with $d$. A natural option would be to use a Transformer (Vaswani et al., 2017); however, because the size of our inputs is $d^{2}$, the self-attention layers would scale as $d^{4}$, and this would severely limit our ability to apply our method to model a distribution over larger DAGs.

We opted for a Linear Transformer (Katharopoulos et al., 2020) instead, which has the advantage to not suffer from this quadratic scaling in the input size. This architecture relies on a linearized attention mechanism, defined as

where ${\bm{x}}$ is the input of the linearized attention layer, $\phi(\cdot)$ is a non-linear feature map, $J$ is the size of the input ${\bm{x}}$ (in our case, $J=d^{2}$), and $Q$, $K$, and $V$ are linear transformations of ${\bm{x}}$ corresponding to the queries, keys, and values respectively, as is standard with Transformers.

Application to Bayesian Structure Learning

We are given a dataset ${\mathcal{D}}=\{{\bm{x}}^{(1)},\ldots,{\bm{x}}^{(N)}\}$ of $N$ observations ${\bm{x}}^{(j)}$, each consisting of $d$ elements. We consider the task of characterizing the posterior distribution $P(G\mid{\mathcal{D}})$ over Bayesian networks that model these observations. We assume that the samples in ${\mathcal{D}}$ are iid. and fully-observed. As an alternative to MCMC (Madigan et al., 1995) or variational inference (Lorch et al., 2021), we approximate the posterior distribution over DAGs using a GFlowNet, as described in the previous section. For any DAG $G$, we will define its reward as the joint probability

where $P(G)$ is a prior over DAGs (Eggeling et al., 2019), and $P({\mathcal{D}}\mid G)$ is the marginal likelihood. In Sec. 3.2, we saw that if the detailed-balance conditions are satisfied for all the states of the GFlowNet, then this yields a sampling process with probability proportional to $R(G)$. Therefore, by Bayes’ theorem, a GFlowNet with the specific reward function in Equation 7 approximates the posterior distribution $P(G\mid{\mathcal{D}})\propto R(G)$. We call our method DAG-GFlowNet.

Following prior works on Bayesian structure learning, we assume that both the priors over parameters $P(\phi\mid G)$ of the Bayesian network (required to compute the marginal likelihood) and over structures $P(G)$ are modular (Heckerman et al., 1995; Chickering et al., 1995). As a consequence the reward $R(G)$ is also modular, and its logarithm can be written as a sum of local scores that only depend on individual variables and their parents in $G$:

Note that with our choice of reward, $\log R(G)$ corresponds to the Bayesian score (Koller and Friedman, 2009). Examples of modular scores include the BDe score (Heckerman et al., 1995) and the BGe score (Geiger and Heckerman, 1994; Kuipers et al., 2014). In order to fit the parameters $\theta$ of the GFlowNet, we will use the detailed-balance loss in Equation 4. We can observe that this loss function only involves the difference in log-rewards $\log R(G^{\prime})-\log R(G)$ between two consecutive states, where $G^{\prime}$ is the result of adding some edge $X_{i}\rightarrow X_{j}$ to the DAG $G$. Using our assumption of modularity, we can therefore compute this difference efficiently, as the terms in Equation 8 remain unchanged for $j^{\prime}\neq j$:

This difference in local scores is sometimes called the delta score, or the incremental value (Friedman and Koller, 2003), and has been employed in the literature to improve the efficiency of search algorithms (Chickering, 2002; Koller and Friedman, 2009).

2 Off-policy learning

As the number of states in DAG-GFlowNet is super-exponential in $d$, the number of nodes in each DAG $G$, it would be impractical to minimize the detailed-balance loss for all possible transitions $G\rightarrow G^{\prime}$. Alternatively, we can minimize this loss in expectation using a distribution $\pi(G\rightarrow G^{\prime})$ with full support over transitions:

This distribution $\pi(G\rightarrow G^{\prime})$ can be arbitrary; for example, we can use $P_{\theta}(G^{\prime}\mid G)$ directly and learn it on-policy (Rummery and Niranjan, 1994), as long as it assigns non-zero probability to any next state $G^{\prime}$.

Taking inspiration from Deep Q-learning (Mnih et al., 2015), we instead learn $P_{\theta}$ using off-policy data. Transitions $G\rightarrow G^{\prime}$ are collected based on $P_{\theta}(G^{\prime}\mid G)$, along with their corresponding delta score Equation 9, and they are stored in a replay buffer. We can also sample some transitions uniformly at random, with probability $\varepsilon$, to encourage exploration. To estimate ${\mathcal{L}}(\theta)$ and update the parameters $\theta$, we can then sample a mini-batch of transitions randomly from the replay buffer. Moreover, again inspired by Deep Q-learning (Van Hasselt et al., 2018), we found it advantageous to evaluate $P_{\bar{\theta}}(s_{f}\mid G^{\prime})$ in Equation 10 with a separate target network—where the parameters $\bar{\theta}$ are updated periodically.

Experimental results

We compared DAG-GFlowNet against 3 broad classes of Bayesian structure learning algorithms: MCMC, non-parametric DAG Bootstrapping (Friedman et al., 1999), and variational inference. We used Structure MCMC (MC3; Madigan et al., 1995) and the recent Gadget (Viinikka et al., 2020) samplers as two representative methods based on MCMC. Following Lorch et al. (2021), we used two variants of Bootstrapping based on the score-based algorithm GES (Bootstrap GES; Chickering, 2002), and the constraint-based algorithm PC (Bootstrap PC; Spirtes et al., 2000), as the internal structure learning routines. Finally for methods based on variational inference, we used DiBS (Lorch et al., 2021) and BCD Nets (Cundy et al., 2021). Throughout this section, we used the BGe score for continuous data, and the BDe score for discrete data, to compute $\log p({\mathcal{D}}\mid G)$.

In order to measure the quality of the posterior approximation returned by DAG-GFlowNet, we want to compare it with the exact posterior distribution $P(G\mid{\mathcal{D}})$. However, the latter requires an exhaustive enumeration of all possible DAGs, which is only feasible for graphs with no more than $5$ nodes. Therefore, we sampled $N=100$ datapoints from a randomly generated (under an Erdős-Rényi model; Erdős and Rényi, 1960) linear-Gaussian Bayesian network over $d=5$ variables. We used the BGe score to compute the reward $R(G)=P(G)P({\mathcal{D}}\mid G)$. The exact posterior distribution $P(G\mid{\mathcal{D}})$ is obtained by enumerating all $29,\!281$ possible DAGs over $5$ nodes and computing their respective rewards $R(G)$ (normalized to sum to $1$).

We evaluated the quality of the approximation based on the probability of various structural features. For example, using samples $\{G_{1},G_{2},\ldots,G_{n}\}$ from the posterior approximation, the marginal probability of an edge feature $X_{i}\rightarrow X_{j}$ can be estimated with

where $\mathbf{1}(\cdot)$ is the indicator function. For the exact posterior, we can obtain the posterior probability of the edge feature by simply marginalizing over $P(G\mid{\mathcal{D}})$. Similarly, we compute (or estimate) the marginal probability of a path feature ${X_{i}\rightsquigarrow X_{j}}$, i.e. of a (directed) path existing from $X_{i}$ to $X_{j}$, and the probability of a Markov feature ${X_{i}\sim_{M}X_{j}}$, i.e. of $X_{i}$ being in the Markov blanket of $X_{j}$ (Friedman and Koller, 2003). These features are computed for all variables $X_{i}$ and $X_{j}$ in the Bayesian network.

In Figure 3, we compare the probabilities of these features for both the exact posterior and the distribution induced by DAG-GFlowNet, where we repeated the experiment above with $20$ different (ground-truth) Bayesian networks. We observe that the probabilities of all structural features estimated by the GFlowNet are strongly correlated with the exact marginal probabilities. This shows that DAG-GFlowNet is capable of learning a very accurate approximation of the posterior distribution over graphs $P(G\mid{\mathcal{D}})$.

2 Simulated data

We follow the experimental setup of Zheng et al. (2018) & Lorch et al. (2021), and sample synthetic data from linear-Gaussian Bayesian networks with randomly generated structures; details about this data generation process are given in Section D.2. To show that DAG-GFlowNet can effectively approximate the posterior distribution over larger graphs, we experimented with Bayesian networks of size $d=20$ (and $d=50$, see Section D.2). Similar to Section 6.1, the ground-truth graphs are sampled according to an Erdős-Rényi model, with $2d$ edges in expectation—a setting sometimes referred to as ER2 (Cundy et al., 2021). For each experiment, we sampled a dataset ${\mathcal{D}}$ of $N=100$ observations, and we used the BGe score to compute $R(G)$.

3 Application: Flow Cytometry Data

We also evaluated DAG-GFlowNet on real-world flow cytometry data (Sachs et al., 2005) to learn protein signaling pathways. The data consists of continuous measurements of $d=11$ phosphoproteins in individual T-cells. Out of all the measurements, we selected the $N=853$ observations corresponding to the first experimental condition of Sachs et al. (2005) as our dataset ${\mathcal{D}}$. Following prior work on structure learning, we used the DAG inferred by Sachs et al. (2005), containing $d=11$ nodes and $17$ edges, as our graph of reference (ground-truth). However, it should be noted that this “consensus graph” may not represent a realistic and complete description of the system being modeled here (Mooij et al., 2020). We standardized the data, and used the BGe score to compute $R(G)$.

Beyond these metrics, we would like to test if the advantages of Bayesian structure learning are also reflected in the distribution induced by DAG-GFlowNet. In particular, we want to study (1) if this distribution covers multiple high-scoring DAGs, instead of being peaked at a single most likely graph, and (2) if the GFlowNet can sample a variety of DAGs from the same Markov equivalence class (MEC), showing the inherent uncertainty over equivalent graphs. In Figure 5, we visualize the MECs of the graphs sampled with DAG-GFlowNet, and two methods based on MCMC (MC3 and Gadget); other baselines were excluded for clarity. The size of each point represents the number of unique DAGs in the corresponding MEC. We observe that DAG-GFlowNet largely follows the behavior of MCMC: the distribution does not collapse to a single most-likely DAG, and covers multiple MECs. Moreover, the GFlowNet is also capable of sampling different equivalent DAGs (corresponding to larger points), showing again that the distribution does not collapse to a single representative of the MECs with higher marginal probability. We also observe that the maximum a posteriori MEC found by DAG-GFlowNet reaches a higher score than the one found with Gadget, but a lower score than MC3; as a point of reference, the score of the best MEC obtained with GES (Chickering, 2002) is $-10,\!716.12$.

4 Application: Interventional data

In addition to the observational data we used in Section 6.3, Sachs et al. (2005) also provided flow cytometry data under different experimental conditions, where the T-cells were perturbed with some reagents; this effectively corresponds to interventional data (Pearl, 2009). Although a molecular intervention may be imperfect and affect multiple proteins (Eaton and Murphy, 2007b), we assume here that these interventions are perfect, and the intervention targets are known. We used a discretized dataset of $N=5,\!400$ samples from $9$ experimental conditions—of which $6$ are interventions. We modified the BDe score to handle this mixture of observational and interventional data (Cooper and Yoo, 1999).

In Table 2, we compare with Eaton and Murphy (2007b), which compute the AUROC of the exact posterior using dynamic programming, therefore working as an upper bound for what a posterior approximation can achieve. They achieve this at the expense of computing only edge marginals, without providing access to a distribution over DAGs. We also use the modified BDe score with MC3, which predicts sparser graphs with higher SHD than DAG-GFlowNet, but lower AUROC. Note that this setup is different from previous works which use continuous data instead (Brouillard et al., 2020; Faria et al., 2022).

Conclusion

We have proposed a new method for Bayesian structure learning, based on a novel class of probabilistic models called GFlowNets, where the generation of a sample graph is treated as a sequential decision problem. We introduced a number of enhancements to the standard framework of GFlowNets, specifically designed for approximating a distribution over DAGs. In cases where the data is limited and measuring the epistemic uncertainty is critical, DAG-GFlowNet offers an effective solution to approximate the posterior distribution over DAGs $P(G\mid{\mathcal{D}})$. However, we also observed that in its current state, DAG-GFlowNet may suffer from some limitations, notably as the size of the dataset ${\mathcal{D}}$ increases; see Appendix A for a discussion.

While DAG-GFlownet operates on the space of DAGs directly, the structure of the GFlowNet may eventually be adapted to work with alternative representations of statistical dependencies in Bayesian networks, such as essential graphs for MECs (Chickering, 2002). Moreover, although we have already shown that DAG-GFlowNet can approximate the posterior using a mixture of observational and interventional data, we will continue to study in future work its applications to causal discovery, especially in the context of learning the structure of models with latent variables.

Appendix A Limitations of DAG-GFlowNet

Although we have shown in the main paper that DAG-GFlowNet is capable of learning an accurate approximation of the posterior distribution $P(G\mid{\mathcal{D}})$ when the size of the dataset ${\mathcal{D}}$ is moderate (a situation where the benefits of a Bayesian treatment of structure learning are larger), we observed that as the size of the dataset increases, fitting the detailed-balance loss in Equation 10 was more challenging. This can be explained by the fact that with a larger amount of data, the posterior distribution becomes very peaky (Koller and Friedman, 2009). As a consequence, in this situation, the delta-score in Equation 9, which is required to calculate the loss, can take a wide range of values: adding an edge to a graph can drastically increase or decrease its score. In turn, the neural network parametrizing $P_{\theta}(G_{t+1}\mid G_{t})$ needs to compensate for these large fluctuations, making it harder to train.

Unfortunately, some of the standard techniques used in Machine Learning to tackle this issue, such as normalization of the inputs, cannot be applied here. Normalizing the delta-score is equivalent to normalizing the rewards $R(G)$ and $R(G^{\prime})$ themselves, and as a consequence it would change the distribution that is being approximated: instead of approximating the posterior distribution $P(G\mid{\mathcal{D}})$, we would approximate a distribution $P(G\mid{\mathcal{D}})^{\tau}$ under some temperature $\tau$. Solutions to this problem include a schedule of temperature, similar to simulated annealing, or a reparametrization of $P_{\theta}(G_{t+1}\mid G_{t})$ to better handle large fluctuations of delta-scores; this exploration is left as future work.

Appendix B Detailed-balance condition with all complete states

In this section, we will prove a special case of the detailed-balance condition introduced by Bengio et al. (2021b) applied to the case where all the states of the GFlowNet are complete (except the terminal state $s_{f}$). To simplify the presentation, we will follow the notations of Bengio et al. (2021b), and denote the forward transition probability by $P_{F}(s_{t+1}\mid s_{t})$—instead of $P_{\theta}(s_{t+1}\mid s_{t})$ in the main paper. Recall that the detailed-balance condition (Bengio et al., 2021b, Def. 17) is given by

In the case where all the states are complete, we also know that (Bengio et al., 2021b, Def. 16)

where $F(s\rightarrow s^{\prime})$ represents the flow from state $s$ to $s^{\prime}$, as described in Section 3.1, $F(s)$ is the total flow through state $s$, and we used Proposition 4 & Equation 34 of Bengio et al. (2021b) to introduce $F(s_{t})$ and $R(s_{t})$ respectively. Replacing $F(\cdot)$ in Equation B.1 yields the expected condition:

The original formulation in Equation B.1 would require us to parametrize both $P_{F}(s_{t+1}\mid s_{t})$ and $F(s)$. On the other hand, using this alternative condition, we only have to parametrize $P_{F}(s_{t+1}\mid s_{t})$ (including when $s_{t+1}=s_{f}$ is the terminal state).

Appendix C Definition and update of the mask over actions

In Section 4.1, we introduced a mask ${\bm{m}}$ associated with a DAG $G$ to indicate which edges could be legally added to $G$ to obtain a new valid DAG $G^{\prime}$. This mask must ignore (1) the edges already present in $G$ (which cannot be added further), and (2) any edge whose addition leads to the introduction of a cycle. The mask ${\bm{m}}$ is constructed using (1) the adjacency matrix of $G$, and (2) the adjacency matrix of the transitive closure of $G^{\top}$, the transpose graph of $G$; recall that $G^{\top}$ is obtained from $G$ by inverting the direction of its edges.

Giudici and Castelo (2003) use a similar construction to efficiently obtain the legal actions their MCMC sampler may take. In particular, they show that this mask ${\bm{m}}$ can be updated very efficiently online as edges are added one by one. In practice, this allows us to circumvent an expensive check for cycles at every stage of the construction of a sample DAG in the GFlowNet. Since the mask can be composed in 2 parts (as explained above), we can simply update each part anytime a new edge is added to a DAG $G$.

In Figure 6, we show how the mask ${\bm{m}}_{t}$ associated with a graph $G_{t}$ can be updated after adding a new edge $C\rightarrow A$ to obtain the mask ${\bm{m}}_{t+1}$. The mask is decomposed in 2 parts: the adjacency matrix of $G_{t}$, and the transitive closure of $G_{t}^{\top}$. After adding $C\rightarrow A$, each component is updated separately:

Adjacency matrix: To update the adjacency matrix, the entry in the adjacency matrix must be set (here, the entry corresponding to the edge $C\rightarrow A$).

Transitive closure: To update the transitive closure of the transpose, we need to compute the outer product of the column corresponding to the target of the edge (here $A$, in blue) with the row corresponding to the source of the edge (here $C$, in red). The outer product is added (more precisely, this is a binary OR) to the initial transitive closure.

These two operations can be done very efficiently in $O(d^{2})$, where $d$ is the number of nodes in the DAG.

Appendix D Additional experimental results

D.2 Simulated data

For our experiments on simulated data, we followed the generation process described in (Lorch et al., 2021). The data was generated in the following way:

We sampled a DAG using an Erdős-Rényi model (Erdős and Rényi, 1960), with $2d$ edges on average; the value of the probability of creating an edge between two nodes was scaled accordingly.

Once the structure of the graph is known, we sampled the parameters of the linear-Gaussian model randomly from a standard Normal distribution ${\mathcal{N}}(0,1)$. The linear-Gaussian model is therefore defined as, $\forall j\in[1,d]$

where $\beta_{ij}\sim{\mathcal{N}}(0,1)$, and $\varepsilon\sim{\mathcal{N}}(0,0.01)$. The defines all the conditional probability distribution of the generative model.

Once the full Bayesian Network is known, we used ancestral sampling to generate $N=100$ datapoints to fill our dataset ${\mathcal{D}}$.

D.3 Flow Cytometry Data

In Section 6.3, we described an application of DAG-GFlowNet to real-world flow cytometry data. In particular, we showed in Figure 5 that DAG-GFlowNet was capable of modeling a distribution that was not only capable of capturing the mode of the posterior distribution (i.e., graphs with a high score), but also had diversity in the graphs sampled, both in terms of the different Markov Equivalence Classes (MECs) those graphs belong to, but also multiple unique DAG instances of the same MEC (depicted by the size of each point in Figure 5).

Out of $1,\!000$ graphs sampled with BCD Nets, those graphs belonged to one of only two MECs (with a BGe score around $-10,\!950$). Furthermore, as shown by the size of each point, those MECs happen to only contain a single unique DAG. Overall, this means that BCD Nets only returned 2 unique DAGs (out of the $1,\!000$ samples), showing the lack of diversity of the posterior approximation learned with BCD Nets.

DiBS sampled a significant number of very low scoring DAGs, with BGe scores as low as $-12,\!600$ (whereas the best MEC obtained with GES (Chickering, 2002) had a score of $-10,\!716.12$).

With our choice of the BGe score, the true posterior distribution would assign the same probability to graphs in the same MEC. However, we can see that DiBS only returned graphs belonging to unique MECs, as opposed to having multiple unique DAGs from the same MEC. This shows that while DiBS has a high diversity in terms of MECs (mainly due to covering low-scoring DAGs), DiBS suffers from a lack of diversity with a single MEC, which would be expected from a faithful approximation of the posterior distribution.