Latent Causal Diffusions for Single-Cell Perturbation Modeling

TL;DR

The latent causal diffusion (LCD), a generative model that frames single-cell gene expression as a stationary diffusion process observed under measurement noise, and bridges generative modeling with causal inference to predict unseen perturbation effects and map the underlying regulatory mechanisms of the transcriptome.

Abstract

Perturbation screens hold the potential to systematically map regulatory processes at single-cell resolution, yet modeling and predicting transcriptome-wide responses to perturbations remains a major computational challenge. Existing methods often underperform simple baselines, fail to disentangle measurement noise from biological signal, and provide limited insight into the causal structure governing cellular responses. Here, we present the latent causal diffusion (LCD), a generative model that frames single-cell gene expression as a stationary diffusion process observed under measurement noise. LCD outperforms established approaches in predicting the distributional shifts of unseen perturbation combinations in single-cell RNA-sequencing screens while simultaneously learning a mechanistic dynamical system of gene regulation. To interpret these learned dynamics, we develop an approach we call causal linearization via perturbation responses (CLIPR), which yields an approximation of the direct causal effects between all genes modeled by the diffusion. CLIPR provably identifies causal effects under a linear drift assumption and recovers causal structure in both simulated systems and a genome-wide perturbation screen, where it clusters genes into coherent functional modules and resolves causal relationships that standard differential expression analysis cannot. The LCD-CLIPR framework bridges generative modeling with causal inference to predict unseen perturbation effects and map the underlying regulatory mechanisms of the transcriptome.

Figures (10)

• Figure 1: The latent causal diffusion (LCD) model and CLIPR. (A) LCDs model single-cell expression data $\mathbf{y}$ (top) as noisy observations of unobserved, noiseless gene states $\mathbf{x}$ (center). States $\mathbf{x}$ are samples from a stationary stochastic process $\mathbf{x}(t)$ evolving under the causal regulatory dynamics $f(\mathbf{x})$ of an SDE (bottom). To train LCDs, we first infer $p(\mathbf{x})$ via empirical Bayes, then learn $f$ in state space via score matching. (B) Dynamics $f$ explicitly model gene regulation and are parameterized by a neural network with shared hidden state across genes. Perturbations modify this hidden state via learned embeddings $\mathbf{e}_{q}$. (C) To interpret the causal dependencies in $f$, we estimate a linear causal matrix that induces the perturbation behavior of $f$ (left). This matrix is computed from the initial ($\mathbf{v}_{f_q}\xspace$) and limit ($\mathbf{l}_{f_q}\xspace$) responses when $f$ is perturbed. The limit response is approximated by following the vector field flow until convergence (right).
• Figure 2: LCDs outperform state-of-the-art models at predicting combinatorial perturbation effects. (A) Maximum mean discrepancy (MMD) between predicted and observed single-cell expressions for held-out double-gene perturbations of the two datasets, computed on top 20 differentially expressed (DE) genes. Each dot represents one perturbation. Two leftmost bars show baseline performance of randomly sampled control and perturbed cells from all conditions. (B) Per-perturbation comparison of LCD with deep-learning baselines on MMD, normalized to $[0,1]$ across methods. Win rate (WR): proportion of perturbations where LCD achieved lower error. (C) Root mean squared error (RMSE) of predicted mean expression of top 20 DE genes (top) and Pearson correlation on full mean expression vector (bottom) (D) MMD stratified by genetic interaction (GI) type, comparing LCD and the additive baseline SALT. Axis labels indicate significant improvement per GI category (percentage reduction in median MMD by LCD) or not (n.s.; one-sided Wilcoxon signed-rank test, $P < 0.05$). (E) Predicted vs. observed expression for top 20 DE genes of a held-out perturbation with neomorphic GI. Grey section shows the perturbed genes (bold labels). All box plots show the median, interquartile range (IQR), and whiskers extend to farthest points within $1.5 \times$ IQR.
• Figure 3: LCD-CLIPR accurately recovers causal effects in linear systems from learned diffusion drifts. (A) AUROC for classifying gene-gene causal effects (positive, negative, or absent) across all gene pairs. Results are shown for systems of 100, 300, and 1000 genes across varying training perturbation counts ($x$-axis), comparing power-law (left) and Erdős-Rényi (right) regulatory dependencies in $\mathbf{A}$. (B) F1 score for classifying gene-gene causal effects (left) and average Pearson correlation between predicted and true causal effects (right) for the top $k$ predicted effects, ranked by absolute magnitude, under power-law dependencies. Markers indicate medians, error bars indicate 10--90th percentiles across test systems. (C) Comparison of causal effects inferred by LCD-CLIPR vs. the ground-truth data-generating system (300 genes). (D) Inferred and true perturbation vectors for the system shown in C. For visual clarity, only 80 perturbations and 150 genes with the largest outgoing effects are displayed.
• Figure 4: LCD-CLIPR infers direct gene-gene causal effects from genome-wide Perturb-seq data. (A) Hierarchical clustering based on shortest-path distances in the CLIPR graph identifies gene modules enriched for specific Gene Ontology terms (top) and canonical pathways (bottom). Clusters with at least one enriched term are shown. (B) Predicted causal effects and associated clusters from the enrichment analysis. (C) Comparison between observed differential expression (DE) (left) and normalized $\mathbf{L}^\top$ (right). For visual clarity, only the 247 genes within identified clusters and their perturbations are shown; see \ref{['fig:full-resolution-causal-replogle']} and \ref{['fig:full-resolution-causal-replogle-mats']} for full matrices. (D) Strongest predicted causal effects (50 total). Edges are colored green where the predicted effect is validated by observed DE, and red otherwise. (E and F) Validation of inferred causal links (C) as a function of predicted effect strength. Stronger predicted causation significantly increases the probability of observing DE (one-sided Fisher's exact test, $P < 0.01$, OR: odds ratio). Results shown for all genes (left) and held-out genes not perturbed during training (right). Error bars indicate 95% CIs (Wilson binomial for E, bootstrap for F).
• Figure S1: Training and evaluation splits of the combinatorial gene perturbation data by Wessels et al. wessels2023efficient and Norman et al. norman2019exploring used for benchmarking (Fig. \ref{['fig:benchmarking']}). (A) Each dataset contains control samples ($q_0$), $u$ single gene perturbations, and $k$ double gene perturbations. Control and single gene perturbation samples are always in the training set. (B) We reserved 20.0 double gene perturbations for hyperparameter tuning of each method (\ref{['sec:methods']}), where 10.0 were included in the training set and 10.0 were used for evaluation. (C) For benchmarking, we evaluated methods on 10.0 disjoint testing folds of the held-out double perturbations. All perturbations used for hyperparameter tuning were included in the training set to avoid information leakage. Combining all testing folds, methods generated held-out predictions for all $k - 20$ held-out double perturbations.

Theorems & Definitions (5)

• Theorem 1
• Corollary 2
• Theorem 3
• Corollary 4
• Lemma 5