Learning Identifiable Factorized Causal Representations of Cellular Responses

$**Readability**$:

Thank you for your feedback regarding the readability of our manuscript. We appreciate your honest assessment and understand the importance of clear and concise writing. We will carefully review the manuscript to eliminate any word redundancy and improve the overall clarity, ensuring that our ideas are communicated more effectively.

$**$p(\mathbf{z**_t|\mathbf{t})$and$p(\mathbf{z}_x|\mathbf{x}) learning}:

Thank you for this question. I’d like to clarify that both $x$ and $t$ are inputs to the neural network, which then encodes them to learn the prior mean and standard deviation for $z_t$ and $z_x$. We use the reparameterization trick to sample $p(z_t|t)$ and $p(z_x|x)$. We will add this explanation under the model figure in the manuscript to provide clearer context.

$**Interpretability**$:

There are generally two approaches to interpreting latent representations, which we also plan to explore in future work:

1. Latent Traversal: Since $z_{tx}$ is element-wise identifiable, we can use latent traversal across all dimensions of $z_{tx}$. This involves fixing all latent variables $z_{tx}^i$ to their inferred values for each cell's gene expression and then varying the value of one latent at a time to observe the reconstructed gene expression. This method allows us to see how each dimension of $z_{tx}$ influences gene expression levels.
2. Interpretable Decoder System: We can also employ an interpretable decoder, such as LDVAE [1], which uses a linear decoder. By analyzing the weights of the latent representations assigned to each gene, we can interpret how different representations influence each gene.
3. On-going Work: In light of the neural additive model [2], we are developing a non-linear interpretable encoder as an extension of LDVAE. Specifically, for each latent dimension, there is a sub-neural network to capture information specific to each latent dimension. A logistic regression layer is added to the end of the neural network. So, by analyzing the logistic regression weights (loading matrix), we can conclude how much one latent variable contributes to the reconstruction of a specific gene. An illustration of this model can be found in the rebuttal PDF file Figure D.

$**Limitation discussions**$:

Thank you for raising this important question. Indeed, this paper has some limitations that we plan to address in future work.

Firstly, the generating function 𝑔 used in our model is either deterministic or incorporates Gaussian noise. Future research will investigate more general scenarios with stochastic generating functions, which are more representative of real-world applications. Secondly, while this paper does not include interpretability mechanisms for FCR, we have proposed potential approaches to enhance interpretability, as discussed in the previous section. Lastly, due to the limitations of the available public datasets, we recognize that testing FCR on datasets with a broader range of drugs and more complex covariate scenarios will be an intriguing direction for future work.

[1] Svensson, V., et al. (2020). Interpretable factor models of single-cell RNA-seq via variational autoencoders. Bioinformatics (Oxford, England), 36(11), 3418–3421.
[2] Agarwal, R., et al. (2021). Neural additive models: Interpretable machine learning with neural nets. Advances in neural information processing systems, 34, 4699-4711.

Thank you for reading our response. Please feel free to let us know if there is anything needs to be clarified after the rebuttal. We value your feedback and want to address your remaining concerns.

$**Hyperparameter selection**$: Thank you for raising this important question, we appreciate that it is a critical point for evaluating how well the method works. Please check General Response to All Reviewer b.

$**Ablation studies**$:

In addition, as the reviewer suggested, we did the components ablation study (without the different loss terms) on the clustering task (sciPlex dataset) with the NMI score reported.
$**Table 1: Benchmark on$\mathbf{z**_x}

$**Table 2: Benchmark on$\mathbf{z**_{tx}}

$**Table 3: Benchmark on$\mathbf{z**_t}

When one of ${w_1, w_2, w_3}$ is set to 0, the remaining weights are set to 1.0. Based on this setup, we draw the following conclusions:

The results of $z_x$ clustering on $x$, {$t, x$}, and $t$ indicate that

• $w_1$ (the weight for $L_{sim}$) significantly enhances the performance of $z_x$ clustering on $x$ and $z_t$ clustering on $t$. This improvement is due to $L_{sim}$ increasing the similarity among $z_x$ and $z_t$.
• $w_2$ (the weight for $L_{ct}$) enhances $z_{tx}$ clustering on {$t$, $x$} and $t$, as well as $z_t$ clustering.
• $w_3$ (the weight for $L_{dist}$) improves $z_{tx}$ clustering on {$t$, $x$}and also enhances the separation within the spaces of $z_{tx}$, $z_t$, and $z_x$.

It's important to note that when ${w_1, w_2, w_3} = {1,1,1}$, the performance of $z_{tx}$ clustering on {$t$, $x$}is not as strong as $z_x$ clustering on {$t$, $x$}. However, by increasing $w_1$, $w_2$, and $w_3$, the performance of $z_{tx}$ surpasses that of $z_x$ (see Appendix Figure 7).
$**Necessary to use$z_{x**^0$instead of$z_{x}$in the function g, if$z_{x}$is independent of$t}?:

Thank you for raising this question. Yes, directly utilizing $z_x$ might achieve higher $R^2$ score, here we followed CPA, VCI and other methods and replace by $z_x^0$ to perform "counterfactual" prediction to further evaluate the learned disentangled representations.

$**Assumption 4.3 is strong**$:

Assumption 4.3 can be better understood through the lens of biological experimental design. If we aim to determine whether there is a covariate $x_1$-specific response to treatment $t_1$, we need to conduct two reference experiments: $(x_0, t_1)$ and $(x_0, t_0)$, where $x_0$ represents reference cells and $t_0$ represents control/reference treatments.

By comparing the treatment effects between the pairs $\{(x_1, t_1), (x_1, t_0)\}$ and $\{(x_1, t_1), (x_0, t_1)\}$, we can assess whether the observed differences are significant and distinct (as described in Assumption 4.4). Without comparing $(x_1, t_1)$ and $(x_1, t_0)$, we cannot determine if $t_1$ has any effect on $x_1$. Similarly, without comparing $(x_1, t_1)$ and $(x_0, t_1)$, we cannot assess whether $t_1$ has a differential effect on $x_1$ than $x_0$. Only when both comparisons are made can we begin to analyze covariate-specific effects. This comparison allows us to confidently identify any $x_1$-specific response to $t_1$.

This is a common approach in biological experimental design. For instance, an increasing number of experiments are being conducted for drug screening across various micro-environment organoid systems (covariates), as well as with different drugs, dosages, and time points [3]. The main goal is to discover these covariate-specific cellular responses.

$**Assumptions empirically arguing that this is the case in common data sets**$:

The substantive changes assumption (iii) means the different cell lines under the same drug treatment, the gene expression is significantly different (non-trivial). It can be seen that there are large parts of gene expression only influenced by cell types, etc. (iv) means under different treatments, the same cell (covariates) have significantly different responses to gene expression.

These phenomena are all observed in biological literatures [1] [2] [3]. The Figures 3.B,C,D ($**Figure D in the rebuttal PDF file**$) in the Trellis paper [3] show evidence that these assumptions are common.

$**Limiations**$: Due to the limited space, please check the response to $**Reviewer WECd**.$

[1] Sanjay R. Srivatsan et al. ,Massively multiplex chemical transcriptomics at single-cell resolution.
[2] McFarland, J.M., Paolella, B.R., Warren, A. et al. Multiplexed single-cell transcriptional response profiling to define cancer vulnerabilities and therapeutic mechanism of action.
[3] Ramos Zapatero, María et al. Trellis tree-based analysis reveals stromal regulation of patient-derived organoid drug responses.

Thank you again for acknowledging our paper strengths and raising the score. We will further polish our manuscript and include the hyperparameter tuning and ablation study in the final version.

$**Dataset size**$:

Due to the limited availability of public datasets and the high cost of single-cell sequencing, single-cell datasets with a large number of drugs are not yet available. However, as demonstrated by the theorems in our paper, incorporating more treatments with well-designed controls can enable the latent representations to capture more meaningful and disentangled information, thereby facilitating the algorithm's learning process. As more datasets become available in the future, we plan to conduct extensive benchmarks to evaluate our approach on a broader scale.

$**Prediction for unseen perturbations**$ :

Predicting responses to novel treatments is a pivotal and fast-evolving field in drug discovery. However, the biological literature indicates that cellular responses are highly context-dependent [1,2,3]. This complexity poses significant challenges for AI-driven drug discovery, which often struggles to achieve success in clinical trials [4].

Our motivation for developing FCR stems from the need to understand how cellular systems react to treatments and identify conditions that can deepen our understanding of these responses. FCR enables the analysis of drug interactions with covariates and contextual variables. Additionally, predicting cellular responses to new treatments necessitates prior knowledge, such as chemical structure and molecular function, and comparisons with known treatments. Without this context, predictions can be unreliable.

In this paper, our primary focus is not on predicting responses to unseen treatments. However, given the relevance of this topic, we conducted pilot experiments to showcase the potential future applications of FCR. The MultiTram-Plex and Multiplex-7 datasets share the same cell lines and Trametinib-24 hours treatment, along with other different treatments. By utilizing these two datasets, we established the following experimental settings for unseen prediction scenario:

• Drug Hold-out Setup: We held out two cell lines, ILAM and SKMEL2, from the Multiplex-Tram dataset, which had been treated with Trametinib for 24 hours. We trained a FCR model using the remaining data, and this model is referred to as $M_h$. We denote the dataset (Multiplex-Tram dataset without Trametinib-24h treated ILAM and SKMEL2) as $D^h$
• Prior Knowledge Model: We trained another FCR model, $M_p$, using the Multiplex-7 dataset, which includes the ILAM and SKMEL2 cell lines treated with Trametinib for 24 hours denoted as $D^p$. We treat this model, $M_p$, as a prior knowledge model.
• Transfer MLP: We extract both $z_{tx}^{p}$ from $M_p$ and $z_{tx}^{h}$ from $M_h$ for $D^p$. Then we trained a 1-layer MLP (the same dimension as $z_{tx}^p$) to transfer $z_{tx}^p$ to $z_{tx}^h$, by minimize the MSE between them.
• Contextual Prior Representation: We extracted the $z_{tx}^p$ representations from model $M_p$ for ILAM and SKMEL2 in holdout set. Then transfer $z_{tx}^p$ by the previous MLP to $\hat z_{tx}^h$ as a prior contextual embedding. For the hold-out cell lines ILAM and SKMEL2 in the Multiplex-Tram dataset, we extracted $z_{x}^h$ from model $M_h$, and Trametinib-24h $z_{t}^h$ from other treated cell lines in $D^h$
• Representation Matching: Then we match the $\hat z_{tx}^h$ by $z_{x}^p$ similarity on ILAM and SKMEL2 across Multiplex-tram and Multiplex-7 in the prior knowledge model space.
• Prediction: We predicted the unseen 24 hours Trametinib responses for the ILAM and SKMEL2 cell lines in holdout dataset using the formula: $\hat y =g(z_{x}^h,\hat z_{tx}^h, z_{t}^h)$,where $\hat z_{tx}^h$ is the corresponding matched prior contextual representation, $z_t^h$ is the tramnib-24 average value in other cell lines.
• Evaluation: We computed the $R^2$ and MSE for the top 20 differentially expressed genes (DEGs) based on the predicted values and compared these results with those from the VCI model.The paper’s OOD prediction setups.)
• $**The illustration of datasets and experiments setups can be found in the rebuttal PDF file, Figure A and B**$.

We have the following results:

$**Table: Unseen Prediction Results**$

[1] Sanjay R. Srivatsan et al. ,Massively multiplex chemical transcriptomics at single-cell resolution.Science.
[2] McFarland, J.M., Paolella, B.R., Warren, A. et al. Multiplexed single-cell transcriptional response profiling to define cancer vulnerabilities and therapeutic mechanism of action. Nat Commun.
[3] Ramos Zapatero, María et al. “Trellis tree-based analysis reveals stromal regulation of patient-derived organoid drug responses.” Cell vol. 186,25 (2023).
[4] Derek Lowe "AI drugs so far".

Dear Reviewer yZVk:

We sincerely appreciate your constructive suggestions on the unseen prediction. During the rebuttal stage, we worked diligently to design and execute the experiment to address your feedback (Please refer to the response above, as well as Figures A and B in the rebuttal FDF file). We would greatly value your feedback on the experiment and would be eager to discuss it further with you before the discussion period ends. Your insights will be crucial in helping us refine our paper.

Thank you,

All the authors

$**Comparing the conditional independence results (6.3) with the other baselines**$:

We thank the reviewer for raising this important question. I would like to clarify the experimental setup. The goal of this experiment is to demonstrate that FCR is capable of learning conditionally independent spaces for $z_x$, $z_t$, and $z_{tx}$. Our algorithm is the first to specifically address the $x$, $t$, and $tx$ subspaces, which makes it challenging to find direct baselines for comparison. However, the baselines such as iVAE, betaVAE, factorVAE, and sVAE are semi-supervised or unsupervised disentangled representation learning methods, they aim to learn representations that should be independent or conditionally independent given $x$ and $t$. Additionally, methods like CPA and VCI enforce certain invariance relations with respect to $t$ and $x$, potentially leading to some level of conditional independence in the learned representations.

However, the challenge in making this comparison lies in not knowing which blocks of latent variables induced by baseline methods satisfy the specific conditional independence relations of interest. To ensure a fair comparison, we exhausted different combinations of the representations for the tests and reported the best results. This comparison is intended to demonstrate that FCR can effectively learn more meaningful and nuanced conditional independence relations in $x$, $t$ and $tx$, compared to the baselines. We will ensure that this clarification is included in the manuscript.

$**Differences are statistically significant**$:

Thank you for raising this concern. We tested for significance of changes in mean between the results from the FCR method and those of the second-best performing method (paired t-test). The p-values, along with the corresponding cell number in the testing set appears in the table below.

Table: p-value for $R^2$ performance

The results indicate that FCR significantly out-performed the baselines across cell lines on sci-Plex, multiPlex-Tram and multiPlex-7. On MultiPlex-9, where FCR is slightly outperformed by CPA, the differences are not significant. We will make sure to include all p-values in the updated manuscript.

To propose an additional metric beyond $R^2$, we also evaluated the mean square error (MSE), for the top 20 Differentially Expressed genes (DEGs). These 20 genes are measured on the largest difference for each cell line with drug treatments compared to control samples (following [1]). Note here, we didn’t compare with CINEMA-OT and scGEN because they are only for binary treatments (case-control data). For detailed results, please check the results in General Response to ALL Reviewers a. Generally speaking, FCR outperforms other baselines in 3 of 4 datasets.

$**Concerns on Equation (6) (the constraint for interaction between$\mathbf{x**$and$\mathbf{t})}:

The reviewer asked the question "How does that guarantee the interactions? The learned embeddings can have empty dimensions, e.g, $k(x) = [k_1(x), \ldots, k_n(x), 1, \ldots, 1]$ and $k(t) = [1, \ldots, 1, k_1(t), \ldots, k_m(t)]$ so that the Hadamard product becomes $[k(x), k(t)]$".

We apologize for this misleading language. Essentially, we did not want to claim any theoretical guarantee in this case, but instead motivate the specific architecture choice. We therefore changed the language to "in order to encourage the prior for $z_{xt}$ to capture interactions between $x$ and $t$, we design the functions $f^p_{\mu}$ and $f^p_{\Sigma}$ to be of the form $k(x) \otimes k(t)$, where $\otimes$ denotes the Hadamard product.

$**x as an argument for$f_{dis**}: We agree with the reviewer that it would be much clearer to put $x$ as an argument of $f_{dist}$. It is already the case that the function $f_{dist}$ implicitly considers $x$ as an argument because we perform a random permutation of the triplet $(z_x, z_{tx}, x)$ under the condition of the same $x$ (i.e., the permutation or shuffle occurs only within a set of $z_x$ and $z_{tx}$ under the same $x$). As a result, the input to $f_{dist}$ is the set of newly generated ($z_x$, $\hat z_{tx}$), where we ensure that $z_x$ and $\hat z_{tx}$ correspond to the same $x$. The outputs of $f_{dist}$ simply are the labels indicating whether a permutation occurred or not. Therefore, we can say that $f_{dist}$ effectively takes $x$ into account. We briefly discuss this process in lines 226-230. We will make sure to update it in our final manuscript.

$**Simulation studies on synthetic data**$: Please check the General Response to ALL Reviewers c.

$**Hyperparameter selection**$: Thank you for raising this important question, we appreciate that it is a critical point for evaluating how well the method works. Please check General Response to All Reviewers b.

Thank you very much for raising the score. For the MSE evaluation, we selected the top 20 differentially expressed genes, identified by conducting a comparison between treated and untreated cells, and ranked them based on t-test p-values, we considered the rank information to some extent in this evaluation. Additionally, we plan to add the Spearman correlation comparison in the final version (The table below shows FCR and VCI results). Again, thank you for your valuable comments which helped us improve our work a lot!
Table: Spearman Correlation (top 500 highly variable genes):

$**Evalution metrics**$: Please check the General Response to All Reviewers a.

For the performance of $R^2$, we assessed the statistical significance of those results between FCR and second best methods using an paired t-test and showed the p-values with the corresponding sample size of the test set.
$**Table : p-value for$R^2$performance**$

The results indicate that FCR statistically significantly out-performed the baselines across datasets on sci-Plex, multiPlex-Tram and multiPlex-7.

We will incorporate the discussion of evaluation and t-test results in our final manuscript.

$**Generalizability**$:

We concur with the reviewer that generalizability is a very important topic in the computational biology field. Given the high cost of data collection, we are still limited to test and evaluate our models on the few public large scale single cell drug screening datasets that are available. As availability improves, there will be improved possibilities for benchmarking, especially in the out-of-distribution scenario.

$**Interpretability**$:

There are generally two approaches to interpreting latent representations, which we also plan to explore in future work:

1. Latent Traversal: Since $z_{tx}$ is element-wise identifiable, we can use latent traversal across all dimensions of $z_{tx}$. This involves fixing all latent variables $z_{tx}^i$ to their inferred values for each cell's gene expression and then varying the value of one latent at a time to observe the reconstructed gene expression. This method allows us to see how each dimension of $z_{tx}$ influences gene expression levels.
2. Interpretable Decoder System: We can also employ an interpretable decoder, such as LDVAE [1], which uses a linear decoder. By analyzing the weights of the latent representations assigned to each gene, we can interpret how different representations influence each gene.
3. On-going Work: In the light of the neural additive model [2], we are developing a non-linear interpretable encoder as an extension of LDVAE. Specifically, for each latent dimension, there is a sub-neural network to capture information specific to each latent dimension. A logistic regression layer is added to the end of the neural network. So, by analyzing the logistic regression weights (loading matrix), we can conclude how much one latent variable contributes to the reconstruction of a specific gene. An illustration of this model can be found in the rebuttal PDF file Figure D.

[1]Svensson, V., et al.(2020). Interpretable factor models of single-cell RNA-seq via variational autoencoders. Bioinformatics (Oxford, England), 36(11), 3418–3421.
[2]Agarwal, R., et al. (2021). Neural additive models: Interpretable machine learning with neural nets. Advances in neural information processing systems, 34, 4699-4711.