SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haotao Wang $^{1}$ ; Ziyu Jiang $^{2}$ ; Yuning You $^{2}$ , Yan Han $^{1}$ , Gaowen Liu $^{3}$ , Jayanth Srinivasa $^{3}$ , Ramana Rao Kompella $^{3}$ , Zhangyang Wang $^{1}$
<sup>1</sup>University of Texas at Austin Texas A&M University<sup>2</sup> Cisco Systems<sup>3</sup>
{htwang, yh9442, atlaswang}@utexas.edu, {jiangziyu, yuning.you}@amu.edu, {gaoliu, jasriniv, rkompell}@cisco.com
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-molbbbpl, when compared to the non-MoE baselines. Our code is publicly available at https://github.com/VITA-Group/Graph-Mixture-of-Experts.
Graph learning has found extensive use in various real-world applications, including recommendation systems [1], traffic prediction [2], and molecular property prediction [3]. Real-world graph data typically exhibit diverse graph structures and heterogeneous nodes and edges. In graph-based recommendation systems, for instance, a node can represent a product or a 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 therefore various graph structures. Moreover, purposefully increasing the diversity of graph data structures in training sets has become a crucial aspect of GNN training. Techniques such as graph data augmentations [4, 5] and large-scale pre-training on diverse graphs [6, 7, 8, 9, 10] have been widely adopted to allow GNNs for extracting more 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.
question: Can one effectively scale a GNN model'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 to share the same aggregation mechanism, regardless of the differences in their node features or neighborhood information.2 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 message-passing 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 neighborhood information are guided towards the same aggregation experts. This fosters specialization within each GMoE 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 effectiveness through a range of comprehensive molecular property prediction tasks, underscoring our commitment to 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-molbbbpl,  $0.95\%$  on ogbn-proteins, and boosts Hits@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.
Graph Neural Networks Graph neural networks (GNNs) [11, 12, 13] have emerged as a powerful approach for learning graph representations. Variants of GNNs have been proposed [11, 12, 13], achieving state-of-the-art performance in different graph tasks. Under the message passing framework [14], graph convolutional network (GCN) adopts mean pooling to aggregate the neighborhood and updates embeddings recursively [11]; GraphSAGE [15] 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-GCN devises a 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 (MoE) [18] has a long history, tracing its origins back to earlier work s[19, 20, 21]. Recently, spurred by advancements in large language models, sparse MoE [22, 23, 24, 25, 26] has regained 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 [25, 27]. Subsequent studies have further refined the stability and efficiency of sparse MoE [28, 29]. The remarkable success of sparse MoE in the realm of language has spurred its adoption in diverse domains, including vision [30, 31], multi-modal [32], and multi-task learning [33, 34, 35].
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 a fixed-weight
There are exceptions. For example, GAT and GraphTransformer adaptively learn 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.
Figure 1: Each row represents the aggregation of a single node, and each column corresponds to a different network or sub-module. Blue dots  $(\bullet)$  denote the input features passed to the red dots  $(\bullet)$  via the colorful edges. On the left, we demonstrate two hop-1 experts with distinct weights, along with 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.
Mixture of Experts (MoE). Pioneering efforts [38, 39] have also investigated the application of MoE to address the well-known issue of imbalance and to develop unbiased classification or generalization algorithms. However, none of these approaches harnessed the potential advantages of sparsity and adaptivity. Recent work [40] introduced the use of a mixture of experts for molecule property prediction. They employed a GNN as a feature extractor and applied a mixture of experts, where each expert is a linear classifier, on top of the extracted features for graph classification. In contrast, each layer of GMoE constitutes a mixture of experts, with each expert being a GCN/GIN layer featuring different aggregation step sizes. Another distinction from [40] lies in their utilization of domain-specific knowledge (specifically, molecule topology) for expert routing, whereas our approach is designed to operate on general graphs without relying on domain-specific assumptions. One more concurrent study by [41] employs MoE to achieve fairness in predictions for GNNs.
Our study takes a significant stride forward by introducing sparse MoE to 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 to 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 a 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 to 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 not rely on (batch) normalization as heavily.
Preliminaries: Graph Neural Networks Taking the classical Graph Convolutional Network (GCN) [11] as an example, the propagation mechanism can be formulated as
$$
h _ {i} ^ {\prime} = \sigma \left(\sum_ {j \in N _ {i}} \frac {1}{\sqrt {| N _ {i} | | N _ {j} |}} h _ {j} W ^ {(i)}\right), \tag {1}
$$
where  $W^{(i)} \in \mathbb{R}^{s \times s}$  is a trainable weight and  $\sigma$  is an element-wise non-linear activation function.  $h_j \in \mathbb{R}^{b \times s}$  denotes the input feature of  $j$  th node while  $h_i' \in \mathbb{R}^{b \times s}$  is its output feature in  $i$  th node.  $b$  and  $s$  are batch size and hidden feature size, respectively.  $N_i$  denotes the collection of neighbors for  $i$  th node including self-connection. The output feature is normalized by  $\frac{1}{\sqrt{|N_i||N_j|}}$ . 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 1. The GMoE layer comprises multiple experts, each utilizing either the hop-1 or hop-2 aggregation function. To determine which experts to use for a 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 a 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-1 and hop-2 experts enables the model to dynamically capture short-range or long-range information aggregation for each node. Formally, a GMoE layer can be written as:
$$
h _ {i} ^ {\prime} = \sigma \left(\sum_ {o = 1} ^ {m} \sum_ {j \in N _ {i}} G \left(h _ {i}\right) _ {o} E _ {o} \left(h _ {j}, e _ {i j}, W\right) + \sum_ {o = m} ^ {n} \sum_ {j \in N _ {i} ^ {2}} G \left(h _ {i}\right) _ {o} E _ {o} \left(h _ {j}, e _ {i j}, W\right)\right), \tag {2}
$$
where  $m$  and  $n$  denote the hop-1 and total experts number, respectively. Hence the number of hop-2 experts is  $n - m$ .  $E_{o}$  and  $e_{ij}$  denote the message function and edge feature between  $i$  th and  $j$  th nodes, respectively. It can represent multiple types of message-passing functions such as one employed by GCN [47] or GIN [13].  $G$  is the gating function that generates multiple decision scores with the input of  $h_i$  while  $G(h_i)_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 \left(h _ {i}\right) = \operatorname {S o f t m a x} \left(\operatorname {T o p K} \left(Q \left(h _ {i}\right), k\right)\right), \tag {3}
$$
$$
Q \left(h _ {i}\right) = h _ {i} W _ {g} + \epsilon \cdot \text {S o f t p l u s} \left(h _ {i} W _ {n}\right), \tag {4}
$$
where  $k$  denotes the number of selected experts.  $\epsilon \in \mathcal{N}(0,1)$  denotes standard Gaussian noise.  $W_{g}\in \mathbb{R}^{s\times n}$  and  $W_{n}\in \mathbb{R}^{s\times n}$  are learnable weights that control clean and noisy scores, respectively.
The proposed GMoE layer can be applied to many GNN backbones such as GIN [13] or GCN [11]. 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-GIN (for GCN and GIN, respectively).
Additional Loss Functions to Mitigate GMoE Collapse Nonetheless, if this model is trained solely using the expectation-maximization loss, it may succumb to a trivial solution wherein only a single group of experts is consistently selected. This arises due to the self-reinforcing nature of the imbalance: the chosen experts can proliferate at a much faster rate than others, leading to their increased frequency of selection. To mitigate this, we implement two additional loss functions to prevent such collapse [22]. The first one is importance loss:
$$
\operatorname {I m p o r t a n c e} (H) = \sum_ {h _ {i} \in H, g \in G \left(h _ {i}\right)} g, L _ {\text {i m p o r t a n c e}} (H) = C V (\operatorname {I m p o r t a n c e} (H)) ^ {2}, \tag {5}
$$
where the importance score  $\mathrm{Importance}(H)$  is defined as the sum of each node's gate value  $g$  across the whole batch.  $CV$  represents the coefficient of variation. The importance loss  $L_{\mathrm{importance}}(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 to different experts. For instance, one expert could receive a few high scores, while another might be selected by many more nodes yet t 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 load-balanced loss to encourage a more even selection probability per expert. Specifically,  $G(h_{i}) \neq 0$  if and only if  $Q(h_{i})_{o}$  is greater than the  $k$ -th largest element of  $Q(h_{i})$  excluding itself. Consequently, the probability of  $G(h_{i}) \neq 0$  can be formulated as:
$$
P \left(h _ {i}, o\right) = P r \left(Q \left(h _ {i}\right) _ {o} > \operatorname {k t h} _ {-} \exp \left(Q \left(h _ {i}\right), k, o\right)\right), \tag {6}
$$
where  $k$ th  $ex()$  denotes the  $k$ -th largest element excluding itself.  $P(h_i, o)$  can be simplified as
$$
P \left(h _ {i}, o\right) = \Phi \left(\frac {h _ {i} W _ {g} - \operatorname {k t h} _ {-} \exp (Q \left(h _ {i}\right) , k , o)}{\text {S o f t p l u s} \left(h _ {i} W _ {n}\right)}\right), \tag {7}
$$
where  $\Phi$  is the CDF of standard normal distribution. The load is then defined as ( $p$  is the node-wise probability in the batch):
$$
L _ {\text {l o a d}} (H) = C V \left(\sum_ {h _ {i} \in H, p \in P \left(h _ {i}, o\right)} p\right) ^ {2}. \tag {8}
$$
The final loss employs both the task-specific loss and two load-balance losses, leading to the overall optimization target ( $\lambda$  is a hand-tuned scaling factor):
$$
L = L _ {E M} + \lambda \left(L _ {\text {l o a d}} (H) + L _ {\text {i m p o r t a n c e}} (H)\right), \tag {9}
$$
where  $L_{EM}$  denotes the task-specific MoE expectation-maximizing loss.
Pre-training GMoE We further discover that GMoE could be combined with and strengthened by the self-supervised graph pre-training techniques. We employ GraphMAE [10] as the self-supervised pre-training technique, defined as
$$
\mathcal {L} (H, M) = D \left(d (f (M \cdot H)), H\right), \tag {10}
$$
where  $f$  and  $d$  denote the encoder and the decoder networks, and  $M$  represents the mask for the input graph  $H$ .  $f$  and  $d$  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 (FLOPs).
The computation cost of a GMoE layer can be defined as
$$
C _ {G M o E} = \sum_ {h _ {i} \in H} \mathcal {F} \left[ \sum_ {o = 1} ^ {m} G \left(h _ {i}\right) _ {o} \sum_ {j \in N _ {i}} E _ {o} \left(h _ {j}, e _ {i j}, W\right) + \sum_ {o = m} ^ {n} G \left(h _ {i}\right) _ {o} \sum_ {j \in N _ {i} ^ {2}} E _ {o} \left(h _ {j}, e _ {i j}, W\right) \right], \tag {11}
$$
where  $\mathcal{F}$  maps functions to its flops number.  $C_{GMoE}$  denotes the computation cost of the whole layer in GMoE-GCN. Given there exists an efficient algorithm that can solve hop-1 and hop-2 functions with matching computational complexity [36], we can further simplify  $C_{GMoE}$  as
$$
C _ {G M o E} = \sum_ {h _ {i} \in H} \sum_ {j \in N _ {i}} C \sum_ {o = 1} ^ {n} \mathbf {1} \left(G \left(h _ {i}\right) _ {o}\right), \tag {12}
$$
$$
C = \mathcal {F} \left[ E _ {o} \left(h _ {j}, e _ {i j}, W\right) \right], \quad \mathbf {1} \left(G \left(h _ {i}\right) _ {o}\right) = \left\{ \begin{array}{l l} 0 & \text {i f} G \left(h _ {i}\right) _ {o} = 0, \\ 1 & \text {o t h e r w i s e .} \end{array} \right. \tag {13}
$$
Given  $\sum_{o = 1}^{n}\mathbf{1}(G(h_i)_o) = k$ ,  $C_{GMOE}$  can be further simplified to
$$
C _ {G M o E} = k \sum_ {h _ {i} \in H} \sum_ {j \in N _ {i}} C. \tag {14}
$$
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 are 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.
Model	ogbg-molbbbpl	ogbg-molhiv ROC-AUC ↑	ogbg-moltoxcast ↑	ogbg-moltox21
GCN	68.87 ± 1.51	76.06 ± 0.97	63.54 ± 0.42	75.29 ± 0.69
GMoE-GCN	70.28 ± 1.36 (+1.41)	77.87 ± 1.03 (+1.81)	64.12 ± 0.61 (+0.58)	75.45 ± 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.
Model	ogbg-molesol	ogbg-molfreesolv
RMSE (↓)
GCN	1.114 ± 0.036	2.640 ± 0.239
GMoE-GCN	1.087 ± 0.043	2.500 ± 0.193
(-0.027)	(-0.140)
Here,  $C$  is the computation cost of a single message passing in GMoE-GCN. Denote the computation cost of a single GCN message passing as  $C_0$ , and the total computational cost in the whole layer as  $C_{GCN}$ . By setting  $C = \frac{C_0}{k}$  in GMoE-GCN, we have  $C_{GMoE} = \sum_{h_i \in H} \sum_{j \in N_i} C_0 = C_{GCN}$ .
In traditional GMoE-GCN or GMoE-GIN, the adjustment of  $C$  can be easily realized by controlling the hidden feature dimension size  $s$ . For instance, GMoE-GCN and GMoE-GIN with hidden dimension size of  $s = \frac{s_0}{\sqrt{k}}$  can have similar FLOPs with its corresponding GCN and GIN with dimension size of  $s_0$ . The computation cost of gating functions in GMoE is meanwhile negligible compared to the cost of selected experts, since both  $W_{g} \in \mathbb{R}^{n \times s}$  and  $W_{n} \in \mathbb{R}^{n \times s}$  is in a much smaller dimension than  $W^{(i)} \in \mathbb{R}^{s \times s}$  given  $n \ll s$ .
In practice, on our NVIDIA A6000 GPU, the inference times for 10,000 samples are  $30.2 \pm 10.6ms$  for GCN-MoE and  $36.3 \pm 17.2ms$  for GCN. The small variances in GPU clock times align with their nearly identical theoretical FLOPs.
In this section, we first describe the detailed settings in Section 4.1. We then show our main results on graph learning in Section 4.2. Ablation studies and analysis are provided in Section 4.3.
Datasets and Evaluation Metrics We conduct experiments on ten graph datasets in the OGB benchmark [48], including graph-level (i.e., ogbg-bbbb, ogbg-hiv, ogbg-moltoxcast, ogbg-moltox21, ogbg-molesol, and ogbg-freesolv), node-level (i.e., ogbn-protein, ogbn-arxiv), and link-level prediction (i.e., ogbl-ddi, ogbl-ppa) tasks. Following [48], we use ROC-AUC (i.e., area under the receiver operating characteristic curve) as the evaluation metric on ogbg-bbbb, ogbg-hiv, ogbg-moltoxcast, ogbg-moltox21, and ogbn-protein; RMSE (i.e., root mean squared error) on ogbg-molesol, and ogbg-freesolv; classification accuracy (Acc) on ogbn-arxiv; Hits@100 score on ogbl-ppa; Hits@20 score on ogbl-ddi.
Model Architectures and Training Details We use the GCN [11] 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 (e.g., learning rates, training epochs, batch size, etc.)
Table 3: Performance Comparison for Graph Classification Tasks using GIN with and without Pretraining. 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.
Model	ogbg-molbbbpl	ogbg-molhiv	ogbg-moltoxcast ROC-AUC (↑)	ogbg-moltox21	Average
GIN	65.5 ± 1.8	75.4 ± 1.5	63.3 ± 1.5	74.3 ± 0.5	69.63
GMoE-GIN	66.93 ± 1.72 (+1.43)	76.14 ± 1.03 (+0.74)	62.86 ± 0.37 (-0.44)	74.76 ± 0.66 (+0.46)	70.17 (+0.54)
GIN+Pretrain	69.94 ± 0.92	76.1 ± 0.8	62.96 ± 0.55	73.85 ± 0.64	70.71
GMoE-GIN+Pretrain	68.62 ± 1.02 (-1.32)	76.9 ± 0.9 (+0.8)	64.48 ± 0.50 (+1.18)	75.25 ± 0.78 (+1.40)	71.31 (+0.6)
Table 4: Node Prediction Task Results. The table header displays the dataset names and the corresponding evaluation metrics, both reported in percentage. Mean and standard deviation values from ten random runs are presented. The most notable results are highlighted in bold, with GMoE's performance gains indicated in parentheses
Model	ogbn-protein
ROC-AUC(↑)	ogbn-arxiv
Acc(↑)
GCN	73.53 ± 0.56	71.74 ± 0.29
GMoE-GCN	74.48 ± 0.58 (+0.95)	71.88 ± 0.32 (+0.14)