Technical Summary: Learning Latent Dynamical Causal Processes for Single-Cell Perturbation Prediction

Problem Statement

Single-cell perturbation prediction aims to model how cells respond to genetic or chemical interventions, specifically to predict responses under unseen perturbations (Out-of-Distribution or OOD generalization). A central challenge is that existing machine learning approaches typically address only one of two inseparable characteristics of perturbation responses:

1. Latency: Perturbations act through unobserved cellular programs (latent variables) rather than directly on observed gene expression.
2. Dynamics: The states of these latent programs evolve over time, generating observed gene expression profiles at different timepoints.

Current methods either focus on learning latent causal mechanisms but treat responses as static snapshots, ignoring temporal evolution; or they model temporal changes in gene expression but fail to explicitly capture the underlying causal generative mechanisms driving these changes. This paper addresses the gap by proposing a unified framework that jointly captures latent cellular programs and their temporal evolution.

Methodology: CITE-VAE

The authors propose CITE-VAE (Causality-Inspired temporal Variational AutoEncoder), a learning framework grounded in a Latent Dynamical Causal Generative Model.

1. Generative Model

The model decomposes the latent cellular state $z_t$ at time $t$ into two components:

• Perturbation-invariant background programs ($z_t^\iota$): Representing stable cellular states unaffected by the intervention.
• Perturbation-responsive mechanisms ($z_t^\nu$): Representing dynamic states driven by the intervention.

The generative process is formalized as:

• Exogenous Noise: Invariant noise $n_t^\iota$ is independent, while responsive noise $n_t^\nu$ is conditionally independent given the perturbation label $u$.
• Latent Dynamics:
  • Invariant states evolve via a time-invariant transition: $z_t^\iota = f_\iota(z_{t-1}^\iota) + n_t^\iota$.
  • Responsive states evolve via a structural causal equation dependent on parents and the intervention: $z_{t,\nu,j}^t = f_{\nu,j}(z_{t-1}^{pa(j)}, u) + n_{t,\nu,j}$.
• Observation: Observed gene expression $x_t$ is generated via a time-invariant, smooth, and invertible mapping $g(z_t)$.

2. Theoretical Identifiability

The paper provides a theoretical analysis proving that under specific assumptions, the latent causal variables are identifiable (recoverable up to a trivial equivalence class). The key assumptions are:

• Environmental Sufficiency: There exist sufficiently diverse perturbations to induce distinct distributional changes in the responsive innovations.
• Optimal Temporal Alignment: The learned invariant representation is stable across environments (minimizing an alignment loss).
• Temporal Intervention Sufficiency: For each responsive coordinate, there exists an environment where parent contributions are removed (hard intervention).

Theorem 4.1 establishes that the invariant block is identifiable up to an invertible linear transform, and the responsive block is identifiable up to permutation and component-wise scaling.

3. Learning Framework

To translate these theoretical guarantees into a practical algorithm, CITE-VAE optimizes a joint objective function comprising three components:

1. Temporal ELBO (Likelihood Maximization): Ensures the latent variables retain sufficient information to reconstruct observed snapshots and account for their temporal evolution under interventions.
2. Invariant Alignment Loss: Enforces that the inferred invariant representation ($z^\iota$) remains consistent across different perturbation environments, preventing the entanglement of invariant and responsive factors.
3. Structural Sparsity Regularization: Encourages sparse dependencies in the responsive transition mechanism (via a learnable lower-triangular adjacency matrix) to mitigate overfitting and promote interpretable causal structures.

The training utilizes pseudo-pairs of snapshots from adjacent timepoints to approximate temporal transitions without requiring cell-wise pairing (which is impossible in destructive single-cell sequencing).

Key Contributions

1. Unified Model: Proposes a latent dynamical causal generative model that simultaneously captures latent cellular programs and their temporal evolution, addressing the limitations of static causal or purely temporal models.
2. Identifiability Analysis: Provides rigorous theoretical guarantees showing that latent causal variables can be uniquely recovered (up to trivial transformations) when combining likelihood-based learning with invariant alignment.
3. CITE-VAE Algorithm: Develops a theory-driven variational autoencoder framework that operationalizes these identifiability conditions through a specific combination of ELBO, alignment, and sparsity objectives.
4. Empirical Validation: Demonstrates the method's effectiveness on both synthetic benchmarks and real-world single-cell CRISPR perturbation data.

Experimental Results

Synthetic Data (Causal-3DIdent)

Experiments on the Causal-3DIdent benchmark validate the theoretical identifiability results.

• Recovery: The model successfully recovers ground-truth latent factors with high accuracy ($R^2 > 0.95$ for invariant factors, $0.959$ for responsive factors).
• Trajectory Visualization: Projected latent trajectories closely track ground-truth physical motion, confirming the model captures consistent state transitions over time.

Real-World Data (hiPSC CRISPR Perturbation)

Evaluated on a longitudinal human induced pluripotent stem cell (hiPSC) dataset with 23 single-gene knockouts using a Leave-One-Intervention-Out (LOO) protocol for OOD generalization.

• Performance: CITE-VAE outperforms state-of-the-art baselines, including latent causal models (SAMS-VAE, Discrepancy-VAE), temporal dynamical models (RENGE, scNODE), and foundation models (scGPT, UCE).
• Metrics: It achieves superior performance on differentially expressed (DE) gene metrics (e.g., $\Delta$ Pearson correlation of 0.691 vs. 0.621 for SAMS-VAE) while maintaining competitive genome-wide reconstruction accuracy ($R^2 > 0.95$).
• Ablation: Removing alignment or sparsity regularizers significantly degrades performance on DE-focused metrics, confirming the necessity of the joint objective for capturing perturbation-responsive mechanisms.

Auxiliary Benchmark (E-MTAB-14065)

On a cardiac differentiation dataset with clone-level variation, CITE-VAE achieved the best RMSE among stable baselines and maintained strong performance on DE metrics, demonstrating robustness across different biological contexts.

Significance and Claims

The paper claims that by explicitly modeling the latent and dynamical nature of perturbation responses, CITE-VAE achieves principled generalization to unseen perturbations. The significance lies in:

• Theoretical Grounding: Moving beyond heuristic approaches by providing identifiability guarantees that link the learning objective to the recovery of true causal mechanisms.
• Practical Robustness: Demonstrating that the theoretical insights hold even when strict assumptions are only approximately satisfied in real biological data, leading to improved prediction of unseen gene knockouts.
• Methodological Insight: Highlighting that effective representation learning in dynamical, perturbed systems requires coupling generative modeling with invariance constraints, rather than relying on either in isolation.

The authors acknowledge that, like most causal representation learning work, the approach relies on assumptions (e.g., additive noise, sufficient environmental diversity) that may not strictly hold in all real-world scenarios. However, empirical results suggest the method remains robust and informative under approximate conditions.