SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In our experiments, we represent perturb-seq datasets as a gene expression matrix  $X \in \mathbf{N}^{N \times D_x}$  and a perturbation dosage matrix  $T \in \{0,1\}^{N \times T}$  for  $N$  cells,  $D_x$  gene transcripts, and  $T$  perturbations. Entries  $X_{i,j}$  represent the number of transcripts from gene  $j$  observed in cell  $i$  and entries  $T_{i,k}$  represent whether cell  $i$  received perturbation  $k$ . We compare SAMS-VAE against baseline models based on their ability to model the distribution of perturb-seq data, generalize to new perturbations in combinatorial settings, and disentangle known biological pathways in their latent variables.
Baselines We consider CPA-VAE, SVAE+, and conditional VAE [21] as baselines. As discussed in Section 2.4, CPA-VAE can be thought of as an extension of CPA [13] to a fully specified generative model and takes advantage of our proposed correlated inference strategy. We additionally consider ablations of the correlated inference strategies for SAMS-VAE and CPA-VAE. Complete details of model choices are provided in the appendix.
Code availability Our code, which includes implementations of all models and experiment configurations, is available at https://github.com/insitro/sams-vae
4.1 Generalization under individual perturbations
Dataset To assess model generalization to held out samples under individual perturbations, we analyze a subset of the genome-wide CRISPR interference (CRISPRi) perturb-seq dataset from Replogle et al. [17], which we call replogle-filtered. CRISPRi is a type of genetic perturbation that represses the expression of selected target genes. Following the preprocessing steps from Lopez et al. [12], replogle-filtered is filtered to contain perturbations that were identified as having strong effects and genes that were associated with these perturbations. We additionally include cells with non-targeting CRISPR guides to use as controls for average treatment effect prediction. All together, replogle-filtered contains 118,461 cells, 1,187 gene expression features per cell, and 722 unique CRISPR guides (perturbations). We randomly sample train, validation, and test splits.
Model	Inference	Test IWELBO	Mask PW. Acc.	ATE-Pearson
Conditional VAE	amortized MF	-1766.10 ± 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. \( z_{basal} \)	-1756.57 ± 0.14	-	0.571
SAMS-VAE	amortized MF	-1757.72 ± 0.14	0.68 ± 0.09	0.302
SAMS-VAE	corr. E	-1758.08 ± 0.07	0.71 ± 0.04	0.319
SAMS-VAE	corr. \( z_{basal} \)	-1756.40 ± 0.06	0.87 ± 0.02	0.718
SAMS-VAE	corr. \( z_{basal} \) and E	-1756.27 ± 0.10	0.89 ± 0.03	0.765
Table 1: Quantitative evaluation of treatment effects on Replogle filtered dataset (100 latent dimensions) 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.
Figure 2: Visualization of inferred latent perturbation masks and embedding means for the best performing checkpoint of each model in replogle-filtered. We visualize the latent variables for the 345 perturbations with pathway annotations from Replogle et al. [17] and group by pathway. The SAMS-VAE and CPA-VAE models were trained with our proposed correlated inference strategy.
Evaluation protocol Each model is trained with a 100 dimensional latent space and MLP encoders and decoders with a single hidden layer of dimension 400 (see Section A.4 for full training details). Based on validation performance and sparsity, a  $\mathrm{Beta}(1,2)$  prior was selected for the SVAE+ mask, and a Bern(0.001) prior was selected for SAMS-VAE. For each model type, we compute the test set importance weighted ELBO as described in Section 3.1 and report the mean and standard deviation across five training runs with different random seeds. We additionally estimate the model average treatment effect with  $K = 10,000$  particles as defined in Section 3.2 for the best model of each type and report the correlation between this quantity and the estimated differential expression from data.
Results Quantitative results are presented in Table1 Comparing first between model types, we observe that SAMS-VAE with fully correlated inference achieves the best test IWELBO and average treatment effect correlation. Interestingly, CPA-VAE 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, we perform an ablation of SAMS-VAE and CPA-VAE inference strategies. We find that the correlated  $z_{basal}$  strategy yields substantial improvements in performance for both SAMS-VAE and CPA-VAE, while the correlated  $E$  strategy improvements are minor.
4.1.1 Recovery of biological mechanisms based on disentangled factors
Evaluation protocol We assess the degree to which the pattern of perturbation effects on inferred latent factors in the SAMS-VAE and SVAE+ models from the previous section are predictive of
Figure 3: We visualize model-estimated treatment effects (ATEsAMS-VAE) and data-estimated differential expression  $(\mathrm{DE}_{\mathrm{Data}})$  for intervention-gene pairs in the Replogle experiment. We observe broad correlation (Pearson  $r = 0.765$ ): for example, perturbations of ribosomal subunits influence on all expression broadly with matching directionality, while other guides exhibit more targeted effects.
known biological pathways as annotated by Repogle et al. [17] (345 of the 722 targeted genetic perturbations are annotated). To do so, we define an inferred binary mask of perturbation effects on latent factors by thresholding the inferred latent mask probabilities in each model at  $p = 0.5$ . For each model, we fit a random forest model using scikit-learn [16] to predict pathway annotations from a subset of the perturbation latent masks and assess pathway prediction accuracy on the remaining perturbations. This evaluation is performed on the best checkpoint for each model type with 10 random splits of perturbations (70% train, 30% test), and the mean and standard deviation of the pathway prediction accuracy is reported.
We also provide a set of visualizations to qualitatively assess the latent structures learned by each model. We plot the inferred masks and embeddings for SAMS-VAE, SVAE+, and CPA-VAE in Figure2 and visualize the SAMS-VAE estimated average treatment effects and estimated differential expression corresponding to these perturbation effects in Figure3. Hierarchical clustering and UMAP projection of the inferred perturbation embeddings are presented in Figure7 in Section A.6
Results We observe that the latent mask inferred by SAMS-VAE is more predictive of the annotated pathways in the replogle-filtered dataset than that inferred by SVAE+. Additionally, we find that performing correlated inference on the perturbation embeddings improved pathway prediction performance for SAMS-VAE. Qualitatively, we observe that both SAMS-VAE and SVAE+ infer sparse masks with distinct patterns between annotated pathways.
4.2 Modeling compositional interventions in a CRISPRa perturb-seq screen
Dataset We analyze the CRISPR activation (CRISPRa) perturb-seq screen from Norman et al. [15] to assess how effectively SAMS-VAE and baselines model the effect of perturbation combinations. This screen was specifically designed with perturbations that have non-additive effects in combination, making this a challenging setting for modeling combinations. We adopt the preprocessing from [6], which contains 105 unique targeting guides applied both on their own and in 131 combinations. In total, the dataset contains 111,255 cells, each with 5,000 gene expression features.
Evaluation protocol We define two tasks using this data. The first, norman-ood, assesses the ability of each model to predict gene expression profiles for held-out cells that have received perturbation combinations that are not included in the training set. Each model is trained on cells that received a single guide, along with [0, 25, 50, 75, 100% of combinations. Held-out cells receiving the final  $25\%$  of combinations are used to evaluate each model. We perform this analysis for 5 random splits of the combinations. The second task, norman-data-efficiency, assesses how efficiently the models can learn combination phenotypes when trained on cells that have received a single guide and increasing numbers of cells sampled uniformly across all combinations. Each model is evaluated
Out-of-Distribution Combinations
Combination Data Efficiency
Figure 4: Results from norman-ood and norman-data-efficiency experiments. Within splits, test IWELBO values are plotted relative to the test IWELBO for SAMS-VAE trained with 0 combinations 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  $\mathrm{Beta}(1,2)$  prior and SAMS-VAE is trained with a Bern(0.01) prior.
Results Quantitative results are presented in Figure 4 with additional inference strategy ablations in Figure 10 in Section A.6. SAMS-VAE and CPA-VAE both achieve strong performance on the norman-ood task across metrics, often within 1 standard deviation of one another. Conditional 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 a new combination without observing it in training. We do observe that the SVAE+ likelihood still improves a small amount as more combinations are included in training set, which may be attributable to improvements in the encoder and decoder (which are shared across perturbations). These results support the utility of the compositional mechanisms in SAMS-VAE and CPA-VAE (and for encoding combinations as defined in  $d_{i}$  for conditional VAE).
In norman-data-efficiency, we observe similar trends. SAMS-VAE, CPA-VAE, and conditional VAE, which can share information across individual and combined perturbations, all achieve better ATE prediction for held out cells when less than  $50\%$  of the available combination cells. However,  $\mathrm{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-VAE again perform the best, with SVAE+. These results further support the utility of the additive composition mechanism from SAMs-VAE and CPA-VAE in low data settings.
Disentangled VAEs Disentangled variational auto-encoders have been proposed as early as in [7], where weak supervision was utilized to learn sparse masks over different subspaces. A popular framework for unsupervised disentanglement was proposed in [3] through a reweighting of the regularizer in the objective, but ignores weak supervision about conditions. A more comprehensive treatise and theoretical analysis of disentanglement was presented in [10]. Finally, the formal link to sparse mechanism shift and explicit causal disentanglement was also established recently in [9]. Our work shares assumptions with some of these works, in that we assume dosage is known leading to specific shifted effects per perturbation that can be used to learn disentangled factors.
Models of Cellular Perturbation A popular generative modeling framework for cellular sequencing data utilizing VAEs has been proposed in [11], a model which inspired our use of their likelihood. Closer to our application on perturbed datasets is the CPA [13]. Similar to this model, we adopt the idea to disentangle cellular latent spaces into basal and perturbation latent variables. However, we pose the resulting model as a joint generative model with a rigorous inference framework and, crucially, a sparsity mechanism to disentangle the perturbation effects into subspaces related to the affected mechanisms. The recent SVAE+ [12] model is an exciting variant of [9] that utilizes disentanglement in a fashion that matches our goals. Our work differs by factorizing variation into basal and perturbation variables, adding a mechanism to compose perturbations, and in terms of inference strategies (SVAE+ learns perturbation representations by optimizing a prior). GEARS [18] leverages prior information of perturbation and gene features to predict the effect of applying new perturbations to unperturbed cells. This work instead focuses on specifying a generative model for perturbation effects with minimal assumptions, though strategies for integrating prior information on perturbations and features in specific use cases is an exciting future direction.
Performing unbiased scientific discovery is an aspirational goal in the field of drug discovery to detect mechanisms of action with intervenable potential. We propose a model that attempts to use few explicit assumptions about the nature of the observed data and relies heavily on a sparsity assumption and decomposition into explicit treatment effects in latent space to learn models of perturbational screening data, making it general enough for application to arbitrary data modalities and perturbation types. In this work, we apply SAMS-VAE to genetic perturbations and single-cell sequencing readouts and observe two key outcomes: improved predictive performance compared to omitting the sparsity assumption, and improved ability to recover factors correlated with real mechanisms in biological data. Our technical contributions cover both a specification of a novel sparse generative model, SAMS-VAE, as well as a suite of inference strategies for improved model fit which also apply to our baseline CPA-VAE. We also propose an evaluation strategy for perturbation models related to posterior predictive checks utilizing average treatment effects and differential expression as test statistics to perform model criticism, and observe our model performing competitively in this metric.
Such models may ultimately be useful to perform experiments in an iterative fashion, and help specify actionable hypothesis spaces for more targeted experiments down the line. Our work falls into a long line of literature on disentanglement and more recently the Sparse Mechanism Shift hypothesis related to causality, and we believe that the specific setup of SAMS-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.
We acknowledge and thank insitro for funding this work.
[17] J. M. Replogle, R. A. Saunders, A. N. Pogson, J. A. Hussmann, A. Lenail, A. Guna, L. Mascibroda, E. J. Wagner, K. Adelman, G. Lithwick-Yanai, N. Iremadze, F. Oberstrass, D. Lipson, J. L. Bonnar, M. Jost, T. M. Norman, and J. S. Weissman. Mapping information-rich genotype-phenotype landscapes with genome-scale perturb-seq. Cell, 185(14):2559-2575.e28, 2022. ISSN 0092-8674. doi: https://doi.org/10.1016/j.cell.2022.05.013. URL https://www.sciencedirect.com/science/article/pii/S0092867422005979
[18] Y. Roohani, K. Huang, and J. Leskovec. Predicting transcriptional outcomes of novel multigene perturbations with gears. Nature Biotechnology, 2023. doi: 10.1038/s41587-023-01905-6.
[19] D. B. Rubin. Bayesianly justifiable and relevant frequency calculations for the applied statistician. The Annals of Statistics, pages 1151-1172, 1984.
[20] B. Scholkopf, 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 X. Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015.