Multi-task Learning of Order-Consistent Causal Graphs

1 Introduction . 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 can not 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 . 2 Background . 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 ) , can not be identified from the equivalent class : { G̃ ( π ) : π ∈ Sp } ( 6 ) based on the distribution of X , even if infinitely many samples are observed . 3 Joint estimation of multiple DAGs . 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 . 3.1 Assumption . 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 .

Summary: 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.