Paper: Multimodal Molecular Pretraining via Modality Blending (ICLR-2024)

Reviewer comment (Experiments):
As highlighted by the authors, earlier studies typically aligned different modalities at a broader molecule level, potentially hindering the capture of detailed molecular structures. Initially, when considering atom-level depictions of a molecule, our focus might be on the atom-level rather than the atomic-relation level. For instance, we might design a model that determines 3D atom coordinates based on a 2D molecular graph. What's the rationale behind emphasizing atomic-relation level multi-modal pretraining? It would be beneficial to provide a thorough reasoning in the methodology section and draw empirical comparisons in the experiments.

Author response:
We express our gratitude to the reviewer for providing valuable feedback. Below we address the comments in detail.

**Q1: Why atomic-relation level rather than atomic-relation level.**

We thank the reviewer for the question.

In our analysis and intuitive insights, different modalities share the identical atoms (C, N, O, etc) and **diverse atomic relations**. Upon closer examination, the modalities mainly differ in how they depict **atom relations**, while the atoms are commonly shared. To this end, to align different modalities, we propose that the key to aligning different modalities lies in aligning the distinct atom relations that encapsulate richer information from various perspectives, rather than focusing solely on the identical atoms.

**Predicting atom coordinates implicitly learns atom relations.** Regarding the coordinates mentioned, we posit that the individual atom coordinates themselves may not hold primary significance. Instead, the **relation/interplay between atoms (i.e., atom distance in 3D space, the structure)** imbues the molecule with diverse properties [1]. While predicting all coordinates collectively offers a degree of insight into structures and implicitly learns the atom relations, we contend that focusing on atom relations, specifically the 3D directed distances, **more directly and accurately convey the structural information**.

At present, we recognize that the effectiveness of predicting atom coordinates lies in its implicit learning of structure and atom relations. The rationale behind the superiority of atom-level predictions over atom-relation level remains elusive. To this end, experiments exploring atom-level alignment falls beyond the scope of our current paper and can be left as future work.

If any aspects require further clarification, please feel free to communicate them to us. We are eager to engage in additional discussions and appreciate your consideration. Thank you.

[1] Structure-based drug design with geometric deep learning.

Reviewer comment (Evaluation):
The experimental findings appear to be underwhelming. MoleBlend's performance, as shown in Tables 1 and 2, doesn't seem to fare well against prior studies.

Author response:
**Q2: The experimental findings appear to be underwhelming.**

Our method achieves state-of-the-art results across 16 out of 23 mainstream tasks evaluated. This is a substantial rather than marginal improvement. Notably, in Table 1, we achieve an average improvement of 1.09 point across all 8 MolNet tasks over the prominent SOTA MoleculeSDE [1] from ICML 2023. This improvement is obvious and far from being considered marginal.

Besides, different from current 2D+3D baselines [1][2][3] that all rely on strong established methods like contrastive learning to enhance their performance, our method is a de novo endeavor -- the first attempt to align modalities at atom-relation level. We think these results are particularly encouraging, considering our distinctive perspective that offers a fresh outlook on the problem.

[1] A group symmetric stochastic differential equation model for molecule multi-modal pretraining. ICML 2023

[2] Pretraining molecular graph representation with 3d geometry. ICLR 2022

[3] 3d infomax improves gnns for molecular property prediction. ICML 2022

Reviewer comment (Evaluation):
The authors did not discuss the training time cost of different pretraining methods.

Author response:
We thank the reviewer for the comment. The computation required are listed below. We can observe that GNN based pretrainig methods have relatively small compute cost. Transformer based pretraining methods cost a lot more computation.

| Method      | Pretraining Time Cost     |
| ----------- | ------------------------- |
| AttrMask    | 9.2h on one V100 GPU      |
| ContextPred | 23.3h on one V100 GPU     |
| MolCLR      | 16.7h on one V100 GPU     |
| InfoGraph   | 10h on one V100 GPU       |
| GROVER      | 2.5 days on 250 V100 GPUs |
| 3D Infomax  | Not reported              |
| GraphMVP    | Not reported              |
| MoleculeSDE | 50h on one V100 GPU       |
| Ours        | 24 hours on 8 A100 GPUs   |

Reviewer comment (Experiments):
A series of pertaining baselines are missed in related works and comparisons. For example:  
[1] Xu M, Wang H, Ni B, et al. Self-supervised graph-level representation learning with local and global structure. ICML 21.  
[2] Zhang Z, Liu Q, Wang H, et al. Motif-based graph self-supervised learning for molecular property prediction. NeurIPS 21.  
[3] Zaidi S, Schaarschmidt M, Martens J, et al. Pre-training via denoising for molecular property prediction. ICLR 23.

Author response:
We thank the reviewer for the reference. [1][2] have been added into our paper for comparison. 

We do not compare with the number in [3] as they use a much stronger model backbone that excels at 3D property prediction, which even outperforms all Transformer-based pretraining method without pretraining. 

However, we implement the same denoising idea [3] on our backbone for comparison. In Table 3, where the first row is denoising (Noisy Node [4] is identical to [3] when applied to pretraining, as discussed in Section 2 and 3.2.2 of [3]), the second and the third row are our methods. The results show that our method achieve better results across a broad range of tasks.

[1] Self-supervised graph-level representation learning with local and global structure. ICML 21.

[2] Motif-based graph self-supervised learning for molecular property prediction. NeurIPS 21.

[3] Pre-training via denoising for molecular property prediction. ICLR 23.

[4] Simple GNN regularisation for 3d molecular property prediction and beyond. ICML 2022

Reviewer comment (Theory):
The theoretical analysis is good. However, could the authors provide more insights into why MOLBLEND overperforms existing methods theoretically?

Author response:
Theoretically, the superiority of MoleBLEND lies in the more comprehensive and fine-grained optimization target for molecular modeling when compared with existing methods.

- More comprehensive: as evidenced by Proposition 3.2, MoleBLEND can be viewed as a comprehensive target that encompasses single-modal masking, cross-modal contrastive and predictive methods, with their conditional versions. Therefore, MoleBLEND is able to learn cross-modal correspondences as well as single-modal context dependencies, which are both essential for multimodal molecular modeling. However, many existing methods focus solely on one aspect. For example, 3D InfoMax does not explicitly model the single-modal context dependencies.
- More fine-grained: In the perspective of mutual information maximization, cross-modal contrastive and predictive methods aim at $\max I(A,B)$, where $A$ and $B$ are two modalities of the input. In contrast, MoleBLEND performs fine-grained modality alignment in Eq.15, i.e. $\max I(A_1,B_1)$, where $A_1$ and $B_1$ denote two modalities of the randomly partitioned input. This also makes the former objective a special case of ours. Fine-grained alignment encapsulates the cross-modal correspondence of molecular substructures, a critical characteristic of molecules that is rarely captured by previous methods.

Reviewer comment (Experiments):
How does the different choice of encodings for 2D/3D modalities influence the pretaining?

Author response:
To the best of our knowledge, 2D shortest path, edge type, 3D euclidean distance are the most commonly used encodings for molecules. The only other encoding we know is molecule shape [1], but this encoding cannot be effectively represented in terms of atoms and their relationships, thus it is not compatible for our method. We acknowledge the potential existence of other encodings and are open to exploring their impact on pretraining. Should you have relevant references, we are willing to delve into further learning and are prepared to conduct swift experiments to assess their influence.

[1] Learning Harmonic Molecular Representations on Riemannian Manifold. ICLR 2023

Reviewer comment (Writing):
The motivations are not clearly claimed or supported. For instance, on Page 2, the authors say that they “observe that although appearing visually distinct … are intrinsically equivalent as they are essentially different manifestations of the same atoms and their relationships”. What does “equivalent” mean here? A lot of 2D-3D pretraining methods start by saying such two modalities are complementary to each other.

Author response:
**Q1: What does 2D and 3D 'equivalent' mean? Motivation to feed both modalities as one unified data structure to one single model.**

They are indeed complementary as existing 2D-3D pretraining methods claim. Furthermore, through our analysis and intuitive insights, we claim that they are essentially equivalent and just appear different. 

Different modalities share the same underlying structures, i.e., atoms and their relations. Different modalities share the identical atoms (C, N, O, etc). The 2D and 3D structures (edge type and 3D distance) are essentially different depictions of the **same** atom relations. To this end, we treat them as essentially equivalent. For example, a C-N bond in a 2D molecular graph and the corresponding 3D distance between C and N atoms essentially represent **the same thing**.

Given they are essentially equivalant and complementary, these modalities should be amalgamated to form a more comprehensive information set. Consequently, we advocate for feeding these modalities as one unified data into one single model, to learn a comprehensive representation of molecules.

If there are any aspects that need additional clarification, please let us know. We would be delighted to engage in further discussions. Thank you.

Reviewer comment (Theory):
On Page 2, what is the motivation to feed both modalities as one unified data structure to one single model in MoleBlend?

Author response:
**Q1: What does 2D and 3D 'equivalent' mean? Motivation to feed both modalities as one unified data structure to one single model.**

They are indeed complementary as existing 2D-3D pretraining methods claim. Furthermore, through our analysis and intuitive insights, we claim that they are essentially equivalent and just appear different. 

Different modalities share the same underlying structures, i.e., atoms and their relations. Different modalities share the identical atoms (C, N, O, etc). The 2D and 3D structures (edge type and 3D distance) are essentially different depictions of the **same** atom relations. To this end, we treat them as essentially equivalent. For example, a C-N bond in a 2D molecular graph and the corresponding 3D distance between C and N atoms essentially represent **the same thing**.

Given they are essentially equivalant and complementary, these modalities should be amalgamated to form a more comprehensive information set. Consequently, we advocate for feeding these modalities as one unified data into one single model, to learn a comprehensive representation of molecules.

If there are any aspects that need additional clarification, please let us know. We would be delighted to engage in further discussions. Thank you.

Reviewer comment (Writing):
Notations are misleading in Sec 3.1. Why is \(R_{spd}\) required? Because they can be derived from \(R_{edge}\)?

Author response:
**Q2: Notations Clarification**

> Why is R_spd required? Because they can be derived from R_edge?

We would like to clarify that R_spd and R_edge are two different types of information. R_spd is the shortest path between two atoms on a molecular graph, reflecting the topology information of the molecule. R_edge is the edge type (e.g. single bond), reflecting the chemistry property. Providing information from different perspectives can enhance model understanding of molecules.

Reviewer comment (Writing):
For Eq 1, the notation should be \(R_{2D3D,S}\) (with S in the subscript).

Author response:
**Q2: Notations Clarification**

> For Eq 1, the notation should be (with S in the subscript).

We thank the reviewer for the advice. We have updated this notation in our paper.

Reviewer comment (Writing):
Is \(R_{2D3D}\) the masked version of \(R_{spd}, R_{edge},\) and \(R_{distance}\)?

Author response:
**Q2: Notations Clarification**

> Is R_2d3d the masked version of R_spd, R_edge, R_distance?

Yes, it is the blended and masked version of R_spd, R_edge, and R_distance.

Reviewer comment (Theory):
The distance modeling is invariant. A more advanced equivariant model is preferred here.

Author response:
**Q3: Equivariant Head**

We thank the reviewer for pointing this out and the advice. We conduct some preliminary attempts and transform our 3D prediction head to an equivariant head.

Originally, to predict the 3D distance between atoms, we perform outer product operations between atoms representation, followed by flattening and projection operations to map them into the target space to predict the $N^2$ atom relations:

$$o_{ij} = \text{G}(W_{l} X_i^{L+1}) \otimes \text{G}(W_{r} X_j^{L+1})^\top \in \mathbb{R}^{m\times m}$$

$$z_{ij} = W_{\text{head}}\cdot \text{Flatten}(o_{ij}) \in \mathbb{R}^{c}$$

Following [1], in order to achieve equivariant 3D relative position prediction, we decompose the flattened relational representation into three directions by multiplying the normalized relative position offset $\Delta_{i,j}=\frac{r_{i,j}}{\|\|r_{i,j}\|\|_2}$ and apply linear projection head to each component of the 3D relational representation in their respective direction:

$$o_{ij} = \text{G}(W_{l} X_i^{L+1}) \otimes \text{G}(W_{r} X_j^{L+1})^\top \in \mathbb{R}^{m\times m}$$

$$z_{ij}^{k} = W_{N}\cdot \Delta_{i,j}^{k} \text{Flatten}(o_{ij}), \quad  k=0,1,2$$

where $\text{Flatten}(o_{ij})$ is the flattened relational representation, $\Delta_{i,j}^k$ is the k-th element of the directional vector $\frac{r_{i,j}}{\|\|r_{i,j}\|\|_2}$ between atom $i$ and atom $j$ and $W_N\in \mathbb{R}^{1\times d}$ is learnable weight matrices.

We pretrain for 200K steps and compare it apple-to-apple with blending. The experiment results are shown below. Current equivariant head brings performance boost on 2 out of 4 tasks evaluated.

| Method | BBBP | BACE | Tox21 | ToxCast |
| --- | --- | --- | --- | --- |
| Blending  | 71.68 | 83.41 | 76.58 | 65.46 |
| Blending w/ Equivariant head | 72.60 | 81.55 | 76.41 | 66.03 |

Please note that this is just a preliminary attempt, and further investigation can be left as future work.

[1] One transformer can understand both 2d & 3d molecular data. ICLR 2023

Reviewer comment (Presentation):
It would be better to explicitly add a column in the result tables on what backbone models are pretrained/comparing.

Author response:
**Q5: Other improvements**

> It would be better to explicitly add a column in the result tables on what backbone models are pretrained/comparing.

> The citation of 3D InfoGraph is wrong. Please fix it.

We thank the reviewer for these suggestions. We have uploaded an updated version of our paper that addresses these questions.

Reviewer comment (Presentation):
The citation of 3D InfoGraph is wrong. Please fix it.

Author response:
**Q5: Other improvements**

> The citation of 3D InfoGraph is wrong. Please fix it.

We thank the reviewer for these suggestions. We have uploaded an updated version of our paper that addresses these questions.

Reviewer comment (Presentation):
I am confused about Fig 1.b. Does this mean the input can be either 2D or 3D, and the output can be both 2D and 3D?

Author response:
**Q4: Explanation of Fig 1 b**

The output is still either 2D or 3D. The work of Fig 1 (b) is [1]. Their multimodal pretraining task involves taking a 2D input and predicting its corresponding 3D information, or taking a 3D input and predicting the associated 2D information, which corresponds to Fig 1 (b).

[1] Unified 2d and 3d pre-training of molecular representations.

Reviewer comment (Novelty):
The information gathered in the relation matrix is quite limited and much information in the original structure is lost, especially those in the 3D structure. The matrix construction is quite similar to the work of "One transformer can understand both 2d & 3d molecular data" published in ICLR 2023.

Author response:
**Q1: The information gathered in the relation matrix is limited and lost a lot.** 

We thank the reviewer for the question. We would like to clarify that no information is lost in our method in both the pretraining and finetuning stage.

Please note that in modality-blending, we are performing self-/un-supervised pretraining. The learning task is to leverage the blended relation matrix to **predict non-blended elements**. Although some information in the original structure is not blended and lost, these elements will be **used as targets for prediction**. This supervision guarantees no information is lost in the process.

During supervised finetuning and inference stage, we **utilize all available modalities without random blending, where no information in the original structure is lost**. The choice of modalities depends on the specific scenario at hand:

- When a large amount of 2D molecular graphs are available while 3D conformations are too expensive to obtain, the model can be finetuned and perform inference with only 2D inputs.
- When it comes to scenarios where 3D information is available, we propose to incorporate both 2D and 3D information as model input, as generating 2D molecular graphs from 3D conformations is free and can bring in useful information from 2D perspective.

Detailed descriptions about how to incorporate each modality under different scenarios are presented in Section 3.3 of the paper. Our modality-blending pretraining technique enables the model to be versatile and adaptable to various downstream scenarios.

If there are any aspects that need additional clarification, please let us know. We would be delighted to engage in further discussions. Thank you.

**Q2: The matrix construction is similar to Transformer-M.**

We thank the reviewer for the comment. It is imperative for us to elucidate that our primary emphasis lies in the development of an innovative "pretraining algorithm", and follow the model backbone of Transformer-M [1] and Graphformer [2]. In Section 3.2, we explicitly cite that we follow Transformer-M and Graphformer to choose 2D Shortest Path distance, 2D Edge Type, and 3D Euclidean distance as atom-relation encodings.

It is crucial to underscore that our focus lies in improving the **incorporation** of the given 2D and 3D information, while the specific choice of 2D and 3D features is not our focus. Our method is generic and can be applied to any given 2D and 3D matrix, orthogonal to the specific choice.

To draw an analogy, our work can be likened to the relationship of BERT to Transformer. In the realm of pretraining, our concentration is the formulation of an effective training algorithm. The design of the model backbone and data encoding falls outside the scope.

[1] One transformer can understand both 2d & 3d molecular data. ICLR 2023
[2] Do transformers really perform badly for graph representation? NIPS 2021

Reviewer comment (Experiments):
Ablation studies on blending two masks should be provided.

Author response:
**Q3: Ablations on blending two masks.**

We thank the reviewer for the suggestion. The ablation results are presented in the table below, showing that blending 3 kinds of atom relations achieves best performance.

| Method | BBBP | Bace | Tox21 | ToxCast |
| --- | --- | --- | --- | --- |
| SPD mask | 68.95  | 80.64  | 75.59  | 62.82 |
| Edge mask  | 69.02  | 81.97  | 76.01  | 63.81 |
| 3D mask   | 67.60  | 80.35  | 75.65  | 63.28 |
| SPD_EDGE | 70.56 | 82.37 | 75.07 | 65.12 |
| EDGE_3D | 71.12 | 81.43 | 76.07 | 65.10 |
| SPD_3D | 71.23 | 79.1 | 75.90 | 65.04 |
| Blending   | **71.68** | **83.41** | **76.58** | **65.46** |

Reviewer comment (Experiments):
Some details of the experimental setup is missing.

Author response:
**Q4: details of the experimental setup.**

We have made revisions to our paper (Section A.1). Below we address the detailed questions respectively.

> Baseline results:

- GraphMVP, MoleSDE, GraphCL, GraphMAE are from their own paper.
- AttrMask, ContextPred, InfoGraph, MolCLR from MoleculeSDE
- MoleBERT, 3D Infomax are from MoleBERT
- The results of GROVER are from Uni-Mol

All compared methods use the same dataset split.

> Table 4, does the single-modality mask-then-predict strategies use the same network as the proposed blending strategy?

Yes, the compared experiments use the same network. All settings are the same, isolating training objective as the varying factor for ablation.

Reviewer comment (Experiments):
Does the author run all baseline methods on the experimented splits or cite some results from other papers? For the results in Table 4, does the single-modality mask-then-predict strategy use the same network as the proposed blending strategy?

Author response:
**Q4: details of the experimental setup.**

We have made revisions to our paper (Section A.1). Below we address the detailed questions respectively.

> Baseline results:

- GraphMVP, MoleSDE, GraphCL, GraphMAE are from their own paper.
- AttrMask, ContextPred, InfoGraph, MolCLR from MoleculeSDE
- MoleBERT, 3D Infomax are from MoleBERT
- The results of GROVER are from Uni-Mol

All compared methods use the same dataset split.

> Table 4, does the single-modality mask-then-predict strategies use the same network as the proposed blending strategy?

Yes, the compared experiments use the same network. All settings are the same, isolating training objective as the varying factor for ablation.

Reviewer comment (Experiments):
In Table 4, I think the author compares blending three masks to using only one mask. What's the effect of blending two masks?

Author response:
**Q3: Ablations on blending two masks.**

We thank the reviewer for the suggestion. The ablation results are presented in the table below, showing that blending 3 kinds of atom relations achieves best performance.

| Method | BBBP | Bace | Tox21 | ToxCast |
| --- | --- | --- | --- | --- |
| SPD mask | 68.95  | 80.64  | 75.59  | 62.82 |
| Edge mask  | 69.02  | 81.97  | 76.01  | 63.81 |
| 3D mask   | 67.60  | 80.35  | 75.65  | 63.28 |
| SPD_EDGE | 70.56 | 82.37 | 75.07 | 65.12 |
| EDGE_3D | 71.12 | 81.43 | 76.07 | 65.10 |
| SPD_3D | 71.23 | 79.1 | 75.90 | 65.04 |
| Blending   | **71.68** | **83.41** | **76.58** | **65.46** |

Reviewer comment (Experiments):
The conclusion in section D.1 is not well supported, since different 3D networks may perform differently and may have different speed of convergence. It's not enough to draw the conclusion by comparing the proposed method to only one 3D model.

Author response:
**Q5: The conclusion in section D.1 is not well supported.**

We thank the reviewer for the comment. We sincerely agree that more experiments are needed to support such conclusions. Due to this is not our primary focus and the convergence comparison with more 3D models needs further detailed investigation, we delete this experiment from our paper and leave it as future work.

Reviewer comment (Experiments):
Why run ablation studies on different datasets and different tasks?

Author response:
**Q6: Why run ablation studies on different datasets and different tasks?**

We thank the reviewer for the question. For our final model, we evaluate it on a broad range of tasks (8 MolNet classification tasks, 3 MolNet regression tasks, and 12 QM9 tasks).

For ablations, we assess partial tasks for two reasons: 1) the conclusion typically shares across different tasks, and can already be apparently established on conducted experiments; 2) conducting ablations on all datasets and tasks would be prohibitively costly.

For ablation of Table 3, we evaluate full 8 molnet classification tasks plus two QM9 tasks. As for the ablations in Table 4 and 5, where the conclusions are clear and easily established due to a substantial performance margin, we opt to evaluate only on a subset of tasks from MolNet and QM9. These selected tasks align with our internal development process and serve as a validation of our methodology. In these instances, we believe it is unnecessary to conduct evaluations on all tasks, given the evident and easily comprehensible conclusions arising from the observed substantial margins.