SlowGuess/NeurIPS_2023_OCR huggingface dataset:
are identical as those in [48]. We show the performance gains brought by their GMoE counterparts: GMoE-GCN and GMoE-GIN. For GMoE models, as described in Section 3, we select  $k$  experts out of a total of  $n$  experts for each node, where  $m$  out of  $n$  experts are hop-1 aggregation functions and the rest  $n - m$  are hop-2 aggregation functions. All three hyper-parameters  $n, m, k$ , together with the loss trade-off weight  $\lambda$  in Eq. (9), are tuned by grid searching:  $n \in \{4,8\}$ ,  $m \in \{0, n/2, n\}$ ,  $k \in \{1,2,4\}$ , and  $\lambda \in \{0.1,1\}$ . The hidden feature dimension would be adjusted with  $k$  (Section 3.2) to ensure the same flops for all comparisons. The hyper-parameter values achieving the best performance on validation sets are selected to report results on test sets, following the routine in [48]. All other training hyper-parameters on GMoE (e.g., batch size, learning rate, training epochs) are kept the same as those used on the single-expert baselines. All experiments are run for ten times with different random seeds, and we report the mean and deviation of the results following [48].
Pre-training Settings Following the transfer learning setting of [49, 6, 8, 10], we pre-train the models on a subset of the ZINC15 dataset [50] containing 2 million unlabeled molecule graphs. For training hyperparameters, we employ a batch size of 1024 to accelerate the training on the large pre-train dataset for both baselines and the proposed method. We follow [10] employing GIN [13] as the backbone, 0.001 as the learning rate, adam as the optimizer, 0 as the weight decay, 100 as the training epochs number, and 0.25 as the masking ratio.
Our evaluation primarily centers around comparing GMoE-GCN with the baseline single-expert GCN using six graph property prediction datasets. These datasets encompass four graph classification tasks and two regression tasks within a supervised learning framework. In all cases, models are trained from scratch, utilizing only the labeled samples in the training set. The classification and regression results are outlined in Table 1 and 2, respectively.
Table 5: Link Prediction Task Results. The table header provides the dataset names and the respective evaluation metrics, both expressed in percentage. Mean and standard deviation values from ten random runs are presented. The most noteworthy results are highlighted in bold, with GMoE's performance improvements indicated in parentheses.
Model	ogbl-ppa
Hits@100(↑)	ogbl-ddi
Hits@20(↑)
GCN	18.67 ± 1.32	37.07 ± 0.051
GMoE-GCN	19.25 ± 1.67
(+0.58)	37.96 ± 0.082
(+0.89)
Table 6: Graph Property Prediction Results employing GMoE-GCN with Varied Hyper-Parameter Configurations. The third column displays the average number of nodes within graphs across different datasets. Mean and standard deviation values from ten random runs are provided. The most noteworthy outcomes are emphasized in bold. ROC-AUC scores are expressed in percentage.
Dataset	Metric	Average #Nodes	n=m=4 k=4	n=m=4 k=1	n=m=8 k=4	n=8,m=0 k=4
ogbg-molfreesolv	RMSE (↓)	8.7	2.856 ± 0.118	2.500 ± 0.193	2.532 ± 0.231	2.873 ± 0.251
ogbg-molhiv	ROC-AUC (↑)	25.5	76.67 ± 1.25	76.41 ± 1.72	77.53 ± 1.42	77.87 ± 1.03
It is worth noting that GMoE-GCN consistently outperforms the baseline across all six datasets. Notably, there are substantial improvements in ROC-AUC, with increases of  $1.81\%$  on the molhiv dataset and  $1.41\%$  on the molbbp dataset. While these enhancements may seem modest, they represent significant progress. Additionally, lifts of  $1.40\%$  on ogbg-moltox21,  $1.43\%$  on ogbg-molbbbq, and  $1.18\%$  on ogbg-moltoxcast have been observed. The uniformity of these improvements across diverse tasks and datasets underscores the reliability and effectiveness of GMoE-GCN.
Leveraging large-scale pretraining on auxiliary unlabeled data has notably enhanced the generalization capabilities of GNNs even further. Expanding on this, our GMoE consistently improves performance when integrated with large-scale pretraining methods. Following the methodology outlined in [10], we employ GIN [13] as the baseline model for comparison with GMoE-GIN. As illustrated in Table 3, GMoE-GIN outperforms GIN on 3 out of 4 datasets, enhancing the average performance by  $0.54\%$ , even without pretraining. This underscores the versatility of GMoE in enhancing various model architectures. When coupled with the pretraining, GMoE-GIN further widens the performance gap with the GIN baseline, achieving a slightly more pronounced improvement margin of  $0.6\%$
Additionally, GMoE showcases its potential in node and link prediction tasks. Experiments conducted on ogbn-protein and ogbn-arxiv for node prediction, as well as ogbl-ddi and ogbl-ppa for link prediction, further validate this observation. The results, detailed in Table 4 and Table 5, demonstrate GMoE-GCN's superiority over the single-expert GCN. Notably, it enhances performance metrics like ROC-AUC by  $0.95\%$  and Hits@20 by  $0.89\%$  on the ogbn-protein and ogbl-ddi datasets.
4.3 Ablation Study and Analysis
Observation 1: Larger graphs prefer larger hop sizes We conducted an ablation study on two prominent molecule datasets, namely ogbg-molhiv and ogbg-molfreesolv, which exhibit the largest and smallest average graph sizes, respectively, among all molecular datasets in the OGB benchmark. As illustrated in the third column of Table 6, the average size of molecular graphs in ogbg-molhiv is approximately three times greater than that of the ogbg-molfreesolv dataset. Remarkably, on ogbg-molfreesolv, if all experts utilize hop-2 aggregation functions (i.e., when  $m = 0$ ), the performance substantially lags behind the scenario where all experts employ hop-1 aggregations (i.e., when  $m = n$ ). In contrast, on ogbg-molhiv, characterized by significantly larger graphs, leveraging all hop-2 experts (i.e., when  $m = 0$ ) yields superior performance compared to employing all hop-1 experts (i.e., when  $m = n$ ). This observation suggests that larger graphs exhibit a preference for aggregation experts with greater hop sizes. This alignment with intuition stems from the notion
Table 7: Results for graph property prediction using GMoE-GCN with different scaling ratios  $\lambda$  for the loading balance losses on moltox21. Mean and standard deviation over ten random runs are reported. The best results are shown in bold. ROC-AUC is reported in percentage.
λ	0	0.1	1
ROC-AUC	72.87±1.04	75.45±0.58	75.27±0.33
that larger molecular graphs may necessitate more extensive long-range aggregation information in contrast to their smaller counterparts.
Observation 2: Sparse expert selection helps not only efficiency, but also generalization As depicted in Table 6, the optimal performance on both datasets is attained with sparse MoEs (i.e., when  $k < n$ ), as opposed to a full dense model (i.e., when  $k = n$ ). This seemingly counter-intuitive finding shows another benefit of sparsity besides significantly reducing inference computational cost: that is, sparsity is also promising to improve model generalization, particularly when dealing with multi-domain data, as it allows a group of experts to learn collaboratively and generalize to unseen domains compositionally. Our finding also echoes prior work, e.g., [51], which empirically shows sparse MoEs are strong domain generalizable learners.
Observation 3: GMoE demonstrates performance gains, though converging over more epochs. We compared the convergence curves of GMoE-GCN and the single-expert GCN, illustrated in Figure 2. Specifically, we plotted the validation and test ROC-AUC at different epochs while training on the protein dataset. As observed, the performance of GMoE-GCN reaches a plateau later than that of GCN. However, GMoE-GCN eventually achieves a substantial performance improvement after a sufficient number of training epochs.
We provide an explanation for this phenomenon below. In the initial stages of training, each expert in GMoE has been updated for  $k$  times fewer iterations compared to the single expert GNN. Consequently, all experts in GMoE are relatively weak, leading to GMoE's inferior performance compared to GCN during this early phase. However, after GMoE has undergone enough epochs of training, all experts have received ample updates, enabling them to effectively model their respective subgroups of data. As a result, the performance of GMoE surpasses that of the single-expert GCN.
Figure 2: Convergence curve of GMoE-GCN and GCN when trained on the ogb-protein dataset. The solid curves and the shaded areas are the mean and standard deviations of ROC-AUC calculated over ten random runs.
Observation 4: Load balancing is essential We conducted an investigation into the scaling factor for the load balancing loss, denoted by  $\lambda$ . As shown in Table 7, utilizing  $\lambda = 0.1$  resulted in significantly improved performance, exceeding  $2.58\%$  in terms of ROC-AUC compared to when  $\lambda = 0$ . This underscores the crucial role
of implementing load-balancing losses. Conversely, the choice of  $\lambda$  exhibits less sensitivity; opting for  $\lambda = 1$  would yield a similar performance of  $75.27\%$ .
Observation 5: GMoE improves other state-of-the-art GNN methods Finally, we explore whether our proposed GMoE can yield an improvement in the performance of other state-of-the-art models listed on the OGB leaderboard. We selected Neural FingerPrints [52] as our baseline. As of the time of this submission, Neural FingerPrints holds the fifth position in the ogbg-molhiv benchmark. It is noteworthy that this methodology has gained significant acclaim, as it serves as the foundation for the top four current methods [53, 54, 55, 56].
After integrating the GMoE framework into the Neural FingerPrints architecture, we achieved an accuracy of  $82.72\% \pm 0.53\%$ , while maintaining the same computational resource requirements. This performance outperforms Neural FingerPrints by a margin of  $0.4\%$ , underlining the broad and consistent adaptability of the GMoE framework.
In this work, we propose the Graph Mixture of Experts (GMoE) model, aiming at addressing the challenges posed by diverse graph structures. By incorporating multiple experts at each layer, each equipped with its own trainable parameters, GMoE introduces a novel approach to modeling graph data. Through intelligent expert selection during training, GMoE ensures nodes with similar neighborhood information are directed towards the same aggregation experts, promoting specialization within each expert for specific subsets of training samples. Additionally, the inclusion of aggregation experts with distinct inductive biases further enhances GMoE's adaptability to different graph structures. Our extensive experimentation and analysis demonstrate GMoE's notable accuracy-efficiency trade-off improvements over baselines. These advancements hold great promise for real-world applications, particularly in scenarios requiring the efficient processing of vast amounts of candidate samples.
Limitations: As an empirical solution to enhance the capability of GNNs to encode diverse graph data, we primarily assess the effectiveness of the proposed method using empirical experimental results. It is important to note that our initial comparison primarily focused on fundamental backbones such as GCN/GIN to elucidate our concept. Recognizing the crucial importance of a comprehensive evaluation, we are actively broadening our comparative scope to include state-of-the-art models derived from GCN and other MoE models. Future iterations will thus incorporate a more exhaustive evaluation. Moreover, given the prevalence of GCN in semi-supervised learning, there is potential in exploring the benefits of our proposed approach in such tasks – an avenue we are eager to explore further. Another open question stemming from this work is whether we can apply GMoE on GNNs with heterogeneous aggregation mechanisms, such as GAT and GraphTransformer, which may potentially further improve the performance on diverse graph data.
The work is sponsored by a Cisco Research Grant (UTAUS-FA00002062).
[1] Shiwen Wu, Fei Sun, Wentao Zhang, Xu Xie, and Bin Cui. Graph neural networks in recommender systems: A survey. ACM Computing Surveys, 55(5):1-37, 2022.
[2] Weiwei Jiang and Jiayun Luo. Graph neural network for traffic forecasting: A survey. Expert Systems with Applications, page 117921, 2022.
[3] Oliver Wieder, Stefan Kohlbacher, Méline Kuenemann, Arthur Garon, Pierre Ducrot, Thomas Seidel, and Thierry Langer. A compact review of molecular property prediction with graph neural networks. *Drug Discovery Today: Technologies*, 37:1-12, 2020.
[4] Xiaotian Han, Zhimeng Jiang, Ninghao Liu, and Xia Hu. G-Mixup: Graph data augmentation for graph classification. arXiv preprint arXiv:2202.07179, 2022.
[5] Songtao Liu, Rex Ying, Hanze Dong, Lanqing Li, Tingyang Xu, Yu Rong, Peilin Zhao, Junzhou Huang, and Dinghao Wu. Local augmentation for graph neural networks. In ICML, 2022.
[6] Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. Advances in Neural Information Processing Systems, 33:5812-5823, 2020.
[7] Yuning You, Tianlong Chen, Zhangyang Wang, and Yang Shen. When does self-supervision help graph convolutional networks? In International Conference on Machine Learning, pages 10871-10880, 2020.
[8] Yuning You, Tianlong Chen, Yang Shen, and Zhangyang Wang. Graph contrastive learning automated. In International Conference on Machine Learning, pages 12121-12132, 2021.
[44] Haotao Wang, Aston Zhang, Shuai Zheng, Xingjian Shi, Mu Li, and Zhangyang Wang. Removing batch normalization boosts adversarial training. In International Conference on Machine Learning, pages 23433-23445, 2022.
[45] Andrew Brock, Soham De, and Samuel L Smith. Characterizing signal propagation to close the performance gap in unnormalized resnets. arXiv preprint arXiv:2101.08692, 2021.
[46] Andy Brock, Soham De, Samuel L Smith, and Karen Simonyan. High-performance large-scale image recognition without normalization. In ICML, pages 1059–1071, 2021.
[47] Guohao Li, Chenxin Xiong, Ali Thabet, and Bernard Ghanem. DeeperGCN: All you need to train deeper GCNs. arXiv preprint arXiv:2006.07739, 2020.
[48] 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.
[49] Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint arXiv:1905.12265, 2019.
[50] Teague Sterling and John J Irwin. Zinc 15-ligand discovery for everyone. Journal of Chemical Information and Modeling, 55(11):2324-2337, 2015.
[51] Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. arXiv preprint arXiv:2206.04046, 2022.
[52] Weibin Li, Shanzhuo Zhang, Lihang Liu, Zhengjie Huang, Jieqiong Lei, Xiaomin Fang, Shikun Feng, and Fan Wang. Molecule representation learning by leveraging chemical information. Technical report, 2021.
[53] Lanning Wei, Huan Zhao, Quanming Yao, and Zhiqiang He. Pooling architecture search for graph classification. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pages 2091-2100, 2021.
[54] Yan Wang, Hao Zhang, Jing Yang, Ruixin Zhang, and Shouhong Ding. Technical report for ogb graph property prediction. In Technical Report, 2021.
[55] Zhuoning Yuan, Yan Yan, Milan Sonka, and Tianbao Yang. Large-scale robust deep auc maximization: A new surrogate loss and empirical studies on medical image classification. In CVPR, pages 3040-3049, 2021.
[56] Hao Zhang, Jiaxin Gu, and Pengcheng Shen. Gman and bag of tricks for graph classification, 2021.