BayesDAG: Gradient-Based Posterior Inference for Causal Discovery
Authors: Yashas Annadani
Nick Pawlowski
Joel Jennings
Stefan Bauer
Cheng Zhang
Wenbo Gong
Venue: NeurIPS 2023 poster (2023)
Keywords: Causal Discovery
Structure Learning
Bayesian Inference
Variational Inference
MCMC
Generative Model
TL;DR: We propose a Bayesian causal discovery method that samples Directed Acyclic Graphs (DAG) from an augmented space with gradient-based MCMC which is applicable to high-dimensional linear and nonlinear causal models.

Abstract: Bayesian causal discovery aims to infer the posterior distribution over causal models from observed data, quantifying epistemic uncertainty and benefiting downstream tasks. However, computational challenges arise due to joint inference over combinatorial space of Directed Acyclic Graphs (DAGs) and nonlinear functions. Despite recent progress towards efficient posterior inference over DAGs,  existing methods are either limited to variational inference on node permutation matrices for linear causal models, leading to compromised inference accuracy, or continuous relaxation of adjacency matrices constrained by a DAG regularizer, which cannot ensure resulting graphs are DAGs. In this work, we introduce a scalable Bayesian causal discovery framework based on a combination of stochastic gradient Markov Chain Monte Carlo (SG-MCMC) and Variational Inference (VI) that overcomes these limitations. Our approach directly samples DAGs from the posterior without requiring any DAG regularization, simultaneously draws function parameter samples and is applicable to both linear and nonlinear causal models. To enable our approach, we derive a novel equivalence to the permutation-based DAG learning, which opens up possibilities of using any relaxed gradient estimator defined over permutations. To our knowledge, this is the first framework applying gradient-based MCMC sampling for causal discovery. Empirical evaluation on synthetic and real-world datasets demonstrate our approach's effectiveness compared to state-of-the-art baselines.

BayesDAG: Gradient-Based Posterior Inference for Causal Discovery
Yashas Annadani 1,3,4 Nick Pawlowski
Joel Jennings
Stefan B
Cheng Zhang? Wenbo Gong" 2
KTH Royal Institute of Technology, Stockholm Microsoft Research 3 Helmholtz AI, Munich 4 TU Munich
Abstract
Bayesian causal discovery aims infer the posterior distribution over causal models from observed data, quantifying epistemic uncertainty and benefiting downstream tasks. However, computational challenges arise due to joint inference over combinatorial space of Directed Acyclic Graphs (DAGs) and nonlinear functions. Despite recent progress towards efficient posterior inference over DAGs, existing methods are either limited to variational inference node permutation matrices for linear causal models, leading to compromised inference accuracy, or continuous relaxation of adjacency matrices constrained by DAG regularizer, which cannot ensure resulting graphs are DAGs. In this work, we introduce scalable Bayesian causal discovery framework based stochastic gradient Markov Chain Monte Carlo (SG-MCMC) and Variational Inference (VI) that overcomes these limitations. Our approach directly samples DAGs from the posterior without requiring any DAG regularization, simultaneously draws function parameter samples and is applicable to both linear and nonlinear causal models. To enable our approach, we derive a novel equivalence to the permutation--asee DAG learning, which opens up possibilities of using any relaxed gradient estimator defined over permutations. To our knowledge, this is the first framework applying gradient-based MCMC sampling for causal discovery. Empirical evaluation synthetic and real-world datasets demonstrate our approach's effectiveness compared to wfae-of-the-wee baselines.
combination of
Introduction
The quest for discovering causal relationships in data-generating processes lies at the heart of empirical sciences and decision-making [56] 59 68). Structural Causal Models (SCMs) [52] and their associated Directed Acyclic Graphs (DAGs) provide robust mathematical framework for modeling such relationships. Knowledge of the underlying SCM and its corresponding DAG permits predictions of unseen interventions and causal reasoning, thus making causal discovery - learning an unknown SCM and its associated DAG from observed data subject of extensive research [54] 60).
In contrast traditional methods that infer single graph or its Markov equivalence class (MFC) 14 60), Bayesian causal discovery [2] 33 [5] aims to infer posterior distribution over SCMs and their DAGs from observed data. This approach encapsulates the epistemic uncertainty, degree of confidence in every causal hypothesis, which is particularly valuable for real-world applications when data scarce. also beneficial for downstream tasks such as experimental design
'Equal contribution. Work done during internship Microsoft Research. venbogongemiccoosof..com
Correspondence
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

The central challenge in Bayesian causal discovery lies in inferring the posterior distribution over the union of the exponentially growing (discrete) DAGs and (continuous) function parameters. Prior works have used Markov Chain Monte Carlo (MCMC) to directly sample DAGs or bootstrap but these methods are typically limited to linear models which admit closed-form marginalization over continuous parameters. Recent advances have begun to utilize gradient information for more efficient inference. These approaches are either: (1) DAG regularizer-basee based methods, c.g. DIBS [43] which use continuous relaxation of adjacency matrices together with DAG regularizer [77]. But DIBS formulation fails to model edge co-dependencies and suffer from inference quality due to its inference engine (Stein variational gradient descent) ( 29] Additionally, all DAG regularizer based methods cannot guarantee DAG generation; (2) permutation based DAG learning. which directly infers permutation matrices and guarantees to generate DAGs. However, existing works focus on using only variational inference [11 16], which may suffer from and is sometimes restricted to only linear models [16].
traditional discovery methods [14 49) GS
inaccurate inference quality 28 61 66
In this work, we introduce BayesDAG, gradient-based Bayesian causal discovery framework that overcomes the above limitations. Our contributions are:
1. We prove that an augmented space of edge beliefs and node potentials (W, p), similar to NoCurl [72]. permits equivalent Bayesian inference in DAG space without the need for any regularizer. (Section 3.1)
2. We derive an equivalence relation from this augmented space to permutation- DAG learning which provides general framework for gradient-based posterior inference. (Sec
tion 3.2)
3. Based on this general framework, we propose a scalable Bayesian causal discovery that is model-agnostic for linear and non-linear cases and also offers improved inference quality.
We instantiate our approach through two formulations: (1) combination of SG-MCMC and VI (2) SG-MCMC with continuous relaxation. (Section 4)
4. We demonstrate the effectiveness of our approach in providing accurate Bayesian inference quality and superior causal discovery performance with comprehensive empirical evaluations various datasets. We also demonstrate that method can be easily scaled to 100 variables with nonlinear relationships. (Sectioon 6)
2
Background
Causal Graph and Structural Causal Model Consider a data generation process with variables X E Rd. The causal relationships among these variables is represented by a Structural Causal Model (SCM) which consists of a set of structural equations [54] where each variable X; is a function of its direct causes Xp and exogenous noise variable e, with distribution P..:
X₁ = fi(Pps
(I)
These equations induce causal graph G = CVV. E), comprising node set V with [V| = d indexing
the variables X and directed edge set E. If directed edge Eij E exists between node pair wa, ua E V (i.e., + U;), we say that X; causes Xj or X; is the parent of Xj. We use the binary adjacency matrix G E {0, 1) dddd to represent the causal graph, where the entry Gūj 0y + Uy. A standard assumption in causality is that the structural assignments are acyclic and the induced causal graph is DAG [2] which we adopt in this work. We further assume that the SCM is causally sufficient i.e. all variables are measurable and exogenous noise variables e, are mutually independent. Throughout this work, we consider special form of SCM called Gaussian additive noise model (ANM):
denotes
x= (Xp + where e, ~ ~1(,00
(2)
If the functions are not linear or constant in any of its arguments, the Gaussian ANM is structurally identifiable 33 55
Bayesian Causal Discovery Given a dataset D = (I) (.V) (M) with i.i.d observations, un-
derlying graph G and SCM parameters O they induce unique joint distribution p(D, O, G) O θ) with the prior p(G, θ) and likelihood G θ) (1]. Under finite data
and/or limited identifiability of SCM (c.g upto MEC), it is desirable to have accurate uncertainty

---

estimation for downstream decision making rather than inferring a single SCM and its graph (for e.g. with maximum likelihood estimate). Bayesian causal discovery therefore aims to infer the posterior O = D, O /) D). However, this posterior is intractable due to the super-exponential growth of the possible DAGs G [58] and continuously valued model parameters O in nonlinear functions. VI [75] or SG-MCMC [24 45 are two types of methods developed to tackle general Bayesian inference problems, but adaptations are required for Bayesian causal discovery.
NoCurl Characterization Inferring causal graphs is challenging due to the DAG constraint. Previous works [22 25 40 43 [1]] directly infer adjacency matrix with the DAG regularizer [77] However, it requires an annealing schedule, resulting in slow convergence, and no guarantees on generating DAGs, Recently, 722 introduced NoCurl, a novel characterization of the weighted DAG space. They define potential p; R for each node i, grouped as potential vector p R" Further, a gradient operator on mapping it to a symmetrie matrix is introduced:
(grad (i. p: Pj
(3)
Based the above operation, a mapping that directly maps from the augmented space (W, .p) to the DAG space -((--)) Rdxd x Rdd + Raxd was proposed:
<W. = W O ReLU(grad
where ReLU(-) is the ReLU activation function and W is skew-symmetric continuously weighted matrix. This formulation is complete (Theorem 2.1 in (22), as any continuously weighted DAG can be represented by (W, p) pair and vice versa. NoCurl translates the learning of a single weighted DAG to a corresponding (W. p) pair. However, direct gradient-based optimization is challenging due to a highly non-convex loss landscape, which leads to the reported failure 772
Although NoCurl appears suitable for our purpose, the failure in directly learning suggests non-trivial optimizations. We hypothesize that this arises from the continuously weighted matrix W. In the following, we introduce our proposed parametrization inspired by NoCurl to characterize the binary DAG adjacency matrix.
3 Sampling the DAGs
In this section, we focus on the Bayesian inference over binary DAGs through novel mapping, r(W, , p)) a modification of NoCurl. We establish the validity of performing Bayesian inference within (W, p) space utilizing (Section 3.1 However, yields uninformative gradient during back- propagation, a challenge we overcome by deriving an equivalent formulation based on permutation- based DAG learning, thereby enabling the use of relaxed gradient estimators (Section 3.2)
3.1 Bayesian Inference in W.p Space
The NoCurl formulation (qquation (4)) focuses on learning single weighted DAG, which is not directly useful for our purpose. We need to address two key questions: (1) considering only binary adjacency matrices without weights; (2) ensuring Bayesian inference in (W, pp) is valid.
We note that the proposed transformation in NoCurl (Equation (4) can be hard to optimize for gives a fully connected DAG. The main purpose of W matrix therefore is to disable the edges. Continuous W requires thresholding to properly disable the edges, since is hard for a continuous matrix to learn exactly 0 during the optimization: (ii) ReLu p) and W are both continuous valued matrices. Thus, learning of the edge weights and DAG structure are not explicitly separated, resulting in complicated non-convex optimizations" Parameterizing the search space in terms of binary adjacency matrices significantly simplifies the optimization complexity as the aforementioned issues are cireumvented. Therefore, we introduce a modification 7: (0, 1) dxd x Rd + {0, [0,1)dxx dxd:;
the following reasons: (i) ReLU
(W, p) W @ Step(grad p)
(5)
where we abuse the term W for binary matrices, and replace ReLU(.) with Step(-). W acts as mask to disable the edge existence. Thus, due to the Step, t can only output binary adjacency matrix.
See discussion below Eq. 3 in (72 for more details.
3

---

Next, we show that performing Bayesian inference in such augmented (W.P) .p) space is valid, i.e., using the posterior p(W, D) to replace D)) This differs from NoCurl, which focuses on single graph rather than the validity for Bayesian inference, requiring a new theory for soundness.
Theorem 3.1 (Equivalence of inference in (W, p) and binary DAG space), Assume graph G is binary adjacency matrix representing a DAG and node potential p does not contain the same values, L.e. Ps # Pj vi, j. Then, with the induced joint observational distribution pl D, G), dataset D, and corresponding prior P(G), we have
P(G = P (I.
dp
(6)
T W) I (WW .P) where pr (W.)) .p) the prior 1(-) is the indicator function, and (p W (D) is the posterior distribution over p, W.
Refer to Appendix B.I for detailed proof.
This theorem guarantees that instead of performing inference directly in the constrained space (i.e. DAG space), we can apply Bayesian inference less complex (W, p) space where W e 10,1) and p IRd without explicit constraints.
For inference of p. we adopt a sampling buneling-based approach, which is asymptotically accurate (45] In particular, we consider SG-MCMC (refer to [Section 4), which avoids the expensive Metropolis- Hastings acceptance step and scales to large datasets. We emphasize that any other suitable sampling algorithms can be directly plugged in, thanks the generality of the framework.
However, the mapping does not provide meaningful gradient information for p due to the piecewise constant Step(-) function, which is required by SG-MCMC.
3.2 Equivalent Formulation
In this section, we address the above issue by deriving an equivalence to a permutation learning prob- lem. This alternative formulation enables various techniques that can approximate the gradient of p.
Intuition The node potential implicitly defines topological ordering through the mapping Step(grad(.))- In particular, grad(.) outputs skew-symmetric adjacency matrix, where each entry specifies the potential difference between nodes. Step (-rad())) zeros out the negative potential differences (i.e. pt S Pý). and only permits the edge direction from higher potential to the lower one (i.c. > Pý). This implicitly defines sorting operation based on the descending node potentials, involving permutation matrix.
which can cast as a particular arg max problem g 37 50 74
Alternative formulation We define L E (0, 1]dxd matrix with lower triangular part to be 1, and vector 0 = [l.....d] .d. We propose the following formulation:
G= W @ (o()))
where o arg max p" (a'o)
o'E
Here, Σg represents the space of all dimensional permutation matrices. The following theorem states the equivalence of this formulation to EEEuuation (5)
Theorem 3.2 (Equivalence to NoCurl formulation). Assuming the conditions in Theorem 3./ are satisfied. Then, for given (W, .p), we have
G W c Step(gradp) -Wo (o(p)
where G is DAG and (p) is defined in Equation (7)
Refer to Appendix B.2 for details.
This theorem translates our proposed operator Step (tep(grad(p)) into finding a corresponding permuta- tion matrix (p). Although this does not directly solve the uninformative gradient, it opens the door for approximating this gradient with the tools from the differentiable permutation literature S 47 50)
For simplicity, we adopt the Sinkhorn approach [47, but we emphasize that this equivalence general enough that any past or future approximation methods can be casily applied.

---

Sinkhorn operator The Sinkhorn operator S(M) on matrix M [ is defined as sequence of row and column normalizations, each called Sinkhorn iteration.
47 showed that the non-differentiable arg max problem
arg max (a', M) a'E
(8)
can be relaxed through entropy regularizer with its solution being expressed by S(-). In particular, they showed that S /t) = arg maxx'e., (a', M) + th(o'), where ((.) is the entropy function. This regularized solution converges to the solution of Equation (8) when 0, i.c. lim,->0 S(M /t).
Since the Sinkhorn operator is differentiable, (MM (t) can be viewed as differentiable approxima- tion Equation (8) which can be used to obtain the solution Equattio (7) Specifically, we have
arg max p" p0 (o'o) arg max(o', po")) = lim S( po"" o'EE.
o'EZa
(-0 t
In practice, we approximate it wth t >> 0, resulting in a doubly stochastic matrix. To get the binary permutation matrix, we apply the Hungarian algorithm 48 During the backward pass, we use
straight-through estimator 7 for p.
Some of the previous works 16 have leveraged the Sinkhorn operator to model variational distributions over permutation matrices. However, they start with a full rank M, which has been reported to require over 1000 Sinkho iterations converge [16.. However, our formulation, based on explicit node potential po", generates a rank-1 matrix, requiring much fewer Sinkhorn steps (around 300) in practice, saving two-thirds of the computational cost.
Bayesian Causal Discovery via Sampling
In this section, we delve into two specific methodologies that are derived from the proposed framework. The first one, which will be our main focus, combines SG-MCMC and VI a Gibbs sampling manner. The second one, which is based entirely SG-MCMC with continuous relaxation, is also derived, but we include its details in Appendix A due to its inferior empirical performance.
4.1 Model Formulation
We build upon the model formulation of [22]. which combines the additive noise model with neural networks to describe the functional relationship. Specificallyy, X; == fe(Xxs: + 6G, where fs adheres to the adjacency relation specified by G. i.. (0) = 0 if no edge exists between nodes and j. We define fa as
[1 -G E (gdd(os))
(10)
where and 4 are neural networks with parameters θ, and G serves a mask disabling non- parent values. To reduce the number of neural networks, we adopt a weight-sharing mechanism: (.(-) ((u,,) and (4--) ((u,,)) with trainable node embeddings u;
Likelihood of SCM The likelihood can be evaluated through the noise e x - f(; θ). showed that if G is DAG, then the mapping from e to æ is invertible with Jacobian determinant of 1. Thus, the observational data likelihood is:
p G) = (x O)) = II (2: - fo fi(Fra)))
(11)
1=1
Prior design joint prior:
We implicitly define the prior p(G) via p(p, W). We propose the following for the
PW +) We IW 0. aWW (xpp
where controls the initialization scale of and X controls the sparseness of G.
5

---

4.2 Bayesian Inference of W, p,
The main challenge lies the binary nature of W E {0. 1],1]xxd which requires discrete sampler. Although recent progress has been made [30 62 73 (6),, these methods either involve expensive Sarropolis-Hastingg (MH) steps or require strong assumptions on the target posterior when handling batched gradients. To address this, we propose a combination of SG-MCMC for p. O and VI for W. should be noted that our framework can incorporate any suitable discrete sampler if needed.
We employ a Gibbs sampling procedure [0], which iteratively applies (I) sampling p, O ~ p(p, θ D, W) with SG-MCMC; (2) updating the variational posterior qøl W p. D)
w
We define the posterior p(p, + D. W) & exp(-U((p, W,O),, where (p, W, ()
pp, D, W @). SG-MCMC in continuous time defines a specific form of Itô diffusion that maintains the target distribu- tion invariant
45 without the expensive computation of the MH step. We adopt the Fuler-Maruyama discretization for simplicity. Other advanced discretization can be easily incorporated [2 57
Figure 1: Graphical
Preconditioning techniques have been shown to accelerate SG-MCMC
model of the inference problem.
convergence [13 28 69 70] We modify the sampler based on (28),
which is inspired by Adam 33 Detailed update equations can found
in(Appendix C
The following proposition specifies the gradients required by SG-MCMC: Pp. O)
Proposition 4.1. Assune the model is defined as above, then we have the following:
VpU = -
lo D (WW
(12)
and
VoU = - Ve logp(()) - Vo
W))
(13)
Refer to Apppendii B.5 for details.
Variational inference for W We use the variational posterior ds (W))) to approximate the true posterior p W (p, O D). Specifically, select an independent Bernuulli distribution with logits
defined by the output of a neural network p (p):
qó
W p) = II Ber(
(14)
i
To train ls. we derive the corresponding evidence lower bound (ELBO):
ELBO( = Eg, (WIp)
W[W]] DKL (P
(15)
where DKL is the Ballback-Leibler divergence. The derivation is in Appendix B.6 summarizes this inference procedure.
Algorithm
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 W, a continuous relaxation of W, facilitating fully sampling-based approach. For detailed formulation, please refer to Appendixx A] We report its performance in [Appendix E.3 which surprisingly is inferior SG-MCMC+VI. We hypothesize that coupling W through μ4 is important since changes in results in changes of the permutation matrix o (p), which should also influence W accordingly during posterior inference. However, through sampling 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 compu- lational cost compared to existing approaches, such DIBS 43 The computational complexity of

---

Algorithm BayesDAG SG-MCMC+VI Inference
Input: dataset D; prior p(p, W),p(() SG-MCMC sampler Sampler; sampler hyperparameters
V; network (+):);; training iteration T.
Output: samples { ( p) and variational posterior 96
Initialize (0) (0)
for ..... do
Sample ( 9a W (t-1))
Evaluate V p.el/ (Equations (12) and (13) with (tt-1) p((-1), w(t-1)
(t), p(() Sampler( Ceappeer(Vp.sii P.
if storing condition met then
{p. ) + p(),0()) (t)
end if'
Maximize ELBO [Equation (15) w.r.t. ø with p(t), ())) end for
our method is OBN,, + N,dd), where B represents the batch size and Np 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 complexity of ONN2N- + N2dd) with N >> 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 result, approach provides linear scalability w.r.t. N₂ with substantially smaller batch size B. Conversely, DIBS exhibits quadratic scaling in terms of Np and lacks support for mini-batch gradients.
5 Related Work
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 [8]] 32,, often requiring restrictions on number of parents [38 Alternative exact
dynamic programming methods are limited 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) samples DAGs from autore- gressive adjacency matrix distributions, while 43 utilizes Stein variational approach 42) for DAGs
and causal model parameters. [16] proposed variational inference framework on node orderings using the gumbel-sinkhorn gradient estimator [47]
[1 employ the GFlowNet framework (6 for
inferring the DAG posterior. Most methods, except are restricted linear models, while 45 has high computational costs and lacks DAG generation guarantees compared our method.
In contrast, quasi-Bayesian methods, such DAG bootstrap 20 formance. 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 [ 22
demonstrate competitive per-
Experiments
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 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 (II) and DIBS 43
7

---

6.1 Evaluation on Synthetic Data
MMD True Posterior (1)
E-CPDAG SHD (1)
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- Rěnyi (ER) [19] or Scale-Free (SF) [5] graphs with per node degree and drawing at random ground truth parameters for linear or nonlinear models. For 5, we use N 500 training. while for higher dimensions, we use N 5000.
Badós-Rénys
0.5
10
0.0
Scale-Frea
0.5
We assess performance on random datasets for each setting.
0.0
Metries 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 Struc- tural Hamming Distance (SHD). For higher- dimensional nonlinear models with intractable posterior, we compute the expected SHD (E- SHD), expected orientation Fl score (Edge F1) and negative log-likelihood of the held-out data
BGES DDS BCD
DIBS
BaDAG
BGES DDS BCD DIBS
BaDAG
Figure 2: Posterior inference linear synthetic datasets with 5. Metrics are computed against the true posterior. + denotes lower is better.
(NLL). Our synthetic data generation and evaluation protocol follows prior work [ 22 43]. All the experimental details, including how we use cross- cross- validation select hyperparameters in Appendix D
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 5-varialle linear SCM with
unequal noise variance (identifiable upto MEC [53). The true posterior over graphs y(G D) can computed using the BGe score [23 39 Results in Figure 2] show that our method out- performs DIBS and DDS in both ER and SF settings. Compared to BCD, we perform bet- ter terms of MMD in ER but worse in SF. We find that BGES performs very well low- dimensional linear settings, but suffers significantly in more realistic nonlinear settings (see below).
6.1.2 Evaluation in Higher Dimensions
Table 1: E-SHD (with 95% CI) for ER graphs in higher dimensional nonlinear causal models. DIBS becomes computationally prohibitive for 50.
We evaluate our method on high dimensional scenarios with nonlinear relations. Our ap- proach the first to attempt full posterior in- ference in nonlinear models using permutation- based methods. Results for d 30 variables in Figure 3 demonstrate that BayesDAG signifi- ap-
70
d = 100
BGES BCD
355.77 1: 18.02
563.02 1: 27.21
217.05 + 9.58
362.66 + 29.18
DIBS
N/A
N/A 295.92 + 24.67
cantly outperforms other permulation-baa proaches and DIBS in most of the metrics. For d = 50, BayesDAG performs comparably to DIBS in ER but little worse in SF. However, our method achieves better NLL on held-out data -0-3000 30, 50, ER and SF settings. Only DDS gives better NLL for 30 ER setting. but this doesn't translate well other metrics and settings. We additionally evaluate on E [70, 100} variables (Table 1b. We find that our method consistently out- performs the baselines with 70 and in terms E-SHD with = 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.
BaDAG
143.70 + 11.61
compared to most baselines including DIBS for d

---

(1)
Edge (f)
(1)
Edge (1)
NLL
0.6
300
0.75
100
Erdds- Rddy
150
200
0.25
0.2
200
Scale-Free
100
150
30
100
50
BCD
pams
BaDAo
Bep
ras
BaDAC
Brags
DEs
BCD
riBs
BaDAd
HoRS
BCD
BaDAo
BoRs
Daw
bp BaDAG
Brag
DEAS
BCD
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 perform significantly worse in terms of EESH IE-SHD and thus has been omitted for clarity. J denotes lower better and f denotes higher is better.
6.2 Ablation Studies
We conduct ablation studies on our method using the nonlinear ER d = 30 dataset.
Initialized p scale Figure investigates the influence of the initialized scale p. We found that the performance is the best with a 0.01 or 10-5, and deteriorates with increasing scales. This because with larger initialization scale, the absolute value 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 EESHD and Edge Fl 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 θ the SG-MCMC algorithm (i.e. 8 of the SG-MCMC formulation in Appendix C Interestingly, the noise level of p does not impact the performance much as the level of O Injecting noise helps improve the performance, but smaller noise level should be chosen for θ to avoid divergence from optima.
6.3 Application 1: Evaluation Semi-Synthetic Data
We evaluate our method on the SynTReN simulator 67 This simulator creates synthetic tran- scriptional regulatory networks and produces simulated gene expression data that approximates real experimental data. We use five different simulated datasets provided by 40 with = 500 samples each. Table 2 presents the results of all the methods. We find that method recovers the true network much better in terms of E-SHD as well as Edge Fl compared to baselines.
6.4 Application 2: Evaluation on Real Data
We also evaluate on real dataset which measures the expression level of diffferent proteins and phospholipids in human cells (called the Sachs Protein Cells Dataset) [59]. The data corresponds to network of protein- -protein interactions There are 853 observational samples in total, from which we bootstrap 800 samples of 5 different datasets. to 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 the baselines even with model misspecification, proving the suitability of the proposed framework for roal-world settings.
11 different proteins with 17 edges in total among them.

---

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 different restarts. + denotes lower better and denotes higher is better.
Syntren (d 20)
Sachs Protein Cells (d 11]
E-SHD (4)
Edge FI (T)
E-SHD (J)
Edge Fl (f)
BGES
66.18 9.47
0.21 0.05
16.61 0.44
0.22 0.02
DDS
134.37 = 4.58
0.13
34.90 + 0.73
0.21 = 0.02
BCD
38.38 7.12
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
T-4HD 11:
Fige (1)
SUL
100
SHD
Edge
NEL
60
0.5
30
HOSTORO
Gai Scali
Nan www.SCNNCCC Chaca
Nazi SG-MCMMC Claso
Nax S0-SEMCC Cheis
(a) Posterior inference with different initialized p scale.
(b) Posterior inference with different number paral- lel SG-MCMC chains.
E SHD
Edec Et
NLL (1)
EHDDD
Edec FI
NEL
100
90
04
60
Samper Notia Scale
Sumper Noiio Scale
Sumppr Nekia Sccil
Sineprr Neisa Scale
Sumppr hoiia Seak d)
Sanper Kcala 1
(c) Posterior inference with different level of injected noise scale for p.
(d) Posterior inference with different level injected noise scale for θ
Figure Ablation study of posterior inference quality of BayesDAG on 30 ER synthetic dataset.
7 Discussion
In this work, we propose BayesDAG, novel, scalable Bayesian causal discovery framework that employs SG-MCMC (and VI) infer causal models. We establish the validity of performing Bayesian inference in the augmented (W, p) space and demonstrate its connection to permutation- basedd DAG learning. Furthermore, we provide two instantiations of the proposed framework that offers direct DAG sampling and model-agnosticimm 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 Ho to capture the complex distributions of W compared to 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 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' 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.
References
[1] Ryan Prescott Adams and Richard $ Zemel. Ranking via sinkhorn propagation. arXiv preprint 1rXiv:110611225 2011.
[2] Raj Agrawal, Chandler Squires, Karren Yang, Karthikeyan Shanmugam, and Caroline Uhler. Abed-strategy: Budgeted experimental design for targeted causal structure discovery. In The
10

---

22nd International Conference on Artificial Intelligence and Statistics, pages 3400-3409. PMLR, 2019.
[3] Yashas Annadani, Jonas Rothfuss, Alexandre Lacoste, Nino Scherrer, Anirudh Goyal, Yoshua Bengio, and Stefan Bauer. Variational causal networks: Approximate bayesian inference over causal structures. arXiv preprint w0w.:::206.077700 07635, 2021.
[4] Yashas Annadani, Panagiotis Tigas, Desi R Ivanova, Andrew Jesson, Yarin Gal, Adam Foster, and Stefan Bauer. Differentiable multi-target causal bayesian experimental design. arXiv preprint arXiv: 10607, 2023.
[5] Álszró-LászIó Barabási and Réka Albert. Emergence of scaling in random networks. science, 286(5439):509-512 1999.
[6] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. arXiv preprint arXiv:211 1.02261,,0026,,. 2021.
[7] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 2013.
[8] Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In International Conference on Machine Learning, pages 950-959. PMLR, 2020.
[9] Stephan Bongers, Patrick Forré, Jonas Peters, and Joris M Mooij. Foundations of structural causal models with cycles and latent variables. The Annals of Statistics, 49(5))2885-2915, 2021.
[10] George Casella and Edward I George. Explaining the gibbs sampler. The American Statistician, 46(3):167-174, 1992.
[11] Bertrand Charpentier, Simon Kibler, and Stephan Günnemann. Differentiable dag sampling. arXiv preprint arXiv:2203.cm 08509, 2022.
[12] Changyou Chen, Nan Ding, and Lawrence Carin. On the convergence of stochastic gradient meme algorithms with high-order integrators. Advances in neural information processing systems, 28, 2015.
[13] Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pages 1683-1691. PMLR, 2014.
[14] David Maxwell Chickering. Optimal structure identification with greedy search. Journal of machine learning research, 3(Nov):507-554, 2002.
[15] Gregory F Cooper and Edward Herskovits. A bayesian method for the induction of probabilistic networks from data. Machine learning, 9:309-344 347, 1992.
[16] Chris Cundy, Aditya Grover, and Stefano Ermon. Bed nets: Scalable variational approaches for bayesian causal discovery. Advances in Neural Information Processing Systems, 34:7095-7110, 2021.
[17] Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon I Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. In Uncertainty in Artificial Intelligence, pages 518-528. PMLR, 2022.
[18] Daniel Eaton and Kevin Murphy, Bayesian structure learning using dynamic programming and meme. arXiv preprint arXiv: 12X6v:1206.22444 5247, 2012.
[19] Paul Erdös, Alfréd Rényi, al. On the evolution of random graphs. Publ. Math. Inst. Hung. Acad. Sci, 5(1):17-60, 1960.
[20] Nir Friedman, Moises Goldszmidr, and Abraham Wyner. Data analysis with bayesian networks: bootstrap approach. arXiv preprint arXiv: 6695,, 2013.
[21] Nir Friedman and Daphne Koller. Being bayesian about network structure. bayesian approach to structure discovery in bayesian networks. Machine learning, 50(1):95-125, 2003,

---

[22] Tomas Geffner, Javier Antoran, Adam Foster, Wenbo Gong, Chao Ma, Emre Kiciman, Amit Sharma, Angus Lamb, Martin Kukla, Nick Pawlowski, et al. Deep end-to-end causal inference. arXiv preprini arXiv:2202.02/95, 2022.
[23] Dan Geiger and David Heckerman. Parameter priors for directed acyclic graphical models and the characterization of several probability distributions. The Annals of Statistics, 30(5):1412- 1440, 2002.
[24] Wenbo Gong, Advances in approximate inference: combining VI and MCMC and improving on Stein discrepancy. PhD thesis, University of Cambridge, 2022.
[25] Wenbo Gong, Joel Jennings, Cheng Zhang, and Nick Pawlowski. Rhino: Deep causal temporal relationship learning with bistory-dependent noise. arXiv preprint arXiv:2210. 147066 2022.
[26] Wenbo Gong, Yingzhen Li, and José Miguel Hernández-Looato. Meta-learning for stochastic gradient memc. arXiv preprint arXiv: 04522, 2018
[27] Wenbo Gong, Yingzhen Li, and José Miguel Hemnáneez-Lobato. Sliced kernelized stein discrepancy. arXiv preprint arXiv:2006. 16531, 2020.
[28] Wenbo Gong, Sebastian Tschiatschek, Sebastian Nowozin, Richard E Turner, José Miguel Hermández-Lobate and Cheng Zhang. Icebreaker: wiss efficient information acquisi- tion with bayesian deep latent gaussian model. Advances in neural information processing systems, 32, 2019.
[29] Wenbo Gong, Kaibo Zhang, Yingzhen Li, and José Miguel Hemández-Lobato. Active slices for sliced stein discrepancy. In International Conference on Machine Learning. pages 3766-3776. PMLR, 2021.
[30] Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris Maddison. Oops took a gradient:
Scalable sampling for discrete distributions. International Conference on Machine Learning, pages 3831-3841. PMLR, 2021.
[31] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773 2012.
[32] Marco Grzegorczyk and Dirk Husmeier. Improving the structure meme sampler for bayesian networks by introducing new edge reversal move. Machine Learning. 71(2-3):265, 2008
[33] David Heckerman, Christopher Meek, and Gregory Cooper. A bayesian approach to causal discovery. Innovations in Machine Learning: Theory and Applications, pages 1 1-28, 2006.
[34] Patrik Hoyer, Dominik Janzing, Joris M Mooij. Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models, Advances in neural information processing systems, 21, 2008.
[35] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[36] Mikko Koivisto. Advances in exact bayesian structure discovery in bayesian networks. arXiv preprint arXiv: 02..12006.888 6828, 2012.
[37] Harold W Kuhn. The hungarian method for the assignment problem. Naval research logisties quarterly, 2(1-2):83-97, 1955.
[38] Jack Kuipers and Giusi Moffa. Partition meme for inference on acyclic digraphs. Journal of the American Statistical Association, 1 11((517):82-229) 2017.
[39] Jack Kuipers, Giusi Moffa, and David Heckerman. Addendum on the scoring of gaussian directed acyclic graphical models. 2014.
[40] Sébastien Lachapelle, Philippe Brouillard, Tristan Deleu, and Simon Lacoste-Julien. Gradient- based neural dag learning. arXiv preprint arXiv:1906.02222. 02226 2019.
12

---

[41] Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin. Preconditioned stochastic gradient langevin dynamics for deep neural networks. In Thirtieth AAAI Conference on Artificial Intelligence, 2016.
[42] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
[43] Lars Lorch, Jonas Rothfuss, Bernhard Schölkopf, and Andreas Krause. Dibs: Differentiable bayesian structure learning, Advances in Neural Information Processing Systems, 34:24111- 24123. 2021.
[44] Chao Ma, Yingzhen Li, and José Miguel Hernández-Looato. Variational implicit processes. In International Conference on Machine Learning, pages 4222-4233. PMLR, 2019.
[45] Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient meme. Advances in neural information processing systems, 28, 2015.
[46] David Madigan, Jeremy York, and Denis Allard. Bayesian graphical models for discrete data. International Statistical Review/Revue Internationale de Statistique, pages 215-232, 1995.
[47] Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permuta- tions with gumbel-sinkhom networks. arXiv preprint arXiv: (rXiv:/802.08663 2018.
[48] James Munkres. Algorithms for the assignment and transportation problems. Journal of the society for industrial and applied mathematies, 5(1):32-38, 1957.
[49] Kevin P Murphy. Active learning of causal bayes structure. Technical report, technical report, UC Berkeley, 2001.
[50] Vlad Niculae, Andre Martins, Mathieu Blondel, and Claire Cardie. Sparsemap: Differentiable sparse structured inference. In International Conference on Machine Learning, pages 3799- 3808. PMLR, 2018
[51] Mizu Nishikawa-Toone Tristan Deleu, Jithendaraa Subramanian, Yoshua Bengio, and Laurent Charlin. Bayesian learning of causal structure and mechanisms with gflownets and variational bayes. arXiv preprini arXiv:2211.00 1.02763, 2022.
[52] Judea Pearl. Causality. Cambridge university press, 2009.
[53] Jonas Peters and Peter Bühlmann. Identifiability of gaussian structural equation models with equal error variances. Biometrika, 101(1):219-228, 2014.
[54] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: founda- tions and learning algorithms. The MIT Press, 2017.
[55] Jonas Peters, Joris M Mooij. Dominik Janzing, and Bernhard Schölkopf. Causal discovery with continuous additive noise models. 2014.
[56] Dana Pe'er, Aviv Regev, Gal Elidan, and Nir Friedman. Inferring subnetworks from perturbed expression profiles. Bioinformatics, 17(suppL 2001.
[57] Eekhard Platen and Nicola Bruti-L.berrati. Numerical solution of stochastic differential equations with jumps in finance, volume 64. Springer Science & Business Media, 2010.
[58] Robert W Robinson. Counting labeled acyclic digraphs. New directions in the theory of graphs, pages 239-273, 1973.
[59] Karen Sachs, Omar Perez. Dana Pe'er, Douglas A Lauffenburger, and Garry P Nolan. Causal protein-signaling networks derived from multiparameter single-cell data. Science, 308(5721)553-529, 2005.
(60] Peter Spirtes, Clark Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search. MIT press, 2000.
13

---

[61] Jost Tobias Springenberg. Aaron Klein, Stefan Falkner, and Frank Hutter. Bayesian optimization with robust bayesian neural networks. Advances 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 2rXiv:2206.4433 2022.
[63] Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational bayesian neural networks. arXiv preprint
2019.
[64] Panagiotis Tigas, Yashas Annadani, Andrew Jesson, Bernhard Schölkopf, 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 1001.v8801.0622 062.00000.....0 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 Tch. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference 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 vexiiv:/73..444 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 12166-12166. PMLR, 2021.
[73] Giacomo Zanella. Informed proposals for local meme in discrete spaces. Journal of the American Statistical Association, 115(530):852-866 865, 2020.
[74] Valentina Zantedeschi, Luca Franceschi, Jean Kaddour, Matt Kusner, and Vlad Niculae. Dag learning the permutahedron. arXiv preprint arXiv:2301.11889 2023.
[75] Cheng Zhang. Judith Büttpagg, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE transactions on pattern analysis and machine intelligence, 41(8):2008-204 2018.
[76] Ruqi Zhang, Xingchao Liu, and Qiang Liu. 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.
14

---

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, θ using SG-MCMC, thereby avoiding the need for variational in- ference for W.
We adopt continuous relaxation of W, similar to [43]. by introducing a latent variable W. The graphical model illustrated in Figure We can define
W
W
e
P( W = II p(W Wyy
(16)
with pW = 1 = o (..)) where a(-) is the sig- moid function. In other words, Wij defines the existence logits of Wij.
Figure 5: Graphical model with latent
variable W.
With the introduction of W, the original posterior expectations of p(p, W, W,O e.g. during evaluation, can be translated using the following proposition.
Proposition A.1 (Equivalence of posterior expectation). Under the generative model Figure S have
E (fw]))/((G,))))) (1
Bpp.W [f(G= r(p W), θ) Mp.. W. W.)
(17)
E b(w.)
where f is the target quantity.
This proof is in Appendix B.3
With this proposition, instead of sampling W, use SG-MCMC to draw W samples. Similar to Section 4.2 use SG-MCMC for p. W, O, we need their gradient information. The following proposition specifies the required gradients.
Proposition A.2. With the generative model defined as Figure 5 have
p.O.w Upp Up, O) Vp log p(p) Ve log)(O)
V logp( V p.e. w DDW
(18)
The proof is in Appendix B.4
With these gradients, we can directly plug in existing SG-MCMC samplers draw samples for p, W. and @ in joint inference (Algorithm 2).
B
Theory
B.1
Proof Theorem 3.1
For completeness, we recite the thcorem here. Theorem 3.7 (Equivalence of inference in (W, and binary DAG space). Assume graph is binary adjacency matrix representing DAG and node potential p does not contain the same values, ie. Ps * Vi.j. Then, with the induced joint observational distribution DD G). dataset D and corresponding prior p(G), we have
f P = pr WID) 1 = T(W, (WW.p))dWddppp dp
(19)
pG) = (p. W) 1 (WWP))dW dp, where p- W p) the prior, 1(.) the indicator
function and p. (p. W D) is the posterior distribution over p. W.
15

---

Algorithm 2 Joint inference
Input: dataset D, prior p(p. W, O), SG-MCMC sampler update Sampler("): sampler hyperpa- rameter Y; training steps T.
Output: posterior samples (p, W,+)
Initialize Po. Wo. 9o
for t 1,..., , -''do do
Evaluate gradient V P:-1, W U based on Equation (18)
Update samples pt, W,, O = Sampler V if storing condition met then
0, U; V)
{p. W, ,) Pc, W,, O end if end for
To prove this theorem, we first prove the following lemma stating the equivalence of r Equation (5) binary DAG space.
Lemma B.1 (Equivalence of T to DAG space), Consider d random variables, a node potential vector p E Idd and a binary matrix {0, 1]dxd, Then the following holds:
(a) For any W 1] p R G TW .p) is DAG.
(b) For any DAG G E D, where D is the space of all DAGs, there exists corresponding (W .)) such that T(W, (W.p) =
Proof. The main proof directly follows the theorem 2.1 in [22]. For (a), show the output from r(W, p) must be a DAG. By leveraging the Lemma 3.4 (72), we can easily obtain that Step(grad p) emits binary adjacency matrix representing DAG. The only difference is that replace the ReLU(.)) with Step(-) but the conclusion can be directly generalized.
For (b), we show that for any DAG G, there exists a (W, p) pair s.t. 7(W, = =G. To see this, we can observe that p implicitly defines a topological order the mapping т. For any p; >> Py. we have j + i after the mapping Step gradd Step(grad p). Thus, by leveraging Theorem 3.7 in (72), obtain that
there exists potential vector p E IRd for any DAG G such that
(grad (i. > 0 when Gij =1
Thus, we can choose W in the following way:
Wi = 0
if Gt 0 if Giz 1
W
W1 =1
Next, let's prove the Theorem 3.1
From Lemma B.1 we see that the mapping is complete. Namely, the (W,, P) space can represent the entire DAG space. Next, we show that performing Bayesian inference in (W.p) space can also correspond to the inference in DAG space.
Proof of Theorem 3.1
Assume have the prior p- PwW p)) Then through mapping implicitly define prior over the DAG G' in the following:
G = pr (W. .)) =TW d
(20)
This basically states that the corresponding prior over G is an accumulation of the corresponding probability associated with (W,y) p) pairs.
Similarly, we can define a corresponding posterior Pr
P- (G = - p, D
1 TW (W.P))dWW dW dp
(21)
16

---

Now, let's show that this posterior Pr (G D) if prior matches, i... p(G) = (G). From Bayes's rule, we can easily write down
pr Pr (w.p) p D = +( p))pp(p W)
(22)
DGEDD(DDG G')
Then, by substituting Equation (22)
into Equation (21) we have
P+ p pDD 1(G r(W, p) dW
pp p) = dp
(23)
DDD G")
GD JP p) 1(G = dp
(24)
(D)))(D,G)
(G)
(G)
(25)
where Equation (2 is from the fact that EGED p(D, G') is independent of (W, p) due to marginal-
ization. Equation (24) is obtained because p( p(D(G) also independent of (W. .p) due to (1)
(G 7(W, p)) and (2) pf G)) is constant when fixing G. Equation (25) obtained by
applying Bayes's rule and pr (G) PGG).
B.2 Proof Theorem 3.2
Theorem 3. (Equivalence of NoCurl formulation). Assuming the conditions in Theorem 3.1 satisfied. Then, for a given (W, p). we have
c Step(grad p) = W o (o'
where G is DAG and o² (p) is defined in Equation (7)
To prove this theorem, we need to first prove the following lemma. Lemma B.2. For any permutation matrix M E Eg. we have
grad Mp) M grad(p) M
where grad is the operator defined in Equation (3)
Proof. By definition of grad(.), (-), we have
grad = (Mp): - (Mp);
-1(1) Mp 1(j)" Mp
M.. M,.D
where 1(i) is one-hot vector with hhh entry 1, and M.. is the 7th row of matrix M. The above is equivalent to computing the grad with new labels obtained by permuting p with M. Therefore, we can see that grad(Mp) can be computed by permuting the original grad(p) by matrix M.
grad(Mp) = M grad(p)
Proof of Theorem 3.2 Since W plays the same role in both formulations, we focus on the equiva- lence Step(grad(())
Define a sorted p = op, where a E Egs such that for < j. we have P. > Pj- Namely, or is a permutation matrix. Thus, we have
grad(p) grad(o T p).
17

---

By Lemma B.2 we have
grad(o" p) grad
Since p is an ordered vector. Therefore, grad(p) skew-symmetric matrix with positive lower half part.
Therefore, have
Step(grad(p)) = Step(o grad(p)o") - a Step(grad(i)o' o
This is true because o is just a permutation matrix that does not alter the sign of grad(p).
Since is permutation matrix that p value in ascending order, from Lemma 1 in [8) we have
arg max (a'o)
o Ed
B.3
Proof of Proposition A.I
Proof.
Bptn.Wo1D r(p w(,
= f W. Wlw...DD.iiip....D/.....eeeeea k) MPde
=E =Bp, @ID)
JWDD W)dW
B): [f(G, O),(D, Olp. W)]
=E plp.p..o)
16 W w) Dp(D, W))
B.4 Proof of Proposition A.2
Proof.
VPU(P, (p, W, - vp logglp,
o, W, D)
V logg((p) - Vp low o, D(p)
V S P(DIW, P. O)p(Opp W W))(WWW))(WWW))
Vp logp(p)
D p. θ ()p(p) p. W WW WW)
V. logp(p) P(W P W) (pl D O)]
V
(W(Ww) P W p. O)]
V, Polgg((p) - Vp log R wW (p( D p. O]]
Other gradient V U and VeU can be derived using the similar approach, which concludes the proof.
B.5
Proof of Proposition 4.1
Proof of Proposition
By definition, we have easily have
V P
V, log.(pp W. W.O,D) O)
V Volog logg((p W) V, log D, 0 rW p))
P(g W) Vp log D p)) + V log 1ogr((177(2 W))
18

---

Similarly, we have
Vol -e D)
= Ve - Ve logp(O)p, w) Ve pp,, W)
Ve logn(
Ve (og
B.6 Derivation of ELBO
logr(p. θ D) = log / p(p, @ D. wyw
9o (W(p))) (p) Lo (W p(p o D. w,Waw
log
Uo W p) pp θ D(W)aW +f (WII)) log P(W)
dW (26)
(o
E ( p) [lp θ D[W) DrL. las (W
where the Equation (26) is obtained by Jensen's 's inequality.
C SG-MCMC Update
Assume we want to draw samples p ~ p(p) W) x exp(-U(p, W 9)) with U(p, W, θ) log p(p, W, θ), we can compute U by
U( = Σ
θ) - log W, W,
(27)
In practice we typically use mini-batches instead of the entire dataset D. Therefore, an approxi- mation is
U(p.W, =
Σ RES
logp(æ..(C G = 7 (W,, W.W.. p). θ) log W. +))
(28)
S
where [S] and DI are the minibatch and dataset sizes, respectively.
28 uses the preconditioning technique on stochastic gradient Hamiltonian Monte Carlo(SG-HMC), similar to the preconditioning technique in [41]. In particular, they use moving-average approxima- tion of diagonal Fisher information to adjust the momentum. The transition dynamics at step with
EM discretization is
1
B -1
V, = - N
9: VI+ VV.
8g.
r: B1T:-1 1g. V (p. W, θ) +1
Op: 8/ 21 1. B)n
P: P:-1 + 1gir:
(29)
where 12 is the learning rate; (,1, 32) controls the preconditioning decay rate, 1j is the Gaussian noise with mean and unit variance, and the hyperparameter controlling the level of injected noise to SG-MCMC. Throughout the paper, use (31, 32) (0.9, 0.99) for all experiments.
19

---

D Experimental Settings
D.1 Baselines
For all the experimental settings, we compare with the following baselines:
Bootstrap GES (BGES) [20 [4 is bootstrap based quasi-Bayesinn Bayesian approach for linear Gaussian models which first resamples with replacement data points at random and then estimates linear SCM using the GES algorithm [4] for each bootstrap set. GES a score based approach to learn point estimate of a linear Gaussian SCM. For all the experimental settings, we use 50 bootstrap sets.
Differentiable DAG Sampling (DDS) is VI based approach to learn distribution over DAGs and a point estimate over the nonlinear functional parameters. DDS performs inference on the node permutation matrices, thus directly generating DAGs. Gumbel-sinkhorn [47] is used for obtaining valid gradients and Hungarian algorithm is used for the straight- gradient estimator 77]. In the author provided implementation, for evaluation, single permutation matrix is sampled and the logits of the edge beliefs are directly thresholded. In this work, in order to make the comaprison fair to Bayesian learning methods, we directly sample the binary adjacency matrix based on the edge logits.
BCD Nets 16 is a VI based fully Bayesian structure learning approach for linear causal models. BCD performs inference on both the node permutations through the Gumbel- sinkhorn 44 operator as well as the model parameters through VI distribution. Both DDS and BCD nets operate directly on full rank initializations to the Gumbel-sinkhorn operator, unlike our rank-1 initialization, which saves computations in practice.
DIBS [43) uses SVGD [42 with the DAG regularizer [77 and bilinear embeddings to perform inference over both linear and nonlinear causal models. As our data generation process involves SCM with unequal noise variance, we extend DIBS framework with an inference noise variance using SVGD, similar the original paper.
While DIBS and DDS can handle nonlinear parameterization, approaches like BGES and BCD, which are primarily designed for linear models still give competitive results when applied on nonlinear data. Given that there are limited number of baselines in the nonlinear case, and DIBS being the only fully Bayesian nonlinear baseline, compare with BGES and BCD for all settings despite their model misspecification.
D.2 Evaluation Metrics
For higher dimensional settings with nonlinear models, the true posterior is intractable. While in general it is hard to evaluate the posterior inference quality in high dimensions, prior work has suggested to evaluate on proxy metrics which we adopt in this work as well 43 In particular, we evaluate the following metrics:
E-SHD: Structural Hamming Distance (SHD) measures the hamming distance between graphs. In particular, it is a measure of number of edges that are to be added, removed or reversed to get the ground truth from the estimated graph. Since we have a posterior distribution (G') over graphs, we measure the expected SHD:
EESHD := (SH GOT)
N. Σ GHD, GCT)] with G(i) ~ (G))
where GGT is the ground-truth causal graph.
Edge F1: is Fl score of each edge being present absent in comparison to the true edge set, averaged over all edges.
NLL:: We also measure the negative log- ligeliieldoood of the held-out data, which is also
typically used as a proxy metric on evaluating the posterior inference quality (26
The first two metrics measure the goodness of the graph posterior while the NLL measures the goodness of the joint posterior over the entire causal model.
20

---

For d = 5 with linear models (unequal noise variance, identifiable upto MEC 53 34 we evaluate the following metries:
MMD True Posterior: Since the true posterior is tractable, we compare the approximation with the ground truth using Maximum Mean Discrepancy (MMD) metric [31]. If P == P[G D) is the marginalized true posterior over graphs and Q is the approximated posterior over graphs, then the MMD between these two distributions is defined as:
MMD2(P Q) = EG-P (K (G) + EG-o (G G)] 2E //
where k(G, G') = H[G.G") is the Hamming kernel, and И is the Hamming distance between G and G'. This requires just the samples from the true posterior and the model. For calculating the true posterior which involves marginalization of the model parameters, appropriate prior over these parameters are required. This is ensured by using BGe score [23 39] which places Gaussian Wishart prior on the parameters. This leads closed form marginal likelihood which is distribution equivalent, i.e. all graphs within the MEC will have the same likelihood. In addition, due to the low dimensonality (d 5), we enumerate all possible DAGs and compute the normalizing constant p( D). We refer to 223 for details. This metric has been used prior work
d
EECPPAG SHD: An MEC can be represented by Completed Partially Directed Acyclic Graph (CPDAG) (54) which contains both directed edges and ares (undirected edges). When causal relations between certain set of variables can be established, directed edge is present. If there is an association between a certain set of variables for which causal direction is not identifiable, an undirected edge is present. For any graph, it has corresponding CPDAG associated to the MEC which it belongs to. Since the ground truth graph is identifiable only upto MEC, we compare the (structural) Hamming distance between the graph posterior and the CPDAG of the ground truth. This is done by computing the E-CPDAG SHD:
N₂ Σ
IE-CPDG SHD = (SDDggggySS(((Geeee G
[SHD GGE G
N
with (+)) ~ (())) and GGT GEPDAG is the ground-truth CPDAG.
D.3 Synthetic Data
As knowledge of ground truth graph is not possible in many real world settings, it is standard across causal discovery to benchmark in synthetic data settings. Following prior work, we generate synthetic data by first sampling DAG at random from either Erdos-Rènp (ER) [19] or Scale-Free (SF) [5] family. For d = 5, we ensure that the graphs have edges in expectation and 2d edges for d > 5. The ground truth parameters for linear functions nonlinear models, the nonlinear functions are defined by randomly initialized Multi-Layer Perceptrons (MLP) with single hidden layer of 5 nodes and ReLU nonlinearity. The variance of the exogenous Gaussian noise variable is drawn from an Inverse Gamma prior with concentration a = 1.5 and rate 3 1. For d 5 linear case, sample at random N 500 samples from the SCM for training and
drawn at random from a fixed range of [0.5. 1.5]. For
N = 100 for held-out evaluation. For higher dimensional settings, we consider = 5000 random samples for training and N 1000 samples for held-out evaluation. For all settings, we evaluate on 30 random datasets.
D.4 Hyperparameter Selection
In this section, we will give the details our how to select the hyperparameters for method and all the baseline models.
We employ cross- cross-validation validation- procedure for hyperparameter tuning in BayesDAG and DIBS to optimize MMD true posterior (for d = 5 linear setting) and EHSHD value (for nonlinear setting). For each ER and SF dataset with varying dimensions, we initially generate five tuning datasets. After determining the optimal hyperparameters, we fix them and evaluate the models on 30 test datasets. For DDS, we adopt the hyperparameters provided in the original paper [II. BCD and BGES do not necessitate hyperparameter tuning since BCD already incorporates the correct prior graph for ER and SF datasets. For semi-synthetic Syntren and real world Sachs protein cells datasets, assume the
21

---

BayeaDAG
A
Scale
Scale 6
Sparse Init.
linear ER
50
0.001
0.001
False
linear SF d =
50
0.01
0.001
False
nonlinear ER 20
300
0.01
0.01
False
nonlinear SF d 20
200
0.1
0.1
False
nonlinear ER = 30
500
0.01
False
nonlinear SF d 30
300
0.01
0.01
False
nonlinear ER = 50
500
0.01
0.01
True
nonlinear SF d 50
300
0.1
0.01
False
nonlinear ER d= 70
700
0.1
0.01
True
nonlinear SF d 70
300
0.01
0.01
False
nonlinear ER = 100
700
0.1
0.01
False
nonlinear SF d 100
700
0.1
0.01
False
SynTren
300
0.1
0.01
False
Sachs Protein Cells
1200
0.1
0.01
False
Table 3: The hyperparameter selection for BayesDAG for each setting.
DIBS
Qt
latent
hg
ho
linear ER d 5
0.02
1000
linear SF d =
0.02
500
nonlinear ER d 20
0.02
1500
10
nonlinear SF d 20
0.2
1500
10
nonlinear ER d 30
0.2
500
nonlinear SF d 30
0.2
5
1000
nonlinear ER d 50
0.2
500
10
nonlinear SF d = 50
0.2
5
1500
SynTren
0.2
500
10
Sachs Protein Cells
0.2
5
500
10
Table The hyperparameter selection for DIBS for each setting.
number edges in the ground truth graphs are known and we tune our hyperparameters to produce roughly correct number of edges. BCD and DIBS also assume access to the ground truth edge number and use the graph prior to enforce the number edges.
Network structure We use one hidden layer MLP with hidden size of max(4 d, 64) for the nonlinear functional relations, where d is the dimensionalilty of dataset. We use LeakyReLU as the activation function. We also enable the LayerNorm and residual connections in the network. In particular, for variational network He in BayesDAG, we apply the LayerNorm on before inputting the network. We use hidden layer MLP with size 48, LayerNorm and residual connections for ui.
Sparse initialization for BayesDAG For BayesDAG, we additionally allow sparse initialization by sampling sparse W from the Ug. This can be achieved by substracting constant from the existing logits (i.e. the output from (a).
Other hyperparameters For BayesDAG, we run 10 parallel SG-MCMC chains for p and O We implement an adaptive sinkhorn iteration where the iteration automatically stops when the sum of rows and columns are closed within the threshold 0.001 (upto maximum of 3000 iterations). Typically, we found this to require only around 300 iterations. We set the sinkhorn temperature t to be 0.2. For the reparametrization of W matrix with Gumbel-softmax trick, we use temperature 0.2. During evaluation, we use 100 SG-MCMC particles extracted from the particle buffer. We use 0.0003 for SG-MCMC learning rate and batch size 512. We run 700 epochs to make sure the model fully converged.
22

---

Table 5: Walltime results (in minutes, rounded to the nearest minute) of the runtime of different approaches on single 40GB A100 NVIDIA GPU. The N/A fields indicate that the corresponding method cannot be run within the memory constraints of single GPU.
d=30
d=50
d=70
d=100
BaDAG (Ours)(Bayesian, Nonlinear)
171
238
261
448
DIBS (Bayesian, Nonlinear)
187
350
N/A
N/A
BGES "Quasi-Bayesian Linear)
2
3
6
11
BCD (Bayesian, Linear)
252
328
418
600
DDS (Quasi-Bayesian, Nonlinear)
92
130
174
N/A
E-SHD (L)
Edge Fl (1)
NLL (L)
0.75
Erdös-Red Rénni
75
40
0.50
50-
30
25
0.25
0.75
Scale-Fre FFee
75
35
0.50
50
30
0.25
25
BGES BCD
DIBS
BaDAG
BGES DDS BCD
DIBS
BaDAG
BGES DDS BCD
DIBS
BaDAG
Figure 6: Posterior inference of both graph and functional parameters synthetic datasets of nonlinear causal models with d = 20 variables. BayesDAG gives best results across all metrics. denotes lower is better and denotes higher is better. For the sake of clarity, DDS has been omitted for E-SHD due to its significantly inferior performance on this metric.
For DIBS, we can only use 20 SVGD particles for evaluation due to the quadratic scaling with the number of particles. We use 0.1 for Gumbel--oftmax temperature. We 10000 epochs for convergence. The learning rate is selected as 0.01.
Table 3 shows the hyperparameter selection for BayesDAG. Table 4 shows the hyperparameter
selection for DIBS.
E
Additional Results
E.1
Walltime Comparison
presents walltime comparison of different methods. Our method converges faster while being scalable w.r.t. DIBS, the nonlinear Bayesian causal discovery baseline. Other methods like BGES and DDS, while faster, perform much worse in terms of uncertainty quantification. In addition BGES limited linear model and DDS not a fully Bayesian method.
Table 5
E.2 Performance with higher dimensional datasets
Full results for the metries for settings d 20, 70 and 100 for nonlinear settings are presented in Figure 6 Figure 7 and Figure
23

---

E-SHD (L)
Edge F1 (1)
NLL (J)
Erdös-Rényi
400
120
0.4
110
200
0.2
100
300
Scale-Free Seale-Free
0.6
110
200
0.4
100
100
90
BGES
BCD
BaDAG
BGES
BCD
BaDAG
BGES
BaDAG
Figure 7: Posterior inference of both graph and functional parameters synthetic datasets of nonlinear causal models with d = 70 variables. BayesDAG gives best results across most metrics. + denotes lower is better and denotes higher is better. As DIBS and DDS are computationally prohibitive to run for this setting, it has been omitted. BCD has been omitted for NLL as we observed that it performs significantly worse.
E-SHD (!)
Edge F1 (1)
NLL (1)
Erdös-Rényi
1000
0.4
180
500
0.2
160
Scale-Fri Free
400
0.4
160
0.2
200
140
BGES
BCD
BaDAG
BGES
BCD
BaDAG
BGES
BaDAG
Figure 8: Posterior inference of both graph and functional parameters synthetic datasets of nonlinear causal models with d = 100 variables. BayesDAG gives best results across E-SHD, comparable across NLL but slightly worse for Edge F1. J. denotes lower better and denotes higher is better. As DIBS and DDS are computationally prohibitive to run for this setting, it has been omitted. BCD has been omitted for NLL as observed that it performs significantly worse.
24

---

E-SHD (L)
Edge F1 (1)
NLL (1)
Erdös-Rényt
Erdós-Re Rényí
20
17.5
0.6
15.0
10
0.4
12.5
0.8
Scale-Fre Ccale-Free Free
20
20
0.6
10
15
0.4
SG-MCMC
SG-MCMC+VI
SG-MCMC
SG-MCMC+VI
SG-MCMC
SG-MCMC+VI
Figure 9: Performance comparison of SG-MCMC+VI v.s. fully SG-MCMC with W for d 10 variables.
E.3 Performance of SG-MCMC with Continuous Relaxation
We compare the performance of SG-MCMC+VI and SG-MCMC with W on d = 10 ER and SF graph settings. Figure 9 shows the performance comparison. We can observe that SG-MCMC+VI generally outperforms its counterpart in most of the metrics. We hypothesize that this is because VI network fe couples and W. This coupling effect crucial since the changes in results in the change of permutation matrix, where the W can immediately respond to this change through Hag On the other hand, W can only respond to this change through running SG-MCMC steps W with fixed p. In theory, this is the most flexible approach since this coupling do not requires parametric form like pg. However in practice, we cannot run many SG-MCMC steps with fixed p for convergence, which results in the inferior performance.
F
Code and License
For the baselines, we use the code from the following repositories:
BGES:
from the https:/ggitubbooom/aaawaaaaaaaa laaaa////cciieeeeeiiggggg (No license included).
We use the code from
repository
DDS:
We
use the code from the official repository
(No license included).
BCD: We use the code from the official repository http::/ghhhbbbccm//mmmoopppooppp (No license included).
DIBS: We use the code from the official repository https://gthabbccom//sssschhhhhi (MIT license).
Additionally for the Syntren 67 and Sachs Protein Cells (59 datasets, use the data provided with repository http://ghhhbb..mmmuuwwvaaaaaaaaa (MIT license).
G Broader Impact Statement
This work is concerned with understanding cause and effects from data, which has potential applica- tions in empirical sciences, economics and machine perception. Understanding causal relationships
25

---

can improve faimess in decision making, understand biases which might be present in the data and answering causal queries. As such, we envision this line of work to not have any significant negative impact.
26


Reviews:
[{"review_id": "tPDMTdnWfQ", "reviewer": "Reviewer_hANm", "summary": "The paper proposes a novel Bayesian causal discovery (BCD) method that infer the posterior distribution $p(G|\\mathcal{D})$ by projecting the DAG $G$ into an equivalent search space. Instead of sampling $G$, the method constructs the posterior distribution by sampling a binary matrix $W$ and potential vector $p$ with via MCMC sampling and variational inference. The Bayesian causal discovery method can scale up to 100 variables and achieves better accuracy on large datasets.", "questions": "My question focus on the experiment section:\n\n- It seems that the BayesDAG achieves better accuracy on cases $d>=70$. I wonder if authors can also show the efficiency (as in runtime) of different methods? I suspect that BayesDAG would take longer to converge than the VI-based approaches due to the Gibbs sampling procedure. The strength of BayesDAG is better demonstrated if it can achieve much better accuracy with a slight compromise on the efficiency. \n- Some synthetic data is generated based on the unidentifiable linear SEM with non-equal variances. How does the identifiability of the employed SEM affect the BayesDAG method (How does the non-equal variances assumption influence the Eq. (11))? \n- The paper only show the empirical results on $d\\leq 100$ variables. I am wondering if authors can show BayesDAG can scale up to data with higher dimensions. The gradient-based causal discovery methods such as GraN-DAG can also scale up to 100 variables and finish within reasonable runtime with the help of GPU.\n- It would be better if the authors can compare to traditional scalable causal discovery methods such as FGES? ", "limitations": "The idea of the paper seems interesting and the theory is sound. I concern that the Gibbs sampling procedure in the propose algorithm may compromise its efficiency and scalability.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The idea of employing the projection framework from DAG-Nocurl paper is interesting. Especially, the difficulties of the sampling based posterior distribution estimation for DAG learning methods lie in the order of parents sampling. The potential function p automatically reserves the causal order.\n\n- The proposed method is a combination of sampling-based method and variational inference method. Compared to the state-of-art BCD methods that adopt VI, the proposed method achieves better SHD, especially on high-dimensional data.", "weaknesses": "- (**Major**) The experiments are not comprehensive. There is a trade-off between efficiency and accuracy compared sampling-based approach to VI approach. Compared to the existing VI-based BCD methods, It is possible that the proposed approaches are more accurate but also suffer from low efficiency. Please refer to the question section for details.\n\n- (**Minor**) The tuning of hyperparameters such as scale of p and theta. Since the original framework of DAG-Nocurl is derived for continuous parameterization, the algorithm requires the tuning of additional hyperparameters, which increases the training difficulty. But I understand this is a minor concern.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedback for our paper. We will try to address the raised concern in the following.\n\n1. **Major concern**: We have included a wall-clock time comparison in the supplementary PDF. From the comparison, we can see our approach, compared to Dibs and BCD, converges faster while obtaining better performance in most cases. DDS, though converges faster, performs significantly worse than our approach. In addition, it is important to note that the BCD net is a **linear method**, while our approach is **fully non-linear**. However, this simplicity in model assumption may be too restrictive for certain problems, limiting the potential for accurate causal discovery. In contrast, our method not only demonstrates performance advantages for dimensions $d > 70$, but also outperforms competing methods, including BCD net, for $d = 30$ and $d = 50$, along with the faster convergence speed. This can be clearly observed in Figure 3 of our paper. These results serve as strong evidence for the benefits of adopting a non-linear model in Bayesian causal discovery.\n\n2. **Minor concern**: In response to the reviewer's concerns about hyper-parameter tuning, we have provided a comprehensive list of hyper-parameters used in our experiments within Appendix D. To address the challenge of tuning additional hyperparameters, we have conducted ablation studies presented in Appendix E.3. These studies examine the sensitivity of the initialized $\\pmb{p}$, the number of MCMC chains used, and the sampler noise scale for $\\pmb{p}$ and $\\pmb{\\Theta}$. These ablation studies offer valuable guidelines for tuning the additional hyperparameters. To further assist the readers, we will include a concise paragraph in the revised paper, explaining how to select these hyperparameters.\n\n3. **Identifiability**: It is important to emphasize that identifiability is an inherent property of the model (and not that of the inference/estimation method, which is what BayesDAG is about), which we use additive noise model to ensure this. With small amount of data, the prior term can affect the search procedure, but it will not affect the identifiability since our prior will put non-zero probability on any DAG and function parameters. On the other hand, with enough data (e.g.~infinite data), the prior term can be negligible, and will also not affect our search procedure (see Theorem 1 in [1]). In the case of experiment 6.1.1, the ground truth data generative mechanism is not identifiable, and only identifiable up to Markov equivalence class; however, this does not affect our inference procedure in any sense. One advantage of our method is that it can recover multiple graphs that fit the data equally well, which can be used to test the model's uncertainty estimation quality.\n\n4. **Higher dimensional problem**: As requested, we have included the performance in even higher dimensions in the supplementary material (with 40GB A100 GPU). We are the only method that is capable of generalising to $d>100$ among the baselines apart from BGES, which is a quasi-Bayesian approach. BCD and Dibs are not scalable enough to run under this high dimensionality, which demonstrates our method is more memory-efficient compared to most of them. Additionally, we would like to clarify that a direct comparison between our method and GraN-DAG may not be entirely appropriate for the following reasons: (i) GraN-DAG is not a Bayesian causal discovery method; it focuses on inferring a single DAG, whereas our method is designed to infer a distribution over DAGs; (ii) GraN-DAG is based on continuous-relaxation, which may result in graphs that do not strictly adhere to the DAG structure. In contrast, one advantage of our method is to infer DAGs directly. \n\n5. **FGES**: Regarding traditional causal discovery baselines, we have indeed included a comparison with the bootstrap GES method in our paper. Please refer to Figure 2, Figure 3, and Tables 1 and 2 for detailed comparisons. It is important to note that Fast GES is a variant of GES, but it is not a Bayesian causal discovery method, and only infer a single graph. \n\nWe would once again like to express our gratitude to the reviewers for their valuable feedback and hope that these clarifications have effectively addressed your concerns.\n\n[1] Geffner, Tomas, et al. \"Deep end-to-end causal inference.\" arXiv preprint arXiv:2202.02195 (2022)."}, {"review_id": "j7V81iHwlu", "reviewer": "Reviewer_3g7c", "summary": "The paper proposes a Bayesian causal discovery method based on a novel parametrization of the binary DAG space and SG-MCMC. The proposed method does not rely on DAG regularization nor restricted to linear models, overcoming the limitations of prior approaches. Experimental results demonstrate the competitive performance of the proposed method compared to existing approaches.\n", "questions": "Details of the weight-sharing mechanism (line 186) are missing. While the size of the networks is very small, how does it impact the performance of the proposed method other than reducing the total network parameters?\n\nTypo: (Line 230) focusonly\n", "limitations": "Some of the assumptions can be viewed as limitations but I am fine with those (they are crucial to yields the proposed methods)", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and easy to understand. It is easy to follow the core idea of the proposed method.\n\nThe proposed method is sound and well-motivated (i.e., there are apparent limitations of previous work but this method overcomes such issues.)\n\nSeveral techniques are employed smoothly to propose the method.\n\nExperimental results demonstrate the effectiveness and scalability of the proposed method.\n", "weaknesses": "For the empirical evaluation, comparison with MCMC approach [1] is missing. Also, AUROC is not reported, which is widely used for the evaluation of uncertainty quantification in Bayesian causal discovery literature.\n\nThe proposed method is claimed to be scalable and computational complexity is analyzed, but the actual computation cost (e.g., wall clock time) is not compared with DiBS. Computation resources they used are also not provided (e.g., CPU, GPU).\n\nSimilarly, “per node degree 2” seems very limiting which results in a very sparse graph for large d. Large d will certainly affects the performance of likelihood computation and posterior sampling of W where more edges may demonstrate dependencies.\n\n[1] Improving markov chain monte carlo model search for data mining, 2003\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedback for our paper. We will address the concerns in the following:\n\n1. **Missing MCMC baselines**: We acknowledge the importance of a comprehensive analysis. Note that all MCMC methods are only defined on linear models, usually where the parameters can be marginalized out. They are also not scalable, limiting the settings which we can compare it with. There has been extensive research showing their lack of scalability, mixing and convergence speed (see related work). In addition, it has been shown that DIBS [1] outperforms all MCMC methods (refer to Figure 2, 3, 4, and Table 1 in [1]). The superior performance of our method, as compared to DIBS, led us to believe that MCMC as a baseline might not be the best choice. However, given your comment, we will include a comparison to an MCMC baseline for the 5 variable setting for camera-ready, even though we believe adding it will not change our conclusion.\n\n2. **Missing AUROC metric**: Our primary goal was to select a diverse set of metrics that provide a comprehensive assessment of the performance, capturing different aspects of the model. While AUROC is a metric for evaluating graph quality, we believe that its utility in our case is somewhat limited, as it shares similarities with the SHD and F1 score. These latter metrics, which are incorporated into our evaluation, also consider the inferred graph's quality, arguably making the inclusion of AUROC somewhat redundant for our purposes. Apart from the quality of graph posterior, posterior over function parameters is also crucial for Bayesian causal discovery. This aspect is not captured by the AUROC metric. To account for this crucial facet, we opted for the held-out likelihood as an evaluation metric. This choice is grounded in its widespread use in the literature as a reliable indicator of Bayesian inference quality [2,3]. Meanwhile, we can include the AUROC in the camera-read version.\n\n3. **Computational efficiency**: We have included the wall-clock comparison and generalization to $d>100$ datasets in the PDF. It is important to emphasize that our approach is capable of handling significantly larger dataset dimensions, scaling up to $d>100$ cases, and achieve faster convergence compared to BCD and Dibs in terms of wall-clock time under all dimensionalities. On the other hand, DIBS is limited to $d\\leq50$ with a single GPU, and BCD is limited to $d<150$.Moreover, our paper includes an in-depth computational complexity analysis (refer to lines 232-241). This analysis further strengthens the claim that our approach exhibits a competitive edge over DIBS.\n\n4. **Per node degree of 2**: Our choice is motivated by its prevalence in the literature, particularly when it comes to synthetic ER and SF datasets [4]. Notably, DIBS [1] also employed this choice for their experiments. By adopting this, we ensure that our experiments are consistent with established benchmarks, thereby facilitating a more meaningful comparison between our method and baselines.\n\n5. **Weight-sharing mechanism and size of networks**: Sorry for the ambiguity. By formulation of Eq.10, each node requires two separate neural networks. Thus, in total, it requires $2d$ number of different networks to train, incurring high computational cost. To avoid that, instead, we have a separate trainable embedding $\\vu_i$ for each node. Therefore, we only need two neural networks with this trainable embedding to differentiable different nodes. This is equivalent to sharing the weights across the $2d$ neural networks [5]. also adopted the same sharing mechanism. We will add a clearer explanation in the revised paper. The reason we choose this network size is that it is already a common choice in the literature [1,5]. In fact, Dibs [1] used even smaller network sizes (2 hidden layers with 5 units). Our method can be easily extended to larger network sizes and it shouldn't impact the performance much, since SG-MCMC is designed to accommodate the large network and dataset sizes [6].\n\nWe would once again like to express our gratitude to the reviewers for their valuable feedback and hope that these clarifications have effectively addressed your concerns.\n\n[1] Lorch, Lars, et al. \"Dibs: Differentiable bayesian structure learning.\" Advances in Neural Information Processing Systems 34 (2021): 24111-24123.\n\n[2] Gong, Wenbo, Yingzhen Li, and José Miguel Hernández-Lobato. \"Meta-learning for stochastic gradient MCMC.\" arXiv preprint arXiv:1806.04522 (2018).\n\n[3] Lorch, Lars, et al. \"Amortized inference for causal structure learning.\" Advances in Neural Information Processing Systems 35 (2022): 13104-13118.\n\n[4] Zheng, Xun, et al. \"Dags with no tears: Continuous optimization for structure learning.\" Advances in neural information processing systems 31 (2018).\n\n[5] Geffner, Tomas, et al. \"Deep end-to-end causal inference.\" arXiv preprint arXiv:2202.02195 (2022).\n\n[6] Chen, Changyou, et al. \"Bridging the gap between stochastic gradient MCMC and stochastic optimization.\" Artificial Intelligence and Statistics. PMLR, 2016."}, {"review_id": "abdCGc4OY5", "reviewer": "Reviewer_zwvL", "summary": "The authors propose a method for the posterior inference of DAG structure *and* function parameters with potential applicability to arbitrary functional relations between nodes. The authors modify a novel characterization of DAGs, and interpret this characterization in terms of a sorting operation which can be relaxed to allow differentiability. The authors define priors on DAGs (in the alternative space) and function parameters and based on a specific model choice characterize likelihood. They use the resulting joint distribution to iteratively sample some parameters and conduct variational inference re. others. The authors examine the performance of their proposed methodology on various synthetic and real datasets.", "questions": "- Are there any potential difficulties with using the authors' method with other SCM model assumptions / likelihoods?\n- What are grounds for baseline selection in experiments? I think this would be an important addition to the final text.\n- In 6.4, were there model misspecification in other methods as well?", "limitations": "I think the authors adequately address the limitations of their work overall, however see Weaknesses section above for some important caveats.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well written. It presents the previous work, motivation for current research, and reasoning behind methodological choices very clearly.\n- The paper utilizes recent, previous research intelligently and presents concrete innovations to solve well-defined problems.\n- Posterior inference in the DAG structure and parameter space without some of the limitations of previous work is valuable and is likely to inspire future work.", "weaknesses": "- DAG model selection results have causal implications given specific model assumptions regarding generative model of the data. ANM is such a model assumption. However, it is unclear whether the identifiability results still apply in this case, given the priors defined on DAG structure and function parameters. I think the authors' work still would be valuable as only a DAG inference method; however, since the authors present their proposal as a causal discovery + inference method, this point needs further discussion.\n- I think the authors' presentation should be modified to make sure their inference method is more clearly understood. Given their initial presentation, including \"posterior sampling\" in the title, and frequent reference to Gibbs sampling throughout the text, leads the reader think that the authors will present results with a correct MCMC algorithm and produce a full posterior distribution. However, most promising results presented by authors include their iterative algorithm that samples from the posterior of some parameters and uses variational inference for others. This is fine as a methodological choice, but their presentation leads the reader to have higher expectations, which can become crucial depending on the use case of the reader.\n- Causal sufficiency assumption prevents using the current method in problems where unobserved confounding is likely. In my opinion this is acceptable given the difficulty of the problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedback for our paper. We will address your concerns in the following:\n\n1. **Identifiability**:  We would like to clarify that identifiability is a property of the SCM (the specific parametric form thereof), and not that of an estimation/inference method, which is what BayesDAG is about. With small amount of data, the prior can affect the search procedure, but it won't affect the identifiability since our priors over graphs are non-zero for any DAG. Namely, for each sampled graph and function parameters, the resulting ANM is identifiable, apart from some exceptions (e.g. linear Gaussian, etc.). With enough data (e.g.infinite data), the prior can be negligible and it will also not affect the search procedure (see Theorem 1 in [1]). Since our method can also perform causal inference task, our model is also causal inference identifiable. For each DAG and function parameters pair, we have a corresponding SEM, which the causal inference quantity can be obtained by manipulating the SEM. \n\n2. **Clarity of the inference method**: Sorry about the ambiguity. But we want to emphasise that although the fully SG-MCMC approach does not give promising empirical results compared to the SG-MCMC+VI, proposing such framework itself is a contribution and requires further investigations regarding its inferior performances. We will make the presentation of the inference method clearer in the revised paper. \n\n3. **Causal sufficiency**: We agree that causal sufficiency is a strong assumption in practice, but under the scope of this paper, we only consider the system without latent confounders. This is also a common assumption adopted by all our baselines and most previous work. Future work is needed to relax this constraint. \n\n4. **Potential difficulty with other SCM**: It depends on which SCM has been chosen. The core assumption required by our method is that the SCM should be structurally identifiable. For example, post-nonlinear model, which is identifiable, can be used to replace the ANM model in our case. \n\n5. **Baseline selection**: The selection criteria for baselines are the following: (1) It should be a Bayesian causal discovery method; (2) the baselines should cover both linear (BCD nets) and non-linear models (Dibs); (3) the baseline should have quasi-Bayesian method adapted from traditional causal discovery approach (BGES). These principles ensure a comprehensive set of baselines to demonstrate the effectiveness of our method.\n\n6. **Model misspecification in 6.4**: Yes, our method, along with other baselines, should have model misspecification, since the ground truth mechanism may not have additive noise structure.\n\nWe would once again like to express our gratitude to the reviewers for their valuable feedback and hope that these clarifications have effectively addressed your concerns. \n\n\n[1] Geffner, Tomas, et al. \"Deep end-to-end causal inference.\" arXiv preprint arXiv:2202.02195 (2022).\n\n"}, {"review_id": "rtPDnMxQlR", "reviewer": "Reviewer_R7H9", "summary": "BayesDAG proposes a hybrid SG-MCMC sampling and variational inference for drawing samples from the posterior distribution of DAGs in the context of Bayesian structure learning. \nThis work is closely related to a recently published \"Yu et al., Dags with no curl, 2021\" [NoCurl] where the space of DAGs is converted to the space of a skew-symmetric matrix, \"W\", and a potential vector,\"p\".\nThe main difference is that [NoCurl] focuses on an optimization setting (to return a most probable DAG) while the focus of the current paper is on Bayesian inference (via sampling). \nThey also mention that (for reasons that are not entirely clear to me) NoCurl approach optimization is challenging (due to uninformative gradients) and that this is due to the fact that the entries of the skew-symmetric matrix, W are continuous. \nTo address this problem, they replace the continuous matrix W with a binary matrix. \nFor this purpose, they slightly modify the theory presented in NoCurl (namely, replacing relu with step function). Then they reformulate the problem as a  matrix permutation setting and finally propose an approximate solution for the latter formulation via Sinkhorn approach to learn latent permutations by the Gumble trick.    ", "questions": "Given that the Sinkhorn approach is approximate, how do you guarantee the DAGness of your graphs ?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This is a well-written paper addressing an important problem i.e. Bayesian structure learning.\n2. The proposed approach is an interesting (albeit sophisticated) combination and modification of several algorithms and recent advances in the field.\n3. Even though I did not follow why the continuous matrix W had to be replaced by a binary matrix in the first place, but I found the way they did it and managed to approximate its gradient, very interesting.  \n ", "weaknesses": "1. As I mentioned in the summary section, a key contribution of this paper is the insight that it is better to replace the continuous matrix W (of NoCurl algorithm) with a binary matrix (see lines 101-104). But their justification for this claim does not seem convincing to me and should be explained better. To be more concrete: \n(a) Why should replacing a continuous matrix with a discrete matrix be helpful when we are relying on gradient information for optimization and sampling? \n(b) In line 100 they mention a \"reported failure\" of NoCurl. It would be great if the authors would provide a reference to where this failure is reported (Or if it is reported in the original NoCurl paper, the relevant section). \n\n2. Given the close link of the present paper with the NoCurl approach, it would be great if the authors would compare their algorithm (with binary W) with an alternative approach where just like NoCurl, a continuous W would be used (and relu instead of step function, etc).    \n\n3. I see references to repositories that the authors have used but no link to their own code. Given that implementation of the proposed algorithm from scratch is by no means trivial, I encourage the authors to provide the code. Both for facilitating other researchers to use their algorithm as well as allowing the reviewers to check the reproducibility of the reported results.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedback for our paper. We will address your concerns in the following:\n\n1. **Advantages compared to no-curl**: Sorry about the ambiguity. We will make it more clearer in the revised paper. There are several reasons on why the binary matrix is better than continuous one: (i) Note that $ReLU(grad \\pmb{p})$ gives a fully connected DAG. The main purpose of W matrix therefore is to disable the edges. Continuous W requires threshold to properly disable the edges, since it is hard for a continuous matrix to learn exactly 0 during the optimization; (ii) $ReLU(grad \\pmb{p})$ already contains continuous values, and $W$ matrix also contains the continuous values. Thus, learning of the edge weights and DAG structure are not explicitly separated, resulting in complicated non-convex optimizations (see discussion below Eq.3 in [1]). With binary matrices (e.g. replacing with $Step$ function and binary $W$), we only focus on learning the graph structure, which significantly simplifies the optimization complexity. \n\n2. **Reported failure of no-curl**: In the original No-curl work [1], they reported that direct optimizing the $W\\cdot ReLU(grad \\pmb{p})$ results in poor graph discovery performance (refer to rand init and rand p in Table 1 [1]). That is exactly the reason they propose to use NoTears as the initialisation and design complicated projection scheme to project initialisation to a DAG space. On the other hand, our approach does not require any projections and sampling with our modified objective can directly lead to a good graph discovery performance. \n\n3. **Comparison of alternative approach**: To the best of our knowledge, our work is the only one that shares some similarities with No-Curl. It is important to note that our method is focused on a Bayesian causal discovery framework, while No-Curl is not explicitly designed for this purpose. Thus, No-Curl might not serve as the most appropriate baseline for our study. We further observed that No-Curl demonstrated very poor performance when directly optimized, as shown in Table 1 of [1]. This observation led us to reasonably assume that our method would outperform it. Regarding the No-Curl with the projection steps, it was found to have similar performance to No-Tears, as illustrated in Figure 1 of [1], but it is hard to adapt the projection step for Bayesian causal discovery. Dibs [2], a Bayesian causal discovery algorithm inspired by No-Tears, serves as a better comparison for our work. In most cases, our method exhibits superior performance compared to Dibs, further solidifying the effectiveness of our proposed approach.\n\n4. **Access to code**: We have provided a link to AC containing the code and will open source it on acceptance.\n\n5. **DAGness with approximate gradient**: Although the Sinkhorn is an approximation method, and can only output bi-stochastic matrix. To obtain a valid permutation matrix, we use Hungarian matching algorithm to ensure a valid permutation can always be obtained from this bi-stochastic matrix (see line 167). With a valid permutation matrix, a DAG structure can be guaranteed. \n\n\nWe would once again like to express our gratitude to the reviewers for their valuable feedback and hope that these clarifications have effectively addressed your concerns.\n\n\n[1] Yu, Yue, et al. \"DAGs with no curl: An efficient DAG structure learning approach.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] Lorch, Lars, et al. \"Dibs: Differentiable bayesian structure learning.\" Advances in Neural Information Processing Systems 34 (2021): 24111-24123."}]

Author rebuttal:
We would like to express our sincere appreciation to all reviewers for their valuable time and efforts in providing constructive feedbacks for our paper. We are delighted to hear that the reviewers generally find our work interesting (hANm, R7H9), sound and well-motivated (3g7c, zwvL), well-written (3g7c, zwvL, R7H9), valuable and is likely to inspire future work  (zwvL). 

We have taken into consideration the common concerns raised, specifically regarding computational efficiency and scalability. In response, we have included a comprehensive comparison of wall-clock times and performance results for dimensions $d > 100$ in the supplementary PDF. In summary, we would like to emphasize that BayesDAG, as a nonlinear Bayesian approach, demonstrates better computational efficiency in comparison to existing state-of-the-art Bayesian causal discovery algorithms, such as Dibs and BCD nets. It is also worth noting that algorithms like Dibs and DDS are unable to run for dimensions $d > 70$ within a single GPU setup, and we are the only one that is capable of generalising to $d>100$ cases (with a single 40GB A100 GPU). These additional experiments effectively demonstrate the computational advantages of our approach compared to the baselines. In applications like gene regulatory network inference, approaches that can deal with several hundred variables are required, and our contribution takes a positive step in this direction. Just like NoTears, our work, based on NoCurl, and our theoretical contributions lay ground work for scalable causal discovery approaches in more complex settings, for example with hidden variables. We also shared our code as requested by R7H9.

Meta-review: This paper proposes a Bayesian approach to directly acyclic graph (DAG) structure learning based on a hybrid SG-MCMC sampling and variational inference scheme for drawing samples from the DAG posterior. This work is related to the NoCurl approach of Yu et al., where the space of DAGs is represented in the space of a skew-symmetric matrix W and a potential vector p. While NoCurl provides a point-estimate, the proposed method focuses on the Bayesian setting (and hence providing a full posterior over DAGs).

I thank the reviewers and the authors for their discussions. I believe all the reviewers are positive about this paper in that it represents a contribution to the problem of DAG estimation (Bayesian and non-linear cases). Most concerns have been addressed by the authors in the rebuttal. In particular, that of considering computational efficiency vs. accuracy.
Decision: Accept (poster)