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

Single-cell perturbation modeling is fundamental for understanding and predicting cellular responses to genetic perturbations. However, existing approaches, from causal representation learning to foundation models, often struggle with an overlooked challenge: gene expression is dominated by perturbation-invariant information, while perturbation-specific signals are intrinsically sparse. As a result, learned representations either entangle invariant and perturbation-specific information, leading to spurious and non-generalizable predictors, or suppress perturbation-specific signals altogether, rendering them ineffective for prediction. To address this, we propose PerturbedVAE, a general framework designed to resolve this signal imbalance. The framework explicitly separates perturbation-specific information from dominant invariant structure and recovers causal representations to effectively utilize such information for prediction. We further provide an identifiability analysis that characterizes the conditions under which sparse perturbation effects can be reliably recovered, thereby clarifying how the framework can be concretely specified under such conditions. Empirically, PerturbedVAE achieves state-of-the-art performance on a widely used benchmark across multiple evaluation settings, yielding significant gains on out-of-distribution combinatorial predictions and uncovering interpretable perturbation-response programs.

This paper studies how ideas from causal representation learning can be used to predict how cells respond when scientists change one or more genes. The goal is not only to predict gene activity after a perturbation, but also to understand which latent factors drive the cellular response. A key challenge is that most measured gene activity reflects a cell’s normal background state, while perturbation-induced signals are often small and hard to separate. We propose a machine learning method that explicitly separates normal cellular activity from perturbation-induced signals, allowing the model to focus on the factors most relevant for prediction. We further study when these latent perturbation-related factors can be reliably recovered from data, connecting practical single-cell perturbation prediction with theoretical questions in causal representation learning. Experiments show that our method improves prediction accuracy, especially for new combinations of gene perturbations, while producing response patterns that are easier to interpret biologically.