Learning Molecular Representation in a Cell

1. Decoding for Multiple Neighbors

   Can I consider the decoding process as the procedure of autoregressively generating the sampled paths?

2. Finetuning

   It does not make sense to freeze the encoders as the original paper of UniMol[1] utilized full finetuning on its downstream tasks (some of the datasets even have less data than those used in the article), so do other molecular pretraining works such as [2]. More importantly, it does not make sense that the Morgan fingerprint outperforms most of the molecular pretraining models, as shown in Table 1 (even outperforming all on ChemBL2k), as this phenomenon challenges the meaning of the entire field of molecular pretraining. The performance of these molecular pretraining models has been severely underestimated. The authors should demonstrate the performance of the fully finetuned version and also update the results in response to "W2, Q4, feature concatenation".

[1] Zhou G, Gao Z, Ding Q, et al. Uni-mol: A universal 3d molecular representation learning framework[J]. 2023.

[2] Rong Y, Bian Y, Xu T, et al. Self-supervised graph transformer on large-scale molecular data[J]. Advances in neural information processing systems, 2020, 33: 12559-12571.

We sincerely appreciate the reviewer's thoughtful suggestions and questions. We have provided point-by-point answers to each weakness and question. We have also revised the main text and appendix to incorporate the reviewer's valuable feedback, with all changes clearly highlighted in blue for ease of reference. Should any concerns remain, we remain fully committed to addressing them promptly and thoroughly.

W1, Q1: Definition of neighbors

As stated in Eq (3) ($\sum_{v\in\mathcal{P}_x}$), the neighbor $v$ refers to the node in the random walk path. We have revised the text for better clarity.

W1, Q2 Decoding for multiple neighbors

Thanks for your question. If multiple neighbors share the same data type, such as cell morphology, the latent representation of the molecule is passed through a single decoder for that data type. Based on Eq. (3), the decoder's output is then aligned with the different features associated with $v_1, v_2, \dots, v_n$, with weights assigned according to $\alpha(v_i \mid \mathcal{P}_x)$ ($i=1, 2, \dots, n$).

W2, Q3: Generation of gene expression and cell morphology

Thanks for your question. Most explanations are provided in the Introduction and Related Work sections. We apologize for not referencing them when introducing the dataset and have revised the paper (Section 6.1.1, Appendix D.1) to address your concern. Below is a brief summary:

Molecules act as perturbations that yield perturbed cell states; those cell states can be measured in two ways relevant here: as gene expression values for a thousand or more genes [1] and/or microscopy Cell Painting images [2], which are represented by a thousand or more morphology features [3]. This is how a gene expression and morphology profile are associated with a given molecule.

Generating cell morphology and gene expression features requires extensive and costly experiments, so downstream tasks like ToxCast and Biogen3K may not always include these features. As described in Appendix D.1, ChEMBL2K is a subset overlapping with the existing JUMP-CP dataset [4], the largest cell painting dataset for cell morphology features. Relevant gene expression data for the molecules in ChEMBL2K are sourced from [6]. The available data for both types are detailed in Table 5 in the appendix.

W2, Q4: Feature concatenation

Thanks for your comments. We appreciate the helpful analysis and have now included it in Appendix C.3 and Table 7. The requested new results are provided in the table below:

We observe that concatenating UniMol representations with cell morphology and gene expression features improves performance in prediction tasks. However, it still does not match the performance of InfoAlign, which achieves the best results by aligning molecular representations with cell morphology and gene expression features during pretraining, rather than in the downstream stage.

W2, Q4: Fine-tuning

Thanks for your comment. For all pre-trained methods, the encoder is frozen during fine-tuning, and only the MLP is trained for prediction. We have now explained this in Section 6.1 and Appendix C.3.

Reference

[1] A Next Generation Connectivity Map: L1000 Platform and the First 1,000,000 Profiles. Cell. 2017.

[2] Three million images and morphological profiles of cells treated with matched chemical and genetic perturbations. Nature Method. 2024.

[3] Optimizing the Cell Painting assay for image-based profiling. Nature Protocols. 2023.

[4] JUMP Cell Painting dataset: morphological impact of 136,000 chemical and genetic perturbations.

[5] Drug-induced adverse events prediction with the LINCS L1000 data. Bioinformatics. 2016.

• The paper focuses on molecular representations, but if such a representation cannot outperform the features obtained from full fine-tuning, then what is the significance of such a kind of representation? If you have collected so much multimodal data and constructed complex pre-training tasks, but they still cannot surpass the work of others that use single-modality pretraining plus full fine-tuning, then why go through all this trouble to collect this data? I only need to use single-modality pre-training plus full fine-tuning, which is more in line with Occam's Razor.
• Moreover, single-molecule data is easier to collect and more abundant than biological data, making it easier for people to scale up their models. So, if your method cannot defeat an approach that only uses molecular pre-training plus full fine-tuning, why would anyone choose to use your method, which lacks both scalability and performance?

Thank you for your thoughtful comments and questions. To clarify, we would like to briefly reiterate the key points previously discussed:

• The primary focus of the paper is representation learning, with an emphasis on the quality of molecular representations.
• We have ensured fair evaluation by using consistent fine-tuning methods with MLP task predictors.

We are happy to continue the discussion on full fine-tuning, at the same time, we wish to maintain the focus on representation learning. To summarize the distinction:

• Full fine-tuning optimizes the encoder parameters for downstream tasks, which may lead to biased evaluation due to the use of stronger encoders.
• In contrast, representation learning aims to achieve high-quality representations even with simpler encoders, such as the GIN used in this paper.

Fully fine-tuning and representations can be studied separately. Advancements in both areas can contribute to the field. Real-world virtual screening is more complex than relying on a single representation or a fully fine-tuned model. Ensemble methods that combine both fine-tuning and representations offer advantages in both areas. However, understanding the rationale behind each component is essential. This paper focuses on representation learning with consistent, fair evaluations. In this context, InfoAlign has shown strong performance through extensive experiments in Table 1/2 and Figure 3, including 27 baselines.

As requested by the reviewer, we provide new results extending beyond the main contribution of representation learning, focusing on fully fine-tuning. The results show that the simpler GIN encoder in InfoAlign outperforms the Transformer used in UniMol, even in fine-tuning scenarios, highlighting the superiority of InfoAlign as a pretraining method.

• The supplementary experimental results in the authors' rebuttal have demonstrated that full-finetuning can significantly enhance the performance of the baselines. However, the paper does not present the performance of fully finetuned baselines. The presented performance of pretrained-GNNs is substantially underestimated. This means the paper lacks sufficient evidence to show that the proposed methods, which require more complex multi-modal data construction, demonstrate a clear advantage over simple single-molecule pretraining.
• Additionally, the supplementary experimental results indicate that after full finetuning, the proposed method only provides a marginal performance improvement on 3/4 of the datasets used in the paper compared to Uni-Mol, not to mention that Uni-Mol is not the most outstanding among all pretrained-GNNs. According to Occam's razor, which is agreed upon by both reviewers and authors, it is difficult to justify the necessity of the proposed complex multi-modal pretraining, especially when the collection of biological experimental data is much more challenging than that of single-molecule data.
• Even if the authors were to update the performance of all pretrained-GNNs, nearly half of the experimental results would need to be modified, and consequently, the experiment analysis would also require extensive revision. If the authors were to undertake such a large-scale rewrite of the paper, the review provided by the reviewer would no longer be applicable to this article, necessitating new reviewers and a new round of peer review, which would violate the ICLR submission process.

For the above three reasons, I will downgrade my score to reject.

Dear reviewer Sxqu,

I am a senior author on the paper. I greatly appreciate your taking the time to discuss our work.

The concerns that my lab member’s “viewpoint appears to dismiss the efforts of biochemistry professionals and is fraught with arrogance and a lack of understanding” are misplaced - our laboratory has actually led the experimental work for a large proportion of the publicly available image data of this type, so we are all aware of the challenges and value of lab work. I hope this is a simple mis-reading of my lab member’s response as I don’t see anything in our note to warrant this response.

Now back to the debate:

I believe your major concern this: if fully tuned models are doing better than our learned representations, then our learned representations are not terribly useful, and are even less useful if the representations are cumbersome to generate, i.e, they need data collected from wet lab experiments.

If that is indeed your primary concern, I think that's a valid one and I am happy to clarify our stance further, below.

Concern 1: If fully tuned models perform better than our learned representations, then our learned representations are not terribly useful

1. Representation learning supports a broader range of use cases. Given the vast number of unlabeled molecules, pre-trained representations can be efficiently stored and applied to various downstream tasks, such as visualization in virtual screening analysis.

2. In fully fine-tuned scenarios, representations from pre-trained models can still enhance performance. For example, on ChEMBL2K and Broad6K, InfoAlign’s representations (81.3±0.6 and 70.0±0.1) outperform fully-tuned UniMol (78.9±0.2 and 65.1±1.0) by 3.0% and 7.7%, respectively. In the NLP community [1], research shows that the usefulness of each paradigm depends on the specific downstream task. Therefore, both paradigms offer distinct advantages and can be combined. Representation learning does not conflict with full fine-tuning.

Concern 2: The representations are even less useful if the they are cumbersome to generate, i.e, they need data collected from wet lab experiments.

I think this is a fair criticism -- it's definitely a non-trivial overhead if one's approach needs data collected from wet lab experiments.

However,

1. The data we used is already publicly available; we simply tapped into existing resources.
2. Profiling assays like Cell Painting are now being routinely run by several academic labs and pharma companies, sometimes for quite large subsets of their compound libraries. Thus, in some contexts this data is freely available to their scientists for compounds of interest. What an amazing opportunity researchers have to tap into that, using InfoAlign-like methods!
3. InfoAlign can produce embeddings for a molecule even if we don't have a Cell Painting/gene expression profile for the molecule; hopefully this bit was already clear.

Let us know what you think, and thank you for remaining engaged!

[1] To tune or not to tune? adapting pretrained representations to diverse tasks. ACL. 2019.

We are pleased that you found our work interesting and appreciate your thoughtful comments and suggestions for improving the paper. We have revised the paper accordingly and provided a point-by-point response. We believe after revision, the current structure of the paper is well-balanced. If any concerns remain, we welcome further discussion to address them. Thank you again for your constructive feedback.

W1, Q3: Model clarity

The model setup is described in Line 307-312: we use a Graph Isomorphism Network (GIN) as the encoder and an MLP as the decoder. All details about model configurations, including layers and hidden dimensions, are provided in the supplementary code, along with a pretrained checkpoint for easy reproducibility. Below, we report the hyperparameters from the code:

The MLP consists of three layers: an input dimension of 300, a hidden layer with a dimension of 4 × 300, and an output layer corresponding to the feature/task dimension. This MLP architecture is generally applied on all representations.

We apologize for the oversight in referencing the code. We have updated Section 5 and Appendix B.3 with the relevant model information. In our attempt to stay anonymous, the complete code is available in the supplementary materials and we will include the code link in the camera-ready version.

W2: Impact of the MI bottleneck

This work is motivated by the principle of the information bottleneck, which leads to the optimization targets in Eq. (3) for molecular representation learning and pretraining.

The effectiveness of the bottleneck was evaluated by comparing InfoAlign with multi-modal contrastive alignment approaches.

For joint use of data in pretraining, Figure 1 and Lines 85–102 show that the proposed architecture is guided by the information bottleneck principle, which includes a single encoder with multiple decoders. Without this principle, jointly using two data types would be contrastive learning approaches, such as InfoCORE and CLOOME (Figure 1(a)). Comparing these baselines in Tables 1 and 2 thus allows us to assess the impact of the information bottleneck. We also note that methods for jointly integrating molecular, cell morphology, and gene expression data in contrastive learning remain underdeveloped.

We focus on applying the information bottleneck principle during the pretraining.

Joint use of data in downstream tasks, such as ToxCast and Biogen3K (Table 2), is often not feasible due to the high cost of obtaining cell morphology and gene expression data from biological experiments, compared to extracting molecular structure features. This is why we apply the information bottleneck during pretraining, rather than in downstream tasks. We agree that training models with MI constraints in downstream tasks, instead of using joint features, could be a promising approach. However, this is beyond the scope of our current focus on pretraining and may be explored in future work.

Q1: Results in table 1

Thank you for the question. There are 32 tasks for Broad6K, and the value 3.1 means that one task (3.125% = 1/32 * 100) is frequently predicted with high AUC (>90%). The task ID is 274_752. Using metadata from Broad6K, this task involves an MLPCN LGR2 assay, which aims to identify compounds targeting the LGR2 GPCR protein. This assay is an antagonist of the LGR2 target which will then compromise survival for pests like ticks or mosquitoes.

Regarding resolution, we observe that 3.1 occurs frequently with a threshold >85%, and the current resolution is sufficient to reflect the high AUC for this specific task.

Q4: Weighting of the paper

Thanks for your suggestions. Currently, the main text consists of 2.8 pages for background (Introduction, Related Work, problem definition), 3 pages for methods (1.7 pages for methods, 0.5 pages for theoretical outcomes, and 0.8 pages for model implementation), and 4.2 pages for experimental results.

For the theory, Section 4.3 presents the key theoretical outcomes in less than half a page, with further details in Appendix B (formerly Appendix A).

We have added more details on model construction, with references to the appendix for additional information. The code provides more clarity than text descriptions alone. We have included the code with checkpoints for easy reproducibility in the supplementary materials and will add the link in the camera-ready version.

We have updated Figures 3, 4, 5, and Table 4 for better clarity. We also included random walk analysis in Appendix D.5 and relocated some related content to the appendix to accommodate page limitations.

The current page allocation for background, methods, and experiments is 2.8:3:4.2, which we believe is well-balanced, with a emphasis on experimental results. We appreciate your comment and welcome any further discussion on adjusting the paper’s content for better presentation.

Q5: Figure 3

We have updated the caption for improved clarity. Figure 3 compares the relative performance of two groups of models: (1) representations using single-modal information (Single Rep.) and (2) molecular representations from multi-modal alignment methods (Aligned Rep.). The top bar compares models using single-modal information from the best baselines in Molecular Structure, Cell Morphology, and Gene Expression. The bottom bar compares three models that use multimodal information.

Q6: Table 3 right part

We have separated the figure and table, resulting in the current Table 3 and Figure 4. We also updated the distribution figure to be a histogram and added x-axis labels. We appreciate your suggestion, and with these changes, we believe Figure 4 now better illustrates the observations in Section 6.2.2.

Q7: Figure 4 (a)

In Figure 4(a) (now Figure 5 (a)), $\beta$ refers to the hyperparameter from the second term in Eq (3), not the learning rate. The figure shows how pretraining losses change with varying regularization strengths, which may not be as clearly represented in tables.

From the figure, we observe that pretraining loss may serve as an indicator for selecting $\beta$. Specifically, for $\beta = 1e-9$ and $\beta = 1e-12$, lower pre-training losses correspond to better downstream performance.

We have updated the figure to highlight the pretraining losses for $\beta = 1e-9$ and $\beta = 1e-12$ in the revision to address your concerns.

Q8: Figure 4 (a): walk length

Thanks for your comments. Each result is based on ten runs, with the points representing the mean value and error bars showing one standard deviation. We have updated the figure and caption for clearer representation.

Figure 4(b) (now Figure 5(b)) is primarily intended to support our claim of robust performance across different hyperparameter choices for $L$, as confirmed by the reviewers. As discussed in our response to W3/Q9, we do not observe any properties indicating that length 8 is an outlier. While the error bar may not fully capture the variance, we believe the variation at length 8 does not affect the comparison and observation with the best baseline.

Further discussions on the other questions about the figure are provided in the responses to W3/Q9.

We sincerely thank the reviewer for their insightful suggestions. We provide point-by-point responses and have revised the text and appendix, highlighting changes in blue.

W1: Random pre-specified graph

We conducted the requested experiments, and the results are shown in the table below. We find that randomly replacing edges significantly impacts prediction performance. The replaced edges lack the biological, chemical, and computational meaning of those used to construct the context graph in the paper.

W2: Genetic perturbation data

In Table 4, we observe a performance drop when excluding gene expression data in loss functions.

We also conducted new ablation studies, presented in Appendix D.4, with results shown in the table below. Removing all nodes related to genetic perturbation data from the context graph further decreases the performance, confirming the importance of genetic perturbation data.

W3: GROVER performance on ToxCast

The GROVER performance on ToxCast is reported as 53.1. We follow the standard splitting from Open Graph Benchmarking [1], which differs from the splitting used in the original GROVER paper [2]. Our results are consistent with those reported in [3].

W4: Ablation studies on removing data

We have conducted the requested ablation studies and have clarified them in Appendix D.4. The results are also shown in the table below. The first two rows display the removal of cell morphology or gene expression-related nodes from the context graph. We observe a further performance drop when these data are removed.

Q1, Q4: Performance drop in ablation studies

The performance of InfoAlign relies on the GNN encoder and different decoders with optimization targets.

As shown in Tables 1/2 and Figure 3, GNN encoders can extract meaningful representations with proper loss designs. Ablation studies in Table 4 demonstrate that, even without one type of data, the remaining two types still form proper targets based on information bottleneck principles, supporting the pretraining of a good GNN encoder.

We do not remove the GNN encoders, which continue to extract meaningful representations from molecular structures.

In the ablation studies, the "absence of molecular features" refers to the removal of fingerprint vectors from the loss functions. In this case, cell morphology and gene expression can still optimize the GNN encoder for meaningful representation, as observed in previous work like InfoCORE [4].

Q2 Relevant literature

We found that they were all published or released this year, including the most recent NeurIPS 2024 [6]. We are happy to discuss them and have updated the related work and appendix accordingly.

Lines 136-138: "CLOOME, MIGA [5], and MoCoP, and MolPhenix [6] contrast cellular images with molecules."

Line 774: "Approximating the mutual information of high-dimensional variables is a challenging task [7]"

Reference:

[1] Open Graph Benchmark: Datasets for Machine Learning on Graphs. NeurIPS. 2020

[2] Self-Supervised Graph Transformer on Large-Scale Molecular Data. NeurIPS 2020.

[3] Evaluating Self-Supervised Learning for Molecular Graph Embeddings. NeurIPS 2023.

[4] Removing Biases from Molecular Representations via Information Maximization. ICLR 2024.

[5] Cross-Modal Graph Contrastive Learning with Cellular Images. Advanced Science 2024.

[6] How Molecules Impact Cells: Unlocking Contrastive PhenoMolecular Retrieval. NeurIPS 2024.

[7] Approximating mutual information of high-dimensional variables using learned representations.

We sincerely appreciate the reviewer's thoughtful suggestions and questions. We have provided point-by-point answers to each weakness and question. We have also revised the main text and appendix to incorporate the reviewer's valuable feedback, with all changes clearly highlighted in blue for ease of reference. Should any concerns remain, we remain fully committed to addressing them promptly and thoroughly.

W1: Biological context

Thank you for recognizing the significance of our research problem. In the revision, we have further clarified the definition of the tasks (Appendix D.1) to reduce the barrier to understanding the biological knowledge prerequisites.

So far, we have: (1) formulated the problem as a machine learning task in Section 2, (2) explained the method and curated the dataset for ease of use in Sections 5.1, 6, and Appendices C and D, (3) provided a clear background on the motivation in the first paragraph of the Introduction, and (4) expanded the discussion of cellular response data in the related work section, Appendices C and D, with relevant references for readers interested in further details.

We hope these efforts will help readers in the ML community better understand the problem and encourage further solutions. We would also be happy to discuss additional strategies to further clarify the biological prerequisites in the main text and appendix.

W2: Edge weight

Edge weights are motivated by tasks in drug discovery.

Sorry for the confusion: We have now clarified this in the revision (Lines 198-201): "For example, edges derived from computational criteria between molecule nodes are assigned weights based on the assumption that structurally similar molecules may exhibit similar biological effects, a concept widely used in drug discovery, such as lead optimization."

Specifically, as introduced in Section 5, edges based on chemical and biological criteria have uniform weights (i.e., edge weights are set to 1). Edge weights are primarily applied to edges introduced by computational criteria, where we compute similarity using features, ranging from 0 to 1. For example, the similarity between Morgan fingerprints constructs weighted edges in the context graph. This approach is motivated by assumptions in tasks like lead optimization, where structurally similar molecules may exhibit similar biological effects. The weights quantify similarity, rather than assuming identical effectiveness for structurally similar molecules.

Edge weights are observed to empirically improve performance robustness.

In pretraining, the cumulative edge weights help avoid aligning a molecule with features from distant nodes on the context graph, especially when using longer walk lengths. In the table below, we present new experiments on ChEMBL2K using an unweighted context graph (all edge weights set to 1). With walk length $L=4$ and weighted edges, InfoAlign achieves 81.3±0.6. As the walk length increases, the performance with weighted edges is more stable and performs better than with unweighted edges.

W3: Remove noisy edges

Thanks for your suggestion. We conducted new experiments without using the mechanism described in Section 5 to remove noisy edges. As shown in the first row of the table below, performance decreases to varying degrees with more noisy edges, compared to using the mechanism with fewer noisy edges.

(The setting for "with more noisy edges": after computing the similarity, we applied a threshold of 0.5 to select the similar edges.)

Q1: Data requirements

InfoAlign does not require cell morphology and gene expression data for downstream tasks but improves InfoAlign pretraining

Cellular response data support the pretraining of InfoAlign. For downstream tasks, where only molecules are used as input and the output is a molecular representation, no additional data types are required.

We have updated Appendix D.4 with new experiments that removed either cell morphology or gene expression-related nodes from the context graph. The results are also given in the table below. The downstream performance drops when these nodes are removed. Although gene expression data is more sparse (about ten times less than cell morphology), it still contributes valuable information to InfoAlign’s performance. In summary, we believe that additional data from cell morphology and gene expression add value, and InfoAlign may perform worse with less data. Fortunately, we have curated multiple sources of cellular response data to ensure diverse and comprehensive data for pre-training.

Q2: Molecular features on context graphs

Thank you for the interesting question. In this work, we primarily use Morgan fingerprints to avoid introducing unnecessary factors that could influence the development and analysis of InfoAlign's pre-training strategy and model. Based on our observations (Table 1/2), Morgan fingerprints remain a competitive and cost-effective representation of molecular structures.

Improving InfoAlign with other pre-trained GNN representations or even their ensemble is indeed a promising direction. However, since this is outside the main focus of this work, we leave it for future exploration.

Q3: Computational Complexity

Thank you for the question. Compared to existing work, InfoAlign has cost dealing with the context graph. Pretraining with the context graph introduces minimal additional computational complexity. We use a sparse matrix to extract the random walk on the context graph. Let $N$ and $M$ denote the number of nodes and edges in the context graph, respectively, with an average node degree $k$, where $k \ll N$ and $M \ll N$. The time and space complexity of the random walk are $\mathcal{O}(k)$ and $\mathcal{O}(M)$, both much smaller than the dense version's $\mathcal{O}(N^2)$.

By using random walks to sample local neighborhoods for pretraining molecules, we can scale the context graph efficiently. Data size is not a major concern; the practical limitation lies in the limited number of cell morphology and gene expression features due to the high cost of generating them. Pretraining is efficient on a V100, using less than 13GB of GPU memory with a large batch size of 3072.