SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michael Bereket
insitro*
mbereket@stanford.edu
Theofanis Karaletsos
theofanis@karaletsos.com
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.
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 (VAEs) [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.
*Research supporting this publication conducted while authors were employed at insitro
In this work, we propose the Sparse Additive Mechanism Shift Variational Autoencoder (SAMS-VAE), 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  $(\pmb{x}_i, \pmb{d}_i)_{i=1}^N$  of observations  $\pmb{x}_i \in \mathbb{R}^{D_x}$  and perturbation dosage vectors  $\pmb{d}_i \in \{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(\pmb{x}|\pmb{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} = z _ {i} ^ {b} + z _ {i} ^ {p}
$$
$$
\boldsymbol {x} _ {i} \sim p (\boldsymbol {x} _ {i} | \boldsymbol {z} _ {i}; \boldsymbol {\theta})
$$
$z_{i} \in \mathbb{R}^{D_{z}}$  is the latent state embedding for sample  $i$ , which is modeled as the sum of a latent basal state embedding  $z_{i}^{b} \in \mathbb{R}^{D_{z}}$  and a latent perturbation effect embedding  $z_{i}^{p} \in \mathbb{R}^{D_{z}}$ . Observations are then sampled from a conditional likelihood  $p(\boldsymbol{x}_{i} | \boldsymbol{z}_{i}; \boldsymbol{\theta})$ . In this paper, we focus on likelihoods for  $p(\boldsymbol{x}_{i} | \boldsymbol{z}_{i}; \boldsymbol{\theta})$ , where parameters are computed from  $z_{i}$  using a neural network with parameters  $\boldsymbol{\theta}$ .
The core modeling assumption of SAMS-VAE relates to the distribution  $p(z_i^p | d_i)$ . We propose to model perturbations as inducing sparse latent offsets that compose additively as follows:
$$
\boldsymbol {z} _ {i} ^ {p} = \sum_ {t = 1} ^ {T} d _ {i, t} \left(\boldsymbol {e} _ {t} \odot \boldsymbol {m} _ {t}\right), \tag {1}
$$
where  $\boldsymbol{e}_t \in \mathbb{R}^{D_z}$  and  $\boldsymbol{m}_t \in \{0,1\}^{D_z}$  are global latent variables that determine the latent offset due to perturbation  $t$ .  $\boldsymbol{m}_t$  is a binary mask that performs feature selection on the latent offset  $\boldsymbol{e}_t$ : when  $\boldsymbol{m}_t$  is sparse  $\boldsymbol{e}_t \odot \boldsymbol{m}_t$  will result in a sparse offset. Importantly, global variables  $\boldsymbol{e}_t$  and  $\boldsymbol{m}_t$  are shared across all samples (corresponding to cells) that receive perturbation  $t$ .
We specify the prior distributions  $p(\boldsymbol{e}_t) \sim \mathcal{N}(0, \beta I)$  and  $p(\boldsymbol{m}_t) = \mathrm{Bern}(\alpha)$  for perturbation effects, where  $\alpha$  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  $\boldsymbol{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(\boldsymbol{z}_i^b) \sim \mathcal{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 \in \mathbb{R}^{N \times D_x}$ ,  $D \in \{0,1\}^{N \times T}$   
for  $t$  from 1 to  $T$  do  
 $e_t \sim \mathcal{N}(0,I)$ $m_t \sim \operatorname{Bern}(\alpha)$   
end for  
for  $i$  from 1 to  $N$  do  
 $z_i^b \sim \mathcal{N}(0,I)$ $z_i^p = \sum_{t=1}^{T} d_{i,t}(e_t \odot m_t)$ $z_i = z_i^b + z_i^p$ $x_i \sim p(x_i | z_i; \theta)$   
end for
Figure 1: SAMS-VAE represented as an generative process (left) and as a graphical model (right).
$$
p (\boldsymbol {X}, \boldsymbol {Z} ^ {b}, \boldsymbol {M}, \boldsymbol {E} | \boldsymbol {D}; \boldsymbol {\theta}) = \left[ \prod_ {t = 1} ^ {T} p (\boldsymbol {e} _ {t}) p (\boldsymbol {m} _ {t}) \right] \left[ \prod_ {i = 1} ^ {N} p (\boldsymbol {z} _ {i} ^ {b}) p (\boldsymbol {x} _ {i} | \boldsymbol {z} _ {i} ^ {b}, \boldsymbol {d} _ {i}, \boldsymbol {M}, \boldsymbol {E}; \boldsymbol {\theta}) \right], \tag {2}
$$
for observations  $\pmb{X} \in \mathbb{R}^{N \times D_x}$ , perturbation dosages  $\pmb{D} \in \{0,1\}^{N \times D_z}$ , latent basal states  $\pmb{Z}^b \in \mathbb{R}^{N \times D_z}$ , latent perturbation embeddings  $\pmb{E} \in \mathbb{R}^{T \times D_z}$ , and latent perturbation masks  $M \in \{0,1\}^{T \times D_z}$ .
2.2 Likelihood choice for scRNA-Seq Data
In the previous section, we used a generic form  $p(\boldsymbol{x}_i|\boldsymbol{z}_i;\boldsymbol{\theta})$  for the observation model over  $\boldsymbol{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  $\boldsymbol{x}_i \in \mathbb{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:
$$
\boldsymbol {\rho} _ {i} = f _ {\theta} \left(\boldsymbol {z} _ {i}\right)
$$
$$
\boldsymbol {\lambda} _ {i} \sim \Gamma (\boldsymbol {\rho} _ {i} \boldsymbol {l} _ {i}, \boldsymbol {\theta} _ {d})
$$
$$
\boldsymbol {x} _ {i} \sim \operatorname {P o i s s o n} \left(\lambda_ {i}\right),
$$
where  $\pmb{\rho}_i\in [0,1]^{D_x}$  represents the expected frequency of each transcript in cell  $i$  and is parameterized by  $f_{\theta}$ , a neural network with a softmax output. Observations are then sampled from a Gamma-Poisson distribution (equivalently, a negative binomial distribution) with mean  $\pmb{\rho}_i l_i\in \mathbb{R}_+^{D_x}$  and inverse dispersion  $\theta_d\in \mathbb{R}_+^{D_x}$ .  $\theta_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(\pmb{X}|\pmb{D})$  by optimizing parameters  $\phi$  and  $\theta$ . We do so by maximizing the evidence lower bound (ELBO) for SAMS-VAE, defined as follows:
$$
\operatorname {E L B O} (\phi , \theta) = \mathbb {E} _ {\boldsymbol {Z} ^ {b}, \boldsymbol {E}, M \sim q (\cdot | \boldsymbol {X}, \boldsymbol {D}; \phi)} \log \frac {p (\boldsymbol {X} , \boldsymbol {Z} ^ {b} , \boldsymbol {M} , \boldsymbol {E} | \boldsymbol {D} ; \theta)}{q (\boldsymbol {Z} ^ {b} , \boldsymbol {M} , \boldsymbol {E} | \boldsymbol {X} , \boldsymbol {D} ; \phi)}. \tag {3}
$$
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:
$$
q \left(\boldsymbol {Z} ^ {b}, \boldsymbol {M}, \boldsymbol {E} | \boldsymbol {X}, \boldsymbol {D}; \phi\right) = \left[ \prod_ {t = 1} ^ {T} q \left(\boldsymbol {m} _ {t}; \phi\right) q \left(\boldsymbol {e} _ {t}; \phi\right) \right] \left[ \prod_ {i = 1} ^ {N} q \left(\boldsymbol {z} _ {i} ^ {b} | \boldsymbol {x} _ {i}; \phi\right) \right]. \tag {4}
$$
We parameterize  $q(\boldsymbol{m}_t; \phi) = \operatorname{Bern}(\hat{\boldsymbol{p}}_t)$  and  $q(\boldsymbol{e}_t; \phi) = \mathcal{N}(\hat{\boldsymbol{\mu}}_t, \hat{\boldsymbol{\sigma}}_t)$  with learnable parameters  $\hat{\boldsymbol{p}}_t, \hat{\boldsymbol{\mu}}_t, \hat{\boldsymbol{\sigma}}_t$ . We define  $q(z_i^b | x_i) = \mathcal{N}(\hat{f}_{enc}(x_i))$ , where  $\hat{f}_{enc}$  is a learnable neural network that predicts mean and standard deviation parameters. During training, gradients are computed for  $q(\boldsymbol{m}_t; \phi)$  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  $\boldsymbol{z}_i^b$  and the global latent perturbation masks and embeddings (correlated encoder). We do so by implementing  $q(\boldsymbol{z}_i^b | \boldsymbol{x}_i, \boldsymbol{d}_i, \boldsymbol{E}, \boldsymbol{M}) = \mathcal{N}(\hat{f}_{enc}([\boldsymbol{x}_i \boldsymbol{z}_i^p]))$  for  $\boldsymbol{z}_i^p$  as defined in equation 1, where  $\hat{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(\boldsymbol{e}_t)$  with  $q(\boldsymbol{e}_t | \boldsymbol{m}_t)$  (correlated embeddings). We implement  $q(\boldsymbol{e}_t | \boldsymbol{m}_t) = \mathcal{N}(\hat{f}_{emb}(\boldsymbol{m}_t, t))$  with a learnable neural network  $\hat{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:
$$
q \left(\boldsymbol {Z} ^ {b}, \boldsymbol {M}, \boldsymbol {E} | \boldsymbol {X}, \boldsymbol {D}\right) = \left[ \prod_ {t = 1} ^ {T} q \left(\boldsymbol {m} _ {t}; \phi\right) q \left(\boldsymbol {e} _ {t} | \boldsymbol {m} _ {t}; \phi\right) \right] \left[ \prod_ {i = 1} ^ {N} q \left(\boldsymbol {z} _ {i} ^ {b} | \boldsymbol {x} _ {i}, \boldsymbol {d} _ {i}, \boldsymbol {E}, \boldsymbol {M}; \phi\right) \right]. \tag {5}
$$
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-VAE, that is identical to SAMS-VAE with all mask components fixed to 1. Thus, in contrast to equation 1, we have that  $z_{i}^{p} = \sum_{t=1}^{T} 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,\theta)$  on held out data, where  $\theta$  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:
$$
\operatorname {I W E L B O} (X | D, \phi , \theta) = \mathbb {E} _ {\boldsymbol {H} ^ {(1)}, \dots , \boldsymbol {H} ^ {(K)} \sim q (\boldsymbol {H} | \boldsymbol {X}, \boldsymbol {D}, \phi)} \log \frac {1}{K} \sum_ {k = 1} ^ {K} \frac {p (\boldsymbol {X} , \boldsymbol {H} _ {k} | \boldsymbol {D} , \boldsymbol {\theta})}{q (\boldsymbol {H} _ {k} | \boldsymbol {X} , \boldsymbol {D} , \phi)}.
$$
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:
$$
\operatorname {I W E L B O} (\boldsymbol {X} | \boldsymbol {D}, \phi , \boldsymbol {\theta}) = \mathbb {E} _ {\boldsymbol {H} ^ {(1)}, \dots , \boldsymbol {H} ^ {(K)} \sim q (\boldsymbol {H} | \boldsymbol {X}, \boldsymbol {D}, \phi)} \log \frac {1}{K} \sum_ {k = 1} ^ {K} w _ {k}, \tag {6}
$$
$$
w _ {k} = \left[ \frac {p \left(\boldsymbol {M} ^ {(k)} , \boldsymbol {E} ^ {(k)} \mid \boldsymbol {\theta}\right)}{q \left(\boldsymbol {M} ^ {(k)} , \boldsymbol {E} ^ {(k)} \mid \phi\right)} \prod_ {i = 1} ^ {N} \frac {p \left(\boldsymbol {x} _ {i} \mid \boldsymbol {z} _ {i} ^ {b (k)} , \boldsymbol {M} ^ {(k)} , \boldsymbol {E} ^ {(k)} , \boldsymbol {d} _ {i} , \boldsymbol {\theta}\right) p \left(\boldsymbol {z} _ {i} ^ {b (k)}\right)}{q \left(\boldsymbol {z} _ {i} ^ {b (k)} \mid \boldsymbol {M} ^ {(k)} , \boldsymbol {E} ^ {(k)} , \boldsymbol {X} , \boldsymbol {d} _ {i} , \boldsymbol {\phi}\right)} \right]. \tag {7}
$$
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:
$$
\mathrm {A T E} = \mathbb {E} _ {i \in \mathcal {D}} \left[ \left[ x _ {i, j} | \mathrm {d o} (\boldsymbol {d} ^ {*}) \right] - \left[ x _ {i, j} | \mathrm {d o} (\boldsymbol {d} _ {0}) \right] \right].
$$
We define the average treatment effect for SAMS-VAE as  $\mathrm{ATE}_{\mathrm{SAMS - VA}}(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$ :
$$
\operatorname {A T E} _ {\text {S A M S - V A E}} \left(\boldsymbol {d} ^ {*} \mid \boldsymbol {D} _ {m}\right) = \mathbb {E} _ {p \left(\boldsymbol {z} _ {j} ^ {b}\right) p (\boldsymbol {E}, \boldsymbol {M} \mid \boldsymbol {D} _ {m})} \left[ T _ {1} - T _ {2} \right],
$$
with  $T_{1}:= \mathbb{E}_{p(x_{i,j}|\mathrm{do}(\boldsymbol{d}^{*}),\boldsymbol{z}_{j}^{b},\boldsymbol{M},\boldsymbol{E})}[x_{i,j}]$  and  $T_{2}:= \mathbb{E}_{p(x_{i,j}|\mathrm{do}(\boldsymbol{d}_{0}),\boldsymbol{z}_{j}^{b},\boldsymbol{M},\boldsymbol{E})}[x_{i,j}]$ .
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_k^b \sim p(z^b)$  and  $E_k$ ,  $M_k \sim 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_{\delta}$  over which the expectation is computed (where  $D_{\delta}^{d}$  denotes the subset of the dataset  $D_{\delta}$  under condition  $d$ ):
$$
\mathrm {D E} _ {\text {D a t a}} \left(\boldsymbol {d} ^ {*} \mid \boldsymbol {D} _ {\delta}\right) = \mathbb {E} _ {x _ {i} \sim \boldsymbol {D} _ {\delta} ^ {\boldsymbol {d} ^ {*}}} \left[ x _ {i, j} \mid \mathrm {d o} \left(\boldsymbol {d} ^ {*}\right) \right] - \mathbb {E} _ {x _ {i} \sim \boldsymbol {D} _ {\delta} ^ {\boldsymbol {d} _ {0}}} \left[ x _ {i, j} \mid \mathrm {d o} \left(\boldsymbol {d} _ {0}\right) \right].
$$
To create the metric ATE-Pearson  $r(\mathrm{ATE}_{\mathrm{SAMS - VAE}}(d^{*}|\mathbf{D}_{m}),\mathrm{DE}_{\mathrm{Data}}(d^{*}|\mathbf{D}_{\delta}))$  we compute the Pearson correlation coefficient  $r$  between differential expression estimates from data  $\mathrm{DE}_{\mathrm{Data}}$  and our model-based estimator  $\mathrm{ATE}_{\mathrm{SAMS - VAE}}$  across all features (commonly genes indexed by  $j$ ).
This also reveals the relationship to a PPC  $p(r(\mathrm{ATE}_{\mathrm{SAMS - VAE}}(d^{*}|D_{m}), \mathrm{DE}_{\mathrm{Data}}(d^{*}|D_{\delta}))|D_{m}, D_{\delta})$ , 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_{\delta}$  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].
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.