Motif-aware Attribute Masking for Molecular Graph Pre-training

Thank you to the reviewer for providing us with their comments and questions for the manuscript. Below we will address the previously listed concerns.

1. We use inter-motif influence to evaluate the long-range message passing of our model. The limitation of previous work is the over-reliance of a node’s direct neighbors during training and the propagation bottleneck. We demonstrate that our model, with no modification to the underlying GNN architecture, is able to better mitigate issues related to feature over-squashing. The pre-training process itself is agnostic to the inter-motif influence measurements, as we can only measure the influence on a pre-trained network.

2. We address this concern using our motif selection criteria in lines 153-154. For criteria 1, all nodes within the masked motif must be within a k-hop distance from a node with visible features. Therefore, it is guaranteed that all masked nodes will receive information on some of its neighbors (excluding intra-motif neighbors). This pre-training task is indeed more difficult, but not excessively difficult. This added challenge encourages the model to learn longer-range dependencies and node interactions. Additionally, in practice, the average motif has ~2.4 nodes. With a 5-layer GIN, the masked nodes will still receive sufficient knowledge of its inter-motif neighbors.

3. GNN message passing is limited by the number of layers used, k. Therefore, the node influence calculations will only need to be performed on neighbors within a k-hop radius of each other. This means that the time complexity of our evaluation is O(n * d^k), where n is the number of nodes in the graph, k is the number of layers of our GNN (k = 5), and d is the average degree of a node. d^k <= n as molecular graphs are sparse, so the evaluation is not nearly as inefficient as O(n^2). Additional related analysis can be found in this work (1).

4. Following the experimental setup of this pioneering graph attribute masking work (2), we report the test ROC-AUC at the best validation epoch.

5. Thank you for this comment! We have opted to move our discussion of inter-motif influence (Sec. 4.3) in order to preserve the logical flow Sec. 4.2 and 4.4.

(1) Nikolentzos, Giannis, George Dasoulas, and Michalis Vazirgiannis. "k-hop graph neural networks." Neural Networks 130 (2020): 195-205.

(2) Hu, Weihua, et al. "Strategies for pre-training graph neural networks." arXiv preprint arXiv:1905.12265 (2019).

Thank you for the suggestion! We have performed additional experiments on the Lipophilicity, ESOL, and FreeSolv datasets. Our method shows comparable average performance to previous methods and performs the best out of the previous methods on the ESOL dataset. We reported the test RMSE for the best validation RMSE across 3 independent runs. For brevity, we have included the average RMSE across all three datasets below.

None: 1.716

AttrMask: 1.591

MGSSL: 1.500

GraphLoG: 1.607

JOAO: 1.786

GraphMAE: 1.530

D-SLA: 1.515

Mole-BERT: 1.494

MoAMa: 1.509

For molecular property prediction, it has become standard practice to separate the training, validation, and test sets according to molecular scaffolds, i.e. fragments. Therefore, the data distribution on the training set will not reflect that of the validation or test sets. This procedure makes the focus of predictive models to be on generalization. However, when training, it is hard to know what knowledge will be most useful for generalization, especially when considering multiple downstream datasets. Hence, the trained knowledge from pre-training and the fine-tuning training set may not always be useful for evaluation on the validation and test sets. Our method shows great generalization on 5 of 8 downstream datasets, but the knowledge transferred for the remaining 3 may not be as useful as the knowledge transferred from other models.

Thank you to the reviewer for their positive comments for our work! We will address the weaknesses and questions below:

1. This was not an issue with our work. For MGSSL, the training process requires a full list of available motifs to generate, so it is vital to create an appropriate motif generation strategy that reduces the number of motif classes to generate, but does not reduce motifs down to single atoms. For our work, we are using motif segmentation to choose which nodes to mask for attribute masking. Therefore, the exact motif label itself is not necessary for training at all, only the segmentation. We considered the issue of finding appropriately sized motifs, but through statistical analysis, we found that motifs on average had 2.4 nodes each. Random attribute masking models the situation where each motif is of size 1, so already our method already takes into consideration higher-level semantic information during training as compared to the baseline AttrMask. Furthermore, we use the motif selection criteria to disallow the masking of nodes in motifs too great in size, written in lines 153-154. For these reasons we were able to use BRICS without encountering the same issues as the authors of MGSSL would have needed to deal with.

2. We have performed additional experiments on the Lipophilicity, ESOL, and FreeSolv datasets. Our method shows comparable average performance to previous methods and performs the best out of the previous methods on the ESOL dataset. We reported the test RMSE for the best validation RMSE across 3 independent runs. For brevity, we have included the average RMSE across all three datasets below.

None: 1.716

AttrMask: 1.591

MGSSL: 1.500

GraphLoG: 1.607

JOAO: 1.786

GraphMAE: 1.530

D-SLA: 1.515

Mole-BERT: 1.494

MoAMa: 1.509

3. These results are compiled from these three recent works (1, 2, 3). If the same baseline was found in two or more works, the better result was reported. The experimental settings for these works were the same as ours, and to demonstrate the integrity of our work, we chose to use their reported numbers rather than our own. We did notice similar results in our own experiments when reproducing the results from the baselines, however.

Lastly, to address the weaknesses brought up by the reviewer, we have conducted case studies on the embedding spaces of pairs of molecules with similar structures and but different properties. The results illuminate further why random masking cannot capture long-range signals effectively and map molecules with similar local structures to similar places in the embedding space. We will include these details in the appendix to our work.

(1) Kim, Dongki, Jinheon Baek, and Sung Ju Hwang. "Graph self-supervised learning with accurate discrepancy learning." Advances in Neural Information Processing Systems 35 (2022): 14085-14098. (2) Xu, Minghao, et al. "Self-supervised graph-level representation learning with local and global structure." International Conference on Machine Learning. PMLR, 2021. (3) Jun Xia, Chengshuai Zhao, Bozhen Hu, Zhangyang Gao, Cheng Tan, Yue Liu, Siyuan Li, Stan Z. Li: Mole-BERT: Rethinking Pre-training Graph Neural Networks for Molecules. ICLR 2023

Thank you to the reviewer for their comments on how we can improve our manuscript! Below we provide our response to the comments raised above.

1. We have briefly mentioned some related works within NLP and CV that utilize semantically related regions for masking (lines 29-32, 53-57), however we agree that a more detailed discussion would be important to contextualize this work to the larger research community. We will include a more detailed discussion of related masking works, such as Masked Language Modeling and Masked Autoencoding in CV in our final version.

2.1 Graph transformers, which have been growing in popularity, have been shown to outperform traditional GNNs on different graph prediction tasks, including those in the material science domain (1, 2, 3). The strength of these transformer models is the leveraging of full-connections among nodes in the graph, allowing the capture of long-range dependencies. Therefore, long-range dependencies are shown to increase performance of prediction models. However, many message passing GNNs struggle to pass long-range information, even when considering the k-hop neighborhood restriction. Many MPNN perform better when only considering a node’s 1-hop neighborhood (4) due to the propagation bottleneck and subsequent over-smoothing of features, even though long distance interactions have been shown to be important during training in chemical domains (5). Through our experiments, we have empirically measured the reliance of inter-motif neighbors on previous GNN based attribute masking models, and we have found that our method is better able to mitigate the propagation bottleneck on a 5-layer network while demonstrating notable downstream predictive results.

2.2 Both inter-motif and intra-motif interaction information are useful for downstream prediction. However, prior work in message passing for quantum chemistry has shown that long-range dependencies are important for downstream prediction in chemical domains (5). Additionally, traditional GNNs have the same expressiveness as the 1-WL test (6). In fact, generationalization to higher order graph structures is an on-going problem. One recent work proposes a message passing strategy that aggregates subgraph features during each message passing iteration (7), moving the expressivity boundary to the 3-WL test. Our work does not make any claims to extend GNNs expressivity beyond the 1-WL test, but we do leverage node features in a strategic way to overcome the propagation bottleneck. With smarter aggregation of features during each 1-hop step.

2.3 Even with only 5 layers, the bottleneck problem can still be observed (8). So even with small molecules, it is possible for over-squashing to remove long-range interactions during propagation. As shown through our experiments in Table 3, our method transferred the most inter-motif knowledge, thus mitigating the over-squashing problem more than previous works.

3.1 Thank you for this comment. We will modify our sentence to be more precise. We took care in selecting these SSL pre-training models because they were highly related to our work and represented SOTA in a diverse range of strategies.

3.2 In Table 1, we demonstrate that MoAMa outperforms previous masking works (Grover, AttrMask, ContextPred, GraphMAE, Mole-BERT). We do not believe that our method is complementary to regular masking as it replaces regular random masking with a motif-aware masking strategy. We will make a clearer distinction between which models are directly comparable (e.g. have the same pre-training data and architecture). All models except MCM and Grover utilize a 5-layer GIN with the ZINC15 dataset for pre-training.

3.3 Our model, even without L_aux, still holds a 1.0% average performance increase as compared to previous methods, which we believe is significant and represents SOTA.

(1) Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Létourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. (2) Grégoire Mialon, Dexiong Chen, Margot Selosse, and Julien Mairal. GraphiT: Encoding graph structure in transformers. Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do transformers really perform badly for graph representation? (4) Dwivedi, Vijay Prakash, et al. "Long range graph benchmark." (5) Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. (6) Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. (7) Feng, Jiarui, et al. "How powerful are k-hop message passing graph neural networks." (8) Alon, Uri, and Eran Yahav. "On the bottleneck of graph neural networks and its practical implications."

Thank you for the detailed responses, I'll consider updating my score after discussing with the other reviewers.

2 If the work focuses on long-range dependencies, then the experiments should maybe include experiments specifically comparing to "deeper" GNNs. My original comment was primarily focusing on the reasoning in terms of chemistry.

3.1 I think it would be indeed helpful to clarify to which of the common/existing masking strategies the paper compares to directly. The missing comparison to the Grover strategy is unlucky, but acceptable. In which terms is ContextPred a masking strategy?

Thank you to the reviewer for providing their insightful comments! We will provide our response to the concerns below.

1. The criteria on the masked motif selection are created directly due to the limitations of GNNs more broadly. Because of the over-smoothing issue of GNNs with many layers, GNNs are restricted in depth. Therefore, to guarantee that a masked node will receive feature signals from its inter-motif neighbors, we must restrict the size of the motifs being masked. We use criteria 1 to guarantee all nodes within a motif are passed relevant inter-motif knowledge. Criteria 2 is used to eliminate the cases where two large masked motifs are adjacent, and thus nodes near the boundary may not receive any propagation signals. Traditional GNNs have the same expressiveness as the 1-WL test and are thus unable to recognize larger order structures within graphs (1). Due to the propagation bottleneck inherent to GNNs, traditional GNN models of size k >= 4 are also unable to fully utilize local graph structures effectively (2). While staying within the constraints of the GNN architecture, our work seeks to mitigate the propagation bottleneck by encouraging the transfer of longer-range signals from a node’s inter-motif neighbors. With the retention of these longer range signals, our model is able to propagate featural and structural information further than other pretrained graph models. Because all nodes within a selected motif are masked, this inherently encourages the learning of local intra-motif structure. Additionally, our motif-aware masking strategy encourages the propagation of inter-motif signals. We develop inter-motif influence measurements in Sec 4.3 to demonstrate the weakness of previous methods and empirically show in Table 3 the increase of inter-motif knowledge transfer using our model as compared to baselines using the same GNN architecture.

2.1 Even without L_aux, our method is still competitive with a 1.0% average AUC-ROC performance increase across the baselines, which we believe still represents SOTA performance.

2.2 Thank you for suggesting this interesting direction for future work! Because the focus of this paper is on using a motif-aware masking strategy to encourage long-range signal transfer and mitigate over-squashing, we have chosen to leave this investigation of global motif knowledge for future work. With that being said, we have performed preliminary experiments using L_aux to complement the loss function of AttrMask and found a 5.6% performance increase from 77.8 to 83.4 on the BACE dataset, a 3.9% increase from 65.2 to 69.1 on the BBBP dataset, a 5.2% increase from 73.5 to 78.7 on the ClinTox dataset, and a 3.4% increase from 60.5 to 63.9 on the SIDER dataset. While our current strategy still outperforms AttrMask with L_aux, this increase in predictive performance is significant. Our L_aux certainly seems to complement this existing masking strategy well, and it would be a fruitful direction for further research.

2.3 Eq. 15 is inspired by eq. 5 from a similarity-based contrastive approach for 3D molecular representation learning (3).

3. See response to 2.1

4. In Figure 1, we represent the motif-aware masking strategy, in which the molecule is decomposed into disjoint motifs. Several motifs are selected according to masking criteria detailed in lines 153-154. For these masked nodes within the selected motifs, we used a graph autoencoder to reconstruct the graph node features, namely atom type, as our pre-training task. We have elucidated the design space of our method in Sec. 4.4 and Sec. 5.4 and provided the relevant experimental results, which should be comprehensive and insightful for understanding the graph autoencoder design.

5. We believe that this work not only contributes to the exploration of graph attribute reconstruction, but also provides a strong basis for future research in graph pre-training strategies. There are only a handful of fragment/motif based graph pretraining works, but the results from our work show that even a simple motif-based pretraining strategy is competitive with existing works. The SOTA results demonstrate that our simple yet effective method is able to utilize domain knowledge from motifs to enhance predictive performance over models which learn graph knowledge randomly. Because our work in motif-based graph pretraining mitigates feature over-squashing without the need to modify standard GNN architecture, we believe more sophisticated models can gain insights from our work to further improve GNN capabilities.

(1) Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. (2) Alon, Uri, and Eran Yahav. "On the bottleneck of graph neural networks and its practical implications." (3) Atsango, Austin, et al. "A 3D-Shape Similarity-based Contrastive Approach to Molecular Representation Learning."