Graph Ladling: Shockingly Simple Parallel GNN Training
without Intermediate Communication

Graph Neural Networks (GNNs) exploit an iterative message passing (MP) scheme to capture the structural information within nodes’ neighborhoods. Let $G=(\mathcal{V},\mathcal{E})$ be a simple and undirected graph with vertex set $\mathcal{V}$, edge set $\mathcal{E}$ and .Let $\boldsymbol{A}\in^{N\times N}$ be the associated adjacency matrix, such that $\boldsymbol{A}_{ij}=1$ if $(i,j)\in\mathcal{E}$, and $\boldsymbol{A}_{ij}=0$ otherwise. $N=\lvert\mathcal{V}\rvert$ is the number of nodes. Let $\boldsymbol{X}\in^{N\times P_{0}}$ be the node feature matrix, whose $i$-th row represents a $P_{0}$-dimension feature vector for the $i$-th node. To facilitate understanding, a unified formulation of MP with K layers is presented as follows:

where $\sigma$ is an activation function (e.g. ReLU), $\boldsymbol{A}^{(l)}$ is the adjacency matrix at $l$-th layer, and $\boldsymbol{W}^{(l)}$ is the feature transformation weight matrix at $l$-th layer which will be learnt for the downstream tasks.

Despite the enormous success of GNNs in learning high-quality node representations, many state-of-the-art GNNs are shallow due to notorious challenges in their trainability, resulting in sub-standard performance. Recently, we have been observing a massive explosion in the topological structure of real-world graph data, which raises demand to increase the model capacity to facilitate capturing long-range dependencies, and relaxing the widely-adopted restrictions to learn from limited neighbors subjected to resource constraints due to the design issues of message passing strategy.

Several works (Li et al., 2021; Jaiswal et al., 2022; Zhou et al., 2021b) have attempted the empirical scaling law from NLP, which suggests that over-parametrized models have better generalization capabilities, and proposed to similarly scale GNNs by stacking layers citing their increased ability to capture long-term dependencies. However, scaling GNNs capacity either by deepening to capture long-range dependencies or by widening to increase neighborhood aggregation is inhibited by the prevalent issues of unhealthy gradients during back-propagation leading to poor trainability, over-smoothening causing feature collapse, and information squashing due to overload and distortion of messages being propagated from distant nodes.

Figure 1 illustrates the effect of state-of-the-art GNNs scaling from the perspective of deepening and widening on OGBN-arxiv. It can be clearly observed that vanilla-GCN, GCNII and GPRGNN, all suffer significant performance from drop with increasing model capacity. Note that this substandard performance still exists in GCNII and GPRGNN, which are equipped with SOTA architectural modifications (eg. skip-connections), regularization, and normalization. Additionally, as shown in Figure 1(right), increasing the breadth of message passing to capture more and more neighborhood information for GraphSAGE on OGBN-arxiv does not necessarily help in improving the performance rather increases the memory footprints and reduce throughput (iterations/sec).

To this end, we argue that model explosion is not necessarily the only solution towards building high-quality generalizable GNNs and propose an orthogonal direction inspired by the recent distributed and parallel training platform of model soups. We propose to build independently and parallelly trained multiple comparatively weaker models without any intermediate communication, and merge their knowledge to create a better generalizable GNN. In the next sections, we will explain how model soups can be easily adapted for many state-of-the-art GNNs, and how can we smoothly extend the current graph sampling and partitioning to create candidate ingredients for the model soup.

In this section, we discuss about the harmonious adaptation of model soups for various state-of-the-art GNNs. Although model soups have been recently receiving ample attention for parallel training of large language pre-trained models, it is still unclear and unexplored if they will work for comparatively much smaller GNNs trained from scratch. We for the first time observed that unlike using a pre-trained initialization in LLMs, GNNs optimized independently from the same random initialization organically lie in the same basin of the error landscape, and they can be smoothly aggregated using linear interpolation across their weight parameters, with significant performance gain.

More specifically, we propose to create $\boldsymbol{K}$ instances (soup ingredients) of the same GNN architecture of interest, sharing a common random initialization to ensure that the optimization trajectory does not deviate significantly during training facilitating smooth mergability. Subject to resource availability, GNN instances are trained independently in isolation across multiple GPUs using the hyperparameter configuration set ${\{h_{1},h_{2},...,h_{K}\}}$ designed with slight variations in the learning rate, weight decay, seed values, etc.

Once the soup ingredients are ready, we start preparing the soup using our greedy interpolation soup procedure as illustrated in Algorithm 1, which in general follows (Wortsman et al., 2022b). The greedy soup sequentially adds each model as a potential ingredient in the soup, and only keeps the model in the soup if it leads to improving performance on the validation set. For merging two ingredients, we perform a search for an optimal interpolation ratio $\alpha\in[0,1]$ that helps in performance gain, else the ingredient is discarded.

In comparison with several state-of-the-art GNNs, we experimentally illustrate that GNN soups prepared using ingredients having exactly the same model configuration, significantly perform better without any requirement of increasing model depth or width. For example, a 4-layer GCNII model soup prepared using 50 candidate ingredients, beats GCNII with a similar configuration on PubMed by $1.12\pm 0.36$. Our experiments across multiple GNN architectures and datasets ranging from small graphs to large graphs {Cora, Citeseer, PubMed, OGBN-ArXiv, OGBN-products} strongly validate the universal effectiveness of model soups for GNNs.

In this section, we discuss our approach for preparing data-centric model soups in scenarios when we have resource constraints to perform message passing with the entire graph data in memory, by leveraging the SOTA graph sampling and partitioning mechanisms. As MP requires nodes aggregating information from their neighbors, the integral graph structures are inevitably preserved during forward and backward propagation, thus occupying considerable running memory and time. Following Equation 1, the key bottleneck relies in $\boldsymbol{A}^{(l)}\boldsymbol{X}^{(l)}$ for vanilla MP, requiring entire sparse adjacency matrix to be present in one GPU during training, which becomes challenging with growing topology of real-world graphs. Recently, numerous efforts with regards to graph sampling (Hamilton et al., 2017; Zou et al., 2019; Chen et al., 2018) and partitioning (Chiang et al., 2019; Zeng et al., 2019) have been proposed for approximating the iterative full-batch MP to reduce the memory consumption for training within one single GPU to mitigate the aforementioned issue and scale up GNNs.

Provided the formulation of Equation 1; graph sampling and partitioning paradigms pursue an optimal way to perform batch training, such that each batch will meet the memory constraint of a single GPU for message passing. We restate a unified illustration to elucidate the formulation of graph sampling and partitioning used for training our model soup ingredient $M_{i}$ as follows:

where $\boldsymbol{\tilde{A}}^{\boldsymbol{(l)}}$ indicate the adjacency matrix for the $l$-th layer sampled from the full graph, $\mathcal{B}_{l}$ is the set of nodes sampled at $l$-th layer, $M_{i}$ and $\boldsymbol{W}^{(l)}[M_{i}]$ indicate the $i$-th candidate ingredient and weight associated with its $l$-th layer. This mini-batch training approach combined with a sampling and partitioning strategy significantly helps in alleviating time consumption and memory usage which grow exponentially with the GNN depth. It is worth noting that we have explored extension to scalable infrastructure topics like distributed training with multiple GPUs to alleviate the expenses of training ingredients for model soup.

Our experiments are conducted on two GPU servers equipped with RTX A6000 and RTX 3090 GPUs. The hyper-parameters for soup ingredients corresponding to different datasets training are selected via our built-in efficient parameter sweeping functionalities from pre-defined ranges nearby our optimal setting in Appendix B. For our small-scale experiments, we use three standard citation network datasets: Cora, Citeseer, and Pubmed, while our large-scale experiments are conducted with four popular large-scale open benchmarks: Flickr, Reddit, OGBN-ArXiv, and OGBN-products (Appendix A). For our evaluation on our chosen datasets, we have closely followed the data split settings and metrics reported by the recent benchmark (Duan et al., 2022; Chen et al., 2021). We consider variations in the key hyperparameters {random seed, batch size, learning rate, weight decay, and dropout rate} for generating the candidate ingredient models in our model soup. Unless explicitly specified, we have used 50 candidate ingredients for small-scale graphs and 30 candidate ingredients for large-scale graphs experiments and our model soup is prepared using interpolation hyperparameter $\alpha\in[0-1]$ with a step size of 0.01 in Algorithm 1. For comparison with SOTA models, we have adopted the official implementation (Appendix C) of JKNet, DAGNN, APPNP, GCNII, SGC, ClusterGCN, and GraphSAINT. Note that while generating our model soup, we make sure to use exactly the same architectural design for our candidate ingredients to ensure a fair comparison.

In this section, we conduct a systematic and extensive study to understand the harmonious adaptation of model soups for state-of-the-art GNNs on popular graph benchmarks Cora, Citeseer, Pubmed, and OGBN-ArXiv. Note that although model soups have recently attracted significant attention for large pre-trained language models, it is still unclear and unexplored if they can work for comparatively much smaller graph neural networks trained from scratch which learn from graph-structured data with relational properties unlike NLP and vision datasets having independent training samples. Model soups provide an orthogonal way of increasing model capacity without deepening or widening GNNs which brings many unwanted trainability issues in GNNs. Table 1 illustrates the performance summarization of our model soup generated using 50 candidate ingredients independently trained on the benchmarking datasets with respect to the vanilla performance of several SOTA GNNs following the exact same architectural design (4 layers with 256 hidden dimension) to ensure a fair comparison. The results reported for vanilla GNNs within Table 1 are averaged across 30 independent runs using different seed values selected at random from $[1-100]$ without replacement.

We first observe that across all datasets and GNN architecture, model soups outputs significantly outperform their vanilla counterpart. More noticeably, it improves GCN performance on Cora by $1.13\%$, GCNII performance on PubMed by $1.38\%$, JKNet, DAGNN, and SGC performance on OGBN-ArXiv by $1.22\%$,$0.79\%$,$1.17\%$ respectively, and APPNP performance on Citeseer by $1.02\%$. Moreover, Table 2 illustrates the current state of various fancy architectural and regularization modifications recently proposed to facilitate deepening of GCNs and help in improving their performance. It can be clearly observed that our model soup prepared by combining the strength of 50 candidate ingredients of 2-layer GCNs can significantly outperform all these fancy methods, bolstering our claim that model explosion by deepening and widening is necessarily not the only and right direction for building high-quality generalizable GNNs.

In this section, we provide experimental results for preparing data-centric model soup in scenarios when we do not have the luxury of resources to perform message passing on the entire graph, by leveraging the SOTA graph sampling (node, edge, layer) and partitioning mechanisms to prepare the candidate ingredients of the soup. Table 4 illustrates the performance comparison of state-of-the-art graph sampling approaches GraphSAGE, FastGCN, and LADIES with respect to our node, edge, and layer-wise sampled soup prepared using Algorithm 2. Results of GraphSAGE, FastGCN, and LADIES are reported as mean across 30 independent runs while our soup results are reported as the mean of 5 independent runs. We additionally report the performance of graph ensemble prepared using the 30 independent runs of best-performing baseline (GraphSAGE). It can be clearly observed that our Node sampled soup performs best and comfortably outperforms all the baselines along with the graph ensemble which has a hefty inference cost, by significant margins. Similar to the sub-standard performance of layer-wise sampling approaches FastGCN and LADIES, our Layer sampled soup has comparatively low (although better than FastGCN and LADIES) performance, possibly because layer-wise induced adjacency matrix is usually sparser than the others, which accounts for its sub-optimal performance.

In Table 6, we attempted to analyze the activation memory usage of activations and the hardware throughput (higher is better) of GraphSAGE (neighborhood sample size of 40) with respect to our data-centric model soup using node sampling equivalent to one-fouth of GraphSAGE (i.e., neighborhood sample size of 10) on OGBN-products. We found that with comparatively less memory requirement, our approach can $~{}\sim 1\%$ better performance, eliciting the high potential in improving GNNs performance by combining the strength of multiple weak models. Next, Table 5 illustrates the performance comparison of our Graph Partition Soup with respect to two well-known graph partitioning GNNs: ClusterGCN and GraphSAINT. Results of ClusterGCN and GraphSAINT are reported as mean across 30 independent runs while our soup results are reported as mean of 5 independent runs each generated using 30 candidate ingredients using Algorithm 3. We also present the performance of graph ensemble prepared using the 30 independent runs of the best-performing baseline (ClusterGCN). Across all benchmarking datasets (Flicket, Reddit, OGBN-ArXiv, and OGBN-products), our data partition soup outperforms both GraphSAINT and ClusterGCN. More noticeably, our method beats ClusterGCN with a similar architecture design by $>1\%$ margin on Flickr and OGBN-ArXiv dataset. In addition, Figure 2 illustrates the individual performance of each participating candidate ingredient of our data-partition soup, and it can be clearly observed that there is orthogonal knowledge stored in the learned weights of these networks. This gets elucidated by merging candidates and thereby improving overall performance which demonstrates the strength of our data-centric model soups.

In this section, we try to understand the strength of increasing ingredient count to our final data-centric model soup performance. Table 7 illustrates the performance comparison of our data partition soup on OGBN-ArXiv dataset, with varying candidate ingredient counts. It can be clearly observed that increasing candidates generally lead to better performance, thanks to the greedy interpolation procedure. However, due to the increase in computational and time cost, we restrict the number of ingredients to 30 for all experiments related to large graphs, and 50 for small graphs.

In this section, we attempt to answer another abaltion question: How does intermediate communication across candidate ingredients (i.e. souping at intervals during training) benefit the performance of the final model soup? To this end, we prepared a data partition model soup of GCNs using OGBN-ArXiv dataset, where we executed souping across candidate ingredients at regular intervals of 100 epochs during training. To our surprise, we found that the final soup performance is $-0.745\%$ less compared to our communication-free approach presumably due to a lack of diversity and orthogonal knowledge across the candidate ingredients. Moreover, intermediate communication incurs additional soup preparation overhead along with a new communication interval hyperparameter for optimization.

Our algorithm is mainly inspired by a GNN ensembling perspective, and its aim is to optimize model accuracy with or without distributed data parallelism, by interpolating individually trained GNN model weights. Meanwhile, several recent or concurrent works have contributed to improving speed and accuracy for distributed data-parallel GNN training with efficient model averaging or randomized partitions.

(Ramezani et al., 2021) proposed a communication-efficient distributed GNN training technique (LLCG), which first trains a GNN on local machine data by ignoring the dependency between nodes among different machines, then sends the locally trained model to the server for periodic model averaging. However, LLCG heavily relies on the global Server Corrections module on the server to refine the locally learned models. CoFree-GNN (Cao et al., 2023) presents another communication-free GNN training framework, that utilizes a Vertex Cut partitioning, i.e., rather than partitioning the graph by cutting the edges between partitions, the Vertex Cut partitions the edges and duplicates the node information to preserve the graph structure.

One concurrent and independent work by (Zhu et al., 2023) proposed a distributed training framework that assembles independent trainers, each of which asynchronously learns a local model on locally-available parts of the training graph. In this way, they only conducted periodic (time-based) model aggregation to synchronize the local models. Note that, unlike their periodic weight averaging of the participating trainers, our work performs averaging only after the complete training of candidates. Meanwhile, our candidates are trained by sampling different graph clusters from the input graph every epoch, to facilitate the diversity demand for model soups from the input graph. In comparison, each candidate model in (Zhu et al., 2023) only trains on localized subgraph assigned by randomized node/super-node partitions, without having access to entire input graph.