Large-Scale Differentiable Causal Discovery of Factor Graphs

A common theme in causal inference is learning causal relationships between observed variables, also known as causal discovery. This is usually a daunting task, given the large number of candidate causal graphs and the combinatorial nature of the search space. Perhaps for this reason, most research has so far focused on relatively small causal graphs, with up to hundreds of nodes. However, recent advances in fields like biology enable generating experimental data sets with thousands of interventions followed by rich profiling of thousands of variables, raising the opportunity and urgent need for large causal graph models. Here, we introduce the notion of factor directed acyclic graphs (f -DAGs) as a way to restrict the search space to non-linear low-rank causal interaction models. Combining this novel structural assumption with recent advances that bridge the gap between causal discovery and continuous optimization, we achieve causal discovery on thousands of variables. Additionally, as a model for the impact of statistical noise on this estimation procedure, we study a model of edge perturbations of the f -DAG skeleton based on random graphs and quantify the effect of such perturbations on the f -DAG rank. This theoretical analysis suggests that the set of candidate f -DAGs is much smaller than the whole DAG space and thus may be more suitable as a search space in the high-dimensional regime where the underlying skeleton is hard to assess. We propose Differentiable Causal Discovery of Factor Graphs (DCD-FG), a scalable implementation of f -DAG constrained causal discovery for high-dimensional interventional data. DCD-FG uses a Gaussian non-linear low-rank structural equation model and shows significant improvements compared to state-of-the-art methods in both simulations as well as a recent large-scale single-cell RNA sequencing data set with hundreds of genetic interventions. 10 2 10 3 10 4 Number of nodes 10 -3 10 -2 10 -1 10 0 10 1 Penalty gradient runtime (s/it) Tr-exp Power Spectral Radius Tr-exp Factor Spectral Factor 10 2 10 3 10 4 Number of nodes 10 -3 10 -2 10 -1 Likelihood gradient runtime (s/it) Linear-dense Linear-LR MLP Dense MLP Factor 10 2 10 3 10 4 Number of nodes 10 -2 10 -1 10 0 Total runtime (s/it) NOTEARS NOBEARS NOTEARS-LR DCDI DCD-FG DCD-FG (linear) A. B. C.