SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 mean-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 ± 0.022	0.067 ± 0.143	0.368 ± 0.027	0.045 ± 0.179	0.341 ± 0.017	0.064 ± 0.217
Gibbs-MC3	0.357 ± 0.022	0.028 ± 0.127	0.367 ± 0.026	0.150 ± 0.162	0.341 ± 0.018	0.062 ± 0.159
B-GES*	0.263 ± 0.070	0.635 ± 0.180	0.302 ± 0.080	0.544 ± 0.230	0.129 ± 0.022	0.955 ± 0.026
B-PC*	0.305 ± 0.057	0.570 ± 0.138	0.349 ± 0.058	0.471 ± 0.154	0.354 ± 0.072	0.821 ± 0.087
DiBS	0.312 ± 0.038	0.737 ± 0.071	0.357 ± 0.041	0.710 ± 0.079	0.504 ± 0.052	0.643 ± 0.093
BCD Nets	0.215 ± 0.055	0.819 ± 0.097	0.266 ± 0.057	0.774 ± 0.109	0.327 ± 0.040	0.850 ± 0.067
VBG	0.237 ± 0.037	0.816 ± 0.064	0.284 ± 0.027	0.799 ± 0.050	0.434 ± 0.058	0.738 ± 0.091
JSP-GFN (diag)	0.018 ± 0.005	0.998 ± 0.001	0.022 ± 0.005	0.998 ± 0.001	0.019 ± 0.006	0.999 ± 0.001
JSP-GFN (full)	0.019 ± 0.007	0.998 ± 0.001	0.021 ± 0.007	0.998 ± 0.002	0.020 ± 0.008	0.999 ± 0.001
the marginal probability of a path  $X_{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
$$
P \left(X _ {i} \rightsquigarrow X _ {j} \mid \mathcal {D}\right) = \mathbb {E} _ {P (G | \mathcal {D})} \left[ \mathbb {1} \left(X _ {i} \rightsquigarrow X _ {j} \in G\right)\right] \tag {D.4}
$$
$$
P \left(X _ {i} \in \operatorname {M B} \left(X _ {j}\right) \mid \mathcal {D}\right) = \mathbb {E} _ {P (G \mid \mathcal {D})} \left[ \mathbb {1} \left(X _ {i} \in \operatorname {M B} _ {G} \left(X _ {j}\right)\right) \right], \tag {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 (D.4) & (D.5) are computed under the posterior approximation, and can be estimated using a Monte Carlo estimate over sample DAGs from the model.
In this experiment, we considered two variants of JSP-GFN. The first one, called JSP-GFN (diag), where  $P_{\phi}(\theta_i \mid G, \text{stop})$  in (12) is parametrized as a Normal distribution with a diagonal covariance matrix; this adds a modeling bias since here the exact posterior  $P(\theta_i \mid G, \mathcal{D})$  is a Normal distribution with full covariance (see Appendix D.5.1). To control for this bias, the second model called JSP-GFN (full) assumes that (12) has a full covariance matrix, as in VBG (Nishikawa-Toomey et al., 2023).
In Figure D.1, we show a similar plot as in Figure 2 (a) for all these features, for both models JSP-GFN (diag) and JSP-GFN (full). We observe that for all features, the approximation of the posterior given by JSP-GFN is very accurate (as confirmed by the Pearson's correlation coefficients). Interestingly, the more expressive model JSP-GFN (full) seems to perform slightly worse than 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.
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)} [ \log P (\theta \mid G, \mathcal {D}) ] \approx - \frac {1}{K} \sum_ {k = 1} ^ {K} \log P \left(\theta^ {(k)} \mid G ^ {(k)}, \mathcal {D}\right), \tag {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 Appendix 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} = \operatorname {M L P} \left(\boldsymbol {M} _ {i} \boldsymbol {X}; \theta_ {i}\right), \tag {D.7}
$$
where  $M_{i} = \mathrm{diag}\bigl (\mathbb{1}(X_{1}\in \mathrm{Pa}_{G}(X_{i})),\ldots ,\mathbb{1}(X_{d}\in \mathrm{Pa}_{G}(X_{i}))\bigr)$  and  $X = (X_{1},\dots,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 Appendix D.2.1.
D.3.2 Estimation of the log-terminating state probability
When the graphs are larger, it becomes impossible to compare the posterior approximation returned by JSP-GFN with the exact joint posterior  $P(G,\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}), \tag {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} \sim \sim G} \prod_ {t = 0} ^ {T - 1} P _ {\phi} \left(G _ {t + 1} \mid G _ {t}\right), \tag {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
$$
\begin{array}{l} P _ {\phi} (G \mid G _ {0}) = \sum_ {\tau : G _ {0} \sim \sim G} P _ {\phi} (\tau) (D.10) \\ = \sum_ {\tau : G _ {0} \sim \sim G} \frac {P _ {B} (\tau)}{P _ {B} (\tau)} P _ {\phi} (\tau) (D.11) \\ = K! \cdot \mathbb {E} _ {\tau \sim P _ {B}} \left[ P _ {\phi} (\tau) \right], (D.12) \\ \end{array}
$$
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)}), \tag {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 (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 (D.10), and complement them with some randomly sampled trajectories as in (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} _ {\text {t o p}}} P _ {\phi} (\tau) + \sum_ {\tau \notin \mathcal {T} _ {\text {t o p}}} P _ {\phi} (\tau) \approx \sum_ {\tau \in \mathcal {T} _ {\text {t o p}}} P _ {\phi} (\tau) + \frac {K ! - B}{M} \sum_ {m = 1} ^ {M} P _ {\phi} (\tau^ {(m)}), \tag {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 (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} \left(G _ {t + 1} \mid G _ {t}\right) = \mathbb {1} (e \in G) P _ {\phi} \left(G _ {t + 1} \mid G _ {t}\right). \tag {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
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,\dots ,G_k\}$  sampled from the posterior approximation as
$$
\mathbb {E} - \mathrm {S H D} \approx \frac {1}{n} \sum_ {k = 1} ^ {n} \mathrm {S H D} \left(G _ {k}, G ^ {\star}\right), \tag {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 Appendix 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 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} = \operatorname {M L P} \left(\boldsymbol {M} _ {i} \boldsymbol {X}; \theta_ {i}\right), \tag {D.17}
$$
(c) Linear Gaussian
(d) Non-linear Gaussian
Figure D.2: (a-b) Comparison of JSP-GFN with other Bayesian structure learning methods in terms of the expected-SHD to the ground truth graphs  $G^{\star}$  used for data generation. (c-d) Comparison in terms of Area Under the ROC curve (AUROC) to the ground truth graphs  $G^{\star}$ .
where  $X$  encodes the discrete inputs as one-hot encoded values, and the MLP has a softmax activation function for the output layer. In total, the model has  $|\theta| = 6,545$  parameters. The priors over parameters  $P(\theta \mid G)$  and over graphs  $P(G)$  follow the ones described in Appendix D.2.1.
Gene regulatory networks from gene expression data. Gene expression data is composed of either non-zero continuous data (when a gene is expressed) or (exactly) zero values (when the gene is inhibited). To capture this type of observations, we model CPDs of the Bayesian Network as zero-inflated Normal distributions:
$$
P \left(X _ {i} \mid \operatorname {P a} _ {G} \left(X _ {i}\right); \theta_ {i}\right) = \alpha_ {i} \delta_ {0} \left(X _ {i}\right) + (1 - \alpha_ {i}) \mathcal {N} \left(\mu_ {i}, \sigma_ {i} ^ {2}\right) \tag {D.18}
$$
where  $\mu_{i}$  is the result of a 2-layer MLP with 16 hidden units, as in (D.17). The parameters of the CPDs contain the parameters of the MLP, as well as the mixture parameter  $\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 Appendix D.2.1.
D.5 Proofs
D.5.1 Posterior of the linear Gaussian model
Recall that the CPD for the linear Gaussian model can be written as
$$
P \left(X _ {i} \mid \operatorname {P a} _ {G} \left(X _ {i}\right); \theta_ {i}\right) = \mathcal {N} \left(\mu_ {i}, \sigma^ {2}\right) \quad \text {w h e r e} \quad \mu_ {i} = \sum_ {j = 1} ^ {d} \mathbb {1} \left(X _ {j} \in \operatorname {P a} _ {G} \left(X _ {i}\right)\right) \theta_ {i j} X _ {j}, \tag {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 \left(\theta_ {i j} \mid G\right) = \left\{\begin{array}{l l}\mathcal {N} \left(\mu_ {0}, \sigma_ {0} ^ {2}\right)&\text {i f} X _ {j} \rightarrow X _ {i} \in G\\\delta_ {0}&\text {o t h e r w i s e ,}\end{array}\right. \tag {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  $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} = \operatorname {d i a g} \left(\mathbb {1} \left(X _ {1} \in \operatorname {P a} _ {G} \left(X _ {i}\right)\right), \dots , \mathbb {1} \left(X _ {d} \in \operatorname {P a} _ {G} \left(X _ {i}\right)\right)\right). \tag {D.21}
$$
We can rewrite the complete model above as
$$
P \left(\theta_ {i} \mid G\right) = \mathcal {N} \left(D _ {i} \mu_ {0}, \sigma_ {0} ^ {2} D _ {i}\right) \tag {D.22}
$$
$$
P \left(X _ {i} \mid \operatorname {P a} _ {G} \left(X _ {i}\right); \theta_ {i}\right) = \mathcal {N} \left(X D _ {i} \theta_ {i}, \sigma^ {2} I _ {N}\right). \tag {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
$$
P \left(\theta_ {i} \mid G, \mathcal {D}\right) = \mathcal {N} \left(\bar {\mu} _ {i}, \bar {\Sigma} _ {i}\right) \quad \text {w h e r e} \quad \begin{array}{l} \bar {\mu} _ {i} = \bar {\Sigma} _ {i} \left[ \frac {1}{\sigma_ {0} ^ {2}} 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} X D _ {i} \end{array} \tag {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$ .