# GNNAutoScale: Scalable and Expressive Graph Neural Networks
via Historical Embeddings

Matthias Fey  Jan Eric Lenssen  Frank Weichert  Jure Leskovec

# Appendix

Matthias Fey  Jan Eric Lenssen  Frank Weichert  Jure Leskovec

###### Abstract

We present _G NNAutoScale_ (GAS), a framework for scaling arbitrary message-
passing GNNs to large graphs. GAS prunes entire sub-trees of the computation
graph by utilizing historical embeddings from prior training iterations,
leading to constant GPU memory consumption in respect to input node size
without dropping any data. While existing solutions weaken the expressive
power of message passing due to sub-sampling of edges or non-trainable
propagations, our approach is provably able to maintain the expressive power
of the original GNN. We achieve this by providing approximation error bounds
of historical embeddings and show how to tighten them in practice.
Empirically, we show that the practical realization of our framework, _PyGAS_
, an easy-to-use extension for PyTorch Geometric, is both fast and memory-
efficient, learns expressive node representations, closely resembles the
performance of their non-scaling counterparts, and reaches state-of-the-art
performance on large-scale graphs.

Machine Learning, Deep Learning, Graph Neural Networks, ICML

##  1 Introduction

_Graph Neural Networks_ (GNNs) capture local graph structure and feature
information in a trainable fashion to derive powerful node representations
suitable for a given task at hand (Hamilton, 2020; Ma & Tang, 2020). As such,
numerous GNNs have been proposed in the past that integrate ideas such as
maximal expressiveness (Xu et al., 2019), anisotropy and attention (Veličković
et al., 2018), non-linearities (Wang et al., 2019), or multiple aggregations
(Corso et al., 2020) into their message passing formulation. However, one of
the challenges that have so far precluded their wide adoption in industrial
and social applications is the difficulty to scale them to large graphs
(Frasca et al., 2020).

While the full-gradient in a GNN is straightforward to compute, assuming one
has access to _all_ hidden node embeddings in _all_ layers, this is not
feasible in large-scale graphs due to GPU memory limitations (Ma & Tang,
2020). Therefore, it is desirable to approximate its full-batch gradient
stochastically by considering only a mini-batch
$\mathcal{B}\subseteq\mathcal{V}$ of nodes for loss computation. However, this
stochastic gradient is still expensive to obtain due to the exponentially
increasing dependency of nodes over layers; a phenomenon framed as _neighbor
explosion_ (Hamilton et al., 2017). Due to neighbor explosion and since the
whole computation graph needs to be stored on the GPU, deeper architectures
can not be applied to large graphs. Therefore, a scalable solution needs to
make the memory consumption constant or sub-linear in respect to the number of
input nodes.

Recent works aim to alleviate this problem by proposing various sampling
techniques based on the concept of dropping edges (Ma & Tang, 2020; Rong et
al., 2020): _Node-wise sampling_ (Hamilton et al., 2017; Chen et al., 2018b;
Markowitz et al., 2021) recursively samples a fixed number of 1-hop neighbors;
_Layer-wise sampling_ techniques independently sample nodes for each layer,
leading to a constant sample size in each layer (Chen et al., 2018a; Zou et
al., 2019; Huang et al., 2018); In _subgraph sampling_ (Chiang et al., 2019;
Zeng et al., 2020b, a), a full GNN is run on an entire subgraph
$\mathcal{G}[\mathcal{B}]$ induced by a sampled batch of nodes
$\mathcal{B}\subseteq\mathcal{V}$. These techniques get rid of the neighbor
explosion problem by sampling the graph but may fail to preserve the edges
that present a meaningful topological structure. Further, existing approaches
are either still restricted to shallow networks, non-exchangeable GNN
operators or operators with reduced expressiveness. In particular, they
consider only specific GNN operators and it is an open question whether these
techniques can be successfully applied to the wide range of GNN architectures
available (Veličković et al., 2018; Xu et al., 2019; Corso et al., 2020; Chen
et al., 2020b). Another line of work is based on the idea of decoupling
propagations from predictions, either as a pre- (Wu et al., 2019; Klicpera et
al., 2019a; Frasca et al., 2020; Yu et al., 2020) or post-processing step
(Huang et al., 2021). While this scheme enjoys fast training and inference
time, it cannot be applied to any GNN, in particular because the propagation
is non-trainable, and therefore reduces model expressiveness. A different
scalability technique is based on the idea of training each GNN layer in
isolation (You et al., 2020). While this scheme resolves the neighbor
explosion problem and accounts for all edges, it cannot infer complex
interactions across consecutive layers.

(a) Mini-batch selection

(b) Original computation graph

(c) GAS computation graph

Figure 1: Mini-batch processing of GNNs with historical embeddings.
$\blacksquare$ denotes the nodes in the current mini-batch and $\blacksquare$
represents their direct 1-hop neighbors. For a given mini-batch (a), GPU
memory and computation costs exponentially increase with GNN depth (b). The
usage of historical embeddings avoids this problem as it allows to _prune_
entire sub-trees of the computation graph, which leads to constant GPU memory
consumption in respect to input node size (c). Here, nodes in the current
mini-batch _push_ their updated embeddings to the history
$\bm{\bar{H}}^{(\ell)}$, while their direct neighbors _pull_ their most recent
historical embeddings from $\bm{\bar{H}}^{(\ell)}$ for further processing.

Here, we propose the _G NNAutoScale_ (GAS) framework that disentangles the
scalability aspect of GNNs from their underlying message passing
implementation. GAS revisits and generalizes the idea of _historical_
embeddings (Chen et al., 2018b), which are defined as node embeddings acquired
in previous iterations of training, _cf._ Figure 1. For a given mini-batch of
nodes, GAS prunes the GNN computation graph so that only nodes inside the
current mini-batch and their direct 1-hop neighbors are retained,
_independent_ of GNN depth. Historical embeddings act as an offline storage
and are used to accurately fill in the inter-dependency information of out-of-
mini-batch nodes, _cf._ Figure 1(c). Through constant memory consumption in
respect to input node size, GAS is able to scale the training of GNNs to large
graphs, while still accounting for _all_ available neighborhood information.

We show that approximation errors induced by historical information are solely
caused by the staleness of the history and the Lipschitz continuity of the
learned function, and propose solutions for tightening the proven bounds in
practice. Furthermore, we connect scalability with expressiveness and
theoretically show under which conditions historical embeddings allow to learn
expressive node representations on large graphs. As a result, GAS is the first
scalable solution that is able to keep the existing expressivity properties of
the used GNN, which exist for a wide range of models (Xu et al., 2019; Morris
et al., 2019; Corso et al., 2020).

We implement our framework practically as _PyGAS_
111https://github.com/rusty1s/pyg_autoscale, an extension for the PyTorch
Geometric library (Fey & Lenssen, 2019), which makes it easy to convert common
and custom GNN models into their scalable variants and to apply them to large-
scale graphs. Experiments show that GNNs utilizing GAS achieve the same
performances as their (non-scalable) full-batch equivalents (while requiring
orders of magnitude less GPU memory), and are able to learn expressive node
representations. Furthermore, GAS allows the application of expressive and
hard-to-scale-up models on large graphs, leading to state-of-the-art results
on several large-scale graph benchmark datasets.

##  2 Scalable GNNs via Historical Embeddings

#### Background.

Let $\mathcal{G}=(\mathcal{V},\mathcal{E})$ or
$\bm{A}\in{\\{0,1\\}}^{|\mathcal{V}|\times|\mathcal{V}|}$ denote a _graph_
with node feature vectors $\bm{x}_{v}$ for all $v\in\mathcal{V}$. In this
work, we are mostly interested in the task of _node classification_ , where
each node $v\in\mathcal{V}$ is associated with a label $y_{v}$, and the goal
is to learn a representation $\bm{h}_{v}$ from which $y_{v}$ can be easily
predicted. To derive such a representation, GNNs follow a _neural message
passing scheme_ (Gilmer et al., 2017). Formally, the $(\ell+1)$-th layer of a
GNN is defined as (omitting edge features for simplicity)

$\displaystyle\leavevmode\resizebox{49.32227pt}{}{$\bm{h}^{(\ell+1)}_{v}$}$
$\displaystyle\leavevmode\resizebox{246.6185pt}{}{
$=\bm{f}^{(\ell+1)}_{\bm{\theta}}\left(\bm{h}_{v}^{(\ell)},{\left\\{\left\\{\bm{h}_{w}^{(\ell)}\right\\}\right\\}}_{w\in\mathcal{N}(v)}\right)$}$
(1) $\displaystyle\leavevmode\resizebox{375.80542pt}{}{
$=\textsc{Update}^{(\ell+1)}_{\bm{\theta}}\left(\bm{h}^{(\ell)}_{v},\,\bigoplus\limits_{\mathclap{w\in\mathcal{N}(v)}}~{}\textsc{Msg}^{(\ell+1)}_{\bm{\theta}}\big{(}\bm{h}_{w}^{(\ell)},\bm{h}_{v}^{(\ell)}\big{)}\right)$}$

where $\bm{h}^{(\ell)}_{v}$ represents the embedding of node $v$ obtained in
layer $\ell$ and $\mathcal{N}(v)$ defines the neighborhood set of $v$. We
initialize $\bm{h}^{(0)}_{v}=\bm{x}_{v}$. Here,
$\bm{f}_{\bm{\theta}}^{(\ell+1)}$ operates on _multisets_ $\\{\\{\ldots\\}\\}$
and can be decomposed into differentiable
$\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions parametrized by weights
$\bm{\theta}$, as well as permutation-invariant aggregation functions
$\bigoplus$, _e.g._ taking the sum, mean or maximum of features (Fey &
Lenssen, 2019; Gilmer et al., 2017; Qi et al., 2017; Wang et al., 2019; Xu et
al., 2019; Kipf & Welling, 2017; Veličković et al., 2018; Hamilton et al.,
2017; Klicpera et al., 2019a; Chen et al., 2020b; Xu et al., 2018). Our
following scalability framework is based on the general message passing
formulation given in Equation (1) and thus is applicable to this wide range of
different GNN operators.

#### Historical Embeddings.

Let $\bm{h}^{(\ell)}_{v}$ denote the node embedding in layer $\ell$ of a node
$v\in\mathcal{B}$ in a mini-batch $\mathcal{B}\subseteq\mathcal{V}$. For the
general message scheme given in Equation (1), the execution of
$\bm{f}^{(\ell+1)}_{\bm{\theta}}$ can be formulated as:

$\bm{h}_{v}^{(\ell+1)}=\bm{f}^{(\ell+1)}_{\bm{\theta}}\Big{(}\bm{h}_{v}^{(\ell)},{\left\\{\left\\{\bm{h}_{w}^{(\ell)}\right\\}\right\\}}_{w\in\mathcal{N}(v)}\Big{)}$
(2)
$=\bm{f}^{(\ell+1)}_{\bm{\theta}}\Big{(}\bm{h}_{v}^{(\ell)},{\left\\{\left\\{\bm{h}_{w}^{(\ell)}\right\\}\right\\}}_{w\in\mathcal{N}(v)\cap\mathcal{B}}\cup{\left\\{\left\\{\bm{h}_{w}^{(\ell)}\right\\}\right\\}}_{w\in\mathcal{N}(v)\setminus\mathcal{B}}\,\Big{)}$
$\approx\bm{f}^{(\ell+1)}_{\bm{\theta}}\Big{(}\bm{h}_{v}^{(\ell)},{\left\\{\left\\{\bm{h}_{w}^{(\ell)}\right\\}\right\\}}_{w\in\mathcal{N}(v)\cap\mathcal{B}}\cup\underbrace{{\left\\{\left\\{\bm{\bar{h}}_{w}^{(\ell)}\right\\}\right\\}}_{w\in\mathcal{N}(v)\setminus\mathcal{B}}}_{\textrm{\footnotesize
Historical embeddings}}\Big{)}$

Here, we separate the neighborhood information of the multiset into _two_
parts: (1) the local information of neighbors $\mathcal{N}(v)$ which are part
of the current mini-batch $\mathcal{B}$, and (2) the information of neighbors
which are not included in the current mini-batch. For out-of-mini-batch nodes,
we approximate their embeddings via historical embeddings acquired in previous
iterations of training (Chen et al., 2018b), denoted by
$\bm{\bar{h}}_{w}^{(\ell)}$. After each step of training, the newly computed
embeddings $\bm{h}^{(\ell+1)}_{v}$ are pushed to the history and serve as
historical embeddings $\bm{\bar{h}}_{w}^{(\ell+1)}$ in future iterations. The
separation of in-mini-batch nodes and out-of-mini-batch nodes, and their
approximation via historical embeddings represent the foundation of our GAS
framework.

A high-level illustration of its computation flow is visualized in Figure 1.
Figure 1(b) shows the original data flow without historical embeddings. The
required GPU memory increases as the model gets deeper. After a few layers,
embeddings for the entire input graph need to be stored, even if only a mini-
batch of nodes is considered for loss computation. In contrast, historical
embeddings eliminate this problem by approximating entire sub-trees of the
computation graph, _cf._ Figure 1(c). The required historical embeddings are
pulled from an offline storage, instead of being re-computed in each
iteration, which keeps the required information for each batch local. For a
single batch $\mathcal{B}\subseteq\mathcal{V}$, the GPU memory footprint for
one training step is given by
$\mathcal{O}(|\bigcup_{v\in\mathcal{B}}\mathcal{N}(v)\cup\\{v\\}|\cdot L)$ and
thus only scales linearly with the number of layers $L$. The majority of data
(the histories) can be stored in RAM or hard drive storage rather than GPU
memory.

In the following, we are going to use $\bm{\tilde{h}}_{v}^{(\ell)}$ to denote
embeddings estimated via GAS (line 3 of Equation (2)) to differentiate them
from the exact embeddings obtained without historical approximation (line 1 of
Equation (2)). In contrast to existing scaling solutions based on sub-sampling
edges, the usage of historical embeddings as utilized in GAS provides the
following additional advantages:

(1) GAS trains over all the data: In GAS, a GNN will make use of all available
graph information, _i.e._  _no_ edges are dropped, which results in lower
variance and more accurate estimations (since
$\|\bm{\bar{h}}^{(\ell)}_{v}-\bm{h}_{v}^{(\ell)}\|\ll\|\bm{h}_{v}^{(\ell)}\|$).
Importantly, for a single epoch and layer, each edge is still only processed
once, putting its time complexity $\mathcal{O}(|\mathcal{E}|)$ on par with its
full-batch counterpart. Notably, more accurate estimations will further
strengthen gradient estimation during backpropagation. Specifically, the model
parameters will be updated based on the node embeddings of _all_ neighbors
since $\partial\bm{\tilde{h}}^{(\ell+1)}_{v}/\partial\bm{\theta}$ also depends
on $\\{\\{\bm{\bar{h}}_{w}^{(\ell)}\colon
w\in\mathcal{N}(v)\setminus\mathcal{B}\\}\\}$.

(2) GAS enables constant inference time complexity: The time complexity of
model inference is reduced to a constant factor, since we can directly use the
historical embeddings of the last layer to derive predictions for test nodes.

(3) GAS is simple to implement: Our scheme does not need to maintain recursive
layer-wise computation graphs, which makes its overall implementation
straightforward and comparable to full-batch training. Only minor
modifications are required to _pull_ information from and _push_ information
to the histories, _cf._ our training algorithm in the appendix.

(4) GAS provides theoretical guarantees: In particular, if the model weights
are kept fixed, $\bm{\tilde{h}}_{v}^{(\ell)}$ eventually equals
$\bm{h}^{(\ell)}_{v}$ after a fixed amount of iterations (Chen et al., 2018b).

##  3 Approximation Error and Expressiveness

The advantages of utilizing historical embeddings $\bm{\bar{h}}_{v}^{(\ell)}$
to compute an approximation $\bm{\tilde{h}}_{v}^{(\ell)}$ of the exact
embedding $\bm{h}_{v}^{(\ell)}$ come at the cost of an approximation error
$\|\bm{\tilde{h}}_{v}^{(\ell)}-\bm{h}_{v}^{(\ell)}\|$, which can be decomposed
into two sources of variance: (1) The _closeness_ of estimated inputs to their
exact values, _i.e._
$\|\bm{\tilde{h}}_{v}^{(\ell-1)}-\bm{h}_{v}^{(\ell-1)}\|\geq 0$, and (2) the
_staleness_ of historical embeddings, _i.e._
$\|\bm{\bar{h}}_{v}^{(\ell-1)}-\bm{\tilde{h}}^{(\ell-1)}_{v}\|\geq 0$. In the
following, we show concrete bounds for this error, which can be then tightened
using specific procedures. Here, our analysis focuses on arbitrary
$\bm{f}^{(\ell)}_{\bm{\theta}}$ GNN layers as described in Equation (1), but
we restrict both $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ to model $k$-Lipschitz continuous
functions due to their potentially highly non-linear nature. Proofs of all
lemmas and theorems can be found in the appendix.

######  Lemma 1.

Let $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ be Lipschitz continuous functions
with Lipschitz constants $k_{1}$ and $k_{2}$, respectively. If, for all
$v\in\mathcal{V}$, the inputs are close to the exact input, _i.e._
$\|\bm{\tilde{h}}^{(\ell-1)}_{v}-\bm{h}^{(\ell-1)}_{v}\|\leq\delta$, and the
historical embeddings do not run too stale, _i.e._
$\|\bm{\bar{h}}^{(\ell-1)}_{v}-\bm{\tilde{h}}^{(\ell-1)}_{v}\|\leq\epsilon$,
then the output error is bounded by

$\|\bm{\tilde{h}}^{(\ell)}_{v}-\bm{h}^{(\ell)}_{v}\|\leq\delta\,k_{2}+(\delta+\epsilon)\,k_{1}\,k_{2}\,|\mathcal{N}(v)|.$

Due to the behavior of Lipschitz constants in a series of function
compositions, we obtain an upper bound that is dependent on $k_{1}$, $k_{2}$
and $|\mathcal{N}(v)|$, as well as dependent on the errors $\delta$ and
$\epsilon$ of the inputs. Interestingly, sum aggregation, the most expressive
aggregation function (Xu et al., 2019), introduces a factor of
$|\mathcal{N}(v)|$ to the upper bound, while we can obtain a much tighter
upper bound for mean or max aggregation, _cf._ its proof. Next, we take a look
at the final output error produced by a $L$-layered GNN:

######  Theorem 2.

Let $\bm{f}^{(L)}_{\bm{\theta}}$ be a $L$-layered GNN, containing only
Lipschitz continuous $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions with Lipschitz constants
$k_{1}$ and $k_{2}$, respectively. If, for all $v\in\mathcal{V}$ and all
$\ell\in\\{1,\ldots,L-1\\}$, the historical embeddings do not run too stale,
_i.e._
$\|\bm{\bar{h}}^{(\ell)}_{v}-\bm{\tilde{h}}^{(\ell)}_{v}\|\leq\epsilon^{(\ell)}$,
then the final output error is bounded by

$\|\bm{\tilde{h}}_{v,j}^{(L)}-\bm{h}_{v,j}^{(L)}\|\leq\sum_{\ell=1}^{L-1}\epsilon^{(\ell)}\,k_{1}^{L-\ell}\,k_{2}^{L-\ell}\,{|\mathcal{N}(v)|}^{L-\ell}.$

Notably, this upper bound does not longer depend on
$\|\bm{\tilde{h}}_{v}^{(\ell)}-\bm{h}_{v}^{(\ell)}\|\leq\delta^{(\ell)}$, and
is instead solely conditioned on the staleness of histories
$\|\bm{\bar{h}}_{v}^{(\ell)}-\bm{\tilde{h}}_{v}^{(\ell)}\|\leq\epsilon^{(\ell)}$.
However, it depends _exponentially_ on the Lipschitz constants $k_{1}$ and
$k_{2}$ as well as $|\mathcal{N}(v)|$ with respect to the number of layers. In
particular, each additional layer introduces a less restrictive bound since
the errors made in the first layers get immediately propagated to later ones,
leading to potentially high inaccuracies for histories in deeper GNNs. We will
later propose solutions for tightening the proven bound in practice, allowing
the application of GAS to deep and non-linear GNNs. Furthermore, Theorem 2
lets us immediately derive an upper error bound of gradients as well, _i.e._

$\|\nabla_{\bm{\theta}}\mathcal{L}(\bm{\tilde{h}}^{(L)}_{v})-\nabla_{\bm{\theta}}\mathcal{L}(\bm{h}^{(L)}_{v})\|\leq\lambda\|\bm{\tilde{h}}^{(L)}_{v}-\bm{h}^{(L)}_{v}\|$

in case $\mathcal{L}$ is $\lambda$-Lipschitz continuous. As such, GAS
encourages low variance and bias in the learning signal as well. However,
parameters are not guaranteed to converge to the same optimum since we
explicitely consider arbitrary GNNs solving non-convex problems (Cong et al.,
2020).

It is well known that the most powerful GNNs adhere to the same
representational power as the _Weisfeiler-Lehman (WL) test_ (Weisfeiler &
Lehman, 1968) in distinguishing non-isomorphic structures, _i.e._
$\bm{h}_{v}^{(L)}\neq\bm{h}_{w}^{(L)}$ in case $c_{v}^{(L)}\neq c_{w}^{(L)}$
(Xu et al., 2019; Morris et al., 2019), where $c_{v}^{(L)}$ denotes a node’s
coloring after $L$ rounds of color refinement. However, in order to leverage
such expressiveness, a GNN needs to be able to reason about structural
differences across neighborhoods directly _during_ training. We now show that
GNNs that scale by sampling edges are not capable of doing so:

######  Proposition 3.

Let $\bm{f}^{(L)}_{\bm{\theta}}\colon\mathcal{V}\to\mathbb{R}^{d}$ be a
$L$-layered GNN as expressive as the WL test in distinguishing the $L$-hop
neighborhood around each node $v\in\mathcal{V}$. Then, there exists a graph
$\bm{A}\in{\\{0,1\\}}^{|\mathcal{V}|\times|\mathcal{V}|}$ for which
$\bm{f}^{(L)}_{\bm{\theta}}$ operating on a sampled variant $\bm{\tilde{A}}$,
$\tilde{a}_{v,w}=\begin{cases}\frac{|\mathcal{N}(v)|}{|\mathcal{\tilde{N}}(v)|},&\textrm{if
}w\in\mathcal{\tilde{N}}(v)\\\ 0,&\textrm{otherwise}\end{cases}$, produces a
non-equivalent coloring, _i.e._
$\bm{\tilde{h}}^{(L)}_{v}\neq\bm{\tilde{h}}^{(L)}_{w}$ while
$c_{v}^{(L)}=c_{w}^{(L)}$ for nodes $v,w\in\mathcal{V}$.

While sampling strategies lose expressive power due to sub-sampling of edges,
scalable GNNs based on historical embeddings are leveraging _all_ edges during
neighborhood aggregation. Therefore, a special interest lies in the question
if historical-based GNNs are as expressive as their full-batch counterpart.
Here, a maximally powerful _and_ scalable GNN needs to fulfill the following
two requirements: (1) It needs to be as expressive as the WL test in
distinguishing non-isomorphic structures, and (2) it needs to account for the
approximation error $\|\bm{\bar{h}}_{v}^{(\ell-1)}-\bm{h}_{v}^{(\ell-1)}\|$
induced by the usage of historical embeddings. Since it is known that there
exists a wide range of maximally powerful GNNs (Xu et al., 2019; Morris et
al., 2019; Corso et al., 2020), we can restrict our analysis to the latter
question. Following upon Xu et al. (2019), we focus on the case where input
node features are from a countable set $\mathbb{P}^{d}\subset\mathbb{R}^{d}$
of bounded size:

######  Lemma 4.

Let $\\{\\{\bm{h}_{v}^{(\ell-1)}\colon v\in\mathcal{V}\\}\\}$ be a countable
multiset such that
$\|\bm{h}_{v}^{(\ell-1)}-\bm{h}_{w}^{(\ell-1)}\|>2(\delta+\epsilon)$ for all
$v,w\in\mathcal{V}$, $\bm{h}_{v}^{(\ell-1)}\neq\bm{h}_{w}^{(\ell-1)}$. If the
inputs are close to the exact input, _i.e._
$\|\bm{\tilde{h}}_{v}^{(\ell-1)}-\bm{h}_{v}^{(\ell-1)}\|\leq\delta$, and the
historical embeddings do not run too stale, _i.e._
$\|\bm{\bar{h}}_{v}^{(\ell-1)}-\bm{\tilde{h}}_{v}^{(\ell-1)}\|\leq\epsilon$,
then there exist $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions, such that

$\|\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{\tilde{h}}_{v}^{(\ell-1)})-\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{h}_{v}^{(\ell-1)})\|\leq\delta+\epsilon$

and

$\|\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{h}_{v}^{(\ell-1)})-\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{h}_{w}^{(\ell-1)})\|>2(\delta+\epsilon+\lambda)$

for all $v,w\in\mathcal{V}$, $\bm{h}_{v}^{(\ell-1)}\neq\bm{h}_{w}^{(\ell-1)}$
and all $\lambda>0$.

Informally, Lemma 4 tells us that if (1) exact input embeddings are
sufficiently far apart from each other and (2) historical embeddings are
sufficiently close to the exact embeddings, there exist historical-based GNN
operators which can distinguish equal from non-equal inputs. Key to the proof
is that $(\delta+\epsilon)$-balls around exact inputs do not intersect each
other and are therefore well separated. Notably, we do not require
$\bm{f}_{\bm{\theta}}^{(\ell)}$ to model strict injectivity since it is
sufficient for $\bm{f}_{\bm{\theta}}^{(\ell)}$ to be
$2(\delta+\epsilon)$-injective (Seo et al., 2019).

Following Xu et al. (2019), one can leverage MLPs to model and learn such
Message and Update functions due to the universal approximation theorem
(Hornik et al., 1989; Hornik, 1991). However, the theory behind Lemma 4 holds
for any maximally powerful GNN operator. Finally, we can use this insight to
relate the expressiveness of scalable GNNs to the WL test color refinement
procedure:

######  Theorem 5.

Let $\bm{f}^{(L)}_{\bm{\theta}}$ be a $L$-layered GNN in which all
$\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions fulfill the conditions of
Lemma 4. Then, there exists a map $\phi\colon\mathbb{R}^{d}\to\Sigma$ so that
$\phi(\bm{\tilde{h}}^{(L)}_{v})=c^{(L)}_{v}$ for all $v\in\mathcal{V}$.

Theorem 5 extends the insights of Lemma 4 to multi-layered GNNs, and indicates
that scalable GNNs using historical embeddings are able to distinguish non-
isomorphic structures (that are distinguishable by the WL test) directly
during training, which is what makes reasoning about structural properties
possible. It should be noted that recent proposals such as DropEdge (Rong et
al., 2020) are still applicable for data augmentation and message reduction.
However, through the given theorem, we disentangle scalability and
expressiveness from regularization via edge dropping.

While sampling approaches lose expressiveness compared to their original
counterparts (_cf._ Proposition 3), Theorem 5 tells us that, in theory, there
exist message passing functions that are as expressive as the WL test in
distinguishing non-isomorphic structures while accounting for the effects of
approximation in stored embeddings. In practice, we have two degrees of
freedom to tighten the upper bounds given by Lemma 1 and Theorem 2, leading to
a lower approximation error and higher expressiveness in return: (1)
Minimizing the _staleness_ of historical embeddings, and (2) maximizing the
_closeness_ of estimated inputs to their exact values by controlling the
Lipschitz constants of Update and Message functions. In what follows, we
derive a list of procedures to achieve these goals:

#### Minimizing Inter-Connectivity Between Batches.

As formulated in Equation (2) in Section 2, the output embeddings of
$\bm{f}_{\bm{\theta}}^{(\ell+1)}$ are exact if
$|\bigcup_{v\in\mathcal{B}}\mathcal{N}(v)\cup\\{v\\}|=|\mathcal{B}|$, _i.e._
all neighbors of nodes in $\mathcal{B}$ are as well part of $\mathcal{B}$.
However, in practice, this can only be guaranteed for full-batch GNNs.
Motivated by this observation, we aim to minimize the inter-connectivity
between sampled mini-batches, _i.e._
$\min|\bigcup_{v\in\mathcal{B}}\mathcal{N}(v)\setminus\mathcal{B}|$, which
minimizes history access, and increases closeness and reduces staleness in
return.

Similar to Cluster-GCN (Chiang et al., 2019), we make use of graph clustering
techniques, _e.g._ , Metis (Karypis & Kumar, 1998; Dhillon et al., 2007), to
achieve this goal. It aims to construct partitions over the nodes in a graph
such that intra-links within clusters occur much more frequently than inter-
links between different clusters. Intuitively, this results in a high chance
that neighbors of a node are located in the same cluster. Notably, modern
graph clustering methods are both fast and scalable with time complexities
given by $\mathcal{O}(|\mathcal{E}|)$, and only need to be applied once, which
leads to an unremarkable computational overhead in the pre-processing stage.
In general, we argue that the Metis clustering technique is highly scalable,
as it is in the heart of many large-scale distributed graph storage layers
such as (Zhu et al., 2019; Zheng et al., 2020) that are known scale to
billion-sized graphs. Furthermore, the additional overhead in the pre-
processing stage is quickly compensated by an acceleration of training, since
the number of neighbors outside of $\mathcal{B}$ is heavily reduced, and
pushing information to the histories now leads to contiguous memory transfers.

#### Enforcing Local Lipschitz Continuity.

To guide our neural network in learning a function with controllable error, we
can enforce its intermediate output layers $\bm{f}_{\bm{\theta}}^{(\ell)}$ to
be invariant to small input perturbations. In particular, following upon Usama
& Chang (2018), we found it useful to apply the auxiliary loss

$\mathcal{L}_{\textrm{reg}}^{(\ell)}=\|\bm{f}_{\bm{\theta}}^{(\ell)}(\bm{\tilde{h}}_{v}^{(\ell-1)})-\bm{f}_{\bm{\theta}}^{(\ell)}(\bm{\tilde{h}}_{v}^{(\ell-1)}+\bm{\epsilon})\|$
(3)

in highly non-linear message passing phases, _e.g._ , in GIN (Xu et al.,
2019). Such regularization enforces equal outputs for small pertubations
$\bm{\epsilon}\sim\mathcal{B}_{\delta}(\bm{0})$ inside closed balls of radius
$\delta$. Notably, we do not restrict $\textsc{Update}^{(\ell)}_{\bm{\theta}}$
and $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ to separately model global
$k$-Lipschitz continuous functions, but rather aim for local Lipschitz
continuity at each $\bm{h}_{v}^{(\ell-1)}$ for $\bm{f}^{(\ell)}_{\bm{\theta}}$
as a whole. For other message passing GNNs, _e.g._ , in GCN (Kipf & Welling,
2017), $L_{2}$ regularization is usually sufficient to ensure closeness of
historical embeddings. Further, we found gradient clipping to be an effective
method to restrict the parameters from changing too fast, regularizing history
changes in return.

##  4 Related Work

Our GAS framework utilizes historical embeddings as an affordable
approximation. The idea of historical embeddings was originally introduced in
VR-GCN (Chen et al., 2018b). VR-GCN aims to reduce the variance in estimation
during neighbor sampling (Hamilton et al., 2017), and avoids the need to
sample a large amount of neighbors in return. Cong et al. (2020) further
simplified this scheme into a _one-shot sampling_ scenario, where nodes no
longer need to recursively explore neighborhoods in each layer. However, these
approaches consider only a specific GNN operator which prevent their
application to the wide range of GNN architectures available. Furthermore,
they only consider shallow architectures and do not account for the increasing
approximation error induced by deeper and expressive GNNs, which is well
observable in practice, _cf._ Section 6.1.

In order to minimize the inter-connectivity between mini-batches, we utilize
graph clustering techniques for mini-batch selection, as first introduced in
the subgraph sampling approach Cluster-GCN (Chiang et al., 2019). Cluster-GCN
leverages clustering in order to infer meaningful subgraphs, while we aim to
minimize history accesses. Furthermore, Cluster-GCN limits message passing to
intra-connected nodes, and therefore ignores potentially useful information
outside the current mini-batch. This inherently limits the model to learn from
nodes nearby. In contrast, our GAS framework makes use of _all_ available
neighborhood data for aggregation, and therefore avoids this downside.

##  5 PyGAS: Auto-Scaling GNNs in PyG

Figure 2: Illustrative runtime performances of a serial and concurrent mini-
batch execution in comparison to a full-batch GNN execution. In the full-batch
approach (a), all necessary data is first transferred to the device via the
Host2Device (H2D) engine, before GNN layers are executed in serial inside the
kernel engine. As depicted in (b), a serial mini-batch execution suffers from
an I/O bottleneck, in particular because each kernel engine has to wait for
memory transfers to complete. The concurrent mini-batch execution (c) avoids
this problem by leveraging an additional worker thread and overlapping data
transfers, leading to two times performance improvements in comparison to a
serial execution, which is on par with the standard full-batch approach.

(a) 2-GCN on Cora

(b) 64-GCNII on Cora

(c) 4-GIN on Cluster

Figure 3: Model performance comparison between full-batch, an unoptimized
history-based baseline and our GAS approach. In contrast to the historical-
based baseline, GAS reaches the quality of full-batch training, especially for
(b) deep and (c) expressive models.

We condense our GAS framework and theoretical findings into a tool named
_PyGAS_ that implements all the presented techniques in
practice.222https://github.com/rusty1s/pyg_autoscale _PyGAS_ is built upon
PyTorch (Paszke et al., 2019) and utilizes the PyTorch Geometric (PyG) library
(Fey & Lenssen, 2019). It provides an easy-to-use interface to convert common
and custom GNN models from PyTorch Geometric into their scalable variants.
Furthermore, it provides a fully deterministic test bed for evaluating models
on large-scale graphs. An example of the interface is shown in the appendix.

#### Fast Historical Embeddings.

Our approach accesses histories to account for any data outside the current
mini-batch, which requires frequent data transfers to and from the GPU.
Therefore, _PyGAS_ optimizes pulling from and pushing to histories via _non-
blocking_ device transfers. Specifically, we immediately start pulling
historical embeddings for each layer asynchronously at the beginning of each
optimization step, which ensures that GPUs do not run idle while waiting for
memory transfers to complete. A separate worker thread gathers historical
information into one of multiple pinned CPU memory buffers (denoted by Pull),
from where it can be transfered to the GPU via the usage of CUDA streams
without blocking any CPU or CUDA execution. Synchronization is done by
synchronizing the respective CUDA stream before inputting the transferred data
into the GNN layer. The same strategy is applied for pushing information to
the history. Considering that the device transfer of $\bm{\bar{H}}^{(\ell-1)}$
is faster than the execution of $\bm{f}_{\bm{\theta}}^{(\ell)}$, this scheme
does not lead to any runtime overhead when leveraging historical embeddings
and can be twice as fast as its serial non-overlapping counterpart, _cf._
Figure 2. We have implemented our non-blocking transfer scheme with custom
C++/CUDA code to avoid Python’s global interpreter lock.

##  6 Experiments

In this section, we evaluate our GAS framework in practice using _PyGAS_ ,
utilizing 6 different GNN operators and 15 datasets. Please refer to the
appendix for a detailed description of the used GNN operators and datasets,
and to our code for hyperparameter configurations. All models were trained on
a single GeForce RTX 2080 Ti (11 GB). In our experiments, we hold all
histories in RAM, using a machine with 64GB of CPU memory.

###  6.1 GAS resembles full-batch performance

Table 1: Full-batch vs GAS performance on small transductive graph benchmark
datasets across 20 different initializations. Predictive performance of models
trained via GAS closely matches those of full-batch gradient descent on all
models for all datasets.

$\dagger$ Results omitted due to unstable performance across different weight
initializations, _cf._ Shchur et al. (2018)
---
Dataset | GCN | GAT | APPNP | GCNII
Full | GAS | Full | GAS | Full | GAS | Full | GAS
Cora | 81.88$\pm$0.75 | 82.29$\pm$0.76 | 82.80$\pm$0.47 | 83.32$\pm$0.62 | 83.28$\pm$0.60 | 83.19$\pm$0.58 | 85.04$\pm$0.53 | 85.52$\pm$0.39
CiteSeer | 70.98$\pm$0.66 | 71.18$\pm$0.97 | 71.72$\pm$0.91 | 71.86$\pm$1.00 | 72.13$\pm$0.73 | 72.63$\pm$0.82 | 73.06$\pm$0.81 | 73.89$\pm$0.48
PubMed | 78.73$\pm$1.10 | 79.23$\pm$0.62 | 78.03$\pm$0.40 | 78.42$\pm$0.56 | 80.21$\pm$0.20 | 79.82$\pm$0.52 | 79.72$\pm$0.78 | 80.19$\pm$0.49
Coauthor-CS | 91.08$\pm$0.59 | 91.22$\pm$0.45 | 90.31$\pm$0.49 | 90.38$\pm$0.42 | 92.51$\pm$0.47 | 92.44$\pm$0.58 | 92.45$\pm$0.35 | 92.52$\pm$0.31
Coauthor-Physics | 93.10$\pm$0.84 | 92.98$\pm$0.72 | 92.32$\pm$0.86 | 92.80$\pm$0.61 | 93.40$\pm$0.92 | 93.68$\pm$0.61 | 93.43$\pm$0.52 | 93.61$\pm$0.41
Amazon-Computer | 81.17$\pm$1.81 | 80.84$\pm$2.26 | —† | —† | 81.79$\pm$2.00 | 81.66$\pm$1.81 | 83.04$\pm$1.81 | 83.05$\pm$1.16
Amazon-Photo | 90.25$\pm$1.66 | 90.53$\pm$1.40 | —† | —† | 91.27$\pm$1.26 | 91.23$\pm$1.34 | 91.42$\pm$0.81 | 91.60$\pm$0.78
Wiki-CS | 79.08$\pm$0.50 | 79.00$\pm$0.41 | 79.44$\pm$0.41 | 79.56$\pm$0.47 | 79.88$\pm$0.40 | 79.75$\pm$0.53 | 79.94$\pm$0.67 | 80.02$\pm$0.43
$\bm{\Delta}$ Mean Accuracy | +0.13 | +0.29 | -0.01 | +0.29

First, we analyze how GAS affects the robustness and expressiveness of our
method. We compare GAS against two different baselines: a regular full-batch
variant and a history baseline, which naively integrates history-based mini-
batch training without any of the additional GAS techniques. To evaluate, we
make use of a shallow 2-layer GCN (Kipf & Welling, 2017) and two recent state-
of-the-art models: a deep GCNII network with 64 layers (Chen et al., 2020b),
and a maximally expressive GIN network with 4 layers (Xu et al., 2019). We
evaluate those models on tasks for which they are well suitable: classifying
academic papers in a citation network (Cora), and identifying community
clusters in Stochastic Block Models (Cluster) (Yang et al., 2016; Dwivedi et
al., 2020), _cf._ Figure 3. Since Cluster is a node classification task
containing multiple graphs, we first convert it into a super graph (holding
all the nodes of all graphs), and partition this super graph using twice as
many partitions as there are initial graphs. It can be seen that especially
for deep (64-GCNII, _cf._ Figure 3(b)) and expressive (4-GIN, _cf._ Figure
3(c)) architectures, the naive historical-based baseline fails to reach the
desired full-batch performance. This can be contributed to the high
approximation error induced by deep and expressive models. In contrast, GAS
shows far superior performance, reaching the quality of full-batch training in
both cases.

In general, we expect the model performances of our GAS mini-batch training to
closely resemble the performances of their full-batch counterparts, except for
the variance introduced by stochastic optimization (which is, in fact, known
to improve generalization (Bottoue & Bousquet, 2007)). To validate, we compare
our approach against full-batch performances on small transductive benchmark
datasets for which full-batch training is easily feasible. We evaluate on four
GNN models that significantly advanced the field of graph representation
learning: GCN (Kipf & Welling, 2017), GAT (Veličković et al., 2018), APPNP
(Klicpera et al., 2019a) and GCNII (Chen et al., 2020b). For all experiments,
we tried to follow the hyperparameter setup of the respective papers as
closely as possible and perform an in-depth grid search on datasets for which
best performing configurations are not known. We then apply GAS mini-batch
training on the _same_ set of hyperparameters. As shown in Table 1, all models
that utilize GAS training perform as well as their full-batch equivalents
(with slight gains overall), confirming the practical effectiveness of our
approach. Notably, even for deep GNNs such as APPNP and GCNII, our approach is
able to closely resemble the desired performance.

Table 2: Relative performance improvements of individual GAS techniques within
a GCNII model. The performance improvement is measured in percentage points in
relation to the corresponding model performance obtained by full-batch
training.

| Cora | CiteSeer | PubMed | Coauthor- | Amazon- | Wiki-CS
---|---|---|---|---|---|---
| CS | Physics | Computer | Photo
Baseline | -3.26 | -5.66 | -3.20 | -0.79 | -0.50 | -5.76 | -4.16 | -3.19
Regularization | -2.12 | -1.03 | -1.24 | -0.46 | -0.24 | -3.02 | -1.19 | -0.74
Metis | -1.57 | -3.12 | -1.50 | -0.47 | +0.13 | -2.75 | -1.02 | -0.24
GAS | +0.48 | +0.83 | +0.47 | +0.07 | +0.18 | +0.01 | +0.18 | +0.08

We further conduct an ablation study to highlight the individual performance
improvements of our GAS techniques within a GCNII model, _i.e._ minimizing
inter-connectivity and applying regularization techniques. Table 2 shows the
relative performance improvements of individual GAS techniques in percentage
points, compared to the corresponding model performance obtained by full-batch
training. Notably, it can be seen that both techniques contribute to
resembling full-batch performance, reaching their full strength when used in
combination. We include an additional ablation study for training an
expressive GIN model in the appendix.

###  6.2 GAS is fast and memory-efficient

Table 3: GPU memory consumption (in GB) and the amount of data used ($\%$)
across different GNN execution techniques. GAS consumes low memory while
making use of all available neighborhood information during a single
optimization step.

| # nodes | 717K | 169K | 2.4M
---|---|---|---|---
| # edges | 7.9M | 1.2M | 61.9M
| Method | Yelp | ogbn- | ogbn-
| arxiv | products
2-layer | Full-batch | 6.64GB/ | 100% | 1.44GB/ | 100% | 21.96GB/ | 100%
GraphSAGE | 0.76GB/ | 9% | 0.40GB/ | 27% | 0.92GB/ | 2%
Cluster-GCN | 0.17GB/ | 13% | 0.15GB/ | 40% | 0.16GB/ | 16%
GAS | 0.51GB/ | 100% | 0.22GB/ | 100% | 0.36GB/ | 100%
3-layer | Full-batch | 9.44GB/ | 100% | 2.11GB/ | 100% | 31.53GB/ | 100%
GraphSAGE | 2.19GB/ | 14% | 0.93GB/ | 33% | 4.34GB/ | 5%
Cluster-GCN | 0.23GB/ | 13% | 0.22GB/ | 40% | 0.23GB/ | 16%
GAS | 0.79GB/ | 100% | 0.34GB/ | 100% | 0.59GB/ | 100%
4-layer | Full-batch | 12.24GB/ | 100% | 2.77GB/ | 100% | 41.10GB/ | 100%
GraphSAGE | 4.31GB/ | 19% | 1.55GB/ | 37% | 11.23GB/ | 8%
Cluster-GCN | 0.30GB/ | 13% | 0.29GB/ | 40% | 0.29GB/ | 16%
GAS | 1.07GB/ | 100% | 0.46GB/ | 100% | 0.82GB/ | 100%

For training large-scale GNNs, GPU memory consumption will directly dictate
the scalability of the given approach. Here, we show how GAS maintains a low
GPU memory footprint while, in contrast to other scalability approaches,
accounts for _all_ available information inside a GNN’s receptive field in a
single optimization step. We compare the memory usage of GCN+GAS training with
the memory usage of full-batch GCN, and mini-batch GraphSAGE (Hamilton et al.,
2017) and Cluster-GCN (Chiang et al., 2019) training, _cf._ Table 3. Notably,
GAS is easily able to fit the required data on the GPU, while memory
consumption only increases linearly with the number of layers. Although
Cluster-GCN maintains an overall lower memory footprint than GAS, it will only
utilize a fraction of available information inside its receptive field, _i.e._
$\approx$23% on average.

Figure 4: Runtime overhead in relation to the inter-/intra-connectivity ratio
of mini-batches, both for serial and concurrent history access patterns. The
overall runtime overhead is further separated into computational overhead
(overhead of aggregating additional messages) and I/O overhead (overhead of
pulling from and pushing to histories). Our concurrent memory transfer reduces
I/O overhead caused by histories by a wide margin.

We now analyze how GAS enables large-scale training due to fast mini-batch
execution. Specifically, we are interested in how our concurrent memory
transfer scheme (_cf._ Section 5) reduces the overhead induced by accessing
historical embeddings from the offline storage. For this, we evaluate runtimes
of a $4$-layer GIN model on synthetic graph data, which allows fine-grained
control over the ratio between inter- and intra-connected nodes, _cf._ Figure
4. Here, a given mini-batch consists of exactly 4,000 nodes which are randomly
intra-connected to $60$ other nodes. We vary the number of inter-connections
(connections to nodes outside of the batch) by adding out-of-batch nodes that
are randomly inter-connected to $60$ nodes inside the batch. Notably, the
naive serial memory transfer increases runtimes up to 350%, which indicates
that frequent history accesses can cause major I/O bottlenecks. In contrast,
our concurrent access pattern incurs _almost no_ I/O overhead _at all_ , and
the overhead in execution time is solely explained by the computational
overhead of aggregating far more messages during message propagation. Note
that in most real-world scenarios, the additional aggregation of history data
may only increase runtimes up to 25%, since most real-world datasets contain
inter-/intra-connectivity ratios between $0.1$ and $2.5$, _cf._ appendix.
Further, the additional overhead of computing Metis partitions in the pre-
processing stage is negligible and is quickly mitigated by faster training
times: Computing the partitioning of a graph with 2M nodes takes only about
20–50 seconds (depending on the number of clusters).

Next, we compare runtimes and memory consumption of GAS to the recent GTTF
proposal (Markowitz et al., 2021), which utilizes a fast neighbor sampling
strategy based on tensor functionals. For this, we make use of a 4-layered GCN
model with equal mini-batch and receptive field sizes. As shown in Table 4,
GAS is both faster and consumes less memory than GTTF. Although GTTF makes use
of a fast vectorized sampling procedure, its underlying recursive neighborhood
construction still scales _exponentially_ with GNN depth, which explains the
observable differences in runtime and memory consumption.

Table 4: Efficiency of GCN with GTTF and GAS.

Dataset |  Runtime (s) |  Memory (MB)
---|---|---
GTTF | GAS | GTTF | GAS
Cora | 0.077 | 0.006 | 18.01 | 2.13
PubMed | 0.071 | 0.006 | 28.79 | 2.19
PPI | 0.976 | 0.007 | 134.86 | 12.37
Flickr | 1.178 | 0.007 | 325.97 | 16.32

###  6.3 GAS scales to large graphs

In order to demonstrate the scalability and generality of our approach, we
scale various GNN operators to common large-scale graph benchmark datasets.
Here, we focus our analysis on GNNs that are notorious hard to scale-up but
have the potential to leverage the increased amount of available data to make
more accurate predictions. In particular, we benchmark deep GNNs, _i.e._ GCNII
(Chen et al., 2020b), and expressive GNNs, _i.e._ PNA (Corso et al., 2020).
Note that it is not possible to run those models in full-batch mode on most of
these datasets as they will run out of memory on common GPUs. We compare with
10 scalable GNN baselines: GraphSAGE (Hamilton et al., 2017), FastGCN (Chen et
al., 2018a), LADIES (Zou et al., 2019), VR-GCN (Chen et al., 2018b), MVS-GNN
(Cong et al., 2020), Cluster-GCN (Chiang et al., 2019), GraphSAINT (Zeng et
al., 2020b), SGC (Wu et al., 2019), SIGN (Frasca et al., 2020) and GBP (Chen
et al., 2020a). Since results are hard to compare across different approaches
due to differences in frameworks, model implementations, weight
initializations and optimizers, we additionally report a shallow GCN+GAS
baseline. GAS is able to train all models on all datasets on a single GPU,
while holding corresponding histories in CPU memory. On the largest dataset,
_i.e._ ogbn-products, this will consume $\approx L\cdot$ 2GB of storage for
$L$ layers, which easily fits in RAM on most modern workstations.

Table 5: Performance on large graph datasets. GAS is both scalable and general
while achieving state-of-the-art performance.

# nodes | 230K | 57K | 89K | 717K | 169K | 2.4M
---|---|---|---|---|---|---
# edges | 11.6M | 794K | 450K | 7.9M | 1.2M | 61.9M
Method | Reddit | PPI | Flickr | Yelp | ogbn- | ogbn-
arxiv | products
GraphSAGE | 95.40 | 61.20 | 50.10 | 63.40 | 71.49 | 78.70
FastGCN | 93.70 | — | 50.40 | — | — | —
LADIES | 92.80 | — | — | — | — | —
VR-GCN | 94.50 | 85.60 | — | 61.50 | — | —
MVS-GNN | 94.90 | 89.20 | — | 62.00 | — | —
Cluster-GCN | 96.60 | 99.36 | 48.10 | 60.90 | — | 78.97
GraphSAINT | 97.00 | 99.50 | 51.10 | 65.30 | — | 79.08
SGC | 96.40 | 96.30 | 48.20 | 64.00 | — | —
SIGN | 96.80 | 97.00 | 51.40 | 63.10 | — | 77.60
GBP | — | 99.30 | — | 65.40 | — | —
Full-batch |  |  |  |  |  |  |
GCN | 95.43 | 97.58 | 53.73 | OOM | 71.64 | OOM
GCNII | OOM | OOM | 55.28 | OOM | 72.83 | OOM
PNA | OOM | OOM | 56.23 | OOM | 72.17 | OOM
GAS | GCN | 95.45 | 98.92 | 54.00 | 62.94 | 71.68 | 76.66
GCNII | 96.77 | 99.50 | 56.20 | 65.14 | 73.00 | 77.24
PNA | 97.17 | 99.44 | 56.67 | 64.40 | 72.50 | 79.91

As can be seen in Table 5, the usage of deep and expressive models within our
framework advances the state-of-the-art on Reddit and Flickr, while it
performs equally well for others, _e.g._ , PPI. Notably, our approach
outperforms the two historical-based variants VR-GCN and MVS-GNN by a wide
margin. Interestingly, our deep and expressive variants reach superior
performance than our GCN baseline on _all_ datasets, which highlights the
benefits of evaluating larger models on larger scale.

##  7 Conclusion and Future Work

We proposed a general framework for scaling arbitrary message passing GNNs to
large graphs without the necessity to sub-sample edges. As we have shown, our
approach is able to train both deep and expressive GNNs in a scalable fashion.
Notably, our approach is _orthogonal_ to many methodological advancements,
such as unifying GNNs and label propagation (Shi et al., 2020), graph
diffusion (Klicpera et al., 2019b), or random wiring (Valsesia et al., 2020),
which we like to investigate further in future works. While our experiments
focus on node-level tasks, our work is technically able to scale the training
of GNNs for edge-level and graph-level tasks as well. However, this still
needs to be verified empirically. Another interesting future direction is the
fusion of GAS into a distributed training algorithm (Ma et al., 2019; Zhu et
al., 2016; Tripathy et al., 2020; Wan et al., 2020; Angerd et al., 2020; Zheng
et al., 2020), and to extend our framework in accessing histories from disk
storage rather than CPU memory. Overall, we hope that our findings lead to the
development of sophisticated and expressive GNNs evaluated on large-scale
graphs.

## Acknowledgements

This work has been supported by the _German Research Association (DFG)_ within
the Collaborative Research Center SFB 876 _Providing Information by Resource-
Constrained Analysis_ , projects A6 and B2.

## References

  * Angerd et al. (2020) Angerd, A., Balasubramanian, K., and Annavaram, M.  Distributed training of graph convolutional networks using subgraph approximation.  _ICLR submission_ , 2020.
  * Bottoue & Bousquet (2007) Bottoue, L. and Bousquet, O.  The tradeoffs of large scale learning.  In _NIPS_ , 2007.
  * Chen et al. (2018a) Chen, J., Ma, T., and Xiao, C.  FastGCN: Fast learning with graph convolutional networks via importance sampling.  In _ICLR_ , 2018a.
  * Chen et al. (2018b) Chen, J., Zhu, J., and Song, L.  Stochastic training of graph convolutional networks with variance reduction.  In _ICML_ , 2018b.
  * Chen et al. (2020a) Chen, M., Wei, Z., Ding, B., Li, Y., Yuan, Y., Du, X., and Wen, J. R.  Scalable graph neural networks via bidirectional propagation.  In _NeurIPS_ , 2020a.
  * Chen et al. (2020b) Chen, M., Wei, Z., Huang, Z., Ding, B., and Li, Y.  Simple and deep graph convolutional networks.  In _ICML_ , 2020b.
  * Chiang et al. (2019) Chiang, W. L., Liu, X., Si, S., Li, Y., Bengio, S., and Hsieh, C. J.  Cluster-GCN: An efficient algorithm for training deep and large graph convolutional networks.  In _KDD_ , 2019.
  * Cong et al. (2020) Cong, W., Forsati, R., Kandemir, M., and Mahdavi, M.  Minimal variance sampling with provable guarantees for fast training of graph neural networks.  In _KDD_ , 2020.
  * Corso et al. (2020) Corso, G., Cavalleri, L., Beaini, D., Liò, P., and Veličković, P.  Principal neighbourhood aggregation for graph nets.  In _NeurIPS_ , 2020.
  * Dhillon et al. (2007) Dhillon, I. S., Guan, Y., and Kulis, B.  Weighted graph cuts without eigenvectors: A multilevel approach.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 29(11):1944–1957, 2007.
  * Dwivedi et al. (2020) Dwivedi, V. P., Joshi, C. K., Laurent, T., Bengio, Y., and Bresson, X.  Benchmarking graph neural networks.  _CoRR_ , abs/2003.00982, 2020.
  * Fey & Lenssen (2019) Fey, M. and Lenssen, J. E.  Fast graph representation learning with PyTorch Geometric.  In _ICLR-W_ , 2019.
  * Frasca et al. (2020) Frasca, F., Rossi, E., Eynard, D., Chamberlain, B., Bronstein, M. M., and Monti, F.  SIGN: Scalable inception graph neural networks.  In _ICML-W_ , 2020.
  * Gilmer et al. (2017) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry.  In _ICML_ , 2017.
  * Hamilton (2020) Hamilton, W. L.  Graph representation learning.  _Synthesis Lectures on Artificial Intelligence and Machine Learning_ , 14(3):1–159, 2020.
  * Hamilton et al. (2017) Hamilton, W. L., Ying, R., and Leskovec, J.  Inductive representation learning on large graphs.  In _NIPS_ , 2017.
  * Hornik (1991) Hornik, K.  Approximation capabilities of multilayer feedforward networks.  _Neural Networks_ , 4(2):251–257, 1991.
  * Hornik et al. (1989) Hornik, K., Stinchcombe, M., and White, H.  Multilayer feedforward networks are universal approximators.  _Neural Networks_ , 2(5):359–366, 1989.
  * Hu et al. (2020) Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J.  Open Graph Benchmark: Datasets for machine learning on graphs.  In _NeurIPS_ , 2020.
  * Huang et al. (2021) Huang, Q., He, H., Singh, A., Lim, S. N., and Benson, A. R.  Combining label propagation and simple models out-performs graph neural networks.  In _ICLR_ , 2021.
  * Huang et al. (2018) Huang, W., Zhang, T., Rong, Y., and Huang, J.  Adaptive sampling towards fast graph representation learning.  In _NeurIPS_ , 2018.
  * Karypis & Kumar (1998) Karypis, G. and Kumar, V.  A fast and high quality multilevel scheme for partitioning irregular graphs.  _SIAM Journal on Scientific Computing_ , 20(1):359––392, 1998.
  * Kipf & Welling (2017) Kipf, T. N. and Welling, M.  Semi-supervised classification with graph convolutional networks.  In _ICLR_ , 2017.
  * Klicpera et al. (2019a) Klicpera, J., Bojchevski, A., and Günnemann, S.  Predict then propagate: Graph neural networks meet personalized PageRank.  In _ICLR_ , 2019a.
  * Klicpera et al. (2019b) Klicpera, J., Weißenberger, S., and Günnemann, S.  Diffusion improves graph learning.  In _NeurIPS_ , 2019b.
  * Ma et al. (2019) Ma, L., Yang, Z., Miao, Y., Xue, J., Wu, M., Zhou, L., and Dai, Y.  NeuGraph: Parallel deep neural network computation on large graphs.  In _USENIX Annual Technical Conference_ , 2019.
  * Ma & Tang (2020) Ma, Y. and Tang, J.  _Deep Learning on Graphs_.  Cambridge University Press, 2020.
  * Markowitz et al. (2021) Markowitz, E., Balasubramanian, K., Mirtaheri, M., Abu-El-Haija, S., Perozzi, B., Ver Steeg, G., and Galstyan, A.  Graph traversal with tensor functionals: A meta-algorithm for scalable learning.  In _ICLR_ , 2021.
  * Maron et al. (2019) Maron, H., Ben-Hamu, H., Serviansky, H., and Lipman, Y.  Provably powerful graph networks.  In _NeurIPS_ , 2019.
  * Mernyei & Cangea (2020) Mernyei, P. and Cangea, C.  Wiki-CS: A wikipedia-based benchmark for graph neural networks.  In _ICML-W_ , 2020.
  * Morris et al. (2019) Morris, C., Ritzert, M., Fey, M., Hamilton, W. L., Lenssen, J. E., Rattan, G., and Grohe, M.  Weisfeiler and Leman go neural: Higher-order graph neural networks.  In _AAAI_ , 2019.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  PyTorch: An imperative style, high-performance deep learning library.  In _NeurIPS_ , 2019.
  * Qi et al. (2017) Qi, C. R., Yi, L., Su, H., and Guibas, L. J.  PointNet++: Deep hierarchical feature learning on point sets in a metric space.  In _NIPS_ , 2017.
  * Rong et al. (2020) Rong, Y., Huang, W., Xu, T., and Huang, J.  DropEdge: Towards deep graph convolutional networks on node classification.  In _ICLR_ , 2020.
  * Sen et al. (2008) Sen, G., Namata, G., Bilgic, M., and Getoor, L.  Collective classification in network data.  _AI Magazine_ , 29, 2008.
  * Seo et al. (2019) Seo, Y., Loukas, A., and Perraudin, N.  Discriminative structural graph classification.  _CoRR_ , abs/1905.13422, 2019.
  * Shchur et al. (2018) Shchur, O., Mumme, M., Bojchevski, A., and Günnemann, S.  Pitfalls of graph neural network evaluation.  In _NeurIPS-W_ , 2018.
  * Shi et al. (2020) Shi, Y., Huang, Z., Wang, W., Zhong, H., Feng, S., and Sun, Y.  Masked label prediction: Unified message passing model for semi-supervised classification.  _CoRR_ , abs/2009.03509, 2020.
  * Tripathy et al. (2020) Tripathy, A., Yelick, K., and Buluc, A.  Reducing communcation in graph neural network training.  _CoRR_ , abs/2005.03300, 2020.
  * Usama & Chang (2018) Usama, M. and Chang, D. E.  Towards robust neural networks with lipschitz continuity.  _CoRR_ , abs/1811.09008, 2018.
  * Valsesia et al. (2020) Valsesia, D., Fracastoro, G., and Magli, E.  Don’t stack layers in graph neural networks, wire them randomly.  _ICLR submission_ , 2020.
  * Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y.  Graph attention networks.  In _ICLR_ , 2018.
  * Wan et al. (2020) Wan, C., Li, Y., Kim, N. S., and Lin, Y.  BDS-GCN: Efficient full-graph training of graph convolutional nets with partition-parallelism and boundary sampling.  _ICLR submission_ , 2020.
  * Wang et al. (2019) Wang, Y., Sun, Y., Liu, Z., Sarma, S. E., Bronstein, M. M., and Solomon, J. M.  Dynamic graph CNN for learning on point clouds.  _ACM Transactions on Graphics (TOG)_ , 2019.
  * Weisfeiler & Lehman (1968) Weisfeiler, B. and Lehman, A. A.  A reduction of a graph to a canonical form and an algebra arising during this reduction.  _Nauchno-Technicheskaya Informatsia_ , 2(9), 1968.
  * Wu et al. (2019) Wu, F., Zhang, T., de Souza Jr., A. H., Fifty, C., Yu, T., and Weinberger, K. Q.  Simplifying graph convolutional networks.  In _ICML_ , 2019.
  * Xu et al. (2018) Xu, K., Li, C., Tian, Y., Sonobe, T., Kawarabayashi, K., and Jegelka, S.  Representation learning on graphs with jumping knowledge networks.  In _ICML_ , 2018.
  * Xu et al. (2019) Xu, K., Hu, W., Leskovec, J., and Jegelka, S.  How powerful are graph neural networks?  In _ICLR_ , 2019.
  * Yang et al. (2016) Yang, Z., Cohen, W., and Salakhutdinov, R.  Revisiting semi-supervised learning with graph embeddings.  In _ICML_ , 2016.
  * You et al. (2020) You, Y., Chen, T., Wang, Z., and Shen, Y.  L2-GCN: Layer-wise and learned efficient training of graph convolutional networks.  In _CVPR_ , 2020.
  * Yu et al. (2020) Yu, L., Shen, J., Li, J., and Lerer, A.  Scalable graph neural networks for heterogeneous graphs.  _CoRR_ , abs/2011.09679, 2020.
  * Zaheer et al. (2017) Zaheer, M., Kottur, S., Ravanbhakhsh, S., Póczos, B., Salakhutdinov, R., and Smola, A. J.  Deep sets.  In _NIPS_ , 2017.
  * Zeng et al. (2020a) Zeng, H., Zhang, M., Xia, Y., Srivastava, A., Kannan, R., Prasanna, V., Jin, L., Malevich, A., and Chen, R.  Deep graph neural networks with shallow subgraph samplers.  _CoRR_ , abs/2012.01.380, 2020a.
  * Zeng et al. (2020b) Zeng, H., Zhou, H., Srivastava, A., Kannan, R., and Prasanna, V.  GraphSAINT: Graph sampling based inductive learning method.  In _ICLR_ , 2020b.
  * Zheng et al. (2020) Zheng, D., Ma, C., Wang, M., Zhou, J., Su, Q., Song, X., Gan, Q., Zhang, Z., and Karypis, G.  DistDGL: Distributed graph neural network for training for billion-scale graphs.  _CoRR_ , abs/2010.05337, 2020.
  * Zhu et al. (2019) Zhu, R., Zhao, K., Yang, H., Lin, W., Zhou, C., Ai, B., Li, Y., and Zhou, J.  AliGraph: A comprehensive graph neural network platform.  In _KDD_ , 2019.
  * Zhu et al. (2016) Zhu, X., Chen, W., Zheng, W., and Ma, X.  Gemini: A computation-centric distributed graph processing system.  In _USENIX Symposium on Operating Systems Designand Implementation_ , 2016.
  * Zou et al. (2019) Zou, D., Hu, Z., Wang, Y., Jiang, S., Sun, Y., and Gu, Q.  Layer-dependent importance sampling for training deep and large graph convolutional networks.  In _NeurIPS_ , 2019.

##  8 Proofs

######  Lemma 1.

Let $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ be Lipschitz continuous functions
with Lipschitz constants $k_{1}$ and $k_{2}$, respectively. If, for all
$v\in\mathcal{V}$, the inputs are close to the exact input, _i.e._
$\|\bm{\tilde{h}}^{(\ell-1)}_{v}-\bm{h}^{(\ell-1)}_{v}\|\leq\delta$, and the
historical embeddings do not run too stale, _i.e._
$\|\bm{\bar{h}}^{(\ell-1)}_{v}-\bm{\tilde{h}}^{(\ell-1)}_{v}\|\leq\epsilon$,
then the output error is bounded by

$\|\bm{\tilde{h}}^{(\ell)}_{v}-\bm{h}^{(\ell)}_{v}\|\leq\delta\,k_{2}+(\delta+\epsilon)\,k_{1}\,k_{2}\,|\mathcal{N}(v)|.$

###### Proof.

By triangular inequality, it holds that
$\|\bm{\bar{h}}^{(\ell-1)}_{v}-\bm{h}^{(\ell-1)}_{v}\|\leq\delta+\epsilon$.
Since both $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ denote Lipschitz continuous functions
with Lipschitz constants $k_{1}$ and $k_{2}$, respectively, it further holds
that for any $\bm{x},\bm{y}$:

$\displaystyle\|\textsc{Message}^{(\ell)}_{\bm{\theta}}(\bm{x})-\textsc{Message}^{(\ell)}_{\bm{\theta}}(\bm{y})\|$
$\displaystyle\leq k_{1}\|\bm{x}-\bm{y}\|$
$\displaystyle\|\textsc{Update}^{(\ell)}_{\bm{\theta}}(\bm{x})-\textsc{Update}^{(\ell)}_{\bm{\theta}}(\bm{y})\|$
$\displaystyle\leq k_{2}\|\bm{x}-\bm{y}\|$

Furthermore, the Lipschitz constants for the aggregations
$\sum_{\bm{x}\in\mathcal{X}}\bm{x}$,
$\frac{1}{|\mathcal{X}|}\sum_{\bm{x}\in\mathcal{X}}\bm{x}$ and
$\max_{\bm{x}\in\mathcal{X}}\bm{x}$ are given as $|\mathcal{X}|$, $1$ and $1$,
respectively. Then,

$\displaystyle\|\textsc{Update}^{(\ell)}_{\bm{\theta}}(\bm{\tilde{h}}^{(\ell-1)}_{v},\bigoplus\limits_{\mathclap{w\in\mathcal{N}(v)}}\textsc{Message}^{(\ell)}_{\bm{\theta}}(\bm{\bar{h}}_{w}^{(\ell-1)}))-\textsc{Update}^{(\ell)}_{\bm{\theta}}(\bm{h}^{(\ell-1)}_{v},\bigoplus\limits_{\mathclap{w\in\mathcal{N}(v)}}\textsc{Message}^{(\ell)}_{\bm{\theta}}(\bm{h}_{w}^{(\ell-1)}))\|$
$\displaystyle\leq\,$ $\displaystyle
k_{2}\,(\delta+|\mathcal{N}(v)|\,(k_{1}\,(\delta+\epsilon)))=\delta\,k_{2}+(\delta+\epsilon)\,k_{1}\,k_{2}\,|\mathcal{N}(v)|.\qed$

######  Theorem 2.

Let $\bm{f}^{(L)}_{\bm{\theta}}$ be a $L$-layered GNN, containing only
Lipschitz continuous $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions with Lipschitz constants
$k_{1}$ and $k_{2}$, respectively. If, for all $v\in\mathcal{V}$ and all
$\ell\in\\{1,\ldots,L-1\\}$, the historical embeddings do not run too stale,
_i.e._
$\|\bm{\bar{h}}^{(\ell)}_{v}-\bm{\tilde{h}}^{(\ell)}_{v}\|\leq\epsilon^{(\ell)}$,
then the final output error is bounded by

$\|\bm{\tilde{h}}_{v,j}^{(L)}-\bm{h}_{v,j}^{(L)}\|\leq\sum_{\ell=1}^{L-1}\epsilon^{(\ell)}\,k_{1}^{L-\ell}\,k_{2}^{L-\ell}\,{|\mathcal{N}(v)|}^{L-\ell}.$

###### Proof.

For layer $\ell=1$, the inputs do not need to be estimated, _i.e._
$\delta^{(0)}=\|\bm{\tilde{h}}_{v}^{(0)}-\bm{h}_{v}^{(0)}\|=0$, and, as a
result, the output is _exact_ , _i.e._
$\delta^{(1)}=\|\bm{\tilde{h}}_{v}^{(1)}-\bm{h}_{v}^{(1)}\|=0$. With
$\|\bm{\bar{h}}^{(1)}_{v}-\bm{\tilde{h}}^{(1)}_{v}\|\leq\epsilon^{(1)}$, it
directly follows via Lemma 1 that the approximation error of layer $\ell=2$ is
bounded by
$\|\bm{\tilde{h}}^{(2)}_{v}-\bm{h}^{(2)}_{v}\|\leq\epsilon^{(1)}\,k_{1}\,k_{2}\,|\mathcal{N}(v)|=\delta^{(2)}$.
Recursively replacing

$\delta^{(\ell)}=\delta^{(\ell-1)}\,k_{2}+(\delta^{(\ell-1)}+\epsilon^{(\ell-1)})\,k_{1}\,k_{2}\,|\mathcal{N}(v)|$

in
$\|\bm{\tilde{h}}_{v}^{(L)}-\bm{h}_{v}^{(L)}\|\leq\delta^{(L-1)}\,k_{2}+(\delta^{(L-1)}+\epsilon^{(L-1)})\,k_{1}\,k_{2}\,|\mathcal{N}(v)|$
(_cf._ Lemma 1) yields

$\|\bm{\tilde{h}}_{v}^{(L)}-\bm{h}_{v}^{(L)}\|\leq\sum_{\ell=1}^{L-1}\epsilon^{(\ell)}\,k_{1}^{L-\ell}\,k_{2}^{L-\ell}\,{|\mathcal{N}(v)|}^{L-\ell}.\qed$

######  Proposition 3.

Let $\bm{f}^{(L)}_{\bm{\theta}}\colon\mathcal{V}\to\mathbb{R}^{d}$ be a
$L$-layered GNN as expressive as the WL test in distinguishing the $L$-hop
neighborhood around each node $v\in\mathcal{V}$. Then, there exists a graph
$\bm{A}\in{\\{0,1\\}}^{|\mathcal{V}|\times|\mathcal{V}|}$ for which
$\bm{f}^{(L)}_{\bm{\theta}}$ operating on a sampled variant $\bm{\tilde{A}}$,
$\tilde{a}_{v,w}=\begin{cases}\frac{|\mathcal{N}(v)|}{|\mathcal{\tilde{N}}(v)|},&\textrm{if
}w\in\mathcal{\tilde{N}}(v)\\\ 0,&\textrm{otherwise}\end{cases}$, produces a
non-equivalent coloring, _i.e._
$\bm{\tilde{h}}^{(L)}_{v}\neq\bm{\tilde{h}}^{(L)}_{w}$ while
$c_{v}^{(L)}=c_{w}^{(L)}$ for nodes $v,w\in\mathcal{V}$.

###### Proof.

Consider the colored graph $\bm{A}$ and its sampled variant $\bm{\tilde{A}}$
as shown on the right. Here, it holds that
$\bm{h}^{(1)}_{v_{1}}=\bm{h}^{(1)}_{v_{4}}$ while
$\bm{\tilde{h}}^{(1)}_{v_{1}}\neq\bm{\tilde{h}}^{(1)}_{v_{4}}$. ∎

######  Lemma 4.

Let $\\{\\{\bm{h}_{v}^{(\ell-1)}\colon v\in\mathcal{V}\\}\\}$ be a countable
multiset such that
$\|\bm{h}_{v}^{(\ell-1)}-\bm{h}_{w}^{(\ell-1)}\|>2(\delta+\epsilon)$ for all
$v,w\in\mathcal{V}$, $\bm{h}_{v}^{(\ell-1)}\neq\bm{h}_{w}^{(\ell-1)}$. If the
inputs are close to the exact input, _i.e._
$\|\bm{\tilde{h}}_{v}^{(\ell-1)}-\bm{h}_{v}^{(\ell-1)}\|\leq\delta$, and the
historical embeddings do not run too stale, _i.e._
$\|\bm{\bar{h}}_{v}^{(\ell-1)}-\bm{\tilde{h}}_{v}^{(\ell-1)}\|\leq\epsilon$,
then there exist $\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions, such that

$\|\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{\tilde{h}}_{v}^{(\ell-1)})-\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{h}_{v}^{(\ell-1)})\|\leq\delta+\epsilon$

and

$\|\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{h}_{v}^{(\ell-1)})-\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{h}_{w}^{(\ell-1)})\|>2(\delta+\epsilon+\lambda)$

for all $v,w\in\mathcal{V}$, $\bm{h}_{v}^{(\ell-1)}\neq\bm{h}_{w}^{(\ell-1)}$
and all $\lambda>0$.

###### Proof.

Define $\bm{\phi}\colon\mathbb{R}^{d}\to\mathbb{R}^{d}$ as the Voronoi
tessellation induced by exact inputs $\\{\bm{h}_{v}^{(\ell-1)}\colon
v\in\mathcal{V}\\}$:

$\bm{\phi}(\bm{x})=\bm{h}_{v}^{(\ell-1)}\quad\textrm{if}\quad\|\bm{x}-\bm{h}_{v}^{(\ell-1)}\|\leq\|\bm{x}-\bm{h}_{w}^{(\ell-1)}\|\quad\textrm{for
all }v\neq w\in\mathcal{V}$ (4)

Furthermore, we know that there exists
$\textrm{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions so that
$\bm{f}^{(\ell)}_{\bm{\theta}}$ is injective for all countable multisets
(Zaheer et al., 2017; Xu et al., 2019; Morris et al., 2019; Maron et al.,
2019). Therefore, it holds that
$\|\bm{f}^{(\ell)}_{\bm{\theta}}\big{(}\bm{\phi}\big{(}\bm{\tilde{h}}_{v}^{(\ell-1)}\big{)}\big{)}-\bm{f}^{(\ell)}_{\bm{\theta}}\big{(}\bm{\phi}\big{(}\bm{h}_{v}^{(\ell-1)}\big{)}\big{)}\|=0\leq\delta$.
Since $\\{\\{\bm{h}_{v}^{(\ell-1)}\colon v\in\mathcal{V}\\}\\}$ is countable
and $\bm{f}^{(\ell)}_{\bm{\theta}}$ is injective, there exists a $\kappa>0$
such that
$\|\bm{f}^{(\ell)}_{\bm{\theta}}\big{(}\bm{\phi}\big{(}\bm{h}_{v}^{(\ell-1)}\big{)}\big{)}-\bm{f}^{(\ell)}_{\bm{\theta}}\big{(}\bm{\phi}\big{(}\bm{h}_{w}^{(\ell-1)}\big{)}\big{)}\|>\kappa$
for all $v,w\in\mathcal{V}$, $\bm{h}_{v}^{(\ell-1)}\neq\bm{h}_{w}^{(\ell-1)}$.
Due to the homogeneity of $\|\cdot\|$, it directly follows that there must
exists $\alpha>0$ so that

$\|\alpha\bm{f}^{(\ell)}_{\bm{\theta}}\big{(}\bm{\phi}\big{(}\bm{h}_{v}^{(\ell-1)}\big{)}\big{)}-\alpha\bm{f}^{(\ell)}_{\bm{\theta}}\big{(}\bm{\phi}\big{(}\bm{h}_{w}^{(\ell-1)}\big{)}\big{)}\|>\alpha\,\kappa\geq
2(\delta+\epsilon+\lambda)$

for all $v,w\in\mathcal{V}$, $\bm{h}_{v}^{(\ell-1)}\neq\bm{h}_{w}^{(\ell-1)}$
and all $\lambda>0$. ∎

######  Theorem 5.

Let $\bm{f}^{(L)}_{\bm{\theta}}$ be a $L$-layered GNN in which all
$\textsc{Message}^{(\ell)}_{\bm{\theta}}$ and
$\textsc{Update}^{(\ell)}_{\bm{\theta}}$ functions fulfill the conditions of
Lemma 4. Then, there exists a map $\phi\colon\mathbb{R}^{d}\to\Sigma$ so that
$\phi(\bm{\tilde{h}}^{(L)}_{v})=c^{(L)}_{v}$ for all $v\in\mathcal{V}$.

###### Proof.

Define $\phi\colon\mathbb{R}^{d}\to\Sigma$ as the Voronoi tessellation induced
by exact outputs $\\{\bm{h}_{v}^{(L)}\colon v\in\mathcal{V}\\}$:

$\phi(\bm{x})=c_{v}^{(L)}\quad\textrm{if}\quad\|\bm{x}-\bm{h}^{(L)}_{v}\|\leq\|\bm{x}-\bm{h}^{(L)}_{w}\|\quad\textrm{for
all }v\neq w\in\mathcal{V}$

Since each GNN layer $\bm{f}^{(\ell)}_{\bm{\theta}}$ is injective for exact
inputs, we know that such a function needs to exist (Xu et al., 2019; Morris
et al., 2019). Therefore, it is sufficient to show that there exists a
$\delta^{(L)}>0$ so that
$\|\bm{\tilde{h}}^{(L)}_{v}-\bm{h}^{(L)}_{v}\|\leq\delta^{(L)}$ and
$\|\bm{h}^{(L)}_{v}-\bm{h}^{(L)}_{w}\|>2\delta^{(L)}$ for all
$v,w\in\mathcal{V}$, $\bm{h}^{(L)}_{v}\neq\bm{h}^{(L)}_{w}$. Following upon
Theorem 2, we know that $\|\bm{\tilde{h}}_{v}^{(1)}-\bm{h}_{v}^{(1)}\|=0$. Due
to Lemma 4, it holds that
$\|\bm{\tilde{h}}_{v}^{(2)}-\bm{h}_{v}^{(2)}\|\leq\epsilon^{(1)}$. The next
layer introduces an increased error, _i.e._
$\|\bm{\bar{h}}_{v}^{(2)}-\bm{h}_{v}^{(2)}\|\leq\epsilon^{(1)}+\epsilon^{(2)}$,
and to compensate, we set $\lambda^{(2)}=\epsilon^{(2)}$ so that
$\|\bm{h}_{v}^{(2)}-\bm{h}_{w}^{(2)}\|>2\,(\epsilon^{(1)}+\epsilon^{(2)})$ for
all $v,w\in\mathcal{V}$, $\bm{h}^{(L)}_{v}\neq\bm{h}^{(L)}_{w}$. By
recursively applying Lemma 4 with $\lambda^{(\ell)}=\epsilon^{(\ell)}$, it
immediately follows that
$\|\bm{\tilde{h}}_{v}^{(L)}-\bm{h}_{v}^{(L)}\|\leq\sum_{\ell=1}^{L-1}\epsilon^{(\ell)}=\delta^{(L)}$,
and
$\|\bm{\tilde{h}}_{v}^{(L)}-\bm{h}_{w}^{(L)}\|>\sum_{\ell=1}^{L-1}2\,\epsilon^{(\ell)}$
for all $v,w\in\mathcal{V}$, $\bm{h}^{(L)}_{v}\neq\bm{h}^{(L)}_{w}$. ∎

##  9 Algorithm

Our GAS mini-batch training algorithm is given in Algorithm 1:

Algorithm 1 GAS Mini-batch Execution

Input: Graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$, input node features
$\bm{H}^{(0)}$, number of batches $B$, number of layers $L$

$\\{\mathcal{B}_{1},\ldots,\mathcal{B}_{B}\\}\leftarrow\textsc{Split}(\mathcal{G},B)$

$\mathcal{V}_{b}\leftarrow\bigcup_{v\in\mathcal{B}_{b}}\mathcal{N}(v)\cup\\{v\\}$
$\forall b\in\\{1,\ldots,B\\}$

$\mathcal{G}_{b}\leftarrow\mathcal{G}[\mathcal{V}_{b}]$ $\forall
b\in\\{1,\ldots,B\\}$

for $\mathcal{B}_{b}\in\\{\mathcal{B}_{1},\ldots,\mathcal{B}_{B}\\}$ do

for $\ell\in\\{1,\ldots,L-1\\}$ do

$\bm{h}_{v}^{(\ell)}\leftarrow\bm{f}^{(\ell)}_{\bm{\theta}}(\bm{h}_{v}^{(\ell-1)},\\{\\{\bm{h}_{w}^{(\ell-1)}:w\in\mathcal{N}(v)\\}\\})$
$\forall v\in\mathcal{B}_{b}$

$\textsc{Push}^{(\ell)}(\bm{h}_{v}^{(\ell)})$ $\forall v\in\mathcal{B}_{b}$

$\bm{h}_{w}^{(\ell)}\leftarrow\textsc{Pull}^{(\ell)}(w)$ $\forall
w\in\mathcal{V}_{b}\setminus\mathcal{B}_{b}$

end for

$\bm{h}_{v}^{(L)}\leftarrow\bm{f}^{(L)}_{\bm{\theta}}(\bm{h}_{v}^{(L-1)},\\{\\{\bm{h}_{w}^{(L-1)}:w\in\mathcal{N}(v)\\}\\})$
$\forall v\in\mathcal{B}_{b}$

end for

##  10 GNN Operators

We briefly recap the details of all graph convolutional layers used in our
experiments. We omit final non-linearities and edge features due to
simplicity.

#### Graph Convolutional Networks (GCN)

use a symmetrically normalized mean aggregation followed by linear
transformation (Kipf & Welling, 2017)

$\bm{h}_{v}^{(\ell)}=\sum_{w\in\mathcal{N}(v)\cup\\{v\\}}\frac{1}{c_{w,v}}\bm{W}\bm{h}_{w}^{(\ell-1)},$

where $c_{w,v}=\sqrt{\deg(w)+1}\sqrt{\deg(v)+1}$.

#### Graph Attention Networks (GAT)

perform an anisotropic aggregation (Veličković et al., 2018)

$\bm{h}_{v}^{(\ell)}=\sum_{w\in\mathcal{N}(v)\cup\\{v\\}}\alpha_{w,v}\bm{W}\bm{h}_{w}^{(\ell-1)},$

where normalization is achieved via learnable attention coefficients

$\alpha_{w,v}=\frac{\exp\left(\textrm{LeakyReLU}\left(\bm{a}^{\top}\left[\bm{W}\bm{h}_{v}^{(\ell-1)},\bm{W}\bm{h}_{w}^{(\ell-1)}\right]\right)\right)}{\sum_{k\in\mathcal{N}(v)\cup\\{v\\}}\exp\left(\textrm{LeakyReLU}\left(\bm{a}^{\top}\left[\bm{W}\bm{h}_{v}^{(\ell-1)},\bm{W}\bm{h}_{k}^{(\ell-1)}\right]\right)\right)}.$

#### Approximate Personalized Propagation of Neural Predictions (APPNP)

networks first perform a graph-agnostic prediction of node labels, _i.e._
$\bm{h}^{(0)}_{v}=\textrm{MLP}(\bm{x}_{v})$, and smooth initial label
predictions via propagation afterwards (Klicpera et al., 2019a)

$\bm{h}^{(\ell)}=\alpha\,\bm{h}^{(0)}+(1-\alpha)\sum_{w\in\mathcal{N}\cup\\{v\\}}\frac{1}{c_{w,v}}\bm{h}^{(\ell-1)}_{w},$

where $\alpha\in[0,1]$ denotes the teleport probability and $c_{w,v}$ is
defined as in GCN. Notably, the final propagation layers are non-trainable,
and predictions are solely conditioned on node features (while gradients of
model parameters are not).

#### Simple and Deep Graph Convolutional Networks (GCNII)

extend the idea of APPNP to a trainable propgation scheme which leverages
initial residual connections (Chen et al., 2020b)

$\bm{h}_{v}^{(\ell)}=\alpha\bm{W}\bm{h}_{v}^{(0)}+(1-\alpha)\sum_{w\in\mathcal{N}(v)\cup\\{v\\}}\frac{1}{c_{w,v}}\bm{W}\bm{h}_{w}^{(\ell-1)},$

and $\bm{W}$ makes use of identity maps, _i.e._
$\bm{W}\leftarrow(1-\beta)\bm{I}+\beta\bm{W}$ for $\beta\in[0,1]$.

#### Graph Isomorphism Networks (GIN)

make use of sum aggregation and MLPs to obtain a maximally powerful GNN
operator (Xu et al., 2019)

$\bm{h}_{v}^{(\ell)}=\textrm{MLP}_{\bm{\theta}}\left((1+\epsilon)\,\bm{h}_{v}^{(\ell-1)}+\sum_{w\in\mathcal{N}(v)}\bm{h}_{w}^{(\ell-1)}\right),$

where $\epsilon\in\mathbb{R}$ is a trainable parameter in order to distinguish
neighbors from central nodes.

#### Principal Neighborhood Aggregation (PNA)

networks leverage mulitple aggregators combined with degree-scalers to capture
graph structural properties (Corso et al., 2020)

$\bm{h}_{v}^{(\ell)}=\bm{W}_{2}\left[\bm{h}_{v}^{(\ell-1)},\bigoplus_{w\in\mathcal{N}(v)}\bm{W}_{1}\left[\bm{h}_{v}^{(\ell-1)},\bm{h}_{w}^{(\ell-1)}\right]\right],$

where

$\bigoplus=\underbrace{\begin{bmatrix}1\\\ s(\deg(v),1)\\\
s(\deg(v),-1)\end{bmatrix}}_{\text{Scalers}}\otimes\underbrace{\begin{bmatrix}\textrm{mean}\\\
\min\\\ \max\end{bmatrix}}_{\text{Aggregators}},$

with $\otimes$ being the tensor product and

$s(d,\alpha)={\left(\frac{\log(d+1)}{\frac{1}{|\mathcal{V}|}\sum_{v\in\mathcal{V}}\log(\deg(v)+1)}\right)}^{\alpha}$

denoting degree-scalers.

##  11 PyGAS Programming Interface

To highlight the ease-of-use of our framework, we showcase the necessary
changes to convert a common GCN architecture (Kipf & Welling, 2017)
implemented in PyTorch Geometric (Fey & Lenssen, 2019) (_cf._ Listing 1) to
its corresponding scalable version (_cf._ Listing 2). In particular, our model
now inherits from ScalableGNN, which takes care of creating all history
embeddings (accessible via self.histories) and provides an efficient
concurrent history access pattern via push_and_pull(). Notably, the forward()
execution method of our model now takes in the additional n_id parameter,
which holds the global node index for each node in the current mini-batch.
This assignment vector is necessary to push and pull the intermediate mini-
batch embeddings to and from the global history embeddings.

⬇

from torch_geometric.nn import GCNConv

class GNN(Module):

def __init__(self, in_channels, hidden_channels, out_channels, num_layers):

super(GNN, self).__init__()

self.convs = ModuleList()

self.convs.append(GCNConv(in_channels, hidden_channels))

for _ in range(num_layers - 2):

self.convs.append(GCNConv(hidden_channels, hidden_channels))

self.convs.append(GCNConv(hidden_channels, out_channels))

def forward(self, x, adj_t):

for conv in self.convs[:-1]:

x = conv(x, adj_t).relu()

return self.convs[-1](x, adj_t)

Listing 1: Full-batch GCN (Kipf & Welling, 2017) model within PyTorch
Geometric (Fey & Lenssen, 2019).

⬇

from torch_geometric.nn import GCNConv

⬇

from torch_geometric_autoscale import ScalableGNN

⬇

w

⬇

class GNN(ScalableGNN):

def __init__(self, num_nodes, in_channels, hidden_channels, out_channels,
num_layers):

super(GNN, self).__init__(num_nodes, hidden_channels, num_layers)

⬇

self.convs = ModuleList()

self.convs.append(GCNConv(in_channels, hidden_channels))

for _ in range(num_layers - 2):

self.convs.append(GCNConv(hidden_channels, hidden_channels))

self.convs.append(GCNConv(hidden_channels, out_channels))

⬇

w

⬇

def forward(self, x, adj_t, n_id):

for conv, history in zip(self.convs[:-1], self.histories):

⬇

x = conv(x, adj_t).relu()

⬇

x = self.push_and_pull(history, x, n_id)

⬇

return self.convs[-1](x, adj_t)

Listing 2: Mini-batch GCN (Kipf & Welling, 2017) model within PyTorch
Geometric (Fey & Lenssen, 2019) and our proposed _PyGAS_ framework.
$\blacksquare$ denotes lines that require changes, while $\blacksquare$ refers
to newly added lines. Only minimal changes are required to auto-scale GCN (or
any other model) to large graphs.

##  12 Addtional Ablation Studies

We report additional ablation studies to further strengthen the motivation of
our GAS framework:

Table 6: Inter-/intra-connectivity ratio for real-world datasets with
different mini-batch sampling strategies. Utilizing Metis heavily minimizes
inter-connectivity between mini-batches, which reduces history accesses and
tightens approximation errors in return.

Sampling | Cora | CiteSeer | PubMed | Coauthor- |  Amazon- | Wiki-CS
---|---|---|---|---|---|---
Scheme | CS | Physics | Computer | Photo
Random | 1.33 | 1.24 | 3.17 | 6.81 | 9.94 | 9.05 | 5.61 | 5.85
Metis | 0.14 | 0.02 | 0.52 | 2.77 | 2.26 | 2.27 | 1.03 | 1.12
| Cluster | Pattern | Reddit | PPI | Flickr | Yelp | ogbn- | ogbn-
| arxiv | products
Random | 36.64 | 51.02 | 6.58 | 6.79 | 1.82 | 6.74 | 3.02 | 26.18
Metis | 1.57 | 1.61 | 2.80 | 1.27 | 1.07 | 2.52 | 0.48 | 1.94

#### Minimizing Inter-Connectivity Between Batches.

We make use of graph clustering methods (Karypis & Kumar, 1998; Dhillon et
al., 2007) in order to minimize the inter-connectivity between batches, which
minimizes history accesses and therefore increases closeness and reduces
staleness in return. To evaluate this impact in practice, Tabel 6 lists the
inter-/intra-connectivity ratio of all real-world datasets used in our
experiments, both for randomly sampled mini-batches as well as for utilizing
Metis partitions as mini-batches. Notably, applying Metis beforehand reduces
the overall inter-/intra-connectivity ratio by a factor of 4 on average, which
results in only a fraction of history accesses. Furthermore, most real-world
datasets come with inter-/intra-connectivity ratios between $0.1$ and $2.5$,
leading to only marginal runtime overheads when leveraging historical
information, as confirmed by our runtime analysis.

Table 7: Ablation study for a 4-layer GIN (Xu et al., 2019) model on the
Cluster dataset (Dwivedi et al., 2020). Combining both GAS techniques help in
resembling full-batch performance for expressive models with highly non-linear
message passing phases.

|  |  | Accuracy
---|---|---|---
|  |  | Training | Validation | Test
| Full-batch Baseline | 60.49 | 58.17 | 58.49
| Minimizing | Enforcing |  |  |
| Inter-Connectivity | Lipschitz Continuity |  |  |
GAS | ✗ | ✗ | 55.66 | 54.86 | 55.15
✔ | ✗ | 58.97 | 57.79 | 57.82
✔ | ✔ | 60.67 | 58.21 | 58.51

#### Analysis of Gains for Obtaining Expressive Node Representations.

Next, we highlight the impacts of minimizing the inter-connectivity between
mini-batches and enforcing Lipschitz continuity of the learned function in
order to derive expressive node representations. Here, we benchmark a 4-layer
GIN model (Xu et al., 2019) on the Cluster dataset (Dwivedi et al., 2020),
_cf._ Table 7. Notably, both solutions achieve significant gains in training,
validation and test performance, and together, they are able to closely
resemble the performance of full-batch training. However, we found that
Lipschitz continuity regularization only helps in non-linear message passing
phases, while it does not provide any additional gains for linear operators
such as GCN (Kipf & Welling, 2017).

##  13 Datasets

We give detailed statistics for all datasets used in our experiments, _cf._
Table 8, which include the following tasks:

Table 8: Dataset statistics.

| Dataset | Task | Nodes | Edges | Features | Classes | Label Rate
---|---|---|---|---|---|---|---
Small-scale | Cora | multi-class | 2,708 | 5,278 | 1,433 | 7 | 5.17%
CiteSeer | multi-class | 3,327 | 4,552 | 3,703 | 6 | 3.61%
PubMed | multi-class | 19,717 | 44,324 | 500 | 3 | 0.30%
Coauthor-CS | multi-class | 18,333 | 81,894 | 6,805 | 15 | 1.64%
Coauthor-Physics | multi-class | 34,493 | 247,962 | 8,415 | 5 | 0.29%
Amazon-Computer | multi-class | 13,752 | 245,861 | 767 | 10 | 1.45%
Amazon-Photo | multi-class | 7,650 | 119,081 | 745 | 8 | 2.09%
Wiki-CS | multi-class | 11,701 | 215,863 | 300 | 10 | 4.96%
Large-scale | Cluster | multi-class | 1,406,436 | 25,810,340 | 6 | 6 | 83.35%
Reddit | multi-class | 232,965 | 11,606,919 | 602 | 41 | 65.86%
PPI | multi-label | 56,944 | 793,632 | 50 | 121 | 78.86%
Flickr | multi-class | 89,250 | 449,878 | 500 | 7 | 50.00%
Yelp | multi-label | 716,847 | 6,977,409 | 300 | 100 | 75.00%
ogbn-arxiv | multi-class | 169,343 | 1,157,799 | 128 | 40 | 53.70%
ogbn-products | multi-class | 2,449,029 | 61,859,076 | 100 | 47 | 8.03%

  1. 1.

classifying academic papers in citation networks (Cora, CiteSeer, PubMed) (Sen
et al., 2008; Yang et al., 2016)

  2. 2.

categorizing computer science articles in Wikipedia graphs (Wiki-CS) (Mernyei
& Cangea, 2020)

  3. 3.

predicting active research fields of authors in co-authorshop graphs
(Coauthor-CS, Coauthor-Physics) (Shchur et al., 2018)

  4. 4.

predicting product categories in co-purchase graphs (Amazon-Computer, Amazon-
Photo) (Shchur et al., 2018)

  5. 5.

identifying community clusters in Stochastic Block Models (Cluster, Pattern)
(Dwivedi et al., 2020)

  6. 6.

predicting communities of online posts based on user comments (Reddit)
(Hamilton et al., 2017)

  7. 7.

classifying protein functions based on the interactions of human tissue
proteins (PPI) (Hamilton et al., 2017)

  8. 8.

categorizing types of images based on their descriptions and properties
(Flickr) (Zeng et al., 2020b)

  9. 9.

classifying business types based on customers and friendship relations (Yelp)
(Zeng et al., 2020b)

  10. 10.

predicting subject areas of arXiv Computer Science papers (ogbn-arxiv) (Hu et
al., 2020)

  11. 11.

predicting product categories in an Amazon product co-purchasing network
(ogbn-products) (Hu et al., 2020)