Paper: Motif-aware Attribute Masking for Molecular Graph Pre-training (ICLR-2024)

Reviewer comment (Experiments):
My first question is that is inter-motif influence considered in the proposed method? In my understanding, inter-motif influence is only considered as an evaluation metric. Is it considered in the pre-training process?

Author response:
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.

Reviewer comment (Theory):
If the inter-motif influence is an evaluation process, then the main contribution of the MoAMa is masking the whole motif instead of masking nodes. It comes up with an issue that completely masking the entire motif might be excessively difficult for the pre-trained GNN to reconstruct. Specifically, if all the nodes within a k-hop range of a motif are masked, then that particular node will never acquire any information about the features of its neighbors.

Author response:
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.

Reviewer comment (Evaluation):
The complexity of calculating inter-motif influence needs to be discussed. For each graph, influence between any two nodes needs to be calculated by applying two modified graphs into a graph neural network, which means it needs to run the graph neural networks at most O(N^2) times. I highly recommend authors discuss the complexity of their proposed evaluation method.

Author response:
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).

Reviewer comment (Reproducibility):
How have the authors presented the test AUC? Is it the test AUC corresponding to the best validation AUC, or is it the final test AUC? It's essential to provide clarity on these experimental configurations to ensure the results can be reproduced.

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

Reviewer comment (Presentation):
The layout of the paper, especially section 4.3, is unclear. I found it challenging to link section 4.3 with its preceding section. Additionally, section 4.4 appears to have a closer relation to section 4.2, given its extended discussion on the pre-training strategy.

Author response:
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.

Reviewer comment (Experiments):
The proposed approach was applied to binary classification problems. It will be interesting to see its performance on regression problem (e.g., the QM9 dataset).

Author response:
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

Reviewer comment (Evaluation):
Experiments show that the proposed approach does not always outperform existing methods. It will be nice if the authors can provide some insights.

Author response:
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.

Reviewer comment (Novelty):
The related work is missing all references to masking approaches in SSL beyond graphs. ICLR is a more general DL conference and graph SSL has clearly been inspired by those, and there are various works on masking, including theory.

Author response:
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.

Reviewer comment (Theory):
Some strong assumptions of the proposed approach are not enough backed by references, and the experiment design largely makes use of them (the proposed MRR metrics, ablation studies).

Author response:
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.

Reviewer comment (Writing):
l. 35: "The presence and interactions between chemical motifs directly influence molecular properties, such as reactivity and solubility (Frechet, 1994; Plaza et al., 2014). Therefore, to capture the interaction information between motifs, it is important to transfer inter-motif structural knowledge and other long-range dependencies during the pre-training of graph neural networks." - while this says that interactions are important, it does not say that neighbor information is not important or, specifically, that it is less important.

Author response:
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.

Reviewer comment (Theory):
"Unfortunately, the random attribute masking strategies used in previous work for graph pre-training were not able to capture the long-range dependencies inherent in inter-motif knowledge (Kipf & Welling, 2016; Hu et al., 2020b; Pan et al., 2019). That is because they rely on neighboring node feature information for reconstruction (Hu et al., 2020a; Hou et al., 2022)." - Usually, GIN is applied with 5 layers and the studied molecules are rather small.

Author response:
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.

Reviewer comment (Experiments):
While the inter-motif information is relevant, the experiments should focus on showing that it is complementary to or improves upon regular masking. However, this is not visible from the experiments reported, and the paper could make more visible which models are directly comparable in terms of masking (e.g., using the same baseline architecture and pre-training data) and which are not (e.g., comparison to Grover is lacking since the baseline model is already very different).

Author response:
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.

Reviewer comment (Evaluation):
This paper lacks visualization results, case studies, and other interpretability analyses.

Author response:
We apologize for any confusion. We have uploaded a revision of the manuscript that contains full details on the regression experiments (A.2) and the case study (A.3).

Reviewer comment (Experiments):
According to MGSSL[1], BRICS alone tends to generate motifs with large numbers of atoms. When the motif segmentation is too fine, many generated motifs are single atoms or bonds, which inhibits GNNs from learning higher-level semantic information through motif generation tasks. I wonder if the authors considered this issue when using only BRICS for motif fragmentation.

Author response:
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.

Reviewer comment (Experiments):
The experimental datasets used in this study are all classification datasets in molecular property prediction. Could the authors validate regression datasets such as ESOL and FreeSolv, to further assess the effectiveness and generalization ability of MoAMa?

Author response:
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

Reviewer comment (Reproducibility):
I have some confusion about the experimental results for the baselines. The experimental results for the baselines in this paper show a certain discrepancy from the reported AUROC in the original baseline papers. Could the authors please clarify the reasons behind this difference?

Author response:
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.  

(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

Reviewer comment (Theory):
The proposed method does not seem solid enough. Many details lack theoretical support. For example, the criteria to restrain the masked motif selection is intuitive without much explanation. The authors claim that the proposed method is able to utilize graph structure, while it is unclear how the motif-aware masking addresses this issue. Another claimed contribution is to capture both local and long-range information, while it is not discussed how the motif masking captures long-range other than random masking. The masked motifs can still be close, and the message passing is still within the k-hop.

Author response:
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.

Reviewer comment (Experiments):
The L_aux is proposed to complement L_rec since attribute masking focuses on local graph structures. Then this loss seems to fit with all the attribute masking methods. The results in Table 1 show that, MoAMa w/o Laux does not show much superiority compared with other baselines. It seems that the proposed motif-aware masking might perform limited. The authors should try Laux on other attribute masking methods to see if there are improvements. Also, it is not well elaborated on why and how Eq. 14 is designed.

Author response:
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).

Reviewer comment (Evaluation):
The results do not seem promising, and as the last point states, the effectiveness of the proposed motif-aware masking is questionable.

Author response:
3. See response to 2.1

Reviewer comment (Presentation):
The paper lacks an overall framework to help understand the detailed implementation.

Author response:
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.

Reviewer comment (Novelty):
The technical novelty of the paper is relatively incremental for an ICLR paper. The paper primarily focuses on a narrow aspect of the pretraining method, switching from random masking to motif-wise random masking, which may have limited impact on the community.

Author response:
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.