Diffusion Models for Causal Discovery via Topological Ordering

Introduction

Understanding the causal structure of a problem is important for areas such as economics, biology (Sachs et al., 2005) and healthcare (Sanchez et al., 2022), especially when reasoning about the effect of interventions. When interventional data from randomised trials are not available, causal discovery methods (Glymour et al., 2019) may be employed to discover the causal structure of a problem solely from observational data. Causal structure is typically modelled as a directed acyclic graph (DAG) ${\mathcal{G}}$ in which each node is associated with a random variable and each edge represents a causal mechanism i.e. how one variable influences another.

However, learning such a model from data is NP-hard (Chickering, 1996). Traditional methods search the DAG space by testing for conditional independence between variables (Spirtes et al., 1993) or by optimising some goodness of fit measure (Chickering, 2002). Unfortunately, solving the search problem with a greedy combinatorial optimisation method can be expensive and does not scale to high-dimensional problems.

In line with previous work (Teyssier & Koller, 2005; Park & Klabjan, 2017; Bühlmann et al., 2014; Solus et al., 2021; Wang et al., 2021; Rolland et al., 2022), we can speed up the combinatorial search problem over the space of DAGs by rephrasing it as a topological ordering task, ordering from leaf nodes to root nodes. The search space over DAGs with $d$ nodes and $(d^{2}-d)/2$ possible edges is much larger than the space of permutations over $d$ variables. Once a topological ordering of the nodes is found, the potential causal relations between later (cause) and earlier (effect) nodes can be pruned with a feature selection algorithm (e.g. Bühlmann et al. (2014)) to yield a graph which is naturally directed and acyclic without further optimisation.

Recently, Rolland et al. (2022) proposed the SCORE algorithm for topological ordering. SCORE uses the Hessian of the data log-likelihood, $\nabla_{\textnormal{x}}^{2}\log p({\mathbf{x}})$.By verifying which elements of $\nabla_{\textnormal{x}}^{2}\log p({\mathbf{x}})$’s diagonal are constant across all data points, leaf nodes can be iteratively identified and removed. Rolland et al. (2022) estimate the Hessian point-wise with a second-order Stein gradient estimator (Li & Turner, 2018) over a radial basis function (RBF) kernel. However, point-wise estimation with kernels scales poorly to datasets with large number of samples $n$ because it requires inverting a $n\times n$ kernel matrix.

Here, we enable scalable causal discovery by utilising neural networks (NNs) trained with denoising diffusion instead of Rolland et al.’s kernel-based estimation. We use the ordering procedure, based on Rolland et al. (2022), which requires re-computing the score’s Jacobian at each iteration. Training NNs at each iteration would not be feasible. Therefore, we derive a theoretical analysis that allows updating the learned score without re-training. In addition, the NN is trained over the entire dataset ($n$ samples) but only a subsample is used for finding leaf nodes. Thus, once the score model is learned, we can use it to order the graph with constant complexity on $n$, enabling causal discovery for large datasets in high-dimensional settings. Interestingly, our algorithm does not require architectural constraints on the neural network, as in previous causal discovery methods based on neural networks (Lachapelle et al., 2020; Zheng et al., 2020; Yu et al., 2019; Ng et al., 2022). Our training procedure does not learn the causal mechanism directly, but the score of the data distribution.

Contributions. In summary, we propose DiffAN, an identifiable algorithm leveraging a diffusion probabilistic model for topological ordering that enables causal discovery assuming an additive noise model: (i) To the best of our knowledge, we present the first causal discovery algorithm based on denoising diffusion training which allows scaling to datasets with up to $500$ variables and $10^{5}$ samples. The score estimated with the diffusion model is used to find and remove leaf nodes iteratively; (ii) We estimate the second-order derivatives (score’s Jacobian or Hessian) of a data distribution using neural networks with diffusion training via backpropagation; (iii) The proposed deciduous score (Section 3) allows efficient causal discovery without re-training the score model at each iteration. When a leaf node is removed, the score of the new distribution can be estimated from the original score (before leaf removal) and its Jacobian.

Preliminaries

The topological ordering (also called causal ordering or causal list) of a DAG ${\mathcal{G}}$ is defined as a non-unique permutation $\pi$ of $d$ nodes such that a given node always appears first in the list than its descendants. Formally, $\pi_{i}<\pi_{j}\iff j\in De_{{\mathcal{G}}}({\textnormal{x}}_{i})$ where $De_{{\mathcal{G}}}({\textnormal{x}}_{i})$ are the descendants of the $ith$ node in ${\mathcal{G}}$ (Appendix B in Peters et al. (2017)).

2 Nonlinear Additive Noise Models

Learning a unique ${\bm{A}}$ from ${\bm{X}}$ with observational data requires additional assumptions. A common class of methods called additive noise models (ANM) (Shimizu et al., 2006; Hoyer et al., 2008; Peters et al., 2014; Bühlmann et al., 2014) explores asymmetries in the data by imposing functional assumptions on the data generation process. In most cases, they assume that assignments take the form ${\textnormal{x}}_{i}\coloneqq f_{i}(Pa({\textnormal{x}}_{i}))+\epsilon_{i}$ with $\epsilon_{i}\sim p^{\epsilon}$. Here we focus on the case described by Peters et al. (2014) where $f_{i}$ is nonlinear. We use the notation $f_{i}$ for $f_{i}(Pa({\textnormal{x}}_{i}))$ because the arguments of $f_{i}$ will always be $Pa({\textnormal{x}}_{i})$ throughout this paper. We highlight that $f_{i}$ does not depend on $i$.

Identifiability. We assume that the SCM follows an additive noise model (ANM) which is known to be identifiable from observational data (Hoyer et al., 2008; Peters et al., 2014). We also assume causal sufficiency, i.e. there are no hidden variables that are a common cause of at least two observed variables. In addition, corollary 33 from Peters et al. (2014) states that the true topological ordering of the DAG, as in our setting, is identifiable from a $p({\mathbf{x}})$ generated by an ANM without requiring causal minimality assumptions.

Finding Leaves with the Score. Rolland et al. (2022) propose that the score of an ANM with distribution $p({\mathbf{x}})$ can be used to find leavesWe refer to nodes without children in a DAG ${\mathcal{G}}$ as leaves.. Before presenting how to find the leaves, we derive, following Lemma 2 in Rolland et al. (2022), an analytical expression for the score which can be written as

Where $Ch({\textnormal{x}}_{j})$ denotes the children of ${\textnormal{x}}_{j}$. We now proceed, based on Rolland et al. (2022), to derive a condition which can be used to find leaf nodes.

Given a nonlinear ANM with a noise distribution $p^{\epsilon}$ and a leaf node $j$; assume that $\frac{\partial^{2}\log p^{\epsilon}}{\partial x^{2}}=a$, where $a$ is a constant, then

Lemma 1 enables finding leaf nodes based on the diagonal of the log-likelihood’s Hessian.

Rolland et al. (2022), using a similar conclusion, propose a topological ordering algorithm that iteratively finds and removes leaf nodes from the dataset. At each iteration Rolland et al. (2022) re-compute the Hessian with a kernel-based estimation method. In this paper, we develop a more efficient algorithm for learning the Hessian at high-dimensions and for a large number of samples. Note that Rolland et al. (2022) prove that Equation 2 can identify leaves in nonlinear ANMs with Gaussian noise. We derive a formulation which, instead, requires the second-order derivative of the noise distribution to be constant. Indeed, the condition $\frac{\partial^{2}\log p^{\epsilon}}{\partial x^{2}}=a$ is true for $p^{\epsilon}$ following a Gaussian distribution which is consistent with Rolland et al. (2022), but could potentially be true for other distributions as well.

3 Diffusion Models Approximate the Score

The process of learning to denoise (Vincent, 2011) can approximate that of matching the score (Hyvärinen, 2005). A diffusion process gradually adds noise to a data distribution over time. Diffusion probabilistic models (DPMs) Sohl-Dickstein et al. (2015); Ho et al. (2020); Song et al. (2021) learn to reverse the diffusion process, starting with noise and recovering the data distribution. The diffusion process gradually adds Gaussian noise, with a time-dependent variance $\alpha_{t}$, to a sample $\mathbf{x}_{0}\sim p_{\text{data}}(\mathbf{x})$ from the data distribution. Thus, the noisy variable $\mathbf{x}_{t}$, with $t\in\left[0,T\right]$, is learned to correspond to versions of $\mathbf{x}_{0}$ perturbed by Gaussian noise following $p\left(\mathbf{x}_{t}\mid\mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t};\sqrt{\alpha_{t}}\mathbf{x}_{0},\left(1-\alpha_{t}\right){\bm{I}}\right)$, where $\alpha_{t}:=\prod_{j=0}^{t}\left(1-\beta_{j}\right)$, $\beta_{j}$ is the variance scheduled between $[\beta_{\text{min}},\beta_{\text{max}}]$ and ${\bm{I}}$ is the identity matrix. DPMs (Ho et al., 2020) are learned with a weighted sum of denoising score matching objectives at different perturbation scales with

where ${\mathbf{x}}_{t}=\sqrt{\alpha_{t}}{\mathbf{x}}_{0}+\sqrt{1-\alpha_{t}}\epsilon$, with ${\mathbf{x}}_{0}\sim p({\mathbf{x}})$ being a sample from the data distribution, $t\sim\mathcal{U}\left(0,T\right)$ and $\epsilon\sim\mathcal{N}\left(0,{\bm{I}}\right)$ is the noise. $\lambda(t)$ is a loss weighting term following Ho et al. (2020).

Throughout this paper, we leverage the fact that the trained model ${\bm{\epsilon}_{\theta}}$ approximates the score $\nabla_{{\textnormal{x}}_{j}}\log p({\mathbf{x}})$ of the data (Song & Ermon, 2019).

The Deciduous Score

Discovering the complete topological ordering with the distribution’s Hessian (Rolland et al., 2022) is done by finding the leaf node (Equation 2), appending the leaf node ${\textnormal{x}}_{l}$ to the ordering list $\pi$ and removing the data column corresponding to ${\textnormal{x}}_{l}$ from ${\bm{X}}$ before the next iteration $d-1$ times. Rolland et al. (2022) estimate the score’s Jacobian (Hessian) at each iteration.

Instead, we explore an alternative approach that does not require estimation of a new score after each leaf removal. In particular, we describe how to adjust the score of a distribution after each leaf removal, terming this a “deciduous score”An analogy to deciduous trees which seasonally shed leaves during autumn.. We obtain an analytical expression for the deciduous score and derive a way of computing it, based on the original score before leaf removal. In this section, we only consider that $p({\mathbf{x}})$ follows a distribution described by an ANM, we pose no additional assumptions over the noise distribution.

Given a ANM which entails a distribution $p({\mathbf{x}})$, we can use Equation 1 to find an analytical expression for an additive residue $\Delta_{l}$ between the distribution’s score $\nabla\log p({\mathbf{x}})$ and its deciduous score $\nabla\log p({\mathbf{x}}_{-l})$ such that

In particular, $\Delta_{l}$ is a vector $\{\delta_{j}\mid\forall j\in\left[1,\dots,d\right]\backslash l\}$ where the residue w.r.t. a node ${\textnormal{x}}_{j}$ can be denoted as

If ${\textnormal{x}}_{j}\notin Pa({\textnormal{x}}_{l})$, $\delta_{j}=0$.

Observing Equation 1, the score $\nabla_{{\textnormal{x}}_{j}}\log p({\mathbf{x}})$ only depends on the following random variables (i) $Pa({\textnormal{x}}_{j})$, (ii) $Ch({\textnormal{x}}_{j})$, and (iii) $Pa(Ch({\textnormal{x}}_{j}))$. We consider ${\textnormal{x}}_{l}$ to be a leaf node, therefore $\nabla_{{\textnormal{x}}_{j}}\log p({\mathbf{x}})$ only depends on ${\textnormal{x}}_{l}$ if ${\textnormal{x}}_{j}\in Pa({\textnormal{x}}_{l})$. If ${\textnormal{x}}_{j}\in Pa({\textnormal{x}}_{l})$, the only term depending on $\nabla_{{\textnormal{x}}_{j}}\log p({\mathbf{x}})$ dependent on ${\textnormal{x}}_{l}$ is one of the terms inside the summation. ∎

However, we wish to estimate the deciduous score $\nabla\log p({\mathbf{x}}_{-l})$ without direct access to the function $f_{l}$, to its derivative, nor to the distribution $p^{\epsilon}$. Therefore, we now derive an expression for $\Delta_{l}$ using solely the score and the Hessian of $\log p({\mathbf{x}})$.

Consider an ANM of distribution $p({\mathbf{x}})$ with score $\nabla\log p({\mathbf{x}})$ and the score’s Jacobian ${\bm{H}}(\log p({\mathbf{x}}))$. The additive residue $\Delta_{l}$ necessary for computing the deciduous score (as in Proposition 2) can be estimated with

Causal Discovery with Diffusion Models

DPMs approximate the score of the data distribution (Song & Ermon, 2019). In this section, we explore how to use DPMs to perform leaf discovery and compute the deciduous score, based on Theorem 1, for iteratively finding and removing leaf nodes without re-training the score.

where $\nabla_{i,j}{\bm{\epsilon}_{\theta}}({\bm{x}},t)$ means the $ith$ output of ${\bm{\epsilon}_{\theta}}$ is backpropagated to the $jth$ input. The diagonal of the Hessian in Equation 7 can, then, be used for finding leaf nodes as in Equation 2.

In a two variable setting, it is sufficient for causal discovery to (i) train a diffusion model (Equation 3); (ii) approximate the score’s Jacobian via backpropagation (Equation 7); (iii) compute variance of the diagonal across all data points; (iv) identify the variable with lowest variance as effect (Equation 2). We illustrate in Appendix C the Hessian of a two variable SCM computed with a diffusion model.

2 Topological Ordering

When a DAG contains more than two nodes, the process of finding leaf nodes (i.e. the topological order) needs to be done iteratively as illustrated in Figure 2. The naive (greedy) approach would be to remove the leaf node from the dataset, recompute the score, and compute the variance of the new distribution’s Hessian to identify the next leaf node (Rolland et al., 2022). Since we employ diffusion models to estimate the score, this equates to re-training the model each time after a leaf is removed.

We hence propose a method to compute the deciduous score $\nabla\log p({\mathbf{x}}_{-l})$ using Theorem 1 to remove leaves from the initial score without re-training the neural network. In particular, assuming that a leaf ${\textnormal{x}}_{l}$ is found, the residue $\Delta_{l}$ can be approximatedThe diffusion model itself is an approximation of the score, therefore its gradients are approximations of the score derivatives. with

where ${\bm{\epsilon}_{\theta}}({\bm{x}},t)_{l}$ is output corresponding to the leaf node. Note that the term $\nabla_{l}{\bm{\epsilon}_{\theta}}({\bm{x}},t)$ is a vector of size $d$ and the other term is a scalar. During topological ordering, we compute $\Delta_{\pi}$, which is the summation of $\Delta_{l}$ over all leaves already discovered and appended to $\pi$. Naturally, we only compute $\Delta_{l}$ w.r.t. nodes ${\textnormal{x}}_{j}\notin\pi$ because ${\textnormal{x}}_{j}\in\pi$ have already been ordered and are not taken into account anymore.

where ${\bm{\epsilon}_{\theta}}$ is a DPM trained with Equation 3. See Appendix E.3 for the choice of $t$. This topological ordering procedure is formally described in Algorithm 1, $score(-\pi)$ means that we only consider the outputs for nodes ${\textnormal{x}}_{j}\notin\pi$

3 Computational Complexity and Practical Considerations

We now study the complexity of topological ordering with DiffAN w.r.t. the number of samples $n$ and number of variables $d$ in a dataset. In addition, we discuss what are the complexities of a greedy version as well as approximation which only utilises masking.

Complexity on $n$. Our method separates learning the score ${\bm{\epsilon}_{\theta}}$ from computing the variance of the Hessian’s diagonal across data points, in contrast to Rolland et al. (2022). We use all $n$ samples in ${\bm{X}}$ for learning the score function with diffusion training (Equation 3). It does not involve expensive constrained optimisation techniquesSuch as the Augmented Lagrangian method (Zheng et al., 2018; Lachapelle et al., 2020). and we train the model for a fixed number of epochs (which is linear with $n$) or until reaching the early stopping criteria. We use a MLP that grows in width with $d$ but it does not significantly affect complexity. Therefore, we consider training to be $O(n)$. Moreover, Algorithm 1 is computed over a batch ${\bm{B}}$ with size $k<n$ instead of the entire dataset ${\bm{X}}$, as described in Equation 9. Note that the number of samples $k$ in ${\bm{B}}$ can be arbitrarily small and constant for different datasets. In Section 5.2, we verify that the accuracy of causal discovery initially improves as $k$ is increased but eventually tapers off.

Complexity on $d$. Once ${\bm{\epsilon}_{\theta}}$ is trained, a topological ordering can be obtained by running $\nabla_{{\bm{x}}}{\bm{\epsilon}_{\theta}}({\bm{x}},t)$ $d$ times. Moreover, computing the Jacobian of the score requires back-propagating the gradients $d-i$ times, where $i$ is the number of nodes already ordered in a given iteration. Finally, computing the deciduous score’s residue (Equation 8) means computing gradient of the $i$ nodes. Resulting in a complexity of $O(d\cdot(d-i)\cdot i)$ with $i$ varying from to $d$ which can be described by $O(d^{3})$. The final topological ordering complexity is therefore $O(n+d^{3})$.

DiffAN Masking. We verify empirically that the masking procedure described in Section 4.2 can significantly reduce the deciduous score’s residue absolute value while maintaining causal discovery capabilities. In DiffAN Masking, we do not re-train the ${\bm{\epsilon}_{\theta}}$ nor compute the deciduous score. This ordering algorithm is an approximation but has shown to work well in practice while showing remarkable scalability. DiffAN Masking has $O(n+d^{2})$ ordering complexity.

Experiments

In our experiments, we train a NN with a DPM objective to perform topological ordering and follow this with a pruning post-processing step (Bühlmann et al., 2014). The performance is evaluated on synthetic and real data and compared to state-of-the-art causal discovery methods from observational data which are either ordering-based or gradient-based methods, NN architecture. We use a 4-layer multilayer perceptron (MLP) with LeakyReLU and layer normalisation. Metrics. We use the structural Hamming distance (SHD), Structural Intervention Distance (SID) (Peters & Bühlmann, 2015), Order Divergence (Rolland et al., 2022) and run time in seconds. See Appendix D.3 for details of each metric. Baselines. We use CAM (Bühlmann et al., 2014), GranDAG (Lachapelle et al., 2020) and SCORE (Rolland et al., 2022). We apply the pruning procedure of Bühlmann et al. (2014) to all methods. See detailed results in the Appendix D. Experiments with real data from Sachs (Sachs et al., 2005) and SynTReN (Van den Bulcke et al., 2006) datasets are in the Appendix E.1.

In this experiment, we consider causal relationships with $f_{i}$ being a function sampled from a Gaussian Process (GP) with radial basis function kernel of bandwidth one. We generate data from additive noise models which follow a Gaussian, Exponential or Laplace distributions with noise scales in the intervals $\{[0.4,0.8],[0.8,1.2],\}$, which are known to be identifiable (Peters et al., 2014). The causal graph is generated using the Erdös-Rényi (ER) (Erdős et al., 1960) and Scale Free (SF) (Bollobás et al., 2003) models. For a fixed number of nodes $d$, we vary the sparsity of the sampled graph by setting the average number of edges to be either $d$ or $5d$. We use the notation [$d$][graph type][sparsity] for indicating experiments over different synthetic datasets. We show that DiffAN performs on par with baselines while being extremely fast, see Figure 3. We also explore the role of overfitting in Appendix E.2, the difference between DiffAN with masking only and the greedy version in Appendix E.4, how we choose $t$ during ordering in Appendix E.3 and we give results stratified by experiment in Appendix E.5.

2 Scaling up with DiffAN Masking

We now verify how DiffAN scales to bigger datasets, in terms of the number of samples $n$. Here, we use DiffAN masking because computing the residue with DiffAN would be too expensive for very big $d$. We evaluate only the topological ordering, ignoring the final pruning step. Therefore, the performance will be measured solely with the Order Divergence metric.

Scaling to large datasets. We evaluate how DiffAN compares to SCORE (Rolland et al., 2022), the previous state-of-the-art, in terms of run time (in seconds) and and the performance (order divergence) over datasets with $d=500$ and different sample sizes $n\in{10^{2},\dots,10^{5}}$, the error bars are results across 6 dataset (different samples of ER and SF graphs). As illustrated in Figure 1, DiffAN is the more tractable option as the size of the dataset increases. SCORE relies on inverting a very large $n\times n$ matrix which is expensive in memory and computing for large $n$. Running SCORE for $d=500$ and $n>2000$ is intractable in a machine with 64Gb of RAM. Figure 4 (left) shows that, since DiffAN can learn from bigger datasets and therefore achieve better results as sample size increases.

Ordering batch size. An important aspect of our method, discussed in Section 4.3, that allows scalability in terms of $n$ is the separation between learning the score function ${\bm{\epsilon}_{\theta}}$ and computing the Hessian variance across a batch of size $k$, with $k<n$. Therefore, we show empirically, as illustrated in Figure 4 (right), that decreasing $k$ does not strongly impact performance for datasets with $d\in{10,20,50}$.

Related Works

Ordering-based Causal Discovery. The observation that a causal DAG can be partially represented with a topological ordering goes back to Verma & Pearl (1990). Searching the topological ordering space instead of searching over the space of DAGs has been done with greedy Markov Chain Monte Carlo (MCMC) (Friedman & Koller, 2003), greedy hill-climbing search (Teyssier & Koller, 2005), arc search (Park & Klabjan, 2017), restricted maximum likelihood estimators (Bühlmann et al., 2014), sparsest permutation (Raskutti & Uhler, 2018; Lam et al., 2022; Solus et al., 2021), and reinforcement learning (Wang et al., 2021). In linear additive models, Ghoshal & Honorio (2018); Chen et al. (2019) propose an approach, under some assumptions on the noise variances, to discover the causal graph by sequentially identifying leaves based on an estimation of the precision matrix.

Hessian of the Log-likelihood. Estimating ${\bm{H}}(\log p({\mathbf{x}}))$ is the most expensive task of the ordering algorithm. Our baseline (Rolland et al., 2022) propose an extension of Li & Turner (2018) which utilises the Stein’s identity over a RBF kernel (Schölkopf & Smola, 2002). Rolland et al.’s method cannot obtain gradient estimates at positions out of the training samples. Therefore, evaluating the Hessian over a subsample of the training dataset is not possible. Other promising kernel-based approaches rely on spectral decomposition (Shi et al., 2018) solve this issue and can be promising future directions. Most importantly, computing the kernel matrix is expensive for memory and computation on $n$. There are, however, methods (Achlioptas et al., 2001; Halko et al., 2011; Si et al., 2017) that help scaling kernel techniques, which were not considered in the present work. Other approaches are also possible with deep likelihood methods such as normalizing flows (Durkan et al., 2019; Dinh et al., 2017) and further compute the Hessian via backpropagation. This would require two backpropagation passes giving $O(d^{2})$ complexity and be less scalable than denoising diffusion. Indeed, preliminary experiments proved impractical in our high-dimensional settings.

We use DPMs because they can efficiently approximate the Hessian with a single backpropagation pass and while allowing Hessian evaluation on a subsample of the training dataset. It has been shown (Song & Ermon, 2019) that denoising diffusion can better capture the score than simple denoising (Vincent, 2011) because noise at multiple scales explore regions of low data density.

Conclusion

We have presented a scalable method using DPMs for causal discovery. Since DPMs approximate the score of the data distribution, they can be used to efficiently compute the log-likelihood’s Hessian by backpropagating each element of the output with respect to each element of the input. The deciduous score allows adjusting the score to remove the contribution of the leaf most recently removed, avoiding re-training the NN. Our empirical results show that neural networks can be efficiently used for topological ordering in high-dimensional graphs (up to $500$ nodes) and with large datasets (up to $10^{5}$ samples).

Our deciduous score can be used with other Hessian estimation techniques as long as obtaining the score and its full Jacobian is possible from a trained model, e.g. sliced score matching (Song et al., 2020) and approximate backpropagation (Kingma & Cun, 2010). Updating the score is more practical than re-training in most settings with neural networks. Therefore, our theoretical result enables the community to efficiently apply new score estimation methods to topological ordering. Moreover, DPMs have been previously used generative diffusion models in the context of causal estimation (Sanchez & Tsaftaris, 2022). In this work, we have not explored the generative aspect such as Geffner et al. (2022) does with normalising flows. Finally, another promising direction involves constraining the NN architecture as in Lachapelle et al. (2020) with constrained optimisation losses (Zheng et al., 2018).

Acknowledgement

This work was supported by the University of Edinburgh, the Royal Academy of Engineering and Canon Medical Research Europe via P. Sanchez’s PhD studentship. S.A. Tsaftaris acknowledges the support of Canon Medical and the Royal Academy of Engineering and the Research Chairs and Senior Research Fellowships scheme (grant RCSRF1819\825).

References

Appendix A Proofs

We re-write Equation 1 here for improved readability:

We start by showing the “$\Leftarrow$” direction by deriving Equation 1 w.r.t. ${\textnormal{x}}_{j}$. If ${\textnormal{x}}_{j}$ is a leaf, only the first term of the equation is present, then taking its derivative results in

Replacing Equation 12 in to Equation 1, we have

Let ${\textnormal{x}}_{c}\in Ch({\textnormal{x}}_{j})$ such that ${\textnormal{x}}_{c}\not\in Pa(Ch({\textnormal{x}}_{j}))$. ${\textnormal{x}}_{c}$ always exist since ${\textnormal{x}}_{j}$ is not a leaf, and it suffices to pick a child of ${\textnormal{x}}_{c}$ appearing at last position in some topological order. If we isolate the terms depending on ${\textnormal{x}}_{c}$ on the RHS of Equation 13, we have

Deriving both sides w.r.t. ${\textnormal{x}}_{c}$, since the LHS of Equation 14 does not depend on ${\textnormal{x}}_{c}$, we can write

Since $g$ does not depend on ${\textnormal{x}}_{j}$, $\frac{\partial f_{c}}{\partial{\textnormal{x}}_{j}}$ does not depend on ${\textnormal{x}}_{j}$ neither, implying that $f_{c}$ is linear in ${\textnormal{x}}_{j}$, contradicting the non-linearity assumption. ∎

A.2 Proof Theorem 1

Using Equation 1, we will derive expressions for each of the elements in Equation 6 and show that it is equivalent to $\delta_{l}$ in Equation 5. First, note that the score of a leaf node ${\textnormal{x}}_{l}$ can be denoted as:

Second, replacing Equation 16 into each element of ${\bm{H}}_{l}(\log p({\mathbf{x}}))\in R^{d}$, we can write

Finally, replacing Equations 16, 17 and 11 into the Equation 5 for a single node ${\textnormal{x}}_{j}$, if $j\neq l$, we have

The last line in Equation 19 is the same as in Equation 5 from Lemma 2, proving that $\delta_{j}$ can be written using the first and second order derivative of the log-likelihood. ∎

Appendix B Score of Nonlinear ANM with Gaussian noise

over the variables ${\mathbf{x}}$ (Peters et al., 2017) By assuming that the noise variables $\epsilon_{i}\sim\mathcal{N}(0,\sigma_{i}^{2})$ and inserting the ANM function, Equation 20 can be written as

The score of $p({\mathbf{x}})$ can hence be written as

Appendix C Visualisation of the Score’s Jacobian for Two Variables

Considering a two variables problem where the causal mechanisms are $B=f_{\omega}(A)+\epsilon_{B}$ and $A=\epsilon_{A}$ with $\epsilon_{A},\epsilon_{B}\sim\mathcal{N}(0,1)$ and $f_{\omega}$ being a two-layer MLP with randomly initialised weights. Note that, in Figure 5, the variance of $\frac{\partial^{2}\log p({\bm{A}},{\bm{B}})}{\partial{\bm{B}}^{2}}$, while not as predicted by Equation 2, is smaller than $\frac{\partial^{2}\log p({\bm{A}},{\bm{B}})}{\partial{\bm{A}}^{2}}$ allowing discovery of the true causal direction.

Appendix D Experiments Details

We now describe the hyperparameters for the diffusion training. We use number of time steps $T=100$, $\beta_{t}$ is a linearly scheduled between $\beta_{\text{min}}=0.0001$ and $\beta_{\text{max}}=0.02$. The model is trained according to Equation 3 which follows Ho et al. (2020). During sampling, $t$ is sampled from a Uniform distribution.

D.2 Neural Architecture

The neural network follows a simple MLP with 5 Linear layers, LeakyReLU activation function, Layer Normalization and Dropout in the first layer. The full architecture is detailed in Table 1.

D.3 Metrics

SHD. Structural Hamming distance between the output and the true causal graph, which counts the number of missing, falsely detected, or reversed edges.

SID. Structural Intervention Distance is based on a graphical criterion only and quantifies the closeness between two DAGs in terms of their corresponding causal inference statements(Peters & Bühlmann, 2015).

Order Divergence. Rolland et al. (2022) propose this quantity for measuring how well the topological order is estimated. For an ordering $\pi$, and a target adjacency matrix $A$, we define the topological order divergence $D_{top}(\pi,A)$ as

If $\pi$ is a correct topological order for ${\bm{A}}$, then $D_{top}(\pi,{\bm{A}})=0$. Otherwise, $D_{top}(\pi,{\bm{A}})$ counts the number of edges that cannot be recovered due to the choice of topological order. Therefore, it provides a lower bound on the SHD of the final algorithm (irrespective of the pruning method).

Appendix E Other Results

We consider two real datasets: (i) Sachs: A protein signaling network based on expression levels of proteins and phospholipids (Sachs et al., 2005). We consider only the observational data ($n=853$ samples) since our method targets discovery of causal mechanisms when only observational data is available. The ground truth causal graph given by Sachs et al. (2005) has 11 nodes and 17 edges. (ii) SynTReN: We also evaluate the models on a pseudo-real dataset sampled from SynTReN generator (Van den Bulcke et al., 2006). Results, in Table 2, show that our method is competitive against other state-of-the-art causal discovery baselines on real datasets.

E.2 Overfitting

The data used for topological ordering (inference) is a subset of the training data. Therefore, it is not obvious if overfitting would be an issue with our algorithm. Therefore, we run an experiment where we fix the number of epochs to $2000$ considered high for a set of runs and use early stopping for another set in order to verify if overfitting is an issue. On average across all 20 nodes datasets, the early stopping strategy output an ordering diverge of $9.5$ whilst overfitting is at $11.1$ showing that the method does not benefit from overfitting.

E.3 Optimal t𝑡t for score estimation

As noted by Vincent (2011), the best approximation of the score by a learned denoising function is when the training signal-to-noise (SNR) ratio is low. In diffusion model training, $t=0$ corresponds to the coefficient with lowest SNR. However, we found empirically that the best score estimate varies somehow randomly across different values of $t$. Therefore, we run the the leaf finding function (Equation 9) $N$ times for values of $t$ evenly spaced in the $[0,T]$ interval and choose the best leaf based on majority vote. We show in Figure 6 that majority voting is a better approach than choosing a constant value for $t$.

E.4 Ablations of DiffAN Masking and Greedy

We now verify how DiffAN masking and DiffAN greedy compare against the original version detailed in the main text which computes the deciduous score. Here, we use the same datasets decribed in Section 5.1 which comprises 4 (20ER1, 20ER5, 20SF1, 20SF5) synthetic dataset types with 27 variations over seeds, noise type and noise scale.

DiffAN Greedy. A greedy version of the algorithm re-trains the ${\bm{\epsilon}_{\theta}}$ after each leaf removal iteration. In this case, the deciduous score is not computed, decreasing the complexity w.r.t. $d$ but increasing w.r.t. $n$. DiffAN greedy has $O(nd^{2})$ ordering complexity.

We observe, in Figure 7, that the greedy version performs the best but it is the slowest, as seen in Figure 8. DiffAN masking

E.5 Detailed Results

We present the numerical results for the violinplots in Section 5.1 in Tables 3 and 4. The results are presented in $\text{mean}_{\text{std}}$ with statistics acquired over experiments with 3 seeds.