SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yashas Annadani $^{\dagger*}$  1,3,4 Nick Pawlowski $^{2}$  Joel Jennings $^{2}$  Stefan Bauer $^{3,4}$  Cheng Zhang $^{2}$  Wenbo Gong $^{*2}$ $^{1}$  KTH Royal Institute of Technology, Stockholm  $^{2}$  Microsoft Research  
 $^{3}$  Helmholtz AI, Munich  $^{4}$  TU Munich
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.
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 a 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 – a subject of extensive research [54, 60].
In contrast to traditional methods that infer a single graph or its Markov equivalence class (MEC) [14, 60], Bayesian causal discovery [21, 33, 65] aims to infer a 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 is scarce. It is also beneficial for downstream tasks such as experimental design [2, 4, 49, 64].
*Equal contribution. † Work done during internship at Microsoft Research. Correspondence to wenbogong@microsoft.com
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 traditional discovery methods [14, 49, 65], 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-based methods, e.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) [27, 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 inaccurate inference quality [28, 61, 66] and is sometimes restricted to only linear models [16].
In this work, we introduce BayesDAG, a 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  $(\pmb{W},\pmb{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-based DAG learning which provides a general framework for gradient-based posterior inference. (Section 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) a combination of SG-MCMC and VI (2) SG-MCMC with a 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 on various datasets. We also demonstrate that our method can be easily scaled to 100 variables with nonlinear relationships. (Section 6)
Causal Graph and Structural Causal Model Consider a data generation process with  $d$  variables  $\mathbf{X} \in \mathbb{R}^d$ . 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_{i}$  is a function of its direct causes  $\mathbf{X}_{\mathrm{Pa}^i}$  and an exogenous noise variable  $\epsilon_{i}$  with distribution  $P_{\epsilon_i}$ :
$$
X _ {i} := f _ {i} \left(\boldsymbol {X} _ {\mathbf {P a} ^ {i}}, \epsilon_ {i}\right) \tag {1}
$$
These equations induce a causal graph  $G = (V, E)$ , comprising a node set  $V$  with  $|V| = d$  indexing the variables  $X$  and a directed edge set  $E$ . If a directed edge  $e_{ij} \in E$  exists between a node pair  $v_i, v_j \in V$  (i.e.,  $v_i \to v_j$ ), we say that  $X_i$  causes  $X_j$  or  $X_i$  is the parent of  $X_j$ . We use the binary adjacency matrix  $G \in \{0, 1\}^{d \times d}$  to represent the causal graph, where the entry  $G_{ij} = 1$  denotes  $v_i \to v_j$ . A standard assumption in causality is that the structural assignments are acyclic and the induced causal graph is a DAG [9, 52], 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  $\epsilon_i$  are mutually independent. Throughout this work, we consider a special form of SCM called Gaussian additive noise model (ANM):
$$
X _ {i} := f _ {i} \left(\boldsymbol {X} _ {\mathbf {P a} ^ {i}}\right) + \epsilon_ {i} \quad \text {w h e r e} \quad \epsilon_ {i} \sim \mathcal {N} \left(0, \sigma_ {i} ^ {2}\right) \tag {2}
$$
If the functions are not linear or constant in any of its arguments, the Gaussian ANM is structurally identifiable [34, 55].
Bayesian Causal Discovery Given a dataset  $D = \{\pmb{x}^{(1)},\dots ,\pmb{x}^{(N)}\}$  with i.i.d observations, underlying graph  $G$  and SCM parameters  $\Theta$ , they induce a unique joint distribution  $p(D,\Theta ,G) = p(D|G,\Theta)p(G,\Theta)$  with the prior  $p(G,\Theta)$  and likelihood  $p(D|G,\Theta)$  [21]. Under finite data and/or limited identifiability of SCM (e.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 a maximum likelihood estimate). Bayesian causal discovery therefore aims to infer the posterior  $p(\pmb{G}, \Theta | \pmb{D}) = p(\pmb{D}, \Theta, \pmb{G}) / p(\pmb{D})$ . However, this posterior is intractable due to the super-exponential growth of the possible DAGs  $G$  [58] and continuously valued model parameters  $\Theta$  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, 71] 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, [72] introduced NoCurl, a novel characterization of the weighted DAG space. They define a potential  $p_i \in \mathbb{R}$  for each node  $i$ , grouped as potential vector  $\pmb{p} \in \mathbb{R}^d$ . Further, a gradient operator on  $\pmb{p}$  mapping it to a skew-symmetric matrix is introduced:
$$
(\operatorname {g r a d} \boldsymbol {p}) (i, j) = p _ {i} - p _ {j} \tag {3}
$$
Based on the above operation, a mapping that directly maps from the augmented space  $(\mathbf{W},\mathbf{p})$  to the DAG space  $\gamma (\cdot ,\cdot):\mathbb{R}^{d\times d}\times \mathbb{R}^d\to \mathbb{R}^{d\times d}$  was proposed:
$$
\gamma (\boldsymbol {W}, \boldsymbol {p}) = \boldsymbol {W} \odot \operatorname {R e L U} (\operatorname {g r a d} \boldsymbol {p}) \tag {4}
$$
where  $\mathrm{ReLU}(\cdot)$  is the ReLU activation function and  $W$  is a skew-symmetric continuously weighted matrix. This formulation is complete (Theorem 2.1 in [72]), as any continuously weighted DAG can be represented by a  $(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 in [72].
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  $\mathbf{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 a novel mapping,  $\tau(W, p)$ , a modification of NoCurl. We establish the validity of performing Bayesian inference within  $(W, p)$  space utilizing  $\tau$  (Section 3.1). However,  $\tau$  yields uninformative gradient during backpropagation, 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 (Equation (4)) focuses on learning a 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  $(\mathbf{W},\mathbf{p})$  is valid.
We note that the proposed transformation in NoCurl  $\gamma$  (Equation (4) can be hard to optimize for the following reasons: (i)  $\mathrm{ReLU}(\mathrm{grad}p)$  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 it is hard for a continuous matrix to learn exactly 0 during the optimization; (ii)  $\mathrm{ReLU}(\mathrm{grad}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<sup>2</sup>. Parameterizing the search space in terms of binary adjacency matrices significantly simplifies the optimization complexity as the aforementioned issues are circumvented. Therefore, we introduce a modification  $\tau : \{0,1\}^{d\times d}\times \mathbb{R}^d\to \{0,1\}^{d\times d}$ :
$$
\tau (\boldsymbol {W}, \boldsymbol {p}) = \boldsymbol {W} \odot \operatorname {S t e p} (\operatorname {g r a d} \boldsymbol {p}) \tag {5}
$$
where we abuse the term  $\mathbf{W}$  for binary matrices, and replace  $\mathrm{ReLU}(\cdot)$  with  $\mathrm{Step}(\cdot)$ .  $\mathbf{W}$  acts as mask to disable the edge existence. Thus, due to the Step,  $\tau$  can only output a binary adjacency matrix.
See discussion below Eq. 3 in [72] for more details.
Next, we show that performing Bayesian inference in such augmented  $(W,p)$  space is valid, i.e., using the posterior  $p(W,p|D)$  to replace  $p(G|D)$ . This differs from NoCurl, which focuses on a 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 a binary adjacency matrix representing a DAG and node potential  $p$  does not contain the same values, i.e.  $p_i \neq p_j \forall i, j$ . Then, with the induced joint observational distribution  $p(D, G)$ , dataset  $D$ , and a corresponding prior  $p(G)$ , we have
$$
p (\boldsymbol {G} | \boldsymbol {D}) = \int p _ {\tau} (\boldsymbol {p}, \boldsymbol {W} | \boldsymbol {D}) \mathbb {1} (\boldsymbol {G} = \tau (\boldsymbol {W}, \boldsymbol {p})) d \boldsymbol {W} d \boldsymbol {p} \tag {6}
$$
if  $p(\pmb{G}) = \int p_{\tau}(\pmb{p}, \pmb{W}) \mathbb{1}(\pmb{G} = \tau(\pmb{W}, \pmb{p})) d\pmb{W} d\pmb{p}$ , where  $p_{\tau}(\pmb{W}, \pmb{p})$  is the prior,  $\mathbb{1}(\cdot)$  is the indicator function, and  $p_{\tau}(\pmb{p}, \pmb{W}|D)$  is the posterior distribution over  $\pmb{p}, \pmb{W}$ .
Refer to Appendix B.1 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 in a less complex  $(\mathbf{W},\mathbf{p})$  space where  $\mathbf{W} \in \{0,1\}^{d\times d}$  and  $\mathbf{p} \in \mathbb{R}^d$  without explicit constraints.
For inference of  $p$ , we adopt a sampling-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 to the generality of the framework.
However, the mapping  $\tau$  does not provide meaningful gradient information for  $\pmb{p}$  due to the piecewise constant Step( $\cdot$ ) 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 problem. This alternative formulation enables various techniques that can approximate the gradient of  $\pmb{p}$ .
Intuition The node potential  $p$  implicitly defines a topological ordering through the mapping  $\operatorname{Step}(\operatorname{grad}(\cdot))$ . In particular,  $\operatorname{grad}(\cdot)$  outputs a skew-symmetric adjacency matrix, where each entry specifies the potential difference between nodes.  $\operatorname{Step}(\operatorname{grad}(\cdot))$  zeros out the negative potential differences (i.e.  $p_i \leq p_j$ ), and only permits the edge direction from higher potential to the lower one (i.e.  $p_i > p_j$ ). This implicitly defines a sorting operation based on the descending node potentials, which can be cast as a particular arg max problem [8, 37, 47, 50, 74] involving a permutation matrix.
Alternative formulation We define  $\pmb{L} \in \{0,1\}^{d \times d}$  as a matrix with lower triangular part to be 1, and vector  $\pmb{o} = [1,\dots,d]$ . We propose the following formulation:
$$
\boldsymbol {G} = \boldsymbol {W} \odot [ \boldsymbol {\sigma} (\boldsymbol {p}) \boldsymbol {L} \boldsymbol {\sigma} (\boldsymbol {p}) ^ {T} ] \quad \text {w h e r e} \boldsymbol {\sigma} (\boldsymbol {p}) = \underset {\boldsymbol {\sigma} ^ {\prime} \in \boldsymbol {\Sigma} _ {d}} {\arg \max } \boldsymbol {p} ^ {T} (\boldsymbol {\sigma} ^ {\prime} \boldsymbol {o}) \tag {7}
$$
Here,  $\Sigma_d$  represents the space of all  $d$  dimensional permutation matrices. The following theorem states the equivalence of this formulation to Equation (5).
Theorem 3.2 (Equivalence to NoCurl formulation). Assuming the conditions in Theorem 3.1 are satisfied. Then, for a given  $(\mathbf{W},\mathbf{p})$ , we have
$$
\boldsymbol {G} = \boldsymbol {W} \odot \operatorname {S t e p} (\operatorname {g r a d} \boldsymbol {p}) = \boldsymbol {W} \odot \left[ \boldsymbol {\sigma} (\boldsymbol {p}) \boldsymbol {L} \boldsymbol {\sigma} (\boldsymbol {p}) ^ {T} \right]
$$
where  $\pmb{G}$  is a DAG and  $\sigma (\pmb {p})$  is defined in Equation (7).
Refer to Appendix B.2 for details.
This theorem translates our proposed operator  $\operatorname{Step}(\operatorname{grad}(\pmb{p}))$  into finding a corresponding permutation matrix  $\sigma(\pmb{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 [8, 47, 50]. For simplicity, we adopt the Sinkhorn approach [47], but we emphasize that this equivalence is general enough that any past or future approximation methods can be easily applied.
Sinkhorn operator The Sinkhorn operator  $S(M)$  on a matrix  $M$  [1] is defined as a sequence of row and column normalizations, each is called Sinkhorn iteration.
[47] showed that the non-differentiable arg max problem
$$
\boldsymbol {\sigma} = \underset {\boldsymbol {\sigma} ^ {\prime} \in \Sigma_ {d}} {\arg \max } \left\langle \boldsymbol {\sigma} ^ {\prime}, M \right\rangle \tag {8}
$$
can be relaxed through an entropy regularizer with its solution being expressed by  $S(\cdot)$ . In particular, they showed that  $S(M / t) = \arg \max_{\sigma' \in \mathcal{B}_d} \langle \sigma', M \rangle + th(\sigma')$ , where  $h(\cdot)$  is the entropy function. This regularized solution converges to the solution of Equation (8) when  $t \to 0$ , i.e.  $\lim_{t \to 0} S(M / t)$ . Since the Sinkhorn operator is differentiable,  $S(M / t)$  can be viewed as a differentiable approximation to Equation (8), which can be used to obtain the solution of Equation (7). Specifically, we have
$$
\underset {\boldsymbol {\sigma} ^ {\prime} \in \boldsymbol {\Sigma} _ {d}} {\arg \max } \boldsymbol {p} ^ {T} \left(\boldsymbol {\sigma} ^ {\prime} \boldsymbol {o}\right) = \underset {\boldsymbol {\sigma} ^ {\prime} \in \boldsymbol {\Sigma} _ {d}} {\arg \max } \left\langle \boldsymbol {\sigma} ^ {\prime}, \boldsymbol {p o} ^ {T} \right\rangle = \lim  _ {t \rightarrow 0} \mathcal {S} \left(\frac {\boldsymbol {p o} ^ {T}}{t}\right) \tag {9}
$$
In practice, we approximate it with  $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 a straight-through estimator [7] for  $p$ .
Some of the previous works [11, 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 rn iterations to converge [16]. However, our formulation, based on explicit node potential  $\boldsymbol{p}\boldsymbol{o}^T$ , generates a rank-1 matrix, requiring much fewer Sinkhorn steps (around 300) in practice, saving two-thirds of the computational cost.
4 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 in a Gibbs sampling manner. The second one, which is based entirely on 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. Specifically,  $X_{i}\coloneqq f_{i}(\pmb{X}_{\mathbf{Pa}^{i}}) + \epsilon_{i}$ , where  $f_{i}$  adheres to the adjacency relation specified by  $\pmb{G}$ , i.e.  $\partial f_i(\pmb {x}) / \partial x_j = 0$  if no edge exists between nodes  $i$  and  $j$ . We define  $f_{i}$  as
$$
f _ {i} (\boldsymbol {x}) = \zeta_ {i} \left(\sum_ {j = 1} ^ {d} G _ {j i} l _ {j} \left(x _ {j}\right)\right), \tag {10}
$$
where  $\zeta_{i}$  and  $l_{i}$  are neural networks with parameters  $\Theta$ , and  $\pmb{G}$  serves as a mask disabling nonparent values. To reduce the number of neural networks, we adopt a weight-sharing mechanism:  $\zeta_{i}(\cdot) = \zeta (\pmb{u}_{i},\cdot)$  and  $l_{i}(\cdot) = l(\pmb{u}_{i},\cdot)$ , with trainable node embeddings  $\pmb{u}_i$
Likelihood of SCM The likelihood can be evaluated through the noise  $\epsilon = x - f(x;\Theta)$ . [22] showed that if  $G$  is a DAG, then the mapping from  $\epsilon$  to  $x$  is invertible with a Jacobian determinant of 1. Thus, the observational data likelihood is:
$$
p (\boldsymbol {x} | \boldsymbol {G}) = p _ {\epsilon} (\boldsymbol {x} - \boldsymbol {f} (\boldsymbol {x}; \boldsymbol {\Theta})) = \prod_ {i = 1} ^ {d} p _ {\epsilon_ {i}} \left(x _ {i} - f _ {i} \left(\boldsymbol {x} _ {\mathbf {P a} _ {G} ^ {i}}\right)\right) \tag {11}
$$
Prior design We implicitly define the prior  $p(\pmb {G})$  via  $p(\pmb {p},\pmb {W})$  . We propose the following for the joint prior:
$$
p (\boldsymbol {W}, \boldsymbol {p}, \boldsymbol {\Theta}) \propto \mathcal {N} (\boldsymbol {\Theta}; \boldsymbol {0}, \boldsymbol {I}) \mathcal {N} (\boldsymbol {p}; \boldsymbol {0}, \alpha \boldsymbol {I}) \mathcal {N} (\boldsymbol {W}; \boldsymbol {0}, \boldsymbol {I}) \exp (- \lambda_ {s} \| \tau (\boldsymbol {W}, \boldsymbol {p}) \| _ {F} ^ {2})
$$
where  $\alpha$  controls the initialization scale of  $\pmb{p}$  and  $\lambda_{s}$  controls the sparseness of  $\pmb{G}$ .
4.2 Bayesian Inference of  $W, p, \Theta$
The main challenge lies in the binary nature of  $W \in \{0,1\}^{d \times d}$ , which requires a discrete sampler. Although recent progress has been made [30, 62, 73, 76], these methods either involve expensive Metropolis-Hasting (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, \Theta$  and VI for  $W$ . It should be noted that our framework can incorporate any suitable discrete sampler if needed.
We employ a Gibbs sampling procedure [10], which iteratively applies (1) sampling  $\pmb{p}, \Theta \sim p(\pmb{p}, \Theta | \pmb{D}, \pmb{W})$  with SG-MCMC; (2) updating the variational posterior  $q_{\phi}(\pmb{W} | \pmb{p}, \pmb{D}) \approx p(\pmb{W} | \pmb{p}, \Theta, \pmb{D})$ .
We define the posterior  $p(\pmb{p}, \Theta | \pmb{D}, \pmb{W}) \propto \exp(-U(\pmb{p}, \pmb{W}, \Theta))$ , where  $U(\pmb{p}, \pmb{W}, \Theta) = -\log p(\pmb{p}, \pmb{D}, \pmb{W}, \Theta)$ . SG-MCMC in continuous time defines a specific form of Itô diffusion that maintains the target distribution invariant [45] without the expensive computation of the MH step. We adopt the Euler-Maruyama discretization for simplicity. Other advanced discretization can be easily incorporated [12, 57].
Preconditioning techniques have been shown to accelerate SG-MCMC convergence [13, 28, 41, 69, 70]. We modify the sampler based on [28], which is inspired by Adam [35]. Detailed update equations can be found in Appendix C.
Figure 1: Graphical model of the inference problem.
The following proposition specifies the gradients required by SG-MCMC:  $\nabla_{\pmb{p},\Theta}U(\pmb {p},\pmb {W},\Theta)$
Proposition 4.1. Assume the model is defined as above, then we have the following:
$$
\nabla_ {\boldsymbol {p}} U = - \nabla_ {\boldsymbol {p}} \log p (\boldsymbol {p}) - \nabla_ {\boldsymbol {p}} \log p (\boldsymbol {D} | \boldsymbol {\Theta}, \tau (\boldsymbol {W}, \boldsymbol {p})) \tag {12}
$$
$$
\nabla_ {\boldsymbol {\Theta}} U = - \nabla_ {\boldsymbol {\Theta}} \log p (\boldsymbol {\Theta}) - \nabla_ {\boldsymbol {\Theta}} \log p (\boldsymbol {D} | \boldsymbol {\Theta}, \tau (\boldsymbol {p}, \boldsymbol {W})) \tag {13}
$$
Refer to Appendix B.5 for details.