Where to Mask: Structure-Guided Masking for Graph Masked Autoencoders

1 Introduction

In domains, such as academic, social, and biological networks, graph-structured data often lacks labels. This problem is especially acute in biochemical graphs due to the expense of wet-lab experiments. To address this challenge, many techniques have been developed to fully exploit the existing massive amounts of unlabeled data, aiming to enhance graph model training. Among these approaches, self-supervised graph pre-training (SSGP) methods are prominent due to their effectiveness, attracting significant interest in academic and industrial realms [Xia et al., 2022b].

∗ Corresponding Author

Currently, SSGP methods are categorized into two primary streams: 1) Contrastive methods, such as GraphCL [You et al., 2020] and SimGRACE [Xia et al., 2022a], which utilize contrastive learning principles to reveal the intrinsic structure and interconnections within graph data. 2) Generative methods, including GraphMAE [Hou et al., 2022] and MaskGAE [Li et al., 2023]. These methods focus on learning node representations through a reconstruction objective. Moreover, generative methods have proved to be simpler and more effective than contrastive approaches that require carefully designed augmentation and sampling strategies. The efficacy of generative methods is further underscored by the enormous successes of models such as BERT and ChatGPT in Natural Language Processing (NLP) [Devlin et al., 2019] and MAE in Computer Vision (CV) [He et al., 2022]. These successes highlight the significant potential of generative approaches across various domains. Accordingly, this paper explores the capabilities of generative methods, specifically graph masked autoencoders (GMAE), in graph learning tasks and recognizes their potential as evidenced in other fields.

GMAE fundamentally involves randomly masking a proportion of input data (i.e., nodes or edges) and leveraging the reconstruction of the removed contents to guide the representation learning. Despite GMAE's promising results, its random masking approach, which assigns equal probability to all nodes in a graph - a universally adopted strategy, presents a suboptimal strategy. Specifically, the masked nodes are sometimes overly simplistic to predict (i.e., the atom C in Figure 1 (a)) with only neighborhood information. In such cases, the model's pre-training phase may not be sufficiently informative, thereby hindering the acquisition of valuable knowledge. However, if we mask a large number of key informative nodes (i.e., SO3 in Figure 1 (b)), the model may fail to perceive the graph's overall structural information. In summary, the indiscriminate nature of random masking, which fails to distinguish nodes of varying informational values, potentially leads to low data efficiency and compromises the quality of the learned graph representations. Therefore, this raises the question: is there a more effective masking strategy than random sampling for enhancing GMAE's pre-training process?

This paper answers the aforementioned open question by introducing StructMAE, which features a novel structure-guided masking strategy designed to enhance GMAE's pre-training process. The key insight of our method is to inject prior graph structure knowledge into the masking process to guide model learning. Specifically, StructMAE comprises two principal components: Structure-based Scoring and StructureguidedMasking. We recognize that the node reconstruction complexity is inherently linked to its structural significance within the graph. Therefore, we derive a scoring method to assess the node significance, distinguishing between informative and less-informative nodes based on structural considerations. In addition, two scoring method variants are proposed: predefined and learnable, which are discussed in detail in Section 4. With node significance scores established, we propose an easy-to-hard masking strategy that gradually increases the difficulty of the self-supervised reconstruction task. This approach commences with the masking of less-informative nodes, progressively shifting towards masking more informative nodes as the model's learning progresses. This strategic progression in masking difficulty is designed to enable the model to gradually and effectively assimilate the graph's structural information.

To evaluate the effectiveness of the StructMAE model, we conduct comprehensive experiments on a range of widelyused datasets, notably the large-scale Open Graph Benchmark (OGB) [Hu et al., 2020a], covering two graph learning tasks: unsupervised and transfer learning. The experimental results consistently demonstrate that StructMAE's performance surpasses that of existing state-of-the-art models in contrastive and generative pre-training domains. This exceptional performance demonstrates our structure-guided masking approach's advantages over conventional randommasking methods. The principal contributions are summarized as follows:

1. We introduce StructMAE, a novel node masking strategy tailored for GMAE. This strategy utilizes the structural information inherent in graphs to gradually and effectively direct the masking process, significantly enhancing the model's capability for representation learning.
2. We evaluate StructMAE through extensive experiments, comparing its performance with generative and contrastive baselines across two graph tasks on various real-world graph datasets, including the OGB dataset. The experimental results consistently validate StructMAE's effectiveness.

2 RelatedWork

Self-supervised Graph Pre-training. Inspired by the success of pre-trained language models like BERT [Devlin et al., 2019], T5 [Raffel et al., 2020], and ChatGPT [Brown et al., 2020], numerous efforts have been directed towards SSGP. Based on model architectures and objective designs, SSGP is naturally divided into contrastive and generative methods. First, contrastive self-supervised learning has dominated graph representation learning in the past two years. Its success is largely due to elaborate data augmentation designs, negative sampling, and contrastive loss. For instance, DGI [Veliˇckovi´c et al., 2019] and InfoGraph [Sun et al., 2020], based on mutual information, leverage corruptions to construct negative pairs. Similarly, models such as SimGRACE [Xia et al., 2022a] and GraphCL [You et al., 2020] utilize in-batch negatives. Differing from previous methods that execute augmentation on graphs, COSTA [Zhang et al., 2022b] implements augmentation within the embedding space, aiding in the mitigation of sampling bias. Moreover, POT [Yu et al., 2024] advances graph contrastive learning (GCL) training through the employment of a node compactness metric, assessing adherence to the GCL principle. Second, generative self-supervised learning focuses on recovering missing parts of the input data. For example, GAE [Kipf and Welling, 2016] is a conventional method that reconstructs the adjacency matrix. Multiple GAE variants utilize graph reconstruction to pre-train Graph Neural Networks (GNNs), including ARVGA [Pan et al., 2018] and SIGVAE [Hasanzadeh et al., 2019]. Recently, a paradigm shift towards GMAE has shown promising results in various tasks [Liu et al., 2023b]. A detailed introduction will be provided in the following section.

GraphMaskedAutoencoders. GMAEmainly involves reconstructing the contents (e.g., nodes and edges) that are randomly masked from the input using autoencoder architecture. A notable example of GMAE is GraphMAE [Hou et al., 2022], which reconstructs randomly masked input node features with several innovative designs, including re-mask decoding and scaled cosine error. In addition, MaskGAE [Li et al., 2023], S2GAE [Tan et al., 2023], and GiGaMAE [Shi et al., 2023] jointly reconstruct the masked edges and node degrees. SimSGT [Liu et al., 2023c], and GCMAE [Wang et al., 2024] combine contrastive learning with GMAE, whereas RARE [Tu et al., 2024] employs self-distillation to enhance GMAE's performance. Unlike the above mentioned GMAE methods, which primarily use message-passing GNNs as backbone models, GMAE-GT [Zhang et al., 2022a] utilizes a graph transformer [Ying et al., 2021;

Liu et al., 2023a] as its encoder backbone. There are also several GMAE applications, including heterogeneous graph representation learning [Tian et al., 2023], protein surface prediction [Yuan et al., 2023], and action recognition [Yan et al., 2023]. However, all the aforementioned methods employ a random masking method when masking graph contents and thus overlook the importance of GMAE mask strategies, potentially inhibiting the model's capabilities. The work most closely related to this study is MoAma [Inae et al., 2023], which similarly focuses on designing superior mask strategies. However, this approach uses motifs which rely on domain knowledge and manual motif pre-definition, limiting its generalizability across various domains.

3 Preliminaries

Notations. A graph G = (V, E) can be represented by an adjacency matrix A ∈ {0, 1} 𝑛×𝑛 and a node feature matrix X ∈ R 𝑛×𝑑 , where V denotes the node sets, E denotes the edge sets, 𝑛 is the number of nodes, 𝑑 is the dimension of the node features, and A[𝑖, 𝑗] = 1 if there exists an edge between nodes 𝑣𝑖 and 𝑣 𝑗 , otherwise, A[𝑖, 𝑗] = 0.

Graph Masked Autoencoders. GMAE operates as a selfsupervised pre-training framework, which focuses on recovering masked node features or edges based on the representations of unmasked nodes. To illustrate this method, we focus on the reconstruction of node features as an example. GMAE comprises two essential components: an encoder ( 𝑓𝐸(·)) and a decoder ( 𝑓𝐷(·)). The encoder maps each unmasked node 𝑣 ∈ Vunmask to a 𝑑-dimensional vector h𝑣 ∈ R 𝑑 , with Vunmask representing the set of unmasked nodes, while the decoder reconstructs the masked node features from these vectors. The entire process can be formally represented as:

where Xunmask and Hunmask denote the features and embeddings of unmasked nodes, respectively, and X ′ represents the reconstructed features of all the nodes. Then, GMAE optimizes the model by minimizing the discrepancy between the reconstructed representation of masked nodes, X ′ mask ⊂ X ′ , and their original features, Xmask ⊂ X.

4 Methodology

This section presents the model architecture of StructMAE in detail. First, we provide a detailed introduction to the masking module in GMAE (§4.1). Subsequently, a comprehensive exploration (§4.2) of the masking strategy is provided. Then, the details of the proposed StrucMAE are elucidated (§4.3). Finally, the overall StructMAE architecture, encompassing training and inference details, is expounded (§4.4).

4.1 Introducing the GMAEMaskingModule

Prior to processing graph data within the GMAE encoder, a subset of nodes Vmask ⊂ V is sampled for masking. According to the methodology described in [Hou et al., 2022], the features of these sampled nodes are replaced with a learnable vector 𝒙[M] ∈ R 𝑑 . Accordingly, for a node 𝑣𝑖 within the masked node subset Vmask, its feature�x𝑖 in the altered feature matrix �X is defined as follows:

Regarding the approach for selecting nodes to mask, most existing works [Hou et al., 2022; Li et al., 2023] utilize a random masking strategy. This method assigns an equal masking probability of masking to each node within a graph. A more detailed exploration of the random masking strategy will be presented in the following section (§4.2).

4.2 Reconsidering the RandomMasking Strategy

In GMAE context, the masking strategy is crucial, as it significantly influences the type of information that the model learns. In previous studies [Hou et al., 2022], nodes within a graph are masked randomly, each with an equal probability. However, this strategy overlooks the varying structural information of different nodes, which has been shown to be crucial in graph learning tasks. Therefore, we conduct a preliminary experiment to explore whether the incorporation of structure priors could augment pre-training efficacy.

In the experiment, models are pre-trained on the MNIST dataset and evaluated in unsupervised settings. Nodes forming numerical values are identified as those possessing rich structural information. Subsequently, themasking probability for these structurally informative nodes is manually increased. The results, presented on the left side of Figure 3, reveal the impact of this modified masking strategy on unsupervised accuracy. The results indicate that: 1) A marginal increase in the masking probability for nodes with rich structural information enhances the model's pre-training learning. Furthermore, up to a probability threshold of 0.2, a corresponding gradual increase in the model's accuracy is observed. 2) Conversely, excessively raising the masking probability of these structurally significant nodes detrimentally affects the model's training. These findings corroborate our initial discussion and highlight the importance of proposing an effective method to integrate structural information into the masking process.

4.3 The ProposedMasking Strategy

Inspired by the preceding discussion, we introduce an innovative structure-guided masking approach for GMAE, named StructMAE (as depicted in Figure 2 (b)). StructMAE involves integrating the graph's structural knowledge into the masking process, thereby directing the model's learning trajectory more effectively. StructMAE is composed of two primary elements: Structure-based Scoring (SBS) and Structureguided Masking (SGM).

Structure-based Scoring

The SBS evaluates the significance of each node based on its structural role within the graph. This evaluation facilitates the identification of nodes that are pivotal for the model to learn, allowing for a more targeted masking approach. To determine the importance of nodes, we introduce two distinct methodologies: the pre-defined and learnable methods.

Pre-defined Structure-based Scoring. The predefined method involves using a set of predetermined criteria, based on the known structural information, to evaluate node importance. Specifically, the computation of importance score S ∈ R 𝑛 is achieved using the PageRank algorithm [Page et al., 1999], a well-established technique for evaluating node significance based on graph structures. Thus, the importance score of node 𝑣𝑖 is defined as:

where 𝑒 denotes the damping factor, 𝐿𝑖 represents the degree of node 𝑣𝑖, and N𝑖 denotes the set of neighboring nodes of node 𝑣𝑖. In addition to PageRank, other prevalent metrics for assessing node importance comprise degree, closeness centrality, and betweenness centrality. Each of these methods is based on different underlying principles, offering diverse perspectives on a node's role and influence within a graph. Although these methods constitute more intricate ways of evaluating node importance, our empirical findings suggest that PageRank serves as a straightforward and effective measure. A detailed discussion and comparison of these methods is presented in the following section.

LearnableStructure-basedScoring. In contrast to the predefined method, the learnable approach dynamically assesses node significance based on the evolving state of the graph during the learning process. Specifically, it integrates the formulation of the assessment metric with masked graph modeling, enabling end-to-end learning of this metric. To accomplish this, we employ a lightweight scoring network, denoted as 𝑓𝑆 (·), which assesses the importance of each node 𝑣𝑖. The scoring network's design is similar to a GNN-style layer, and it effectively captures graph structural information. Thus, the importance score 𝑠𝑖 ∈ S for node 𝑣𝑖 is calculated as follows:

A higher score 𝑠𝑖 indicates greater importance of the corresponding node 𝑣𝑖. Following the scoring, node features are sorted in descending order based on these scores. The ordered node features and their respective scores are represented as � x ′ 𝑖 � and � 𝑠 ′ 𝑖 � , respectively, where 𝑖 = 1, . . . , 𝑛. To facilitate learning the scoring network 𝑓𝑆 (·), the predicted scores are multiplied by the node features to serve as modulating factors. This operation is formally expressed as:

where ˆX denotes the set of node features after scoring. This mechanism ensures that the scoring network is continually updated and refined throughout the model's training process.

In this section, we explore the two distinct SBS methods proposed for StructMAE. First, the pre-defined SBS method offers simplicity and is particularly effective when the structural characteristics are well-understood and can be explicitly defined beforehand. Conversely, the learnable SBS is more adaptable and can cater to complex and variable graph structures, making it suitable for scenarios where the node significance cannot be easily predetermined.

Structure-guided Masking

The SGM component utilizes the scores generated by the SBS to guide its masking decisions. It selectively and progressively masks nodes in an easy-to-hard manner, thereby enhancing the model's capacity to effectively learn and represent graph structures. Specifically, in the initial training stage, a subset of easy nodes with lower scores is masked, making it easier for the model to predict them using basic neighboring information. As training progresses, the masking strategy evolves to encompass more challenging nodes. This enables the model to capture intricate structural information and, consequently, enhances its learning capabilities [Liu et al., 2017; Liu et al., 2018].

This masking strategy relies on the importance scoring matrix S. It is hypothesized that nodes with higher S scores are more informative and significant. Consequently, we gradually increase the masking probabilities for these high-scoring nodes during masked graph modeling. To implement this, we rank the nodes based on their scores S = {𝑠1, 𝑠2, · · · , 𝑠𝑛} and identify the top 𝐾 indices that constitute the set Y. The masking probability for each node is then determined as:

where 𝜖 denotes random noise drawn from a uniform distribution (𝜖 ∼ 𝑈(0, 1)) and 𝛽 represents the increased probability assigned tonodeswithhigher scores. In this instance, nodes in the set Y are considered more informative, and consequently, the model is anticipated to prioritize these nodes. The number of masked nodes, denoted by 𝐾, is dynamically adjusted throughout the training process. Initially set at zero, 𝐾 progressively increases with epoch according to the following formula:

where 𝐾(𝑡) represents the 𝐾 value at epoch 𝑡, 𝑛 denotes the total number of nodes in the graph, 𝑝 is the predefined mask ratio, and 𝑇 represents the total number of training epochs. This approach enables the model to progressively concentrate on more challenging nodes, thereby enhancing its acquisition of complex structural information.

4.4 Overall StructMAE Architecture

Training Process. The StructMAE training process begins with an input graph from which a specified proportion of nodes is chosen based on our selective masking strategy. The selected nodes are subsequently masked using a masktoken. The graph, now with partially masked features, is subsequently fed into the encoder, which generates encoded representations of the nodes. After that, the decoder module is responsible for predicting and reconstructing the features of masked nodes. For optimization, we adopt the scaled cosine error as utilized in GraphMAE [Hou et al., 2022].

Inference andDownstreamTasks. StructMAE is designed to cater to two distinct downstream applications: unsupervised and transfer learning. In unsupervised learning, the encoder processes the input graph without masking during the inference stage. The node embeddings generated by the encoder are then utilized for graph classification tasks with linear classifiers or support vector machines. In the transfer learning context, the pre-trained models are fine-tuned on different datasets. This fine-tuning enables the model to adjust to new data domains, leveraging the foundational knowledge gained during its initial training on the source dataset. Each of these downstream tasks emphasizes the versatility and applicability of StructMAE in diverse graph learning scenarios.

5 Experiment

5.1 Unsupervised Representation Learning

Objective. To assess the efficacy of the pre-trained model in its feature extraction capability, we subject it to a series of unsupervised tasks. Achieving success in these tasks will highlight the model's proficiency in learning high-quality and informative representations, which are crucial for various downstream graph analytics tasks.

Settings. We employ seven real-world datasets, including MUTAG, IMDB-B, IMDB-M, PROTEINS, COLLAB, REDDIT-B, and NCI1, involving diverse domains and sizes. To demonstrate the effectiveness of our proposed method, we compare StructMAE with the following 10 baseline models: a) Two supervised models: GIN [2019] and DiffPool [2018]; b) Six contrastivemodels: Infograph [2020], GraphCL[2020], JOAO [2021], GCC [2020], InfoGCL [2021a], and SimGRACE [Xia et al., 2022a]; c) Two generative models: GraphMAE [Hou et al., 2022] and S2GAE [Tan et al., 2023]. We report the results from previous papers according to graph classification research norms. In the evaluation protocol, we initially generate graph embeddings using the encoder and readout function. Then, the encoded graph-level representations are fed into a downstream LIBSVM [Chang and Lin, 2011] classifier for label prediction, consistent with other baseline models. The performance is assessed by measuring the mean accuracy obtained from a 10-fold cross-validation, and this evaluation is repeated five times to ensure robustness.

Results. The results are detailed in Table 1, where StructMAE-P denotes StructMAE with predefined SBS, and StructMAE-L is with learnable SBS. Analyzing these results enables deriving several observations. 1) State-of-theart Performance: StructMAE-L outperforms existing selfsupervised baselines on four out of seven datasets. Furthermore, it attains state-of-the-art performance considering the average rank across these datasets. Meanwhile, StructMAEP maintains competitive performance against other selfsupervised methods. These results emphatically demonstrates the efficacy of the proposed StructMAE approach. Please note that, StructMAE focuses solely on masking nodes, whereas S2GAE extends its masking strategy to include edges. The differing approach of S2GAE underscores a potential avenue for further development in StructMAE. 2) Comparison with SupervisedMethods: Remarkably, the StructMAE-L, though self-supervised, attain comparable or superior performance on certain datasets, including PROTEINS, IMDB-B, and COLLAB. This finding indicates that the representations learned by StructMAE are high-quality and informative, aligning with supervised learning benchmarks. 3) Comparison with GraphMAE: In comparison to GraphMAE, which employs a random node masking strategy, StructMAE consistently outperforms it, providing further evidence to support our hypothesis that incorporating structural knowledge into the masking process can significantly enhance the model's learning capabilities. 4) Predefined vs. Learnable SBS: A notable trend is that the StructMAE-L's performance generally surpasses that of StructMAE-P, particularly on complex datasets such as COLLAB and REDDIT-B. This trend indicates the greater adaptability and effectiveness of the learnable SBS in handling complex and variable graph structures, as opposed to the predefined method. These results collectively validate the core principles behind StructMAE and its components, highlighting its potential as a powerful tool for unsupervised representation learning in graph-structured data.

5.2 Transfer Learning

Objective. The primary goal of the transfer learning task is to evaluate the transferability of the pre-training scheme utilized in StructMAE. This involves pre-training the model on a specific dataset and fine-tuning it with different datasets.

Settings. During the initial pre-training phase, StructMAE is trained on a dataset comprising two million unlabeled molecules obtained from the ZINC15 [Sterling and Irwin, 2015] dataset. Subsequently, the model is fine-tuned on eight classification benchmark datasets featured in the MoleculeNet dataset [Wu et al., 2018]. In our evaluation, we adopt a scaffold-split approach for splitting the datasets, as outlined in [Hou et al., 2022]. To demonstrate the effectiveness of our proposed method, we compare StructMAE with the following 10 baseline models: a) Three unsupervised models: Infomax, AttrMasking and ContextPred [Hu* et al., 2020b]; b) Four contrastive models: GraphCL [2020], JOAO [2021], GraphLOG [Xu et al., 2021b], and RGCL [Li et al., 2022]; c) Three generative models: GraphMAE [Hou et al., 2022], GraphMAE2 [Hou et al., 2023], and Mole-BERT [Xia et al., 2023]. We report the results of baseline models from previous papers according to research norms. Experiments are conducted 10 times, and the mean and standard deviation of the ROC-AUC scores are reported. According to the default settings used in prior research [Hou et al., 2022], a 5-layer GIN model [Xu et al., 2019] is employed as the encoder and a single-layer GIN as the decoder in StructMAE.

Results. The detailed results in Table 2 offer insightful observations into StructMAE's performance within the transfer learning context. 1) State-of-the-Art Performance Across Datasets: StructMAE-P and StructMAE-L demonstrate stateof-the-art performance across an ensemble of eight datasets. Specifically, StructMAE-P and StructMAE-L achieved 1.4% and 1.5% improvements in average performance metrics, respectively. Additionally, each method individually achieves top-tier performance on several datasets. These results validate the StructMAE's ability to effectively generalize learned representations across diverse datasets. 2) Superiority of StructMAE-L: The results reveal the superior performance of StructMAE-L over StructMAE-P on the average metrics. This trend further highlights the enhanced adaptability and effectiveness of the learnable SBS method, similar to observations made previously in unsupervised learning.

5.3 Ablation Study

A detailed study is conducted to evaluate the impact of different components within StructMAE. It is important to note that, except for the components under analysis, all other aspects of the model remain consistent with the comprehensive StructMAE. The findings, as outlined below, provide valuable insights into the significance of each component: 1) Efficacy of Pre-defined Metrics: As shown in Table 3a, the PageRank metric consistently demonstrates superior performance compared to other pre-defined metrics. However, other metrics also demonstrate commendable performances, suggesting their potential applications in specific scenarios. 2) Scoring Function Comparison: As displayed in Table 3b, we observe that the performance of GNN scoring function outperforms the Multilayer Perceptron (MLP). This outcome emphasizes the importance of incorporating structural information into the scoring process. Furthermore, the combined use of GNN and MLP consistently yields superior performance compared to using either one independently. 3) Masking Strategy Efficiency: The results are detailed in Table 3c, where Top, Middle, and Bottom denote the masking of the top, middle, and bottom 𝑝 ∗ 𝑛 nodes, respectively, based on the masking probability. Analysis of the results indicates that the easy-tohard masking strategy consistently surpasses the performance of other methods. This finding supports our perspective on the effectiveness of gradually increasing the learning challenge, validating the strategic design of our training approach. 4) Role of Random Noise: Notably, a significant drop in performance is observed in the absence of random noise, as indicated in Table 3d. This indicates that the inclusion of randomness in the model enables access to a broader range of information and simultaneously strengthens the learning process robustness.

6 Conclusion

This study proposes the StructMAE model, a novel structureguided masking strategy that incorporates prior structural knowledge into the masking process, thereby enhancing the pre-training model's learning efficiency. The StructMAE framework consists of two pivotal steps: SBS and SGM. Extensive experiments, encompassing two distinct graph learning tasks, demonstrate that StructMAE significantly outperforms existing self-supervised pre-training methods. These results highlight our approach's effectiveness in leveraging structural information for improved model performance. Despite its competitive performance, StructMAE still has room for improvement. For instance, 1) devising more effective scoring methods to fully exploit the structural information, 2) extending the structure-guided masking strategy to encompass edge masking, and 3) expanding structure-guided masking to a broader spectrum of tasks (e.g., node classification).

Acknowledgments

This work was partially supported by the Major Science and Technology Innovation 2030 "New Generation Artificial Intelligence" key project (No. 2021ZD0111700). The work of Wenbin Hu was supported by the National Key Research and Development Program of China (2023YFC2705700). This work was supported in part by the Natural Science Foundation of China (No. 82174230), Artificial Intelligence Innovation Project of Wuhan Science and Technology Bureau (No. 2022010702040070), Natural Science Foundation of Shenzhen City (No. JCYJ20230807090211021).

Contribution Statement

Chuang Liu and Yuyao Wang contributed equally to this work.