DCDILP: a distributed learning method for large-scale causal structure learning

The divide-and-conquer strategy is used by Gu and Zhou (2020) in their method named PEF (Partition, Estimation and Fusion), which consists of three phases: (i) partitioning the set of all variables into clusters, (ii) estimating causal structures cluster by cluster; and (iii) finally producing a fusion of all local learning results. Zhang et al. (2020) proposed a divide-and-conquer approach based on constraint-based methods. In earlier works of Gao et al. (2017), a strategy similar to divide-and-conquer is used, which consists of first Markov blanket learning for each variable and then a procedure for learning and fusion in a from-local-to-global manner.

A different but related approach is the recursive methods (Mokhtarian et al., 2021; Rahman et al., 2021). Mokhtarian et al. (2021) proposed the recursive variable elimination algorithm named MARVEL that involves operations similar to the divide step based on Markov boundary discovery. To obtain Markov blanket information for MARVEL as well as our proposed method (Phase-1), many different algorithms can be used including Grow-Shrink (GS) Margaritis and Thrun (1999), IAMB Tsamardinos et al. (2003), precision matrix-based methods (e.g., Loh and Bühlmann (2014)), and more recent methods by Wu et al. (2020, 2023) such as KMB (Kernel MB learning). The choice of which algorithm to use for computing the Markov boundaries should be made according to the data and application.

Jaakkola et al. (2010) tackle the problem of learning Bayesian network structures using linear programming relaxations. Cussens et al. (2017); Cussens (2023) introduce integer programming and linear programming methods for Bayesian network learning, shedding light on polytopes and facets regarding the search of DAGs.

The above framework, as detailed in Algorithm 1, can be realized in different ways depending on how Phase-1 and Phase-2 are carried out. In the present work, we consider two different algorithms—GES (Chickering, 2002b) and DAGMA (Bello et al., 2022)—for learning the preliminary subgraphs $A^{(i)}$ in Phase-2 (line 3). For Phase-1, a discussion about our choices is given in Section 3.3.

Each subproblem in line 3 takes as input an $n\times|\mbox{\bf S}_{i}|$ submatrix of the whole data matrix $\mathcal{X}$. Subsequently, the output of Phase-2, $\widehat{B}^{(i)}$ (line 4) consists of only the direct causes and effects of $X_{i}$. Note that the estimation of $\widehat{B}^{(i)}$ for each $X_{i}$ is mutually independent given the Markov blankets. Therefore, their computations can be distributed to a number of different CPUs in parallel.

Phase-3 is concerned with building a global causal graph from all partial relations in $\widehat{B}^{(i)}$ constructed in Phase-2.

In contrast to the naive merge $\widehat{B}$ defined in (1), we propose a reconciliation process that combines the edges in matrices $\widehat{B}^{(i)}$ in a selective manner. Since the global causal discovery problem is divided into the separate subproblems according to the Markov blankets, the local solutions $\widehat{B}^{(i)}$ from Phase-2 generally have overlapping answers regarding the causal relations. In a general sense, we refer to the overlapping answers (of two graphs) as merge conflicts defined as follows.

The proof is given in Appendix A. An illustration about the above proposition is shown in Figure 2.

Any pair of matrices $\widehat{B}^{(i)}$ an $\widehat{B}^{(j)}$ by Phase-2 can only have the above three types of conflicts because their edge sets intersect at most on the two directed edges between $X_{i}$ and $X_{j}$; indeed, many more conflicts can otherwise be observed among $\widehat{A}^{(i)}$ and $\widehat{A}^{(j)}$ (in line 3 of Algorithm 1).

Considering all conflicts among the output matrices of Phase-2, DCDILP delegates their resolution to an integer linear programming problem (Wolsey, 2020) by formulating all constraints on the sought solution as follows.

This problem involves: (i) binary variables noted $B_{ij}$ and $S_{ij}$ for all pairs $(i,j)\in[d]\times[d]$ such that $X_{i}\in\mbox{\bf MB}(X_{j})$; (ii) binary variables noted $V_{ijk}$ for all triples $(i,j,k)\in[d]\times[d]\times[d]$ such that $X_{i}$ (respectively, $X_{j}$ and $X_{k}$) belongs to both Markov blankets $\mbox{\bf MB}(X_{j})$ and $\mbox{\bf MB}(X_{k})$ (respectively, $\mbox{\bf MB}(X_{i})$ and $\mbox{\bf MB}(X_{k})$; or $\mbox{\bf MB}(X_{i})$ and $\mbox{\bf MB}(X_{j})$):

The constraints on the above variables express the fact that the sought solution denoted $B$ must be consistent with the given Markov blankets while the objective function of the problem is defined as the similarity of $B$ with the naive concatenation of matrices $\widehat{B}^{(i)}$, namely $\widehat{B}=\sum_{i=1}^{d}\widehat{B}^{(i)}$:

More precisely, the constraints of the above ILP are motivated by the following reasons:

• •

  Sparsity. The constraints (6)–(8) enable us to discard all pairs in $B,S$ or triplets in $V$ that are not involved with the given Markov blanket information.

• •

  2-cycle exclusiveness. For any pair $(i,j),i\neq j$ and $X_{j}\in\mbox{\bf MB}(X_{i})$ (and vice-versa), the first constraint in (9), $B_{ij}+B_{ji}\leq 1$, excludes the 2-cycle between $X_{i}$ and $X_{j}$ given that the entries of $B$ are binary variables.

• •

  Markov blanket membership. The second constraint in (9), $B_{ij}+B_{ji}+S_{ij}\geq 1$, dictates that, when $X_{j}\in\mbox{\bf MB}(X_{i})$ (and vice-versa), they must be either directly causally related or spouses.

• •

  V-structures. For any $k\in[d]$, $\widehat{B}_{ik}\neq 0$ and $\widehat{B}_{jk}\neq 0$, there is a chance that $(X_{i},X_{j},X_{k})$ form a v-structure. Therefore $V_{ijk}$ is encoded (or created) as a binary variable. The constraints in (10) and (11) are necessary conditions for $(V_{ijk},B_{ik},B_{jk},S_{ij})$ to be consistent with the Markov blanket information.

In a classical setting of Gaussian graphical models the sparsity pattern of the inverse covariance matrix of X encodes conditional independence relations between the variables. Consider the usual covariance matrix $\Sigma=\operatorname*{cov}(\mbox{\bf X})$. It is a well-known consequence of the Hammersley–Clifford theorem that the entries of the precision matrix $\Theta=\Sigma^{-1}$ correspond to rescaled conditional correlations Loh and Wainwright (2012). In DCDILP, the empirical inverse covariance estimator is used to achieve Phase-1, that is, $\mbox{\bf MB}(X_{i})\cup\{i\}$ is set as the support of $\Theta_{i,:}$; details are given in Section B.1.

Apart from Phase-1, the computational cost of DCDILP is dominated by Phase-2, which is distributed on a number of parallel workers. Note that the wall time of DCDILP depends mostly on the maximal running time among the parallel workers. Therefore, the total running time of DCDILP is in principle dominated by the running time of the chosen causal discovery routine on the largest Markov blanket during Phase-2. The computational time of the three phases is shown empirically under different settings in Section 4.

The observational data are generated from linear SEMs following the usual settings of Zheng et al. (2018), where the causal structure $B$ is drawn from random DAGs with the Erdős–Rényi (ER) model. The coefficients (edge weights) of $B$ are drawn from the uniform distribution $\text{Unif}([-2,-0.5]\cup[0.5,2])$.

The baselines of the causal discovery benchmarks are GES Chickering (2002b) and DAGMA Bello et al. (2022). More precisely, the GES implementation used in the benchmark—labeled as GES (pcalg)—is from the R package pcalg (https://cran.r-project.org/web/packages/pcalg/index.html), which is a highly efficient implementation.

In these benchmarks, the proposed DCDILP method is represented by DCDILP-GES and DCDILP-DMA, which refers to the specific implementations of DCDILP using GES and DAGMA, respectively, during Phase 2 (line 3). The estimated graphs are evaluated by usual metrics (SHD, TPR, FDR and FPR) for DAG learning (details in Section C.1). The other algorithms are run on one CPU of Intel(R) Xeon(R) Gold 5120 14 cores @ 2.2GHz. The computations in Phase 2 of DCDILP are distributed on a maximal of 300 CPU cores.

The DCDILP method is tested in two scenarios. One scenario is to assess the proposed causal learning methodology irrespective of the statistical performance of Markov blanket discovery, in which case DCDILP (MB^∗) will be used as the label, meaning that the Markov blankets are the ground truth ones. The other scenario is causal discovery, in which case DCDILP takes observational data as input and proceeds by following exactly the framework specified in Algorithm 1.

The implementation of DCDILP is made available at https://github.com/shuyu-d/dcdilp-exp.

In this experiment, DCDILP-GES represents the algorithm of DCDILP that uses GES during its Phase 2 (Algorithm 1, line 3), and is evaluated in causal discovery tasks in comparison with GES. The number $d$ of nodes varies in $\{$50, 100, 200, 400, 800, 1000, 1600$\}$ and the number of samples of X varies as $n=50d$. The choice of the sample sizes is discussed in Section 4.3. The computations in Phase 2 of DCDILP-GES is distributed on $N=\min(2d,300)$ CPU cores. The results are shown in Figure 3, Figure 4 and in Appendix C in the supplementary material.

Figure 3 presents the learning scores of the two algorithms in (TPR, FDR, SHD), depending on the problem dimension $d$. The results in this figure shows that DCDILP-GES outperforms GES in all three learning scores for under all problem dimensions.

It is worth noting that DCDILP-GES uses partially the subgraphs obtained by GES during the local learning subproblems of Phase 2. It is therefore surprising to some extent that DCDILP-GES achieves significant gains over GES in the DAG learning accuracy. We believe that such gains are due to the following reasons: (i) the causal learning subproblems in Phase 2 of DCDILP enjoys a higher effective sample size, in the sense that the ratio $n/|\mbox{\bf S}_{i}|$—for each subproblem on $\mbox{\bf S}_{i}$—is effectively higher than $n/d$ (for the global problem); (ii) the reconciliation by the ILP method of DCDILP-GES is a selection process that removes spurious directions (in terms of DAG learning) from the subgraphs given by GES during Phase 2, while GES is a method for producing CPDAGs.

Figure 4 presents the running time (wall time) of the two algorithms depending on the problem dimension $d$. The results in this figure shows that (i) GES (pcalg) enjoys a highly competitive time efficiency on small- to mid-side problems (for $d$ ranging from $50$ to around $700$), which is superior or comparable to DCDILP-GES; (ii) however, the running time of DCDILP-GES grows more slowly than GES with increasing $d$. The trend in this figure shows that when $d$ surpasses $800$ (on ER1 data) or $700$ (on ER2 data), DCDILP-GES starts to gain speedups over GES with an increasing rate. When $d=1600$, for example, DCDILP-GES achieves around $4\times$ and $5\times$ speedups over GES on ER1 data and ER2 data respectively.

In this experiment, DCDILP-DMA represents the algorithm of DCDILP that uses DAGMA during Phase 2 (Algorithm 1, line 3), and is evaluated in comparison with DAGMA. The number $d$ of nodes varies in $\{$50, 100, 200, 400, 800, 1000, 1600$\}$ and the number of samples of X varies as $n=50d$.

The computation in Phase 2 of DCDILP-DAGMA is distributed on $N=\min(2d,300)$ CPU cores. The results are shown in Figure 5 and Figure 6.

The results in Figure 5 and Figure 6 show that DCDILP-DMA achieves significant gains over DAGMA in running time—under almost all problem dimensions—within a reasonable compromise in learning accuracy. More precisely, on ER1 data, the median TPRs of both DCDILP-DMA and DAGMA close to 100% and the median FDRs of DCDILP-DMA are under 10% under all dimensions $d$, while DAGMA has FDRs that are close to zero (indicating almost exact recovery of the underlying DAGs). On ER2 data, DCDILP-DMA has median TPRs slightly above 90% while DAGMA stays close to 100%; and DCDILP-DMA has median FDRs between 10% and 20%, which is reasonably low despite being higher than DAGMA.

On the other hand, the time efficiency gains of DCDILP-DMA over DAGMA are even more significant than the case with GES (pcalg). Despite that DCDILP-DMA uses DAGMA for the local learning subproblems in Phase 2, its overall wall time becomes $100\times$ lower than DAGMA when $d$ surpasses $400$. As discussed in Section 3.3, such speedups are largely due to the parallelization of the local learning subproblems in Phase 2. Despite that the parallelization of subproblem tasks in the present experiment actually encounters congestion (when $d$ grows) given the predefined limitation on the maximal number of CPU cores, the speedups obtained by DCDILP are already considerable for problems with more than 1000 variables.

In the same settings as the above experiments, we showcase the effects of reconciliation process via the ILP-based method (formulated in Section 3.2, (5)–(11)) in comparison with the naive merge result $\widehat{B}$ (1).

Figure 7 shows their comparisons in the representative case with ER2 data. The scores in this figure shows that for both DCDILP-GES and DCDILP-DMA, the reconciliation process by the ILP method achieves significant gains in learning accuracy than the naive merge result $\widehat{B}$ (1) under all problem dimensions. More precisely, the median TPRs of DCDILP are slightly below those of the naive merge while staying around 90% or above; more interestingly, the median FDRs of DCDILP are significantly reduced than those of the naive merge. These comparisons provide solid validation to the effects of the ILP method.

In the same settings as the above experiments, we evaluate learning accuracy of DCDILP under different sample sizes. The sample size of observational data varies between $5d$ and $50d$ for causal discovery tasks of $d=100$ variables.

The comparative results in Figure 8 give an empirical insight in the sample requirement of DCDILP-GES under the given choices for Phase 1 (MBs by inverse covariance estimation) and Phase 2 (using GES). The learning accuracy of DCDILP-GES improves with increasing sample size. For Gaussian SEMs, the accuracy attains a level comparable to the case with oracle MBs for $n$ reaches $20d$; and for Gumbel-noised SEMs, the accuracy of DCDILP-GES becomes comparable to DCDILP-GES (MB*) when $n$ reaches $40d$. Future work should include more sophisticated or more sample-efficient methods for Markov blanket discovery.