Title: Modelling Cellular Perturbations with the Sparse Additive Mechanism Shift Variational Autoencoder
Venue: NeurIPS 2023
Accepted: yes
Keywords: Disentagled representation learning

VAE

generative models

sparse mechanism shift

perturbation modeling

cellular modeling
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.

Modelling Cellular Perturbations with the Sparse Additive Mechanism Shift Variational Autoencoder

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-V AE 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.

1 Introduction

Scientific discovery often involves observation and intervention on systems with the aim of eliciting a mechanistic understanding. For example, in biology, large cellular perturbation screens with high-dimensional readouts have become increasingly popular as an approach to investigate biological mechanisms, their regulatory dependencies, and their responses to drugs. As technology enables both richer and finer grained measurements of these systems, there is an increasing need and opportunity for machine learning methods to help generate predictive insights of growing complexity.

Generative models such as variational auto-encoders (V AEs) [8] are commonly used to learn representations of complex datasets and their underlying distributions. A common goal in generative modeling is disentanglement, whereby latent structures should factorize into semantic subspaces to facilitate generalization and discovery. A desirable outcome consists of these subspaces learned by models being indicative of latent mechanisms, while sparsely varying according to the underlying latent factors of variation in the true data distribution [10]. This goal has recently been formalized under the Sparse Mechanism Shift framework [20, 9] which connects disentanglement to the causal inference field through the identification of causal graphs. Concomitantly, recent models such as the Compositional Perturbation Autoencoder [13] and SVAE+ [12] have successfully applied disentangled deep learning to scientific problems in single-cell RNA-sequencing under perturbation.

In this work, we propose the Sparse Additive Mechanism Shift Variational Autoencoder (SAMSVAE), a model which extends prior work by capturing interventions and their sparse effects as explicit additive latent variables. Compared to previous approaches for modeling disentanglement in VAEs applied to cellular data, our model explicitly combines sparse perturbation-specific latent effects, perturbation-independent natural variation of cells, and additive composition of perturbation effects in a joint model. We also introduce CPA-VAE, which ablates the sparsity mechanism we propose, yielding a generative model with similar assumptions as the popular perturbation model CPA. To perform approximate inference, we propose rich variational families for these models and showcase how sophisticated inference facilitates identifying predictive factors of variation. We additionally introduce a lens on evaluation of perturbation models for biology based on model-based average treatment effects and differential expression, which we link to posterior predictive checks.

In our experiments we showcase SAMS-VAE in various tasks across cellular sequencing data. We observe that SAMS-VAE achieves superior predictive capability over baselines across two popular single-cell sequencing datasets in tasks related to in-distribution- and out-of-distribution generalization, including combinatorial generalization when multiple perturbations are applied. We furthermore examine the interpretability of the model's disentangled structures and demonstrate significantly improved ability to recover factors predictive of known molecular pathways as compared to recently proposed models. Finally, we show that our best models also excel in the treatment effect estimation evaluation we propose.

2 The Sparse Additive Mechanism Shift Variational Autoencoder

We consider datasets ( x i , d i ) N i =1 of observations x i ∈ R D x and perturbation dosage vectors d i ∈ { 0 , 1 } T , where d i,j is 1 if sample i received perturbation j and 0 otherwise. We aim to develop generative models of p ( x | d ) , representing the distribution of features of a target system conditional on perturbations. In the following sections, we will introduce the details of our proposed modeling strategy, the Sparse Additive Mechanism Shift Variational Autoencoder (SAMS-VAE).

2.1 Generative model

We consider generative models with the following basic structure:

z i ∈ R D z is the latent state embedding for sample i , which is modeled as the sum of a latent basal state embedding z b i ∈ R D z and a latent perturbation effect embedding z p i ∈ R D z . Observations are then sampled from a conditional likelihood p ( x i | z i ; θ ) . In this paper, we focus on likelihoods for p ( x i | z i ; θ ) , where parameters are computed from z i using a neural network with parameters θ .

The core modeling assumption of SAMS-VAE relates to the distribution p ( z p i | d i ) . We propose to model perturbations as inducing sparse latent offsets that compose additively as follows:

where e t ∈ R D z and m t ∈ { 0 , 1 } D z are global latent variables that determine the latent offset due to perturbation t . m t is a binary mask that performs feature selection on the latent offset e t : when m t is sparse e t /circledot m t will result in a sparse offset. Importantly, global variables e t and m t are shared across all samples (corresponding to cells) that receive perturbation t .

We specify the prior distributions p ( e t ) ∼ N (0 , β I ) and p ( m t ) = Bern ( α ) for perturbation effects, where α is chosen to be small to induce sparsity. While we focus on the Bernoulli prior for the mask, we also provide a Beta-Bernoulli prior in our code for mask m t as an easy plug-in replacement. We omit additional prior assumptions regarding the structure of perturbation effects in this work. We specify the prior distribution p ( z b i ) ∼ N (0 , I ) for latent basal states.

Using this latent structure, we define the full generative model for SAMS-VAE as in Figure 5. The joint probability distribution over our observed and latent variables is defined as:

Algorithm 1 SAMS-VAE generative process

Require:

X

∈

R

N

×

D

x

,

D

∈

{

0

,

1

}

N

×

T

for

t

from

1

to

T

do

e

t

∼

N

(0

, I

)

m

t

∼

Bern

(

α

)

end for

for

i

from

1

to

N

do

z

b

i

∼

N

(0

, I

)

z

p

i

=

∑

T

t

=1

d

i,t

(

e

t

/circledot

m

t

)

z

i

=

z

b

i

+

z

p

i

x

i

∼

p

(

x

i

|

z

i

;

θ

)

end for

x

+

dot

z

b

e

t

m

t

θ

N

T

for observations X ∈ R N × D x , perturbation dosages D ∈ { 0 , 1 } N × D z , latent basal states Z b ∈ R N × D z , latent perturbation embeddings E ∈ R T × D z , and latent perturbation masks M ∈ { 0 , 1 } T × D z .

2.2 Likelihood choice for scRNA-Seq Data

In the previous section, we used a generic form p ( x i | z i ; θ ) for the observation model over x to show the generality of the approach. Below, we describe the observation model we use to apply SAMS-VAE to single cell RNA-sequencing data (scRNA-seq) in more detail.

We represent scRNA-seq observations as x i ∈ N D x , where each value x i,j is the number of measured transcripts in cell i that correspond to gene j , and follow the likelihood introduced by Lopez et al. [11, 12] to model the elaborate noise structure of scRNA-seq data. An additional utility quantity library size l i , the total number of transcripts measured in cell i , is included as an observed variable in the conditioning set. This is useful because the library size is largely determined by technical factors that we are not interested in modeling. The full likelihood function is then defined as follows:

where ρ i ∈ [0 , 1] D x represents the expected frequency of each transcript in cell i and is parameterized by f θ , a neural network with a softmax output. Observations are then sampled from a Gamma-Poisson distribution (equivalently, a negative binomial distribution) with mean ρ i l i ∈ R D x + and inverse dispersion θ d ∈ R D x + . θ d is a learned parameter that is shared across cells.

2.3 Inference

We perform inference on SAMS-VAE using stochastic variational inference [4, 8] to approximate the marginal likelihood log p ( X | D ) by optimizing parameters φ and θ . We do so by maximizing the evidence lower bound (ELBO) for SAMS-VAE, defined as follows:

A key question when performing variational inference is the choice of variational family to approximate the posterior distribution. As a baseline inference strategy, we consider the following amortized mean-field inference scheme for SAMS-VAE:

We parameterize q ( m t ; φ ) = Bern ( ˆ p t ) and q ( e t ; φ ) = N ( ˆ µ t , ˆ σ t ) with learnable parameters ˆ p t , ˆ µ t , ˆ σ t . We define q ( z b i | x i ) = N ( ˆ f enc ( x i )) , where ˆ f enc is a learnable neural network that predicts mean and standard deviation parameters. During training, gradients are computed for q ( m t ; φ ) with a Gumbel-Softmax straight-through estimator [5].

We propose two improvements to the mean-field inference scheme that aim to more faithfully invert the SAMS-VAE generative model. First, we model possible correlations between sample latent basal states z b i and the global latent perturbation masks and embeddings ( correlated encoder ). We do so by implementing q ( z b i | x i , d i , E , M ) = N ( ˆ f enc ([ x i z p i ])) for z p i as defined in equation 1, where ˆ f enc is a neural network that takes as input both the observations and the estimated latent perturbation effect embeddings for a given sample. Second, we model possible correlations between the latent perturbation masks and embeddings by replacing q ( e t ) with q ( e t | m t ) ( correlated embeddings ). We implement q ( e t | m t ) = N ( ˆ f emb ( m t , t )) with a learnable neural network ˆ f emb that predicts the embedding from a mask and a one-hot encoding of the treatment index. Applying both of these modifications, we define the correlated variational family for SAMS-VAE as:

This richer variational family posits a joint infinite mixture variational distribution between the global and local variables in the model to finely capture their interdependencies and we evaluate its components separately in our experiments. We elaborate on the objective per minibatch and other details in our supplemental Section A.1.

2.4 CPA-VAE

To directly assess the effect of the sparsity inducing masks, we define an ablated model, CPA-V AE, that is identical to SAMS-VAE with all mask components fixed to 1. Thus, in contrast to equation 1, we have that z p i = ∑ T t =1 d i,t e t . We call this model CPA-VAE because it directly incorporates the additive latent composition assumption from CPA [13], and CPA-VAE can be thought of as an extension of CPA to a fully specified generative model. We note that CPA-VAE inherits the benefits of the inference improvements to the variational families we propose in this work and will assess the contributions of better inference and sparse masking separately.

3 Quantitative Evaluation of Perturbation Models

In this section we discuss two quantitative strategies to rigorously evaluate our perturbation models. First, we discuss how to estimate the marginal likelihood for our model on held-out data, a common strategy employed across generative modeling to assess density estimation. Second, we define a posterior predictive check for model predictions of average treatment effects.

3.1 Marginal Likelihood

We consider the marginal log likelihood of held out data under an inferred generative model, estimated via the importance weighted ELBO (IWELBO) [1], as our primary evaluation metric (a similar metric was used in Lopez et al. [12]) Specifically, we estimate log P ( X | D , θ ) on held out data, where θ denotes decoder parameters. Let H = ( Z b , M , E ) represent the set of latent variables for SAMS-VAE. Then we can write the importance weighted ELBO with K particles as:

The importance weighted ELBO can be used to holistically compare the generalization of generative models such as SAMS-VAE, CPA-VAE, and SVAE+. We note, however, that a marginal likelihood cannot be computed for models that are not fully specified as probabilistic models. In practice, we estimate IWELBO as follows:

for

3.2 A Posterior Predictive Check for Average Treatment Effects of Perturbation Models

As a second category of metrics, we consider posterior predictive checks (PPC) [2, 19]: we query test statistics of interest in the predictive distribution of learned models and compare these statistics against estimates from the data. These types of assessments can be useful when critiquing models for specific use cases, such as predicting the mean of some measurement under different perturbations. However, these assessments only characterize narrow aspects of the predictive distribution, providing a less complete assessment than the marginal likelihood.

As a test statistic for our PPC we choose the population average treatment effect of a perturbation relative to a control perturbation on each measurement x i,j for sample i and gene j , given given as:

We define the average treatment effect for SAMS-VAE as ATESAMS-VAE ( d ∗ | D m ) for an applied treatment d ∗ and conditioning data D m (the training data) as the difference between the expected predictive value of output variable x i,j given a treatment d ∗ and the expected predictive value of x i,j given control treatment d 0 :

Both of the inner expectations share global and local latent variables and only differ in the treatments, while marginalizing over observation noise. We thus disentangle between noise differences caused by a treatment, since observation noise is marginalized out. We also marginalize over the prior basal state p ( z b ) in the outer expectation, which simulates populations of different cells varying by natural variation. In practice we draw K samples z b k ∼ p ( z b ) and E k , M k ∼ p ( E , M | D m ) for the outer expectation and evaluate the inner expectations by a small amount of S samples. In cases where observations contain multiple features (i.e. genes j ), this quantity yields a vector per feature. Using this approach we can generate perturbed cells using the dosages d . We note that other models are treated equivalently when feasible by handling their global and local variables analogously.

Because we cannot directly observe counterfactuals, we must identify a related observed quantity to evaluate our model estimated average treatment effects. We reach for differential expression (DE), a commonly chosen metric to study sequencing data collected under different conditions. A key difference between differential expression and average treatment effects is differential expression's computation based on differences of population averages.

A second key difference is that the model-based ATE marginalizes out observation noise per sample, while DE cannot distinguish noise from perturbation effects. We note that differential expression as such takes the form of the following expression and is computed over a dataset D δ over which the expectation is computed (where D d δ denotes the subset of the dataset D δ under condition d ):

To create the metric ATE-Pearson r ( ATESAMS-VAE ( d ∗ | D m ) , DEData ( d ∗ | D δ )) we compute the Pearson correlation coefficient r between differential expression estimates from data DEData and our model-based estimator ATESAMS-VAE across all features (commonly genes indexed by j ).

This also reveals the relationship to a PPC p ( r ( ATESAMS-VAE ( d ∗ | D m ) , DEData ( d ∗ | D δ )) | D m , D δ ) , considering ATE as the diagnostic statistic which is approximated by DE in the observed sample. The dataset D m for conditioning or training the model and a dataset D δ for estimating differential expression may be the same or different, depending on the use case. We note that the utilization of a separate dataset for a PPC is unconventional, but has previously been used in HPCs [14].

4 Experiments

Overview We compare SAMS-VAE with baseline models through a series of applications to perturb-seq datasets. Perturb-seq is a type of biological experiment in which cells are individually perturbed and subsequently profiled with single cell RNA sequencing (scRNA-seq). Single cell RNA sequencing measures the count of messenger RNA (mRNA) transcripts (also called gene expression) for thousands of genes in each cell, providing a rich, high-dimensional characterization of cellular state. Common perturbation types for perturb-seq experiments include genetic knockouts, which disable the expression of target genes through gene editing, and chemical compounds.

In our experiments, we represent perturb-seq datasets as a gene expression matrix X ∈ N N × D x and a perturbation dosage matrix T ∈ { 0 , 1 } N × T for N cells, D x gene transcripts, and T perturbations. Entries X i,j represent the number of transcripts from gene j observed in cell i and entries T i,k represent whether cell i received perturbation k . We compare SAMS-VAE against baseline models based on their ability to model the distribution of perturb-seq data, generalize to new perturbations in combinatorial settings, and disentangle known biological pathways in their latent variables.

Baselines We consider CPA-VAE , SVAE+ , and conditional VAE [21] as baselines. As discussed in Section 2.4, CPA-VAE can be thought of as an extension of CPA [13] to a fully specified generative model and takes advantage of our proposed correlated inference strategy. We additionally consider ablations of the correlated inference strategies for SAMS-VAE and CPA-VAE. Complete details of model choices are provided in the appendix.

Code availability Our code, which includes implementations of all models and experiment configurations, is available at https://github.com/insitro/sams-vae .

4.1 Generalization under individual perturbations

Dataset To assess model generalization to held out samples under individual perturbations, we analyze a subset of the genome-wide CRISPR interference (CRISPRi) perturb-seq dataset from Replogle et al. [17], which we call replogle-filtered . CRISPRi is a type of genetic perturbation that represses the expression of selected target genes. Following the preprocessing steps from Lopez et al. [12], replogle-filtered is filtered to contain perturbations that were identified as having strong effects and genes that were associated with these perturbations. We additionally include cells with non-targeting CRISPR guides to use as controls for average treatment effect prediction. All together, replogle-filtered contains 118,461 cells, 1,187 gene expression features per cell, and 722 unique CRISPR guides (perturbations). We randomly sample train, validation, and test splits.

Perturbations (Guides)

Embeddings

SAMS-VAE

SVAE+

Mask

Probabilities

Perturbations (Guides)

Annotated Pathways

CPA-VAE

Exosome

Mediator Complex

MT Protein Translocation

Nucleotide Excision Repair

39S Ribosomal Subunit

40S Ribosomal Subunit

60S Ribosomal Subunit

Spliceosome

-101-10°0100

101

-101-10°0100

101

-101-10°0100

101

0

50

50

50

.00

100

100

.50

150

150

200

200

200

250

250

250

800

300

300

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

0.0

0.5

1.0

0.0

0.5

1.0

0

Q

50

50

.00

100

.50

150

200

200

250

250

800

300

0

20

40

60

80

100

0

20

40

60

80

100

Evaluation protocol Each model is trained with a 100 dimensional latent space and MLP encoders and decoders with a single hidden layer of dimension 400 (see Section A.4 for full training details). Based on validation performance and sparsity, a Beta (1 , 2) prior was selected for the SV AE+ mask, and a Bern (0 . 001) prior was selected for SAMS-VAE. For each model type, we compute the test set importance weighted ELBO as described in Section 3.1 and report the mean and standard deviation across five training runs with different random seeds. We additionally estimate the model average treatment effect with K = 10 , 000 particles as defined in Section 3.2 for the best model of each type and report the correlation between this quantity and the estimated differential expression from data.

Results Quantitative results are presented in Table 1. Comparing first between model types, we observe that SAMS-VAE with fully correlated inference achieves the best test IWELBO and average treatment effect correlation. Interestingly, CPA-V AE with correlated inference achieves strong test IWELBO performance but falls behind on average treatment effect prediction, while conditional V AE has weak IWELBO performance but achieves strong average treatment effect prediction. SVAE+ does not perform well on either metric in this setting.

In addition to comparing model types, we perform an ablation of SAMS-VAE and CPA-VAE inference strategies. We find that the correlated z basal strategy yields substantial improvements in performance for both SAMS-VAE and CPA-VAE, while the correlated E strategy improvements are minor.

4.1.1 Recovery of biological mechanisms based on disentangled factors

Evaluation protocol We assess the degree to which the pattern of perturbation effects on inferred latent factors in the SAMS-VAE and SVAE+ models from the previous section are predictive of

ATEofSAMS-VAEperinterventionpergene

DifferentialExpressionperinterventionpergene

0

50

100

100

100

Perturbations

150

F10-1

F10-1

0

0

200

-10-1

-10-1

250

300

-100

-100

0

200

400

600

800

1000

0

200

400

600

800

1000

Genes

Genes

known biological pathways as annotated by Replogle et al. [17] (345 of the 722 targeted genetic perturbations are annotated). To do so, we define an inferred binary mask of perturbation effects on latent factors by thresholding the inferred latent mask probabilities in each model at p = 0 . 5 . For each model, we fit a random forest model using scikit-learn [16] to predict pathway annotations from a subset of the perturbation latent masks and assess pathway prediction accuracy on the remaining perturbations. This evaluation is performed on the best checkpoint for each model type with 10 random splits of perturbations (70% train, 30% test), and the mean and standard deviation of the pathway prediction accuracy is reported.

We also provide a set of visualizations to qualitatively assess the latent structures learned by each model. We plot the inferred masks and embeddings for SAMS-VAE, SVAE+, and CPA-VAE in Figure 2, and visualize the SAMS-VAE estimated average treatment effects and estimated differential expression corresponding to these perturbation effects in Figure 3. Hierarchical clustering and UMAP projection of the inferred perturbation embeddings are presented in Figure 7 in Section A.6.

Results We observe that the latent mask inferred by SAMS-VAE is more predictive of the annotated pathways in the replogle-filtered dataset than that inferred by SV AE+. Additionally, we find that performing correlated inference on the perturbation embeddings improved pathway prediction performance for SAMS-VAE. Qualitatively, we observe that both SAMS-VAE and SVAE+ infer sparse masks with distinct patters between annotated pathways.

4.2 Modeling compositional interventions in a CRISPRa perturb-seq screen

Dataset We analyze the CRISPR activation (CRISPRa) perturb-seq screen from Norman et al. [15] to assess how effectively SAMS-VAE and baselines model the effect of perturbation combinations. This screen was specifically designed with perturbations that have non-additive effects in combination, making this a challenging setting for modeling combinations. We adopt the preprocessing from [6], which contains 105 unique targeting guides applied both on their own and in 131 combinations. In total, the dataset contains 111,255 cells, each with 5,000 gene expression features.

Evaluation protocol We define two tasks using this data. The first, norman-ood , assesses the ability of each model to predict gene expression profiles for held-out cells that have received perturbation combinations that are not included in the training set. Each model is trained on cells that received a single guide, along with [0, 25, 50, 75, 100]% of combinations. Held-out cells receiving the final 25% of combinations are used to evaluate each model. We perform this analysis for 5 random splits of the combinations. The second task, norman-data-efficiency , assesses how efficiently the models can learn combination phenotypes when trained on cells that have received a single guide and increasing numbers of cells sampled uniformly across all combinations. Each model is evaluated

Test combinations not

observed in training

Combination cells uniformly subsampled

from full training set

All combinations

observed in training

Relative IWELBO

Avg. Treatment

Effect Correlation

Fraction combinations

in training set

Subsample fraction of cells receiving

combinations in training set

Out-of-Distribution Combinations

Combination Data Efficiency

12

12

10

10 -

8

8

9

6

4

4

2

2

0

0

-2

-2

4

0.0

0.25

0.5

0.75

1.0

0.0

0.25

0.5

0.75

1.0

0.8-

0.8

0.6

0.6-

0.4-

0.4 -

0.2

0.2

0.0

0.0

0.0

0.25

0.5

0.75

1.0

0.0

0.25

0.5

0.75

1.0

SAMS-VAE

CPA-VAE

ConditionalVAE

SVAE+

based on the IWELBO and ATE-Pearson on the held out test set. To compare model performance across different data splits, within each split we analyze the test IWELBO of each model relative to the test IWELBO of SAMS-VAE trained with no combinations on that split (relative IWELBO). Average treatment effects are predicted with 2,500 particles, and IWELBO values with 100 particles.

We train each model with latent dimension 200 and single hidden layer MLP encoders and decoders for 30,000 training steps. Based on validation performance, SV AE+ is trained with a Beta (1 , 2) prior and SAMS-VAE is trained with a Bern (0 . 01) prior.

Results Quantitative results are presented in Figure 4, with additional inference strategy ablations in Figure 10 in Section A.6. SAMS-VAE and CPA-VAE both achieve strong performance on the norman-ood task across metrics, often within 1 standard deviation of one another. Conditional V AE achieves similarly strong performance for average treatment effect prediction, though is weaker on the IWELBO metric. Unsurprisingly, SVAE+, which models combinations as totally new treatments, is unable to predict the effect of a new combination without observing it in training. We do observe that the SVAE+ likelihood still improves a small amount as more combinations are included in training set, which may be attributable to improvments in the encoder and decoder (which are shared across perturbations). These results support the utility of the compositional mechanisms in SAMS-V AE and CPA-VAE (and for encoding combinations as defined in d i for conditional V AE).

In norman-data-efficiency , we observe similar trends. SAMS-VAE, CPA-VAE, and conditional VAE, which can share information across individual and combined perturbations, all achieve better ATE prediction for held out cells when less than 50% of the available combination cells. However, SVAE+ achieves similar ATE prediction correlations on this dataset when presented with sufficient

combination samples in training. Looking at the relative IWELBO values, we observe that SAMSVAE and CPA-VAE again perform the best, with SVAE+. These results further support the utility of the additive composition mechanism from SAMS-VAE and CPA-VAE in low data settings.

5 Related Work

Disentangled VAEs Disentangled variational auto-encoders have been proposed as early as in [7], where weak supervision was utilized to learn sparse masks over different subspaces. A popular framework for unsupervised disentanglement was proposed in [3] through a reweighting of the regularizer in the objective, but ignores weak supervision about conditions. A more comprehensive treatise and theoretical analysis of disentanglement was presented in [10]. Finally, the formal link to sparse mechanism shift and explicit causal disentanglement was also established recently in [9]. Our work shares assumptions with some of these works, in that we assume dosage is known leading to specific shifted effects per perturbation that can be used to learn disentangled factors.

Models of Cellular Perturbation A popular generative modeling framework for cellular sequencing data utilizing V AEs has been proposed in [11], a model which inspired our use of their likelihood. Closer to our application on perturbed datasets is the CPA [13]. Similar to this model, we adopt the idea to disentangle cellular latent spaces into basal and perturbation latent variables. However, we pose the resulting model as a joint generative model with a rigorous inference framework and, crucially, a sparsity mechanism to disentangle the perturbation effects into subspaces related to the affected mechanisms. The recent SVAE+ [12] model is an exciting variant of [9] that utilizes disentanglement in a fashion that matches our goals. Our work differs by factorizing variation into basal and perturbation variables, adding a mechanism to compose perturbations, and in terms of inference strategies (SV AE+ learns perturbation representations by optimizing a prior). GEARS [18] leverages prior information of perturbation and gene features to predict the effect of applying new perturbations to unperturbed cells. This work instead focuses on specifying a generative model for perturbation effects with minimal assumptions, though strategies for integrating prior information on perturbations and features in specific use cases is an exciting future direction.

6 Conclusion

Performing unbiased scientific discovery is an aspirational goal in the field of drug discovery to detect mechanisms of action with intervenable potential. We propose a model that attempts to use few explicit assumptions about the nature of the observed data and relies heavily on a sparsity assumption and decomposition into explicit treatment effects in latent space to learn models of perturbational screening data, making it general enough for application to arbitrary data modalities and perturbation types. In this work, we apply SAMS-VAE to genetic perturbations and single-cell sequencing readouts and observe two key outcomes: improved predictive performance compared to omitting the sparsity assumption, and improved ability to recover factors correlated with real mechanisms in biological data. Our technical contributions cover both a specification of a novel sparse generative model, SAMS-VAE, as well as a suite of inference strategies for improved model fit which also apply to our baseline CPA-VAE. We also propose an evaluation strategy for perturbation models related to posterior predictive checks utilizing average treatment effects and differential expression as test statistics to perform model criticism, and observe our model performing competitively in this metric.

Such models may ultimately be useful to perform experiments in an iterative fashion, and help specify actionable hypothesis spaces for more targeted experiments down the line. Our work falls into a long line of literature on disentanglement and more recently the Sparse Mechanism Shift hypothesis related to causality, and we believe that the specific setup of SAMS-V AE will be useful in practical scenarios while being quantitatively performant across relevant tasks.

The deliberately generic assumptions we make about perturbations pose opportunities for future inquiry into more detailed aspects of such models. In specific cases, we may have prior knowledge about the nature of perturbations and their effects on the system we observe. An interesting future direction is posed in studying how perturbations may interact and compose in more complex fashion, and incorporating different forms of prior knowledge into such systems, while maintaining the ability of the system to discover knowledge and factors of variations that can be used downstream.

References

Y. Burda, R. Grosse, and R. Salakhutdinov. Importance weighted autoencoders. arXiv preprint arXiv:1509.00519 , 2015.

I. Guttman. The use of the concept of a future observation in goodness-of-fit problems. Journal of the Royal Statistical Society: Series B (Methodological) , 29(1):83-100, 1967.

I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International conference on learning representations , 2017.

M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley. Stochastic variational inference. Journal of Machine Learning Research , 2013.

E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144 , 2016.

Y. Ji, M. Lotfollahi, F. A. Wolf, and F. J. Theis. Machine learning for perturbational singlecell omics. Cell Systems , 12(6):522-537, 2021. ISSN 2405-4712. doi: https://doi.org/ 10.1016/j.cels.2021.05.016. URL https://www.sciencedirect.com/science/article/ pii/S2405471221002027 .

T. Karaletsos, S. Belongie, and G. Rätsch. Bayesian representation learning with oracle constraints. 4th International Conference on Learning Representations (ICLR) , 2016.

D. P. Kingma and M. Welling. Stochastic gradient vb and the variational auto-encoder. In Second international conference on learning representations, ICLR , volume 19, page 121, 2014.

S. Lachapelle, P. Rodriguez, Y. Sharma, K. E. Everett, R. Le Priol, A. Lacoste, and S. LacosteJulien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica. In Conference on Causal Learning and Reasoning , pages 428-484. PMLR, 2022.

F. Locatello, S. Bauer, M. Lucic, G. Raetsch, S. Gelly, B. Schölkopf, and O. Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning , pages 4114-4124. PMLR, 2019.

R. Lopez, J. Regier, M. B. Cole, M. I. Jordan, and N. Yosef. Deep generative modeling for single-cell transcriptomics. Nature Methods , 2018. doi: 10.1038/s41592-018-0229-2.

R. Lopez, N. Tagasovska, S. Ra, K. Cho, J. Pritchard, and A. Regev. Learning causal representations of single cells via sparse mechanism shift modeling. In Conference on Causal Learning and Reasoning , pages 662-691. PMLR, 2023.

M. Lotfollahi, A. K. Susmelj, C. De Donno, Y. Ji, I. L. Ibarra, F. A. Wolf, N. Yakubova, F. J. Theis, and D. Lopez-Paz. Learning interpretable cellular responses to complex perturbations in high-throughput screens. Molecular Systems Biology , 2023. doi: 10.1101/2021.04.14.439903. URL https://www.biorxiv.org/content/early/2021/05/18/2021.04.14.439903 .

G. E. Moran, D. M. Blei, and R. Ranganath. Holdout predictive checks for bayesian model criticism. Journal of the Royal Statistical Society Series B: Statistical Methodology , page qkad105, 2023.

T. M. Norman, M. A. Horlbeck, J. M. Replogle, A. Y. Ge, A. Xu, M. Jost, L. A. Gilbert, and J. S. Weissman. Exploring genetic interaction manifolds constructed from rich singlecell phenotypes. Science , 365(6455):786-793, 2019. doi: 10.1126/science.aax4438. URL https://www.science.org/doi/abs/10.1126/science.aax4438 .

F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011.

J. M. Replogle, R. A. Saunders, A. N. Pogson, J. A. Hussmann, A. Lenail, A. Guna, L. Mascibroda, E. J. Wagner, K. Adelman, G. Lithwick-Yanai, N. Iremadze, F. Oberstrass, D. Lipson, J. L. Bonnar, M. Jost, T. M. Norman, and J. S. Weissman. Mapping informationrich genotype-phenotype landscapes with genome-scale perturb-seq. Cell , 185(14):25592575.e28, 2022. ISSN 0092-8674. doi: https://doi.org/10.1016/j.cell.2022.05.013. URL https://www.sciencedirect.com/science/article/pii/S0092867422005979 .

Y. Roohani, K. Huang, and J. Leskovec. Predicting transcriptional outcomes of novel multigene perturbations with gears. Nature Biotechnology , 2023. doi: 10.1038/s41587-023-01905-6.

D. B. Rubin. Bayesianly justifiable and relevant frequency calculations for the applied statistician. The Annals of Statistics , pages 1151-1172, 1984.

B. Schölkopf, F. Locatello, S. Bauer, N. R. Ke, N. Kalchbrenner, A. Goyal, and Y. Bengio. Toward causal representation learning. Proceedings of the IEEE , 109(5):612-634, 2021.

K. Sohn, H. Lee, and X. Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems , 28, 2015.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a generative model for single cell RNA sequence data that models perturbations as sparse linear offsets in latent space. They derive a variational inference algorithm for their model and evaluate it on a couple of perturb-seq datasets where the perturbations are CRISPR knockouts of genes. They find that the inferred latent mask can be used for predicting known biological pathways, and they get some improvement in out-of-distribution prediction of gene expression.\n\nSTRENGTHS:\n* The paper does a good job of articulating appropriate assumptions for expression data.\n* The empirical results look good relatively to a variant on Lotfollahi et al.'s CPA and Lopez et al.'s sVAE+ in the sense that they show that you can enforce sparsity while maintaining or slightly improving likelihood (or at least an importance weighted bound on likelihood).\n\nWEAKNESSES:\nThis paper is tricky to evaluate because it leans heavily on the sparsity story (because otherwise it is a very small modification of Lotfollahi et al.'s CPA), so they motivate the work by connecting it to ideas from causal representation learning (sparse mechanism shift hypothesis) where the goal is to recover latent variables. But I think that's a risky strategy: making a model sparse does not automatically make it causal - especially when there's no reason to believe that the sparsity corresponds to how you would expect latent variables to behave. Given cascading nonlinear feedback of cell biology, I don't think anyone seriously thinks that the underlying biological mechanisms are well described by sparse linear offsets. But, from the empirical results, it does seem like sparsity is a useful way of constraining the model's capacity, and it provides the added benefit of making the model more interpretable.\n\nI think to properly evaluate the utility of sparsity, I would want to see how a biologist could use the output of this model to discover a pathway. I know that's what the annotated pathway experiments were meant to illustrate, but they used labeled data to train the random forest, and when discovering a novel pathway, you wouldn't have these labels. The plots in figure 2 suggest that this may be possible (there's a lot of shared structure within pathways), but I'm not sure how much of that structure is a function of the fact that the probabilities are order by pathway. Unfortunately, I think doing this kind of an evaluation really well would probably turn this into a paper better suited to a bio venue because they could better evaluate how realistic the proposed procedure is.\n\nQUESTIONS:\n* What do the clusters in the hierarchical clustering correspond to? Do they have a biological interpretation? And if so, could this method be used for discovery via clustering?\n\nLIMITATIONS:\nGiven the causal framing of the introduction, I would have liked to see some more discussion about what biological processes / features the authors believe the latent variables should correspond to, and what simplifying assumptions their model makes with reference to these beliefs.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents an extension of the VAE model by incorporating sparse perturbations of the latent variables, based on the sparse mechanism shift assumption. The proposed method models the latent state of perturbed samples as a combination of a local latent variable capturing sample-specific variation and sparse global variables of latent intervention. These global variables are modeled as the pointwise product between latent perturbation variables and a binary mask. The approach is evaluated using a single cell sequencing dataset to assess its generalization capability under perturbation and its ability to recover latent variations.\n\nSTRENGTHS:\n- The approach exhibits a good effectiveness and outperforms baselines such as SVAE+ significantly.\n- The paper is written in a overall clear manner.\n- The paper presents an original combination of previous models, such as the CPA-VAE and SVAE+ models.\n\nWEAKNESSES:\n- The major weakness in my opinion is that is not entirely clear what are the benefits of the proposed model to the  CPA-VAE  baseline. To elaborate the performance differences compared to the CPA-VAE model appear to be somewhat insignificant in all experiments. It would be beneficial to clearly state the major benefits of employing the proposed model over the CPA-VAE model. Additionally, it would be useful to perform experiments that demonstrate the advantages of the proposed model over CPA-VAE.\n- The experimental section requires some polishing, such as referencing Figure 2 and Figure 3 in the text. Furthermore, the \"todo\" annotations in the Figure 4 caption should be removed. \n- Accessibility could be improved:  it would be useful to include ( e.g. in the appendix) an introduction and more in depth explanations to the biological tasks addressed in order to make the paper self contained and accessible to a non biology experts readers.\n\n*Minor:*\n\n I noticed a few typos:\n\nLine 91: \"functiuon\" should be \"function.\"\n\nLine 105: \"that aim\" should be \"that aims.\"\n\nQUESTIONS:\n- Could the authors provide further elaboration on the rationale behind choosing the additive structure of perturbations in the latent space? Are there alternative choices that could be considered?\n- What motivated the preference for the importance weighted ELBO as an evaluation metric over the standard ELBO measure?\n- What are other practical problems or applications that could be addressed by the proposed model? If the model is thought mainly for solving single cell predicition problems, wouldn't it better to incorporate domain knowledge in it?\n\nLIMITATIONS:\nThe manuscript addresses the limitations adequately.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new VAE-based model for learning disentangled representation, dubbed Sparse Additive Mechanism Shift Variational Autoencoder (SAMS-VAE). The model is specifically designed for datasets under perturbation; the dataset has the form of $(x_i,d_i)$ where $d_i \\in \\{0,1\\}^T$ indicates the perturbation for sample $i$. Though the generative model can be general, here the authors design it for a specific application which is scRNA-sequential data. The model consistent of a set of global variables that correspond to the perturbations as well as their sparsity as well as a set of local variables correspond to noise for each sample. The authors then perform a set of experiments on the genome-wide CRISPR interference (CRISPRi) perturb-seq dataset from Replogle et al and the one from Norman et al. The authors then evaluate their model based on loglikelihood, interpretability of the latent space (based on productiveness of annotated biological pathways), and generalization to unseen data\n\nSTRENGTHS:\n- The paper was easy to read for the most part. The training details as well the model design is relatively clear. I applaud the authors for making the paper as reproducible as possible.\n\n- Though scRNA-sequential data is not my area, I feel that it is an understudied subject for generative models that could benefit from more research and it is something that will be of interest to the NeurIPS community. \n\n- To the best of my knowledge , the paper contains no technical flaws.\n\nWEAKNESSES:\n- While the language of the paper is clear, I don't think the paper is currently self-contained particularly with regards to dismantlement. There are many definitions and notions of disentanglement out there (e.g. independence, equivariance, permutation, linear disentanglement, ...). If you are basing yours on sparse mechanism shift, then it would be helpful to provide some background where you clearly define these concepts rather than immediately jumping into defining your model. \n\n- Relating to the previous point, I missed the motivation for some of the design choices. For example, why use addition to combine the latent? Should the sparsity variable $m_t$ be independent for each dimension? \n\n- Some parts of the results is not very clear for me to draw conclusions from. Neither Figure 2 or 3 are discussed in the main paper. What are the variables in Figure 2? Is the top $e_t$? Why the text says there are 722 perturbation but the figure only shows around 350? At the bottom, the mask is applied to what exactly? Also based on this figure, I see no qualitative difference between SAMS-VAE and SVAE+. As for figure 3, you should also show umap of raw the data and the baselines for comparisons. \n\n- The title for this paper is currently abit too general. This a very specific model that targets perturbation datasets and is not a general objective that can be applied to any VAE. I would consider changing the title. \n\n- Please make sure that IWELBO equation that is not outside the margins on page 4\n\nQUESTIONS:\nSee Weaknesses\n\nLIMITATIONS:\nTo the best of my knowledge, this paper poses no potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new method, SAMS-VAE for learning disentangled representations leveraging the sparsity mechanism assumption. SAMS-VAE is bridging the gap between two existing models - CPA and SVAE+, hence resembles both. In simulated and real (single-cell) data experiments, the authors claim improvement over baselines with regards to ood prediction and interpretability of the latent space.\n\nSTRENGTHS:\n- Originality: SAMS-VAE fills the gap between two AE-based models, CPA which includes learning basal state of a cell without specifying a likelihood model over the latents, and SVAE+ which has a learnable model over the latents but no flexibility for representing basal states. Having both, makes the contribution of SAMS-VAE quite nice.\n\n- Quality: The paper includes both technical and empirical results which with (some exceptions, see below) seem solid.\n\n- Clarity: The manuscript is well written, although dense at places and the reader could benefit from more details in the supplementary.\n\n- Significance: SAMS-VAE can be a quite helpful, relevant tool for the genomics community.\n\nWEAKNESSES:\n1. The title reads too broad, the model and method proposed indeed targets single cell applications, and without further extensions, and elaborations I don't see it as a general method that can extend to other data types. The authors should clarify this.\n2. It would be helpful if the authors included the derivation of Eq 4 , 5.\n3. I may have missed it, but where are the results for significantly improved ability to recover factors predictive of known molecular pathways, as stated in the introduction.\n\nQUESTIONS:\nIn addition to my notes stated in the section above, it would help if the authors can answer the following:\n1. There is no UMAPs for the latent representation of the baselines, could you include that for supporting the statement about improved interpretability?\n2. What is the configuration you use for the simulated study? [12] proposes different setups (latent dimensions, strength of causal effect, sparsity level), which one in particular did you use? Also, why are other metrics avoided, especially the disentanglement ones R2, MCC?\n3. Are there any identifiability guarantees of your model? In particular, how does your model cover modeling a discrete case such as gene expression data.\n4. Is comparison wrt MSE or absolute error possible for the OOD setup, IWELBO is a bit limiting for comparison in applications  and also model specific.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-24"}