[Paper Review] Modelling Cellular Perturbations with the Sparse Additive Mechanism Shift Variational Autoencoder

Michael Bereket, Theofanis Karaletsos|arXiv (Cornell University)|Nov 5, 2023

Single-cell and spatial transcriptomicsBiochemistry, Genetics and Molecular Biology8 citations

TL;DR

SAMS-VAE adds sparse additive perturbation mechanisms to a VAE for modeling cellular perturbations, enabling disentangled, interpretable latent subspaces and strong generalization in perturb-seq data. It also introduces correlated inference and an evaluation framework based on average treatment effects.

ABSTRACT

Generative models of observations under interventions have been a vibrant topic of interest across machine learning and the sciences in recent years. For example, in drug discovery, there is a need to model the effects of diverse interventions on cells in order to characterize unknown biological mechanisms of action. We propose the Sparse Additive Mechanism Shift Variational Autoencoder, SAMS-VAE, to combine compositionality, disentanglement, and interpretability for perturbation models. SAMS-VAE models the latent state of a perturbed sample as the sum of a local latent variable capturing sample-specific variation and sparse global variables of latent intervention effects. Crucially, SAMS-VAE sparsifies these global latent variables for individual perturbations to identify disentangled, perturbation-specific latent subspaces that are flexibly composable. We evaluate SAMS-VAE both quantitatively and qualitatively on a range of tasks using two popular single cell sequencing datasets. In order to measure perturbation-specific model-properties, we also introduce a framework for evaluation of perturbation models based on average treatment effects with links to posterior predictive checks. SAMS-VAE outperforms comparable models in terms of generalization across in-distribution and out-of-distribution tasks, including a combinatorial reasoning task under resource paucity, and yields interpretable latent structures which correlate strongly to known biological mechanisms. Our results suggest SAMS-VAE is an interesting addition to the modeling toolkit for machine learning-driven scientific discovery.

Motivation & Objective

• Learn a generative model of cellular observations under perturbations that disentangles perturbation effects from basal variation.
• Introduce sparse perturbation latent offsets that are additively composed and shared across samples receiving the same perturbation.
• Develop and compare correlated variational inference strategies to improve latent disentanglement and predictive performance.
• Provide an evaluation framework for perturbation models using marginal likelihood (IWELBO) and posterior predictive checks via average treatment effects.
• Demonstrate improved generalization and interpretability on perturb-seq datasets compared to baselines.

Proposed method

• Define z_i = z_i^b + z_i^p as the latent state with z_i^p = sum_t d_{i,t} (e_t ⊙ m_t).
• Model e_t ~ N(0, βI) and m_t ~ Bern(α) to induce sparse, perturbation-specific offsets.
• Use a neural network to parameterize p(x_i|z_i; θ) and a scRNA-seq likelihood via a Gamma-Poisson (Negative Binomial) model with library size l_i.
• Infer with stochastic variational inference using a mean-field or correlated variational family that ties z^b, E, M across samples receiving perturbation t.
• Introduce CPA-VAE as an ablated variant with no sparsity mask (m_t fixed to 1).
• Provide two improved inference schemes: correlated z_basal and correlated embeddings E, for a richer variational family.

Figure 2 : Visualization of inferred latent perturbation masks and embedding means for the best performing checkpoint of each model in replogle-filtered . We visualize the latent variables for the 345 perturbations with pathway annotations from Replogle et al. [ 17 ] and group by pathway. The SAMS-V

Experimental results

Research questions

• RQ1 Can SAMS-VAE accurately model perturbation effects as sparse additive latent offsets?
• RQ2 Do correlated inference strategies improve latent disentanglement, interpretability, and predictive performance over baselines?
• RQ3 How does SAMS-VAE generalize to in-distribution and out-of-distribution perturbations and to combinatorial perturbations?
• RQ4 Is the proposed evaluation framework via average treatment effects and differential expression informative for assessing perturbation models?
• RQ5 How does SAMS-VAE compare to CPA-VAE, SVAE+, and conditional VAE on perturb-seq datasets?

Key findings

• SAMS-VAE with fully correlated inference achieves the best test IWELBO and ATE correlation on replogle-filtered data.
• Correlated z_basal inference provides substantial gains across SAMS-VAE and CPA-VAE; correlated E yields modest gains.
• SAMS-VAE with both correlated z_basal and E achieves the highest Mask PW. Acc and ATE-Pearson among the tested configurations.
• SAMS-VAE and CPA-VAE demonstrate strong generalization in combinatorial perturbation settings (norman-ood) and data-efficiency scenarios.
• Latent perturbation masks from SAMS-VAE are more predictive of annotated biological pathways than those from SVAE+.
• Model-based ATE correlates well with data-driven differential expression, supporting the PPC framework.