SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We argue that the iterative procedure used to restrict the initial search space in Kuipers et al. [2022] can be particularly useful in our case to efficiently approximate the PEPs before running the chain. Following the same notation as in Kuipers et al. [2022], let  $\mathcal{H}$  be the adjacency matrix underlying a skeleton graph obtained from any DAG learning algorithm (e.g. PC, GES, etc.). For node  $X_{j}$ , the adjacency matrix  $H$  defines the permissible parent set  $h^{j} = \{X_{i}: H_{ij} = 1\}$ . All possible parent combinations of node  $X_{j}$  are included in the collection of permissible parent set which is defined by
$$
\mathbf {h} ^ {j} = \{m \mid m \subseteq h ^ {j} \}. \tag {10}
$$
The collection of permissible parent set  $\mathbf{h}^j$  can be then extended to a bigger collection  $\mathbf{h}_{+}^{j}$  as described in Section 4.3 of Kuipers et al. [2022]. The extended collection of permissible parent set  $\mathbf{h}_{+}^{j}$  essentially includes one additional node from outside the permissible parent set  $h^j$  as an additional parent, otherwise the candidate parents of  $X_{j}$  are restricted to the set  $h^j$ . After deriving  $\mathbf{h}_{+}^{j}$  via this procedure as in Kuipers et al. [2022], we then approximate  $\tilde{\pi}_{ij}$  through the following procedure:
1. We remove the acyclicity constraint from the prior (2) and compute the unconstrained posterior distribution  $\pi^u (\gamma)$ , approximating the marginal probability  $\pi^u (\gamma_{ij})$  under  $\mathbf{h}_{+}^{j}$ .
2. We approximate the joint probability of a pair edges by the product
$$
\tilde {\pi} \left(\gamma_ {i j}, \gamma_ {j i}\right) \approx \pi^ {u} \left(\gamma_ {i j}\right) \times \pi^ {u} \left(\gamma_ {j i}\right). \tag {11}
$$
3. We impose the simplest acyclicity constraint which avoids creating two edges with opposite directions, i.e. let  $\tilde{\pi} (\gamma_{ij} = 1,\gamma_{ji} = 1) = 0$ . Approximate
$$
\tilde {\pi} _ {i j} = \frac {\tilde {\pi} \left(\gamma_ {i j} = 1 , \gamma_ {j i} = 0\right)}{\tilde {\pi} \left(\gamma_ {i j} = 1 , \gamma_ {j i} = 0\right) + \tilde {\pi} \left(\gamma_ {i j} = 0 , \gamma_ {j i} = 1\right) + \tilde {\pi} \left(\gamma_ {i j} = 0 , \gamma_ {j i} = 0\right)}. \tag {12}
$$
The full details for the calculations of  $\tilde{\pi}_{ij}$  are given in the supplementary material (Appendix C).
3.3 Introducing the "reversal" neighborhood
In the original PARNI proposal for Bayesian variable selection, each  $K_{r}$  in the intermediate proposals corresponds to exactly one position  $(i,j)$ , such that  $k_{ij} = 1$  and that each intermediate neighborhood  $\mathcal{N}(\gamma (r - 1),K_r)$  only contains two DAGs:  $\gamma (r - 1)$  itself and the DAG obtained by flipping  $K_{r}$  in  $\gamma (r - 1)$ . This sequential neighborhood construction is clearly not appropriate for the problem of DAG sampling and might result in very slow mixing. In addition, it is natural to assume that a pair of edges  $i\longleftrightarrow j$  are highly correlated and thus equally likely to be included in a DAG. It is then necessary to specify a reversal move that achieves good mixing between pairs of edges. The neighborhood construction in the original PARNI proposal does not support a reversal move in one single sub-proposal. To elaborate on this, suppose that currently  $\gamma_{ij} = 1$  and  $\gamma_{ji} = 0$ , to propose a reversal move according to the original PARNI proposal, we should first propose to remove  $\gamma_{ij}$  from the DAG, then include  $\gamma_{ji}$  in the following sub-proposals. These flips might be running into issues if the posterior model probability for excluding both edges,  $\gamma_{ij} = \gamma_{ji} = 0$ , is low, so that these moves are rarely proposed and the chain gets stuck at  $\gamma_{ij} = 1$ .
Therefore, we construct a bigger neighborhood containing the edge reversal move of a DAG. This modification allows for an edge reversal move within a single sub-proposal  $q_{g,K_r}$  without going through the low posterior probability models and avoids the chain getting stuck. To clarify this, suppose  $k$  is sampled at each iteration of the chain. For all those pairs which satisfy  $k_{ij} = 1$  and  $k_{ji} = 1$  in  $k$ , we construct the reversal neighborhoods for edges  $\gamma_{ij}$  and  $\gamma_{ji}$ . For other positions such that  $k_{ij} = 1$ , but  $k_{ji} = 0$ , we construct the original neighborhood which only contains the current model and model with  $\gamma_{ij}$  flipped. This reversal neighborhood specification is included into the main PARNI-DAG proposal given in the supplementary material (Appendix D).
3.4 New adaptive scheme for better computational efficiency
Although the point-wise implementation reduces the per-iteration costs from exponential to linear in the neighborhood size (equivalent to the number of sub-neighborhoods in PARNI), the PARNI proposal is still computationally expensive when the sampled DAGs are not very sparse. For example, considering the gsim100 simulated data of Suter et al. [2023], featuring 161 true edges, we would be often required to compute 320 posterior DAG probabilities on average in a single MCMC iteration, and the computation thus becomes 320 times more expensive than the standard ADR proposal in this case. To reduce the computational cost of the PARNI-DAG proposal even more, we make use of a 'neighborhood thinning' parameter  $\omega$  to control the number of neighborhoods to be evaluated.
In the original PARNI proposal, the thinning parameter  $\omega$  represents the random walk jumping probability, so it does not take the posterior inclusion probabilities into account and does not affect the number of neighborhoods evaluated. For the new PARNI-DAG proposal instead, with probability  $1 - \omega$ , we will skip evaluation of some of the neighborhoods. Within each iteration, for each point  $r$  in the sequence, given the current sub-proposal  $\gamma(r - 1)$  and sub-neighborhood indicator  $K_r$ , PARNI-DAG:
1. with probability  $\omega$ , proposes  $\gamma(r) \sim q_{K_r}(\gamma(r - 1), \cdot)$  as in (5).
2. with probability  $1 - \omega$ , does not evaluate the posterior probabilities of other DAGs in  $\mathcal{N}(\gamma(r-1), K_r)$ , and directly propose  $\gamma(r) = \gamma(r-1)$ .
Varying  $\omega$  can vary the number of posterior DAG probability evaluations. The larger the value of  $\omega$ , the more the DAG probabilities evaluated. If  $\omega = 1$ , all neighborhoods are required to be evaluated to propose a new DAG.
The parameter  $\omega$  is updated adaptively. Here, we consider the Robins-Monro adaptation scheme to update  $\omega$  on the fly. At time  $t$ , the new  $\omega^{(t + 1)}$  parameter is updated via
$$
\operatorname {l o g i t} _ {\epsilon} \omega^ {(t + 1)} = \operatorname {l o g i t} _ {\epsilon} \omega^ {(t)} - \psi_ {t} \left(\mathcal {N} _ {t} - \tilde {\mathcal {N}}\right) \tag {13}
$$
where  $\mathcal{N}_t$  is number of sub-neighborhoods evaluated at time  $t$  and  $\tilde{\mathcal{N}}$  is a pre-specified expected number of evaluations. We set  $\psi_t = t^{-0.7}$ , meaning that the adaptation of  $\omega$  diminishes at the rate of  $\mathcal{O}(t^{-0.7})$ . From the empirical results, we find that  $\tilde{\mathcal{N}} = 10$  evaluated neighborhoods on average seems to achieve optimal balance between mixing performance and computational cost. Therefore, we choose to set  $\tilde{\mathcal{N}} = 10$  for most of the numerical studies in the next section and in the supplementary material.
4.1 Convergence and mixing efficiency
Protein data. We first consider the real-world protein-signalling dataset [Sachs et al., 2005], found also in Cundy et al. [2021], to test PARNI-DAG's mixing. The dataset consists of  $n = 11$  nodes and  $N = 853$  observations and the ground-truth DAG (17 edges) is provided by expert knowledge. We compare the performance of PARNI-DAG to ADR, Order MCMC and Partition MCMC. We provide the full details on the implementation of the compared models in the supplementary materials (Appendix E). For a fair comparison of the models, we re-define the score functions for Order and Partition MCMC such that all schemes target the same posterior distribution as specified in Section 2. Notice that for this dataset, it is feasible for all the four MCMC schemes to use the full skeleton, without the need to restrict the initial space [Kuipers et al., 2022]. For ADR and PARNI-DAG, we use prior parameters  $g = 10$  and  $h = 1/11$ .
We compare trace plots of log posterior DAG probabilities for ADR, PARNI-DAG and Partition MCMC in Figure 1. Since we could not extract the full score trace of Order MCMC using the BiDAG package [Suter et al., 2023], Order MCMC is not included in Figure 1. We ran ADR, PARNI-DAG and Partition MCMC for 480,000, 60,000 and 20,000 iterations respectively which take about the same CPU time $^3$ , and thin the output from ADR and PARNI-DAG so that 20,000 measures are kept. From Figure 1, all these three algorithms mix equally well, the only difference is that the Partition
3Using Intel i7 2.80 GHz processor
Figure 1: Protein dataset: trace plots of log posterior model probabilities. All algorithms were ran for the same CPU-time and thinned to 20,000 measures.
algorithm	MSE
ADR	8.84
PARNI-DAG	19.44
Order MCMC	145.17
Partition MCMC	121.83
Table 1: Protein dataset: time-normalised median MSE  $(\times 10^{-8})$  on estimating posterior edge probabilities. (Lower is better.)
MCMC algorithm explores more DAGs with lower probabilities. We then compare the median mean squared errors (MSE) on estimating PEPs compared to a ground-truth estimate. The MSE assesses both the convergence of the chain and the bias in the Monte Carlo estimates. The ground-truth PEPs were obtained by running Partition MCMC for approximately 5 hours. We ran all algorithms for 20 replications. Each individual chain was run for 3 minutes. The MSE estimates are presented in Table 1. The structure MCMC schemes (ADR and PARNI-DAG) generally outperform Order MCMC and Partition MCMC and return estimates with lower MSE. Notice that ADR in this specific case also outperforms PARNI-DAG, since the dataset is very low-dimensional (11 nodes), with PARNI-DAG not trailing very far behind. We will demonstrate in the next experiments how PARNI-DAG instead performs better in higher-dimensional settings.
gsm100 data. We also study the performance of ADR, PARNI-DAG, Order MCMC and Partition MCMC on a more complex graph, the gsm100 simulated dataset, found in the BiDAG package [Suter et al., 2023]. This features a randomly generated DAG with  $n = 100$  nodes and  $N = 100$  observations, with 161 true edges. For ADR and PARNI-DAG we set prior parameters to  $g = 10$  and  $h = 1/100$ . Since  $n > 20$ , it is not feasible to use the full skeleton as in the protein dataset. So we consider two skeletons to restrict the initial space: PC derived skeleton  $\mathcal{H}_{\mathrm{PC}}$  and the skeleton derived with the iterative procedure of Kuipers et al. [2022]  $\mathcal{H}_{\mathrm{iter}}$ , featuring more edges than  $\mathcal{H}_{\mathrm{PC}}$ .
We examine the trace plot of log posterior DAG probabilities from ADR, PARNI-DAG and Partition MCMC using both skeletons  $\mathcal{H}_{\mathrm{PC}}$  and  $\mathcal{H}_{\mathrm{iter}}$  respectively (Figure 2). We ran ADR, PARNI-DAG and Partition MCMC for 80,000, 20,000 and 20,000 iterations respectively. Notice that, under both  $\mathcal{H}_{\mathrm{PC}}$  and  $\mathcal{H}_{\mathrm{iter}}$ , PARNI-DAG converges very quickly to the high-probabilistic region, whereas ADR takes longer. Partition MCMC under  $\mathcal{H}_{\mathrm{PC}}$  cannot reach the same high-probability region as it targets a posterior distribution restricted to  $\mathcal{H}_{\mathrm{PC}}$ , which is biased if  $\mathcal{H}_{\mathrm{PC}}$  excludes true positive edges. Under the more populated  $\mathcal{H}_{\mathrm{iter}}$  instead it converges, but still slower than PARNI-DAG. When comparing median MSEs on PEPs on this dataset, we considered 20 replications with each individual chain being run for 15 minutes. As for Order and Partition MCMC,  $\mathcal{H}_{\mathrm{PC}}$  seems to lead to slightly better performance than  $\mathcal{H}_{\mathrm{iter}}$ . This might be due to the fact that in this high-dimensional example  $\mathcal{H}_{\mathrm{iter}}$  includes considerably more DAGs than  $\mathcal{H}_{\mathrm{PC}}$ , so Order MCMC, and even more so Partition MCMC, take longer to explore the space. As for ADR and PARNI-DAG, the choice of  $\mathcal{H}$  does not have a big impact, as in their case this is used only to warm-start the chain and not to restrict the space. PARNI-DAG here outperforms ADR and it is competitive to Order MCMC as well, without introducing biases stemming from search space restriction and sampling.
4.2 DAG learning accuracy
Lastly, we compare performance of PARNI-DAG with other main methods in recovering the underlying true DAG on four different real-world graph structures. These graph structures are taken from the
Figure 2: gsim100 dataset: trace plots of log posterior model probabilities. All algorithms were ran for the same CPU-time and thinned to 20,000 measures.
skeleton	algorithm	MSE
HC	ADR	2.91
PARNI-DAG	1.83
Order MCMC	1.79
Partition MCMC	12.75
Hiter	ADR	2.37
PARNI-DAG	1,72
Order MCMC	1.84
Partition MCMC	31.16
Table 2: table gsim100 dataset: time-normalised median MSE  $(\times 10^{-7})$  on estimating posterior edge probabilities. (Lower is better.)
R package bnlearn<sup>4</sup> and include: i) ecoli70 dataset, with  $n = 46$  nodes and 70 edges; ii) magic-niab dataset, with  $n = 44$  nodes and 66 edges; iii) magic-irri dataset, with  $n = 64$  nodes and 102 edges; iv) arth150 dataset, with  $n = 107$  nodes and 150 edges. Given the graph structure and its associated parameters, we simulate  $N = 100$  i.i.d. observations for each of the four cases. The models we compare include: i) PC algorithm; ii) GES; iii) LiNGAM; iv) DiBS, a differentiable variational method introduced by Lorch et al. [2021]; v) DiBS+, a version of DiBS with the weighted particle mixture; vi) DECI, another non-linear additive noise model that uses variational inference introduced by Geffner et al. [2022]; vii) Order MCMC, on a restricted space defined by the skeleton of the PC algorithm; viii) Iterative MCMC, the version of Order MCMC run on a search space defined by starting with the PC skeleton, and iteratively expanding the possible parent sets until the score cannot be improved [Kuipers et al., 2022]; ix) Partition MCMC, on the same space defined by the iterative procedure in Iterative MCMC; x) ADR scheme and xi) PARNI-DAG, both with pre-tuning of the  $\eta$ 's using the Iterative-derived skeleton. All the score-based methods (GES, Order, Partition, Iterative) are assigned an edge-penalty parameter equal to  $2\log(n)$ , as suggested in [Kuipers et al., 2022], to adjust to potentially sparse graphs. Performance is measured via the Structural Hamming Distance between the true DAG and the estimated DAG, averaged over 20 replications of the experiment, for each dataset. Results are collected in Table 3. Notice how PARNI-DAG is either the best or the second best method across all four experiments. In particular, it statistically matches the performance of the resulting best MCMC scheme in the first three experiments featuring a lower number of nodes, while it significantly outperforms all the methods in the case with high number of nodes (arth150), as expected. All variational methods (DiBS, DiBS+ and DECI) struggle to detect edges in these low-sample setups as they require a stronger signal (or a larger sample, as featured in the experiments in Lorch et al. [2021], Geffner et al. [2022]) to perform well. Additionally, they are significantly slow to converge and run into numerical instability issues in the last large- $n$  dataset (arth150).
Although the PARNI-DAG proposal has demonstrated significant improvements over the state-of-the-art algorithms considered in the previous section, it still comes with a few limitations that can be addressed. The first limitation lies in the linearity assumption. In fact, in order to derive the closed-form marginal likelihood featured in the locally informed proposal, PARNI-DAG is currently constrained to the linear model assumption described in (1). This can potentially result in model misspecification in presence of non-linear relationships, which is a threat in cases where this hampers edges identifiability. The recent work of Liang et al. [2023] has shed light on how to efficiently extend the PARNI proposal to more general non-linear models, and can possibly be adjusted to Bayesian structure learning settings that employ Additive Noise Models (ANMs) [Hoyer et al., 2008a], where a closed-form marginal likelihood does not exist. The second limitation of the PARNI-DAG proposal is the computational cost associated with the likelihood evaluation, which scales at least linearly in the
4More information on: http://bnlearn.com/.
ecoli70	magic-niab	magic-irri	arth150
PC	62.65 ± 1.64	66.95 ± 0.82	103.75 ± 0.91	126.70 ± 2.10
GES	47.95 ± 1.72	65.50 ± 0.42	97.60 ± 0.99	129.70 ± 3.47
LiNGAM	107.4 ± 3.96	71.60 ± 1.34	111.25 ± 2.34	-
DiBS	71.05 ± 0.17	65.80 ± 0.24	102.35 ± 0.17	-
DiBS+	71.25 ± 0.28	64.70 ± 0.36	100.45 ± 0.40	-
DECI	70.95 ± 0.02	66.95 ± 0.02	102.00 ± 0.00	-
Order MCMC	42.85 ± 1.30	62.05 ± 1.52	92.90 ± 1.46	110.00 ± 1.96
Iterative MCMC	39.55 ± 2.15	65.00 ± 2.43	96.15 ± 2.11	119.45 ± 3.43
Partition MCMC	48.85 ± 2.76	66.20 ± 1.62	103.35 ± 1.81	140.75 ± 4.04
ADR	39.40 ± 1.27	64.25 ± 0.64	95.20 ± 1.16	143.15 ± 1.62
PARNI-DAG	37.55 ± 1.59	64.14 ± 0.55	91.85 ± 0.97	97.60 ± 2.16
Table 3: SHD averaged over 20 replications with associated  $95\%$  confidence intervals of the various models compared on the four datasets  $(n = 100)$ . All algorithms were ran for approximately the same CPU-time and the significantly best performing models are given in bold. (Lower is better.)
number of datapoints. The issue arises as the likelihood features both in the informed proposals and in the MH step. A viable solution to this problem consists in the possibility of coupling PARNI-DAG with optimal sub-sampling MCMC procedures, such as the ones presented in Korattikara et al. [2014] and Maclaurin and Adams [2014] (FireflyMC). These limitations associated with the PARNI-DAG proposal provide us with interesting future research directions, aimed at improving its scalability and capability to handle more complex data.
In this work, we proposed a novel MCMC sampler, PARNI-DAG, for a fully Bayesian approach to the problem of structure learning. PARNI-DAG samples directly from the space of DAGs and introduces improvements in terms of computational complexity and MCMC mixing, stemming from the nature of its adaptive random neighborhood informed proposal. This proposal facilitates moves to higher probability regions, and is particularly useful when dealing with high-dimensional settings with a high number of nodes, as demonstrated in the experimental section. In future work, we will consider adding global moves to the PARNI-DAG proposal, which re-generates a new DAG model every few iterations, to specifically tackle settings with highly correlated structures.
In addition, we are also interested in studying PARNI's theoretical mixing time bounds. In fact, although the mixing time bound for random walk proposals has been largely studied, similar results for the class of locally informed proposals are relatively under-developed, due to the complexities arising in the proposal distributions. Most results on mixing time in discrete sample spaces focus on the problem of Bayesian variable selection. It has been shown by Yang et al. [2016] that a random walk proposal (specifically the add-delete-swap proposal) can achieve polynomial-time mixing under mild conditions on the posterior distribution. Under similar conditions, Zhou et al. [2022] has shown that the mixing time of the Locally Informed and Thresholded proposal (LIT) does not depend on the number of covariates. The mixing time bound of MCMC samplers in Bayesian structure learning settings is notoriously harder to study, due to the higher complexity of the sample space of DAGs. The recent work of Zhou and Chang [2023] has shed light on the analysis of mixing time bounds for Bayesian structure learning as a generalisation of the results in Yang et al. [2016]. Zhou and Chang [2023] has proven that the mixing time of the Random Walk Greedy Equivalent Search (RW-GES) proposal is at most linear in the number of covariates and the number of datapoints. Moreover, they also presented the necessary conditions for posterior consistency in Bayesian structure learning. It has not been formally proven yet that the informed proposal achieves faster theoretical mixing time compared to the random walk proposal in the context of Bayesian structure learning, although the empirical experiments we conducted suggest so. Similarly, it is highly likely then that the PARNI-DAG proposal can also achieve dimension-free mixing, such as the LIT proposal. We leave the topic of theoretical mixing time bounds of the PARNI proposal in different applications (e.g., Bayesian variable selection, Bayesian structure learning) to future research.
SL is supported by EPSRC research grant EP/V055380/1. AC's research is funded by the Defence Science and Technology Laboratory (Dstl) which is an executive agency of the UK Ministry of Defence providing world class expertise and delivering cutting-edge science and technology for the benefit of the nation and allies. The research supports the Autonomous Resilient Cyber Defence (ARCD) project within the Dstl Cyber Defence Enhancement programme.