Learning Sparse Nonparametric DAGs
Xun Zheng, Chen Dan, Bryon Aragam, Pradeep Ravikumar, Eric P. Xing
Introduction
Learning DAGs from data is an important and classical problem in machine learning, with a diverse array of applications in causal inference (Spirtes et al., 2000), fairness and accountability (Kusner et al., 2017), medicine (Heckerman et al., 1992), and finance (Sanford and Moosa, 2012). In addition to their undirected counterparts, DAG models offer a parsimonious, interpretable representation of a joint distribution that is useful in practice. Unfortunately, existing methods for learning DAGs typically rely on specific model assumptions (e.g. linear or additive) and specialized algorithms (e.g. constraint-based or greedy optimization). As a result, the burden is on the user to choose amongst many possible models and algorithms, which requires significant expertise. Thus, there is a need for a general framework for learning different DAG models—subsuming, for example, linear, parametric, and nonparametric—that does not require specialized algorithms. Ideally, the problem could be formulated as a conventional optimization problem that can be tackled with general purpose solvers, much like the current state-of-the-art for undirected graphical models (e.g. Suggala et al., 2017; Yang et al., 2015; Liu et al., 2009; Hsieh et al., 2013; Banerjee et al., 2008).
Our main contributions can be summarized as follows:
We develop a generic optimization problem that can be applied to nonlinear and nonparametric SEM and discuss various special cases including additive models and index models. In contrast to existing work, we show how this optimization problem can be solved to stationarity with generic solvers, eliminating the necessity for specialized algorithms and models.
We extend the existing smooth characterization of acyclicity from Zheng et al. (2018) to general nonparametric models, show that the linear parametrization is a special case of the general framework, and apply this to several popular examples for modeling nonlinear dependencies (Section 3).
We consider in detail two classes of nonparametric estimators defined through 1) Neural networks and 2) Orthogonal basis expansions, and study their properties (Section 4).
We run extensive empirical evaluations on a variety of nonparametric and semiparametric models against recent state-of-the-art methods in order to demonstrate the effectiveness and generality of our framework (Section 5).
As with all score-based approaches to learning DAGs, ours relies on a nonconvex optimization problem. Despite this, we show that off-the-shelf solvers return stationary points that outperform other state-of-the-art methods. Finally, the algorithm itself can be implemented in standard machine learning libraries such as PyTorch, which should help the community to extend our approach to richer models moving forward.
Related work
The problem of learning nonlinear and nonparametric DAGs from data has generated significant interest in recent years, including additive models (Bühlmann et al., 2014; Voorman et al., 2014; Ernest et al., 2016), generalized linear models (Park, 2018; Park and Raskutti, 2017; Park and Park, 2019; Gu et al., 2018), additive noise models (Hoyer et al., 2009; Peters et al., 2014; Blöbaum et al., 2018; Mooij et al., 2016), post-nonlinear models (Zhang and Hyvärinen, 2009; Zhang et al., 2016) and general nonlinear SEM (Monti et al., 2019; Goudet et al., 2018; Kalainathan et al., 2018; Sgouritsa et al., 2015). Recently, Yu et al. (2019) proposed to use graph neural networks for nonlinear measurement models and Huang et al. (2018) proposed a generalized score function for general SEM. The latter work is based on recent work in kernel-based measures of dependence (Gretton et al., 2005; Fukumizu et al., 2008; Zhang et al., 2012). Another line of work uses quantile scoring (Tagasovska et al., 2018). Also of relevance is the literature on nonparametric variable selection (Bertin et al., 2008; Lafferty et al., 2008; Miller et al., 2010; Rosasco et al., 2013; Gregorová et al., 2018) and approaches based on neural networks (Feng and Simon, 2017; Ye and Sun, 2018; Abid et al., 2019). The main distinction between our work and previous work is that our framework is not tied to a specific model—as in Yu et al. (2019); Bühlmann et al. (2014); Park (2018)—as our focus is on a generic formulation of an optimization problem that can be solved with generic solvers (see Section 2 for a more detailed comparison). This also distinguishes this paper from concurrent work by Lachapelle et al. (2019) that focuses on neural network-based nonlinearities in the local conditional probabilities. Furthermore, compared to Huang et al. (2018) and Yu et al. (2019), our approach can be much more efficient (Section 5.1; Appendix C). As such, we hope that this work is able to spur future work using more sophisticated nonparametric estimators and optimization schemes.
Notation
Background
Existing papers approach this problem as follows: 1) Assume a specific model for (1), 2) Prove identifiability for this specific model, and 3) Develop a specialized algorithm for learning this specific model. By contrast, our approach is generic: We do not assume any particular model form or algorithm, and instead develop a general framework that applies to any model that is identifiable. By now, there is a well-catalogued list of identifiability results for various linear, parametric, and nonlinear models, which we review briefly below (see also Section 3.3).
When the model (1) holds, the graph is not necessarily uniquely defined: A well-known example is when is jointly normally distributed, in which case the are linear functions, and where it can be shown that the graph is not uniquely specified. Fortunately, it is known that this case is somewhat exceptional: Assuming additive noise, as long as the are linear with non-Gaussian errors (Kagan et al., 1973; Shimizu et al., 2006; Loh and Bühlmann, 2014) or the functions are nonlinear (Hoyer et al., 2009; Zhang and Hyvärinen, 2009; Peters et al., 2014), then the graph is generally identifiable. We refer the reader to Peters et al. (2014) for details. Another example are so-called quadratic variance function models, which are parametric models that subsume many generalized linear models (Park and Raskutti, 2017; Park, 2018). In the sequel, we assume that the model is chosen such that the graph is uniquely defined from (1), and this dependence will be emphasized by writing . Similarly, any collection of functions defines a graph in the obvious way. See Section 3.3 for specific examples with discussion on identifiability.
2 Comparison to existing approaches
It is instructive at this point to highlight the main distinction between our approach and existing approaches. A common approach is to assume the are easily parametrized (e.g. linearity) (Zheng et al., 2018; Aragam and Zhou, 2015; Gu et al., 2018; Park and Raskutti, 2017; Park, 2018; Chen et al., 2018; Ghoshal and Honorio, 2017). In this case, one can easily encode the structure of via, e.g. a weighted adjacency matrix, and learning reduces to a parametric estimation problem. Nonparametric extensions of this approach include additive models (Bühlmann et al., 2014; Voorman et al., 2014), where the graph structure is easily deduced from the additive structure of the . More recent work (Lachapelle et al., 2019; Yu et al., 2019) uses specific parametrizations via neural networks to encode . An alternative approach relies on exploiting the conditional independence structure of , such as the post-nonlinear model (Zhang and Hyvärinen, 2009; Yu et al., 2019), the additive noise model (Peters et al., 2014), and kernel-based measures of conditional independence (Huang et al., 2018). Our framework can be viewed as a substantial generalization of these approaches: We use partial derivatives to measure dependence in the general nonparametric model (1) without assuming a particular form or parametrization, and do not explicitly require any of the machinery of nonparametric conditional independence (although we note in some places this machinery is implicit). This allows us to use nonparametric estimators such as multilayer perceptrons and basis expansions, for which these derivatives are easily computed. As a result, the score-based learning problem is reduced to an optimization problem that can be tackled using existing techniques, making our approach easily accessible.
Characterizing acyclicity in nonparametric SEM
In this section, we discuss how to extend the trace exponential regularizer from Zheng et al. (2018) beyond the linear setting, and then discuss several special cases.
The key insight from Zheng et al. (2018) is replacing the combinatorial constraint with the continuous constraint . Our goal is to define a suitable surrogate of for general nonparametric models, so that the same continuous program can be used to optimize (2).
2 A notion of nonparametric acyclicity
precisely encodes the dependency structure amongst the . Thus the program (2) is equivalent to
This implies an equivalent continuous formulation of the program (2). Moreover, when the functions are all linear, is the same as the weighted adjacency matrix defined in Section 3.1. Thus, (5) is a genuine generalization of the linear case (3).
3 Special cases
In addition to applying to general nonparametric models of the form (2) and linear models, the program (5) applies to a variety of parametric and semiparametric models including additive noise models, generalized linear models, additive models, and index models. In this section we discuss these examples along with identifiability results for each case.
The nonparametric additive noise model (ANM) (Hoyer et al., 2009; Peters et al., 2014) assumes that
Generalized linear models
Clearly, (6) is a special case of (LABEL:eq:glm). Furthermore, for linear mean functions, if and only if , recovering the parametric approach in Zheng et al. (2018). Several special cases of GLMs are known to be identifiable: Linear Gaussian with equal variances (Peters and Bühlmann, 2013), linear non-Gaussian models (Shimizu et al., 2006), Poisson models (Park and Park, 2019), and quadratic variance function models (Park and Raskutti, 2017).
Polynomial regression
Additive models
Index models
Among these examples, both polynomial regression and GLMs with linear mean function are nonlinear but finite-dimensional, and hence the problem (5) is straightforward to solve (see Section 4.3).
Optimization
In general, the program (5) is infinite-dimensional. In this section we discuss different ways to reduce this to a tractable, finite-dimensional optimization problem. One of the advantages of encoding dependence via is that it provides a plug-and-play framework for plugging in various nonparametric estimators whose derivatives can be computed. We will illustrate two examples using multilayer perceptrons and orthogonal basis expansions, however, we emphasize that it is straightforward to implement other differentiable models for the . These flexible nonparametric estimators will help reduce (5) to a straightforward optimization problem, as we discuss at the end of this section.
Choose a suitable family of approximations (e.g. neural networks, orthogonal series, etc.);
Translate the loss function and constraint into parametric forms and using the approximating family;
Solve the resulting finite-dimensional problem.
Step 3 above is the key step that enables transforming (5) into a tractable optimization problem. By approximating the with a flexible family of functions parametrized by , we can replace the infinite-dimensional quantity with the simpler . As is standard in the literature on nonparametric estimation, the dimension of is allowed to depend on , although this dependence will be suppressed.
First, we must determine under what conditions is independent of —this is important both for enforcing acyclicity and sparsity. It is not hard to see that if the th column of consists of all zeros (i.e. for all ), then will be independent of . In fact, we have the following proposition, which implies that this constraint precisely identifies the set of MLPs that are independent of :
Consider the function class of all MLPs that are independent of and the function class of all MLPs such that the th column of consists of all zeros. Then .
This important proposition provides a rigorous way to enforce that an MLP approximation depends only on a few coordinates. Indeed, it is clear that constraining for each will remove the dependence on , however, there is a concern that we could lose the expressivity of multiple hidden layers in doing so. Fortunately, this proposition implies that there is in fact no loss of expressivity or approximating power. Furthermore, it follows that if . This result enables us to characterize acyclicity independent of the depth of the neural network, as opposed to handling individual paths through the entire neural network as in Lachapelle et al. (2019), which depends linearly on the depth.
Let denote the parameters for the th MLP and . Define . The problem (2) thus reduces to
2 Basis expansions
As long as the coefficients decay sufficiently fast, can be well-approximated by the finite series . Similar claims are true for one-dimensional Sobolev functions, which applies to both additive (i.e. for ) and index (i.e. for ) models.
We illustrate here an application with additive models and one-dimensional expansions. It is straightforward to extend these ideas to more general models using a tensor product basis, though this quickly becomes computationally infeasible. For more on high-dimensional orthogonal series, see Lee et al. (2016). Thus,
Given integers and assuming is sufficiently smooth, we have (Efromovich, 2008), so that the overall approximation error is on the order . Furthermore, for all . Since we are discarding terms for , in practice it suffices to check that for , or .
Letting denote the parameters for all , it thus suffices to define for the purposes of checking acyclicity. Let be the matrix . To estimate the coefficients , we solve
3 Solving the continuous program
where is a penalty parameter and is a dual variable.
where is a vector of all ones. We note that as in Zheng et al. (2018), (12) is a nonconvex program, and at best can be solved to stationarity. Our experiments indicate that this nonetheless leads to competitive and often superior performance in practice.
Experiments
We study the empirical performance of two instances of the general framework: MLP (4.1) and Sobolev expansions (4.2), denoted by and . For we use an MLP with one hidden layer with 10 hidden units and sigmoid activation function. For we use Sobolev basis , (). Complete details on all baselines and simulations, including a discussion of computational complexity and runtimes, can be found in the appendix. Code implementing our method is available at https://github.com/xunzheng/notears.
For comparison, the following methods are chosen as baselines: fast greedy equivalence search (Ramsey et al., 2017), greedy equivalence search with generalized scores (Huang et al., 2018), DAG-GNN (Yu et al., 2019), NOTEARS (Zheng et al., 2018) for linear SEM, and causal additive models (Bühlmann et al., 2014). To summarize, and are specialized at linear models, whereas , , and targets general nonlinear dependencies. Comparisons with other score-based methods (KGV score (Bach and Jordan, 2003), Spearman correlation (Sokolova et al., 2014)) and constraint-based methods (PC (Spirtes et al., 2000), MM-MB (Aliferis et al., 2010)) can be found in previous work (Huang et al., 2018), hence are omitted.
Simulation
The ground truth DAG is generated from two random graph models: Erdos-Renyi (ER) and scale-free (SF). We use ER2 to denote an ER graph with edges, likewise for SF. Given the ground truth DAG, we simulate the SEM for all in topological order, and each . To evaluate the performance under different data generation mechanisms, we consider four models for the : 1) Additive models with Gaussian processes (GPs) for each , 2) Index models (), 3) ANM with MLPs, and 4) ANM with GPs.
Metrics
We evaluate the estimated DAG structure using the following common metrics: false discovery rate (FDR), true positive rate (TPR), false positive rate (FPR), and structural Hamming distance (SHD). Note that both and return a CPDAG that may contain undirected edges, in which case we evaluate them favorably by assuming correct orientation for undirected edges whenever possible, similar to (Zheng et al., 2018).
1 Structure learning
In this experiment we examine the structure recovery of different methods by comparing the DAG estimates against the ground truth. We simulate {ER1, ER2, ER4, SF1, SF2, SF4} graphs with nodes. For each graph, data samples are generated. The above process is repeated 10 times and we report the mean and standard deviations of the results. For and , are used for respectively.
Figure 1 shows the SHD in various settings; the complete set of results for the remaining metrics are deferred to the supplement. Overall, the proposed method attains the best SHD (lower the better) across a wide range of settings, particularly when the data generating mechanism is an MLP or an index model. One can also observe that the performance of stays stable for different graph types with varying density and degree distribution, as it does not make explicit assumptions on the topological properties of the graph such as density or degree distribution. Not surprisingly, performs well when the underlying SEM is additive GP. On the other hand, when the ground truth is not an additive model, the performance of degrades as expected. Finally, we observe that outperforms and on GP, which is a nonparametric setting in which a kernel-based dependency measure can excel, however, we note that the kernel-based approach accompanies an time complexity, compared to linear dependency on in and . Also, with by properly tuning the regularization parameter, the performance of for each individual setting can be improved considerably, for example in the GP setting. Since such hyperparameter tuning is not the main focus of this paper, we fix a reasonable for all settings (see Appendix C for more discussion).
With respect to runtime and scalability, we note that the computational complexity of our approach depends on the choice of nonparametric estimator. For example, requires flops per iteration of L-BFGS-B. In terms of runtime, the average runtime of on ER2 with , is over 90 minutes, whereas takes less than five minutes on average (see Appendix C for more discussion).
Figure 2 shows the SHD compared with . We first observe that outperforms in multiple index models and MLP models, on the other hand, achieves better accuracy on additive GP and the full GP setting. Recall that the algorithm involves three steps: 1) Preliminary neighborhood search (PNS), 2) Order search by greedy optimization of the likelihood, and 3) Edge pruning. By comparison, our methods effectively only perform the second step, and can easily be pre- and post-processed with the first (PNS) and third (edge pruning) steps. To further investigate the efficacy of these additional steps, we applied both preliminary neighborhood selection and edge pruning to and on additive GP and GP settings, denoted as and . Noticeably, the output from PNS simply translates to a set of constraints in the form of that can be easily incorporated into the L-BFGS-B algorithm for (13), demonstrating the flexibility of the proposed approach. The performance improves in both cases, matching or improving vs. .
2 Sensitivity to number of hidden units
We also investigated the effect of number of hidden units in the estimate. It is well-known that as the size of the hidden layer increases, the functions representable by an MLP become more flexible. On the other hand, larger networks require more samples to estimate the parameters. Indeed, Figure 3 confirms this intuition. We plot the SHD with varying number of hidden units ranging from zero (i.e. linear function) to 100 units, using and samples generated from the additive GP model on SF2 graph with nodes. One can first observe a sharp phase transition between zero and very few hidden units, which suggests the power of nonlinearity. Moreover, as the number of hidden units increases to 20, the performance for both and steadily improves, in which case the increased flexibility brings benefit. However, as we further increase the number of hidden units, while SHD for remains similar, the SHD for deteriorates, hinting at the lack of samples to take advantage of the increased flexibility.
3 Real data
Finally, we evaluated on a real dataset from Sachs et al. (2005) that is commonly used as a benchmark as it comes with a consensus network that is accepted by the biological community. The dataset consists of continuous measurements of expression levels of proteins and phospholipids in human immune system cells for cell types. We report an SHD of 16 with 13 edges estimated by . In comparison, NOTEARS predicts 16 edges with SHD of 22 and predicts 18 edges that attains SHD of 19. (Due to the large number of samples, we could not run on this dataset.) Among the 13 edges predicted by , 7 edges agree with the consensus network: raf mek, mek erk, PLCg PIP2, PIP3 PLCg, PIP3 PIP2, PKC mek, PKC jnk; and 3 edges are predicted but in a reversed direction: raf PKC, akt erk, p38 PKC. Among the true positives, 3 edges are not found by other methods: mek erk, PIP3 PLCg, PKC mek.
Discussion
We present a framework for score-based learning of sparse directed acyclic graphical models that subsumes many popular parametric, semiparametric, and nonparametric models as special cases. The key technical device is a notion of nonparametric acyclicity that leverages partial derivatives in the algebraic characterization of DAGs. With a suitable choice of the approximation family, the estimation problem becomes a finite-dimensional differentiable program that can be solved by standard optimization algorithms. The resulting continuous optimization algorithm updates the entire graph (i.e. all edges simultaneously) in each iteration using global information about the current state of the network, as opposed to traditional local search methods that update one edge at a time based on local information. Notably, our approach is generally more efficient and more accurate than existing approaches, despite relying on generic algorithms. This out-of-the-box performance is desirable, especially when noting that future improvements and specializations can be expected to improve the approach substantially.
We acknowledge the support of NSF via IIS-1909816, OAC-1934584, ONR via N000141812861, NSF IIS1563887 and DARPA/AFRL FA87501720152. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation, Defense Advanced Research Projects Agency, or Air Force Research Laboratory.
References
Appendix A Proofs
In this Appendix, we prove Proposition 1. For completeness, note that
We omit the bias terms in each layer as it does not affect the statement.
We will show that and .
(1) : for any , we have , where for all . Hence the linear function is independent of . Therefore,
is also independent of , which means .
Appendix B Experiment details
Fast greedy equivalence search https://github.com/bd2kccd/py-causal (Ramsey et al., 2017) is based on greedy search and assumes linear dependency between variables.
Greedy equivalence search with generalized scores https://github.com/Biwei-Huang/Generalized-Score-Functions-for-Causal-Discovery/ (Huang et al., 2018) is also based on greedy search, but uses generalized scores without assuming a particular model class.
DAG-GNN https://github.com/fishmoon1234/DAG-GNN (Yu et al., 2019) learns a (noisy) nonlinear transformation of a linear SEM using neural networks.
NOTEARS https://github.com/xunzheng/notears (Zheng et al., 2018) learns a linear SEM using continuous optimization.
Causal additive model https://cran.r-project.org/package=CAM (Bühlmann et al., 2014) learns an additive SEM by leveraging efficient nonparametric regression techniques and greedy search over edges.
For all experiments, default parameter settings are used, except for where both preliminary neighborhood selection and pruning are applied.
Simulation
Given the graph , we simulate the SEM for all in the topological order induced by . We consider the following instances of :
Additive GP: , where each is a draw from Gaussian process with RBF kernel with length-scale one.
Index model: , where , , , and each is drawn uniformly from range .
MLP: is a randomly initialized MLP with one hidden layer of size 100 and sigmoid activation.
GP: is a draw from Gaussian process with RBF kernel with length-scale one.
In all settings, is i.i.d. standard Gaussian noise.
Appendix C Additional results
We show {SHD, FDR, TPR, FPR} results on all {ER1, ER2, ER4, SF1, SF2, SF4} graphs in Figure 4, 5, 6, 7 respectively. Similarly, see Figure 8, 9, 10, 11 for full comparison with . As in Figure 1, each row is a random graph model, each column is a type of SEM. Overall has low FDR/FPR and high TPR, and same for on additive GP. Also observe that in most settings has low FDR as well as low TPR, which is a consequence of only predicting a small number of edges.
Complexity and runtime
Recall that numerical evaluation of matrix exponential involves solving linear systems, hence the time complexity is typically for a dense matrix. Taking with one hidden layer of units as an example, it takes time to evaluate the objective and the gradient. If , this is comparable to the linear case , except for the inevitable extra cost from using a nonlinear function. This highlights the benefit of Proposition 1: the acyclicity constraint almost comes for free. Furthermore, we used a quasi-Newton method to reduce the number of calls to evaluate the gradient, which involves computing the matrix exponential. Table 1 contains runtime comparison of different algorithms on ER2 graph with samples. Recall that the kernel-based approach of comes with a computational complexity, whereas and has dependency on . This can be confirmed from the table, which shows has a significantly longer runtime.
Comments on hyperparameter tuning
The experiments presented in this paper were conducted under a fixed (and therefore suboptimal) value of and weight threshold across all graph types, sparsity levels, and SEM types, despite the fact that each configuration may prefer different regularization strengths. Indeed, we observe substantially improved performance by choosing different values of hyperparameters in some settings. As our focus is not on attaining the best possible accuracy in all settings by carefully tuning the hyperparameters, we omit these results in the main text and only include here as a supplement. For instance, for ER4 graph with variables and samples, when the SEM is additive GP and MLP, setting and threshold = 0.5 gives results summarized in Table 2.