Removing Biases from Molecular Representations via Information Maximization

We appreciate the time and effort that the reviewer invested in reviewing our work. We are glad the reviewer valued the empirical evidence and technical innovation presented in our paper and recognized the key contributions of our work. Below, we offer detailed explanations in response to the reviewer’s questions.

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.

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

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.

We thank the reviewer for their valuable feedback. We are pleased that the reviewer appreciated the presentation and technical novelty of our paper and recognized the key contributions of our paper. In response to the reviewer’s comments, we updated the manuscript providing the following clarifications:

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.

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.

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.

Thanks for the response. Let me clarify.

1. I just realized that I misunderstood the setup and $p(z_g, z_d, x_b)$ is actually the distribution over constructed latents rather than true latens. In this case I agree it is tractable. I was saying justification as in the reasoning for the factorization, not technical justification. But again, such justification is not required if $L_{CLF}$ is tractable.

2. Your clarification is correct but does not address my concern. The weighting is different but that's a modification you made as a direct result of contribution 2, i.e. in order for contribution 2. to work, you must have reweighting. I do not agree that it is a standalone contribution just because it is different. What is the advantage of having an estimated weight instead of equal weight if you only consider samples from the same batch?

We thank the reviewer for carefully reading and helping us improve our manuscript.

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.

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

• 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.

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.

• Evidence of influence of batch identifier on the representation

The representation learned by CLIP is distorted by the batch effects (for example, the green batch cluster in the top right corner). The batch identifier's impact is also evident in the new prediction accuracy metric below.

As suggested by reviewer g7wR, we added several quantitative metrics in Figure 3 and Figure 4, with a detailed explanation of the metrics and results depicted in Appendix I. The prediction accuracy (%) from the latent representation to predict either batch identifier or real effect category is shown in the following table. InfoCORE has a clear advantage compared to CLIP and CCL, and the prediction accuracies to the batch identifier for CLIP and CCL are significantly larger than random guessing (4%), showing that they are influenced by the batch identifier.

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{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)$.

[1] Mutual information neural estimation. ICML 2018.

[2] Understanding the limitations of variational mutual information estimators. ICLR 2020.

[3] High-order conditional mutual information maximization for dealing with high-order dependencies in feature selection. Pattern Recognition 2022.

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.

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.

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.

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.

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.

Thank the authors for the detailed explanations, which resolve most of my concerns. Nevertheless, I still feel that the system of notations can be better defined. The novelty remains limited, given the limited performance gains. Therefore, I have adjusted my rating accordingly.

We thank the reviewer for their reassessment and acknowledgement of our detailed explanations, along with increasing their score. We are pleased to hear that these clarifications have addressed most of their concerns. We appreciate the suggestion about further improving our notation and will add more additional clarifications in the revised manuscript.

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.

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.

4. Question 1

Can the authors comment on Proposition 3? Is this bound tight?

The bound is tight when $\hat{p}(x_b^1|z_g^1)=\hat{p}(x_b^1|z_d^1)$.

5. Question 2

Empirical evidence of InfoCORE's adaptive adjusting procedure.

We observed evidence of InfoCORE’s adaptive reweighting schedule in practice.

• We tracked the entropy of the (normalized) reweighting factors during training. A higher entropy indicates more uniform reweighting. For example, in the GE dataset, the entropy of the reweighting factors increases from 9.99 to 10.26 (in comparison, CLIP, i.e. uniform, has entropy 11).
• We tracked the cross-entropy loss of the batch classifiers. A higher loss implies fewer batch-related features. For example, in the GE dataset, the cross entropy loss of the classifier based on drug structure representation first quickly decreases from 4.8 (random initialization) to 2.7 as we optimize the classifier’s parameters, and then slowly increases to 3.9.

6. Question 3

Performance under rare conditions.

We would like to clarify that in both drug screening datasets, experiments are done using 384-well plates. Thus there are only about 300 drugs in each batch (i.e., plate, since the remaining wells are used as controls) and hundreds of batches in total. Thus, every batch condition is “rare”. This is the reason for the poor retrieving accuracy of CCL over the drug library “whole”.

We thank the reviewer for proposing an additional experiment for further validation. We experimented with the CP dataset and downsampled the data to a third in 5 (out of 97) randomly selected batches that are correlated with other batches. The average retrieving accuracy within these rare batches is shown below. As expected, the performance drop of InfoCORE is about 3%-4% less than that of CCL after downsampling.

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.

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.

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.

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).

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.

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.

We appreciate the reviewer's detailed feedback, which helped improve our paper. The reviewer’s focus on our experiments inspired expanded discussions and enhancements in our revised manuscript and appendix. Detailed explanations of these changes are provided below.

1. About drug representation experiments

Experiments are underwhelming and benchmarks are generic approaches.

Although molecular representation, especially unimodal representation, is a well-explored area, multimodal molecular representation with drug screening data, which can be generalized to new molecular structures without drug screens, is a fairly new field. To the best of our knowledge, there are only a few prior works utilizing cell imaging data to learn molecular representations [1,2], and a few utilizing the gene expression [3,4].

[1] directly applied the generic CLIP method to drug structure and cell imaging data, and [2] applied a similar CLOOME criterion for contrastive learning with the additional masked graph modeling loss. [3] used an over-parametrized autoencoder, and [4] used the standard triplet loss on the L1000 drug screening gene expression data. While these datasets have significant batch effects, e.g., batch number can be predicted from gene expression data, none of these works considered the important batch confounding issue.

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

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

[3] Causal network models of SARS-CoV-2 expression and aging to identify candidates for drug repurposing. Nature Communications 2021.

[4] Predicting mechanism of action of novel compounds using compound structure and transcriptomic signature coembedding. Bioinformatics 2021.

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.

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.

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.

For clarity, my comment about the significance of the improvement was not meant to be understood as a judgement on how large the estimated improvement is compared to prior art, but as whether they are statistically significant (i.e. is the improvement really due to better performance of the method or due to the randomness of the evaluation / estimation process). In many cases, the estimated improvement is smaller than the standard deviation of the estimator, hence the lack of significance. I do agree that it can be very difficult to achieve a significant improvement in such a crowded task, though.

While I agree that 9 cell lines is insufficient to infer a general model that would extrapolate to new cell lines would be unreasonable, this is exactly my concern about the general approach that "learn(s) a generic drug representation that can be generalized to new drug structures". By not taking into account the perturbed cells in the modelling, you are implicitly trying to learn a representation that is universal, or at least agnostic to the cell line the drug is applied on. This is a much harder task that modelling the cell line and limiting the scope to predicting perturbations for those 9 cell lines.

Otherwise, I want to thank the authors for their careful reply and addressing the comments I made in my review. I think that the quality of the manuscript is now significantly better and allows for better reproduction of the results. I will improve my rating of the paper accordingly.

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.