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