Multimodal Molecular Pretraining via Modality Blending

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.

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

We extend our sincere gratitude to the reviewers for your invaluable efforts and dedicated time invested in reviewing our paper.

Let us organize the discussion thread to establish a clear logical progression.

Our primary motivation is that aligning fine-grained atom-relations is crucial for aligning different modalities. This is well established in our introduction and the response to Reviewer TmmD and 3A6t. We thank Reviewers TmmD and 3A6t for the assistance in elucidating this point. Further, we have conducted theoretical analysis and experiments to validate the effectiveness of this fine-grained alignment, thereby supporting our motivation and claim.

[1] explores the effectiveness of atom-level alignment, which is another form of fine-grained alignment of modalities. Reviewer 3A6t recommends that the comparison to [1] can enhance the quality of our paper. However, a notable challenge arises due to their adoption of a distinct evaluation protocol from mainstream studies and the absence of a provided pretrained checkpoint, making the evaluation of their method challenging. This lack of a standardized basis hinders a fair and quick comparison.

It is imperative to underscore, however, that whether atom-level alignment is crucial extends beyond the current scope of our research and is less directly connected to our motivation. While we acknowledge the significance of this aspect, we believe that its thorough exploration can be left for future investigations.

[1] Unified 2D and 3D Pre-Training of Molecular Representations. KDD 2022

We thank Reviewer 3A6t for the further comment and address the further concerns below.

Q1: Why the key of modality alignment exists in the atom relationships.

The reviewer's major concern lies in that atom coordinates play a key role. We argue that when considering coordinates as attributes, what genuinely matters is the relative relation between the coordinates, i.e., the atom relation.

The coordinate themselves does not matter, as supported by the fact that when we rotate a molecule, its properties remain constant and invariant while all coordinates changed. The true determinant of a molecule's property is the invariant atom relations during rotation.

To this end, the employment of coordinates is typically coupled with equivariant constraints [1][2]. We posit that in essential, coordinate + equivariant constraints captialize on atom relations. Both the equivariant operation for processing coorinates (Figure 1 of [1], [2]) and the equivariant head used for predicting coordinates (Eq 6 of [3]) leverages atom relations to guarantee equivariance [1][2][3]. It is crucial to reiterate that coordinates must be coupled with equivariance for modeling molecules. A theoretical study [4] deduces that equivariance in essential can be expressed by relative scalars (e.g. relative distance), underscoring the significance of relative atom relations (or the equivalance of atom relations to equivariance and invariance).

In summary, equivariance and invariance are indispensable for molecule properties, and atom relations are the key for ensuring equivariance and invariance. Notably, atom coordinates themselves lack inherent invariance and equivariance.

We sincerely thank the reviewer for the question, and the relevant reference and rationale will be organized and included in our updated paper. We seek clarification on whether the reviewer agrees with this point.

[1] SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks. NIPS 2020

[2] E(n) Equivariant Graph Neural Networks. ICML 2021

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

[4] Scalars are universal: Equivariant machine learning, structured like classical physics. NIPS 2021

Q2: Baseline [1]

In our experiments, we did not compare with [1], as we find problems existing in their evaluation procedure, and their MoleculeNet finetuning practice in the official repository differs from mainstream methods [2][3]. For example, they finetune and evaluate ClinTox on only one sub-task (https://github.com/teslacool/UnifiedMolPretrain/blob/7f299f65433f9f1b7ceef5a003315640ae0cec50/pretrain3d/data/DCGraphPropPredDataset/deepchem_dataloader.py\#L16), whereas mainstream evaluation trains on two sub-tasks simultaneously with a single head, which is much more difficult. We also struggled to evaluate it under the commonly adopted evaluation protocol as it did not provide a pretrained checkpoint. To this end, we opt not to include its numbers as our baseline. However, we have discussed their method in related work in our initial submission.

The state-of-the-art work MoleculeSDE [2] from ICML 2023 also does not include its results as a baseline.

[1] Unified 2D and 3D Pre-Training of Molecular Representations. KDD 2022

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

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

Hello reviewer TmmD!

Thank you for pointing out that the paper does not satisfy the SE(3)-equivariance property. However, I believe the inclusion of this baseline (or even a training approach as an ablation study) can definitely show the importance of learning from the relationship between atoms is importance, and improve the quality of the work.

I'm happy to further listen to your opinion on this.

We sincerely appreciate the prompt response from Reviewer 3A6t. We have incorporated the results from [1] on partial MolNet tasks in Table 1. However, it is important to note that these results are 1) incomplete MolNet results and 2) have been evaluated under a distinct protocol. Therefore, the comparability of these results to other baselines and our work may be limited.

Concretely, one most obvious deviation is that the common protocal [2][3][4] of finetuning on MolNet is multitask training (as shown in Table 8 of GraphMVP [2] and Table 6 of Mole-BERT [4], for example, Tox21 and ClinTox are trained on 12 and 2 sub-tasks jointly). However, as per [1]'s official github repo, they focus only on one sub-task: ClinTox (https://github.com/teslacool/UnifiedMolPretrain/blob/7f299f65433f9f1b7ceef5a003315640ae0cec50/pretrain3d/data/DCGraphPropPredDataset/deepchem_dataloader.py#L16), Tox21 (https://github.com/teslacool/UnifiedMolPretrain/blob/7f299f65433f9f1b7ceef5a003315640ae0cec50/pretrain3d/data/DCGraphPropPredDataset/deepchem_dataloader.py#L34).

We follow the [2][3][4]'s protocal for evaluation, which is also the mainstream protocal adopted by the community. To this end, our results are not comparable to [1]. However, we have discussed the methodology of this work in the related work section.

We welcome the feedback on whether these adjustments adequately address the reviewer's concerns.

Thank you.

[1] Unified 2D and 3D Pre-Training of Molecular Representations. KDD 2022

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

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

[4] Mole-BERT: Rethinking Pre-Training Graph Neural Networks for Molecules. ICLR 2023

Dear reviewer TmmD

What I mean by relationship here is the relationship between atoms defined in this paper, i.e., distance, spd, edge type.

As my understanding, GraphMVP and 3D Infomax are training the model by learning the representation of molecule by molecule-level alignment, while this paper aims to learn the representation by aligning atom-level representation.

To the best of my knowledge, the KDD paper is the only paper with atom-level alignment, and this paper learns the representation by directly generating in data space without regarding relationship(i.e., distance, spd, edge type). Therefore, I believe the key difference between this paper and the KDD paper is regarding the relationship between the atoms, and the inclusion of this would improve the paper.

I'm happy to be corrected if there is something wrong. Feel free to share your thoughts!

Also, for authors, if there is something wrong in my knowledge, please don't hesitate to correct it.

We thank the reviewer for the helpful comments. We have uploaded a revision of our paper. Below we address the detailed questions.

Q1: Pretraining time cost

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.

Q2: Missed baselines.

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

Q3: More insights into why MOLBLEND overperforms existing methods theoretically.

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.

Q4: How does the different choice of encodings for 2D/3D modalities influence the pretaining?

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

We express our gratitude to the reviewer for dedicating time to review our paper. We have provided comprehensive clarifications to all the questions including 1) the pretraining time cost, 2) missed baselines, 3) theoretical insights of why MoleBLEND performs better, and 4) influences of different 2D/3D encodings. As the discussion deadline looms within a day, we would like to inquire if our responses have adequately addressed your questions. We anticipate your participation in the Reviewer-Author discussion phase, as your insights are invaluable for refining and enhancing the quality of our paper. Should you have any additional queries or require further clarification, please do not hesitate to inform us. We are more than willing to address any concerns and ensure a comprehensive resolution. Thank you for your time and consideration.

We thank the reviewer for the positive feedback. Below we address your comments in detail.

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.

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.

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.

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.

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.

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

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.

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.

Thank you. My questions have been addressed.

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.

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.

We thank the reviewer for the helpful comments. Below we address each question in detail.

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

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.

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.

We express our gratitude to the reviewer for dedicating time to review our paper. We have provided comprehensive clarifications to all the questions, e.g. 1) the information gathered in the relation matrix is limited, 2) the similarity to Transformer-M. As the discussion deadline looms within a day, we would like to inquire if our responses have adequately addressed your questions. We anticipate your participation in the Reviewer-Author discussion phase, as your insights are invaluable for refining and enhancing the quality of our paper. Should you have any additional queries or require further clarification, please do not hesitate to inform us. We are more than willing to address any concerns and ensure a comprehensive resolution. Thank you for your time and consideration.