Representation Learning for Distributional Perturbation Extrapolation

We consider the problem of modelling the effects of perturbations such as gene knockdowns or drug combinations on low-level measurements like RNA sequencing data. Specifically, given data collected under some perturbations, we aim to predict the distribution of measurements for new perturbations. To address this challenging extrapolation task, we posit that perturbations act additively in a suitable, unknown embedding space. More precisely, we formulate the generative process underlying the observed data as a latent variable model, in which perturbations amount to mean shifts in latent space. We prove that the representation and perturbation effects are identifiable up to affine transformation and use this to characterize the class of unseen perturbations for which we obtain extrapolation guarantees. To estimate the model from data, we propose the perturbation distribution autoencoder (PDAE) which is trained by maximising the distributional similarity between true and predicted perturbation distributions The trained model can then be used to predict previously unseen perturbation distributions. Preliminary empirical evidence suggests that PDAE compares favourably to CPA (Lotfollahiet al., 2023) and other baselines at predicting the effects of unseen perturbations.