Learning Identifiable Factorized Causal Representations of Cellular Responses

Summary

This paper proposes FCR, a causal VAE-based model that aims to decompose and cluster covariates, treatments, and their interactions in the latent space.

Strengths

This paper has a solid and rigorous mathematical foundation, and the assumptions are validated after the model is trained.

Weaknesses

This paper is limited in readability, mainly because of word redundancy, making it hard for the reader to follow the authors’ ideas. Besides, Figure 2 outlines the FCR model, but the encoders and decoders are disconnected, so it is not intuitive to understand how we train components such as $p(z_t|t)$ and $p(z_x|t)$.

Questions

Is it possible to enhance other model’s interpretability using the factorized $z_x,z_t,z_{tx}$?

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

I cannot find a section discussing the limitations of this work, but I believe the noise and hidden variables in the single-cell data will cap the model performance.

Summary

The authors propose a method for disentangling into factorized representations dependent on only covariates, only treatment and the interaction between treatment and covariates. Introducing a set of assumptions, the authors prove the identifiability of these disentangled variables based on previous proofs on non-linear ICA. The authors implement the model using variational inference and a set of regularizers and discriminators to ensure the needed independence and disentangling. The authors validate the clustering performance, conditional independencies and conditional cellular response prediction against commonly used methods for perturb-seq prediction and show generally very good performance.

Strengths

- Non-linear interactions between covariates and treatment are very likely and have to be modeled - Given a set of additional assumptions, the method has some guarantees of disentangling, which is also shown empirically - Understanding (e.g. through clustering) interactions between covariates and treatment is important for precision medicine applications - Good performance as shown by authors

Weaknesses

- Complicated implementation with many loss functions and minmax training, a discussion on how easy to train the method is would be good (dependence of performance on hyperparameters and hyperparameters on data set). - The ablation study should include a comparison without the different loss terms to be able to see the importance of the different parts of the model - Assumption 4.3 could be quite strong in many settings, if for every covariate every drug response has to be available. A comment on sufficient experiments would be good and how common this is. - For some of the assumptions (e.g. substantive changes of the distribution) empirically arguing that this is the case in common data sets would make the real world application of the proof (instead of the purely academic pursuit) more convincing.

Questions

- How were the hyperparameters chosen for the different data sets? - Maybe I am misunderstanding "control", but in the conditional cellular response prediction, why is it necessary to use $z_x^0$ instead of $z_x$ in the function $g$ if $z_x$ is independent of $t$? - For prediction validation, what is train/val/test protocol? Is the correlation of y in the train set?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors could discuss the limitations (based on assumptions) more directly.

Summary

The authors present a novel method for causal representation learning in cellular perturbation settings, called Factorized Causal Representation, in which the authors learn disentangled latents for the cellular covariates, the treatment, and the interactions between them. They provide identifiability results following the methods of recent papers, and demonstrate how an implementation of the method, which uses a number of novel regularizers to enforce various causal constraints, achieves top results amongst comparable methods.

Strengths

- novel causal representation learning method explicitly learning interactions between covariates and treatments - novel regularizers for enforcing conditional independence between sets of variables - achieves top cellular response prediction ($R^2$) over comparable conditions

Weaknesses

- the datasets used have a relatively small number of similar treatments (16 at the most), so it is unclear how well the approach will work when applied to large diverse chemical libraries typical of real-life settings - the decision not to evaluate on unseen treatments or cell contexts is surprisingly conservative, since the promise of causal models is in generalizing to unseen domains

Questions

- why not evaluate on unseen treatments or cell contexts? Did you try and find the results to be comparatively poor?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

The paper presents a novel method, the Factorized Causal Representation (FCR), which leverages identifiable deep generative models to understand cellular responses to genetic and chemical perturbations across multiple cellular contexts. This method improves upon prior models by learning disentangled representations that delineate treatment-specific, covariate-specific, and interaction-specific factors, which are shown to be theoretically identifiable. The effectiveness of FCR is demonstrated on four single-cell datasets. The performance, measured by the R^2 score of conditional cellular responses, often marginally outperforms state-of-the-art methods, showcasing FCR’s potential in accurately modeling cellular dynamics.

Strengths

- The paper successfully demonstrates the identifiability of treatment-specific, covariate-specific, and interaction-specific representations. - The inclusion of the cell type classifier $f_{ct}$ and permutations discriminator $f_{dist}$ are innovative loss components. - Sections 6.1 and 6.2 are well-designed to rigorously test the proposed method's ability to maintain the integrity of dimension reduction, decomposition, and conditional independence, validating the model's effectiveness in designing and training.

Weaknesses

Choice of Evaluation Metric: A significant concern arises from the choice of the $R^2$ metric in section 6.3 for evaluating the performance of the Factorized Causal Representation (FCR) model. $R^2$ is suitable to assess whether a model adequately explains the variation in the response variable. However, it is not necessarily an indicator of the model's predictive accuracy. Moreover, the near-identical $R^2$ reported for FCR, VCI, CPA, and sVAE (as shown in Table 1) suggest that $R^2$ gives limited insights into the comparative effectiveness of these methods. There are some minor misformats, which I noticed while reading. - Lines 35, 54: citation format issue - Line 294: subscript format issue

Questions

Can you elaborate more about the design choice of the evaluating metric?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

- Generalizability: The paper does not show the method’s applicability to apply on new dataset. This limitation raises concerns about how well the FCR can adapt to other biological datasets or experimental setups that deviate from the four testing datasets, which restricts the method’s utility. - Interpretability: The model lacks mechanisms for clear interpretation, particularly in linking gene expression profile changes directly to specific treatments. So the decomposition does not provide insights into how gene expression is affected by specific treatments. I understand that both limitations are significantly challenging and it is unrealistic to tackle both of them in one paper while they highlight critical areas for future research.

Summary

The authors develop a modification of the identifiable nonlinear ICA algorithm of Khemakhem et al. 2020, suited for scenarios with two disjoint groups of auxilliary variables x and t and their interactions, with an application of treatment effect on cell gene expression. The authors propose a learning objective designed to enforce the conditional independence assumptions of the model, and evaluate it against baselines from the disentanglement, ICA, and cellular response estimation.

Strengths

- the proposed method is justified theoretically - the paper is written in a clear and concise manner - extensive selection of baselines

Weaknesses

- I fail to see the reasoning of comparing the conditional independence results (6.3) with the other baselines, which were not designed to have separate latent subspaces for x,t, and tx, and thus it does not seem surprising that the conditional independence does not hold for random subsets of them? - Results of the cellular response tasks are barely better than the baselines - did the authors conduct any tests to see if the differences are statistically significant?

Questions

- Equation (6) (the constraint for interaction between x and t) - how does that guarantee the interactions? The learned embeddings can have empty dimensions, e.g, k(x) = [k_1(x), …, k_n(x), 1, …, 1] and k(t) = [1, …, 1, k_1(t), …, k_m(t)] so that the Hadamard product becomes [k(x), k(t)] - Equation (7) - shouldn’t f_dis also take x as an argument since you want to enforce a conditional indpendence? Or do you train a separate f_dis for each possible value of x that you condition on (probably not feasible)? - If the model is supposed to be fully identifiable wrt. z_xt, then I would suggest evaluating this with the appropriate metrics on synthetically generated data in addition to the clustering objective (e.g., MCC as in Khemakhem et al. 2020, or any of the disentanglement metrics [Locatello et al. 2019]) - How were the hyperparameters selected for the cellular response task?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

- the algorithm has 3 hyperparameters, which might make training difficult. Details on hyperparameter selection were ommited - the evaluation is limited - results of the downstream task of predicting cellular response yield limited improvements, and the identifiability of z_{t,x} is not properly evaluated (with e.g., the MCC metric using synthetic data)