Modelling Cellular Perturbations with the Sparse Additive Mechanism Shift Variational Autoencoder
Authors: Michael Bereket
Theofanis Karaletsos
Venue: NeurIPS 2023 poster (2023)
Keywords: Disentagled representation learning
VAE
generative models
sparse mechanism shift
perturbation modeling
cellular modeling
TL;DR: We propose a new VAE structure, SAMS-VAE, which captures explicit perturbation latent variables and sparsity effects to model sparse mechanism shift in latent space for applications such as compositional perturbation modeling of cells.

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
Michael Bereket
Theofanis Karaletsos
insitro* mbereket@@stanford.cco
insitro
choofanis@@kraleetsss.co
Abstract
Generative models of observations under interventions have been vibrant topic of interest across machine learning and the sciences in recent years. For example, drug discovery, there is need to model the effects of diverse interventions cells in order to characterize unknown biological mechanisms of action. We propose the Sparse Additive Mechanism Shift Variational Autoencoder, SAMS- VAE, combine compositionality, disentanglement, and interpretability for perturbation models. SAMS- models the latent state of 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-sseeific latent subspaces that are flexibly composable. We evaluate SAMS-VAL both quantitatively and qualitatively on range of tasks using two popular single cell sequencing datasets. In order to measure perturbation-speciie we also introduce framework for evaluation of perturbation models based on average treatment effects with links to posterior predictive checks. SAMS- -AEE-AA VAE outperforms comparable models in terms of generalization across in-distribution and out-of-distribution tasks, including a combinatorial reasoning task under re- source paucity, and yields interpretable latent structures which correlate strongly to known biological mechanisms. Our results suggest SAMS-VAL is an interesting addition to the modeling toolkit for machine learning driven scientific discovery.
model-properties,
Introduction
Scientific discovery often involves observation and intervention 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 mechanisms, their regulatory dependencies, and their responses to drugs. As technology enables both richer and finer grained measurements of these systems, there an increasing need and opportunity for machine learning methods to help generate predictive insights of growing complexity.
approach investigate biological
Generative models such variational auto-encoders (VAEs) are commonly used to learn represen- tations of complex datasets and their underlying distributions. A common goal in generative modeling disentanglement, whereby latent structures should factorize into semantic subspaces to facilitate generalization and discovery. A desirable outcome consists of these subspaces learned by models be- ing indicative of latent mechanisms, while sparsely varying according to the underlying latent factors of variation in the true data distribution [TO]. 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 the Compositional Perturbation Autoencoder [13] and SVAE+ [12) have successfully applied disentangled deep learning scientific problems in single-cell RNA-sequencing under perturbation.
Reseearch supporting this publication conducted while authors were emplayed at insitro
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this work, we propose the Sparse Additive Mechanism Shift Variational Autoencoder (SAMS- VAE), 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-specife latent effects, 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 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 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-distribttion generalization, including combinatorial generalization when multiple perturbations are applied. We furthermore examine the interpretability of the model's 's disentangled structures and demonstrate significantly improved ability 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 (æ;; d of observations RR and perturbation dosage vectors d, E {0, 13", where di,j is 1 if sample i received perturbation j and O otherwise. We aim to develop generative models of p( d)), representing the distribution of features of 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- -VEE).
2.1 Generative model
We consider generative models with the following basic structure:
2, 2 + x
ax ~ p ( Zz;
2; D. is the latent state embedding for sample i, which is modeled the sum of a latent basal state embedding 2 E R= and a latent perturbation effect embedding 2 E DR Observations are then sampled from conditional likelihood p(2; z;; 0). In this paper, focus likelihoods for p(z,|2,;;0 0), where parameters are computed from x, using a neural network with parameters O.
The core modeling assumption of SAMS- VAE relates the distribution p(=]'d,) [d.). We propose to model perturbations as inducing sparse latent offsets that compose additively as follows:
2 =
da ④ m.),
where e, E Rx and m, E P) are global latent variables that determine the latent offset due to perturbation L. m, is binary mask that performs feature selection the latent offset e,: when m, is sparse e, @ m. will result in sparse offset. Importantly, global variables e; and m, are shared across all samples (corresponding to cells) that receive perturbation t.
We specify the prior distributions p(ez) ~ NO, B1) and p(m)) Bern(a) for perturbation effects, where a 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, as an easy plug-in in replacement. We omit additional prior assumptions regarding the structure of perturbation effects in this work. We specify the prior distribution pl p(z²) NO0 /) for latent basal states.
Using this latent structure, we define the full generative model for SAMS- VAE as Figure The joint probability distribution over observed and latent variables is defined as:

---

Algorithm SAMS-VAE generative process
Ct
m
Require: X E &RNDD= D E {0, 1 for /. from 1 to T do
T
~ (N,
my Bern(a))
dot
end for for from 1 to N do
2 ~ N0, I)
-1=1 -i=ddi di o m.)
X
2
0; p 2z;
N
end for
Figure 1: SAMS- represented as generative process (left) and graphical model (right).
T p Z" M,,, = I1 p(e.)(mm.)
I Hr<z))(a de M, E;
(2)
for observations NxD=, perturbation dosages (0. (0,1)xx:: latent basal states
2 NND: latent perturbation embeddings E RTx,, and latent perturbation masks
M E (0.1)TxD.
2.2 Likelihood choice for seRNA-Seq Data
In the previous section, used generic form p(=,|z,; 0) for the observation model over to show the generality of the approach. Below, we describe the observation model we use to apply SAMS-VAE single cell RNA-sequencing data (scRNA-seq) in more detail.
We represent scRNA-seq observations as æ; E NP=, where each value Vijj is the number of measured transcripts in cell that correspond gene and follow the likelihood introduced by Lopez al. 12 to model the elaborate noise structure of scRNA-sec data. An additional utility quantity library size 4, the total number of transcripts measured cell 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 then defined as follows:
P: = folzi)
A ~ l'pple 0a)
02 Poisson Ax);
where p; E [0, P represents the expected frequency of each transcript cell and is parameterized by fe. neural network with softmax output. Observations are then sampled from a Gamma-Poisson distribution (equivalently, a negative binomial distribution) with mean p.l, TR and inverse dispersion 0a BBD. 0₂ is learned parameter that is shared across cells.
2.3 Inference
We perform inference on SAMS-VAE using stochastic variational inference g to approximate the marginal likelihood log p( X D) by optimizing parameters Φ and 0. We do so by maximizing the evidence lower bound (ELBO) for SAMS-VAE defined follows:
P(X. 2d D
ELBO 0) Ezz M.P.E.MM( log ( M, M
(3)
3

---

A key question when performing variational inference is the choice of variational family to approxi- mate the posterior distribution. As a baseline inference strategy, we consider the following amortized mean-field inference scheme for SAMS-VAE:
T
NN [ [(
(Z', E D = I (f (ammm))asssss Ø))
We parameterize q(m;; ()) Bern(p:) and q(eri)) (N(r;; σ1) with learnable parameters Pt, fi, o₁. We define <(=,(=;)= = (()eno(x:)) where fene is a learnable neural network that predicts mean and standard deviation parameters. During training, gradients are computed for q(m;; o) with a Gumbel-Softmal straight-through estimator [5]]
We propose two improvements to the mean- -feld inference scheme that aim to more faithfully invert the SAMS-VAS VAE generative model. First, we model possible correlations between sample latent basal states x and the global latent perturbation masks and embeddings (correlated encoder). We do so by implementing q(=]] z;; d,, E, M) [((fene((a; zf())) for z defined in equation 1, where Senc is neural network that takes as input both the observations and the estimated latent perturbation effect embeddings for given sample. Second, we model possible correlations between the latent perturbation masks and embeddings by replacing g(e,) with ((e,)m.) (correlated embeddings).
We implement q( (m+) = ( (m, 1)) with a learnable neural network femb that predicts the embedding from mask and one-hot encoding of the treatment index. Applying both of these modifications, we define the correlated variational family for SAMS-VAB as:
O M,,,,, II amm ()q(er Φ)
[[])aa, di, E,,
(5)
This richer variational family posits joint infinite mixture variational distribution between the global and local variables in the model 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.T
2.4 CPA-VAE
To directly assess the effect of the sparsity inducing masks, we define an ablated model, CPA-VAE, that is identical to SAMS-VAL VAE with all mask components fixed to 1. Thus, in contrast to equation
1, we have that 2 = Ei de. We call this model CPA-VAR because it directly incorporates the additive latent composition assumption from CPA [3], and CPA-VAE can be thought of as an extension of CPA to 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 quantitative strategies 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 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) II, primary evaluation metric (a similar metric was used in Lopez et al. [2] Specifically, we estimate Ix 0) on held out data, where denotes decoder parameters. Let H (Z', M, EE) represent the set of latent variables for SAMS- -AE. Then we can write the importance weighted ELBO with K particles as:
p(X, H D H D.D
IWELBO(X D, Φ. 0) E H(0) H.X, xD) log

---

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 marginal likelihood cannot be computed for models that are not fully specified as probabilistic models. In practice, we estimate IWELBO as follows:
IWELBO D. Φ EHO ....,.KK H log
Wk:
(6)
for
(M E (MM() E
N
b(R) p(z =
M (k), E(h), d. (2; b(R)
II
M E(k) X d +)
3.2 A Posterior Predictive Check for Average Treatment Effects of Perturbation Models
we query test statistics of interest in the predictive distribution of learned models and compare these statistics against estimates from the data. These types assessments can be useful when critiquing models for specific cases, such predicting the mean of some measurement under different perturbations. However, these assessments only characterize narrow aspects of the predictive distribution, providing less complete assessment than the marginal likelihood.
As a second category of metrics, we consider posterior predictive checks (PPC) [2] 19]:
As test statistic for our PPC we choose the population average treatment effect of a perturbation relative to control perturbationn each measurement Tı, for sample and gene j, given given
ATE = EiED [Iris/do(d")) (f)(dd)) [mx, dd(())]].
We define the average treatment effect for SAMS-VAL VAE ATE (dr(dd[DD,,, for an applied treatment d' and conditioning data D, (the training data) as the difference between the expected predictive value of output variable Viij given treatment d' and the expected predictive value of zi.j given control treatment do:
ATESAMS-VAEdd" (d" Dm. -P(E))D(E,MIID.. [2i 72],
with Ti := Ep(x.., do(d"). 2.2.. M,,, [Z and Ta := B p(r.., :; [xs.g]-
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 treatment, since observation noise is marginalized out. We also marginalize over the prior basal state p(z") the outer expectation, which simulates populations of different cells varying by natural variation. In practice we draw K samples = p(z") and E,, M2 p(E, M Dm)) for the outer expectation and evaluate the inner expectations by 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 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 dataset D; over which the expectation is computed (where DÅ denotes the subset of the dataset Ds under condition d):
DE (E Ds) E Dg (Ii (do(d")) E. D, (2i.) do(dn))
5

---

To create the metric ATE-Pearson r r(ATESAMS-VAL (d' Dm)) DEDat (d) we compute the Pear- son correlation coefficient r between differential expression estimates from data DEpata and our model-based estimator ATFSAMS VAE. across all features (commonly genes indexed by j).
This also reveals the relationship to PPC (AATESSMMS Dm), DEpata D (ss) Ds). 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 dataset D3 for estimating differential expression may be the same different, depending the use case. We note that the utilization of 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 series of applications perturb-seq datasets. Perturb-seq is a type of biological experiment in which cells are individually perturbed and subsequently profiled with single cell RNA sequeneing (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 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-seg datasets gene expression matrix X E N and a perturbation dosage matrix
{0, N for N cells, D2 gene transcripts, and 7 perturbations. Entries Xi, represent the number of transcripts from gene j observed in cell i and entries T, represent whether cell received perturbation We compare SAMS- VAE against baseline models based on their ability to model the distribution of perturb-seq data, generalize 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-VAB can be thought of as extension of CPA 13 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 config- urations, available at
setm/saamsaaaaa vae
4.1 Generalization under individual perturbations
Dataset To assess model generalization to held out samples under individual perturbations, we analyze subset of the genome- wide CRISPR interference (CRISPRi) perturb-seq dataset from Replogle et al. [17. which we call replogle-filtered filteedd CRISPRi is a type of genetic perturbation that represses the expression of selected target genes. Following the preprocessing steps from Lopez 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- fi teree 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.
et al. [12), replogle- lilggeeffeeereddd
Model
Inference
Test IWELBO
Mask PW. Acc.
ATErPearson
Conditional VAE
amortized MF
- 1766.10=4 + 0.18
0.765
SVAE+
amortized MF
1761.42 0.06
0.78 0.04
0.605
CPA- VAE
amortized MF
- 1760.14 = 0.20
0.523
CPA-VAE
corr. Zbasal
1756.57 + 0.14
0.571
SAMS-VAL
amortized MF
1757.72 = 0.14
0.68 0.09
0.302
SAMS- VAM
COTт. E
1758.08 0.07
0.71 + 0.04
0.319
SAMS-VAL VAE
COTT. Zbasal
1756.40 0.06
0.87 0.02
0.718
SAMS-VAL
coTr. 2basal and E 1756.27 + 0.10
0.89 + 0.03
0.765
Table I: Quantitative evaluation of treatment effects on Replogle filtered dataset (100 latent dimen-
sions) using K 10.000 samples. We find that inference strategies utilizing correlated variational families lead to better quantitative results, and that ATE and Mask Recovery are correlated.

---

SAMS-VAE
SVAE+
CPA-VAE
Embeddings
50
50
50
Perturbations (Guides)
100
150
150
200
200
200
300
300
300
100
80
Mask Probabilities
Latent Dimensions C/6
Latent Dimensions
Latent Dimensions
1.0
Annotated Pathways
Exosome
Mediator Complex
50
50
MT Protein Translocation
Perturbation (Guldes)
Nucleotide Excision Repair
39S Ribosomal Subunit
150
150
200
40S Ribosomal Subunit
250
60S Ribosomal Subunit
300
200
Spliceosome
Latent Dimensions
Latent Dimensions
Figure 2: Visualization of inferred latent perturbation masks and embedding means for the best performing checkpoint of each model in gepllgle-filteree We visualize the latent variables for the 345 perturbations with pathway annotations from Replogle et al. [17] and group by pathway. The SAMS- MAS- and CPA-VAL VAE models were trained with proposed correlated inference strategy.
Evaluation protocol Each model is trained with 100 dimensional latent space and MLP encoders and decoders with single hidden layer of dimension 400 (see Section A.4 for full training details). Based validation performance and sparsity, Beta(1, 2) prior selected for the SVAE+ mask, and Bern Bern(0.001, (0.001) prior was selected for SAMS- VAE, For each model type, we compute the test set importance weighted ELBO as described Section 3. 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 Comparing first between model types, we observe that SAMS- BAMS-VAE with fully correlated inference achieves the best test IWELBO and average treatment effect correlation. Interestingly, CPA-VAL PRAE with correlated inference achieves strong test IWELBO performance but falls behind on average treatment effect prediction, while conditional VAE 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, perform an ablation of SAMS-VAR and CPA- VAE inference strategies. We find that the correlated Abasal strategy yields substantial improvements performance for both SAMS- VAE and CPA-VAE while the correlated E strategy improvements are minor.
4.1.1 Recovery of biological mechanisms based disentangled factors
Evaluation protocol We assess the degree to which the pattern of perturbation effects on inferred latent factors in the SAMS-VAY and SVAE+ models from the previous section are predictive of

---

per intervention per gene
Differential Exoression per intervention per gene
10
10
Perturbettors
10
10
-00-1
250
300
-10"
-100
200
400
600
800
1000
200
400
600
800
1000
Genes
Figure 3: We visualize model-estimated treatment effects (ATESAMS-VAE) and data-estimated differ- ential expression (DEpas) for intervention-gene pairs in the Replogle experiment. We observe broad correlation (Pearson = 0.765): for example, perturbations ribosomal subunits influence on all expression broadly with matching directionality, while other guides exhibit more targeted effects.
known biological pathways as annotated by Replogle al. 117 (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 each model, we fit random forest model using scikit-learn [16] to predict pathway annotations from 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.
0.5. For
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- SVAE+, and CPA-VA in Figure 2 and visualize the SAMS- VAE estimated average treatment effects and estimated differential expression corresponding these perturbation effects in Figure Hierarchical clustering and UMAP projection of the inferred perturbation embeddings are presented Figure in Section A.6
Results We observe that the latent mask inferred by SAMS-VAL is more predictive of the annotated Filtered pathways in the replogle- ffltered dataset than that inferred by SVAE+. Additionally, we find that performing correlated inference on the perturbation embeddings improved pathway prediction performance for SAMS- VAF. Qualitatively, we observe that both SAMS-VAL and SVAF+ infer sparse masks with distinct patters between annotated pathways.
4.2 Modeling compositional interventions CRISPRa perturb-seq screen
Dataset We analyze the CRISPR activation (CRISPRa) perturb- sequurrb-equueeee screen from Norman et al. to assess how effectively SAMS- VAE and basclines model the effect of perturbation combinations. This screen was specifically designed with perturbations that have non-additive effects in combination, making this challenging setting for modeling combinations. We adopt the preprocessing from (6), which contains 105 unique targeting guides applied both on their own and 131 combinations. In total, the dataset contains 11,255 cells, each with 5,000 gene expression features.
15
Evaluation protocol We define two tasks using this data. The first, norman-ool ood, assesses the abil- ity 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 trained on cells that received single guide, along with [0, 25, 50. 75, 100]% of combinations. Held-out cells receiving the final 25% combinations are used to evaluate each model. We perform this analysis for random splits ciency, assesses how efficiently the
of the combinations. The second task, norman-data-ee models can learn combination phenotypes when trained on cells that have received single guide and increasing numbers of cells sampled uniformly across all combinations. Each model is evaluated

---

Out-of-Distriuttion Combinations
Combination Data Efficiency
Test combinations not observed training
combinations abserved in training
Combination uniformly subsampled from full training set
12-
10-
Relative IWELBO
8-
6-
4-
0.0
0.25
0.5
0.75
1.0
0.0
0.25
0.75
0.8
Avg. Trealment Effect Correlation
0.8
0.5
0.1
0.2
0.0
0.25
0.5
0.75
1.b
0.0
0.25
0.5
0.75
1.0
Fraction combinations in training set
Subsample fraction of cells receiving combinations in training set
SAMS-VAE
CPA-VAE
Conditional VAE
SVAE+
Figure 4: Results from norman- ood and norman data- effieiffiiiiooe experiments. Within splits, test IWELBO values are plotted relative to the test IWELBO for SAMS-VAL trained with 0 combina- tions on that split (relative IWELBO) to enable comparison across splits. SAMS- VAE and CPA-VAE models are trained with the correlated inference schemes described in methods.
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, SVAE+ is trained with a Beta(1, 2) prior and SAMS-VAL is trained with (0rn(0.01) prior.
Results Quantitative results are presented in Figure 4 with additional inference strategy ablations in Figure in Section A.6 SAMS- and CPA-VAB both achieve strong performance the norman-odd task across metrics, often within standard deviation of one another. Conditional VAE 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 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-VAL and CPA-VAE (and for encoding combinations defined in d, for conditional VAE).
In norman-data-en ciieicycicc....... 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 SAMS- VAE and CPA- 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. more comprehensive treatise and theoretical analysis of disentanglement was presented in [TO. Finally, the formal link to sparse mechanism shift and explicit causal disentanglement was also established recently in 19. 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 sequenc- ing data utilizing VAEs has been proposed in [I]] a model which inspired our use of their likelihood. Closer to our application on perturbed datasets is the CPA [33 Similar to this model, we adopt the idea to disentangle cellular latent spaces into basal and perturbation latent variables. However, pose the resulting model joint generative model with rigorous inference framework and, crucially, a sparsity mechanism to disentangle the perturbation effects into subspaces related to the affected mechanisms. The recent SVAE+ [2) 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 mechanism to compose perturbations, and in terms of inference strategies (SVAE+ learns perturbation representations by optimizing a prior). GEARS [18] leverages prior information of perturbation and gene features 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 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 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 novel sparse generative model, SAMS- 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 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-VAE 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.
A Acknowledgements
We acknowledge and thank insitro for funding this work.
10

---

References
[1] Y. Burda, R. Grosse, and R. Salakhutdinov. Importance weighted autoencoders. arXiv preprint arXiv: 1509.00599, 2015.
[2] Guttman. The use of the concept of future observation in goodness-of-fi problems. Journal of the Royal Statistical Society: Series B (Methodological), 29(1):83-100, 1967.
[3] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Ler- chner. beta-vae: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2017.
[4] M. D. Hoffman, D. M. Blei, C. Wang, and Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
[5] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax, arXiv preprint arXiv: 111111111111,,...111111111111111 2016.
[6] Y. M. Lotfollahi, cell omics. Cell Systems, 12(6):522 537, 2021. ISSN 2405-4712. doi: https///oi..or/
Wolf, and F. J. Theis. Machine learning for perturbational single-
10.1016ff.ells2221..00 URL wethttts//wwrt scwencedirect. com wwt.s/ciice//ratii...oooooooo p11/52406471122002220
[7] Karaletsos, Belongie, and G. Rätsch. Bayesian representation learning with oracle constraints. 4th International Conference on Learning Representationss (ICLR), 2016.
[8] D. P. Kingma and M. Welling. Stochastic gradient vb and the variational auto-encoder. In Second international conference learning representations, ICLR, volume 19, page 121, 2014.
[9] S. Lachapelle, P. Rodriguez, Y. Sharma, K. E. Everett, R. Le Priol, A. Lacoste, and S. Lacoste- Julien. Disentanglement via mechanism sparsity regularization: new principle for nonlinear ica. In Conference on Causal Learning and Reasoning, pages 428-484. PMLR, 2022.
[10] Locatello, S. Bauer, M. Lucic, G. Raetsch, S. Gelly, B. Schölkopf, and O. Bachem. Chal- lenging common assumptions in the unsupervised learning of disentangled representations, In international conference on machine learning, pages 41 4114-4124 PMLR, 2019.
[11] 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. 00.103//41592-088-o00 592-018-0029-22
[12) R. Lopez, N. Tagasovska, S. Ra, K. Cho, J. Pritchard, and A. Regev. Learning causal represen- tations of single cells via sparse mechanism shift modeling. Conference on Causal Learning and Reasoning, pages 662-691. PMLR, 2023.
[13] M. Lotfollahi, A. K. Susmelj, C. De Donno, 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.
URL https::/s///ww biorxiv. wwagg//ooeett//eee http://aaarr///1////// 18/ /0202222222222 04. 14. 439903
[14] 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.
[15] T. M. Norman, M. A. Horlbeck, J. M. Replogle, Y. Ge, Xu, M. Jost, L. A. Gilbert, and J. S. Weissman. Exploring genetic interaction manifolds constructed from rich single- cell phenotypes. Science. 365(6455):786 793, 2019. doi: 10. htt.///ceec...aa.////// https:///wwr http:://ww.sciinceeee
URL
science. ottp://////a///....///////ioooo. wwt.:..i1...e.eec... aax4438
[16] Pedregosa, G. Varoquaux, A. Gramfort, Michel, B. Thirion, Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn: Machine learning Python, Journal of Machine Learning Research, 12:2825-2834 2.1

---

[17] J. M. Replogle, R. A. Saunders, A. N. Pogson, J. A. Hussmann, A. Lenail, A. Guna, L. Masci- broda, E. J. Wagner, K. Adelman, G. Lithwick- Manni,, N. Iremadze, F. Oberstrass, Lip-
son, J. L. Bonnar, M. Jost, T. M. Norman, and J. S. Weissman. Mapping information- rich teaitype-phenttppp landscapes with genome-scale perturb-seq. Cell. 185(14):2559- 2575.e28, 2022. ISSN 0092-8674. doi: http://doo.ogggoooooloii.......d httpp://wwe
URL
[18] Y. Roohani, K. Huang, and J. Leskovee. Predicting transcriptional outcomes of novel multigene perturbations with gears. Nature Biotechnology, 2023. doi: 10.1008/441533--02--1955---
[19] D. B. Rubin. Bayesianly justifiable and relevant frequency calculations for the applied statisti- cian. The Annals of Statistics, pages 1151-1172, 1984.
[20] 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.
[21] K. Sohn, H. Lee, and Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015,
12


Reviews:
[{"review_id": "yJdoZ1hwkG", "reviewer": "Reviewer_9yRM", "summary": "This 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.", "questions": "In 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.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- 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. ", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive review and time spent on out paper.\n\n**Derivation of Eq.4,5:**\nEqs 4 and 5 indicate choices of variational families we design to perform inference in the model. As such, they are proposed by us and form a key part of our contribution in modeling the posterior of both SAMS-VAE and CPA-VAE significantly better than a naive variational inference approach would do, as indicated by our results ablating them in Table 1. We ablate CPA-VAE to isolate the sparsity contribution from SAMS-VAE from our inference improvements.\nIn CPA, SAMS-VAE, SVAE+, we have multiple variables interacting, and in previous works their approximate posteriors are handled as independent (i.e. mean field). We introduce increasing degrees of correlation in the approximate posterior (Eq. 4 being the simple mean field, 5 using correlated embeddings and correlated encoders, and an intermediate step we describe in the paragraph modeling q(e_t) without the dependence on m), showing that increasing their complexity yields significantly better results in modeling cells by allowing the different interacting variables to “coordinate”. \nSection 2.3 explains a big part of our ML contributions for these classes of models regarding how we perform inference more successfully.\n\n**Title:**\nWe thank the author for the comment. While our method is applied to single cell sequencing and a simulation taken from prior ML literature (see Appendix), we agree that we could easily add more specificity to the title to show we model cell perturbations. \nWe want to add that nothing prevents exactly this model from being applied also to microscopy images or any other readout of cells; We leave that to future work as the practicalities of engineering high dimensional readouts with a VAE require bespoke work in the decoder and likelihood for each modality here exceeding the scope of this paper, but the key mechanisms of SAMS-VAE would remain very much intact. If the authors have another idea for a common ML-Benchmarking system that corresponds to the setup of our model we’d be happy to describe it here, as well.\n\n\n**Results & recovery of molecular pathways:**\nIn Table 1, we summarize quantitative results for our model. On the rightmost column, we compare the ability of the model to disentangle cell embeddings according to whether dimensions correspond to a perturbation effect or not (versus being natural variation of a cell independent of perturbation). We then train classifiers on subsets of these and predict pathways which exist in the annotation of the underlying data, but the model is unaware of. We perform this experiment for both SVAE+ by Lopez et al 2023, and for various ablations of SAMS-VAE (with different inference strategies). The results show that SVAE+ gets an accuracy of 0.69 vs 0.9 for SAMS-VAE, indicating that the disentanglement corresponds much tighter to abstractions of pathways in embedding space. We hope this experiment makes sense to the reviewer.\nOne practical utility of being able to do this is that one can inspect which subspaces of variation (and corresponding decoded features when varying those dimensions) are for instance relevant to a particular perturbation, and when electing perturbations to correspond to disease-modulators, being able to drill down in more detail into involved mechanisms of action. \nThis is the task commonly performed in target discovery when performing high throughput screens, whereby looking for “signatures” of variation is the key task, and our model aids the discovery of more specific signatures as subspaces in an embedding containing many different effects relevant to the state of a cell/system.\n\n\n**Identifiability Guarantees:**\nIn previously published models there are only identifiability guarantees (i.e. in the SVAE paper) when the decoder is ensured to have specific Lipschitz-properties, they do not exist in full generality for this model class. The SVAE+ paper by Lopez et al. also does not provide them. This is why we explicitly go for the term “disentanglement” and not “causality” in our paper. \nWe believe that showing identifiability (and causality) is enormously difficult in such a general model class as a general VAE (with nonlinearities) without significant caveats and limitations, and indeed has not been done in any comparable papers we are aware of as of submission of this work. \nWe thus prioritize generalization and measure it via the loglikelihood in in-distribution data, out-of-distribution data, and by measuring recovery of pathways successfully compared to other models. Our results suggest that getting inference to work better gets us a long way towards better generalization and practical utility.\n\n**OOD MSE:**\nOur likelihood is over discrete count data, as commonly done in the probabilistic literature for genomic data following numerous papers. We are afraid that MSE is not compatible either with the preprocessing or assumptions in our data. \nWe consciously and laboriously elected to use the same likelihood used in SVAE+ over all our baselines to make IWELBO fully comparable in Table 1, including a carefully built generative version of CPA to make it fully comparable. \nHowever, to make models more generally comparable, we also perform an analysis in the appendix capturing average treatment effect, which models correlations of the predictive distribution of our model given a perturbation compared to differential expression. Our models do well in that metric, too, but we want to point out that this metric is flawed since differential expression and indeed absolute error severely underestimate aleatoric uncertainty in the count data, which is a huge factor in the present counts, and the probabilistic models are more sensitive to.\n\n**Umaps:**\nWe are happy to add UMAPs and visualizations for other relevant models in the appendix of an updated paper. We want to point out that we did add visualizations of the latent factors already in Figure 2.\n"}, {"review_id": "4FaondWu6C", "reviewer": "Reviewer_y8vh", "summary": "This 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", "questions": "See Weaknesses", "limitations": "To the best of my knowledge, this paper poses no potential negative societal impact.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive review and their detailed suggestions. We will address their feedback in the following:\n\n\n**Disentanglement/Self-Contained:**\nWhile we believe that we cite a rich set of papers explaining the setting we are based on, in particular the recent SVAE+ by Lopez et al which also models disentanglement in cells and was published in 2023, is inspired by Sparse Mechanism Shift, but underperforms our work, we are happy to take the feedback and will include more explanation over our assumptions and the setting in the paper. For now, we explicitly do not make strong causal statements or limiting assumptions of the nature of linear disentanglement etc., We propose a VAE-style model with a specific, novel structure and stick to a subspace inference task with a binary mask over the latent space per treatment (where each dimension is independent in the prior, but some of our inference methods correlate the dimensions in the posterior). \nWe model the existence of multiple treatment as additive perturbations, allowing superpositions of disentangled effects, which show quantitative value in our experiments in Sec. 3.2 .\n\n**Combinations:**\nIn some screens combinations of treatments exist, i.e. multiple CRISPR interventions changing the genetics, or multiple chemicals having been applied to a cell. \n\n**722 perturbations:**\nThank you for catching this. There are indeed 722 perturbations and we train the model with all of them. However, when we made this visualization we dropped unannotated treatments to make it more legible and easier to annotate/sort by pathway. Our quantitative results on IWELBO contain results over the full set, we just made the figure more compact for visualization purposes and for the pathway prediction task which is not testable for treatments without annotations. We will of course explain and clarify this better and apologize for the confusion. We assure you that no ill intent or hidden performance problems exist with this discrepancy, it is simply omitted in the figure explanation.\nWe are sharing the composition over the 722 here, showing that we filtered the 377 with label “none”. These labels are obtained from the replogle dataset preprocessing accompanying the SVAE+ paper.\n\nNone                          377;\nS40_RIBOSOMAL_UNIT            101;\nS60_RIBOSOMAL_UNIT             56;\nS39_RIBOSOMAL_UNIT             44;\nSPLICEOSOME                    38;\nMT_PROTEIN_TRANSLOCATION       37;\nMEDIATOR_COMPLEX               26;\nNUCLEOTIDE_EXCISION_REPAIR     23;\nEXOSOME                        20;\n\n**Mask application and assumptions:**\nThe mask is applied as the model shows element-wise to the treatment variable that is an inferred embedding before that is added to the basal state variable to model a cell.\nThe model takes the shape z = e_t * m_t + z^b , where each of these objects has the same dimensionality of the embedding space D_z. As such our model determines through m_t for each treatment which dimension of the cell embedding gets shifted and using e_t for each one that gets shifted separately by how much. This allows us to model natural variation in an object that does not depend on a perturbation separately in the unperturbed dimensions, while also allowing for multiple different mechanisms to exist in superposition since each one has its own index and can be applied.\nWe hope this also explains the additive nature of the model better: if these mechanisms are independent of each other we can apply them separately and we observe that our model outperforms every baseline in the combinatorial modeling task with that inductive bias, in particular the SVAE+ baseline which is the most recent paper, while the CPA-VAE version we compare to is imbued with our own inferential advances to perform better than its baseline version would and still underperforms the sparsified SAMS-VAE slightly.\n\n\n**Title too broad:**\nWe agree that we primarily consider perturbations on cells in this manuscript and a simulation, since this is a lucid scenario where we have data of the form p(y|d). We could model images, proteomics, sequencing, any variable imaginable for this problem with our proposed model.\nIf the reviewer is aware of any other scenario with public data of this very general structure we’d be happy to account for it. We want to note that the structure we propose could also be used for RL purposes with disentanglement over actions/interventions, but we do not aim to perform such a variant any time soon ourselves. As such, we’d be happy to adjust the title to reflect the focus on biology since this is a domain where reasoning over interventions makes sense, but we also want to note that our model has sufficiently general structure for a category of problems, way outside genomics.\n\n**IWELBO equation going over margins:**\nThank you for catching this, we are adjusting this with a linebreak and apologize for the visual disturbance, we fully agree this is aesthetically unnecessary and suboptimal.\n"}, {"review_id": "4JVrvvyKwg", "reviewer": "Reviewer_Aahc", "summary": "The 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.", "questions": "- 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", "limitations": "The manuscript addresses the limitations adequately.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and hope to convince them to improve their score:\n\n**Benefits compared to baselines:**\nWe apologize for the confusion, we must have not explained the results well enough in our experimental section. CPA-VAE is a model we propose to turn CPA into a generative model, imbue it with the same likelihood and make its quantitative evaluation comparable. \nAt baseline, similar to the structure of CPA, CPA-VAE uses mean-field amortized inference (standard VAE inference) and obtains a Test IWELBO of -1760.06. One of our core contributions is the inference machinery we invent in Section 2.3, and when we utilize correlated inference for CPA-VAE we can push it to -1756.72. That is based on improvements to the inference we made for SAMS-VAE, that we -for fairness- also throw into a model without sparsity so we can test the value of sparsity entirely separately, but already that performance did not previously exist in this field without us even looking into sparsity. For comparison, the most recent competitor SVAE+ (published in early 2023) obtains a score of -1759.23. \nWe then note that with the improved inference SAMS-VAE utilizing the extra sparsity still outperforms both on raw predictive power with an IWELBO of -1756.11, while adding a considerable amount of structure and interpretability to CPA-VAE. Separately, we evaluated the models that perform disentanglement in the pathway prediction task and see that compared to SVAE+ at 0.69 our model reaches 0.9 accuracy, which is a dramatic improvement on a task related to semantics and interpretability for a directly comparable model class while also improving predictive accuracy on the data.\nLastly, we evaluated our model on combinatorial tasks in Fig. 4, and show that SAMS-VAE outperforms CPA-VAE (both with the best possible inference strategy we propose, not at baseline for CPA-VAE to maximize the quality of the baseline) slightly, and SVAE+ soundly.\nWe again remind the reviewer that the best version of CPA-VAE utilizes a major contribution of ours by using correlated inference, and at baseline that model is significantly worse.\nWe hope this clarifies the convincing nature of our quantitative improvements, gains in interpretability, and overall superior performance of our contributions, and we will absolutely take a pass at the experimental section to make this more apparent and explain it better.\nWe thank the reviewer for sharing their confusion about that, we underestimated the need to explain the separate contributions in improved approximate inference and improved modeling by introducing sparsity and hope our explanation clarifies that the combination of both is what makes this performance so strong.\nIn the same spirit we thank the reviewer for the feedback on Fig. 2 and 4, and emphasize that we agree we will improve the explanations in the experimental section to make the results “pop” more as we hope we demonstrated they do in their present form and link to the Figures better and polish writing overall.\n\n**More biology background:**\nWe believe this is a fantastic suggestion and will follow it, the appendix offers plenty of space for such an explanation.\nWe will also share a few lines for a general rebuttal, explaining the setting of a screen in drug discovery as a core example where such a model would be directly applicable in practice, and think our readers would benefit from a tighter intuitive explanation of a few of these terms. \nWhile the core contribution of this work is the inference machinery and the new model utilizing sparsity, understanding the setting of our experiments with more clarity could be of further secondary utility to a broader audience that might be interested in further methodological advancements that might benefit such a setting and requires some more intuition to access it.\n\n**Additive structure in latent space:**\nIn some of our our experiments, in particular in experiment 3.2 using compositions of treatments, multiple perturbations are applied to cells at the same time. \nWhile baseline models such as SVAE+ treat combinations as a new “one-hot treatment” with no sharing with individual treatments, we can explicitly model each treatment and its masks and embeddings separately and establish composition via addition. An underlying assumption is here that addition is a good way to model composition, which in the case of independent perturbations may be the case, and in the case of interactions between treatments less so. In our discussion we note that we leave it for future work to study richer mechanisms in more detail for modeling composability, and also note that our proposed method already works better than other existing prior work we compare to and forms a sound baseline for further inquiry into this specific problem. Interesting dependent correlations could occur when considering drug cocktails with interacting effects while trying to optimize some fitness response.\n\n**IWELBO vs ELBO:**\nThe ELBO (evidence lower bound) is not a quantity that permits strict comparison between models. While one can say “model A is no worse than ELBO_A, and model B is no worse than ELBO_B”, the nature of a lower bound is that this can be arbitrarily far from the true performance of the model and gives no guarantees for relative rankings of actual performance. As such, in an effort to report approximations to the actual marginal likelihood of the data logp(y|d), we utilize importance weighting to obtain the IWELBO to report quantities which are much closer to a comparable quantity and approximate the marginal loglikelihood stochastically. This was also performed in the SVAE+ paper, and is common practice in variational inference papers. \n\n**Other problems vs domain knowledge:**\nThis is an excellent question, thank you for posing it. Since R4 posed a similar question, we will refer to it in the general overview, please read our response there, and please feel free to read our response to R4."}, {"review_id": "lwlv4h2Ojc", "reviewer": "Reviewer_TMYT", "summary": "This 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. ", "questions": "* 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?", "limitations": "Given 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. ", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* 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).", "weaknesses": "This 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their deep feedback and thoughtful response about the core of our proposed model. We believe we can address your concerns resoundingly and hope to convince you to raise your score. Please read our general response since it addresses your main concern as well.\n\n**Results vs. CPA and SVAE+:**\nWe thank the reviewer for pointing this out as a positive. However, as we discussed with R3, we believe our results are not just slightly better, but significantly so. CPA-VAE only is close to SAMS-VAE when we apply correlated inference strategies to that model, which is also one of our contributions from Sec 2.3. The baseline version of the model using mean-field amortized inference, which is a fair statement is the baseline, performs significantly worse as we show in Table 1. In Fig 4 SAMS-VAE also is the best performer, although it is compared to CPA-VAE with correlated inference we propose in this paper. We did not report mean field there to ablate the effect of the sparsity itself.\nWe put a lot of work into getting the interactions in the posterior between the basal state and the treatment embedding right, and later adding the mask to that and maintaining that performance. Out of the box these objects are worse behaved, in performance and interpretability as also the pathway prediction task in in Table 1 shows, and as such we would like to highlight our work on approximate inference in Sec 2.3 by proposing a suite of scalable improvements to the variational family for this model as an important contribution to this model class. \nSVAE+ is the most recent published comparator here as well and underperforms compared to SAMS-VAE which also models sparsity.\nWe thank you for pointing out that we managed to improve results compared to CPA-VAE and SVAE+ while incorporating the best inductive biases of both models, it absolutely benefits from the inference work.\n\n**Evaluation of Sparsity vs Causality:**\nWe also discuss this with R3 and will write a general overview, but we wanted to give you a personalized response here as well.\nWe believe this topic is probably the most important reason why we built this model and want to convince you of its relevance.\nFirst of all, we do not make claims about causality in our work. We explicitly call it disentanglement, where comparable papers including SVAE+ call their model causal, because we do not believe we can assure causality in a narrow sense. And we do not intend to.\nWe fully agree that the nonlinear and obscure nature of cell biology cannot be captured easily in a mask. Some works attempt to model it by regulatory network inference, which is a hard task and reduces to (causal) structure learning.\nWhat we are trying to achieve in our work is to provide an unbiased way to characterize mechanisms of action as subspaces in an embedding and disentangle them from each other and from the natural variation of a cell (i.e. cell-cycle etc.) that is not perturbation specific. \nOur model manages to capture a diverse set of perturbations, and we show that the masks it learns are highly predictive of pathways in Table 1 and Sec. 3.1.1. \nHowever, calling the masks pathways would be inaccurate, since some treatments in this model could easily be chemical in nature, and would not strictly correspond to one pathway.\nHowever, as we see these “mechanism” do have interesting structure:  when observing Fig. 2, we would highlight that it becomes clear that our model separately models activity of a perturbation on an embedding dimension from the direction/effect of that perturbation.\nFor instance, different ribosomal units have largely overlapping masks, interestingly shifting in some subspaces for each unit but sharing in some, while the embeddings for some shared subspaces have the opposite sign, for example guides 100-150 modeling S39 ribosomal unit have the opposite sign in treatment embedding 63, but a consistent activation of mask with other ribosomal units below them.\nThis is a neurips paper, so we do not emphasize biological interpretability here and wanted to highlight the ability of the model to perform strongly and to disentangle, ideally in rigorous quantifiable ways. Using such structures for insight generation for a biology journal is of course of interest to future work, but in this manuscript we wanted to establish its predictive power and the many details in inference and modeling necessary to achieve those.\nBut from a practical perspective, this model is designed to help study biology in a screen, so that specific effects can be surfaced and quantified. One can then also vary the specific subspaces belonging to a perturbation and observe what effects they have in observed space.\nMoreover, the model in its current form is applicable to any readout, be it images, transcriptomics, or whatever, given any type of perturbation, chemical or genetic or otherwise, and can point to subspaces of variation that uniquely capture the relevant bits of information.\nFor instance, we are interested in analyzing the CPJump dataset ( https://jump-cellpainting.broadinstitute.org/ ) which would require an unbiased model that does not strictly just capture pathways or other pre-specified prior knowledge, and may want to study the overlap of effects between genetic and chemical perturbations, which in our model are trivial to jointly incorporate without any changes.\nReturning to our manuscript, our pathway prediction task convincingly shows that the disentanglement in our paper is a lot better at recovering previously known pathways, but in practice we agree we would not care about that. However, in order to gain confidence in a model that we use for novel discovery as hinted at above, we would clearly want to evaluate it on positive controls such as known pathways to establish its utility. \n"}]

Author rebuttal:
We thank our 4 reviewers for their feedback and questions.

The reviewers agreed that the paper is sound, interesting, and mostly clear and well-written.
While we addressed many of the individual concerns, we also want to summarize some of the shared discussion.

(I) we received questions about how much stronger our results are compared to baselines.
(II) reviewers are worried about identifiability and causality and how we think about that.
(III) the semantics and biological relevance of the disentangling and out overall set-up are not clear to all the reviewers.

We hope to give a general response to all three that supplements your individual responses and that you will find your questions positively resolved.

**Our contributions:**
Quick summary of our contributions.
First, a minor contribution is that we propose a generative version of the CPA model and coin it CPA-VAE, imbuing it with a count-based likelihood to make it comparable to other works in this line and affording it a generative interpretation and evaluation.
Second, our key change is that we propose adding a disentangling mechanism to that model by adding a mask m_t (same dimension as the embeddings) per treatment variables which allows the model to perform variable selection for which dimensions are useful to perturb to model cells under perturbation p(y|d) where d indications a dose vector/indicator.
Third, in Sec 2.3 we propose 2 changes to the basic VAE inference machinery tailored to our models, by capturing correlations in the posterior between the treatment variable and the basal state (effectively allowing the model to explain away treatment influences relative to basal state); and correlations between the masks and the treatment variables, thus allowing the model to not spend capacity on "switching on"  perturbation for dimensions that do not carry relevant information. 

**Performance:**
We apply our improved inference machinery also to CPA-VAE in our experiments, and still show that SAMS-VAE outperforms both CPA-VAE and SVAE+, its most recent comparable model in the literature, in terms of modeling cells in-distribution, out-of distribution generalization in combinatorial settings with multiple treatments, and pathway recovery tasks based on the quality of the disentanglement.
We note that CPA-VAE mean field is the fairest baseline here, since improving the variational family is a contribution of our work, and as such believe Table 2 and Fig 4 show conclusively improved performance compared to prior work across the board.
In addition we share a simulation task inspired by previous papers in the appendix which shows that the model improves its empirical ability to recover/identify latent variables as its likelihood improves with better inference, and an evaluation framework using average treatment effect that is model agnostic and can be used across generative and non-generative models for cellular perturbations, also showing our models to perform very well.

**Causality/identifiability:**
We do not make claims about causality, but we are encouraged about the fact that our model generalizes better than its comparators in-distribution, out-of-distribution, and in a quantitative interpretation task where we infer previously known pathways.
We cannot guarantee causality, but raw predictive performance across all those by making a model more structured sounds quite promising to us as an intermediate step that is useful to inform further scientific inquiry.
Our simulation also shows that the model improves identification of latents as its NLLK improves with increasing sample-sizes available for training.

**Biological Interpretation:**
There are works in Genomics which utilize knowledge of specific pathways etc. in the model to improve performance. We want to explain why we do not consider that in this work and focus on more general ML.
In the setting we consider, which is a screen for cell responses given perturbations, we oftentimes may use chemicals as inputs or CRISPr interventions, and in the case of chemicals we do not know a priori which mechanisms they affect. In fact, the goal of a screen is to perform an unbiased study to surface the mechanisms of action affected by perturbations. It is useful to have an unbiased approach beyond known pathways and genetic perturbations, permitting for chemical perturbations and the existence of unknown mechanisms, as is the case in most practical screens.
Further, while the transcriptomics readout we have is a practical one for evaluating a paper given the existence of baseline work and comparable benchmarks in public, many biology screens also use other readouts given the relative expense of single cell sequencing and its lack of scalability. For instance, a popular readout is cell imaging, where cells are taken images of and used as readouts. An example of this is the CPJump dataset ( https://jump-cellpainting.broadinstitute.org/ ), which perturbs cells at scale using chemistry and genetic perturbations and visualizes cell-painting using microscopy. 
Our model would be directly applicable to such a scenario using modifications for images, which we leave as future work. 
We wrote this paper anchored on the experiments we elected because there is comparable work in ML that we can harden the modeling side against to set up the machinery to be rigorous, but we very consciously keep the model structure more general to account for broader applicability to modern quantitative biology.


We believe this work captures the best of CPA and SVAE+, adds inference tricks and quantitative rigor, and the resulting model shows a lot of promise for being useful in a variety of tasks strongly corresponding to the needs in screening in modern quantitative biology, while also being of general use to the community by showing what tricks were needed to make it work and how we can disentangle VAEs in more structured ways than the previous papers published in ML along the lines of our work.

Meta-review: Although this paper had overall borderline reviews initially, the author response and subsequent discussion period did well to clarify many issues, with one reviewer raising the score accordingly. Two other reviewers did not raise their scores, but argued strongly that this paper should be accepted during private discussion between reviewers and the AC.

As AC, I would note that the final revision of the paper should reflect the discussion with the reviewers during the rebuttal period, and include the proposed edits.
Decision: Accept (poster)