Latent Causal Diffusions for Single-Cell Perturbation Modeling

The gist: Latent causal diffusion (LCD) is a generative model that frames single-cell gene expression as a stationary diffusion process observed under measurement noise, allowing it to predict distributional shifts of unseen perturbation combinations while learning an explicit, gene-level dynamical system of gene regulation.

How it works

1. The LCD model frames the joint expression of d genes in a cell as a stochastic process x(t) evolving over time t, where each gene g is represented by a random state xg(t) that evolves according to a stochastic differential equation (SDE) dxg(t) = f(x(t))gdt + σdWg(t), where the drift function f: R d → R d models the causal regulatory mechanisms among the genes, and the diffusion term captures stochastic fluctuations with scale σ<ref:3pg1.

2. LCDs model this process x(t) as stationary, meaning the SDEs induce a time-independent distribution p(x) over gene states, which enables inference of the causal dynamics f directly from the state distribution p(x), without requiring time-resolved observations of x(t) itself<ref:3pg1.

3. The drift function f is parameterized by a neural network that shares parameters across the causal mechanisms of each gene, and perturbations modify this hidden state via learned embeddings eq, where q is the index of a perturbation, as gene expressions tend to co-vary across perturbations, and perturbations cluster by their shared transcriptomic effects<ref:3pg1.

4. To interpret the causal dependencies in f, Causal Linearization via Perturbation Responses (CLIPR) is developed to approximate the direct gene-gene causal effects modeled in the diffusion process, which provably identifies causal effects under a linear drift assumption and recovers causal structure in both simulated systems and a genome-wide perturbation screen<ref:3pg1.

Key Findings on Performance

** LCDs predicted the distributional shifts of unseen perturbation combinations more accurately than established approaches while learning an explicit, gene-level dynamical system as the generative model<ref:3pg1. **

** In evaluating against various methods for predicting post-perturbation responses in Perturb-seq assays, LCDs outperformed established approaches in terms of overall distributional accuracy of the predicted single-cell transcriptomes (Fig. 2A)<ref:6pg6. **

** Per-perturbation comparisons on each dataset show that LCDs made more accurate predictions as measured by MMD than all other methods in more than roughly 80% of test perturbations<ref:6pg7. **

** LCDs achieved significant accuracy gains by modeling a full single-cell count distribution, as they capture interaction effects that unroll through the nonlinear dynamics, unlike additive models which only sum differential expression profiles<ref:6pg7. **

Causal Inference and Interpretation

1. CLIPR is used to estimate a linear causal matrix Ab that replicates the same response vectors as the original perturbed drifts, and if the drift function is linear (fq = feq) and L has full row rank, then Ab = A<ref:4pg4.

2. CLIPR distinguishes between direct causal effects and differential expression (DE), showing that DE corresponds to the limit perturbation responses L, while direct causal effects A encode precisely these direct causal dependencies<ref:5pg5.

3. LCD-CLIPR can infer direct gene-gene causal effects among genes in Perturb-seq screens by training on a genome-wide Perturb-seq assay, and it uncovered regulatory mechanisms in K562 cells that remain hidden under standard differential expression analysis<ref:4pg8.

4. Hierarchical clustering of genes based on shortest path distances in the CLIPR graph identifies gene modules enriched for specific Gene Ontology terms and canonical pathways, suggesting that tight dependencies in the LCD-CLIPR graph organize genes into functionally coherent modules<ref:5pg11.

Validation and Generalization

** In simulated systems with linear SDEs, LCD-CLIPR achieved high accuracy in predicting both causal links and their effect signs across all sparsity structures, numbers of genes, and numbers of perturbations<ref:7pg8. **

** When applying LCD-CLIPR to a genome-wide Perturb-seq screen, it uncovered regulatory mechanisms in K562 cells that remain hidden under standard differential expression analysis, providing a scalable framework for mapping causal dependencies in the cell<ref:4pg8. **

** LCD-CLIPR predicted nontrivial causal structure even for unperturbed genes (Fig. 3 and Fig. 4E and F), showing that LCDs can make inferences about unperturbed components of the system<ref:5pg11. **

** CLIPR estimates a full causal structure Ab λ irrespective of whether all or only some genes were perturbed, making predictions about the causal effects of genes not originally perturbed in the screen<ref:5pg12. **

** Stronger predicted causal effects learned by the LCD drift fq aligned strongly with those of the true system, indicating that strong LCD-CLIPR links are consistent with observed downstream effects<ref:7pg8. **

Limitations and Future Directions

LCDs performed best at predicting perturbation effects when the genetic interaction (GI) type is unexpected, synergistic, or suppressive<ref:6pg11. However, the model relies on embeddings learned during training and thus only allows predicting the outcomes of combinations of observed perturbations<ref:6pg7.

** The framework can be extended to drug perturbation screens since perturbation targets need not be known to learn LCDs or compute CLIPR<ref:11pg10. **

** Future research could explore approaches for merging the state density and drift inference steps without dropping established noise model assumptions or the separation of biological and measurement noise<ref:12pg11. **

Methods Summary

The inference process involves two consecutive steps: first, learning likelihood parameters π and state densities pq(x) that maximize evidence of the observed Yq using an empirical Bayes approach<ref:6pg12. Second, learning the drift f and embeddings eq jointly to fit pq(x)<ref:12pg13.

The inference of the implicit prior p(x; ϕq) is achieved by maximizing the evidence Z(Yq, ϕq,π) using a temperature parameter τ to interpolate between fitting full joint distributions and marginal distributions<ref:12pg14. The drift f is inferred by learning to denoise smoothed data xe ∼ N (xe; x, α2I), where the loss for drift f(x; θ, eq, α) is LD(Xq, θ, eq, α)<ref:13pg15.

CLIPR is computed by numerically approximating the fixed point defining the perturbation response lfq by integrating the flow of f conditioned on eq as lfq ≈ limt→∞ x(t), where dx(t)/dt = f(x(t); θ, eq, α1) and x(0) = 0<ref:14pg15. The final causal matrix Ab λ is obtained using Tikhonov regularization: Ab λ = −VL⊤LL⊤ + λI−1<ref:22pg22.

The training involves learning the priors by minimizing LIP (Eq. 11) and then training the drift f and perturbation embeddings eq by minimizing LD (Eq. 12)<ref:3pg14. The model parameters are learned using Adam optimizer with various learning rates and weight decays<ref:3pg14.

The evaluation involves comparing predicted sets of N samples Yb q for perturbation q to the Nq observed post-perturbation samples Yq using the maximum mean discrepancy (MMD) metric<ref:20pg25. The accuracy of Ab λ relative to the ground-truth A is evaluated by classifying predicted effects into positive, negative, and absent using thresholds ±0.5<ref:3pg16.

The study employed various datasets including Perturb-seq data from Norman et al. (4), Wessels et al. (5), and Replogle et al. (1)<ref:15pg27. The analysis focused on d = 1000 highly variable genes<ref:3pg25.

The study also included simulations of linear systems with single-gene shift perturbations to validate the CLIPR estimator, achieving high accuracy in predicting both causal links and their effect signs across all sparsity structures<ref:7pg8. The final analysis involved hierarchical clustering based on shortest path distances in the CLIPR graph to identify gene modules enriched for specific Gene Ontology terms and canonical pathways<ref:4pg11.

The study concluded that LCD-CLIPR can learn biologically plausible causal structure among genes, even when genes were not perturbed in the original screen<ref:5pg12. The framework has the prospect of integrating with different biological modalities and future modeling approaches<ref:16pg33.

The study was supported by several funding agencies including the Swiss National Science Foundation under NCCR Automation, grant agreement 51NF40 180545 and the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany’s Excellence Strategy – EXC number 2064/1<ref:3pg33. The authors declare no competing financial interests<ref:3pg16. All code was implemented in Python, and the source code will be released upon publication of this work<ref:17pg16.

The study has been published on arXiv as 2601.15341v1<ref:2pg20.

1. Latent Causal Modeling for Unseen Combinations The improved system can predict unseen perturbation combinations more accurately than established approaches by framing gene expression as a stationary diffusion process observed under measurement noise, allowing it to model the effects of combinatorial perturbations that are not directly observed in training data.

2. Mechanistic Causal Structure Learning The system learns an explicit, gene-level dynamical system as the generative model and uses causal linearization via perturbation responses (CLIPR) to recover causal structure in both simulated systems and a genome-wide perturbation screen, enabling the identification of direct causal effects between all genes modeled by the diffusion.

3. Disentangling Biological Signal from Technical Noise The framework explicitly models biological stochasticity and technical noise separately, enabling the model to decouple biological stochasticity from technical noise which is crucial for accurate prediction in single-cell data.

4. Inferring Direct Causal Effects Beyond Differential Expression Unlike standard differential expression analysis, the improved system can infer direct causal effects among genes in Perturb-seq screens, distinguishing these from downstream cascades, as it models direct causal effects A rather than just downstream expression changes.

5. Mapping Functional Gene Modules The system can organize genes into coherent functional modules by analyzing the inferred causal graph and performing Gene Set Enrichment Analysis (GSEA) on predicted clusters, revealing how these modules organize themselves according to known Gene Ontology terms and canonical pathways.

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.