Graph Mixture of Experts: Learning on Large-Scale Graphs with Explicit Diversity Modeling
Authors: Haotao Wang
Ziyu Jiang
Yuning You
Yan Han
Gaowen Liu
Jayanth Srinivasa
Ramana Rao Kompella
Zhangyang Wang
Venue: NeurIPS 2023 poster (2023)
Keywords: graph neural networks

Abstract: Graph neural networks (GNNs) have found extensive applications in learning from graph data. However, real-world graphs often possess diverse structures and comprise nodes and edges of varying types. To bolster the generalization capacity of GNNs, it has become customary to augment training graph structures through techniques like graph augmentations and large-scale pre-training on a wider array of graphs. Balancing this diversity while avoiding increased computational costs and the notorious trainability issues of GNNs is crucial. This study introduces the concept of Mixture-of-Experts (MoE) to GNNs, with the aim of augmenting their capacity to adapt to a diverse range of training graph structures, without incurring explosive computational overhead. The proposed Graph Mixture of Experts (GMoE) model empowers individual nodes in the graph to dynamically and adaptively select more general information aggregation experts. These experts are trained to capture distinct subgroups of graph structures and to incorporate information with varying hop sizes, where those with larger hop sizes specialize in gathering information over longer distances. The effectiveness of GMoE is validated through a series of experiments on a diverse set of tasks, including graph, node, and link prediction, using the OGB benchmark. Notably, it enhances ROC-AUC by $1.81\%$ in ogbg-molhiv and by $1.40\%$ in ogbg-molbbbp, when compared to the non-MoE baselines. Our code is publicly available at https://github.com/VITA-Group/Graph-Mixture-of-Experts.

Graph Mixture of Experts: Learning on Large-Scale Graphs with Explicit Diversity Modeling
Haotao Wang'; Ziiu Jiang"; Yuning You?, Yan Han', Gaowen Liu Jayanth Srinivasa
Ramana Rao Kompella", Zhangyang Wang'
UUniversity of Texas Austin [htwang, yh9442, atulassa edu, {jiangziyu, yuning. yaaangyyoo@tttm.eeeee edu, {gaoliu, jasriniv, @xompell]@cisco.. com
Texas A&M University²
Cisco Systems
Abstract
Graph neural networks (GNNs) have found extensive applications learning from graph data. However, real-world graphs often possess diverse structures and comprise nodes and edges of varying types. To bolster the generalization capacity of GNNs, it has become customary to augment training graph structures through techniques like graph augmentations and large-scale pre-training on wider array of graphs. Balancing this diversity while avoiding increased computational costs and the notorious trainability issues of GNNs crucial. This study introduces the concept of fixture-oo--Epeert (MoE) to GNNs, with the aim of augmenting their capacity to adapt to diverse range of training graph structures, without incurring explosive computational overhead. The proposed Graph Mixture of Experts (GMoE) model empowers individual nodes in the graph to dynamically and adaptively select more general information aggregation experts. These experts are trained to capture distinct subgroups of graph structures and to incorporate information with varying hop sizes, where those with larger hop sizes specialize in gathering information over longer distances. The effectiveness of GMoE is validated through series of experiments graph, node, and link prediction, using the OGB benchmark. Notably, it enhances ROC-AUC by 1.81% in ogbg-molhiv and by 1.40% in ogbg- molbbbp, when compared to the non-MoE baselines. Our code is publicly available at https; //github. com Citon/VITA-Gaan wwt..ooo//VVVAAoppppa SAApTTA--rrro//GGraa-iiittttteee of of-xxertss
diverse set of tasks, including
Introduction
Graph learning has found extensive use in various real-world applications, including recommendation systems [1, traffic prediction (2). and molecular property prediction [ Real-world graph data typically exhibit diverse graph structures and heterogeneous nodes and edges. In graph- recommendation systems, for instance, node can represent product or customer, while an edge can indicate different interactions such as view, like, or purchase. Similarly in biochemistry tasks, datasets can comprise molecules with various biochemistry properties and thherefore various graph structures. Moreover, purposefully increasing the diversity of graph data structures in training sets has become crucial aspect of GNN training. Techniques such graph data augmentations [4] 5] and large-scale pre- pre-training on diverse graphs 6 10) have been widely adopted to allow GNNs for extracting robust and generalizable features. Meanwhile, many real-world GNN applications, such as recommendation systems and molecule virtual screening, usually involve processing a vast number of candidate samples and therefore demand computational efficiency. That invites the key
"The first two authors contributed equally.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

question: Can one effectively scale GNN model's 's capacity to leverage larger-scale, more diverse graph data, without compromising its inference efficiency?
A common limitation of many GNN architectures is that they are essentially "homogeneous" across the whole graph, i.e., forcing all nodes share the same aggregation mechanism, regardless of the differences in their node features or neighborhood information That might be suboptimal when training on diverse graph structures, e.g, when some nodes may require information aggregated over longer ranges while others prefer shorter-range local information. Our solution is the proposal of a novel GNN architecture dubbed Graph Mixture of Experts (GMoE). It comprises multiple "experts" at each layer, with each expert being an independent paassigg function with its own trainable parameters. The idea establishes a new base to address the diversity challenges residing in graph data,
Throughout the training process, GMoE is designed to intelligently select aggregation experts tailored to each node. Consequently, nodes with similar neighhorhood information are guided towards the same aggregation experts. This fosters specialization within each GMoF expert, focusing on specific subsets of training samples with akin neighborhood patterns, regardless of range or aggregation levels. In order to harness the full spectrum of diversity, GMoE also incorporates aggregation experts with distinct inductive biases. For example, each GMoE layer is equipped with aggregation experts of varying hop sizes. Those with larger hop sizes cater to nodes requiring information from more extended ranges, while the opposite holds true for those with smaller hop sizes.
We have rigorously validated GMoE's 's effectiveness through a range of comprehensive molecular property prediction tasks, underscoring our commitment deliberate diversity modeling. Moreover, our analysis demonstrates that GMoE surpasses other GNN models in terms of inference efficiency, even when they possess similar-sized parameters, thanks to the dynamic expert selection. This efficiency proves crucial in real-world scenarios, such as virtual screening in libraries of trillion-scale magnitude or beyond. The potency of our approach is corroborated by extensive experiments on ten graph learning datasets within the OGB benchmark. For instance, GMoE enhances the ROC-AUC by 1.81% on ogbg-molhiv, 1.40% on ogbg-molbbbp, 0.95% on ogbn-proteins, and boosts Hits fiits@2 20 score by 0.89% on ogbl-ddi, when compared to the single-expert baseline. To gain deeper insights into our method, we conduct additional ablation studies and comprehensive analyses.
2 Related Work
Graph Neural Networks Graph neural networks (GNNs) [T] 12 13
have emerged as powerful
approach for learning graph representations. Variants of GNNs have been proposed [1] 12 K
achieving state-of-the-an performance in different graph tasks. Under the message passing framework
11), graph convolutional network (GCN) adopts mean pooling to aggregate the neighborhood and updates embeddings recursively Mi; GraphSAGE [5] adopts sampling and aggregation schemes to eliminate the inductive bias in degree; and graph attention network (GAT) utilizes the learnable attention weights to adaptively aggregate. To capture long-range dependencies in disassortative graphs, Geom-Gor GCNN devises geometric aggregation scheme [16] to enhance the convolution, benefiting from a continuous space underlying the graph. Lately, Graphormer [17] proposes a novel graph transformer model that utilizes attention mechanisms to capture the structural information.
Mixture of Experts The concept of Mixture of Experts (MoB) [88] has a long history, tracing 21 Recently, spurred by advancements in large language 26 has re-gained prominence. This variant selectively activates only a small subset of experts for each input, significantly enhancing efficiency and enabling the development of colossal models with trillions of parameters. This breakthrough has revolutionized the learning process, particularly on vast language datasets [75] Subsequent studies have further refined the stability and efficiency of sparse MoE (28 29]. The remarkable success of sparse MoE
its origins back to earlier work s 19
20
models, sparse MoE [22 23
24
25
in the realm of language has spurred its adoption in diverse domains, including vision 30 31]
multi-modal
(32 and multi-task task learning 33 34 33).
MoE for GNNs In the domain of graph analysis, prior research has assembled knowledge from multiple ranges by combining various GNNs with different scopes [36] 37 akin to fixed-weight
'There are exceptions. For example, GAT and GraphTrranffrme adaptively leam the aggregation function for each node. This paper focuses on discussing whether GMoE can improve common homogeneous GNNs such as GCN and GIN. Extending GMoE over heterogeneous GNNs is left for future research.

---

Node
Node
Node 3
Expert (hop-1)
Expert (hop-1)
Expert (hop-2)
Graph Mixture of Experts (Mix hop-1&hops
Figure 1: Each represents the aggregation of single node, and each column corresponds different network or sub-module. Blue dots (+) denote the input features passed to the red dots (+) via the colorful edges. On the left, demonstrate one hop-2 expert. On the right, GMoE is depicted. In this instance, the proposed GMoE selectively chooses one expert for each node while masking the others. Best viewed in color.
hop-1 experts with distinct weights, along with
Mixture of Experts (MoE). Pioneering efforts [38 39] have also investigated the application of MoE address the well-known issue of imbalance and develop unbiased classification generalization algorithms. However, none of these approaches harnessed the potential advantages of sparsity and adaptivity. Recent work [40] introduced the use of mixture of experts for molecule property prediction. They employed a GNN as a feature extractor and applied a mixture of experts, where contrast, each expert is linear classifier, on top the extracted features for graph classification.
each layer of GMoE constitutes a mixture of experts, with each expert being GCN/GIN layer featuring different aggregation step sizes. Another distinction from [40 lies their utilization of domain-specific knowledge (specifically, molecule topology) for expert routing, whereas approach is designed to operate on general graphs without relying on domain -ppaifi-peciiic assumptions. One more concurrent study by [41] employs MoE to achieve fairness in predictions for GNNs.
Our study takes significant stride forward by introducing sparse MoE scale graph neural networks in an end-to-end fashion, enabling efficient learning on datasets featuring diverse graph structures. We incorporate experts with varying scopes, allowing the gating function dynamically select neighbors with the desired range.
Adapting Deep Architectures for Training Data Diversity Several prior studies have delved into enhancing the capacity of generic deep neural networks to effectively leverage wide array of training samples without incurring additional inference costs. For instance. [42) suggested employing
two distinct batch normalization (BN) layers for randomly and adversarially augmented training
samples, based on the observation that these two sets of augmented samples originate from different distributions. Building upon this concept, [43] extended it by introducing an auxiliary instance normalization layer further reduce the heterogeneity of input features before reaching the BN layers. More recently, [44] demonstrated that normalizer- free convolutional networks (CNNs) 45 46 exhibit significantly greater capability in accommodating diverse training sets compared to conventional BN-based CNNs. However, these prior works have primarily concentrated on devising improved normalization strategies for CNNs, while Graph Neural Networks (GNNs) often do rely on (batch) normalization as heavily.

---

3 Method
Preliminaries: Graph Neural Networks Taking the classical Graph Convolutional Network (GCN) [1 as an example, the propagation mechanism can be formulated
K1 = \ JEN, VIN.NNN)
Σ
h h
where W() E Rex is trainable weight and is element-wise non-linear activation function. hy E IRxx denotes the input feature of th node while h E R its output feature in th node. and are batch size and hidden feature size, respectively. N; denotes the collection of neighbors for i th node including self-connection. The output feature is normalized by VIN,I, A canonical GCN
layer only aggregates the information from immediately adjacent neighbors (hop-1).
3.1 Graph Mixture of Experts
The general framework of GMoE is outlined in Figure The GMoE layer comprises multiple experts, each utilizing either the hop-1 hop-2 aggregation function. To determine which experts to for given node, a gating function is employed. This allows for similar nodes to be assigned to the same experts when learning with diverse graph structures, thereby enabling each expert to specialize in particular structure type. By doing so, the model can more effectively capture diverse graph structures present within the training set. The GMoE layer's adaptive selection between the hop-I and hop-2 experts enables the model to dynamically capture short-range or long-range information aggregation for each node. Formally, GMoE layer can be written as:
M4 =4= =
Σ Σ JEN,
m
G(hi)aFo (1i B4y, +)+ Σ Σ O JEN?
a
G(hi)oFo (ley, Bij, W)),
(2)
where m and n denote the hop-1 and total experts number, respectively. Hench the number of hop-2 experts is n m. Fg and eij denote the message function and edge feature between i th and th nodes, respectively.
can represent multiple types of messsae-passsnng functions such as one employed by GCN 447 GIN [13. G is the gating function that generates multiple decision scores with the input of h; while G(h;)o o denotes the o th item in the output vector of G. We employ the noisy top-k gating design for G following (22), which can be formalized with
G(h,) Softmaax((op((((()))))))))
(3)
Q(h.) = W,W, + e Softplus W.,,
where k denotes the number of selected experts. E N (0,1) 1) denotes standard Gaussian noise. Wy E gaxn and Wn E g=xx are learnable weights that control clean and noisy scores, respectively.
The proposed GMoE layer can be applied to many GNN backbones such as GIN I3 or GCN [T]) In practice, we replace every layer of the backbone with its corresponding GMoE layer. For simplicity, we name the resultant network GMoE-GCN or GMoE-Gr M-N (for GCN and GIN, respectively).
simplicity
Additional Loss Functions to Mitigate GMoE Collapse Nonetheless, this model is trained solely using the expectation-mantt innectatioonmaxiiizzaionn loss, it may succumb to trivial solution wherein only single group experts is consistently selected. This arises due to the self-reinforeing nature of the imbalance: the chosen experts can proliferate at much faster rate than others, leading to their increased frequency of selection. To mitigate this, we implement two additional loss functions to The first one is importance loss:
prevent such collapse 22
Importance(H: H)
Σ
9. Limportance (H) = CV (Importtancc(HH))
(5)
by GEE()))
where the importance score Importance( H) is defined as the sum of each node's gate value y across the whole batch. CV represents the coefficient of variation. The importance loss Timportance (H) hence measures the variation of importance scores, enforcing all experts to be "similarly important".

---

While the importance score enforces equal scoring among the experts, there may still be disparities in the load assigned different experts. For instance, one expert could receive few high scores, while another might be selected by many more nodes ye all with lower scores. This situation can potentially lead to memory or efficiency issues, particularly on distributed hardware setups. To address this, we introduce an additional bad-balanced loss to encourage more even selection probability per expert. Specifically, G(h+) of 0 if and only if Q(h;), ) is greater than the k-th largest element of (hh,) excluding itself. Consequently, the probability of G(ha) o 0 can be formulated as:
P(h; o) = Pr(Q(hs). > klh_ee(Q(hh)) o)),
(6)
where kth_exx( denotes the k-th largest element excluding itself. P(h;, o) can be simplified as
P(h,, o) h, kth (x(((()),, Softplus(h, Wn.
where & is the CDF of standard normal distribution. The load is then defined as (p is the node-wise probability in the batch):
(H) CV M.. p)?
(8)
The final loss employs both the task-specific loss and two saad-blaacee losses, leading to the overall optimization target (A is a hand-tuned scaling factor):
L LEAs + X(Lioad (H) Limportance (H)),,
where LEM denotes the task specifcc MoE ponecctatio-mmxximmiin loss.
We further discover that GMoE could be combined with and strengthened by the self-supeervised graph pre-training techniques. We employ GraphMAE [0 as the self-supervised pre-training technique, defined
Pre-training GMoE
L (H, D( H)),
(10)
where f and d denote the encoder and the decoder networks, and M represents the mask for the input graph H. f and collaborative conduct the reconstruction task from the corrupted input, whose quality is measured by the distance metric D. We later will experimentally demonstrate and compare GMoE performance with and without pre-training.
3.2 Computational Complexity Analysis
We show that GMoE-GNN brings negligible overhead on the inference cost compared with its GNN counterpart. We measure computational cost. using the number of floating point operations (FL.Os).
The computation cost GMoE layer can be defined as
CaMoE
Σ KEH
f
Σ G(hi)o Σ
E. (bg W)+
Σ G(bi)o T=
Σ JEN?
E. (hyj W)
JEN,
(11)
where FF maps functions its flops number. CGMOB denotes the computation cost of the whole layer in GMoE- Given there exists an efficient algorithm that can solve hop-1 and hop-2 functions with matching computational complexity [36]. we can further simplify CGMOE
CGMOE =
Σ
Σ
(G(()o))
(12)
TEEijEN,
0=1
[Eo 1(G(h)))) so
if Gf = O, otherwise.
(13)
Given Ed ((G()))) CGMOE can be further simplified to
CoMoE = k ΣΣO C.
(14)
5

---

Table 1: Performance Summary for Graph Classification Tasks. The table header lists the dataset names and corresponding evaluation metrics. Mean and standard deviation values from ten random runs presented. The most outstanding results are highlighted in bold. The improvements achieved by GMoE are indicated in parentheses. ROC-AUC scores are reported in percentage.
ogbg-molbbbp
ogbg-molhiy ogbg-moltoxcast ROC-AUC (t)
ogbg-moltox21
Model
GCN
68.87 + 1.51
76.06 0.97
63.54 + 0.42
75.29 + 0.69
70.28 + 1.36
77.87
1.03
64.12 + 0.61
75.45 + 0.58
GMoB-GCN
(+1.41)
(+1.81)
(+0.58)
(+0.16)
Table 2: Summary of Graph Regression Task Results. The table header displays the dataset names and the corresponding evaluation metrics. Mean and standard deviation values from ten random runs are presented. The most outstanding results are highlighted in bold. GMoE's performance improvements are indicated in parentheses. All evaluation metrics are reported in percentage.
ogbg-molesol
ogbg-molfreesolv
Model
RMSE (1)
GCN
1.114 + 0.036
2.640 + 0.239
1.087 +: 0.043
2.500 +: 0.193
GMoE-GCN
(-0.027)
(-0.140)
Here, C is the computation cost of single message passing in GMoE-GCN Denote the computation cost of a single GCN message passing as Co. and the total computational cost in the whole layer as CGCN. By setting C' Ca in GMoE-GGNN we have CGMoE = Σi, EH ΣjEN, Co = CGCN-
k
In traditional GMoE-GCN GMoE-GIN, the adjustment of C' be easily realized by controlling the hidden feature dimension size 8. For instance, GMoE-GCN and GMoE-GIN with hidden dimension size of s So can have similar FLOPs with its corresponding GCN and GIN with vk
dimension size of sp. The computation cost of gating functions in GMoE is meanwhile negligible compared to the cost of selected experts, since both W, R and Wn. E Raxs is in much smaller dimension than W R given <
In practice, on our NVIDIA A6000 GPU, the inference times for 10, 000 samples are 30.2 10.6ms for GCN-MoE and 36.3 + 17.2ms for GCN. The small variances in GPU clock times align with their nearly identical theoretical FLOPs.
Experimental Results
In this section, we first describe the detailed settings in Section 4.1 We then show main results on graph learning in Section 4.2 Ablation studies and analysis are provided in Section 4.3
4.1 Experimental Settings
Datasets and Evaluation Metrics We conduct experiments on ten graph datasets in the OGB benchmark 48). including graph-level (i.e., ogbg-bbby ogbg-hiv, ogbg- moltoxcast, moltox21, ogbg-molesol, and ggbg-freesolv), node-level (i.e., ogbn-protein, ogbn-arxiv), and link-level predic- tion (i.e., ogbl-ddi, ogbl-ppa) tasks. Following [88) we use ROC-AUC (i.e., area under the receiver operating characteristic curve) the evaluation metric ogbg-bbbp, ogbg-hiv, agbg-moltox211 , and ogbn-protein; RMSE (i.c., root mean squared error) on agbg-molesol, and ogbg-freesolv; classification accuracy (Acc) ogbn-arxiv; Hits @ 100 score on ogbl-ppa; Hits @ 20 score on ogbl-ddi.
We use the GCN and GIN 13 provided by OGB benchmark [48] as the baseline models. All model settings (e.g., number of layers, hidden feature dimensions, etc.) and training hyper-parameters (c.g., learning rates, training epochs, batch size, etc.)
Model Architectures and Training Details

---

Table 3: Performance Comparison for Graph Classification Tasks using GIN with and without Pre- training. We employ GraphMAE [10] for pre-training. The table lists four datasets along with the evaluation metric. The final column displays the average performance across these datasets. Mean and standard deviation values from ten random runs are provided. The most notable results are highlighted in bold.
oghg-molbbbp
ogbg-molhiv
ogbg-moltoxeast ROC-AUC (1)
oghg-moltox21
Average
Model
GIN
65.5 + 1.8
75.4 + 1.5
63.3 + 1.5
74,3 = 0.5
69.63
66.93 1.72
76.14 1.03
62.86 0.37
74.76 0.66
70.17
GMoE-GIN
(+1.43)
(+0.74)
(-0.44)
(+0.46)
(+0.54)
GIN+Pretrain
69.94 0.92
76. 0.8
62.96 + 0.55
73.85 0.64
70.71
68.62 + 1.02 (-1.32)
76.9 + 0.9
64.48 + 0.50
75.25
0.78
71.31
CMot-GGNN-rrerriin
(+0.8)
(+1.18)
(41.40)
(+0.6)
Table 4: Node Prediction Task Results. The table header displays the dataset names and the corresponding evaluation metrics, both reported in percentage. Mcan and standard deviation values from ten random runs presented. The most notable results are highlighted in bold, with GMoE's performance gains indicated in parentheses
ogbn-protein ROC-AUC(1)
ogbn-arxiv Acc(t)
Model
GCN
73.53 + 0.56
71.74 + 0.29
74.48 + 0.58
71.88 + 0.32
GMoE-GCN
(+0.95)
(+0.14)
identical 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 we select & experts out of total of experts for each node, where m out of experts hop-1 aggregation functions and the rest are hop-2 aggregation functions. All three hyper-parameten n, m. k. together with the loss trade-off weight in Eq. (9). are tuned by grid searching: n E {4,8), m E n/2, n). k {1, 4}. and À E {0.1, The hidden feature dimension would be adjusted with & (Section 3.2) to ensure the same flops for all comparisons. The on validation sets are selected to report results on test sets, following the routine in [48]. All other training
values achieving the best performance
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 449 10], we pre-train the
models on 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 [ 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 the masking ratio.
4.2
Main Results
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 supervised learning framework. In all cases, models are trained from scratch, utilizing only the labeled samples the training set. The classification and regression results are outlined in Table and respectively.
7

---

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.
ogbl-ppa Hits @ 100(1)
ogbl-ddi Hits@20(?
Model
GCN
18.67
1.32
37.07
0.051
19.25 1.67
37.96 + 0.082
GMoE-GCN
(+0.58)
(+0.89)
Table 6: Graph Property Prediction Results employing GMoE-GCN with Varied 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 expressed in percentage.
Dataset
Metrie
Average #Nodes
T1
N
" 8,m
==4
k=
k=
k=
oglg-molfreesol
RMSE (1)
8.7
2,856
0.118
2.500 1 0.193
2.532 +: 0.231 2.873
0.251
ogbg-molhiv
ROC-AUC (1)
25,5
76.67 1 1.25
76.41 1 1.72 77.53 - 1.42
77.87 1. 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 significant progress. Additionally, lifts of 1.40% on ogbg-moltox2 1.18% on ogbg-moltoxcan have been observed. The uniformity of these improvements across diverse tasks and datasets underscores the reliability and effectiveness of GMoE- -CN..
modest, they represent 1.43% on ogbg-molbbbo and
Leveraging large-scale pretraining 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 GMoE-GIN outperforms GIN out of 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%
A 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 4a GMoE-GCN's s superiority over the single-expert GCN. Notably, it enhances performance metries like ROC-AUC by 0.95% and Hits @ 20 by 0.89% on the ogbn-proteis protein and ogbl-ddi datasets.
Table 5] demonstrate
4.3 Ablation Study and Analysis
Larger graphs prefer larger hop sizes We conducted an ablation study on two prominent molecule datasets, namely ogbg-molhiv and gebe-molfreesoly, 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-mo molfreesoly dataset. Remarkably, on ogbg- molfreesolv, 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). 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
Observation 1:

---

Table 7: Results for graph property prediction using GMoE-GCN with different scaling ratios À 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.1
ROC-AUC 72.87
75.45:10.58
75.27:10.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 Mořss (i.e., when & < 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 group of experts to learn collaboratively and generalize unseen domains compositionally. Our finding also echoes prior work, e.g., [51], which empirically shows sparse MoEs strong domain generalizable learners.
Observation GMoE demonstrates performance gains, though converging over more epochs We compared the convergence curves of GMoE-GGO and the single-expert GCN, illustrated in Figure Specifically, we plotted the validation and test ROC-AUC different epochs while training on the protein dataset. As observed, the performance of GMoE-GCN reaches plateau later than that of GCN. However, GMoE-GCN eventually achieves 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 per- formance compared to GCN during this early phase. However, after GMoE has undergone enough epochs training, all experts have re-
70
ROC AUC (%)
60
55
ceived ample updates, enabling them to effec- tively model their respective subgroups data. As result, the performance of GMoE surpasses that the single-expert GCN.
45
40
GMoE GCN
GCN
0
250
500
750
1000 Epuchs
1250
1500
1750 2000
Observation 4: Load balancing is essential We conducted an investigation into the sealing factor for the load balancing loss, denoted by X. As shown in Table 7 utilizing X = 0.1 resulted in significantly improved performance, exceed- ing 2.58% in terms ROC-AUC compared to when X 0. This underscores the crucial role of implementing loal-balancing losses. Conversely, the choice of À exhibits less sensitivity; opting for X = 1 would yield similar performance of 75.27%.
Figure 2: Convergence curve of GMoE-GCU and GCN when trained 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 5: GMMoE 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 of- 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. is noteworthy that this methodology has gained significant acclaim, as serves as the
foundation for the top four current methods [53 [4 55 56).

---

After integrating the GMoE framework into the Neural FingerPrints architecture, we achieved an accuracy of 82.72% = 0.53%, while maintaining the same computational resource requirements. This performance outperforms Neural FingerPrints by margin of 0.4%, underlining the broad and consistent adaptability of the GMoE framework.
5
Conclusion
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 neighbor- hood 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-efficiens 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.
trade-off
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 GCN/GIN to elucidate evaluation, we are actively broadening our comparative scope to include state-of-the-an models derived from GCN and other MoE models. Future iterations will thus incorporate a more exhaustive evaluation. Moreover, given the prevalence GCN in semi-supervised learning, there is potential in exploring the benefits of our proposed approach in such tasks avenue we eager 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.
concept. Recognizing the crucial importance of comprchensive
Acknowledgement
The work is sponsored by a Cisco Research Grant (UTAUS-FA0000000)
References
[1] Shiwen Wu, Fei Sun, Wentao Zhang, Xu Xie, and Bin Cui. Graph neural networks in recom- mender 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élaine 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 arX1v:2202.07119, 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 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.
10

---

[9] Yuning You, Tianlong Chen, Zhangyang Wang, and Yang Shen. Bringing your own view: Graph contrastive learning without prefabricated data augmentations. In Proceedings of the Fifteenth ACM International Conference on Web Search and Data Mining, pages 1300-1309, 2022.
[10] Zhenyu Hou, Xiao Liu, Yuxiao Dong, Chunjie Wang, Jie Tang, et al. GraphMAE: Self- supervised masked graph autoencoders. arXiv preprint arXiv:2205 2022.
[11] Thomas N Kipf and Max Welling. Supi-supe classification with graph convolutional networks. arXiv preprint arXiv:1609.029000 2016.
[12] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv: 10903, 2017.
[13] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint trXiv:/810.00055 2018.
[14] 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, 2017.
[15] Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in neural information processing systems, 2017.
[16) Hongbin Pei, Bingzhe Wei, Kevin Chen-Chuan Chang, Yu Lei, and Bo Yang. Geom-GCN Geometric graph convolutional networks. arXiv preprint arXiv:2002.05287 2020.
[17] 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? In NeurIPS, 2021.
[18] Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79 87, 1991.
[19] Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-214, 1994.
[20] Ke Chen, Lei Xu, and Huisheng Chi. Improved learning algorithms for mixture of experts in multiclass classification. Neural networks, 12(9): 100011222-222210 1252, 1999.
[21] Seniha Esen Yuksel, Joseph N Wilson, and Paul D Gader. Twenty years of mixture of experts. IEEE Transactions on Neural Networks and Learning Systems, 237 2012.
[22] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 1701.06588, 2017.
[23] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. arXiv preprint arXiv:2006.16663 16668, 2020.
[24] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.
[25] Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. ICML, pages 4057-4086, 2022.
[26] Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. In Advances in Neural Information Processing Systems, pages 17555-17566, 2021.
[27] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
11

---

[28] Mike Lewis, Shruti Bhosale, Tim Detimers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning. pages 6265-6274 2021.
[29] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. Designing effective sparse expert models. arXiv preprint arXiv:2202.08900 2022.
[30] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In Advances in Neural Information Processing Systems, volume 34, pages 8583-8595, 2021.
[31] Xin Wang, Fisher Yu, Lisa Dunlap, Yi-An Ma, Ruth Wang, Azalia Mirhoseini, Trevor Darrell, and Joseph E Gonzalez. Deep mixture of experts via shallow embedding. In Uncertainty in Artificial Intelligence, pages 552-562, 2020.
[32] Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multi- modal contrastive learning with limoe: the language-image mixture of experts. arXiv preprint MrXiv:2206.02770, 2022.
[33] Jiaqi Ma. Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H Chi. Modeling task relationships in multi-task learning with multi-gate mixture-of-experts In ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 1930-1939, 2018.
[34] Jinguo Zhu, Xizhou Zhu, Wenhai Wang, Xiaohua Wang, Hongsheng Li, Xiaogang Wang, and Jifeng Dai. Mo:: Learning sparse generalist models with conditional moes. arXiv preprint arXivv2206.04677 2022.
[35] Hanxue Liang, Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, Yu Cheng, Cong Hao, Zhangyang Wang, et al. vit Mixture-of-experts vision transformer for efficient multi-task learning with model-accelerator co-design. In Advances in Neural Information Processing Systems, 2022.
[36] Sami Abu-El-Haija, Bryan Perozzi, Amol Kapoor, Nazanin Alipourfard, Kristina Lerman, Hrayr Harutyunyan, Greg Ver Steeg, and Aram Galstyan. MixHop: Higher-order graph convolutional architectures via sparsified neighborhood mixing. In ICML, pages 21-29, 2019.
[37] Sami Abu-El-Haija, Amol Kapoor, Bryan Perozzi, and Joonseok Lee. N-GCN: Multi-scale graph convolution for semi-supervised node classification. In Uncertainty in Artificial Intelligence, pages 841 851, 2020.
[38] Fenyu Hu, W Liping, L Qiang, Shu Wu, Liang Wang, and Tieniu Tan. Graphdive: graph classifcation by mixture of diverse experts. In ICCAI, 2022.
[39] Liguang Zhou, Yuhongze Zhou, Tin Lun Lam, and Yangsheng Xu. CAME: Context-aware mixture-of-experts for unbiased scene graph generation. arXiv preprint arXiv:2208.07109 2022.
[40] Suyeon Kim, Dongha Lee, SeongKu Kang, Seonghyeon Lee, and Hwanjo Yu. Learning topology-specifs experts for molecular property prediction. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pages 8291-8299, 2023.
[41] Zheyuan Liu, Chunhui Zhang, Yijun Tian, Erchi Zhang, Chao Huang, Yanfang Ye, and Chuxu Zhang. Fair graph representation learning via diverse mixture of experts. In The Web Conference, 2023.
[42] Cihang Xie, Mingxing Tan, Boqing Gong, Jiang Wang, Alan L Yuille, and Quoc V Le. Ad- versarial examples improve image recognition. In IEEE Conference on Computer Vision and Pattern Recognition, pages 819-828, 2020.
[43] Haotao Wang, Chaowei Xiao, Jean Kossaifi, Zhiding Yu, Anima Anandkumar, and Zhangyang Wang. AugMax Adversarial composition of random augmentations for robust training. In Advances in Neural Information Processing Systems, volume 34, pages 237-250, 2021.
12

---

[44] Haotao Wang, Aston Zhang, Shuai Zheng, Xingjian Shi, Mu Li, and Zhangyang Wang. Remov ing batch normalization boosts adversarial training. In International Conference 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 Xxiv::0006007739 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 DeXivv1905. 12265, 2019.
[50] Teague Sterling and John J Irwin. Zinc 15-ligand discovery for everyone. Journal of Chemical Information and Modeling, 55(11 25(11):2324-333, 2015.
[51] Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse extuure-of-epperss are domain generalizable learners. arXiv preprint 2rXiv: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. 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.
13


Reviews:
[{"review_id": "JDGA3MtWj9", "reviewer": "Reviewer_k4PS", "summary": "This paper tracks the problem of training a good GNN. To enhance the generalization ability of GNNs, it has become common to increase the diversity of training graph structures by incorporating graph augmentations and/or performing large-scale pre-training on more graphs. Therefore, it is essential for a GNN to simultaneously model diverse graph structures. However, naively increasing the GNN model capacity will suffer from both higher inference costs and the notorious trainability issue of GNNs. To address this issue, this paper proposes Graph mixture of experts, borrowing the idea from recent development of mixture of experts. Based on this new method, each node can select their own optimal information aggregation experts. Experiments are solid and verify the effectiveness of the proposed method.", "questions": "1. This paper can be improved further in terms of writing. You introduce two interesting modules but without good motivations. You can verify your motivation via conducting experiments to show that considering different location for each node will make things different. Thus, you can motivate your idea: we need to choose optimal local information (hop 1 or hop 2). \n\n2. The paper structure should be improved. I think the preliminaries should be considered as a separate section, instead of subsection. It can make readers obtain essential knowledge before the method.\n\n\n3. For formula, there should always be “,”, “.” or “;” in the end of it. Both of formula and text consist of your paper, so formula is just text to be displayed. Besides, text inside the formula should be with /textnormal{}.\n\n4. There is one plus point if the authors can address it well. Because you select experts, how could you ensure that the algorithm converges? The selection results will always converge to a certain selection result? It is very interesting to discuss this point. What is the main drawback to obtain a corresponding convergence result?\n\n5. Some in-text citation is not completed. For example, in line 197, reference 10.\n", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. This paper focuses on an important problem: learning under graphs, which is a common scenario in the real world. \n\n2. This paper designs a new and interesting framework to explore how to use mixture of experts on graphs. The way to mixture graphs is interesting and make a lot of sense. For example, for different node, its optimal local information should be different. This paper considers a novel way to address this problem.\n\n3. The experiments are conducted on many datasets, and the proposed method performs well. It is very hard to explore a useful way to mixture graphs. This paper makes the first try, and the proposed method is novel and verified by solid experiments, which is appreciated.\n", "weaknesses": "1. This paper can be improved further in terms of writing. You introduce two interesting modules but without good motivations. You can verify your motivation via conducting experiments to show that considering different location for each node will make things different. Thus, you can motivate your idea: we need to choose optimal local information (hop 1 or hop 2). \n\n2. The paper structure should be improved. I think the preliminaries should be considered as a separate section, instead of subsection. It can make readers obtain essential knowledge before the method.\n\n\n3. For formula, there should always be “,”, “.” or “;” in the end of it. Both of formula and text consist of your paper, so formula is just text to be displayed. Besides, text inside the formula should be with /textnormal{}.\n\n4. There is one plus point if the authors can address it well. Because you select experts, how could you ensure that the algorithm converges? The selection results will always converge to a certain selection result? It is very interesting to discuss this point. What is the main drawback to obtain a corresponding convergence result?\n\n5. Some in-text citation is not completed. For example, in line 197, reference 10.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive comments and insightful questions. We addressed each point in your comment carefully. \n\nQuestion: This paper can be improved further in terms of writing. You introduce two interesting modules but without good motivations. You can verify your motivation via conducting experiments to show that considering different location for each node will make things different. Thus, you can motivate your idea: we need to choose optimal local information (hop 1 or hop 2).\n\n__Answer__: Thank you for your great suggestion. The results to verify our motivations are listed below. \nOn ogbg-molhiv dataset, if we use all hop-1 experts (all 4 experts are hop-1), the resulting ROC-AUC (↑) is 76.41 ± 1.72. \nIn contrast, if we allow selection between hop-1 and hop-2 experts (2 hop-1 experts and 2 hop-2 experts), the resulting ROC-AUC (↑) is 77.72  ± 0.75.\nThis shows allowing selection from different types of experts achieves performance gains over the ensemble of one single type of experts.\nWe will put these results in the introduction section to verify our motivation in the final version. \n\nQuestion: The paper structure should be improved. I think the preliminaries should be considered as a separate section, instead of subsection. It can make readers obtain essential knowledge before the method.\n\n__Answer__: Thanks for your kind suggestion on the paper’s structure. We agree that organizing preliminaries as a separate section to provide essential knowledge is a good idea. We will re-structure the paper accordingly to enhance its readability.\n\nQuestion: For formula, there should always be “,”, “.” or “;” in the end of it. Both of formula and text consist of your paper, so formula is just text to be displayed. Besides, text inside the formula should be with /textnormal{}.\n\n__Answer__: Thank you for pointing out the formatting issues with the formulas and text. We will fix the formulas as per your suggestion in our final version.\n\nQuestion: There is one plus point if the authors can address it well. Because you select experts, how could you ensure that the algorithm converges? The selection results will always converge to a certain selection result? It is very interesting to discuss this point. What is the main drawback to obtain a corresponding convergence result?\n\n__Answer__: We appreciate your feedback! In our proposed Graph Mixture of Experts (GMoE), the inherent weight balancing loss acts as a built-in regularizer, ensuring convergence of the selection process. This aspect indeed provides stability to the selection results. However, we understand the significance of discussing this aspect in detail. We will make sure to elaborate on it in our final version. The main drawback to ensuring convergence of the selection process lies in striking a balance between the information aggregation experts. While the inherent weight balancing loss aids convergence to a certain extent, achieving convergence without compromising the diverse expertise captured by different experts remains an area of ongoing exploration.\n\nQuestion: Some in-text citation is not completed. For example, in line 197, reference 10.\n\n__Answer__: Finally, we apologize for any incomplete in-text citations. We will review and correct all references in our final version.\n"}, {"review_id": "8lMN9Y4pTE", "reviewer": "Reviewer_7Ua4", "summary": "This paper proposes a novel Graph Mixture of Experts (GMoE) architecture that enhances the ability of Graph Neural Networks (GNNs) to leverage and benefit from the structural diversity present in training graphs. The GMoE model comprises multiple \"experts\" at each layer, with each expert being an independent aggregation function with its own trainable parameters. During training, the model learns to select the most appropriate aggregation experts for each node, resulting in nodes with similar neighborhood information being routed to the same aggregation experts.", "questions": "The paper lacks a detailed analysis of the computational cost of the proposed GMoE model compared to other GNN architectures. How does the computational cost of the model compare to other state-of-the-art models, and how practical is it for real-world applications?\n\nThe paper does not provide a detailed explanation of the selection mechanism that the proposed GMoE model uses to choose the most appropriate aggregation experts for each node. How does the model make these decisions, and how can the selection mechanism be improved?\n\nThe paper could be improved by including more details on the process used to tune the hyperparameters of the compared baselines. What process was used to tune these hyperparameters, and how can a fair comparison be made between the proposed GMoE model and other state-of-the-art models?", "limitations": "The paper lacks a detailed analysis of the computational cost of the proposed GMoE model compared to other GNN architectures and does not provide a thorough explanation of the selection mechanism used by the model to choose the most appropriate aggregation experts for each node. To improve the paper, the authors could include more details on the process used to tune the hyperparameters of the compared baselines, enabling a fair comparison with other state-of-the-art models. Additionally, providing a more in-depth explanation of the selection mechanism used by the GMoE model would facilitate a better understanding of the model's decision-making process and enable its improvement.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed Graph Mixture of Experts (GMoE) architecture is a new approach to enhancing the ability of Graph Neural Networks (GNNs) to leverage and benefit from the structural diversity present in training graphs.\n\nThe proposed GMoE model is evaluated on several benchmark datasets and achieves state-of-the-art performance on node classification and graph classification tasks. This demonstrates the effectiveness of the proposed architecture in improving the performance of GNNs on diverse graph structures.\n\nThe proposed GMoE model is flexible and can be easily integrated with existing GNN architectures. This makes it easy to incorporate the proposed model into existing GNN-based applications and frameworks.", "weaknesses": "While the proposed GMoE model is designed to avoid additional computational overhead, the paper does not provide a detailed analysis of the computational cost of the model compared to other GNN architectures. This makes it difficult to assess the practicality of the proposed model in real-world applications.\n\nWhile the paper describes how the model learns to select the most appropriate aggregation experts for each node, it does not provide a detailed explanation of the selection mechanism. This makes it difficult to understand how the model makes decisions and how it can be improved.\n\nOne suggestion for improving this paper would be to include more details of parameter tuning for the compared baselines. This would help to ensure a fair comparison between the proposed GMoE model and other state-of-the-art models. The paper could also describe the process used to tune these hyperparameters, such as grid search or random search.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your recognition of the strengths in our paper and your valuable insights shared in your comments. In response, we have carefully addressed each of your points. \n\nQuestion. While the proposed GMoE model is designed to avoid additional computational overhead, the paper does not provide a detailed analysis of the computational cost of the model compared to other GNN architectures. This makes it difficult to assess the practicality of the proposed model in real-world applications.\n\n__Answer__\nWe provided complexity analysis in the section 3.3. Specifically, we measure computational cost using the number of floating-point operations (FLOPs) and show that GMoE-GNN brings negligible overhead on computational cost compared with its traditional GNN counterpart.\n\n\nQuestion. The paper does not provide a detailed explanation of the selection mechanism that the proposed GMoE model uses to choose the most appropriate aggregation experts for each node. How does the model make these decisions, and how can the selection mechanism be improved?\n\n__Answer__\nWe provided the expert selection mechanism in the section 3.2. Specifically, each node in the graph selects its own optimal information aggregation experts based on the similarity between the node's local graph structure and the structures modeled by the experts, where the similarity is measured using a gating mechanism that computes the dot product between the node's feature vector and the expert's feature vector. We also describe how the gating mechanism can be improved by incorporating additional information about the node's local graph structure, such as the degree of the node and the degree of its neighbors.\n\n\nQuestion. The paper could be improved by including more details on the process used to tune the hyperparameters of the compared baselines. What process was used to tune these hyperparameters, and how can a fair comparison be made between the proposed GMoE model and other state-of-the-art models?\n\n__Answer__\nThank you for your kind suggestion! We provided the expert selection mechanism in the section 4.1. Specifically, we describe the hyperparameter tuning process and range for each baseline model and report the best hyperparameters found during the tuning process. Furthermore, we compare the performance of the proposed GMoE model with other state-of-the-art models on several benchmark datasets, using the same hyperparameters for each model. This ensures a fair comparison between the models and allows for a meaningful evaluation of the effectiveness of the proposed GMoE model. We will add the above description to the final version. \n"}, {"review_id": "x1y6ANf3Cv", "reviewer": "Reviewer_BUC6", "summary": "The paper claims that increasing the diversity of graph structures plays an important role in GNN training, which makes GNNs more robust and generalizable. To make GNN models leverage the structural diversity better, the paper proposes the Graph Mixture of Experts (GMoE). In each layer of GMoE, the representation of each node will be computed based on multiple experts with independent aggregation functions. Each expert will utilize either the hop-1 or hop-2 aggregation function to increase the diversity further. The paper conducts extensive experiments on OGB benchmark datasets and improves the performance compared to the single-expert baselines. \n", "questions": "1. Is there any evidence that GMoE will find the optimal experts?\n2. From Equation (3), it seems that the selection of experts only relies on the target node i and do not rely on the node j. Is it reasonable?\n3. The authors might give more discussions of the difference or connections between GMoE and GAT with 2-hop neighbors.", "limitations": "There is no negative societal impact of the work.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper claims the importance of graph data diversity in GNN training. \n2. The paper proposes a new architecture called GMoE, which combines MoE in neighborhood aggregation. \n3. The paper also tests whether the GMoE can help the graph self-supervised learning. \n", "weaknesses": "1. In the abstract, the paper argues that GMoE enables each node to select its own optimal information aggregation experts. But we do not see evidence why GMoE will find optimal experts. \n2. From my view, the formulation of the proposed GMoE is similar to GAT with 2-hop neighbors, limiting the work's novelty. \n3. Although GMoE improves the results compared to single-expert GCN/GIN models, the results are still far from the state-of-the-art results on the OGB leaderboard (e.g., top-1 result of ogbg-molhiv is 84.20). \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and constructive comments. We carefully addressed your comments below.\n\nQuestion. In the abstract, the paper argues that GMoE enables each node to select its own optimal information aggregation experts. But we do not see evidence why GMoE will find optimal experts.\n\n\n__Answer__\nThe use of the word “optimal” in the abstract was a typographical error. We would revise lines 11-13 to accurately reflect our intended statement in the future version: “Our new Graph Mixture of Expert (GMoE) model enables each node in the graph to dynamically and adaptively select more generalizable information aggregation experts.”\n\n\nQuestion. From my view, the formulation of the proposed GMoE is similar to GAT with 2-hop neighbors, limiting the work's novelty.\n\n__Answer__\nOur approach is distinct from the Graph Attention Network (GAT). As detailed in the abstract, our network is designed to enhance the capacity of the Graph Neural Network (GNN) model while preserving similar computational overheads. In contrast, an increase in the capacity for GAT would result in a linear increasing in computational complexity.\n\n\nQuestion. Although GMoE improves the results compared to single-expert GCN/GIN models, the results are still far from the state-of-the-art results on the OGB leaderboard (e.g., top-1 result of ogbg-molhiv is 84.20).\n\n__Answer__\nWe follow the common practice [1][2] to study our technique on standard backbone like GCN/GIN for proof-of-concept exploration. We are actively working on extending our technique to more networks. Due to the limited time frame allocated for rebuttal, we will be updating the results in a subsequent version of this work.\n \n[1] You, Yuning, et al. \"Graph contrastive learning with augmentations.\" Advances in neural information processing systems 33 (2020): 5812-5823.\n[2] Hu, Weihua, et al. \"Open graph benchmark: Datasets for machine learning on graphs.\" Advances in neural information processing systems 33 (2020): 22118-22133.\n\nQuestion. From Equation (3), it seems that the selection of experts only relies on the target node i and do not rely on the node j. Is it reasonable?\n\n__Answer__\nOur experiments have confirmed that we can select suitable experts, thereby enhancing efficiency based on the feature of node i. Additionally, integrating the information from node j is an intriguing area for further exploration, which we would explore in future work.\n"}, {"review_id": "4aWGaZmQBq", "reviewer": "Reviewer_Deyj", "summary": "This paper introduces the Graph Mixture of Expert (GMoE) model, enhancing Graph Neural Networks' (GNNs) ability to handle diverse graph structures without additional computational costs. GMoE allows nodes to select their information aggregation experts dynamically, catering to different subgroups of graph structures. Additionally, it incorporates experts with varying aggregation hop sizes for broader information capture. Its effectiveness is validated through various tasks in the OGB benchmark, with improved results compared to non-MoE baselines. ", "questions": "please refer to weaknesses.", "limitations": "please refer to weaknesses.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "This article uses MoE in GNN, which is a relatively new idea.", "weaknesses": "1. After reading the entire article, the motivation behind it remains unclear. What issue pertaining to GNN does the article aim to address? It would be beneficial if the author could include an \"preliminary\" section that clearly delineates the problem this article seeks to solve.\n\n2. Certain assertions in the article are questionable. For instance, the statement: “forcing all nodes to share the same aggregation mechanism” does not align with the functioning of GAT and GraphTransformer, where each node's aggregation is adaptively carried out via the attention mechanism.\n\n3. The article doesn't clarify how the model's improved generalization ability is assured. How can the model learn a more diverse graph structure based on observation data from a single domain? This is an important point that requires further clarification.\n\n4. The algorithm's design seems more akin to an engineering solution or a skill rather than a scientifically robust methodology. The article lacks theoretical discussion explaining why the algorithm is designed in its particular way. This absence of theoretical framework underpinning the algorithm design makes it difficult to fully appreciate or understand the design choices.\n\n5. There are also many articles [1-3] that use MoE to GNN, and the author does not add these works to the discussion or compare them.\n\n6. According to the guidelines provided by the NIPS official website, authors are required to include sections on 'Limitations' and 'Broader Impacts'. However, these essential components appear to be missing in this article.\n\n[1] Explore Mixture of Experts in Graph Neural Networks\n\n[2] Graph Classification by Mixture of Diverse Experts, IJCAI 2022\n\n[3] Fair Graph Representation Learning via Diverse Mixture-of-Experts, WWW 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and insightful comments. We carefully addressed every point raised in your feedback. We hope this can address any concerns you may have had and foster a more positive perspective on our paper.\n\nQuestion: What is the motivation of the paper? What issue does the article aim to address?\n\n__Answer__: We agree it is important to emphasize the motivations. Our motivations are highlighted in the introduction section. The problem we try to solve is: How to “create a GNN model architecture that can more effectively leverage and benefit from diverse training samples without any additional computational overhead?” (lines 35-37). This is an important research question, since the graph structures in real-world training data are typically diverse and heterogeneous. For example, in graph-based recommendation systems, a node can represent a product or a customer, while an edge can indicate different interactions such as viewed, liked, or purchased. Similarly, in biochemistry tasks, datasets often comprise molecules with various biochemistry properties and, thus, various graph structures (lines 25-29). Moreover, purposefully increasing the diversity of graph data structures in training sets has become a crucial aspect of GNN training. Examples include random data augmentation (e.g., edge and node dropping) and large-scale pre-training (where the pretraining datasets usually have distributional gaps from the downstream datasets) (lines 30-34).\n\nQuestion: The statement: “forcing all nodes to share the same aggregation mechanism” does not align with the functioning of GAT and GraphTransformer\n\n__Answer__: You are correct that GAT and GraphTransformer adaptively learn the aggregation function for each node.  However, we want to humbly point out that we claimed “many” (line 38), instead of all, GNN architectures used the same aggregation mechanism for all nodes. We apologize that our wording may cause such confusion, and will add a footnote at this place in the final version to better clarify that not all GNNs are designed in such a way. We would like to thank you again for your careful review and pointing this out. \n\nQuestion: The article doesn't clarify how the model's improved generalization ability is assured. How can the model learn a more diverse graph structure based on observation data from a single domain?\n\n__Answer__: According to both sampling complexity theories [1] and empirical observations in deep learning, the model’s generalization ability depends mainly on two factors: the diversity of the training samples (ie., the “estimation error” in Section 5.2 in [1]. Intuitively this is how well the sampled training data represents the underlying data distribution), and the expressiveness of the hypothesis space (ie., the “Approximation Error” in Section 5.2 in [1]. Intuitively this is how much capacity the deep model has in order to fit complex data distributions). \n\nMany techniques, such as random data augmentation and large-scale pretraining, have been used to increase the generalization ability of GNNs, since they increase the diversity of the training samples. _Equipped with these techniques boosting training data diversity, we aim to increase the model’s generalization ability from the other end: To design GNN structures with larger capacity to fit such diverse training data._ One straightforward way to do this is to increase the depth or width of the GNNs. However, these solutions will harm the inference efficiency, which is crucial for may GNN applications such as molecule virtual screening (see footnote on page 2 for more references). To this end, we propose our solution GMoE, which increases the model capacity to fit diverse training data without computational overhead during inference. The intuition why GMoE can increase the model’s capability to fit diverse training data is explained in lines 47-54 of our paper and also in the answer to your next question. \n\n[1] Understanding Machine Learning: From Theory to Algorithms. Cambridge press 2014.\n\nQuestion:  The article lacks theoretical discussion explaining why the algorithm is designed in its particular way\n\n__Answer__: Being an empirical endeavor, our work's primary focus does not lie in furnishing theoretical assurances for the design. Nevertheless, we are content to reaffirm the underlying intuition behind our approach here. \nThe goal of GMoE is to utilize different experts to better fit the training set with diverse graph structures.  Each expert in our GMoE architecture is trained to focus on its own subgroup of training samples that share similar neighborhood information patterns (e.g., shorter or longer range, less or more to aggregate). Intuitively speaking, when the training data distribution is complex (e.g., when the training graphs are diverse in structures), a single expert (i.e., traditional GCN or GIN) has limited capability to learn all those diverse data within a single expert. However, for GMoE, there are multiple experts, where each expert  can focus on its own expertise. In this way, GMoE can better fit diverse training data, and thus achieve better generalization performance.\n\nQuestion: Difference between other GNN papers with MoE\n\n__Answer__: Due to space limitations we only briefly summarize the difference here. We will add more details in final version.\n\n[1] uses linear mixing while we do non-linear mixing. [1] do not use any weight-balancing regularizations as done in ours, which makes their training prone to collapse [14]. Also [1] only uses one type of expert while our framework uses experts with different hops to further increase the capability to fit diverse data. \n[2] is cited in Sec. 2.2 as [41].\n[3] uses MoE to achieve fair predictions for GNN, which differs from our work in both objectives and technical approaches.\n\nQuestion: Missing Limitations and Broader Impacts\n\n__Answer__: Thank you for pointing out. Due to space limit, we will add both sections in the final version. \n"}, {"review_id": "p9RLxN7nui", "reviewer": "Reviewer_NqGe", "summary": "This paper considers the problem of graph neural networks. After discussing the drawback of existing methods, this paper presents a new model, called  Graph Mixture of Expert (GMoE), which can dynamically select its own optimal information aggregation experts. Extensive experiments conducted on several graph datasets show the benefit of the proposed method.", "questions": "Please see the weaknesses.", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The authors have carefully reviewed the drawback of existing graph models and then propose a novel and effective model, called Graph Mixture of Expert (GMoE). The GMoE enables the model to dynamically select its own optimal information aggregation experts.\n\n+ The paper is well-written and easy to follow. The reviewer can easily understand the motivation of this paper and the details of the proposed method.\n\n+ Experiments have shown the benefit of the proposed method over the traditional GCN.", "weaknesses": "- Although the authors have shown the benefit of the proposed method over GCN, the improvement seems somewhat limited. For example, in several tables, the improvement over GCN is only 0.*%.\n\n- It seems that the comparison to other methods is limited. The authors have only compared to the traditional GCN. However, the comparison to other state-of-the-art models modified based on GCN is missed. In addition, it is very important to compare with other MoE models.\n\n- As GCN is also popular in semi-supervised learning. I wonder if the proposed can also benefit the semi-supervised learning task.\n\n- The authors stated that the proposed method is efficient. However, the computation cost about training time and testing time is not provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive comments and insightful questions. We addressed each point in your comment carefully. \n\nQuestion: Although the authors have shown the benefit of the proposed method over GCN, the improvement seems somewhat limited. For example, in several tables, the improvement over GCN is only 0.*%.\n\n__Answer__: Thank you for your observation. While the numerical improvements you mentioned may seem modest in terms of percentage points, it's important to note that these values represent meaningful enhancements in the context of our study. Specifically, our technique demonstrated a lift of 1.43% on ogbg-molbbbq, 1.18% on ogbg-moltoxcast, and 1.40% on ogbg-moltox21. These improvements are consistent across a range of tasks and datasets, reinforcing the reliability and effectiveness of our proposed method.\n\n\nQuestion: It seems that the comparison to other methods is limited. The authors have only compared to the traditional GCN. However, the comparison to other state-of-the-art models modified based on GCN is missed. In addition, it is very important to compare with other MoE models.\n\n__Answer__: We appreciate your observation regarding the comparison scope of our study. As you rightly pointed out, our focus in this initial work has been on standard backbones like GCN/GIN for the purpose of establishing a proof-of-concept. We are indeed actively engaged in the process of extending our technique to a wider range of network architectures. We understand the importance of a comprehensive evaluation, including comparisons with state-of-the-art models derived from GCN and other MoE models. While the time frame for the rebuttal is limited, we are committed to addressing this gap in our subsequent work and updating the results accordingly.\n\nQuestion: As GCN is also popular in semi-supervised learning. I wonder if the proposed can also benefit the semi-supervised learning task.\n\n__Answer__: Thanks for your valuable feedback! Certainly, our proposed approach has the potential to benefit semi-supervised learning tasks as well. We consider this an important avenue for exploration and plan to include it in our future research.\n\nQuestion: The authors stated that the proposed method is efficient. However, the computation cost about training time and testing time is not provided.\n\n__Answer__: We only claimed inference efficiency instead of training efficiency. The definition of efficiency here is: GCN-MoE can more effectively leverage and benefit from diverse training samples than GCN, without increased computational overhead during inference. We do not claim GCN-MoE to be more efficient than GCN. We apologize if this wording causes any confusion. \nSpecifically, in section 3.3, we demonstrate through mathematical analysis of the inference FLOPs that GMoE-GNN introduces minimal additional computational cost during inference (measured in FLOPs) in comparison to its conventional GNN counterpart. The real on-device inference time is dependent on hardware specifications and may vary dramatically on different devices. On our NVIDIA A6000 GPU, the total inference time for 10000 samples are 30.2 ± 10.6ms  (GCN-MoE) v.s. 36.3 ± 17.2ms (GCN). Note that these clock times measured on GPUs can be noisy and we do not claim GCN-MoE to be more efficient than GCN. Instead, theoretically they have almost identical FLOPs. We do not pursue training efficiency in this paper. \nWe will add the above numbers and explanations in the final version. \n"}]

Author rebuttal:

Meta-review: This paper tracks the problem of training a good GNN. To enhance the generalization ability of GNNs, it has become common to increase the diversity of training graph structures by incorporating graph augmentations and/or performing large-scale pre-training on more graphs. Therefore, it is essential for a GNN to simultaneously model diverse graph structures. However, naively increasing the GNN model capacity will suffer from both higher inference costs and the notorious trainability issue of GNNs. To address this issue, this paper proposes Graph mixture of experts, borrowing the idea from recent development of mixture of experts. Based on this new method, each node can select their own optimal information aggregation experts. Experiments are solid and verify the effectiveness of the proposed method.

To be specific, this paper focuses on an important problem--learning under graphs, which is a common scenario in the real world. This paper designs a new and interesting framework to explore how to use mixture of experts on graphs. The way to mixture graphs is interesting and make a lot of sense. For example, for different node, its optimal local information should be different. This paper considers a novel way to address this problem. The experiments are conducted on many datasets, and the proposed method performs well. It is very hard to explore a useful way to mixture graphs. This paper makes the first try, and the proposed method is novel and verified by solid experiments, which is appreciated.

Although the reviewers had some concerns and follow-up questions, the authors did a point-to-point rebuttal in details. Meanwhile, the clarity and novelty are above the bar of NeurIPS. Therefore, this paper can be accepted as a poster.
Decision: Accept (poster)