[Paper Review] Large-Scale Differentiable Causal Discovery of Factor Graphs

Romain Lopez, Jan-Christian Hütter+2 more|arXiv (Cornell University)|Jun 15, 2022

Bayesian Modeling and Causal Inference9 citations

TL;DR

This paper introduces Differentiable Causal Discovery of Factor Graphs (DCD-FG), a scalable method for large-scale causal discovery using factor directed acyclic graphs (f-DAGs) that enforce low-rank, non-linear structural equations. By combining differentiable optimization with a low-rank constraint on causal interactions, DCD-FG enables efficient causal structure learning on thousands of variables, outperforming state-of-the-art methods in simulations and single-cell RNA-seq data with hundreds of interventions.

ABSTRACT

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

Motivation & Objective

• Address the computational intractability of causal discovery in high-dimensional settings with thousands of variables.
• Overcome the limitations of existing methods that scale poorly beyond ~100 nodes due to cubic complexity in acyclicity constraints.
• Develop a scalable, differentiable framework for causal discovery that leverages low-rank structural assumptions to reduce search space complexity.
• Enable accurate causal structure learning from large-scale interventional data, particularly in high-throughput biological applications.
• Provide theoretical and empirical justification for the use of factor graphs as a tractable and biologically plausible search space for causal discovery.

Proposed method

• Propose factor directed acyclic graphs (f-DAGs) as a low-rank structural constraint, where causal interactions are mediated through a small number of latent factors (m ≪ d).
• Model the joint distribution using a Gaussian non-linear low-rank structural equation model with factor-based parent and child relationships.
• Integrate the f-DAG constraint into a differentiable optimization framework using a continuous relaxation of acyclicity, enabling end-to-end training.
• Use a differentiable penalty for acyclicity with O(md) complexity, significantly reducing computational cost compared to O(d³) methods.
• Apply an augmented Lagrangian method to enforce acyclicity constraints during optimization, enabling scalable training on large graphs.
• Leverage neural networks with leaky ReLU activations and Xavier initialization for flexible, non-linear modeling of gene-to-factor and factor-to-gene relationships.

Experimental results

Research questions

• RQ1 Can a low-rank structural assumption (f-DAGs) significantly reduce the search space of possible causal graphs while preserving expressivity in high-dimensional settings?
• RQ2 How does the differentiable optimization framework DCD-FG perform compared to state-of-the-art methods in terms of structural accuracy on large-scale interventional data?
• RQ3 To what extent is DCD-FG robust to statistical noise and model mis-specification, such as non-Gaussian noise or observational data?
• RQ4 How does the choice of the rank hyperparameter (m) affect the performance and generalization of DCD-FG in practice?
• RQ5 Can f-DAGs effectively model real-world biological systems, such as gene regulatory networks in single-cell RNA-seq data?

Key findings

• DCD-FG achieves state-of-the-art performance in causal structure discovery on synthetic data with up to 1000 variables, significantly outperforming NOTEARS, NOTEARS-LR, and NOBEARS in terms of F1-score and structural accuracy.
• On a large-scale single-cell RNA-seq dataset with 300 genetic interventions and ~1000 genes, DCD-FG successfully inferred a causal graph with m=20 latent factors and 196,303 edges, demonstrating scalability to real-world biological data.
• The method shows strong robustness to model mis-specification: performance degrades only slightly under uniform noise (vs. Gaussian), and the relative ranking of methods remains unchanged.
• DCD-FG maintains high performance even with observational data (no interventions), though accuracy decreases, indicating utility in both interventional and observational regimes.
• The validation likelihood is strongly correlated with downstream performance, justifying the use of hold-out likelihood for hyperparameter selection, including the rank m.
• Theoretical analysis shows that the set of candidate f-DAGs is significantly smaller than the full DAG space, especially under edge perturbations, supporting the use of f-DAGs as a tractable search space in high-dimensional settings.