Multi-task Learning of Order-Consistent Causal Graphs
Year: 2021

Abstract:
We consider the problem of discoveringK related Gaussian directed acyclic graphs (DAGs), where the involved graph structures share a consistent causal order and sparse unions of supports. Under the multi-task learning setting, we propose a l1/l2regularized maximum likelihood estimator (MLE) for learning K linear structural equation models. We theoretically show that the joint estimator, by leveraging data across related tasks, can achieve a better sample complexity for recovering the causal order (or topological order) than separate estimations. Moreover, the joint estimator is able to recover non-identifiable DAGs, by estimating them together with some identifiable DAGs. Lastly, our analysis also shows the consistency of union support recovery of the structures. To allow practical implementation, we design a continuous optimization problem whose optimizer is the same as the joint estimator and can be approximated efficiently by an iterative algorithm. We validate the theoretical analysis and the effectiveness of the joint estimator in experiments.

heading: 1 Introduction
text: Estimating causal effects among a set of random variables is of fundamental importance in many disciplines such as genomics, epidemiology, health care and finance [1, 2, 3, 4, 5, 6]. Therefore, designing and understanding methods for causal discovery is of great interests in machine learning.
Causal discovery from finite observable data is often formulated as a directed acyclic graph (DAG) estimation problem in graphical models. A major class of DAG estimation methods are score-based, which search over the space of all DAGs for the best scoring one. However, DAG estimation remains a very challenging problem from both the computational and statistical aspects [7]. On the one hand, the number of possible DAG structures grows super-exponentially in the number of random variables, whereas the number of observational sample size is normally small. On the other hand, some DAGs are non-identifiable from observational data even with infinitely many samples.
Fortunately, very often multiple related DAG structures need to be estimated from data, which allows us to leverage their similarity to improve the estimator. For instance, in bioinformatics, gene expression levels are often measured over patients with different subtypes [8, 9] or under various experimental conditions [10]. In neuroinformatics, fMRI signals are often recorded for multiple subjects for studying the brain connectivity network [11, 12]. In these scenarios, multiple datasets will be collected, and their associated DAGs are likely to share similar characteristics. Intuitively, it may be beneficial to estimate these DAGs jointly.
∗Work done partially during the visit at MBZUAI (Mohamed bin Zayed University of Artificial Intelligence)
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
In this paper, we focus on the analysis of the multi-task DAG estimation problem where the DAG structures can be related by a consistent causal order and a (partially) shared sparsity pattern, but allowed to have different connection strengths and edges, and differently distributed variables. In this setting, we propose a joint estimator for recovering multiple DAGs based on a group norm regularization. We prove that the joint l1/l2-penalized maximum likelihood estimator (MLE) can recover the causal order better than individual estimators.
Intuitively, it is not surprising that joint estimation is beneficial. However, our results provide a quantitative characterization on the improvement in sample complexity and the conditions under which such improvement can hold. We show that:
• For identifiable DAGs, if the shared sparsity pattern (union support) size s is of order O(1) in K where K is the number of tasks (DAGs), then the effective sample size for order recovery will be nK where n is the sample size in each problem. Furthermore, as long as s is of order o( √ K) in K,
the joint estimator with group norm regularization leads to an improvement in sample complexity. • A non-identifiable DAG cannot be distinguished by single-task estimators even with indefinitely
many observational data. However, non-identifiable DAGs can be recovered by our joint estimator if they are estimated together with other identifiable DAGs.
Apart from the theoretical guarantee, we design an efficient algorithm for approximating the joint estimator through a formulation of the combinatorial search problem to a continuous programming. This continuous formulation contains a novel design of a learnable masking matrix, which plays an important role in ensuring the acyclicity and shared order for estimations in all tasks. An interesting aspect of our design is that we can learn the masking matrix by differentiable search over a continuous space, but the optimum must be contained in a discrete space of cardinality p! (reads p factorial, where p is the number of random variables).
We conduct a set of synthetic experiments to demonstrates the effectiveness of the algorithm and validates the theoretical results. Furthermore, we apply our algorithm to more realistic single-cell expression RNA sequencing data generated by SERGIO [13] based on real gene regulatory networks.
The remainder of the paper is organized as follows. In Section 2, we introduce the linear structural equation model (SEM) interpretation of Gaussian DAGs and its properties. Section 3 is devoted to the statement of our main results, with some discussion on their consequences and implications. In Section 4, we present the efficient algorithm for approximating the joint estimator. Section 5 summarizes related theoretical and practical works. Experimental validations are provided in Section 6.

heading: 2 Background
text: A substantial body of work has focused on the linear SEM interpretation of Gaussian DAGs [14, 15, 16, 17]. Let X = (X1, · · · , Xp) be a p-dimensional random variable. Then a linear SEM reads
X = G̃>0 X +W, W ∼ N (0,Ω0), (1)
where Ω0 is a p×p positive diagonal matrix which indicates the variances of the noiseW . G̃0 ∈ Rp×p is the connection strength matrix or adjacency matrix. Each nonzero entry G̃0ij represents the direct causal effect of Xi on Xj . This model implies that X is Gaussian, X ∼ N (0,Σ), where
Σ := (I − G̃0)−>Ω0(I − G̃0)−1. (2)
Causal order π0. The nonzero entries of G̃0 defines its causal order (also called topological order), which informs possible “parents” of each variable. A causal order can be represented by a permutation π0 over [p] := (1, 2, · · · , p). We say G̃0 is consistent with π0 if and only if
G̃0ij 6= 0⇒ π0(i) < π0(j). (3)
There could exist more than one permutations that are consistent with a DAG structure G̃0, so we denote the set of permutations that satisfy Eq. 3 by Π0. Once a causal order π0 is identified, the connection strengths G̃0 can be estimated by ordinary least squares regression which is a comparatively easier problem.
Identifiability and equivalent class. However, estimating the true causal order π0 is very challenging, largely due to the existence of the equivalent class described below.
Let Sp be the set of all permutations over [p]. For every π ∈ Sp, there exists a connection strength matrix G̃(π), which is consistent with π, and a diagonal matrix Ω(π) = diag [ σ1(π) 2, · · · , σp(π)2 ] such that the variance Σ in Eq. 2 equals to
Σ = (I − G̃(π))−>Ω(π)(I − G̃(π))−1, (4) and therefore the random variable X in Eq. 1 can be equivalently described by the model [15]:
X = G̃(π)>X +W (π), W (π) ∼ N (0,Ω(π)). (5)
Without further assumption, the true underlying DAG, G̃0 = G̃(π0), cannot be identified from the
equivalent class: {G̃(π) : π ∈ Sp} (6) based on the distribution of X , even if infinitely many samples are observed.

heading: 3 Joint estimation of multiple DAGs
text: In the multi-task setting, we consider K linear SEMs:
X(k) = G̃ (k) 0 >X(k) +W (k), for k = 1, · · · ,K.
The superscript notation (k) indicates the k-th model. As mentioned in Sec 2, each model defines a random variable X(k) ∼ N (0,Σ(k)) with Σ(k) = (1− G̃(k)0 )−>Ω (k) 0 (1− G̃ (k) 0 ) −1.

heading: 3.1 Assumption
text: Similarity. In Condition 3.1 and Definition 3.1, we make a few assumptions about the similarity of these models. Condition 3.1 ensures the causal orders in the K DAGs do not have conflicts. Definition 3.1 measures the similarity of the sparsity patterns in the K DAGs. If the sparsity patterns are strongly overlapped, then the size of the support union s will be small. We do not enforce a strict constraint on the support union, but we will see in the later theorem that a smaller s can lead to a better recovery performance. Condition 3.1 (Consistent causal orders). There exists a nonempty set of permutations Π0 ⊆ Sp such that ∀π0 ∈ Π0, it holds for all k ∈ [K] that G̃0(k)ij 6= 0⇒ π0(i) < π0(j). Definition 3.1 (Support union). Recall G̃(π) in Eq. 6. The support union of the K DAGs associated with permutation π is denoted by S(π) := { (i, j) : ∃k ∈ [K] s.t. G̃(k)ij (π) 6= 0 } . The support union of the K true DAGs G̃(k)0 is S0 := S(π0). We further denote s0 := |S0| and s := supπ∈Sp |S(π)|.
Identifiability. To ensure the consistency of the estimator based on least squared loss, we first assume that the DAGs to be recovered are minimum-trace DAGs. Condition 3.2 (Minimum-trace). Recall the equivalent class defined in Eq. 4 to Eq. 4. Assume for all k = 1, · · · ,K, trace(Ω(k)0 ) = minπ∈Sp trace(Ω(π)(k)).
However, the minimum-trace DAG may not be unique without further assumptions, making the true DAG indistinguishable from other minimum-trace DAGs. Therefore, we consider the equal variance condition in Condition 3.3 which ensures the uniqueness of the minimum-trace DAG. In this paper, we assume the first K ′ ≤ K models satisfy this condition, so they are identifiable. We do not make such an assumption on the other K −K ′ models, so the K −K ′ models may not be identifiable. Condition 3.3 (Equal variance). For all k = 1, · · · ,K ′ with K ′ ≤ K, the noise W (k) ∼ N (0,Ω(k)0 ) has equal variance with Ω(k)0 = σ (k) 0 Ip.
3.2 l1/l2-penalized joint estimator
Denote the sample matrix by X(k) whose row vectors are n i.i.d. samples fromN (0,Σ(k)). Based on the task similarity assumptions, we propose the following l1/l2-penalized joint maximum likelihood estimator (MLE) for jointly estimating the connection strength matrices {G̃(k)0 }Kk=1:
π̂, {Ĝ(k)} = arg min π∈Sp,{G(k)∈D(π)} K∑ k=1 1 2n ‖X(k) −X(k)G(k)(π)‖2F + λ‖G(1:K)(π)‖l1/l2 . (7)
Similar to the notation in Eq. 4, G(k)(π) indicates its consistency with π. D(π) denote the space of all DAGs that are consistent with π. It is notable that a single π shared across all K tasks is optimized in Eq. 7, which respects Condition 3.1. The group norm over the set of K matrices is defined as
‖G(1:K)(π)‖l1/l2 := ∑p i=1 ∑p j=1 ‖G (1:K) ij ‖2, where G (1:K) ij := [G (1) ij (π), · · · , G (K) ij (π)].
It will penalize the size of union support in a soft way. When K = 1, this joint estimator will be reduced to the l1-penalized maximum likelihood estimation. Remark 3.1. The optimization in Eq. 7 is used for analysis only. A continuous program with the same optimizer will be discussed in Section 4 for practical implementation.

heading: 3.3 Main result: causal order recovery
text: We start with a few pieces of notations and definitions. Then the theorem statement follows.
Definition 3.2. Let g̃(k)j (π) denote the j-th column of G̃ (k)(π). Let
d := supj∈[p],π∈Sp ∣∣ ∪k∈[K] supp(g̃(k)j (π))∣∣, gmax := supπ∈Sp,(i,j)∈S(π) ‖G̃(1:K)ij (π)‖2/√K.
In Definition 3.2, d is the maximal number of parents (in union) in the DAGs G̃(k)(π), which is also a measure of the sparsity level. gmax is bounded by the maximal entry value in the matrices G̃(k)(π). Condition 3.4 (Bounded spectrum). Assume for all k = 1, · · · ,K, the covariance matrix Σ(k) is positive definite. There exists constants 0 < Λmin ≤ Λmax <∞ such that for all k = 1, · · · ,K,
(a) all eigenvalues of Σ(k) are upper bounded by Λmax; (b) all eigenvalues of Σ(k) are lower bounded by Λmin.
Condition 3.5 (Omega-min). There exists a constant ηw > 0 such that for any permutations π /∈ Π0, 1 pK′ ∑K′ k=1 ∑p j=1(σ (k) j (π) 2 − σ(k)0 2)2 > 1ηw .
Condition 3.5 with K ′ = K = 1 is called ‘omega-min’ condition in [14], so we follows this terminology. In some sense, when ηw is larger, σj(π) with π /∈ Π0 is allowed to deviate less from the true variance σj(π0) = σ0 with π0 ∈ Π0, which will make it more difficult to separate the set Π0 from its complement in a finite sample scenario. Ideally, we should allow ηw to be large, so that the recovery is not only restricted to easy problems.
Now we are ready to present the recovery guarantee for causal order. Theorem 3.1 is a specific statement when the regularization parameter λ follows the classic choice in (group) Lasso problems. A more general statement which allows other λ is given in Appendix B along with the proof. Theorem 3.1 (Causal order recovery). Suppose we solve the joint optimization in Eq. 7 with specified regularization parameter λ = √ p log p n for a set of K problems that satisfy Condition 3.1, 3.4 (a), 3.2, 3.3 and 3.5. If the following conditions are satisfied:
θ(n,K,K ′, p, s) := p
s
√ n
p log p
K ′2
K > κ1ηw, (8)
n ≥ c1 logK + c2(d+ 1) log p, (9) K ≤ κ2p log p, (10)
then the following statements hold true:
(a) With probability at least 1− c3 exp (−κ4(d+ 1) log p)− exp (−c4n), it holds that π̂ ∈ Π0.
(b) If in addition, n satisfies n ≥ κ5d̂(logK + log p) with d̂ := maxj∈[p],k∈[K] ‖ĝ (k) j − g̃ (k) j ‖0,
and Condition 3.4 (b) holds, then with probability at least 1 − c3 exp (−κ4(d+ 1) log p) − exp (−c4n)− exp(− log p− logK),
1 K ∑K k=1 ‖Ĝ(k) − G̃ (k) 0 ‖2F = O ( s0 √ p log p nK ) .
In this statement, c1, c2, c3, c4 are universal constants (i.e., independent of n, p,K, s,Σ(k)), κ1, κ2, κ3, κ4 are constants depending on gmax and Λmax, and κ5 is a constant depending on Λmin.
Discussion on Eq. 8-Eq. 10. (i) In Eq. 8, Theorem 3.1 identifies a sample complexity parameter θ(n,K,K ′, p, s) . Following the terminology in [18], our use of the term “sample complexity” for θ reflects the dominant role it plays in our analysis as the rate at which the sample size much grow in order to obtain a consistent causal order. More precisely, for scalings (n,K,K ′, p, s) such that θ(n,K,K ′, p, s) exceeds a fixed critical threshold κ1ηw, we show that the causal order can be correctly recovered with high probability.
(ii) The additional condition for the sample size n in Eq. 9 is the sample requirement for an ordinary linear regression problem. It is in general much weaker than Eq. 8, unless K grows very large.
(iii) The last condition in Eq. 10 on K could be relaxed if a tighter analysis on the distribution properties of Chi-squared distribution is available. However, it is notable that this restriction on the size of K has already been weaker than many related works on multi-task `1 sparse recovery, which either implicitly treat K as a constant [18, 19] or assume K = o(log p) [20, 21].
Recovering non-identifiable DAGs. A direct consequence of Theorem 3.1 is that as long as the number of identifiable DAGs K ′ is non-zero, the joint estimator can recover the causal order of non-identifiable DAGs with high probability. This is not achievable in separate estimation even with infinitely many samples. Therefore, we show that how the information of identifiable DAGs helps recover the non-identifiable ones.
Effective sample size. As indicated by θ(n,K,K ′, p, s) in Eq. 8, the effective sample size for recovering the correct causal order is nK ′2
K if the support union size s is of orderO(1) in K. To show the improvement in sample complexity, it is more fair to consider the scenario when the DAGs are identifiable, i.e., K ′ = K. In this case, it is clear that the parameter θ(n,K,K, p, s) indicates a lower sample complexity relative to separate estimation as long as s is of order o( √ K).
Separate estimation. Consider the special case of a single task estimation with K = K ′ = 1, in which the joint estimator reduces to `1-penalized MLE. We discuss how our result can recover previously known results for single DAG recovery. Unfortunately, existing analyses were conducted under different frameworks with different conditions. [14] and [22] are the most comparable ones since they were based on the same omega-min condition (i.e., Condition 3.5), but they chose a smaller regularization parameter λ. In our proof, the sample complexity parameter is derived from θ(n,K,K ′, p, s) = pK ′/ ( sλ √ K ) and Eq. 8 is pK ′/ ( sλ √ K ) > κ1ηw. When K = K ′ = 1, this
condition matches what is identified in [14] and [22] for recovering the order of a single DAG.
Error of Ĝ(k). To compare the estimation of Ĝ(k) to the true DAG G̃(k)0 , Theorem 3.1 (b) says the averaged error in F-norm goes to zero when nK →∞. It decreases in K as long as s = o(K). To summarize, Theorem 3.1 analyzes the causal order consistency for the joint estimator in Eq. 7. Order recovery is the most challenging component in DAG estimation. After π̂ ∈ Π0 has been identified, the DAG estimation becomes p linear regression problems that can be solved separately. Theorem 3.1 (b) only shows the estimation error of connection matrices in F-norm. To characterize the structure error, additional conditions are required.

heading: 3.4 Support union recovery
text: Theorem 3.1 has shown that π̂ = π0 ∈ Π0 holds with high probability. Consequently, the support recovery analysis in this section is conditioned on this event. In fact, given the true order π0, what remains is a set of p separate l1/l2-penalized group Lasso problems, in each of which the order π0 plays a role of constraining the support set by the set Sj(π0) := {i : π0(i) < j}. However, we need to solve p such problems simultaneously where p is large. A careful analysis is required, and directly combining existing results will not give a high recovery probability.
In the following, we impose a set of conditions and definitions, which are standard in many l1 sparse recovery analyses [18, 19], after which theorem statement follows. See Appendix C for the proof.
Definition 3.3. The union support of j-th columns of {G̃(k)0 }k∈[K] is denoted by RSj := {i ∈ [p] : ∃k ∈ [K] s.t. G̃(k)0ij 6= 0}. The maximal cardinality is rmax := supj∈[p] ∣∣RSj∣∣.
Definition 3.4. ρu := supk∈[K],j∈[p],S=RSj maxi∈Sc ( Σ (k) ScSc|S ) ii is the maximal diagonal entry of the conditional covariance matrices, where Σ(k)ScSc|S := Σ (k) ScSc − Σ (k) ScS(Σ (k) SS) −1Σ (k) SSc .
Definition 3.5. gmin := inf(i,j)∈S0 ‖G̃ (1:K) 0ij ‖2/
√ K represents the signal strength.
Condition 3.6 (Irrepresentable condition). There exists a fixed parameter γ ∈ (0, 1] such that sup
j∈[p],S=RS(g̃(1:K)0j ) ‖A(S)‖∞ ≤ 1− γ, where A(S)ij := supk∈[K] ∣∣(Σ(k)ScS(Σ(k)SS)−1)ij∣∣. Theorem 3.2 (Union support recovery). Assume on the subset of probability space where {π̂ ∈ Π0} holds, and assume Condition 3.6. Assume the following conditions are satisfied
n ≥ κ6rmax log p, (11) K ≤ c7 log p, (12)√
8Λmax log p
Λminn +
2
Λmin
√ p log p
n rmax K = o(gmin), (13)
where κ6 is a constant depending on γ,Λmin, ρu, σmax, and c7 is some universal constant. Then w.p. at least 1− rmax exp (−c5K log p)− exp (−c6 log p)− c8K exp (−c9(n− rmax − log p)), the support union of Ĝ(1:K) is the same as that of G̃(1:K)0 , and that ∥∥Ĝ(1:K) − G̃(1:K)∥∥ l∞/l2 / √ K = o(gmin).
Discussion on Eq. 11 and Eq. 12. (i) Eq. 11 poses a sample size condition. The value rmax is the sparsity overlap defined in Definition 3.3 (i). It takes value in the interval [d,min{s0, p, dK}], depending on the similarity in sparsity pattern.
(ii) The restriction on K in Eq. 12 plays a similar role as Eq. 10 in Theorem 3.1. This a stronger restriction, but also guarantees the stronger result of support recovery. Existing analyses on l1/l2penalized group Lasso were not able to relax this constraint, neither, so some of them treated K as a constant in the analysis [18, 19]. Recall that in Theorem 3.1, we were able to allow K = O(p log p). Technically, this was achieved because in the proof of Theorem 3.1, we avoid analyzing the general recovery of group Lasso, but only its null-consistency (i.e., the special case of true structures having zero support), where tighter bound can be derived and it is sufficient for order recovery.
Benefit of joint estimation. Eq. 13 plays a similar role as Eq. 8 in Theorem 3.1. It specifies a rate at which the sample size must grow for successful union support recovery. As long as rmax is of order o( √ K), K will effectively reduce the second term in Eq. 8. Apart from that, the recover probability specified in Theorem 3.2 grows in K.

heading: 4 Algorithm
text: Solving the optimization in Eq. 7 by searching over all permutations π ∈ Sp is intractable due to the large combinatorial space. Inspired by the smooth characterization of acyclic graph [16], we propose a continuous optimization problem, whose optimizer is the same as the estimator in Eq. 7. Furthermore, we will design an efficient iterative algorithm to approximate the solution.

heading: 4.1 Continuous program
text: We convert Eq. 7 to the following constrained continuous program
min T∈Rp×p G(1),··· ,G(K)∈Rp×p K∑ k=1 1 2n ∥∥∥X(k) −X(k)G(k)∥∥∥2 F + λ‖G(1:K)‖l1/l2 + ρ‖1p×p − T‖ 2 F (14)
subject to h(T ) := trace(eT◦T )− p = 0, (15)
whereG(k) := G(k)◦T is element-wise multiplication betweenG(k) and T , and 1p×p is a p×pmatrix with entries equal to one. Eq. 15 is a smooth ‘DAG-ness’ constraint proposed by NOTEARS [16], which ensures T is acyclic. One can also use h(T ) := trace((I + T ◦ T )p)− p proposed in [23]. We would like to highlight the novel and interesting design of the matrix T in Eq. 14. What makes Eq. 7 difficult to solve is the requirement that {G(k)} must be DAGs and share the same order. A straightforward idea is to apply the smooth acyclic constraint to every G(k), but it is not clear how to enforce their consistent topological order. Our formulation realizes this by a single matrix T .
Algorithm 1: Joint Estimation Algorithm Hyperparameters :ρ, α, λ, t, δ Initialize G(1:K), T randomly; for itr = 1, · · · ,M do
for itr′ = 1, · · · ,M ′ do [G(1:K), T ]← GradOptStep ( f ;G(1:K), T, β ) ; . Gradient-based update on f
∀i, j ∈ [p], G(1:K)ij ← G
(1:K) ij
‖G(1:K)ij ‖2 max
{ 0, ‖G(1:K)ij ‖2 − tλ|Tij | } ; . Proximal step
∀i, j ∈ [p], Tij ← sign(Tij) max { 0, |Tij | − tλ‖G(1:K)ij ‖2 } ; . Proximal step
β ← β + τh(T ); . Dual ascent α← α · (1 + δ); . Typical rule [23]
π = (1, 2, 3, 4)
⇔
T =  0 0 0 01 0 0 01 1 0 0 1 1 1 0  π = (4, 2, 1, 3) ⇔ T =  0 0 1 01 0 1 00 0 0 0 1 1 1 0  To better understand the design rationale of T , recall in Eq. 3 that a matrix G is a DAG of order π if and only if its support set is in {(i, j) : π(i) < π(j)}. The matrix T plays a role of restricting the support set of G(k) by masking its entries. Two examples are shown above. However, unlike the learning of masks in other papers which allows T to have any combinations of nonzero entries, here we need T to exactly represent the support set for each π. That is T is from a space Tp with p! elements: Tp := {T ∈ {0, 1}p×p : Tij = 1 ⇔ π(i) < π(j)}. Now a key question arises: How to perform a continuous and differentiable search over Tp? The following finding motivates our design:
T ∈ Tp ⇐⇒ T ∈ arg minT∈Rp×p { ‖1p×p − T‖2F subject to h(T ) = 0 } .
In other words, T is a continuous projection of 1p×p to the space of DAGs. We can then optimize the mask T in the continuous space Rp×p but the optimal solution must be an element in the discrete space Tp. This observation also naturally leads to the design of Eq. 14. Finally, we want to emphasize that it is important for the optimal T to have binary entries. Without this property, any nonzero value c can scale the (G,T ) pair to give an equivalent masked DAG, i.e., G ◦ T = (cG) ◦ ( 1cT ). This scaling equivalence will make the optimization hard to solve in practice. Proofs for the above arguments and the equivalence between Eq. 7 and Eq. 14 are in Appendix F.

heading: 4.2 Iterative algorithm
text: We derive an efficient iterative algorithm using the Lagrangian method with quadratic penalty, which converts Eq. 14 to an unconstrained problem:
min T,G(1),··· ,G(K)∈Rp×p max β≥0
L(G(1:K), T ; β) := f(G(1:K), T ; β) + λ‖G(1:K)‖l1/l2 ,
where f(G(1:K), T ; β) := K∑ k=1 1 2n ∥∥∥X(k) −X(k)G(k)∥∥∥2 F + ρ‖1p×p − T‖2F + βh(T ) + αh(T )2,
β is dual variable, α is the coefficient for quadratic penalty, and f is the smooth term in the objective.
We can solve this min-max problem by alternating primal updates on (G(1:K), T ) and dual updates on β. Due to the non-smoothness of group norm, the primal update is based on proximal-gradient method, where the proximal-operator with respect to ‖ · ‖l1/l2 has a closed form:[
arg minZ(1:K)∈RK×p×p 1 2 ∑K k=1 ‖Z(k) −X(k)‖2F + c‖Z(1:K)‖l1/l2 ](1:K) ij
= X
(1:K) ij
‖X(1:K)ij ‖2 max{0, ‖X(1:K)ij ‖2 − c},
which is a group-wise soft-threshold. Since G(k) and T are multiplied together element-wisely inside the group norm, the proximal operator will be applied to both of them separately. Together with the
dual update for β, β ← β + τh(T ) with τ as the step size, we summarize the overall algorithm for solving Eq. 14 in Algorithm 1.

heading: 5 Related work
text: Single DAG estimation. Unlike the large literature of research on undirected graphical models [24, 25, 26, 27], statistical guarantees for score-based DAG estimator have been available only in recent years. [14, 15, 17] have shown the DAG estimation consistency in high-dimensions, but they do not consider joint estimation. Nevertheless, some techniques in [14, 15] are useful for our derivations.
Multi-task learning. (i) Undirected graph estimation. There have been extensive studies on the joint estimation of multiple undirected Gaussian graphical models [28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39]. (ii) DAG estimation. In contrast, not much theoretical work has been done for joint DAG learning. A few pieces of recent works addressed certain statistical aspects of multi-task DAG estimation [9, 20, 21], but [20, 21] tackle the fewer task regime with K = o(log p), and [9] assumes the causal order is given. Another related work that we notice after the paper submission is [40], which also assumes the DAGs have a consistent causal order, but it focuses on estimating the difference between two DAGs. (iii) Linear regression. Multi-task linear regression is also a related topic [41, 42, 43, 44, 45, 46, 47], because the techniques for analyzing group Lasso are used in our analysis [18, 19].
Practical Algorithms. Works on practical algorithm design for efficiently solving the score-based optimization are actively conducted [48, 49, 50, 51]. Our algorithm is most related to recent methods exploiting a smooth characterization of acyclicity, including NOTEARS [16] and several subsequent works [23, 52, 53, 54], but they only apply for single-task DAG estimation. Although algorithms for the multi-task counterpart were proposed a decade ago [11, 55, 56, 57], none of them leverage recent advances in characterizing DAGs and providing theoretically guarantees.

heading: 6 Experiments

heading: 6.1 Synthetic data
text: The set of experiments is designed to reveal the effective sample size predicted by Theorem 3.1, and demonstrate the effectiveness of the proposed algorithm. In the simulations, we randomly sample a causal order π and a union support set S0. Then we randomly generate multiple DAGs that follow the order π and have edges contained in the set S0. For each DAG, we construct a linear SEM with standard Gaussian noise, and sample n data points from it. On tasks with different combinations of (p, n, s,K), we exam the behavior of the joint estimator, estimated by Algorithm 1, on 64 tasks and report the statistics in the following for evaluation. In this experiment, we take K ′ = K so that all the DAGs are identifiable. This simpler case will make it easier to verify the proposed algorithm and the rates in the theorem.
Fig 2 shows the success probability in the form of Heat Maps, where the rows indicate an increase in per task sample size n, and the columns indicate an increase in the number of tasks K. The results show that the increases in these two quantities have similar effect to the success probability.
Effectiveness of K in structure recovery. In this experiment, we aim at verifying the effectiveness of the joint estimation algorithm for recovering the structures. For brevity, we only report the numbers for false discovery rate (FDR) and true positive rate (TPR) of edges in Fig 3, but figures for additional metrics can be found in Appendix G. In Fig 3, we can observe consistent improvements when increasing the number of tasks K. When per task sample size is small, this improvement reveals to be more obvious.
Comparison with other joint estimator. In this experiment, we compare our method MultiDAG with JointGES [21] on models with p = 32 and p = 64. Results in Table 1 show that two algorithms have similar performance when K = 1. However, when K increases, our method returns consistently better structures in terms of SHD.

heading: 6.2 Recovery of gene regulatory network
text: We investigate how our joint estimator works on more realistic models, by conducting a set of experiments on realistic gene expression data generated by SERGIO [13], which models the additive
effect of cooperative transcription factor binding across multiple gene regulators in parallel with protein degradation and noisy expression.
We conduct this experiment on the E. coli 100 network, which includes 100 known genes and 137 known regulatory interactions. To evaluate our algorithm, we generate multiple networks by rearranging and re-weighting 5 edges at random in this network without violating the topological order. We simulate the gene expression from each network using SERGIO with a universal non-cooperative hill coefficient of 0.05, which works well with our linear recovery algorithm.
Fig 4 provides a visual comparison of the recovered structures. It can be seem from the true network that there are a few key transcription factors that highlight several rows in the figure. These transcription factors are better identified by the two structures in (b) and (c), but not that clear in (d). Combining this observation with the more quantitative results in Table 2, we see that the combination (K = 10, n = 100) achieves comparable performance to K = 1 with the same total number of samples, and outperforms the single task estimation with n = 100.

heading: 7 Conclusion and discussion
text: In this paper, we have analyzed the behavior of l1/l2-penalized joint MLE for multiple DAG estimation tasks. Our main result is to show that its performance in recovering the causal order is governed by the sample complexity parameter θ(n,K,K ′, p, s) in Eq. 8. Besides, we have proposed an efficient algorithm for approximating the joint estimator via formulating a novel continuous programming, and demonstrated its effectiveness experimentally. The current work applies to DAGs that have certain similarity in sparsity pattern. It will be interesting to consider whether the joint estimation without the group-norm (and without the union support assumption) can also lead to similar improvement in causal order recovery.

heading: Acknowledgement
text: Xinshi Chen is supported by the Google PhD Fellowship. This works is done partially during a visit at MBZUAI. We are grateful for the computing resources provided by the Partnership for an Advanced Computing Environment (PACE) at the Georgia Institute of Technology, Atlanta. We are thankful for PACE Research Scientist Fang (Cherry) Liu’s excellent HPC consulting.

Review summaries:
Authors study the problem of learning the causal ordering and union support of DAGs from K linear SEMs. Authors propose a regularized joint estimator and show that this joint estimation enjoys a better sample complexity than single separate estimations. Moreover, guarantees for causal order and union support recovery are provided. For experiments, authors propose a continuous optimization problem that outputs the same causal order and union support than the joint estimator.

The paper considers structure discovery of multiple but related DAGs in the sense that they share a common causal order. Authors propose an l1/l2-regularized MLE for multiple linear Gaussian models and design a continuous optimization method to obtain the estimates. Sample complexity analysis are also provided regarding causal order and union support recovery, which is interesting and novel. Experiments also validate the theoretical finding and the effectiveness of joint estimation. The paper is generally well-written, though some details seem missing in the experiment part.

The paper proposes a multi-task DAG learning under equal variance and Gaussian assumptions. Consistency results are provided and new algorithm is proposed to solve the formulation. Empirically, results are shown to be consistent with the theoretical analysis.