Modelling Cellular Perturbations with the Sparse Additive Mechanism Shift Variational Autoencoder

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.

Strengths

* The paper does a good job of articulating appropriate assumptions for expression data. * 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. I 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.

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?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

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.

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.

Strengths

- The approach exhibits a good effectiveness and outperforms baselines such as SVAE+ significantly. - The paper is written in a overall clear manner. - The paper presents an original combination of previous models, such as the CPA-VAE and SVAE+ models.

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. - 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. - 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. *Minor:* I noticed a few typos: Line 91: "functiuon" should be "function." Line 105: "that aim" should be "that aims."

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? - What motivated the preference for the importance weighted ELBO as an evaluation metric over the standard ELBO measure? - 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?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The manuscript addresses the limitations adequately.

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

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. - 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. - 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. - 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? - 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. - 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. - Please make sure that IWELBO equation that is not outside the margins on page 4

Questions

See Weaknesses

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

To the best of my knowledge, this paper poses no potential negative societal impact.

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.

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. - Quality: The paper includes both technical and empirical results which with (some exceptions, see below) seem solid. - Clarity: The manuscript is well written, although dense at places and the reader could benefit from more details in the supplementary. - 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. 2. It would be helpful if the authors included the derivation of Eq 4 , 5. 3. 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.

Questions

In addition to my notes stated in the section above, it would help if the authors can answer the following: 1. There is no UMAPs for the latent representation of the baselines, could you include that for supporting the statement about improved interpretability? 2. 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? 3. Are there any identifiability guarantees of your model? In particular, how does your model cover modeling a discrete case such as gene expression data. 4. 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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Yes.

Respond to the Rebuttal

Thank you for your response. **Disentanglement**: Thank you for the clarifications. I understand the contributions but this is why I also mentioned that the title can miss-lead people. The title contains "...For Disentangled Representation Learning" therefore, and as there are different notions of disentanglement, you want to make it clear that the your notion is closer to causal discovery. **Figure 2: I understand the additive choice a bit better now. Thank you for your explanation regarding Figure 2. I am happy to take your word for it.

Disentanglement / causal representation learning.

I wanted to respond to the second section on disentanglement vs causality. I think your goals of disentangling "mechanisms of action as subspaces in an embedding" make a lot of sense and is likely to be useful. That said, the trouble with connecting with causal representation learning too strongly is that a key message of that literature is that achieving disentanglement is far more subtle than one might expect. For example, BetaVAE aims to disentangle data into independent factors of variation --- which is a reasonable goal --- but the main theorem in Locatello et al [2019] (which you cite in the intro) shows that even though BetaVAE's goals are reasonable (and we have empirical examples of it working on some datasets), this form of disentanglement is impossible from IID data without further assumptions. The main role for identifiability proof is to characterize sufficient conditions for when such reasonable sounding methods will actually succeed (in the asymptotic limit!). I think adding a limitations section that makes explicit that, while this approach seems to be successful at disentangling, more theoretical work is need to characterize when this will be successful would address my concern in this regard.

Response to rebuttal

I thank the authors for answering my questions and concerns, and I apologize for the misunderstanding with the baseline CPA-VAE. I updated my score accordingly. Concerning the relation with methods that incorporate domain knowledge: - I think that the paper would benefit a lot in reporting some comparisons with methods that incorporates domain knowledge in the method, or at least a in depth discussion on the expected performance of them. This is because: - It is expected that the method with domain knowledge will perform better, but the gap in performance is informative. - It is very relevant to check in which cases the data driven method outperform the domain informed one, giving evidence for the former to find patterns that allows to make good predictions, which the domain informed model fails to find. - It is important to incorporate the discussion on data driven vs domain informed anyway, highlighting cases where the former can be better and the latter could be too expensive or limited. In general I agree with the other reviewers and authors considerations that, while the formulation is general, the method is targeted for biological applications and this should be reflected in the paper. Without results (theoretical or experimental on the right benchmarks) on disentanglement and/or causality, not much can be said about it. Having a general formulation is for sure a plus, but I think that the focus should be on biological tasks and for this reason, I think that the paper would benefit from the points raised above.