Paper: Removing Biases from Molecular Representations via Information Maximization (ICLR-2024)

Reviewer comment (Experiments):
It is not clear what is driving the algorithmic adjustments described in the Computational Considerations section, and how the choice of hyperparameters will impact the final results.

Author response:
>What is driving the algorithmic adjustments described in the Computational Considerations section?

The algorithmic adjustments are driven by the decomposition of the gradient of $L_{\text{CLIP}}$ with respect to the latent representations. By analyzing the gradient formulas, which we added in Appendix D, the gradient of $L_{\text{CLIP}}$ is driven by two components: one standard component containing $\frac{\partial f(z_g^1,z_d^i)}{\partial z_g^1}$ and $\frac{\partial f(z_g^1,z_d^i)}{\partial z_d^i}$ (and anagously $\frac{\partial f(z_d^1,z_g^i)}{\partial z_d^1}$ and $\frac{\partial f(z_d^1,z_g^i)}{\partial z_g^i}$) as denoted in the second line of each of the gradient formulas in Appendix D; and one competing component containing $\frac{\partial \hat{p}(x_b^1|z_g^i)}{\partial z_g^i}$ and $\frac{\partial \hat{p}(x_b^1|z_d^i)}{\partial z_d^i}$ as denoted in the first line of each of the formulas in Appendix D. 

The standard component optimizes the representations in a similar manner to CLIP, which makes the anchor and positive closer to each other and the anchor and negatives farther away in the latent space, but to different extents according to the reweighting factors in InfoCORE. The competing component optimizes the representations so that their predictive power of the batch identifier becomes lower. As demonstrated in Section 2.2.3, both gradient components drive the representations towards reduced batch effect. We use a hyperparameter $\lambda$ to regulate the extent to which each component contributes to the gradient update. In practice, since the competing component leads to a hard-to-optimize objective, we foud that using a smaller value of $\lambda$ generally works better. We added these explanations in Appendix D.

To further understand the effect of $\lambda$ on performance, we report the retrieving accuracy in the CP dataset with varying $\lambda$ in the table below.

|            |      | *whole* |       |       | *batch* |       |
|------------|------|-------|-------|-------|-------|-------|
| Top N Acc  | N=1  | N=5   | N=10  | N=1   | N=5   | N=10  |
| $\lambda=0$   | 7.19 | 20.47 | 28.34 | 13.38 | 38.77 | 53.13 |
| $\lambda=0.01$ | 7.29 | 19.97  | 27.41 | 13.28 | 38.00 | 52.66 |
| $\lambda=1/4$ | 2.47 | 8.10  | 12.45 | 12.29 | 37.29 | 52.96 |
| $\lambda=1/2$ | 1.44 | 4.93  | 7.34  | 11.90 | 34.87 | 50.07 |
| $\lambda=3/4$ | 0.65 | 2.17  | 3.44  | 8.77  | 28.89 | 44.65 |
| $\lambda=1$   | 0.60 | 2.36  | 3.57  | 8.62  | 28.57 | 44.83 |

When $\lambda$ is large, the model performance drops due to the unstable competing objectives.

Reviewer comment (Evaluation):
In the simulation data experiment section, it can be relatively difficult to evaluate the quality of the learned representation by examining only the first two principal components since the variance in other dimensions can be concealed. Visualization through manifold learning techniques such as t-SNE or UMAP, or evaluation using various metrics, can provide more reliable insights.

Author response:
>Visualization in the simulation experiment.

We would like to clarify that in the simulation study the representations are actually 2-dimensional and no PCA is used. The first two principal components were only applied to the first column of Figure 3 since the original input is 10-dimensional. To avoid confusion, we removed the first column in Figure 3.

In addition, as suggested by reviewer g7wR, we also added several quantitative metrics in Figure 3 and Figure 4, which complement our qualitative analysis. We added a detailed explanation of the metrics and results in Appendix I, and we provide the results of our analysis in the table below. The table shows the prediction accuracy (%) from the latent representation to predict either batch identifier or real effect category; InfoCORE clearly outperforms CLIP and CCL.

|                |              | CLIP       | CCL        | InfoCORE       |
|----------------|--------------|------------|------------|----------------|
| Drug screens representation   | real effect  | 44.2(7.0) | 45.1(4.6) | **73.3(5.0)** |
|  | batch id     | 8.6(2.6) | 7.8(1.9) | **5.6(3.6)** |
| Drug structure representation | real effect  | 61.4(11.7) | 48.0(10.2) | **77.8(6.0)** |
|  | batch id     | 14.4(4.5) | 11.7(3.4) | **5.4(2.4)** |

Reviewer comment (Experiments):
What is driving the algorithmic adjustment described in the computational considerations, what is the risk of doing so, and how the choice of the hyperparameter will impact the final results?

Author response:
>What is driving the algorithmic adjustments described in the Computational Considerations section?

The algorithmic adjustments are driven by the decomposition of the gradient of $L_{\text{CLIP}}$ with respect to the latent representations. By analyzing the gradient formulas, which we added in Appendix D, the gradient of $L_{\text{CLIP}}$ is driven by two components: one standard component containing $\frac{\partial f(z_g^1,z_d^i)}{\partial z_g^1}$ and $\frac{\partial f(z_g^1,z_d^i)}{\partial z_d^i}$ (and anagously $\frac{\partial f(z_d^1,z_g^i)}{\partial z_d^1}$ and $\frac{\partial f(z_d^1,z_g^i)}{\partial z_g^i}$) as denoted in the second line of each of the gradient formulas in Appendix D; and one competing component containing $\frac{\partial \hat{p}(x_b^1|z_g^i)}{\partial z_g^i}$ and $\frac{\partial \hat{p}(x_b^1|z_d^i)}{\partial z_d^i}$ as denoted in the first line of each of the formulas in Appendix D. 

The standard component optimizes the representations in a similar manner to CLIP, which makes the anchor and positive closer to each other and the anchor and negatives farther away in the latent space, but to different extents according to the reweighting factors in InfoCORE. The competing component optimizes the representations so that their predictive power of the batch identifier becomes lower. As demonstrated in Section 2.2.3, both gradient components drive the representations towards reduced batch effect. We use a hyperparameter $\lambda$ to regulate the extent to which each component contributes to the gradient update. In practice, since the competing component leads to a hard-to-optimize objective, we foud that using a smaller value of $\lambda$ generally works better. We added these explanations in Appendix D.

To further understand the effect of $\lambda$ on performance, we report the retrieving accuracy in the CP dataset with varying $\lambda$ in the table below.

|            |      | *whole* |       |       | *batch* |       |
|------------|------|-------|-------|-------|-------|-------|
| Top N Acc  | N=1  | N=5   | N=10  | N=1   | N=5   | N=10  |
| $\lambda=0$   | 7.19 | 20.47 | 28.34 | 13.38 | 38.77 | 53.13 |
| $\lambda=0.01$ | 7.29 | 19.97  | 27.41 | 13.28 | 38.00 | 52.66 |
| $\lambda=1/4$ | 2.47 | 8.10  | 12.45 | 12.29 | 37.29 | 52.96 |
| $\lambda=1/2$ | 1.44 | 4.93  | 7.34  | 11.90 | 34.87 | 50.07 |
| $\lambda=3/4$ | 0.65 | 2.17  | 3.44  | 8.77  | 28.89 | 44.65 |
| $\lambda=1$   | 0.60 | 2.36  | 3.57  | 8.62  | 28.57 | 44.83 |

When $\lambda$ is large, the model performance drops due to the unstable competing objectives.

Reviewer comment (Theory):
The computational considerations for $L_{CLF}$ should be illustrated. $L_{CLF}$ by itself is not tractable without the true sampling distribution $p(x_b | z_d)$ and $p(x_b | z_g)$, and is not exactly a classification loss. Without any clear statement that you're replacing $L_{CLF}$ with conditional density learning through classification loss, there isn't any direct justification for factoring $ \mathbb E_{p(z_d^1, z_g^1, x_b^1)} \log \frac{(\hat{p}(x_b^1 | z_d^1, z_g^1))^2}{p(x_b | z_d^1)p(x_b | z_g^1)}$ into $C - L_{CLF}$.

Author response:
>1. Computational considerations for $L_{\text{CLF}}$.

>$L_{\text{CLF}}$ by itself is not tractable and is not exactly a classification loss

Since $D_\text{KL} \left (p\left(x_b^1|z_d^1\right) \Vert \hat{p}\left(x_b^1|z_d^1\right) \right )=H(p(x_b^1|z_d^1))+CE(p(x_b^1|z_d^1), \hat{p}(x_b^1|z_d^1))$, $L_{\text{CLF}}$ can be factorized into the cross entropy loss of the classifiers $\hat{p}(x_b^1|z_d^1)$ and $\hat{p}(x_b^1|z_g^1)$, and an additional constant. 

>Justification for factorizing $\mathbb{E}_{p(z_d^1,z_g^1,x_b^1)} \log \frac{(\hat{p}(x_b^1|z_d^1,z_g^1))^2}{p(x_b|z_d^1)p(x_b|z_g^1)}$ into $C-L\_{\text{CLF}}$.

Note that using the following transformation, $\mathbb{E}_{p(z_d^1,z_g^1,x_b^1)} \log \frac{(\hat{p}(x_b^1|z_d^1,z_g^1))^2}{p(x_b|z_d^1)p(x_b|z_g^1)}$ can be factorized into $C-L\_{\text{CLF}}$:

$\mathbb{E}_{p(z_d^1,z_g^1,x_b^1)} \log \frac{(\hat{p}(x_b^1|z_d^1,z_g^1))^2}{p(x_b|z_d^1)p(x_b|z_g^1)} = \mathbb{E}\_{p(z_d^1,z_g^1,x_b^1)} \log \frac{(\hat{p}(x_b^1|z_d^1,z_g^1))^2}{\hat{p}(x_b|z_d^1)\hat{p}(x_b|z_g^1)} \cdot \frac{\hat{p}(x_b|z_d^1)}{p(x_b|z_d^1)} \cdot \frac{\hat{p}(x_b|z_g^1)}{p(x_b|z_g^1)}$

$= \mathbb{E}\_{p(z_d^1,z_g^1,x_b^1)} \log \frac{(\hat{p}(x_b^1|z_d^1,z_g^1))^2}{\hat{p}(x_b|z_d^1)\hat{p}(x_b|z_g^1)} - \mathbb{E}\_{p(z_d^1,x_b^1)} \log \frac{p(x_b|z_d^1)}{\hat{p}(x_b|z_d^1)} - \mathbb{E}\_{p(z_g^1,x_b^1)} \log \frac{p(x_b|z_g^1)}{\hat{p}(x_b|z_g^1)}$

The last two terms are the expectation of KL-divergence between the posterior probability and the estimated one:
$\mathbb{E}\_{p(z_d^1,x_b^1)} \left [ \log \frac{p(x_b^1|z_d^1)}{\hat{p}(x_b^1|z_d^1)}\right] = \mathbb{E}\_{p(z_d^1)} \left [ \mathbb{E}\_{p(x_b^1|z_d^1)} \left [ \log \frac{p(x_b^1|z_d^1)}{\hat{p}(x_b^1|z_d^1)}\right]\right] = \mathbb{E}\_{p(z_d^1)} \left [ \displaystyle D\_\text{KL}\left (p\left(x_b^1|z_d^1\right) \Vert \hat{p}\left(x_b^1|z_d^1\right) \right ) \right]$

$\mathbb{E}\_{p(z_g^1,x_b^1)} \left [ \log \frac{p(x_b^1|z_g^1)}{\hat{p}(x_b^1|z_g^1)}\right] = \mathbb{E}\_{p(z_g^1)} \left [ \mathbb{E}\_{p(x_b^1|z_g^1)} \left [ \log \frac{p(x_b^1|z_g^1)}{\hat{p}(x_b^1|z_g^1)}\right]\right] = \mathbb{E}\_{p(z_g^1)} \left [ \displaystyle D\_\text{KL}\left (p\left(x_b^1|z_g^1\right) \Vert \hat{p}\left(x_b^1|z_g^1\right) \right ) \right]$

Thus, $\mathbb{E}_{p(z_d^1,z_g^1,x_b^1)} \log \frac{(\hat{p}(x_b^1|z_d^1,z_g^1))^2}{p(x_b|z_d^1)p(x_b|z_g^1)} = 2\cdot(C-L\_{\text{CLF}})$.

To clarify these points, we updated the detailed derivations in Appendix B accordingly.

Regarding point 1, $p(z_g, z_d, x_b)$ is indeed the distribution over the constructed latents and thus $L_{\text{CLF}}$ is tractable. We added a sentence in the manuscript to clarify this further.

Reviewer comment (Presentation):
Line 2 and line 3 of Table 1 should be merged, or line 3 should be just removed. Reweighting of samples is an adjustment you have to make for large supply of negative samples to work, it is not a standalone contribution or an advantage. The current formulation of the table is slightly misleading.

Author response:
>2. Line 2 and line 3 of Table 1 should be merged.

We would like to clarify that Line 3 (Per-sample weighting) is not exactly the same as Line 2 (Large Supply of Negatives) since even if we only consider the samples from the same batch as in CCL, InfoCORE will weight each sample differently but CCL will weight them equally.

Regarding point 2, we understand the reviewer’s perspective and decided to remove line 3 from Table 1 in the updated manuscript.

Reviewer comment (Presentation):
Figure 1. does not demonstrate the classification workflow clearly. The figure and the yellow coloring on the right grid makes it seem like you're optimizing $M_{ii}$ w.r.t. batch identifiers.

Author response:
>3. Figure 1 does not demonstrate the classification workflow clearly.

We thank the reviewer for this comment. To clarify the figure, we removed the yellow coloring on the right grid in Figure 1.

Reviewer comment (Writing):
The paper is hard to follow.

Author response:
We thank the reviewer for taking the time to evaluate our manuscript and for the comments. We have carefully considered the reviewer’s comments and made revisions to clarify and improve the paper, especially regarding the clarity and exposition of our key arguments and methodologies. In summary, our paper focuses on learning molecular representations from multiple information sources, specifically drug screening data. InfoCORE is introduced to improve representations by explicitly countering confounding batch effects from non-biological sources. Such effects are typical also in other biological experimental data. This is a very different problem from supervised graph classification illustrated in [1,2,3,4,5] that the reviewer referred to. We provide point-by-point explanations to all of the reviewer’s comments below.

>**1. The paper is hard to follow:**

>Definition of the variables.

The variables $X_d, X_g, Z_d, Z_g, X_b$ are defined in Section 2.1. We added the definition of $Z_d^i, Z_g^i, X_b^i, \ i \in 1,...,K$ in Section 2.2.2. Variables with superscript $i, \ i \in 1,...,K$ are the multi-sample analog of those in the single-sample case. Similar to the definition in the multi-sample InfoNCE objective in contrastive learning, $i=1$ denotes the variables associated with the anchor-positive pair, and $i>1$ denotes the variables for the negative samples (drugs other than the one in the anchor-positive pair).

To further clarify the definition of all variables used in our paper, we added a lookup table in Appendix J.

>Definition of the problem.

The problem is defined in detail in the introduction as well as in Section 2.1). Briefly, the overall goal of our paper is to learn better molecular representations from different information sources. These representations can then be applied to various downstream tasks in drug discovery related prediction tasks. More precisely, we use the data from drug screening experiments to learn bioactivity related representations of drugs. Since batch effects dominate many biological experimental data, including drug screens, batch-related spurious features severely bias the learned representations. We translate the goal of learning de-biased representations into a conditional mutual information objective, and propose a lower bound criterion as an effective training objective. 

>Training objective and the algorithm of InfoCORE.

- The exact objective of InfoCORE is described in Equation 4, i.e. $L_{\text{InfoCORE}} = L_{\text{CLIP}} + L_{\text{CLF}}$, where $L_{\text{CLIP}}$ and $L_{\text{CLF}}$ are defined in Proposition 2. $L_{\text{CLIP}}$ resembles the symmetrical InfoNCE loss in CLIP, but with negative samples reweighted according to the posterior batch distribution estimated by the two batch classifiers ($\bar{D}$ and $\bar{G}$ in Figure 1). $L_{\text{CLF}}$ is the classification loss for training the classifiers to predict the batch identifier.

- The algorithm of InfoCORE is described in Section 2.2.2 and Section 2.2.3 and illustrated in Figure 1. Briefly, during training, the encoders and classifiers are updated iteratively. In each iteration, we first update the encoder parameters (i.e. the drug structure encoder 
$\text{Enc}\_{d}(.;\theta_d)$ and the drug screen encoder $\text{Enc}\_g(.;\theta_g)$) to optimize the latent representations using the new $L_{\text{CLIP}}$ objective, while keeping the classifier parameters fixed. Then we fix the encoders and update the parameters of the auxiliary batch classifiers $\hat{p}(x_b^1|z_g^1)$ and $\hat{p}(x_b^1|z_d^1)$ via optimizing the batch classification loss $L_{\text{CLF}}$. 

Based on the reviewer's comments, we updated the introduction and methods section in the revised manuscript to improve its readability and clarity. We welcome additional comments to further improve the clarity.

Reviewer comment (Writing):
All of the variables in the paper have not been defined.

Author response:
Definition of the variables.

The variables $X_d, X_g, Z_d, Z_g, X_b$ are defined in Section 2.1. We added the definition of $Z_d^i, Z_g^i, X_b^i, \ i \in 1,...,K$ in Section 2.2.2. Variables with superscript $i, \ i \in 1,...,K$ are the multi-sample analog of those in the single-sample case. Similar to the definition in the multi-sample InfoNCE objective in contrastive learning, $i=1$ denotes the variables associated with the anchor-positive pair, and $i>1$ denotes the variables for the negative samples (drugs other than the one in the anchor-positive pair).

To further clarify the definition of all variables used in our paper, we added a lookup table in Appendix J.

Reviewer comment (Writing):
The problem is undefined.

Author response:
Definition of the problem.

The problem is defined in detail in the introduction as well as in Section 2.1). Briefly, the overall goal of our paper is to learn better molecular representations from different information sources. These representations can then be applied to various downstream tasks in drug discovery related prediction tasks. More precisely, we use the data from drug screening experiments to learn bioactivity related representations of drugs. Since batch effects dominate many biological experimental data, including drug screens, batch-related spurious features severely bias the learned representations. We translate the goal of learning de-biased representations into a conditional mutual information objective, and propose a lower bound criterion as an effective training objective.

Reviewer comment (Writing):
It is unclear what is meant by claiming X_b as an irrelevant attribute.

Author response:
>Realistic examples for bias issue caused by the batching effects. What does it mean by claiming X_b as an irrelevant attribute?

Batch effects are common and challenging in biological experimental data. Substantial research has gone into mitigating such effects (e.g., [6,7,8,9,10]). Please also refer to the second paragraph in the related work section. Broadly speaking, batch effect is a change in the data that is caused by non-biological factors in an experiment [6], such as temperature, humidity, the person conducting the experiment, etc. Thus this effect is irrelevant to the real biological effect of a drug that we hope to identify in the representations $Z_d$ and $Z_g$.

To better explain this, we added the following sentence in the introduction: “Batch effects refer to non-biological effects introduced into the data through the experimental measurement process.” Specifically, in the context of drug screening, each “batch” refers to a (384-well) plate used in the experiments (384 drugs are tested at once in one plate; experiments on different plates are performed e.g. by different people, on different days, etc. and thus show non-biological variation, termed “batch effect”; however, complicating the situation is the fact that variation between plates is not always non-biological as we discussed in Section 2.1; this is due to the fact that similar drugs are often tested on the same plate and different cell lines are tested in different plates).

[6] Adjusting batch effects in microarray expression data using empirical Bayes methods. Biostatistics, 2007.

[7] Fast, sensitive and accurate integration of single-cell data with harmony. Nature Methods, 2019.

[8] Efficient integration of heterogeneous single-cell transcriptomes using scanorama. Nature Biotechnology, 2019.

[9] Batch effects in single-cell rna-sequencing data are corrected by matching mutual nearest neighbors. Nature Biotechnology, 2018.

[10] Deep learning enables accurate clustering with batch effect removal in single-cell rna-seq analysis. Nature Communications, 2020.

Reviewer comment (Theory):
It is unclear why equation (1) “emphasizes the shared features of the two modalities that are unrelated to batch” and “thus emphasizing the bioactivity of a drug.”

Author response:
> Why does equation (1) “emphasizes the shared features of the two modalities that are unrelated to batch” and “thus emphasizing the bioactivity of a drug”?

When maximizing the conditional mutual information objective, we are optimizing the representations $Z_d$ and $Z_g$ so that their batch-irrelevant shared information is maximized. Since we condition on $X_b$, we disallow any features tied to $X_b$ to contribute to the mutual information we are maximizing. Moreover, since the representations are kept in a fixed dimension, there is an explicit bias against incorporating any irrelevant attributes that do not contribute to mutual information. 

$X_g$ corresponds to drug screening data (for example, the gene expression change of cells after applying a drug) and thus reflects the bioactivity of a drug (i.e. the effect of a drug in the cellular context). Thus, after removing non-biological batch effects $X_b$, the shared information captures the bioactivity of a drug.

Reviewer comment (Writing):
The exact objective of InfoCORE is not clearly stated.

Author response:
>Training objective and the algorithm of InfoCORE.

- The exact objective of InfoCORE is described in Equation 4, i.e. $L_{\text{InfoCORE}} = L_{\text{CLIP}} + L_{\text{CLF}}$, where $L_{\text{CLIP}}$ and $L_{\text{CLF}}$ are defined in Proposition 2. $L_{\text{CLIP}}$ resembles the symmetrical InfoNCE loss in CLIP, but with negative samples reweighted according to the posterior batch distribution estimated by the two batch classifiers ($\bar{D}$ and $\bar{G}$ in Figure 1). $L_{\text{CLF}}$ is the classification loss for training the classifiers to predict the batch identifier.

Reviewer comment (Reproducibility):
The algorithm of InfoCORE is not described.

Author response:
>Training objective and the algorithm of InfoCORE.

- The algorithm of InfoCORE is described in Section 2.2.2 and Section 2.2.3 and illustrated in Figure 1. Briefly, during training, the encoders and classifiers are updated iteratively. In each iteration, we first update the encoder parameters (i.e. the drug structure encoder 
$\text{Enc}\_{d}(.;\theta_d)$ and the drug screen encoder $\text{Enc}\_g(.;\theta_g)$) to optimize the latent representations using the new $L_{\text{CLIP}}$ objective, while keeping the classifier parameters fixed. Then we fix the encoders and update the parameters of the auxiliary batch classifiers $\hat{p}(x_b^1|z_g^1)$ and $\hat{p}(x_b^1|z_d^1)$ via optimizing the batch classification loss $L_{\text{CLF}}$.

Reviewer comment (Presentation):
The information and caption of Figure 3 are hard to read. It seems the batching identifier has no influence on the learned representations for all methods.

Author response:
> The information and the caption of Figure 3 is hard to read. It seems the batching identifier has no influence on the learned representations for all methods.

Detailed descriptions of the simulation setup are given in the main text of Section 3.1. We updated the caption of Figure 3 and Section 3.1 to better explain the setting and results as detailed in response to the comments below.

- *On the Simulation setup*

In a nutshell, the scatter plots show the 2-dimensional latent representations of the simulated input data learned by each method. Each data point is labeled with one of five real effect categories and one of 25 batch identifiers. Top row figures color dots by real effect category, and bottom row figures by batch identifier. Therefore, an ideal representation should separate real effect categories well but not be impacted by the batch identifier.

Reviewer comment (Novelty):
The novelty is limited. It seems the main objective of InfoCORE comes from CCL. The differences between InfoCORE and CCL are unclear.

Author response:
>**2. The novelty is limited.** Differences between InfoCORE and CCL.

We respectfully disagree with the reviewer. Although both InfoCORE and CCL aim to maximize the conditional mutual information objective, computationally tractable estimation of conditional mutual information is an area that has been extensively researched and is known to be quite hard [1,2,3]. InfoCORE and CCL are based on **very different training objectives**. 

* CCL uses the InfoNCE loss with negative samples drawn from the conditional distribution to approximate the conditional mutual information:
$$L_{\text{CCL}} = \frac{1}{2} \left[ \mathbb{E}\_{p(z_d^1,z_g^1) p(z_d^{2:K}|x_b^1)}\left[- \log \frac{e^{f(z_g^1, z_d^1)}}{\frac{1}{K}\sum\_{i=1}^K e^{f(z_g^1, z_d^i)}}\right] + \mathbb{E}\_{p(z_d^1,z_g^1) p(z_g^{2:K}|x_b^1)}\left[-\log \frac{e^{f(z_g^1, z_d^1)}}{\frac{1}{K}\sum\_{i=1}^K e^{f(z_g^i, z_d^1)}} \right]\right]$$
However, it does not work well when there are limited observations for each condition (which is the case in our biological experiments, where there are only about 300 drugs in each batch and there are hundreds of batches in total).

* Instead, InfoCORE establishes a different lower bound for conditional mutual information, with $L_{\text{InfoCORE}} = L_{\text{CLIP}}+L_{\text{CLF}}$ as the training objective: 
$$
L_{\text{CLIP}} = \frac{1}{2}\left[ \mathbb{E}\_{p(z_d^1,z_g^1,x_b^1)p(z_d^{2:K})} \left [ -\log \frac{e^{f(z_g^1,z_d^1)}}{\frac{1}{K} \sum\_{i=1}^K e^{f(z_g^1,z_d^i)} \cdot \hat{p}(x_b^1|z_g^1,z_d^i)} \right ] \\ + \mathbb{E}\_{p(z_d^1,z_g^1,x_b^1)p(z_g^{2:K})} \left [ -\log \frac{e^{f(z_g^1,z_d^1)}}{\frac{1}{K} \sum\_{i=1}^K e^{f(z_g^i,z_d^1)} \cdot \hat{p}(x_b^1|z_g^i,z_d^1)} \right ] \right]
$$
$$L_{\text{CLF}} = \frac{1}{2} \left[ \mathbb{E}\_{p(z_d^1)} \left [ D\_\text{KL} \left (p\left(x_b^1|z_d^1\right) \Vert \hat{p}\left(x_b^1|z_d^1\right) \right ) \right] + \mathbb{E}\_{p(z_g^1)} \left [ D\_\text{KL} \left( p\left(x_b^1|z_g^1\right) \Vert \hat{p}\left(x_b^1|z_g^1\right) \right) \right] \right]$$
This objective is based on information theoretic derivation and can make use of all the negative samples by adaptively reweighting them according to reweighting factors $\hat{p}(x_b^1|z_g^1,z_d^i)$ and $\hat{p}(x_b^1|z_g^i,z_d^1)$.

Reviewer comment (Evaluation):
The improvements are limited. In Table 2, the improvements of InfoCORE over CLIP or CCL are marginal while no standard deviations are given. Even the vanilla objective CLIP can outperform InfoCORE.

Author response:
>**3. Performance.**

>- In Table 2, the improvements of InfoCORE are marginal while no standard deviations are given.

In Table 2, we calculate retrieving accuracy over two retrieval libraries, "*whole*" and "*batch*". It is important to note that accuracies **over both libraries as a whole** reflect the model’s ability of molecule-phenotype retrieval for drug discovery and drug repurposing. InfoCORE is **the only** model that has strong performance **in both libraries**. In comparison, CCL has very poor performance in "*whole*", and CLIP has significantly worse performance in "*batch*" (for example, InfoCORE has 15% gain in top 1 accuracy and 11% gain in top 5 accuracy over batch in the GE dataset compared with CLIP). We also highlighted this in the revised manuscript.

We also added the standard deviation of the retrieving accuracies based on 3 random seeds in Table 6 in Appendix I.

Reviewer comment (Evaluation):
In Table 3, InfoCORE has little-to-no improvements over baselines when considering the standard deviations.

Author response:
>- In Table 3, InfoCORE has little-to-no improvements over baselines.

It is known that the property prediction task in Table 3 is hard and even small gains on these benchmark datasets are significant. For example in reference [1], the proposed model's gains over the baselines are more marginal than in our paper, and in [2][3], the proposed model does not show consistent gains in most datasets.

[1] Uni-Mol: A Universal 3D Molecular Representation Learning Framework. ICLR 2023.

[2] Graph contrastive learning with augmentations. NeurIPS 2020.

[3] Let Invariant Rationale Discovery Inspire Graph Contrastive Learning. ICML 2022.

Reviewer comment (Novelty):
Some related works, including debiasing or invariant molecular representation learning such as [1,2,3,4,5], have not been compared or discussed.

Author response:
>**4. Related works.**
>Some related works including debiasing or invariant molecular representation learning such as [1,2,3,4,5] have not been compared or discussed.

These related works focus on learning _invariant_ or interpretable features in the _supervised_ graph classification task. However, as discussed in our problem setting, InfoCORE learns unbiased molecular representations in an _unsupervised setting_, where the only supervision comes from another paired high-dimensional data modality and is not explicitly related to the targets in the downstream tasks. We added a brief explanation contrasting InfoCORE and these related works in Section 4 summarizing the following discussion:

In [1-5], the label $Y$ is available during training and is the same as the label targeted in OOD generalization, while in InfoCORE, labels in the downstream tasks are not available during pretraining. Besides, [1-5] focus on separating the spurious features and invariant features of a graph structure (eg. certain subgraphs), while the spurious features in InfoCORE are caused by the batch effect in biological experiments, which is unrelated to graph structure. 

Although it is conceivable that the information objective in [2,3,4] could be adapted to our setting by replacing $Y$ with $X_g$, the adaption would be highly non-trivial. This is because $Y$ is the classification label while $X_g$ is a high-dimensional experimental data vector, and the goal of InfoCORE as a pretrained model is different from [2,3,4]. Additionally, one of the main contributions of InfoCORE is the tractable lower bound of conditional mutual information as the training objective.

In the following, we provide a detailed comparison of InfoCORE to each of the referenced works. 

[1] uses a distribution intervener on a graph to generate multiple interventional distributions. However, it is not applicable to our case since generating an interventional or counterfactual phenotype of a drug in a new batch is non-trivial and very hard to evaluate.

[2] optimizes the graph features to be independent with $E$. However, as shown in our graphical model in Figure 2, $X_b$ is a confounder between $X_d$ and $X_g$, thus the optimal features could still be correlated with $X_b$.

[3] optimizes a similar objective as [2], but constrains the conditional independence between $Y$ and $E$ conditioned on the learned graph features. However, in our graphical model, there exists a direct dependence between both modalities and $X_b$, and the correlation between $X_g$ and $X_b$ via non-biological batch effect remains even after taking the real biological effect into account.

[4] focuses on the interpretability of graph classifiers, which is a different problem. It selects the substructure of an input graph that is most relevant to the label $Y$, which minimizes the mutual information between the graph feature and the original graph. However, in our problem, we hope the representations maintain as much shared information between the two modalities as possible, as long as they are not through the backdoor path via $X_b$.

[5] considers a supervised learning setting similar to [2][3], and the algorithm is specifically designed for a graph as input data. Similar to [1], it requires counterfactual unbiased sample generation which is not applicable in our case. Also, the spurious correlation that we want to minimize in the representation is explicitly defined as the batch effect $X_b$, while that in [5] are implicit subgraphs and generally unrelated to batch effect.

Reviewer comment (Reproducibility):
The benchmarks only feature other generic approaches (i.e. CLIP and CCL) and the experimental details are often lacking to ensure a proper reproducibility of the results.

Author response:
>Experimental details are lacking.

We updated and significantly expanded Appendix F and Appendix H to better explain the experimental details. Specifically, in Appendix F, we now explain in detail each of the four adaptations we made to apply prior benchmarking models such as CLIP to drug screening data and we showcased the effectiveness of these adaptations as compared to the vanilla CLIP model via an ablation study. In Appendix H, we provided detailed explanations on the pretraining setup and how accuracy is calculated in the molecule-phenotype retrieval task, as well as the significance of accuracy across each retrieval library.

Reviewer comment (Experiments):
The initial simulation study has very little to do with drug effect, as the data used there is Gaussian noise (from various, controlled mixtures).

Author response:
**2. About simulation experiments**
> Simulation study has very little to do with drug effect.

We agree that the Gaussian distribution does not gene expression data or features from a cell imaging screen. The main purpose of the simulation experiment was to offer an intuitive illustration in a 2-dimensional representation space to show how InfoCORE works to remove biases and extract the real effects (the real effect categories are available since this is a simulation study).

Reviewer comment (Evaluation):
The results of this first part are purely qualitative (plots of low-dimensional representations of the learnt embedding) and are not particularly convincing compared to CLIP or CCL. Quantitative metrics demonstrating that InfoCORE representations preserve the original structure of the data (e.g. neighbourhoods) within batches but mix them well across batches after integration would have strengthened the case (see for instance the scANVI paper from Wu et al. 2021).

Author response:
> Quantitative metrics should be added.

We thank the reviewer for this suggestion. We added the following quantitative metrics in addition to the qualitative illustrations. Following the metrics used in scANVI (Wu et al. 2021), we calculated the prediction accuracy (from the latent representation to predict either batch identifier or real effect category), KNN purity, and entropy of batch mixing for each model. The prediction accuracy (%) is as follows. 

|                |              | CLIP       | CCL        | InfoCORE       |
|----------------|--------------|------------|------------|----------------|
| Drug screens representation   | real effect  | 44.2(7.0) | 45.1(4.6) | **73.3(5.0)** |
|  | batch id     | 8.6(2.6) | 7.8(1.9) | **5.6(3.6)** |
| Drug structure representation | real effect  | 61.4(11.7) | 48.0(10.2) | **77.8(6.0)** |
|  | batch id     | 14.4(4.5) | 11.7(3.4) | **5.4(2.4)** |

We updated the results with all three metrics in Figure 3 and Figure 4. We also updated Section 3.1 and Appendix I to better explain the metrics and results. InfoCORE outperforms CLIP and CCL, especially in prediction accuracy of real effect category and KNN purity.

Reviewer comment (Experiments):
The second part tackles a specific (somewhat unusual maybe?) task of identifying drugs that are likely to induce a given effect, but the details describing this task are lacking even after going through appendices.

Author response:
>Details of task in Table 2 are lacking and results are not significantly better.

In Table 2, the molecule-phenotype retrieval task is a common and essential task corresponding to drug repurposing, which is also evaluated in references [1-3]. Here, we calculate retrieving accuracy over two retrieval libraries, "*whole*" and "*batch*"; "*whole*" comprises all molecules in the held-out set, reflecting the model's overall capacity to match drugs with their effects; in contrast, "*batch*" contains only those molecules from the held-out set sharing the same batch identifier as the target drug, gauging the model's discriminatory ability without batch confounder features. For instance, a model solely reliant on batch-related features might perform well in "*whole*" by randomly choosing drugs from the correct batch, but would likely falter in "*batch*".

Therefore, accuracies **over both libraries as a whole** reflect the model’s ability of molecule-phenotype retrieval for drug discovery and drug repurposing. We updated Section 3.2 and Appendix H accordingly to better explain the setting.

It is important to note that InfoCORE is **the only** model that has strong performance **in both libraries**. In comparison, CCL has very poor performance in "*whole*", and CLIP has significantly worse performance in "*batch*" (for example, in GE, InfoCORE outperforms CLIP over “*batch*” with a 15% increase in top 1 accuracy and 11% in top 5 accuracy). 

We also added the standard deviation of the retrieving accuracies over 3 random seeds in Table 6 in Appendix I. 

[1] Molecule-Morphology Contrastive Pretraining for Transferable Molecular Representation. bioRxiv, 2023.

[2] Cross-modal Graph Contrastive Learning with Cellular Images. bioRxiv, 2022.

[3] CLOOME: contrastive learning unlocks bioimaging databases for queries with chemical structures. bioRxiv, 2022.

Reviewer comment (Experiments):
The quantitative results in the second part do not seem to be significantly better than other methods, and the experimental setting seems a bit convoluted as it requires adapting the methods benchmarked against in a non-trivial way.

Author response:
> The experimental setting seems convoluted and it requires adapting the benchmark methods in a non-trivial way. 

This is correct and is due to the fact that there are to our knowledge no practically effective multimodal contrastive learning frameworks designed for drug screening data. 

While it is not the main focus of our paper, we had to make several novel non-trivial and empirically effective modifications to the generic CLIP and CCL models to adapt and enhance them for drug screening applications, including data augmentation ("*aug*"), cell line specific projection head ("*cellproj*"), training batch sampled by cell line ("*train bycell*"), and adapting MoCo to drug screens data ("*MoCo*"). We updated Appendix F to better explain these adaptations. We also performed an ablation study in the GE dataset to demonstrate the effectiveness of these adaptations.

| Retrieval Library       |      | _whole_ |       |       | _batch_ |       |
|-------------------------|------|---------|-------|-------|---------|-------|
| Top N Acc (%)           | N=1  | N=5     | N=10  | N=1   | N=5     | N=10  |
| InfoCORE                | **6.48** | **19.13**   | **27.53** | **14.16** | **33.93**  | **47.15** |
| CLIP (reported)         | 6.01 | 18.64   | 27.16 | 12.40 | 30.47   | 42.77 |
| CLIP (w/o *MoCo*)         | 5.71 | 17.45   | 26.14 | 11.57 | 28.82   | 41.32 |
| CLIP (w/o *aug*)          | 5.25 | 16.00   | 23.99 | 12.25 | 29.74   | 41.88 |
| CLIP (w/o *cellproj*)     | 5.39 | 17.36   | 25.77 | 11.62 | 29.58   | 41.78 |
| CLIP (w/o *train bycell*) | 5.02 | 16.45   | 24.90 | 10.74 | 28.16   | 40.66 |
| vanilla CLIP            | 4.64 | 14.86   | 22.16 | 11.32 | 28.46   | 40.55 |

Note that the results of CLIP and CCL reported in Table 2 and Table 3 use these adaptations. Compared with vanilla CLIP, InfoCORE has even larger improvements.

Reviewer comment (Evaluation):
In all tables, most values written in bold do not seem to be significantly better, which is problematic.

Author response:
> Results in Table 3 are not significantly better.

It is known that the property prediction task in Table 3 is hard and even small gains on these benchmark datasets are significant. For example in reference [1], the proposed model's gains over the baselines are more marginal than in our paper, and in [2][3], the proposed model does not show consistent gains in most datasets.

[1] Uni-Mol: A Universal 3D Molecular Representation Learning Framework. ICLR 2023.

[2] Graph contrastive learning with augmentations. NeurIPS 2020.

[3] Let invariant rationale discovery inspire graph contrastive learning. ICML 2022.

Reviewer comment (Writing):
The paper avoids addressing what exactly is intended by "batch effect" or "confounding factors".

Author response:
**3. About problem setting**
> The paper avoids addressing what exactly is intended by "batch effect" or "confounding factors".

We added the following sentence in the introduction: “Batch effects refer to non-biological effects introduced into the data through the experimental measurement process.” See e.g. [1] as well as our explanations in Section 2.1. Specifically, in the context of drug screening, each “batch” refers to a (384-well) plate used in the experiments (384 drugs are tested at once in one plate; experiments on different plates are performed e.g. by different people, on different days, etc. and thus show non-biological variation, termed “batch effect”; however, complicating the situation is the fact that variation between plates is not always non-biological as we discussed in Section 2.1; this is due to the fact that similar drugs are often tested on the same plate and different cell lines are tested in different plates).

[1] Adjusting batch effects in microarray expression data using empirical Bayes methods. Biostatistics, 2007.

Reviewer comment (Presentation):
The treatment of the fact that LINCS data features drug screening of different cell lines (and whether it is considered a batch effect) is relegated to the appendices, despite being a crucial question.

Author response:
> The different cell lines screened in LINCS data is a crucial question.

We agree that identifying cell line specific drug effects and their interaction is an important question. However, the main focus of this paper is to learn a generic drug representation that can be generalized to new drug structures without performing additional drug screens. Also, note that screening was only performed in one cell line in the CP dataset, and only 9 cell lines have sufficient data in the GE dataset. This is not sufficient to extrapolate to new cell lines and we therefore do not consider this task.

Reviewer comment (Evaluation):
Methods that predict drug perturbation without modelling explicitly the state of the cells the perturbations are measured on are of limited impact, as it is impossible to extrapolate to cells from other cell lines not present in LINCS or to real tissues.

Author response:
We appreciate the reviewer's constructive feedback and their increased score. We are glad our revisions have addressed most of their concerns. The reviewer's point about considering perturbed cells for understanding drug mechanisms in different contexts is valid, and we will clarify our model's scope in relation to cell lines in our discussion.