What Makes a Representation Good for Single-Cell Perturbation Prediction?

Technical Summary: What Makes a Representation Good for Single-Cell Perturbation Prediction?

Problem Statement
Single-cell perturbation modeling aims to predict how gene expression changes in response to genetic perturbations (e.g., CRISPR-based knockouts). While existing approaches range from causal representation learning to foundation models (FMs), they face a fundamental, often overlooked challenge: signal imbalance. In single-cell data, perturbation-invariant information (background cellular programs) dominates the gene expression space, while perturbation-specific signals are intrinsically sparse.

The authors identify a phenomenon termed the Perturbation Suppression Hypothesis: when invariant variation dominates the data, standard learning approaches fail systematically.

1. Causal Representation Learning: Methods often entangle invariant background information with perturbation-specific representations. This occurs because theoretical identifiability results typically assume all latent variables are intervened upon, whereas real data involves partial interventions (only a subset of genes are perturbed). Consequently, invariant information is absorbed into causal representations, leading to spurious predictors with poor generalization.
2. Foundation Models (FMs): FMs are trained to fit the overall data distribution, prioritizing the dominant invariant structure. Empirical linear probing experiments show that FM representations exhibit weaker decodability of perturbation labels compared to simple PCA baselines, suggesting that perturbation-specific signals are suppressed or rendered inaccessible during pre-training.

Methodology: PerturbedVAE
To address signal imbalance, the authors propose PerturbedVAE (Perturbation-aware Variational Autoencoders), a framework designed to explicitly separate and utilize perturbation-specific information. The framework is built on two core principles: Explicit Extraction and Effective Utilization.

1. Latent Causal Generative Model:
   The model decomposes the latent space into two distinct components:

   • $z_\iota$ (Perturbation-invariant variables): Represents stable background programs shared across all cells, independent of the perturbation condition $u$.
   • $z_\nu$ (Perturbation-responsive variables): Represents latent factors that change in response to genetic perturbations. These are modeled to follow a Directed Acyclic Graph (DAG) structure to capture causal dependencies among perturbation effects.

2. Variational Inference with Contrastive Alignment:
   Standard variational autoencoders (VAEs) might reconstruct data by modeling the dominant invariant component, leaving $z_\nu$ under-recovered. To counteract this, PerturbedVAE introduces a contrastive alignment objective:

   • For every perturbed sample $(x, u)$, a control (unperturbed) sample $x(u_0)$ is sampled.
   • The model enforces that the inferred invariant representations $z_\iota$ for both the perturbed and unperturbed samples are identical (minimizing $\|z_\iota - z_\iota^{(u_0)}\|^2$).
   • This forces the dominant background variation to be explained entirely by $z_\iota$, freeing $z_\nu$ to capture the residual, perturbation-induced changes.

3. Training Objective:
   The total loss combines the Evidence Lower Bound (ELBO) for reconstruction and causal structure learning with the contrastive alignment term:
   $$\mathcal{L} = -\mathcal{L}_{ELBO} + \alpha \mathcal{L}_{contrast}$$
   where $\alpha$ controls the strength of the separation between invariant and responsive factors.

Theoretical Contributions: Identifiability Analysis
The paper provides a rigorous identifiability analysis characterizing the conditions under which sparse perturbation effects can be reliably recovered from partial-intervention data.

• Theorem 1: Under assumptions of invertibility, environmental sufficiency (sufficient diversity in perturbation conditions), optimal alignment, and intervention sufficiency, the framework guarantees:
  • $z_\nu$ (Responsive): Identified up to permutation and scaling (component-wise identifiability).
  • $z_\iota$ (Invariant): Identified up to a linear block transformation.
• This analysis clarifies that the contrastive alignment is not merely heuristic but a necessary condition to prevent invariant information from being absorbed into the causal representations, thereby bridging the gap between theoretical guarantees and the reality of partial interventions.

Empirical Results
The authors evaluate PerturbedVAE on synthetic data and real-world Perturb-seq datasets (Norman et al., 2019; Replogle et al., 2022).

• Synthetic Data: Experiments confirm that contrastive alignment significantly improves the Mean Correlation Coefficient (MCC) for recovering $z_\nu$ and the $R^2$ for recovering $z_\iota$, validating the theoretical identifiability claims.
• Real Data Benchmarks:
  • vs. Foundation Models: PerturbedVAE significantly outperforms FMs (scFoundation, UCE, Geneformer) and simple baselines (RandomForest, ElasticNet) in predicting double-gene (combinatorial) perturbations. FMs degrade in zero-shot combinatorial settings, whereas PerturbedVAE maintains high $R^2$ and low RMSE.
  • vs. Causal Representation Methods: PerturbedVAE outperforms existing latent causal models (Discrepancy-VAE, SENA, sVAE+, SAMS-VAE). Crucially, an ablation study removing the contrastive alignment term ("PerturbedVAE w/o Align") shows performance degradation comparable to existing methods, confirming the necessity of the alignment mechanism.
  • vs. Additive Linear Baselines: While additive linear models perform competitively on condition-level pseudobulk metrics (averaged gene expression), PerturbedVAE achieves superior performance at the single-cell level. Additive models and other deep baselines yield negative cell-level $R^2$, indicating they fail to capture single-cell variability, whereas PerturbedVAE preserves cell-level structure ($R^2 \approx 0.98$).
• Interpretability: The learned causal graph recovers known biological regulatory pathways (e.g., TGFBR2→SNAI1, TP73→CDKN1A), demonstrating that the model learns biologically plausible, structured representations rather than arbitrary statistical associations.

Significance and Claims
The paper argues that the failure of foundation models in perturbation prediction is not primarily a matter of scale but stems from a mismatch between generic training objectives and the sparse, intervention-driven structure of perturbation data.

The significance of PerturbedVAE lies in its demonstration that a "good" representation for single-cell perturbation prediction must be perturbation-aware and causally structured. By explicitly disentangling invariant background programs from perturbation-specific signals via contrastive alignment, the framework enables:

1. Reliable Generalization: Successful zero-shot prediction of unseen combinatorial perturbations.
2. Interpretability: Recovery of causal regulatory mechanisms and perturbation-response programs.
3. Theoretical Grounding: A concrete instantiation of causal representation learning that adheres to identifiability conditions even under partial interventions, a setting where previous theoretical results often fail.

The authors conclude that their approach provides a principled alternative to both purely statistical factorization and generic foundation models, offering a robust solution for understanding cellular responses to genetic interventions.