Masked Gradient-Based Causal Structure Learning

Introduction

Causal graphical models defined on Directed Acyclic Graphs (DAGs) find applications in many sciences including economics and biology . Although controlled experiments can discover the structures effectively, they are often expensive or practically impossible. Learning causal structures from observational data is hence appealing and has been made possible under proper conditions .

Existing approaches for structure learning roughly fall into two classes: constraint- and score-based methods. Constraint-based methods, such as PC and fast causal inference , use conditional independence tests to find initial skeleton and then determine the edge directions according to certain orientation rules . Score-based methods, on the other hand, evaluate the candidate graphs with a predefined score function and search for the ones with the optimal score. Due to the combinatorial nature of acyclicity constraint , exact search algorithms like dynamic programming only work for small problems and therefore local greedy strategies are often adopted to perform the search, such as Greedy Equivalence Search (GES) . We refer the reader to Glymour et al. 2019, Spirtes and Zhang 2018 for more details and also a review of other methods.

More recently, Zheng et al. 2018 have proposed NOTEARS, a score-based method that formulates the structure learning problem of linear Structural Equation Models (SEMs) as a continuous one using a smooth characterization of acyclicity. Subsequent works such as DAG-GNN and GraN-DAG have extended it to handle nonlinear cases. With smooth score functions, these methods utilize gradient-based optimization to learn DAGs by estimating some weighted graph adjacency matrices. NOTEARS and DAG-GNN assume specific forms of SEMs where weighted adjacency matrices naturally exist; their performance usually degrades when the data model does not follow these forms. GraN-DAG uses path products to construct an equivalent weighted matrix representing the estimated graph, which generally requires much effort to obtain explicitly given a particular model function. Notice that choosing a right model function to fit the underlying relationships plays an important role in these methods, since the adopted score functions rely on reconstructed observations. For example, it was shown in Zhu et al. 2020 that GraN-DAG performed poorly on linear SEMs but could be much improved by modifying it with linear model functions. In addition, numeric methods usually result in a number of estimated entries in a contiguous region near zero; it becomes key to picking a proper threshold to obtain exact DAGs while keeping most, if not all, true positives, but current gradient-based approaches mostly rely on prior knowledge or human experience.

In this work, we develop a gradient-based optimization framework for structure learning, called Masked gradient-based Causal Structure Learning (MCSL), which (1) flexibly includes different model functions (e.g., neural networks, polynomial functions) with little extra effort, and (2) is easily thresholded to identify the edges. To achieve these, we reformulate the SEM with additive noises in a form parameterized by binary adjacency matrix that exists for any data model and is called mask in the paper—the binary valued entry can be used to remove the effect of a variable on the output of any function. We characterize the structure identification issue and show that if the original SEM is identifiable, the binary adjacency matrix can be identified as super-graphs of the true causal graph. To enable an efficient learning of binary valued entries, we leverage the Gumbel-Softmax approach and develop a gradient-based structure learning method with carefully devised training procedure. The resulting entries of the estimate are mostly near either zero or one, so that the edges are easily identified by thresholding at $0.5$. We conduct experiments on synthetic and real datasets, and show that the proposed method outperforms other methods on most tasks, such as a nonlinear dataset with quadratic functions and a vector-valued dataset, while being competitive with the best method on other datasets.

Background and Related Work

Let $X$ be the vector concatenating all the variables $X_{i}$ and $P(X)$ the marginal distribution induced by the SEM defined on DAG $\mathcal{G}$. Then $P(X)$ is Markov w.r.t. $\mathcal{G}$, and $\mathcal{G}$ and $P(X)$ are said to form a causal Bayesian network. The problem of causal structure learning is to use the observational data $\{x^{(k)}\}_{k=1}^{n}$, with $x^{(k)}$ being the $k$-th independent sample from the distribution $P(X)$, to infer the causal graph $\mathcal{G}$.

In general, however, it is impossible to recover $\cal G$ using only observational data from $P(X)$, because the underlying DAG is generally not identifiable without further assumption on the SEM . Hence, causal structure learning methods generally estimate, or whose output has to be converted to, the Markov equivalence class in order to have causal interpretation. This fact relates to the identifiability issue in causal structure learning: given an SEM defined on a DAG $\cal G$ with distribution $P(X)$, we say that $\mathcal{G}$ is identifiable if no other SEMs can induce the same distribution $P(X)$ with a different DAG. Fortunately, Peters et al. 2014 have shown that if we consider only a subclass of SEMs, the restricted ANMs where $f_{i}$’s and the density functions of $\epsilon_{i}$’s and $X$ do not solve a system of three-order differential equations (see Peters et al. 2014), then the true DAG is identifiable. Other identifiable models include linear non-Gaussian model , linear Gaussian model with equal noise variances , post-nonlinear model , etc. In this work, we will focus on the restricted ANMs for identifiability issues.

2 Gradient-Based Structure Learning Methods

A more related work is SAM that also considered binary structural gates and the smooth acyclicity constraint. While certain similarities exist, we believe that there are more fundamental differences: (1) SAM studied functional causal models and discussed identification of Markov equivalence class. We formulate SEMs in a form parameterized by binary matrix and characterize the identification of super-graphs using the identifiability results of ANMs. (2) SAM considered a specific score function and used MLPs to fit causal relationships, which were trained in an adversarial way. We do not limit the form of score and model functions, leading to a more general framework and training procedure. (3) running SAM, we find that its output matrix has many entries not near the boundary and may not be a DAG even after thresholding. We devise carefully the training procedure and stopping criterion for our method, resulting in an adjacency matrix guaranteed to be acyclic after thesholding. As shown in Section 5.1, these strategies make our method outperform SAM by a large margin.

Structural Equation Model with Binary Adjacency Matrix

Both NOTEARS and DAG-GNN rely on a notion of weighted adjacency matrix which does not exist for many SEMs, e.g., when the causal relationships are quadratic functions or functions sampled from Gaussian processes. Noticing that every directed graph corresponds to a binary adjacency matrix and vice versa, we consider another form of SEM that is explicitly parameterized by a binary adjacency matrix.

where $\circ$ is the element-wise product. In this formulation, if $A_{ji}=0$, then $X_{i}$ does not depend on the input $X_{j}$ anymore for any function $g_{i}$. Thus, $A_{i}$ can be viewed as a mask to remove non-causal inputs.

Similar to existing gradient-based methods, it appears that we could use a parametric or nonparametric model, together with a $d\times d$ binary matrix, to fit the observed data under the acyclicity constraint. With the fitted model $\hat{g}_{i}$ and $\hat{A}_{i}$, we may output the graph indicated by $[\hat{A}_{1}|\hat{A}_{2}|\cdots|\hat{A}_{d}]$ as our estimated structure. Unfortunately, even if we can fit the data perfectly, it is not clear whether the estimated binary matrix indicates the true causal graph: there may exist $g_{i}$, $A_{i}$ and $g_{i}^{\prime}$, $A_{i}^{\prime}$ that induce the same distribution $P(X)$, but $A_{i}$ and $A_{i}^{\prime}$ disagree for some $i$. In the following, we show that such a formulation results in identification of a super-graph of the true graph under suitable conditions.

$\mathcal{H}$ is a super-graph of $\widetilde{\cal G}$, i.e., all the edges in $\widetilde{\cal G}$ also exist in $\mathcal{H}$.

We proceed to the second question. Recall that the observed data are generated from a distribution induced by the SEM in Eq. (2.1). We may further assume a restricted ANM (see Peters et al. 2014) for the data generating procedure so that the true causal graph $\cal G$ is identifiable. We then obtain the following proposition, with a proof given in Appendix A.2.

Assume a restricted ANM with graph $\cal G$ and distribution $P(X)$ so that the original SEM is identifiable. If the parameterized SEM in the form of Eq. (3.2) with graph $\mathcal{H}$ induces the same $P(X)$, then $\cal H$ is a super-graph of $\cal G$.

With Proposition 3.1, we may then apply parametric or nonparametric model functions and a binary adjacency matrix to fitting the SEM in the form of Eq. (3.2), subject to the acyclicity constraint. As in Lachapelle et al. 2020, if the causal relationships fall into the chosen model functions and we can obtain the exact solution that minimizes the negative log-likelihood given infinite samples, the resulting SEM has the same distribution. Consequently, we obtain an acyclic super-graph, from which existing nonlinear variable selection methods can be used to learn the parental sets and hence the causal graph. Such an idea has been exploited in many ordering based methods where fully-connected DAGs are first learned (see, e.g., Teyssier and Koller 2005, Peters et al. 2014). In practice, we can only apply approximate model functions and also need an efficient optimization procedure that can handle the binary adjacency matrix, which is the topic of next section.

Masked Causal Structure Learning

While we can adopt smooth model and score functions, the binary entries prohibit gradient-based optimization. A direct approach is to apply logistic sigmoid functions parameterized by real valued variables, which however lead to estimated entries lying in a small range near zero, making it hard to apply thresholding to identify edges (see Appendix D.7 for an illustration). We would like each estimated entry to be either close to zero so that it almost removes a non-causal input, or near one so that we can easily apply a thresholding at $0.5$.

We leverage the Gumbel-Softmax approach often used to approximate samples from a categorical distribution . This approach has lower variances of gradient estimates than the straight-through method and also outperforms several REINFORCE based methods on some structured output prediction tasks. For a random variable defined on $\{0,1\}$ with class probabilities $\pi_{0}\in(0,1)$ and $\pi_{1}=1-\pi_{0}$, its binary sample can be approximated by

We thus use Gumbel-Sigmoid to approximate the binary adjacency matrix. Consider that an edge helps produce a better reconstruction and does not violate the acyclicity constraint. Intuitively, if we repeatedly apply Gumbel-Sigmoid and estimate the score function between observed and reconstructed samples, gradient-based optimization will push the logit so that the Gumbel-Sigmoid output is close to one in the expected sense. If the acyclicity constraint is violated, some entries will be pushed towards zero to meet the constraint.

2 Acyclicity Constraint and Optimization

where expectations are taken w.r.t. $\operatorname{Logistic}(0,1)$ samples and $s(U)\coloneqq\operatorname{tr}\left(e^{\,\mathsf{g}_{\tau}(U)}\right)-d\geq 0$. A sparsity-inducing term $\|\mathsf{g}_{\tau}(U)\|_{1}$ is also incorporated.

Following Zheng et al. 2018, the above constrained optimization problem can be solved using augmented Lagrangian method when the score function $\mathcal{L}(\cdot,\cdot)$ and model functions $h_{i}$’s are chosen properly. For example, we may use the least squares loss or negative log-likelihood as the score function. The choices of $h_{i}$ include polynomial functions and MLPs. Augmented Lagrangian method consists of optimizing a sequence of subproblems where the exact solutions converge to a stationary point of the original constrained problem under some regularity conditions . The augmented Lagrangian of the above constrained problem is

The updating rules for augmented Lagrangian method yield:

with $\beta>1$ and $0<\gamma<1$ being the hyperparameters. Different from the optimization procedure in NOTEARS, the updating rules of Eqs. (4.4), (4.5), and (4.6) involve additional expectations w.r.t. $\operatorname{Logistic}(0,1)$ distribution and their closed forms are not easy to obtain. Here we propose to use sample estimates as approximations and details are described below:

The subproblem in Eq. (4.4) can be approximately solved using a gradient-based optimization algorithm, e.g., Adam , where at each iteration we draw independently $\mathbf{g}_{1},\mathbf{g}_{2},\ldots,\mathbf{g}_{b}$ and approximate the gradient by $\nabla_{U,\phi}\left(\frac{1}{b}\sum_{i=1}^{b}L_{\rho}(U,\phi,\alpha;\mathbf{g}_{i})\right)$.

As shown by Jang et al. 2017, the Gumbel-Softmax approach is effective for single sample gradient estimation, so we simply pick $b=1$ which is found to work well.

Augmented Lagrangian method usually stops once the constraint is satisfied. In our case, while a single sample estimation of gradient works well for training, it may render the optimization to stop too early if a sample ‘accidentally’ makes the constraint below the tolerance. Thus, we choose the tolerance $\xi$ to be small and additionally use $\operatorname{tr}(e^{\sigma(U/\tau)})-d<\xi$ as our stopping criterion to lower the probability of stopping the algorithm early. Moreover, notice that the latter criterion is satisfied only when the entries $\sigma(U_{ji}/\tau)$, where $i,j$ are such that the edge $X_{j}\to X_{i}\notin E$ for some DAG with $E$ as edge set, are nearly zeros. If an edge indeed helps minimize the score function, the logit will be pushed to a relatively large positive value so that the score function is minimized in the expected sense. For edges that do not violate acyclicity nor help minimize the score function, the corresponding entries may have intermediate values. Nevertheless, these edges are treated as spurious edges and will be further processed. Thus, as a byproduct of introducing a second stopping criterion, we can readily use $\sigma(U/\tau)$ as our learned matrix, which would indicate a DAG after thresholding at $0.5$.

More details regarding the parameter choices, stopping criterion, and overfitting issue can be found in Appendix C. We remark that such a stopping criterion, together with the devised optimization procedure, leads to a much better performance than the related method SAM ; see experimental results in Section 5.1.

3 Final Output

Experiments

We compare the proposed method against several baselines, including PC , PC-KCI (with kernel-based conditional independence tests) , GES , GES-GS (with generalized score) , CAM , NOTEARS , DAG-GNN , SAM , NOTEARS-MLP , and GraN-DAG . Here, our method uses $4$-layer MLPs as the model functions and is denoted as MCSL-MLP. Detailed setting of the MLP and hyperparameters is described in Appendix D.1. We use least squares as our score function in all experiments. An implementation of MCSL has been released in the gCastle package . https://github.com/huawei-noah/trustworthyAI

For synthetic data, experiments are conducted on data models varying along graph size, degree, and causal relationship. We sample a DAG $\mathcal{G}$ using the Erdös–Rényi (ER) model and generate data in the causal order indicated by $\mathcal{G}$. We consider $d$-node ER graphs with on average $d$ and $4d$ edges, denoted as ER1 and ER4, respectively. Unless otherwise stated, we set the sample size to $3000$ and the additive noises to standard Gaussians. The causal relationships include functions sampled from Gaussian processes (Section 5.1), quadratic functions (Section 5.2), and post-nonlinear models (Section 5.4). A vector-valued case is considered in Section 5.3. For real data, we conduct experiments on a protein and a telecommunication dataset in Section 5.5.

We further conduct empirical studies with scale-free graphs and different sample sizes; the detailed results are provided in Appendix D.5 and D.6, respectively. A discussion of the training time and computational complexity is available in Appendix D.9.

We report Structural Hamming Distance (SHD) and True Positive Rate (TPR) to evaluate the learned graphs, averaged over five random seeds. For GES (GES-GS) and PC (PC-KCI), we treat undireted edges as true positives if the true graph has a directed edge in place of the undirected ones.

We first consider the data model previously used by Peters et al. 2014, Lachapelle et al. 2020, Zheng et al. 2020: each function $f_{i}$ is sampled from a Gaussian Process (GP) with RBF kernel of bandwidth one. This setting is known to be identifiable . For MCSL-MLP, GraN-DAG, CAM, and NOTEARS-MLP, an additional CAM pruning step is used to remove spurious edges. We did not apply this pruning step to other baselines as they have much lower TPRs, especially on ER4 graphs.

The empirical results are reported in Table 1 with graph sizes $d\in\{10,50\}$ and the results for $d\in\{20,100\}$ are given in Appendix D.2. We observe that MCSL-MLP, GraN-DAG, and CAM outperform the other methods across most settings. MCSL-MLP has a better performance with small graphs while CAM performs the best for $100$-node graphs. Nevertheless, their differences are minor, compared with the performance of other methods. NOTEARS-MLP is on par with MCSL-MLP, GraN-DAG, and CAM on ER1 graphs, but performs poorly on ER4 graphs. Both DAG-GNN and NOTEARS have poor performance, possibly because they cannot model this type of causal relationships, and moreover, they operate on the notion of weighted adjacency matrix which is not obvious here. Kernel-based methods have a cubic computational complexity w.r.t. sample size, so we pick $500$ samples for each dataset. Yet they may still be too slow for large ER4 graphs ($>12$ hours) and we do not report their results for these cases. We observe that GES-GS and PC-KCI perform similarly to GES and PC, respectively. A possible reason is that although more powerful score function and conditional independence test are utilized, GES-GS and PC-KCI are limited by their computational complexity and only part of the samples can be used. Finally, we find that the output graph of SAM is not guaranteed to be acyclic and the performance is outperformed by a large margin; in particular, it learns very few true edges with $50$-node ER4 graphs.

2 Quadratic Functions

We consider nonlinear relationships with quadratic functions, as in Zhu et al. 2020. This data model is also identifiable . With prior knowledge that the causal relationship follows a quadratic function, Zhu et al. 2020 used a similar idea from GraN-DAG to modify NOTEARS by constructing an equivalent weighted adjacency matrix. We refer to this method as NOTEARS-quad in this work and omit its details here; see Zhu et al. 2020 for a detailed description. A Quadratic Regression based pruning method (QR pruning) has also been shown to perform well for this data model, so we use it for all methods here. Details for experiment setup and QR pruning can be found in Appendix D.3.

We may also utilize the prior knowledge to apply quadratic regressions to modeling causal relationships and denote the resulting method as MCSL-quad. For NOTEARS-quad, there is extra effort to derive explicitly the equivalent adjacency matrix and further re-implement the method. By contrast, our method only requires replacing the MLPs with quadratic functions.

Here we only consider methods that perform well in the experiment from Section 5.1. For better visualization, we report the average SHDs in Figure 4, with detailed results including standard deviations and TPRs given in Appendix D.3. MCSL-MLP and MCSL-quad have a similar performance and outperform the other methods by a large margin. This also shows that MCSL-MLP has a strong expressive power, as its performance is on par with MCSL-quad that utilizes the knowledge of the form of causal relationships. Although NOTEARS-quad is specifically designed for quadratic functions, it performs poorly here. We believe that it is because the equivalent adjacency matrix makes the optimization problem more complicated and difficult to solve.

3 Vector-Valued Case

4 Post-Nonlinear Causal Models

5 Real Data

We consider to estimate a protein signaling network based on expression levels of proteins and phospholipid , containing both interventional and observational data. Since our method is based on passively observed data, we use only the observational data with $853$ samples. The ground truth proposed by Sachs et al. 2005 has $11$ nodes and $17$ edges. MCSL-MLP and CAM achieve the best SHD $12$, and GraN-DAG has an SHD $13$. DAG-GNN and NOTEARS are on par with MCSL-MLP and CAM w.r.t. true positives, but have more false discoveries, resulting in SHDs $16$ and $19$, respectively. GES obtains $7$ undirected edges while PC estimates $7$ undirected and $1$ directed edges. By contrast, the inferred graphs from MCSL-MLP, CAM, and GraN-DAG consist of only directed edges.

Identifying root cause in telecommunication networks

We proceed to apply the proposed approach to identifying root cause of anomaly variables with time series data. Here each variable represents a certain key performance indicator (like number of users per minute) monitoring the state of telecommunication system. In practice, it is commonly observed that a single system fault may first happen and incur a variable to behave abnormally, and then the variable would cause several others to be abnormal. Thus, we would like to find the root cause one among all abnormal variables, which is believed to be more informative to identify the system fault. This task can be done by learning a causal DAG and then choose the root node. If there are more than one root nodes, our strategy is to return the one with the most descendants.

For time series data, several methods first choose a time lag $p\geq 0$ and then perform causal discovery among $d(p+1)$ variables, where many edges can be excluded as the variable at the current time cannot be a cause of the one of past time. Thus, GraN-DAG, CAM, and NOTEARS-MLP do not easily apply here and we can only consider instantaneous data, i.e., treat time series data as i.i.d. data. Besides, prior knowledge from human experience indicates that the summary graph is acyclic and also some variables cannot be the root cause ones. Most existing approaches, however, do not take into account this condition and directly learn the graph among $d(p+1)$ variables, followed by compressing the graph into a summary graph of $d$ variables.

For this task, we can easily adapt our method as in the vector-valued case. Notice that this problem may also contain hidden confounders and in fact is very challenging. Nonetheless, in all our test cases containing $2$ to $60$ variables, our modified method can identify about $80\%$ correct root causes. For the reasons mentioned above, all the other methods can only find correctly at most $30\%$ cases; in particular, constraint-based methods usually result in many undirected edges, making it hard to determine a root node.

Concluding Remarks

In this work, we reformulate the SEM to incorporate binary adjacency matrix and investigate the structure identification issue. A gradient-based method is proposed by leveraging the recently developed Gumbel-Softmax approach and the smooth acyclicity constraint. Experiments validate the effectiveness as well as the flexibility to include different model functions to fit the underlying causal relationships. A future direction is to use preliminary variable selection and/or second-order optimization to improve both the efficiency and efficacy.

References

Appendix A Proofs

A.2 Proof of Proposition 3.1

Appendix B Derivation of Logistic Distribution

We provide a derivation that if two independent variables $X,Y\sim\operatorname{Gumbel}(0,1)$, then $Z=X-Y\sim\operatorname{Logistic}(0,1)$. We will show that the CDF of $Z$ matches the CDF of $\operatorname{Logistic}(0,1)$.

which indicates $e^{X}$ follows the exponential distribution with the rate parameter being $1$.

which is exactly the CDF of $\operatorname{Logistic}(0,1)$.

Appendix C Further Details on Optimization with Augmented Lagrangian Method

As shown by Jang et al. 2017, the Gumbel-Softmax approach is effective for single sample gradient estimation, so we simply pick $b=1$ which is found to work well in our experiments. We pick $\gamma=0.25$, same as in NOTEARS and DAG-GNN, and initialize the estimate of Lagrange multiplier $\alpha$ to $0$, i.e., $\alpha^{0}=0$. We set the initial value of $\rho$ by $\rho^{0}=10^{-\left\lceil 0.3d\right\rceil}$ for $d$-node graphs, with minor fine tuning. This is because we initialize the logits $U_{ji}$’s to be $0$ and consequently $s(U^{0})$ would be very large with high probability for large graphs. Since both $\alpha$ and $\rho$ do not decrease during training, a large $\rho^{0}$ may make the optimization somewhat ‘omit’ the score function and the resulting graph would be a DAG with high score. However, a small initial value $\rho^{0}$ would result in a long training time, so we choose larger $\beta$ for larger graphs to accelerate training. In particular, we find that $\beta=5$, $15$, $300$ and $8000$ work well for graphs with $10$, $20$, $50$ and $100$ nodes, respectively. These choices are found on synthetic datasets with causal additive model and known true causal graphs (see Appendix D.1 for further discussions). For graphs with other sizes, a linear interpolation on the logarithm scale of the graph sizes can be used as the choice for $\beta$, and one may also consider fine tuning using some synthetic datasets with known true graphs. We use the Adam optimizer with a learning rate of $3\times 10^{-2}$ and $1000$ training iterations to approximately solve the subproblem in Eq. (4.4). Note that for $100$-node graphs, we set the number of training iterations to $2500$ as it may take longer to converge on larger graphs.

C.2 Stopping Criterion

Simply using a single sample estimate $s(U^{t})$ for stopping criterion may make the optimization stop too early. We consider further steps:

we choose the tolerance $\xi$ to be small to lower the probability of stopping the algorithm early. Particularly, we pick $\xi=10^{-10}$ in our experiment which is much smaller than that used in NOTEARS and GraN-DAG.

we set another stopping criterion $\operatorname{tr}\big(e^{\sigma(U^{t}/\tau)}\big)<\xi$ which uses only the logits. This criterion will be satisfied only when the entries $\sigma(U^{t}_{ji}/\tau)$, where $i,j$ are such that the edge $X_{j}\to X_{i}\notin E$ for some DAG with $E$ as edge set, are nearly zeros. If an edge indeed helps minimize the score function, then the logit will be pushed to a relatively large positive value so that the score function is minimized in the expected sense. Thus, as a byproduct, we can then readily use $\sigma(U^{t}/\tau)$ as our learned matrix which indicates a DAG after a hard thresholding at threshold $0.5$.

notice that $s(U^{t})$ and $\operatorname{tr}\left(e^{\sigma(U^{t}/\tau)}\right)$ are evaluated at the end of step $t$ of the augmented Lagrangian method. This also helps avoid terminating the optimization too early.

In Figure 5a, we provide an empirical study regarding $s(U)$ and $\operatorname{tr}\left(e^{\sigma(U/\tau)}\right)$, on the GP dataset with $50$-node ER1 graph. Figure 5a and similar results on other datasets validate the effectiveness of the use of both stopping criteria. As discussed in Section 4.2, the matrix $\sigma(U^{t}/\tau)$ at the last step will be further processed to output our inferred causal graph.

C.3 Overfitting

Unlike DAG-GNN and GraN-DAG that require a held-out dataset to avoid overfitting, we use the whole dataset during training. To investigate whether MCSL-MLP incurs overfitting, we generate a validation dataset of $3000$ samples and monitor the augmented Lagrangian on both training and validation datasets. Here the validation dataset is not used for structure learning. A typical example on a GP dataset with $50$-node ER1 graph is plotted in Figure 5.

We observe two phases: 1) the augmented Lagrangian is optimized to decrease on both datasets in the first $10000$ iterations; and 2) with increased estimate of Lagrange multiplier and penalty parameter, the acyclicity term has a larger effect on the augmented Lagrangian, making the acyclicity term decrease towards the predefined tolerance and the least squares losses tend to increase. The augmented Lagrangian also oscillates in the second phase, due to the varying logits $U_{ji}$’s and the increasing estimate of Lagrange multiplier and penalty parameter. The similar behaviors in both phases indicate that MCSL-MLP can use all the observed data for training and does not need a held-out dataset for validation, at least in our experiments.

Appendix D Further Experiment Details and Results

In practice, however, one could not perform hyperparameter tuning directly on the observed data as the ground truth is not available. Similar to GraN-DAG, we conduct experiments on synthetic data with known causal graphs to search for the hyperparameters, and then use these hyperparameters for all the experiments. In particular, we choose a causal additive data model and the CAM algorithm , which is specifically designed for this type of data models, can serve as a good benchmark for our method.

D.2 Results for Nonlinear SEMs with Gaussian Processes

This section provides additional results for the nonlinear SEMs with Gaussian processes in Section 5.1, as shown in Table 2. Here N/A means that the corresponding experiments could not finish within $12$ hours; since the corresponding methods have been shown to perform poorly in other settings, we simply stopped them.

D.3 Experiment Setup and Results for Nonlinear SEMs with Quadratic Functions

This section provides further experiment details and results for the nonlinear SEMs with quadratic functions in Section 5.2.

For this data model, there may exist very large variable values that would lead to numerical issues for the followed gradient-based structure learning methods including NOTEATS, GraN-DAG, and MCSL-MLP. We therefore limit the diameter of each underlying causal graph to $3$ and normalize each variable’s values (before incorporating the corresponding additive noises) by dividing the total number of corresponding first- and second-order terms. The normalized data are then used as the observed data in the experiment.

Quadratic regression based pruning method

Empirical results

The empirical results are presented in Table 3.

D.4 Results for Post-Nonlinear Causal Models

This section provides additional results for the post-nonlinear models in Section 5.4, as shown in Table 4.

D.5 Scale-Free Graphs

We conduct an empirical study with $50$-node scale-free graphs, by considering the GP dataset from Section 5.1. CAM performs the best with SHDs $12.6\pm 8.0$ and $78.4\pm 9.1$ for graphs with $50$ and $200$ edges on average, respectively, while GraN-DAG has SHDs $14.8\pm 8.4$ and $126.8\pm 4.8$, respectively. Our method is still competitive to the best method CAM with SHDs $15.0\pm 8.2$ and $82.0\pm 4.7$.

D.6 Different Sample Sizes

We also conduct an empirical analysis of MCSL-MLP with different sample numbers. We pick the GP dataset with $20$ nodes and $n=\{300,1000,3000,10000\}$ samples. As shown in Figure 8, a larger number of samples generally leads to a better performance. Another interesting observation from ER4 graphs is that the CAM pruning method seems to remove more true positives than false discoveries and somewhat degrades the performance, when the true causal graphs are dense.

D.7 Sigmoid vs. Gumbel-Sigmoid

We provide an example to show the effectiveness of applying Gumbel-Sigmoid to approximating the binary adjacency matrix. We consider a dataset with $10$-node ER1 graph from the experiment in Section 5.1. Notice that this dataset is only for illustration purpose and similar results are found with other datasets as well. We apply logistic sigmoid and Gumbel-Sigmoid with MCSL-MLP and visualize the resulting adjacency matrices in Figure 6.

D.8 Pruning

We investigate the effect of the additional pruning step on the proposed method. A useful approach is the CAM pruning Bühlmann et al. 2014 that applies significance testing of covariates using generalized additive models and declares significance if the $p$-values do not exceed a predefined value. Though the true causal relationships may not follow causal additive assumption, the CAM pruning usually performs well.

We use the GP dataset from the experiment in Section 5.1 and compare MCSL-MLP with or without pruning on ER1 and ER4 graphs with $d\in\{10,20,50,100\}$ nodes. Notice that this pruning step is necessary for CAM and GraN-DAG. CAM first estimates a topological order of the variables and the graph has all the possible edges that does not violate the acyclicity constraint. It then uses pruning to remove spurious edges. GraN-DAG estimates an equivalent adjacency matrix and then removes an edge if the corresponding entry has the smallest value in the absolute Jacobian matrix until a DAG is obtained. Note that this DAG typically contains many spurious edges and pruning is needed to reduce false discoveries, as shown by the ablation study in Lachapelle et al. 2020.

Figure 7 reports the empirical results in terms of SHD, TPR and also False Discovery Rate (FDR). We observe that the additional pruning step reduces the SHD and FDR much, and has little effect on the TPR on ER1 graphs. For ER4 graphs that are denser, CAM pruning reduces both FDR and TPR. Nevertheless, the overall metric SHD is improved by CAM pruning, and the difference grows as the graph size increases. This experiment demonstrates the practical importance of applying an additional pruning step for MCSL-MLP.

D.9 Computational Complexity and Training Time

Similar to NOTEARS and GraN-DAG, MCSL-MLP requires evaluations of matrix exponential with $\mathcal{O}(d^{3})$ cost per iteration. To reduce the number of $\mathcal{O}(d^{3})$ iterations required to converge, NOTEARS adopts the proximal quasi-Newton algorithm. Although GraN-DAG uses a gradient-based method like ours, it has been observed that GraN-DAG performs fewer iterations than NOTEARS in practice and that the evaluation of matrix exponential does not dominate the total cost at each iteration for graphs with $100$ nodes or less .

To validate whether MCSL-MLP behaves similarly, we simply compare the training time of MCSL-MLP with GraN-DAG using the GP dataset from Section 5.1, as shown in Figure 9. The experiments are run on a standard NC6 instance on Azure cloud, with 6-core Intel Xeon 2.6GHz CPU and one-half Nvidia Tesla K80 GPU. MCSL-MLP takes a shorter time than GraN-DAG across all graph sizes, showing that the $\mathcal{O}(d^{3})$ evaluation costs of matrix exponential are not a problem to MCSL-MLP, at least for problems with $100$ nodes or less. There is an increase in the training time of MCSL-MLP on $100$-node graphs, as we use an increased number of training iterations to solve the subproblem in Eq. (4.4). We believe that a faster algorithm, e.g., a second-order algorithm as in , can be adopted to further increase the efficiency of MCSL-MLP for larger problems, which is treated as a future work.