SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Variational inference for  $W$  We use the variational posterior  $q_{\phi}(\pmb{W}|\pmb{p})$  to approximate the true posterior  $p(\pmb{W}|\pmb{p},\Theta ,\pmb{D})$ . Specifically, we select an independent Bernoulli distribution with logits defined by the output of a neural network  $\mu_{\phi}(\pmb{p})$ :
$$
q _ {\phi} (\boldsymbol {W} | \boldsymbol {p}) = \prod_ {i j} \operatorname {B e r} \left(\mu_ {\phi} (\boldsymbol {p}) _ {i j}\right) \tag {14}
$$
To train  $q_{\phi}$ , we derive the corresponding evidence lower bound (ELBO):
$$
\operatorname {E L B O} (\phi) = \mathbb {E} _ {q _ {\phi} (\boldsymbol {W} | \boldsymbol {p})} [ \log p (\boldsymbol {D}, \boldsymbol {p}, \boldsymbol {\Theta} | \boldsymbol {W}) ] - D _ {\mathrm {K L}} \left[ q _ {\phi} (\boldsymbol {W} | \boldsymbol {p}) \| p (\boldsymbol {W}) \right] ]. \tag {15}
$$
where  $D_{\mathrm{KL}}$  is the Kullback-Leibler divergence. The derivation is in Appendix B.6. Algorithm 1 summarizes this inference procedure.
SG-MCMC with continuous relaxation Furthermore, we explore an alternative formulation that circumvents the need for variational inference. Instead, we employ SG-MCMC to sample  $\tilde{W}$ , a continuous relaxation of  $W$ , facilitating a fully sampling-based approach. For a detailed formulation, please refer to Appendix A. We report its performance in Appendix E.3, which surprisingly is inferior to SG-MCMC+VI. We hypothesize that coupling  $W, p$  through  $\mu_{\phi}$  is important since changes in  $p$  results in changes of the permutation matrix  $\sigma(p)$ , which should also influence  $W$  accordingly during posterior inference. However, through sampling  $\tilde{W}$  with few SG-MCMC steps, this change cannot be immediately reflected, resulting in inferior performance. Thus, we focus only on the performance of SG-MCMC+VI for our experiments.
Computational complexity Our proposed SG-MCMC+VI offers a notable improvement in computational cost compared to existing approaches, such as DIBS [43]. The computational complexity of
Algorithm 1 BayesDAG SG-MCMC+VI Inference
Input: dataset  $D$  ; prior  $p(\pmb {p},\pmb {W}),p(\Theta)$  ; SG-MCMC sampler Sampler; sampler hyperparameters   
 $\Psi$  network  $\mu_{\phi}(\cdot)$  ; training iteration  $T$    
Output: samples  $\{\Theta ,p\}$  and variational posterior  $q_{\phi}$    
Initialize  $\Theta^{(0)},\pmb{p}^{(0)},\phi$    
for  $t = 1\dots ,T$  do Sample  $W^{(t - 1)}\sim q_{\phi}(W|p^{(t - 1)})$  Evaluate  $\nabla_{\pmb{p},\Theta}U$  (Equations (12) and (13)) with  $\Theta^{(t - 1)},\pmb{p}^{(t - 1)},\pmb{W}^{(t - 1)}$ $\Theta^{(t)},\pmb{p}^{(t)} = \mathrm{Samp}$  ler  $(\nabla_{\pmb{p},\Theta}U;\Psi)$  if storing condition met then  $\{p,\Theta \} \gets p^{(t)},\Theta^{(t)}$  end if Maximize ELBO (Equation (15)) w.r.t.  $\phi$  with  $\pmb {p}^{(t)},\Theta^{(t)}$    
end for
our method is  $O(BN_{p} + N_{p}d^{3})$ , where  $B$  represents the batch size and  $N_{p}$  is the number of parallel SG-MCMC chains. This former term stems from the forward and backward passes, and the latter comes from the Hungarian algorithm, which can be parallelized to further reduce computational cost. In comparison, DIBS has a complexity of  $O(N_{p}^{2}N + N_{p}d^{3})$  with  $N \gg B$  being the full dataset size. This is due to the kernel computation involving the entire dataset and the evaluation of the matrix exponential in the DAG regularizer [77]. As a result, our approach provides linear scalability w.r.t.  $N_{p}$  with substantially smaller batch size  $B$ . Conversely, DIBS exhibits quadratic scaling in terms of  $N_{p}$  and lacks support for mini-batch gradients.
Bayesian causal discovery literature has primarily focused on inference in linear models with closed-form posteriors or marginalized parameters. Early works considered sampling directed acyclic graphs (DAGs) for discrete [15, 46, 33] and Gaussian random variables [21, 65] using Markov chain Monte Carlo (MCMC) in the DAG space. However, these approaches exhibit slow mixing and convergence [18, 32], often requiring restrictions on number of parents [38]. Alternative exact dynamic programming methods are limited to low-dimensional settings [36].
Recent advances in variational inference [75] have facilitated graph inference in DAG space, with gradient-based methods employing the NOTEARS DAG penalty [77].[3] samples DAGs from autoregressive adjacency matrix distributions, while [43] utilizes Stein variational approach [42] for DAGs and causal model parameters. [16] proposed a variational inference framework on node orderings using the gumbel-sinkhorn gradient estimator [47]. [17, 51] employ the GFlowNet framework [6] for inferring the DAG posterior. Most methods, except[43] are restricted to linear models, while [43] has high computational costs and lacks DAG generation guarantees compared to our method.
In contrast, quasi-Bayesian methods, such as DAG bootstrap [20], demonstrate competitive performance. DAG bootstrap resamples data and estimates a single DAG using PC [60], GES [14], or similar algorithms, weighting the obtained DAGs by their unnormalized posterior probabilities. Recent neural network-based works employ variational inference to learn DAG distributions and point estimates for nonlinear model parameters [11, 22].
In this section, we aim to study empirically the following aspects: (1) posterior inference quality of BayesDAG as compared to the true posterior when the causal model is identifiable only upto Markov Equivalence Class (MEC); (2) posterior inference quality of BayesDAG in high dimensional nonlinear causal models (3) ablation studies of BayesDAG and (4) performance in semi-synthetic and real world applications. The experiment details are included in Appendix D.
Baselines. We mainly compare BayesDAG with the following baselines: Bootstrap GES (BGES) [14, 20], BCD Nets [16], Differentiable DAG Sampling (DDS [11]) and DIBS [43].
6.1 Evaluation on Synthetic Data
Synthetic data. We evaluate our method on synthetic data, where ground truth graphs are known. Following previous work, we generate data by randomly sampling DAGs from Erdos-Renyi (ER) [19] or Scale-Free (SF) [5] graphs with per node degree 2 and drawing at random ground truth parameters for linear or nonlinear models. For  $d = 5$ , we use  $N = 500$  training, while for higher dimensions, we use  $N = 5000$ . We assess performance on 30 random datasets for each setting.
Metrics For  $d = 5$  linear models, we compare the approximate and true posterior over DAGs using Maximum Mean Discrepancy (MMD) and also evaluate the expected CPDAG Structural Hamming Distance (SHD). For higher-dimensional nonlinear models with intractable posterior, we compute the expected SHD (E
SHD), expected orientation F1 score (Edge F1) and negative log-likelihood of the held-out data (NLL). Our synthetic data generation and evaluation protocol follows prior work [3, 22, 43]. All the experimental details, including how we use cross-validation to select hyperparameters is in Appendix D.
Figure 2: Posterior inference on linear synthetic datasets with  $d = 5$ . Metrics are computed against the true posterior.  $\downarrow$  denotes lower is better.
6.1.1 Comparison with True Posterior
Capturing equivalence classes and quantifying epistemic uncertainty are crucial in Bayesian causal discovery. We benchmark our method against the true posterior using a 5-variable linear SCM with unequal noise variance (identifiable upto MEC [53]). The true posterior over graphs  $p(\pmb{G} \mid \pmb{D})$  can be computed using the BGe score [23, 39]. Results in Figure 2 show that our method outperforms DIBS and DDS in both ER and SF settings. Compared to BCD, we perform better in terms of MMD in ER but worse in SF. We find that BGES performs very well in low-dimensional linear settings, but suffers significantly in more realistic nonlinear settings (see below).
6.1.2 Evaluation in Higher Dimensions
We evaluate our method on high dimensional scenarios with nonlinear relations. Our approach is the first to attempt full posterior inference in nonlinear models using permutation-based methods. Results for  $d = 30$  variables in Figure 3 demonstrate that BayesDAG significantly outperforms other permutation-based approaches and DIBS in most of the metrics. For  $d = 50$ , BayesDAG performs comparably to
DIBS in ER but a little worse in SF. However, our method achieves better NLL on held-out data compared to most baselines including DIBS for  $d = 30, 50$ , ER and SF settings. Only DDS gives better NLL for  $d = 30$  ER setting, but this doesn't translate well to other metrics and settings. We additionally evaluate on  $d \in \{70, 100\}$  variables (Table 1). We find that our method consistently outperforms the baselines with  $d = 70$  and in terms of E-SHD with  $d = 100$ . Full results are presented in Appendix E.2. Competitive performance for  $d > 50$  in nonlinear settings further demonstrates the applicability and computational efficiency of the proposed approach. In contrast, the only fully Bayesian nonlinear method, DIBS, is not computationally efficient to run for  $d > 50$ .
Table 1: E-SHD (with  $95\%$  CI) for ER graphs in higher dimensional nonlinear causal models. DIBS becomes computationally prohibitive for  $d > 50$ .
d=70	d=100
BGES	355.77 ± 18.02	563.02 ± 27.21
BCD	217.05 ± 9.58	362.66 ± 29.18
DIBS	N/A	N/A
BaDAG	143.70 ± 11.61	295.92 ± 24.67
$(d = 30)$
$(d = 30)$
Figure 3: Posterior inference of both graph and functional parameters on synthetic datasets of nonlinear causal models with  $d = 30$  and  $d = 50$  variables. BayesDAG gives best results across most metrics and outperforms other permutation based approaches (BCD and DDS). We found DDS to perform significantly worse in terms of E-SHD and thus has been omitted for clarity.  $\downarrow$  denotes lower is better and  $\uparrow$  denotes higher is better.
$(d = 50)$
6.2 Ablation Studies
We conduct ablation studies on our method using the nonlinear ER  $d = 30$  dataset.
[ \text{Initialization } p \text{ scale} ] Figure 4a investigates the influence of the initialized scale of  $p$ . We found that the performance is the best with  $\alpha = 0.01$  or  $10^{-5}$ , and deteriorates with increasing scales. This is because with larger initialization scale, the absolute value of the  $p$  is large. Longer SG-MCMC updates are needed to reverse the node potential order, which hinders the exploration of possible permutations, resulting in the convergence to poor local optima.
Number of SG-MCMC chains We examine the impact of the number of parallel SG-MCMC chains in Figure 4b. We observe that it does not have a significant impact on the performance, especially with respect to the E-SHD and Edge F1 metrics.
Injected noise level for SG-MCMC In Figures 4c and 4d, we study the performance differences arising from various injected noise levels for  $p$  and  $\Theta$  in the SG-MCMC algorithm (i.e.  $s$  of the SG-MCMC formulation in Appendix C). Interestingly, the noise level of  $p$  does not impact the performance as much as the level of  $\Theta$ . Injecting noise helps improve the performance, but a smaller noise level should be chosen for  $\Theta$  to avoid divergence from optima.
6.3 Application 1: Evaluation on Semi-Synthetic Data
We evaluate our method on the SynTReN simulator [67]. This simulator creates synthetic transcriptional regulatory networks and produces simulated gene expression data that approximates real experimental data. We use five different simulated datasets provided by [40] with  $N = 500$  samples each. Table 2 presents the results of all the methods. We find that our method recovers the true network much better in terms of E-SHD as well as Edge F1 compared to baselines.
6.4 Application 2: Evaluation on Real Data
We also evaluate on a real dataset which measures the expression level of different proteins and phospholipids in human cells (called the Sachs Protein Cells Dataset) [59]. The data corresponds to a network of protein-protein interactions of 11 different proteins with 17 edges in total among them. There are 853 observational samples in total, from which we bootstrap 800 samples of 5 different datasets. It is to be noted that this data does not necessarily adhere to the additive noise and DAG assumptions, thereby having significant model misspecification. Results in Table 2 demonstrate that our method performs well as compared to the baselines even with model misspecification, proving the suitability of the proposed framework for real-world settings.
Table 2: Results (with  $95\%$  confidence intervals) on Syntren (semi-synthetic) and Sachs Protein Cells (real-world) datasets. For Syntren, results are averaged over 5 different datasets. For Sachs, results are averaged over 5 different restarts.  $\downarrow$  denotes lower is better and  $\uparrow$  denotes higher is better.
Syntren (d=20)	Sachs Protein Cells (d=11)
E-SHD (↓)	Edge F1 (↑)	E-SHD (↓)	Edge F1 (↑)
BGES	66.18 ± 9.47	0.21 ± 0.05	16.61 ± 0.44	0.22 ± 0.02
DDS	134.37 ± 4.58	0.13 ± 0.02	34.90 ± 0.73	0.21 ± 0.02
BCD	38.38 ± 7.12	0.15 ± 0.07	17.05 ± 1.93	0.20 ± 0.08
DIBS	46.43 ± 4.12	0.16 ± 0.02	22.3 ± 0.31	0.20 ± 0.01
BaDAG	34.21 ± 2.82	0.20 ± 0.02	18.92 ± 1.0	0.26 ± 0.04
(a) Posterior inference with different initialized  $\pmb{p}$  scale.
(b) Posterior inference with different number of parallel SG-MCMC chains.
(c) Posterior inference with different level of injected noise scale for  $\pmb{p}$ .
Figure 4: Ablation study of posterior inference quality of BayesDAG on  $d = 30$  ER synthetic dataset.
(d) Posterior inference with different level of injected noise scale for  $\Theta$
In this work, we propose BayesDAG, a novel, scalable Bayesian causal discovery framework that employs SG-MCMC (and VI) to infer causal models. We establish the validity of performing Bayesian inference in the augmented  $(W,p)$  space and demonstrate its connection to permutation-based DAG learning. Furthermore, we provide two instantiations of the proposed framework that offers direct DAG sampling and model-agnosticism to linear and nonlinear relations. We demonstrate superior inference accuracy and scalability on various datasets. Future work can address some limitations: (1) designing better variational networks  $\mu_{\phi}$  to capture the complex distributions of  $W$  compared to the simple independent Bernoulli distribution; (2) improving the performance of SG-MCMC with continuous relaxation (Appendix A), which currently does not align with its theoretical advantages compared to the SG-MCMC+VI counterpart.
Acknowledgements. The authors would like to thank Colleen Tyler, Maria Defante, and Lisa Parks for conversations on real-world use cases that inspired this work. YA and SB are thankful for the Swedish National Computing's Berzelius cluster for providing resources that were helpful in running some of the baselines of the paper. In addition, the authors would like to thank the anonymous reviewers for their feedback.
[1] Ryan Prescott Adams and Richard S Zemel. Ranking via sinkhorn propagation. arXiv preprint arXiv:1106.1925, 2011.
[2] Raj Agrawal, Chandler Squires, Karren Yang, Karthikeyan Shanmugam, and Caroline Uhler. Abcd-strategy: Budgeted experimental design for targeted causal structure discovery. In The
[61] Jost Tobias Springenberg, Aaron Klein, Stefan Falkner, and Frank Hutter. Bayesian optimization with robust bayesian neural networks. Advances in neural information processing systems, 29, 2016.
[62] Haoran Sun, Hanjun Dai, Bo Dai, Haomin Zhou, and Dale Schuurmans. Discrete Langevin sampler via Wasserstein gradient flow. arXiv preprint arXiv:2206.14897, 2022.
[63] Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational bayesian neural networks. arXiv preprint arXiv:1903.05779, 2019.
[64] Panagiotis Tugas, Yashas Annadani, Andrew Jesson, Bernhard Scholkopf, Yarin Gal, and Stefan Bauer. Interventions, where and how? experimental design for causal models at scale. Advances in neural information processing systems, 36, 2022.
[65] Simon Tong and Daphne Koller. Active learning for structure in bayesian networks. In International joint conference on artificial intelligence, volume 17, pages 863-869. CiteSeer, 2001.
[66] Brian Trippe and Richard Turner. Overpruning in variational bayesian neural networks. arXiv preprint arXiv:1801.06230, 2018.
[67] Tim Van den Bulcke, Koenraad Van Leemput, Bart Naudts, Piet van Remortel, Hongwu Ma, Alain Verschoren, Bart De Moor, and Kathleen Marchal. Syntren: a generator of synthetic gene expression data for design and analysis of structure learning algorithms. BMC bioinformatics, 7:1-12, 2006.
[68] Chikako Van Koten and AR Gray. An application of bayesian network for predicting object-oriented software maintainability. Information and Software Technology, 48(1):59-67, 2006.
[69] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688, 2011.
[70] Nanyang Ye, Zhanxing Zhu, and Rafal K Mantiuk. Langevin dynamics with continuous tempering for training deep neural networks. arXiv preprint arXiv:1703.04379, 2017.
[71] Yue Yu, Jie Chen, Tian Gao, and Mo Yu. Dag-gnn: Dag structure learning with graph neural networks. In International Conference on Machine Learning, pages 7154-7163. PMLR, 2019.
[72] Yue Yu, Tian Gao, Naiyu Yin, and Qiang Ji. Dags with no curl: An efficient dag structure learning approach. In International Conference on Machine Learning, pages 12156-12166. PMLR, 2021.
[73] Giacomo Zanella. Informed proposals for local mcmc in discrete spaces. Journal of the American Statistical Association, 115(530):852-865, 2020.
[74] Valentina Zantedeschi, Luca Franceschi, Jean Kaddour, Matt J Kusner, and Vlad Niculae. Dag learning on the permutahedron. arXiv preprint arXiv:2301.11898, 2023.
[75] Cheng Zhang, Judith Butepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE transactions on pattern analysis and machine intelligence, 41(8):2008-2026, 2018.
[76] Ruqi Zhang, Xingchao Liu, and Qiang Liu. A Langevin-like sampler for discrete distributions. In International Conference on Machine Learning, pages 26375-26396. PMLR, 2022.
[77] Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing. Dags with no tears: Continuous optimization for structure learning. Advances in Neural Information Processing Systems, 31, 2018.
Appendix - BayesDAG: Gradient-Based Posterior Inference for Causal Discovery
A Joint Inference with SG-MCMC
In this section, we propose an alternative formulation that enables a joint inference framework for  $p, W, \Theta$  using SG-MCMC, thereby avoiding the need for variational inference for  $W$ .
We adopt a continuous relaxation of  $\mathbf{W}$ , similar to [43], by introducing a latent variable  $\tilde{\mathbf{W}}$ . The graphical model is illustrated in Figure 5. We can define
$$
p (\boldsymbol {W} | \tilde {\boldsymbol {W}}) = \prod_ {i, j} p \left(W _ {i j} | \tilde {W} _ {i j}\right) \tag {16}
$$
with  $p(W_{ij} = 1|\tilde{W}_{ij}) = \sigma (\tilde{W}_{ij})$  where  $\sigma (\cdot)$  is the sigmoid function. In other words,  $\tilde{W}_{ij}$  defines the existence logits of  $W_{ij}$ .
Figure 5: Graphical model with latent variable  $\tilde{W}$ .
With the introduction of  $\tilde{W}$ , the original posterior expectations of  $p(\boldsymbol{p},\boldsymbol{W},\Theta|\boldsymbol{D})$ , e.g. during evaluation, can be translated using the following proposition.
Proposition A.1 (Equivalence of posterior expectation). Under the generative model Figure 5, we have
$$
\mathbb {E} _ {p (\boldsymbol {p}, \boldsymbol {W}, \boldsymbol {\Theta} | \boldsymbol {D})} [ \boldsymbol {f} (\boldsymbol {G} = \tau (\boldsymbol {p}, \boldsymbol {W}), \boldsymbol {\Theta}) ] = \mathbb {E} _ {p (\boldsymbol {p}, \tilde {\boldsymbol {W}}, \boldsymbol {\Theta})} \left[ \frac {\mathbb {E} _ {p (\boldsymbol {W} | \tilde {\boldsymbol {W}})} [ f (\boldsymbol {G} , \boldsymbol {\Theta}) p (\boldsymbol {D} , \boldsymbol {\Theta} | \boldsymbol {p} , \boldsymbol {W}) ]}{\mathbb {E} _ {p (\boldsymbol {W} | \tilde {\boldsymbol {W}})} [ p (\boldsymbol {D} , \boldsymbol {\Theta} | \boldsymbol {p} , \boldsymbol {W}) ]} \right] \tag {17}
$$
where  $\pmb{f}$  is the target quantity.
This proof is in Appendix B.3.
With this proposition, instead of sampling  $W$ , use SG-MCMC to draw  $\tilde{W}$  samples. Similar to Section 4.2, to use SG-MCMC for  $p, \tilde{W}, \Theta$ , we need their gradient information. The following proposition specifies the required gradients.
Proposition A.2. With the generative model defined as Figure 5, we have
$$
\begin{array}{l} \nabla_ {\pmb {p}, \Theta , \tilde {\pmb {W}}} U (\pmb {p}, \tilde {\pmb {W}}, \Theta) = - \nabla_ {\pmb {p}} \log p (\pmb {p}) - \nabla_ {\Theta} \log p (\Theta) \\ - \nabla_ {\tilde {\boldsymbol {W}}} \log p (\tilde {\boldsymbol {W}}) - \nabla_ {\boldsymbol {p}, \boldsymbol {\Theta}, \tilde {\boldsymbol {W}}} \log \mathbb {E} _ {p (\boldsymbol {W} | \tilde {\boldsymbol {W}})} [ p (\boldsymbol {D} | \boldsymbol {W}, \boldsymbol {p}, \boldsymbol {\Theta}) ] \tag {18} \\ \end{array}
$$
The proof is in Appendix B.4.
With these gradients, we can directly plug in existing SG-MCMC samplers to draw samples for  $\pmb{p}$ ,  $\tilde{\pmb{W}}$ , and  $\Theta$  in joint inference (Algorithm 2).
B Theory
For completeness, we recite the theorem here.
Theorem 3.1 (Equivalence of inference in  $(\mathbf{W},\mathbf{p})$  and binary DAG space). Assume graph  $\mathbf{G}$  is a binary adjacency matrix representing a DAG and node potential  $\mathbf{p}$  does not contain the same values, i.e.  $p_i\neq p_j\forall i,j$ . Then, with the induced joint observational distribution  $p(D,G)$ , dataset  $\mathbf{D}$  and a corresponding prior  $p(G)$ , we have
$$
p (\boldsymbol {G} | \boldsymbol {D}) = \int p _ {\tau} (\boldsymbol {p}, \boldsymbol {W} | \boldsymbol {D}) \mathbb {1} (\boldsymbol {G} = \tau (\boldsymbol {W}, \boldsymbol {p})) d \boldsymbol {W} d \boldsymbol {p} \tag {19}
$$
if  $p(\pmb{G}) = \int p_{\tau}(\pmb{p}, \pmb{W}) \mathbb{1}(\pmb{G} = \tau(\pmb{W}, \pmb{p})) d\pmb{W} d\pmb{p}$ , where  $p_{\tau}(\pmb{W}, \pmb{p})$  is the prior,  $\mathbb{1}(\cdot)$  is the indicator function and  $p_{\tau}(\pmb{p}, \pmb{W}|D)$  is the posterior distribution over  $\pmb{p}$ ,  $\pmb{W}$ .
Algorithm 2 Joint inference
Input: dataset  $D$  , prior  $p(\pmb {p},\tilde{\pmb{W}},\Theta)$  , SG-MCMC sampler update Sampler(  $\cdot$  ); sampler hyperpa  
rameter  $\Psi$  ; training steps  $T$    
Output: posterior samples  $\{p,\tilde{W},\Theta \}$    
Initialize  $p_0,\tilde{W}_0,\Theta_0$    
for  $t = 1,\dots ,T$  do Evaluate gradient  $\nabla_{\pmb{p}_{t - 1},\tilde{\pmb{W}}_{t - 1},\Theta_{t - 1}}U$  based on Equation (18). Update samples  $\pmb {p}_t,\tilde{\pmb{W}}_t,\Theta_t = \mathrm{Samp}$  ler  $(\nabla_{\pmb {p}_{t - 1},\tilde{\pmb{W}}_{t - 1},\Theta_{t - 1}}U;\Psi)$  if storing condition met then  $\{p,\tilde{W},\Theta \} \gets p_t,\tilde{W}_t,\Theta_t$  end if   
end for
To prove this theorem, we first prove the following lemma stating the equivalence of  $\tau$  (Equation (5)) to binary DAG space.
Lemma B.1 (Equivalence of  $\tau$  to DAG space). Consider  $d$  random variables, a node potential vector  $\pmb{p} \in \mathbb{R}^d$  and a binary matrix  $\pmb{W} \in \{0, 1\}^{d \times d}$ . Then the following holds:
(a) For any  $\pmb{W} \in \{0,1\}^{d \times d}$ ,  $\pmb{p} \in \mathbb{R}^d$ ,  $\pmb{G} = \tau(\pmb{W}, \pmb{p})$  is a DAG.
(b) For any DAG  $G \in \mathbb{D}$ , where  $\mathbb{D}$  is the space of all DAGs, there exists a corresponding  $(W, \pmb{p})$  such that  $\tau(W, \pmb{p}) = G$ .
Proof. The main proof directly follows the theorem 2.1 in [72]. For (a), we show the output from  $\tau(W, p)$  must be a DAG. By leveraging the Lemma 3.4 in [72], we can easily obtain that Step(grad  $p$ ) emits a binary adjacency matrix representing a DAG. The only difference is that we replace the  $\mathrm{ReLU}(\cdot)$  with  $\mathrm{Step}(\cdot)$  but the conclusion can be directly generalized.