SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(b) Ablating the tokenizer. Darker blue indicates higher performance. The encoder is GTS and decoder is GTS-Small. Remask-v2 is used.
Table 3: Average ROC-AUC (%) scores on the eight classification datasets in MoleculeNet.
Encoder	Decoder	Remask	Avg.
GINE	Linear	-	73.2
GINE	GINE-Small	-	73.0
GINE	GINE-Small	v1	74.4
GTS	Linear	-	74.1
GTS	GTS-Small	-	74.1
GTS	GTS-Small	v1	75.2
GTS	GTS-Small	v2	75.8
Tokenizer	GNN tokenizer 's depth
-	1	2	3	4	5
Node	74.7
Motif, MGSSL	75.2
Motif, RelMole	74.4
Pretrain, GIN, GraphCL	75.1	74.5	74.2	74.0	74.6
Pretrain, GIN, VQ-VAE	75.1	74.9	74.4	75.6	75.1
Pretrain, GIN, GraphMAE	75.1	74.9	74.9	75.4	75.2
SGT, GIN	75.8	75.5	75.6	75.3	74.9
(a) MSE loss for pretraining
Figure 5: GTS encoder. Tokenizer is a single-layer SGT of GIN.
(b) Probe encoder for masked atoms.
Figure 6: MGM with SGT of different GNNs. GTS encoder and GTS-Small decoder.
4 Rethinking Masked Graph Modeling for Molecules
Experimental setting. In this section, we perform experiments to assess the roles of tokenizer and decoder in MGM for molecules. Our experiments follow the transfer learning setting in [12, 9]. We pretrain MGM models on 2 million molecules from ZINC15 [42], and evaluate the pretrained models on eight classification datasets in MoleculeNet [28]: BBBP, Tox21, ToxCast, Sider, ClinTox, MUV, HIV, and Bace. These downstream datasets are divided into train-valid/test sets by scaffold split to provide an out-of-distribution evaluation setting. We report the mean performances and standard deviations on the downstream datasets across ten random seeds. Throughout the experiments, we use random node masking of ratio 0.35. More detailed experimental settings can be found in Appendix D.
4.1 Rethinking Decoder
We investigate the impact of an expressive decoder on MGM's performance. A single-layer GIN-based SGT tokenizer is utilized in these experiments. Table 3a summarizes the results.
Finding 1. A sufficiently expressive decoder with remask decoding is crucial for MGM. Table 3a shows that using a sufficiently expressive decoder with remask decoding can significantly improve downstream performance. This can be attributed to the disparity between molecule reconstruction and MRL tasks: the last several layers in the autoencoder will be specialized on molecule reconstruction, while losing some representation ability [15, 38]. When using a linear decoder, the last several layers in the encoder will be reconstruction-specialized, which can yield suboptimal results in fine-tuning.
A sufficiently expressive decoder is essential to account for the reconstruction specialization. As shown by Figure 5a, employing an expressive decoder results in significantly lower reconstruction loss than a linear decoder. However, increasing the decoder's expressiveness without remask decoding cannot improve the downstream performance (Table 3a). This leads to our exploration on remask.
Finding 2. Remask constrains the encoder's efforts on graph reconstruction for effective MRL. Figure 5b shows that the linear probing accuracy of the masked atom types on an encoder pretrained with remask is significantly lower than the accuracy on an encoder pretrained without remask. This shows that remask makes the encoder spend less effort on predicting the corrupted information. Moreover, remask only slightly sacrifices the autoencoder's reconstruction ability, when paired with the GTS-Small decoder (Figure 5a). Combining the
Table 4: Testing decoder's size. Remask-v2 is used.
Encoder	Decoder	Avg.
GTS	GTS-Tiny	74.7
GTS	GTS-Small	75.8
GTS	GTS	74.9
observation that remask improves downstream performances (Table 3a), it indicates that remask constrains the encoder's efforts on graph reconstruction, allowing it to focus on MRL.
In addition, remask-v2 outperforms remask-v1 with GTS architecture. This improvement can be attributed to remask-v2's ability to prevent the Transformer layers from processing the masked nodes, avoiding the gap of using masked nodes in pretraining but not in fine-tuning. Finally, we test decoder's sizes in Table 4. When the encoder is GTS, GTS-Small decoder provides the best performance.
4.2 Rethinking Tokenizer
We investigate the impact of tokenizers on MGM's performance. In the following experiments, the graph autoencoder employs the GTS encoder and the GTS-Small decoder with remask-v2, given their superior performances in the previous section. The results are summarized in Table 3b.
Compared tokenizers. We use the node tokenizer as the baseline. For motif-based tokenizers, we employ the leading fragmentation methods: MGSSL [1] and RelMole [18]. For GNN-based tokenizers, we compare the prevalent pretraining strategies – GraphCL [4], GraphMAE [9], and VQ-VAE [10, 43] – to pretrain tokenizers on the 2 million molecules from ZINC15.
Finding 3. Reconstructing subgraph-level tokens can give rise to MRL. Table 3b shows that, given the appropriate setup, incorporating motif-based tokenizers or GNN-based tokenizers in MGM can provide better downstream performances than the node tokenizer. This observation underscores the importance of applying a subgraph-level tokenizer in MGM.
Finding 4. Single-layer SGT outperforms or matches other tokenizers. Table 3b shows that a single-layer SGT, applied to GIN, delivers comparable performance to a four-layer GIN pretrained by VQ-VAE, and surpasses other tokenizers. Further, Figure 6 shows that SGT can transform GCN and GraphSAGE into competitive tokenizers. We attribute SGTs' good performances to the effectiveness of their graph operators in extracting structural information. It has been shown that linear graph operators can effectively summarize structural patterns for node classification [39, 40].
Finding 5. GNN-based tokenizers have achieved higher performances than motif-based tokenizers. We hypothesize that this is due to GNN's ability of summarizing structural patterns. When using GNN's representations as reconstruction targets, the distance between targets reflects the similarity between their underlying subgraphs – a nuanced relationship that the one-hot encoding of motifs fails to capture. We leave the potential incorporation of GNNs into motif-based tokenizers for future works.
Finally, we show that although GNN-based tokenizers are agnostic to chemistry knowledge, incorporating them in MGM can improve the recognition of FGs. In Figure 7, we use a linear classifier to
probe the encoder's mean pooling output to predict the FGs within the molecule. We use RDkit [30] to extract 85 types of FG. Details are in Appendix D. It can be seen that, incorporating a single-layer SGT in MGM improves the encoder's ability to identify FGs in comparison to the node tokenizer.
Figure 7: Linear probing encoder's output for FGs.
5 Comparison to the State-of-the-Art Methods
In this section, we compare SimSGT to the leading molecule SSL methods for molecular property prediction and a broader range of downstream tasks. For fair comparison, we report the performances of SimSGT and its variant that uses the GINE architecture. This variant employs a GINE as encoder and a GINE-Small as decoder (Table 2), and implements remask-v1 decoding.
Molecular property prediction. The molecular property prediction experiment follows the same transfer learning setting as Section 3. Table 5 presents the results. It can be observed that SimSGT outperforms all the baselines in average performances with both the GTS and GINE architectures. Notably, SimSGT with GTS establishes a new state-of-the-art of  $75.8\%$  ROC-AUC. It exceeds the second method by an absolute  $1.8\%$  points in average performance and achieves the best performance in five out of the eight molecular property prediction datasets. SimSGT with GINE outperforms the baselines by  $0.4\%$  points in average performance. These improvements demonstrate the effectiveness of our proposed tokenizer and decoder to improve MGM performance for molecules.
Table 5: Transfer learning ROC-AUC (%) scores on eight MoleculeNet datasets. Bold denotes the best performance. Baselines are reproduced using source codes. All methods use the GTS encoder.
Dataset	BBBP	Tox21	ToxCast	Sider	ClinTox	MUV	HIV	Bace	Avg.
No Pre-train	68.7±1.3	75.5±0.5	64.5±0.8	58.0±1.0	71.2±4.9	69.7±1.5	74.2±2.2	77.2±1.6	69.9
Infomax	69.2±0.7	74.6±0.5	61.8±0.8	60.1±0.7	74.8±2.7	74.8±1.5	75.0±1.3	76.3±1.8	70.8
ContextPred	68.5±1.1	75.3±0.6	63.2±0.5	61.3±0.8	74.9±3.2	70.9±2.4	76.8±1.1	80.2±1.8	71.4
GraphCL	70.4±1.1	73.8±1.0	63.1±0.4	60.4±1.3	77.8±3.0	73.8±2.0	75.6±0.9	78.3±1.1	71.6
JOAO	70.8±0.6	75.2±1.5	63.8±0.8	61.0±0.9	78.7±3.0	75.7±2.6	77.0±1.2	77.5±1.2	72.5
ADGCL	67.9±1.0	73.0±1.2	63.2±0.4	60.2±1.0	78.8±1.6	75.6±2.5	75.5±1.3	73.4±1.6	71.0
BGRL	72.7±0.9	75.8±1.0	65.1±0.5	60.4±1.2	77.6±4.1	76.7±2.8	77.1±1.2	74.7±2.6	72.5
GraphLOG	67.6±1.6	76.0±0.8	63.6±0.6	59.8±2.1	79.1±3.2	72.8±1.8	72.5±1.6	83.6±2.0	71.9
RGCL	71.4±0.8	75.7±0.4	63.9±0.3	60.9±0.6	80.0±1.6	75.9±1.2	77.8±0.6	79.9±1.0	73.2
S2GAE	67.6±2.0	69.6±1.0	58.7±0.8	55.4±1.3	59.6±1.1	60.1±2.4	68.0±3.7	68.6±2.1	63.5
Mole-BERT	70.8±0.5	76.6±0.7	63.7±0.5	59.2±1.1	77.2±1.4	77.2±1.1	76.5±0.8	82.8±1.4	73.0
GraphMAE	71.7±0.8	76.0±0.9	65.8±0.6	60.0±1.0	79.2±2.2	76.3±1.9	75.9±1.8	81.7±1.6	73.3
GraphMAE2	71.6±1.6	75.9±0.8	65.6±0.7	59.6±0.6	78.8±3.0	78.5±1.1	76.1±2.2	81.0±1.4	73.4
SimSGT	72.2±0.9	76.8±0.9	65.9±0.8	61.7±0.8	85.7±1.8	81.4±1.4	78.0±1.9	84.3±0.6	75.8
Table 6: Transfer learning performance for molecular property prediction (regression) and drug target affinity (regression). Bold indicates the best performance and underline indicates the second best. * denotes reproduced result using released codes. Other baseline results are borrowed from [ICLR'22].
Molecular Property Prediction (RMSE ↓)	Drug-Target Affinity (MSE ↓)
ESOL	Lipo	Malaria	CEP	Avg.	Davis	KIBA	Avg.
No Pre-train	1.178±0.044	0.744±0.007	1.127±0.003	1.254±0.030	1.076	0.286±0.006	0.206±0.004	0.246
ContextPred	1.196±0.037	0.702±0.020	1.101±0.015	1.243±0.025	1.061	0.279±0.002	0.198±0.004	0.238
AttrMask	1.112±0.048	0.730±0.004	1.119±0.014	1.256±0.000	1.054	0.291±0.007	0.203±0.003	0.248
JOAO	1.120±0.019	0.708±0.007	1.145±0.010	1.293±0.003	1.066	0.281±0.004	0.196±0.005	0.239
GraphMVP	1.064±0.045	0.691±0.013	1.106±0.013	1.228±0.001	1.022	0.274±0.002	0.175±0.001	0.225
Mole-BERT*	1.192±0.028	0.706±0.008	1.117±0.008	1.078±0.002	1.024	0.277±0.004	0.210±0.003	0.243
SimSGT, GINE	1.039±0.012	0.670±0.015	1.090±0.013	1.060±0.011	0.965	0.263±0.006	0.144±0.001	0.204
SimSGT, GTS	0.917±0.028	0.695±0.012	1.078±0.012	1.036±0.022	0.932	0.251±0.001	0.153±0.001	0.202
Broader range of downstream tasks. We report the transfer learning performances on regressive molecular property prediction and drug-target affinity (DTA) tasks. DTA aims to predict the affinity scores between molecular drugs and target proteins [44, 29]. Specifically, we substitute the molecule encoder in [44] with a SimSGT pretrained encoder to evaluate DTA performances. Following the experimental setting in [45], we pretrain SimSGT on the 50 thousand molecule samples from the GEOM dataset [46] and report the mean performances and standard deviations across three random seeds. We report the RMSE for the molecular property prediction datasets with scaffold splitting, and report the MSE for the DTA datasets with random splitting. The results are summarized in Table 6. It can be observed that SimSGT achieves significant improvement over the baseline models.
Table 7: Mean Average Error (MAE) performance on the QM datasets. GTS encoder is used.
#Tasks	QM7 1	QM8 12	QM9 12
GraphCL	80.4±3.3	0.0200±0.0004	5.76±0.37
GraphMAE	78.4±2.3	0.0190±0.0003	5.84±0.16
Mole-BERT	79.8±2.6	0.0190±0.0003	5.75±0.16
SimSGT	75.4±0.7	0.0183±0.0003	5.53±0.25
Quantum chemistry property prediction. We report performances of predicting the quantum chemistry properties of molecules [47]. Following [48], we divide the downstream datasets by scaffold split. This experiment reuses the model checkpoints from Table 5, which are pretrained on 2 million molecules from ZINC15. Specifically, we attach a two-layer MLP after the pretrained molecule encoders and fine-tune the models for property prediction. We report average performances
Table 8: Time spent for pretraining 100 epochs on ZINC15 when using the GTS encoder.
Model	GraphMAE	Mole-BERT	S2GAE	GraphMAE2	SimSGT
Pretrain Time	527 min	2199 min	1763 min	1195 min	645 min
and standard deviations across 10 random seeds. The performances are reported in Table 7. We can observe that SimSGT consistently outperforms representative baselines of GraphCL, GraphMAE, and Mole-BERT.
Computational cost. We compare the wall-clock pretraining time for SimSGT and key baselines in Table 8. We can observe that: 1) SimSGT's pretraining time is on par with GraphMAE [9]. This efficiency is largely attributed to the minimal computational overhead of our SGT tokenizer; 2) In comparison to Mole-BERT [10], the prior benchmark in molecule SSL, SimSGT is approximately three times faster. The computational demands of Mole-BERT can be attributed to its combined approach of MGM training and contrastive learning.
In this work, we extensively investigate the roles of tokenizer and decoder in MGM for molecules. We compile and evaluate a comprehensive range of molecule fragmentation functions as molecule tokenizers. The results reveal that a subgraph-level tokenizer gives rise to MRL performance. Further, we show by empirical analysis that a sufficiently expressive decoder with remask decoding improves the molecule encoder's representation quality. In light of these findings, we introduce SimSGT, a novel MGM approach with subgraph-level tokens. SimSGT features a Simple GNN-based Tokenizer capable of transforming GNNs into effective graph tokenizers. It further adopts the GTS architecture for its encoder and decoder, and incorporates a new remask strategy. SimSGT exhibits substantial improvements over existing molecule SSL methods. For future works, the potential application of molecule tokenizers to joint molecule-text modeling [3], remains an interesting direction.
This research is supported by the National Natural Science Foundation of China (9227010114) and the University Synergy Innovation Program of Anhui Province (GXXT-2022-040). This material is based upon work supported by the Google Cloud Research Credit program with the award (6NW8-CF7K-3AG4-1WH1). This research is supported by NExT Research Center.
[1] Zaixi Zhang, Qi Liu, Hao Wang, Chengqiang Lu, and Chee-Kong Lee. Motif-based graph self-supervised learning for molecular property prediction. In NeurIPS, pages 15870-15882, 2021.
[2] Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. In NeurIPS, 2020.
[3] Zhiyuan Liu, Sihang Li, Yanchen Luo, Hao Fei, Yixin Cao, Kenji Kawaguchi, Xiang Wang, and Tat-Seng Chua. Molca: Molecular graph-language modeling with cross-modal projector and uni-modal adapter. In EMNLP, 2023.
[4] Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. In NeurIPS, 2020.
[5] Laurianne David, Amol Thakkar, Rocio Mercado, and Ola Engkvist. Molecular representations in ai-driven drug discovery: a review and practical guide. Journal of Cheminformatics, 12(1):1-22, 2020.
[6] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In ICML, pages 1263-1272. PMLR, 2017.
[7] Chaochao Yan, Qianggang Ding, Peilin Zhao, Shuangjia Zheng, Jinyu Yang, Yang Yu, and Junzhou Huang. Retroxpert: Decompose retrosynthesis prediction like A chemist. In NeurIPS, 2020.
[8] Umit V Ucak, Islambek Ashyrmamatov, Junsu Ko, and Juyong Lee. Retrosynthetic reaction pathway prediction through neural machine translation of atomic environments. Nature communications, 13(1):1186, 2022.
[69] Petar Velickovic, William Fedus, William L Hamilton, Pietro Liò, Yoshua Bengio, and R Devon Hjelm. Deep graph infomax. In ICLR, 2019.
[70] Fan-Yun Sun, Jordan Hoffmann, Vikas Verma, and Jian Tang. Infograph: Unsupervised and semi-supervised graph-level representation learning via mutual information maximization. In ICLR, 2020.
[71] Yuning You, Tianlong Chen, Yang Shen, and Zhangyang Wang. Graph contrastive learning automated. In ICML, Proceedings of Machine Learning Research, pages 12121-12132, 2021.
[72] Susheel Suresh, Pan Li, Cong Hao, and Jennifer Neville. Adversarial graph augmentation to improve graph contrastive learning. In NeurIPS, pages 15920-15933, 2021.
[73] Minghao Xu, Hang Wang, Bingbing Ni, Hongyu Guo, and Jian Tang. Self-supervised graph-level representation learning with local and global structure. In ICML, volume 139, pages 11548-11558, 2021.
[74] Sihang Li, Xiang Wang, An Zhang, Yingxin Wu, Xiangnan He, and Tat-Seng Chua. Let invariant rationale discovery inspire graph contrastive learning. In ICML, pages 13052-13065, 2022.
[75] Qiaoyu Tan, Ninghao Liu, Xiao Huang, Soo-Hyun Choi, Li Li, Rui Chen, and Xia Hu. S2GAE: self-supervised graph autoencoders are generalizable learners with graph masking. In WSDM, pages 787-795. ACM, 2023.
[76] Zhenyu Hou, Yufei He, Yukuo Cen, Xiao Liu, Yuxiao Dong, Evgeny Kharlamov, and Jie Tang. Graphmae2: A decoding-enhanced masked self-supervised graph learner. In WWW, volume abs/2304.04779, 2023.
[77] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint arXiv:2005.00687, 2020.
[78] Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. In AAAI, pages 3538-3545. AAAI Press, 2018.
[79] Joshua David Robinson, Ching-Yao Chuang, Suvrit Sra, and Stefanie Jegelka. Contrastive learning with hard negative samples. In ICLR. OpenReview.net, 2021.
[80] Yoav Levine, Barak Lenz, Opher Lieber, Omri Abend, Kevin Leyton-Brown, Moshe Tennenholtz, and Yoav Shoham. Pmi-masking: Principled masking of correlated spans. In ICLR. OpenReview.net, 2021.
Our results and analysis on the graph tokenizer and graph decoder are confined to the task of MGM pretraining. Different tokenizers and decoders might offer advantages for other generative modeling methods, such as autoregressive modeling [32].
SGTs [39, 40] are limited in expressive power for graph structures compared to standard GNNs, like GINs [37]. Theoretically, the separation of expressiveness power between SGTs and standard GNNs grows exponentially in the GNN's depth [40]. However, SGTs exhibit comparable, if not better, performances to pretrained GNN-based tokenizers, as demonstrated in Table 3b. We attribute this intriguing observation to two key factors. Firstly, SGTs (i.e., simple GNNs) are still powerful and can "distinguish almost all non-isomorphic graphs" [40]. They have shown decent results in practice [39, 49]. Secondly, we conjecture that a better pretraining method for GNN-based tokenizers could exist, but current pretraining techniques do not fully harness the potential of GNNs in their roles as effective tokenizers. Indeed, the significant difference in performance between GraphCL and VQ-VAE (Table 3b) emphasizes the impact of pretraining methods on the tokenizer's performance. We leave the investigation of how to effectively pretrain GNN-based tokenizers as future works.
We have included the literature review of MGM in the main body of the paper. Here we elaborate on the literature review in the following areas.
Molecule SSL with motifs. Motifs are statistically significant subgraph patterns [32, 50], and have been applied in existing molecule SSL methods. Autoregressive pretraining methods [32, 1, 50] generate motifs instead of nodes at each generation step, in order to improve the validity of the generated molecules. Motifs are also used in contrastive learning [17, 51, 18]. Sun et al. [17] substitute motifs within a molecule with their chemically similar counterparts to create high-quality augmentations. [51, 18] construct molecules' views at the motif-level to supplement the original views at the atom-level. In predictive pretraining, Rong et al. [2] pretrain a graph encoder to predict the FGs inside the molecule. These previous works have developed extensive molecule fragmentation methods for motif discovery. However, these fragmentation methods have been overlooked as tokenizers in MGM pretraining. Our work addresses this gap by summarizing the common fragmentation rules and examining the performances of the selected fragmentation methods in MGM pretraining.
Data tokenization. Tokenization is a data pre-processing technique that divides the original data into smaller elements and converts them to tokens. It is widely used in NLP to split sentences into word-level units [13, 52, 53]. Due to the surging interests in Transformers [54], tokenization is also applied on images [15, 16] and audios [55, 56]. Tokenization fragments these data into sequences of patches to fit the shapes of transformer's input and output. A tokenizer can be designed by heuristics [57], incorporating domain knowledge [58], and pretraining on the target dataset [59, 19, 43]. In this work, we study graph tokenizers, which are less explored in previous works.
Relations to contrastive learning. When using a GNN-based tokenizer, MGM involves minimizing the distances between the outputs from two network branches (i.e., the tokenizer branch and the autoencoder branch). At first glance, this design might seem similar to the contrastive learning methods of BYOL [60], SimSiam [61], and BGRL [62], which also minimize the output differences between two network branches. However, a closer inspection reveals several critical distinctions between MGM and these methods. Firstly, MGM feeds uncorrupted data to the tokenizer branch and feeds corrupted data to the autoencoder branch, encouraging the autoencoder to reconstruct the missing information. In contrast, BYOL, SimSiam, and BGRL use corrupted data in both of their branches, constituting different training objectives. Secondly, while BYOL, SimSiam, and BGRL employ nearly identical architectures for their two branches, MGM can adopt distinctly different architectures for its autoencoder and tokenizer. In our best-performing experiment, the autoencoder has more than ten layers of GNNs and Transformers, while the tokenizer is a shallow single-layer network (Table 3). Finally, MGM employs remask decoding to constrain the encoder's ability on reconstruction, which is not used in contrastive learning methods [60, 61, 62].
Subgraph-enhanced Graph Neural Network. Subgraph-enhanced GNN [63, 64, 65] refers to an emerging class of GNNs that fragments a graph into subgraphs before encoding, in order to improve the GNNs' expressiveness [66, 67, 68]. The common graph fragmentation method is node-wise, such
that each fragmented subgraph is associated with a unique node in the original graph. For example, ESAN [63] obtains subgraphs by sampling ego-networks or deleting one node from the original graph. Given the subgraphs, subgraph-enhanced GNNs generate node embeddings in every subgraph by applying a series of equivariant message-passing layers [63, 64]. Finally, these embeddings are pooled to output the graph embedding. Our work is related to subgraph-enhanced GNNs that we also study graph fragmentation. The major distinction is that we focus on using the tokens derived from these fragmented graphs as the reconstruction targets in MGM for molecules.
C Pseudo Code
We present the pseudocode of SimSGT. This code uses a single-layer SGT of GIN as an example.
Algorithm 1 Pytorch style pseudocode of SimSGT
# # phi: graph encoder
# # rho: graph decoder
def SGT(g, embed):
    # # SGT: a single-layer GIN tokenizer
    x, edge_index = g
    # message passing
    x = propagate(embed(x), edge_index) + (1+eps) * embed(x)
    # batch normalization layer
    x = batchnorm(x)
    return x
for g in loader:
    # random masking
    g_hat, m_pos = random_masking(g) # m_pos: mask positions
    # tokenization. embed is a linear layer
    y = SGT(g, phi_embedding).detach() # detach: stop-gradient
    # autoencoder forward
    y_hat = rho(remask phi(g_hat, m_pos))
    # minimize loss
    loss = distance_loss(y_hat[m_pos], y[m_pos])
    loss_backward()