Using graph neural networks for large graphs is challenging since there is no clear way of constructing mini-batches. To solve this, previous methods have relied on sampling or graph clustering. While these approaches often lead to good training convergence, they introduce significant overhead due to expensive random data accesses and perform poorly during inference. In this work we instead focus on model behavior during inference. We theoretically model batch construction via maximizing the influence score of nodes on the outputs. This formulation leads to optimal approximation of the output when we do not have knowledge of the trained model. We call the resulting method influence-based mini-batching (IBMB). IBMB accelerates inference by up to 130x compared to previous methods that reach similar accuracy. Remarkably, with adaptive optimization and the right training schedule IBMB can also substantially accelerate training, thanks to precomputed batches and consecutive memory accesses. This results in up to 18x faster training per epoch and up to 17x faster convergence per runtime compared to previous methods. =-1

§ INTRODUCTION

Creating mini-batches is highly non-trivial for connected data, since it requires selecting a meaningful subset despite the data's connectedness. When the graph does not fit into memory, the mini-batching problem is equally relevant for both inference and training. However, mini-batching methods have so far mostly been focused on training, despite the major practical importance of inference. Once a model is put into production, it continuously runs inference to serve user queries. On AWS, more than 90 of infrastructure cost is due to inference, and less than 10 is due to training [Hutt et al., 2019]. Even during training, inference is necessary for early stopping and performance monitoring. A training method thus has rather limited utility by itself.

Selecting mini-batches for inference is distinctly different from training. Instead of averaging out stochastic sampling effects over many training steps, we need to ensure that every prediction is as accurate as possible. To achieve this, we propose a theoretical framework for creating mini-batches based on the expected influence of nodes on the outputs. Selecting nodes according to this formulation provably leads to an optimal approximation of the output.
The resulting optimization problem shows that we need to distinguish between two classes of nodes: Output nodes and auxiliary nodes. Output nodes are those for which we compute a prediction in this batch, for example a set of validation nodes. Auxiliary nodes provide inputs and define the batch's subgraph. This distinction allows us to choose a meaningful neighborhood for every prediction, while ignoring irrelevant parts of the graph. Note that output nodes in one batch can be auxiliary nodes in another batch.

This distinction furthermore splits mini-batching into two problems: 1. How do we partition output nodes into efficient mini-batches? 2. How do we choose the auxiliary nodes for a given set of output nodes? Having split the problem like this, we see that most previous works either focus exclusively on the first question by only using graph partitions [Chiang et al., 2019] or on the second question and choose a uniformly random subset of nodes as output nodes [Hamilton et al., 2017, Zou et al., 2019]. Jointly considering both aspects with an overarching theoretical framework allows for substantial synergy effects. For example, batching nearby output nodes together allows one output node to leverage another one's auxiliary nodes.

We call this overall framework influence-based mini-batching (IBMB). On the practical side, we propose two instantiations of IBMB by approximating the influence between nodes via personalized PageRank (PPR). We use fast approximations of PPR to select auxiliary nodes by their highest PPR scores. Accordingly, we partition output nodes using PPR-based node distances or via graph partitioning. We then use the subgraph induced by these nodes as a mini-batch. IBMB accelerates inference by up to 130x compared to previous methods that achieve similar accuracy.

Remarkably, we found that IBMB also works well for training, despite being derived from inference. This is due to the computational advantage of precomputed mini-batches, which can be loaded from a cache to ensure efficient memory accesses. We counteract the negative effect of the resulting sparse mini-batch gradients via adaptive optimization and batch scheduling. Overall, IBMB achieves an up to 18x improvement in time per training epoch, with similar final accuracy. This fast runtime more than makes up for any slow-down in convergence per step. Its speed advantage grows even further for the common setting of low label ratios, since our method avoids computation on irrelevant parts of the graph. Our implementation is available online[<https://www.cs.cit.tum.de/daml/ibmb>].
In summary, our core contributions are:

* Influence-based mini-batching (IBMB): A theoretical framework for selecting mini-batches for GNN inference based on influence scores.

* Practical instantiations of IBMB that work for a variety of GNNs and datasets. They substantially accelerate inference and training without sacrificing accuracy, especially for small label ratios.

* Methods for mitigating the impact of fixed, local mini-batches and sparse gradients on training.

§ BACKGROUND AND RELATED WORK

Graph neural networks. We consider a graph $\gG = (\gV, \gE)$ with node set $\gV$ and (possibly directed) edge set $\gE$. $N = |\gV|$ denotes the number of nodes, $E = |\gE|$ the number of edges, and $\mA \in \R^{N \times N}$ the adjacency matrix. GNNs use one embedding per node $\vh_u \in \R^H$ and edge $\ve_{(uv)} \in \R^{H_\text{e}}$ of size $H$ and $H_\text{e}$, and update them in each layer via message passing between neighboring nodes. We denote the embedding in layer $l$ as $\vh_u^{(l)}$ and its $i$'th entry as $\vh_{ui}^{(l)}$. Most GNNs can be expressed via the following equations:

\begin{align}
\vh_u^{(l+1)} &= f_{\text{node}}(\vh_u^{(l)}, \Agg_{v \in \gN_u}[f_{\text{msg}}(\vh_u^{(l)}, \vh_v^{(l)}, \ve_{(uv)}^{(l)})]),\\ \label{eq:gnn}
\ve_{(uv)}^{(l+1)} &= f_{\text{edge}}(\vh_u^{(l+1)}, \vh_v^{(l+1)}, \ve_{(uv)}^{(l)}).
\end{align}

The node and edge update functions $f_{\text{node}}$ and $f_{\text{edge}}$, and the message function $f_{\text{msg}}$ can be implemented using e.g. linear layers, multi-layer perceptrons (MLPs), and skip connections. The node's neighborhood $\gN_u$ is usually defined directly by the graph $\gG$ [Kipf & Welling, 2017], but can be generalized to consider larger or even global neighborhoods [Gasteiger et al., 2019, Alon & Yahav, 2021], or feature similarity [Deng et al., 2020]. The most common aggregation function $\Agg$ is summation, but multiple other alternatives have also been explored [Corso et al., 2020, Geisler et al., 2020]. Edge embeddings $\ve_{(uv)}$ are often not used in GNNs, but some variants rely on them exclusively [Chen et al., 2019].

Scalable GNNs. Multiple works have proposed massively scalable GNNs that leverage the peculiarities of message passing to condense it into a single step, akin to label or feature propagation [Bojchevski et al., 2020, Frasca et al., 2020]. Our work focuses on general, model-agnostic scalability methods. =-1

Scalable graph learning. Classical graph learning faced issues similar to GNNs when scaling to large graphs. Multiple frameworks for distributed graph computations were proposed to solve this without approximations or sampling [Gonzalez et al., 2012, Low et al., 2012, Malewicz et al., 2010, Kyrola et al., 2012]. Other works scaled to large graphs via stochastic variational inference, e.g. by sampling nodes and node pairs [Gopalan et al., 2012]. Interestingly, this approach is quite similar to sampling-based mini-batching for GNNs.

Mini batching for GNNs. Previous mini-batching methods can largely be divided into three categories: Node-wise sampling, layer-wise sampling, and subgraph-based sampling [Liu et al., 2022]. In node-wise sampling, we obtain a separate set of auxiliary nodes for every output node, which are sampled independently for each message passing step. Each output node is treated independently; if two output nodes sample the same auxiliary node, we compute its embedding twice [Hamilton et al., 2017, Ying et al., 2018, Liu et al., 2020]. Layer-wise sampling jointly considers all output nodes of a batch to compute a stochastic set of activations in each layer. Computations on auxiliary nodes are thus shared [Chen et al., 2018, Huang et al., 2018, Zou et al., 2019]. Subgraph-based sampling selects a meaningful subgraph and then runs the GNN on this subgraph as if it were the full graph. This method thus computes the outputs and intermediate embeddings of all nodes in that subgraph [Chiang et al., 2019, Zeng et al., 2020]. Our method most closely resembles the subgraph-based sampling approach. However, IBMB considers both output and auxiliary nodes, resulting in better batches, and only computes the output of predetermined output nodes, similar to node-wise sampling.
Two particularly related methods are Cluster-GCN [Chiang et al., 2019] and shaDow [Zeng et al., 2021]. However, both of these methods cover only part of the IBMB framework.
Cluster-GCN resembles the graph partitioning variant of output node partitioning (see <ref>). It does not select the most relevant auxiliary nodes and cannot ignore irrelevant parts of the graph. IBMB thus exhibits better accuracy for output nodes close to the partition boundary and is significantly faster when training on small training sets.
ShaDow is related to IBMB's node-wise auxiliary node selection variant (see <ref>). ShaDow does not address output node partitioning, resulting in worse runtimes and inconsistent accuracy in both inference and training.

Note that mini-batch generation is an orthogonal problem to training frameworks such as GNNAutoScale [Fey et al., 2021]. We can use IBMB to provide mini-batches as part of GNNAutoScale.

§ INFLUENCE-BASED MINI-BATCHING

Influence scores. To effectively create graph-based mini-batches we must first quantify how important one node is for another node's prediction. As proposed by Xu et al., 2018, we can do this via the influence score, which determines the local sensitivity of the output at node $u$ on the input at node $v$ as: =-1

\begin{equation}
I(v, u) = \sum_i \sum_j \left| \frac{\partial \vh^{(L)}_{ui}}{\partial \mX_{vj}} \right|,
\end{equation}

where $\vh^{(L)}_{ui}$ is the $i$'th entry in the embedding of node $u$ in the last layer $L$ and $\mX_{vj}$ is feature $j$ of node $v$. Analyzing the expected influence score can provide a crisp understanding of how to select nodes for inference when we only have knowledge of the graph, not the model or the node features. To formally prove this connection, we consider a slightly limited class of GNNs and model our lack of knowledge via a randomization assumption of ReLU activations, similar to Choromanska et al., 2015, and by assuming that all nodes have the same expected features, yielding (proof in <ref>):

Given a GNN with linear, graph-dependent aggregation and ReLU activations. Assume that all paths in the model's computation graph are activated with the same probability $\rho$ and nodes have features with expected value $\E[X_{v,i}] = \chi_i$. If we restrict the model input features to a set of auxiliary nodes $\gS_\textnormal{aux} \subseteq \gV$, then the error

\begin{equation}
\| \tilde{\vh}^{(L)}_{u} - \vh^{(L)}_{u} \|_1
\end{equation}

between the approximate logits $\tilde{\vh}^{(L)}_{u}$ and the true logits $\vh^{(L)}_{u}$ is minimized, in expectation, by selecting the nodes $v \in \gS_\textnormal{aux}$ with maximum influence score $I(v, u)$.

Formalizing mini-batching. We can leverage this insight by formalizing mini-batching as the optimization problem

\begin{equation}
\underbrace{%
\max_{\substack{P_\text{out} \in \sP(\gV_\text{out}) \\ |P_\text{out}| = b}} \sum_{\gS_\text{out} \in P_\text{out}}%
_{\text{Output node partitioning}}
\underbrace{%
\max_{\substack{\gS_\text{aux} \subseteq \gV \\ |\gS_\text{aux}| \le B}} \sum_{u \in \gS_\text{out}} \sum_{v \in \gS_\text{aux}}%
\vphantom{\max_{\substack{P_\text{out} \in \sP(\gV_\text{out}) \\ |P_\text{out}| = b}}}%
_{\text{Auxiliary node selection}}
\underbrace{%
I(v, u),%
\vphantom{\max_{\substack{P_\text{out} \in \sP(\gV_\text{out}) \\ |P_\text{out}| = b}}}%
_{\text{Influence score}}
\label{eq:avg_influence}
\end{equation}

where $\sP(\gV_\text{out})$ denotes the set of partitions of the output nodes $\gV_\text{out}$, $b$ the number of batches, and $B$ the maximum batch size. This optimization yields two results: The output node partition $P_\text{out}$ and the auxiliary node set for each batch of output nodes, $\gS_\text{aux}$. The hyperparameter $B$ is determined by the available (GPU) memory, while $b$ trades off runtime and approximation quality.
This formulation optimizes the average approximation across all outputs. This might not be ideal since some nodes might already be approximated well with a lower number of auxiliary nodes. We can instead focus on the worst-case approximation by optimizing the minimum aggregate influence score as

\begin{equation}
\underbrace{%
\max_{\substack{P_\text{out} \in \sP(\gV_\text{out}) \\ |P_\text{out}| = b}} \min_{\gS_\text{out} \in P_\text{out}}%
_{\text{Output node partitioning}}
\underbrace{%
\max_{\substack{\gS_\text{aux} \subseteq \gV \\ |\gS_\text{aux}| \le B}} \min_{u \in \gS_\text{out}} \sum_{v \in \gS_\text{aux}}%
\vphantom{\max_{\substack{P_\text{out} \in \sP(\gV_\text{out}) \\ |P_\text{out}| = b}}}%
_{\text{Auxiliary node selection}}
\underbrace{%
I(v, u).%
\vphantom{\max_{\substack{P_\text{out} \in \sP(\gV_\text{out}) \\ |P_\text{out}| = b}}}%
_{\text{Influence score}}
\label{eq:min_influence}
\end{equation}

Both <ref> split the mini-batching problem into three parts: Output node partitioning, auxiliary node selection, and influence score computation. We call this approach influence-based mini-batching (IBMB).

Computing influence scores. The model's influence score depends on various model details, especially when considering exact, trained models. In many cases we can calculate the expected influence score by making simplifying assumptions, similar to <ref>. This allows tailoring the mini-batching method to the exact model of interest.
For the remainder of this work we will focus our analysis on the broad class of models that use the average as an aggregation function, such as graph convolutional networks (GCN) [Kipf & Welling, 2017]. In this case, we can make similar assumptions on the GNN as in <ref> to prove that the influence score is proportional to a slightly modified random walk with $L$ steps [Xu et al., 2018]. To remove the influence score's dependence on the number of layers $L$, we can furthermore take the limit $L \to \infty$. Unfortunately, this would result in a limit distribution that is independent of node $v$. To avoid this we add restarts to the random walk, as proposed by Gasteiger et al., 2019. The limit $L \to \infty$ then becomes equivalent to personalized PageRank (PPR), which we can thus use an approximation of the influence score. Notably, PPR even works well for models with more complex, data-dependent influence scores, such as GAT (see <ref>). The PPR matrix is given by =-1

\begin{equation}
\mPi^{\text{ppr}} = \alpha (\mI_N - (1 - \alpha) \mD^{-1} \mA)^{-1},
\end{equation}

with the teleport probability $\alpha \in (0, 1]$ and the diagonal degree matrix $\mD_{ii} = \sum_k \mA_{ik}$. The entry $\mPi^{\text{ppr}}_{uv}$ then provides a measure for the influence of node $v$ on $u$. Calculating the above inverse is obviously infeasible for large graphs. However, we can approximate $\mPi^{\text{ppr}}$ with a sparse matrix $\tilde{\mPi}^{\text{ppr}}$ in time $\gO(\frac{1}{\varepsilon\alpha})$ per row, with error $\varepsilon \deg(v)$ [Andersen et al., 2006]. Importantly, this approximation uses only the node's local neighborhood, making its runtime independent of the overall graph size and thus massively scalable. Furthermore, the calculation is deterministic and model-independent, so we only need to perform this computation once during preprocessing.

§.§ Auxiliary node selection

Node-wise selection. Selecting auxiliary nodes on large graphs requires a method that efficiently yields nodes with highest expected influence. Fortunately, there is a well-developed literature of methods for finding the top-k PPR nodes. The classic approximate PPR method [Andersen et al., 2006] is guaranteed to provide all nodes with a PPR value $\mPi^{\text{ppr}}_{uv} > \varepsilon\deg(v)$ the root (output) node $u$. Optimizing auxiliary nodes by the worst-case influence score (<ref>) thus equates to separately running approximate PPR for each output node in a batch $\gS_\text{out}$, and then merging them.

Batch-wise selection. Considering each output node separately does not take into account how one auxiliary node jointly affects multiple output nodes, as required for the average-case formulation in <ref>. Fortunately, PPR calculation can be adapted to use a set of root nodes. To do so, we use a set of nodes in the teleport vector $\vt$ instead of a single node, by leveraging the underlying recursive equation for a PPR vector $\vpi_{\text{ppr}}(\vt) = (1 - \alpha) \mD^{-1} \mA \vpi_{\text{ppr}}(\vt) + \alpha \vt$. $\vt$ is a one-hot vector in the node-wise setting, while for batch-wise PPR it is $1/|\gS_\text{out}|$ for all nodes in $\gS_\text{out}$. This variant is also known as topic-sensitive PageRank. We found that batch-wise PPR is significantly faster than node-wise PPR. However, it can lead to cases where one outlier node receives almost no neighbors, while others have excessively many. Whether node-wise or batch-wise selection performs better thus often depends on the dataset and model.

Subgraph generation. Creating mini-batches also requires selecting a subgraph of relevant edges. We do so by using the subgraph induced by the selected output and auxiliary nodes in a batch. Note that the above node selection methods ignore how these changes to the graph affect the influence scores. This is a limitation of these methods. However, PPR is a local clustering method and we can thus expect auxiliary nodes to be well-connected.

§.§ Output node partitioning

Optimal partitioning. Finding the optimal node partition in <ref> would require trying out every possible partition since a change in $\gS_\text{out}$ can unpredictably affect the optimal choice of auxiliary nodes. Doing so is clearly intractable since the number of partitions grows exponentially with $N$ for a fixed batch size. We thus need to approximate the optimal partition via a scalable heuristic. The implicit goal of this step is finding output nodes that share a large number of auxiliary nodes. One good proxy for these overlaps is the proximity of nodes in the graph.

Distance-based partitioning. We propose two methods that leverage graph locality as a heuristic. The first is based on node distances. In this approach we first compute the pairwise node distances between nodes that are close in the graph. We can use PPR for this as well, since it is also commonly used as a node distance. If we select auxiliary nodes with node-wise PPR, we thus only need to calculate PPR scores once for both steps.

Next, we greedily construct the partition $P_\text{out}$ from $\tilde{\mPi}^{\text{ppr}}$. To do so, we start by putting every node $u$ into a separate batch $\{u\}$. We then sort all elements in $\tilde{\mPi}^{\text{ppr}}$ by magnitude, independent of their row or column. We scan over these values in descending order, considering the value's indices $(u, v)$ and merging the batches containing the two nodes. Afterwards we randomly merge any small leftover batches. We stay within memory constraints by only merging batches that stay below the maximum batch size $B$. This method achieves well-overlapping batches and can efficiently add incrementally incoming out nodes, e.g. in a streaming setting. Our experiments show that this method achieves a good compromise between well-overlapping batches and good gradients for training (see <ref>). Note that the resulting partition is unbalanced, some sets will be larger than others.

Graph partitioning. For our second method, we note that partitioning output nodes into overlapping mini-batches is closely related to partitioning graphs. We can thus leverage the extensive research on this topic by using the METIS graph partitioning algorithm [Karypis & Kumar, 1998] to find a partition of output nodes $P_\text{out}$. We found that graph partitioning yields roughly a two times higher overlap of auxiliary nodes than distance-based partitioning, thus leading to significantly more efficient batches. However, it also results in worse gradient samples, which we found to be detrimental for training (see <ref>).

Computational complexity. Since IBMB ignores irrelevant parts of the graph, inference and training scale linearly in the number of output nodes $\gO(N_\text{out})$. Preprocessing runs in $\gO(\frac{N_\text{out}}{\epsilon\alpha})$ for node-wise PPR-based steps, $\gO(\frac{b}{\epsilon\alpha})$ for batch-wise PPR, and in $\gO(E)$ for graph partitioning. The runtime of IBMB is thus independent of the graph size if we use distance-based partitioning. <ref> gives an overview of the full practical IBMB process.

Practical example of influence-based mini-batching (IBMB). The output nodes are indicated by pentagons. These nodes are first partitioned into batches, by grouping nearby nodes together. We then use influence scores to select the auxiliary nodes of each batch, neighbors with top-$k$ personalized PageRank (PPR) scores. Finally, we generate a batch using the induced subgraph of all selected nodes, but only calculate the outputs of the output nodes we chose when partitioning. Batches can overlap and do not need to cover the whole graph.

§ TRAINING WITH IBMB

Computational advantages. The above analysis focused on node outputs, not gradient estimation and training. However, IBMB also has inherent advantages for training, since we need to perform mini-batch generation only once during preprocessing. We can then cache each mini-batch in consecutive blocks of memory, thereby allowing the data to be stored where it is needed and circumventing expensive random data accesses. This significantly accelerates training, allows efficient distributed training, and enables more expensive node selection procedures. In contrast, most previous methods select both output and auxiliary nodes randomly in each epoch, which incurs significant overhead. Our experiments show that IBMB's more efficient memory accesses clearly outweigh the slightly worse gradient estimates (see <ref>). This seems counter-intuitive since the deterministic, fixed mini-batches in IBMB only provide sparse, fixed gradient samples. In this section we discuss these aspects and how adaptive optimization and batch scheduling counteract their effects.

Sparse gradients. Partitioning output nodes based on proximity effectively correlates the gradients sampled in a batch. The model thus sees a sparse gradient sample, which does not cover all aspects of the dataset. Fortunately, adaptive optimization methods such as Adagrad and Adam were developed exactly for such sparse gradients [Duchi et al., 2011, Kingma & Ba, 2015]. We furthermore ensure an unbiased training process by using every output (training) node exactly once per epoch. =-1

Fixed batches. Using a fixed set of batches can lead to problems with basic stochastic gradient descend (SGD) as well. Imagine training with two fixed batches whose loss functions have different minima. If training has “converged” to one of these minima, SGD would start to oscillate: It would take one step towards the other minimum, and then back, and so forth. To counteract this oscillation, we could add a “consensus constraint” to enforce a consensus between the weights after different batches, akin to distributed optimization [Olfati-Saber et al., 2007]. We can solve this constraint using a primal-dual saddle-point algorithm with directed communication [Gharesifard & Cortés, 2014]. The resulting dynamics are $\dot{x}^{(t)} = - \nabla \tilde{f}^{(t)}(x^{(t)}) - \alpha \lambda \dot{x}^{(t-1)} - \lambda^2 \dot{x}^{(t-2)}$, with the weights $x^{(t)}$ at time step $t$, the learning rate $\lambda$ and the dual variable $\alpha$. These dynamics resemble SGD with momentum, and fit perfectly into the framework of adaptive optimization methods [Reddi et al., 2018]. Indeed, momentum and adaptive methods suppress the oscillations in the above example with two minima. Accordingly, prior works have also found benefits in deterministically selecting fixed mini-batches [Wang et al., 2019, Banerjee & Chakraborty, 2021]. We further improve convergence by adaptively reducing the learning rate when the validation loss plateaus, which ensures that the step size decreases consistently.

Batch scheduling. While Adam with learning rate scheduling consistently ensures convergence, we still observe downward spikes in accuracy during training. To illustrate this issue, consider a sequence of mini-batches. In regular training every mini-batch is similar and the order of these batches is irrelevant. In our case, however, some of the mini-batches might be very similar. If the optimizer sees a series of similar batches, it will take increasingly large steps in a suboptimal direction, which leads to the observed downward spikes in accuracy.
We propose to prevent these suboptimal batch sequences by optimizing the order of batches. To quantify batch similarity we measure the symmetrized KL-divergence of the label distribution between batches. In particular, we use the normalized training label distribution $p_i = c_i / \sum_j c_j$, where $c_i$ is the number of training nodes of class $i$. This results in the pairwise batch distance $d_{ab}$ between batches $a$ and $b$. We propose two ways to use this for improving the batch schedule: (i) Find the fixed batch cycle that maximizes the batch distances between consecutive batches. This is a traveling salesman problem for finding the maximum distance loop that visits all batches. It is therefore only feasible for a small number of batches. (ii) Sample the next batch weighted by the distance to the current batch. Both scheduling methods improve convergence and increase final accuracy, at almost no cost during training. Overall, our training scheme leads to consistent convergence. Even accumulating gradients over the whole epoch does not significantly change convergence or final accuracy (see <ref>).

§ EXPERIMENTS

Experimental setup. We show results for two variants of our method: IBMB with PPR distance-based batches and node-wise PPR clustering (node-wise IBMB), and IBMB with graph partition-based batches and batch-wise PPR clustering (batch-wise IBMB). We also experimented with the two other combinations of the output node partitioning and auxiliary node selection variants, but found these two to work best. We compare them to four state-of-the-art mini-batching methods: Neighbor sampling [Hamilton et al., 2017], Layer-Dependent Importance Sampling (LADIES) [Zou et al., 2019], GraphSAINT-RW [Zeng et al., 2020], shaDow [Zeng et al., 2021], and Cluster-GCN [Chiang et al., 2019]. We use four large node classification datasets for evaluation: ogbn-arxiv <cit.>, ogbn-products <cit.>, Reddit [Hamilton et al., 2017], and ogbn-papers100M <cit.>. While these datasets use the transductive setting, IBMB makes no assumptions about this and can equally be applied to the inductive setting. We skip the common, small datasets (Cora, Citeseer, PubMed) since they are ill-suited for evaluating scalability methods. We do not strive to set a new accuracy record but instead aim for a consistent, fair comparison based on three standard GNNs: GCN [Kipf & Welling, 2017], graph attention networks (GAT) [Veličković et al., 2018], and GraphSAGE [Hamilton et al., 2017]. We use the same training pipeline for all methods, giving them access to the same optimizations. Since full inference is too slow to execute every epoch we use the mini-batching method used for training to also approximate inference during training. We run each experiment 10 times and report the mean and standard deviation in all tables and the bootstrapped mean and 95 confidence intervals in all figures. We fully pipeline data loading and batch creation by prefetching batches in parallel. We found that using more than one worker for this does not improve runtime, most likely because data loading is limited by memory bandwidth, which is shared between workers. We keep GPU memory usage constant between methods, and tune all remaining hyperparameters for both IBMB and the baselines. See <ref> for full experimental details. =-1

Test accuracy and log. inference time averaged over a fixed set of 10 pretrained GNNs. IBMB provides the best accuracy versus time trade-off (top-left corner) in all settings.

Inference. <ref> compares inference accuracy and time of different batching methods, using the same pretrained model and varying computational budgets (number of auxiliary nodes/sampled nodes) at a fixed GPU memory budget. IBMB provides the best trade-off between accuracy and time in all settings. Node-wise IBMB performs better than graph partitioning, except on ogbn-products. IBMB provides a significant speedup over chunking-based full-batch inference on GPU, being 10900 times faster at comparable accuracy. All previous methods are either significantly slower or less accurate. =-1

Training convergence of validation accuracy in log. time. Average and 95 confidence interval of 10 runs. GraphSAINT-RW does not reach the shown accuracy range in some settings due to its bad validation performance. IBMB converges the fastest in 9 out of 10 settings.

Training. IBMB performs significantly better in training than previous methods, converging up to 17x faster than all baselines (see <ref>). This is despite the fact that we always prefetch the next batch in parallel. Note that GAT is slower to compute than GCN and GraphSAGE, limiting the positive impact of a fast batching method. Compute-constrained models like GAT are less relevant in practice since data access is typically the bottleneck for GNNs on large, often even disk-based datasets [Bojchevski et al., 2020]. <ref> in the appendix furthermore shows that IBMB's time per epoch is significantly faster than all sampling-based methods. Cluster-GCN has a comparable runtime, which is expected due to its similarity with IBMB. However, it converges more slowly than IBMB and reaches substantially lower final accuracy. Neighbor sampling achieves good final accuracy, but is extremely slow. GraphSAINT-RW only achieves good final accuracy with prohibitively expensive full-batch inference. Node-wise IBMB achieves the best final accuracy with a scalable inference method in 7 out of 10 settings.
On ogbn-papers100M, IBMB has a substantially faster time per epoch and lower memory consumption than previous methods demonstrating IBMB's favorable scaling with dataset size. It even performs as well as SIGN-XL (66.1 ±0.2) [Frasca et al., 2020], while using 30x fewer parameters and no hyperparameter tuning. Notably, we were unable to evaluate GraphSAINT-RW and Cluster-GCN on this dataset, since they use more than 256 of main memory.

Preprocessing. IBMB requires more preprocessing than previous methods. However, since IBMB is rather insensitive to hyperparameter choices (see <ref>, <ref>), preprocessing rarely needs to be re-run. Instead, its result can be saved to disk and re-used for training different models. Just considering our 10 training seeds, preprocessing of node-wise IBMB only took 1.3 of the training time for GCN and 0.25 for GAT on ogbn-arxiv.

Training set size. The ogbn-arxiv and ogbn-products datasets both contain a large number of training nodes (91k and 197k, respectively). However, labeling training samples is often an expensive endeavor, and models are commonly trained with only a few hundred or thousand training samples. GraphSAINT-RW and Cluster-GCN are global training methods, i.e. they always use the whole graph for training. They are thus ill-suited for the common setting of a large overall graph containing a small number of training nodes (resulting in a small label rate). In contrast, the training time of IBMB purely scales with the number of training nodes. To demonstrate this, we reduce the label rate by sub-sampling the training nodes of ogbn-products and compare the convergence in <ref>. As expected, the gap in convergence speed between IBMB and both Cluster-GCN and GraphSAINT-RW grows even larger for smaller training sets.

Ablation studies. We ablate our output node partitioning schemes by instead batching together random sets of nodes. We use fixed batches since we found that resampling incurs significant overhead without benefits – which is consistent with our considerations on gradient samples and contiguous memory accesses. <ref> shows that this method (“Fixed random”) converges more slowly and does not reach the same level of accuracy as our partition schemes. Node-wise IBMB converges the fastest, which suggests a trade-off between full gradient samples (random batching) and maximum batch overlap (graph partitioning). <ref> shows that random batching (“IBMB, rand batch.”) is also substantially slower and often less accurate for inference. This is due to the synergy effects of output node partitioning: If output nodes have similar auxiliary nodes, they benefit from each other's neighborhood.
We ablate auxiliary node selection by comparing IBMB to Cluster-GCN, since it just uses the graph partition as a batch instead of smartly selecting auxiliary nodes. We use the graph partition size as the number of auxiliary nodes for batch-wise IBMB to allow for a direct comparison. As discussed above, Cluster-GCN consistently performs worse, especially in terms of final accuracy, for inference, and for small label rates.
Finally, <ref> compares the proposed batch scheduling methods. Optimal and weighted sampling-based scheduling improve convergence and prevent or reduce downward spikes in accuracy.

Training convergence in log. time for GCN on ogbn-products with smaller training sets. The gap in convergence speed between IBMB and the baselines grows larger for small training sets, since IBMB scales with training set size and not with overall graph size.

Trained accuracy for node-wise IBMB, depending on the output nodes per batch (GCN, ogbn-arxiv). The impact of this choice is rather minor.

Convergence per time for training GCN on ogbn-arxiv. Both batch-wise and node-wise IBMB lead to faster convergence than fixed random batches.

Batch scheduling for GAT on ogbn-arxiv. Optimal batch order prevents downward spikes in accuracy and leads to higher final accuracy.

Gradient accumulation for batch-wise IBMB on GCN, ogbn-arxiv. The difference is minor, even when accumulating over the full epoch.

Sensitivity analysis. Different IBMB is largely insensitive to different local clustering methods and hyperparameters for selecting auxiliary nodes (see <ref>). Even increasing the number of output nodes per batch with a fixed number of auxiliary nodes per output node only has a minor impact on accuracy, especially above 1000 output nodes per batch, as shown by <ref>. IBMB performs well even in extremely constrained settings with small batches of 100 output nodes per batch. In practice, IBMB only has one free hyperparameter: The number of auxiliary nodes per output node, which allows optimizing for accuracy or speed. The number of output nodes per batch is then given by the available GPU memory, while the local clustering method and other hyperparameters are not important. =-1

Gradient accumulation. Accumulating gradients across multiple batches is a method for smoothing batches if gradient noise is too high. We might expect this to happen in IBMB due to the sparse gradients caused by primary node partitioning. However, <ref> shows that gradient accumulation in fact only has an insignificant effect on IBMB, demonstrating its stability during training.

§ CONCLUSION

We propose influence-based mini-batching (IBMB), a method for extracting batches for GNNs. IBMB formalizes creating batches for inference by maximizing the influence score on the output nodes. Remarkably, with an adaptive optimizer and batch scheduling IBMB even performs well during training. It improves training convergence by up to 17x and inference time by up to 130x compared to previous methods that reach similar accuracy. IBMB performs especially well for sparse labels, large datasets, and when the pipeline is constrained by data access speed.

§.§.§ Acknowledgments

This work was supported by the Deutsche Forschungsgemeinschaft (DFG) through TUM International Graduate School of Science and Engineering (IGSSE) (GSC 81).

[Alon & Yahav, 2021]
Uri Alon and Eran Yahav.
On the Bottleneck of Graph Neural Networks and its
Practical Implications.
In ICLR, 2021.

[Andersen et al., 2006]
R. Andersen, F. Chung, and K. Lang.
Local Graph Partitioning using PageRank Vectors.
In FOCS, 2006.

[Banerjee & Chakraborty, 2021]
Subhankar Banerjee and Shayok Chakraborty.
Deterministic Mini-batch Sequencing for Training Deep
Neural Networks.
In AAAI, 2021.

[Bojchevski et al., 2020]
Aleksandar Bojchevski, Johannes Gasteiger, Bryan Perozzi, Amol Kapoor, Martin
Blais, Benedek Rózemberczki, Michal Lukasik, and Stephan Günnemann.
Scaling Graph Neural Networks with Approximate PageRank.
In KDD, 2020.

[Chen et al., 2018]
Jie Chen, Tengfei Ma, and Cao Xiao.
FastGCN: Fast Learning with Graph Convolutional Networks
via Importance Sampling.
In ICLR, 2018.

[Chen et al., 2019]
Zhengdao Chen, Lisha Li, and Joan Bruna.
Supervised Community Detection with Line Graph Neural
In ICLR, 2019.

[Chiang et al., 2019]
Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh.
Cluster-GCN: An Efficient Algorithm for Training Deep and
Large Graph Convolutional Networks.
In KDD, 2019.

[Choromanska et al., 2015]
Anna Choromanska, Yann LeCun, and Gérard Ben Arous.
Open Problem: The landscape of the loss surfaces of multilayer
In COLT, 2015.

[Corso et al., 2020]
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar
Principal Neighbourhood Aggregation for Graph Nets.
In NeurIPS, 2020.

[Deng et al., 2020]
Chenhui Deng, Zhiqiang Zhao, Yongyu Wang, Zhiru Zhang, and Zhuo Feng.
GraphZoom: A Multi-level Spectral Approach for Accurate
and Scalable Graph Embedding.
In ICLR, 2020.

[Dréo et al., 2006]
Johann Dréo, Alain Pétrowski, Patrick Siarry, and Eric Taillard.
Metaheuristics for Hard Optimization: Methods and Case

[Duchi et al., 2011]
John C. Duchi, Elad Hazan, and Yoram Singer.
Adaptive Subgradient Methods for Online Learning and
Stochastic Optimization.
Journal of Machine Learning Research, 12:0 2121–2159,

[Fey et al., 2021]
Matthias Fey, Jan E. Lenssen, Frank Weichert, and Jure Leskovec.
GNNAutoScale: Scalable and Expressive Graph Neural
Networks via Historical Embeddings.
In ICML, 2021.

[Frasca et al., 2020]
Fabrizio Frasca, Emanuele Rossi, Davide Eynard, Ben Chamberlain, Michael
Bronstein, and Federico Monti.
SIGN: Scalable Inception Graph Neural Networks.
arXiv, 2004.11198, 2020.

[Gasteiger et al., 2019]
Johannes Gasteiger, Aleksandar Bojchevski, and Stephan Günnemann.
Predict then Propagate: Graph Neural Networks Meet
Personalized PageRank.
In ICLR, 2019.

[Gasteiger et al., 2019]
Johannes Gasteiger, Stefan Weißenberger, and Stephan Günnemann.
Diffusion Improves Graph Learning.
In NeurIPS, 2019.

[Geisler et al., 2020]
Simon Geisler, Daniel Zügner, and Stephan Günnemann.
Reliable Graph Neural Networks via Robust Aggregation.
In NeurIPS, 2020.

[Gharesifard & Cortés, 2014]
Bahman Gharesifard and Jorge Cortés.
Distributed Continuous-Time Convex Optimization on
Weight-Balanced Digraphs.
IEEE Transactions on Automatic Control, 590
(3):0 781–786, 2014.

[Gonzalez et al., 2012]
Joseph E. Gonzalez, Yucheng Low, Haijie Gu, Danny Bickson, and Carlos Guestrin.
PowerGraph: Distributed Graph-Parallel Computation on
Natural Graphs.
In OSDI, 2012.

[Gopalan et al., 2012]
Prem K. Gopalan, Sean Gerrish, Michael Freedman, David Blei, and David Mimno.
Scalable Inference of Overlapping Communities.
In NeurIPS, 2012.

[Hamilton et al., 2017]
William L. Hamilton, Zhitao Ying, and Jure Leskovec.
Inductive Representation Learning on Large Graphs.
In NeurIPS, 2017.

[Hu et al., 2020]
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
In NeurIPS, 2020.

[Huang et al., 2018]
Wenbing Huang, Tong Zhang, Yu Rong, and Junzhou Huang.
Adaptive Sampling Towards Fast Graph Representation
In NeurIPS, 2018.

[Hutt et al., 2019]
Gadi Hutt, Vibhav Viswanathan, and Adam Nadolski.
Deliver high performance ML inference with AWS Inferentia,

[Karypis & Kumar, 1998]
George Karypis and Vipin Kumar.
A Fast and High Quality Multilevel Scheme for
Partitioning Irregular Graphs.
SIAM Journal on Scientific Computing, 200
(1):0 359–392, 1998.

[Kingma & Ba, 2015]
Diederik P. Kingma and Jimmy Ba.
Adam: A Method for Stochastic Optimization.
In ICLR, 2015.

[Kipf & Welling, 2017]
Thomas N. Kipf and Max Welling.
Semi-Supervised Classification with Graph Convolutional
In ICLR, 2017.

[Kyrola et al., 2012]
Aapo Kyrola, Guy Blelloch, and Carlos Guestrin.
GraphChi: Large-Scale Graph Computation on Just a PC.
In OSDI, 2012.

[Liu et al., 2022]
Xin Liu, Mingyu Yan, Lei Deng, Guoqi Li, Xiaochun Ye, and Dongrui Fan.
Sampling Methods for Efficient Training of Graph
Convolutional Networks: A Survey.
IEEE/CAA Journal of Automatica Sinica, 90
(2):0 205–234, 2022.

[Liu et al., 2020]
Ziqi Liu, Zhengwei Wu, Zhiqiang Zhang, Jun Zhou, Shuang Yang, Le Song, and Yuan
Bandit Samplers for Training Graph Neural Networks.
In NeurIPS, 2020.

[Low et al., 2012]
Yucheng Low, Danny Bickson, Joseph Gonzalez, Carlos Guestrin, Aapo Kyrola, and
Joseph M. Hellerstein.
Distributed GraphLab: a framework for machine learning and data
mining in the cloud.
In VLDB, 2012.

[Malewicz et al., 2010]
Grzegorz Malewicz, Matthew H. Austern, Aart J.C Bik, James C. Dehnert, Ilan
Horn, Naty Leiser, and Grzegorz Czajkowski.
Pregel: a system for large-scale graph processing.
In SIGMOD, 2010.

[Olfati-Saber et al., 2007]
Reza Olfati-Saber, J. Alex Fax, and Richard M. Murray.
Consensus and Cooperation in Networked Multi-Agent Systems.
Proceedings of the IEEE, 950 (1):0 215–233,

[Reddi et al., 2018]
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar.
On the Convergence of Adam and Beyond.
In ICLR, 2018.

[Veličković et al., 2018]
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro
Liò, and Yoshua Bengio.
Graph Attention Networks.
In ICLR, 2018.

[Wang et al., 2020]
Kuansan Wang, Zhihong Shen, Chiyuan Huang, Chieh-Han Wu, Yuxiao Dong, and
Anshul Kanakia.
Microsoft Academic Graph: When experts are not enough.
Quantitative Science Studies, 10 (1):0
396–413, 2020.

[Wang et al., 2019]
Shengjie Wang, Wenruo Bai, Chandrashekhar Lavania, and Jeff Bilmes.
Fixing Mini-batch Sequences with Hierarchical Robust
In AISTATS, 2019.

[Xu et al., 2018]
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi,
and Stefanie Jegelka.
Representation Learning on Graphs with Jumping Knowledge
In ICML, 2018.

[Ying et al., 2018]
Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L. Hamilton, and
Jure Leskovec.
Graph Convolutional Neural Networks for Web-Scale
Recommender Systems.
In KDD, 2018.

[Zeng et al., 2020]
Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor
GraphSAINT: Graph Sampling Based Inductive Learning
In ICLR, 2020.

[Zeng et al., 2021]
Hanqing Zeng, Muhan Zhang, Yinglong Xia, Ajitesh Srivastava, Andrey Malevich,
Rajgopal Kannan, Viktor Prasanna, Long Jin, and Ren Chen.
Decoupling the Depth and Scope of Graph Neural Networks.
In NeurIPS, 2021.

[Zou et al., 2019]
Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu.
Layer-Dependent Importance Sampling for Training Deep and
Large Graph Convolutional Networks.
In NeurIPS, 2019.