Representation Learning for Extrapolation in Perturbation Modeling

We consider the problem of modeling the effects of perturbations, such as gene knockdowns or drugs, on measurements, such as single-cell RNA or protein counts. Given data for some perturbations, we aim to predict the distribution of measurements for new combinations of perturbations. To address this challenging extrapolation task, we posit that perturbations act additively in a suitable, unknown embedding space. We formulate the data-generating process as a latent variable model, in which perturbations amount to mean shifts in latent space and can be combined additively. We then prove that, given sufficiently diverse training perturbations, the representation and perturbation effects are identifiable up to orthogonal transformation and use this to characterize the class of unseen perturbations for which we obtain extrapolation guarantees. We establish a link between our model class and shift interventions in linear latent causal models. To estimate the model from data, we propose a new method, the perturbation distribution autoencoder (PDAE), which is trained by maximizing the distributional similarity between true and simulated perturbation distributions. The trained model can then be used to predict previously unseen perturbation distributions. Through simulations, we demonstrate that PDAE can accurately predict the effects of unseen but identifiable perturbations, supporting our theoretical results.

💡 Research Summary

The paper tackles the challenging problem of predicting the distribution of cellular measurements (e.g., single‑cell RNA‑seq or proteomics) under novel perturbations, such as unseen combinations of gene knock‑downs or drug dosages. Rather than predicting only conditional means, the authors formulate a distributional regression task: learn a map a → P_X|a that can extrapolate beyond the support of the training perturbation vectors a₀,…,a_M.

Their central modeling assumption is that perturbations act as mean shifts in a latent space Z of dimension d_Z. For each experimental condition e with perturbation label a_e, a latent basal state Z_base ∼ P_Z is shifted by W a_e (W ∈ ℝ^{d_Z×K}) to obtain Z_pert = Z_base + W a_e. A (possibly stochastic) decoder f maps (Z_pert, ε) to the observed data X, where ε captures variation unrelated to the perturbations. This yields a hierarchical latent variable model that can generate any conditional distribution P_X|a via the push‑forward of (Z, ε).

The authors first prove an identifiability theorem (Theorem 4.1). Assuming (i) the decoder f is a C²‑diffeomorphism, (ii) the latent base distribution is standard Gaussian (centered appropriately), and (iii) the matrix of relative perturbation vectors A =