LMC: F AST T RAINING OF GNN S VIA SUBGRAPHWISE SAMPLING WITH P ROVABLE C ONVERGENCE
A BSTRACT
The message passing-based graph neural networks (GNNs) have achieved great success in many real-world applications.
However, training GNNs on large-scale graphs suffers from the well-known neighbor explosion problem, i.e., the exponentially increasing dependencies of nodes with the number of message passing layers.
Subgraph-wise sampling methods—a promising class of mini-batch training techniques—discard messages outside the mini-batches in backward passes to avoid the neighbor explosion problem at the expense of gradient estimation accuracy.
This poses significant challenges to their convergence analysis and convergence speeds, which seriously limits their reliable real-world applications.
To address this challenge, we propose a novel subgraph-wise sampling method with a convergence guarantee, namely L ocal M essage C ompensation (LMC).
To the best of our knowledge, LMC is the first subgraph-wise sampling method with provable convergence.
The key idea of LMC is to retrieve the discarded messages in backward passes based on a message passing formulation of backward passes.
By efficient and effective compensations for the discarded messages in both forward and backward passes, LMC computes accurate mini-batch gradients and thus accelerates convergence.
We further show that LMC converges to first-order stationary points of GNNs.
Experiments on large-scale benchmark tasks demonstrate that
LMC significantly outperforms state-of-the-art subgraph-wise sampling methods in terms of efficiency.
I NTRODUCTION
Graph neural networks (GNNs) are powerful frameworks that generate node embeddings for graphs via the iterative message passing (MP) scheme (Hamilton, 2020).
At each MP layer, GNNs aggregate messages from each node’s neighborhood and then update node embeddings based on aggregation results.
Such a scheme has achieved great success in many real-world applications involving graph-structured data, such as search engines (Brin & Page, 1998), recommendation systems (Fan et al., 2019), materials engineering (Gostick et al., 2016), molecular property prediction (Moloi & Ali, 2005; Kearnes et al., 2016), and combinatorial optimization (Wang et al., 2023).
However, the iterative MP scheme poses challenges to training GNNs on large-scale graphs.
One commonly-seen approach to scale deep models to arbitrarily large-scale data with limited GPU memory is to approximate full-batch gradients by mini-batch gradients.
Nevertheless, for the graphstructured data, the computational costs for computing the loss across a mini-batch of nodes and the corresponding mini-batch gradients are expensive due to the well-known neighbor explosion problem.
Specifically, the embedding of a node at the k -th MP layer recursively depends on the embeddings of its neighbors at the ( k − 1) -th MP layer.
Thus, the complexity grows exponentially with the number of MP layers.
To deal with the neighbor explosion problem, recent works propose various sampling techniques to reduce the number of nodes involved in message passing (Ma & Tang, 2021).
For example, nodewise (Hamilton et al., 2017; Chen et al., 2018a) and layer-wise (Chen et al., 2018b; Zou et al., 2019;
Huang et al., 2018) sampling methods recursively sample neighbors over MP layers to estimate node embeddings and corresponding mini-batch gradients.
Unlike the recursive fashion, subgraph-wise sampling methods (Chiang et al., 2019; Zeng et al., 2020; Fey et al., 2021; Zeng et al., 2021) adopt a cheap and simple one-shot sampling fashion, i.e., sampling the same subgraph constructed based on a mini-batch for different MP layers.
By discarding messages outside the mini-batches, subgraphwise sampling methods restrict message passing to the mini-batches such that the complexity grows linearly with the number of MP layers.
Moreover, subgraph-wise sampling methods are applicable to a wide range of GNN architectures by directly running GNNs on the subgraphs constructed by the sampled mini-batches (Fey et al., 2021).
Because of these advantages, subgraph-wise sampling methods have recently drawn increasing attention.
Despite the empirical success of subgraph-wise sampling methods, discarding messages outside the mini-batch sacrifices the gradient estimation accuracy, which poses significant challenges to their convergence behaviors.
First, recent works (Chen et al., 2018a;
Cong et al., 2020) demonstrate that the inaccurate mini-batch gradients seriously hurt the convergence speeds of GNNs.
Second, in
Section 7.3, we demonstrate that many subgraph-wise sampling methods are difficult to resemble full-batch performance under small batch sizes, which we usually use to avoid running out of GPU memory in practice.
These issues seriously limit the real-world applications of GNNs.
In this paper, we propose a novel subgraph-wise sampling method with a convergence guarantee, namely L ocal M essage C ompensation (LMC), which uses efficient and effective compensations to correct the biases of mini-batch gradients and thus accelerates convergence.
To the best of our knowledge, LMC is the first subgraph-wise sampling method with provable convergence.
Specifically, we first propose unbiased mini-batch gradients for the one-shot sampling fashion, which helps decompose the gradient computation errors into two components: the bias from the discarded messages and the variance of the unbiased mini-batch gradients.
Second, based on a message passing formulation of backward passes, we retrieve the messages discarded by existing subgraph-wise sampling methods during the approximation to the unbiased mini-batch gradients.
Finally, we propose efficient and effective compensations for the discarded messages with a combination of incomplete up-to-date messages and messages generated from historical information in previous iterations, avoiding the exponentially growing time and memory consumption.
An appealing feature of the resulting mechanism is that it can effectively correct the biases of mini-batch gradients, leading to accurate gradient estimation and the speed-up of convergence.
We further show that LMC converges to first-order stationary points of GNNs.
Notably, the convergence of LMC is based on the interactions between mini-batch nodes and their 1 -hop neighbors, without the recursive expansion of neighborhoods to aggregate information far away from the mini-batches.
Experiments on largescale benchmark tasks demonstrate that LMC significantly outperforms state-of-the-art subgraphwise sampling methods in terms of efficiency.
Moreover, under small batch sizes, LMC outperforms the baselines and resembles the prediction performance of full-batch methods.
R ELATED W ORK
In this section, we discuss some works related to our proposed method.
Subgraph-wise Sampling Methods.
Subgraph-wise sampling methods sample a mini-batch and then construct the same subgraph based on it for different MP layers (Ma & Tang, 2021).
For example, Cluster-GCN (Chiang et al., 2019) and GraphSAINT (Zeng et al., 2020) construct the subgraph induced by a sampled mini-batch.
They encourage connections between the sampled nodes by graph clustering methods (e.g., METIS (Karypis & Kumar, 1998) and Graclus (Dhillon et al., 2007)), edge, node, or random-walk-based samplers.
GNNAutoScale (GAS) (Fey et al., 2021) and MVS-GNN (Cong et al., 2020) use historical embeddings to generate messages outside a sampled subgraph, maintaining the expressiveness of the original GNNs.
Recursive Graph Sampling Methods.
Both node-wise and layer-wise sampling methods recursively sample neighbors over MP layers and then construct different computation graphs for each
MP layer.
Node-wise sampling methods (Hamilton et al., 2017;
Chen et al., 2018a) aggregate messages from a small subset of sampled neighborhoods at each MP layer to decrease the bases in the exponentially increasing dependencies.
To avoid the exponentially growing computation, layer-wise sampling methods (Chen et al., 2018b; Zou et al.,
; Huang et al., 2018) independently sample nodes for each MP layer and then use importance sampling to reduce variance, resulting in a constant sample size in each MP layer.
Pre-Processing Methods.
Another line for scalable graph neural networks is to develop preprocessing Methods.
They aggregate the raw input features and then take the pre-processing features as input into subsequent models.
As the aggregation has no parameters, they can use stochastic gradient descent to train the subsequent models without the neighbor explosion problem.
While they are efficient in training and inference, they are not applicable to powerful GNNs with a trainable aggregation process.
Historical Values as an Affordable Approximation.
The historical values are affordable approximations of the exact values in practice.
However, they suffer from frequent data transfers to/from the GPU and the staleness problem.
For example, in node-wise sampling, VR-GCN (Chen et al., 2018a) uses historical embeddings to reduce the variance from neighbor sampling (Hamilton et al., 2017).
GAS (Fey et al., 2021) proposes a concurrent mini-batch execution to transfer the active historical embeddings to and from the GPU, leading to comparable runtime with the standard full-batch approach.
GraphFM-IB and GraphFM-OB (Yu et al., 2022) apply a momentum step on historical embeddings for node-wise and subgraph-wise sampling methods with historical embeddings, respectively, to alleviate the staleness problem.
Both LMC and GraphFM-OB use the node embeddings in the mini-batch to alleviate the staleness problem of the node embeddings outside the mini-batch.
We discuss the main differences between LMC and GraphFM-OB in Appendix C.1.
P RELIMINARIES
We introduce notations and graph neural networks in Sections 3.1 and 3.2, respectively
. 3. N OTATIONS
A graph G = ( V , E ) is defined by a set of nodes V = { v 1 , v 2 , . .
. , v n } and a set of edges E among these nodes.
The set of nodes consists of labeled nodes V L and unlabeled nodes V U :=
V \ V L .
Let ( v i , v
j ) ∈
E denote an edge going from node v i ∈ V to node v j ∈ V , N ( v i ) = { v j ∈ V| ( v i , v j ) ∈ E} denote the neighborhood of node v i , and N ( v i ) denote N ( v i ) ∪{ v i } .
We assume that G is undirected, i.e., v j ∈ N ( v i ) ⇔ v i ∈ N ( v j ) .
Let N ( S )
=
{ v ∈ V| ( v i , v j ) ∈ E , v i ∈ S} denote the neighborhoods of a set of nodes S and N ( S ) denote N ( S ) ∪ S .
For a positive integer L , [ L ] denotes { 1 , . . .
, L } .
Let the boldface character x i ∈ R
d
x denote the feature of node v i with dimension d x .
Let h i ∈
R
d be the d -dimensional embedding of the node v i .
Let X =
( x 1 , x 2 , . . .
, x n ) ∈ R
d x × n and H = ( h 1 , h 2 , . . .
, h n ) ∈ R d × n .
We also denote the embeddings of a set of nodes S
= { v i k } |S| k =1 by H S
= ( h i k ) |S| k =1 ∈ R d ×|S| .
For a p × q matrix A ∈ R p × q , ⃗ A ∈ R pq denotes the vectorization of A , i.e.,
A ij = ⃗ A i +( j − 1) p .
We denote the j -th columns of A by A j .
3. G RAPH N EURAL N ETWORKS
For the semi-supervised node-level prediction, Graph Neural Networks (GNNs) aim to learn node embeddings H with parameters Θ by minimizing the objective function L = 1 |V L | (cid:80) i ∈V L ℓ w ( h i , y i )
such that H = GNN ( X , E ; Θ) , where ℓ w is the composition of an output layer with parameters w and a loss function.
GNNs follow the message passing framework in which vector messages are exchanged between nodes and updated using neural networks.
An L -layer GNN performs L message passing iterations with different parameters Θ
= ( θ l ) Ll =1 to generate the final node embeddings H
= H L as
H l = f θ l ( H l − 1 ; X ) , l ∈ [ L ] , (1) where H
= X and f θ l is the message passing function of the l -th layer with parameters θ l .
The message passing function f θ l follows an aggregation and update scheme, i.e.,
[Equation]
where g θ l is the function generating individual messages for each neighbor of v i in the l -th message passing iteration, ⊕ θ l is the aggregation function mapping a set of messages to the final message m l − 1 N ( v i ) , and u θ l is the update function that combines previous node embedding h l − 1 i , message m l − 1 N ( v i ) , and features x i to update node embeddings.
M ESSAGE P ASSING IN B ACKWARD P ASSES
In Section 4.1, we introduce the gradients of GNNs and formulate the backward passes as message passing.
Then we propose backward SGD, which is an SGD variant, in Section 4.2.
4. B ACKWARD P ASSES AND M ESSAGE PASSING F ORMULATION
The gradient ∇ w L is easy to compute and we hence introduce the chain rule to compute ∇ Θ L in this section, where Θ = ( θ l ) Ll =1 .
Let V l ≜ ∇ H l L for l ∈
[ L ] be auxiliary variables.
It is easy to compute
⃗ V L = ∇
[Equation]
H L L = ∇ ⃗ H L .
By the chain rule, we iteratively compute V l based on V l +1 as
[Equation]
and
[Equation]
Then, we compute the gradient ∇ θ l L = ( ∇ θ l ⃗f θ l ) ⃗ V l ,
l ∈
[ L ] by using autograd packages for vectorJacobian product.
We formulate backward passes, i.e., the processes of iterating Equation (3), as message passing.
To see this, we need to notice that Equation (3) is equivalent to
[Equation]
where V lk is the k -th column of V l and m l N ( v j ) is a function of h li defined in Equation (2).
Equation (5) uses (cid:16) ∇ h li u θ l +1 ( h lj , m l N ( v
[Equation]
) , x j ) (cid:17) V l +1 j , sum aggregation, and the identity mapping as the generation function, the aggregation function, and the update function, respectively.
4. B ACKWARD SGD
In this section, we develop an SGD variant—backward SGD, which provides unbiased gradient estimations based on the message passing formulation of backward passes.
Backward SGD is the basis of our proposed subgraph-wise sampling method, i.e., LMC, in Section 5.
Given a sampled mini-batch V B , suppose that we have computed exact node embeddings ( H l V B ) Ll = and auxiliary variables ( V l V B ) Ll =1 of nodes in V B .
To simplify the analysis, we assume that V B is uniformly sampled from V and the corresponding set of labeled nodes V L B := V B ∩ V L is uniformly sampled from V L .
When the sampling is not uniform, we use the normalization technique (Zeng et al., 2020) to enforce the assumption (please see Appendix A.3.1).
First, backward SGD computes the mini-batch gradient g w ( V B ) for parameters w by the derivative of mini-batch loss L
[Equation]
B =
[Equation]
Then, backward SGD computes the mini-batch gradient g θ l ( V B ) for parameters θ l as
[Equation]
Note that the mini-batch gradients g θ l ( V B ) for different l ∈
[ L ] are based on the same mini-batch
V B , which facilitates designing subgraph-wise sampling methods based on backward SGD.
Another appealing feature of backward SGD is that the mini-batch gradients g w ( V B ) and g θ l ( V B ) ,
l ∈
[ L ] are unbiased, as shown in the following theorem.
Please see Appendix D.1 for the detailed proof.
Theorem 1.
Suppose that a mini-batch V B is uniformly sampled from V and the corresponding labeled nodes V L B = V B ∩ V L is uniformly sampled from V L .
Then the mini-batch gradients g w ( V B ) and g θ l ( V B ) ,
l ∈ [ L ] in Equations (6) and (7) are unbiased.
L OCAL M ESSAGE C OMPENSATION
The exact mini-batch gradients g w ( V B ) and g θ l ( V B ) , l ∈ [ L ] computed by backward SGD depend on exact embeddings and auxiliary variables of nodes in the mini-batch V B rather than the whole graph.
However, backward SGD is not scalable, as the exact ( H l V B ) Ll
=1 and ( V l V B )
L l =1 are expensive to compute due to the neighbor explosion problem.
In this section, to deal with the neighbor explosion problem, we develop a novel and scalable subgraph-wise sampling method for GNNs, namely L ocal M essage C ompensation (LMC).
LMC first efficiently estimates ( H l V B ) Ll
=1 and ( V l V B ) Ll =1 by convex combinations of the incomplete upto-date values and the historical values , and then computes the mini-batch gradients as shown in v 4 vvvv(
[Figure]
Original graph vvvvvvvvvvvv 2 Layer
Layer
Layer(b)
Forward passes of GAS vvvvvvvvvvvv 2 Layer
Layer
Layer(c) Forward passes of LMC
[Figure]
vvvvvvvvvvvv
[Figure]
GAS vvvvvvvvvvvv
[Figure]
Equations (6) and (7).
We show that LMC converges to first-order stationary points of GNNs in
Section 6.
In Algorithm 1 and Section 6, we denote a value in the l -th layer at the k -th iteration by ( · ) l,k , but elsewhere we omit the superscript k and denote it by ( · ) l .
In each training iteration, we sample a mini-batch of nodes V B and propose to approximate values, i.e., node embeddings and auxiliary variables, outside V B by convex combinations of historical values , denoted H l V\V B and V l V\V B , and incomplete up-to-date values , denoted (cid:101)
H l V\V B and (cid:101) V l V\V B .
In forward passes, we initialize the temporary embeddings for l = 0 as (cid:98) H
=
X and update historical embeddings of nodes in V B , i.e., H l V B , in the order of l = 1 , 2 , . . .
, L .
Specifically, in the l -th layer, we first update the historical embedding of each node v i ∈ V B as
[Equation]
[Equation]
Then, we compute the temporary embedding of each neighbor v i ∈ N ( V B ) \ V B as
[Equation]
where β i ∈ [0 , 1] is the convex combination coefficient for node v i , and
[Equation]
We call C lf ≜ ⊕ θ l (cid:16)(cid:110) g θ l ( (cid:98) h l − 1 j ) | v j ∈ N ( v i ) \ V B (cid:111)(cid:17) the local message compensation in the l -th layer in forward passes.
For l ∈ [ L ] , h l i is an approximation to h li computed by Equation (2).
Notice that the total size of Equations (8)–(10) is linear with |N ( V B ) | rather than the size of the whole graph.
Suppose that the maximum neighborhood size is n max and the number of layers is L , then the time complexity in forward passes is O ( L ( n max |V B | d + |V B | d 2 )) .
In backward passes, we initialize the temporary auxiliary variables for l = L as (cid:98) V L = ∇ H L and update historical auxiliary variables of nodes in V B , i.e., V l V B , in the order of l = L − 1 , . . .
Specifically, in the l -th layer, we first update the historical auxiliary variable of each v i ∈
V B as
[Equation]
where h lj , m l N ( v j ) , and (cid:98) h lj are computed as shown in Equations (8)–(10).
Then, we compute the temporary auxiliary variable of each neighbor v i ∈ N ( V B ) \ V B as
[Equation]
where β i is the convex combination coefficient used in Equation (9)
[Equation]
and
[Equation]
We call C
[Equation]
the local message compensation in the l -th layer in backward passes.
For l ∈ [ L ] , V li is an approximation to V li computed by Equation (3).
Similar to forward passes, the time complexity in backward passes is O
( L ( n max |V B | d + |V B | d 2 )) , where n max is the maximum neighborhood size and L is the number of layers.
1: Input: The learning rate η and the convex combination coefficients ( β i ) ni =1 .
2: Partition V into B parts ( V b )
Bb = 3: for k = 1 , . .
. , N do 4:
Randomly sample V b k from ( V b )
Bb = 5:
Initialize H 0 ,k = (cid:98) H 0 ,k
= X 6: for l = 1 , . . .
, L do 7:
Update H l,k V bk
▷ (8) 8:
Compute (cid:98) H l,k N ( V bk )
\V bk ▷ (9) and (10) 9: end for 10:
Initialize V L,k = (cid:98) V L,k = ∇ H L L 11: for l = L − 1 , . .
, 1 do 12:
Update V l,k V bk
▷ (11) 13:
Compute (cid:98) V l,k N ( V b k ) \V
b k ▷ (12) and (13) 14: end for 15:
Compute (cid:101) g kw and (cid:101) g kθ l , l ∈
[ L ] ▷ (6) and (7) 16:
Update parameters by 17: w k = w k − 1 − η (cid:101) g kw 18: θ l,k = θ l,k − 1 − η (cid:101) g kθ l ,
l ∈
[ L ] 19: end for
LMC additionally stores the historical node embeddings H l and auxiliary variables V l for l ∈
[ L ] .
As pointed out in (Fey et al., 2021), we can store the majority of historical values in RAM or hard drive storage rather than GPU memory.
Thus, the active historical values in forward and backward passes employ O ( n max L |V B | d ) and O ( n max L |V B | d ) GPU memory, respectively (see Appendix B).
As the time and memory complexity are independent of the size of the whole graph, i.e., |V| , LMC is scalable.
We summarize the computational complexity in Appendix B.
Figure 1 shows the message passing mechanisms of GAS (Fey et al., 2021) and
LMC.
Compared with GAS, LMC proposes compensation messages between inbatch nodes and their 1-hop neighbors simultaneously in forward and backward passes.
This corrects the biases of minibatch gradients and thus accelerates convergence.
Algorithm 1 summarizes LMC.
Unlike above, we add a superscript k for each value to indicate that it is the value at the k -th iteration.
At preprocessing step, we partition V into B parts ( V b )
Bb =1 .
At the kth training step, LMC first randomly samples a subgraph constructed by V b k .
Notice that we sample more subgraphs to build a large graph in experiments whose convergence analysis is consistent with that of sampling a single subgraph.
Then, LMC updates the stored historical node embeddings H l,k V bk in the order of l = 1 , . . .
, L by Equations (8)–(10), and the stored historical auxiliary variables V l,k V bk in the order of l = L − 1 , . . .
, 1 by Equations (11)–(13).
By the randomly updating, the historical values get close to the exact up-to-date values.
Finally, for l ∈ [ L ] and v j ∈ V b k , by replacing h l,kj , m l,k N ( v j ) and V l,k j in Equations (6) and (7) with h l,kj , m l,k N ( v j ) , and V l,kj , respectively, LMC computes mini-batch gradients (cid:101) g w , (cid:101)
[Table]
T HEORETICAL A NALYSIS
In this section, we provide the theoretical analysis of LMC.
Theorem 2 shows that the biases of mini-batch gradients computed by LMC can tend to an arbitrarily small value by setting a proper learning rate and convex combination coefficients.
Then, Theorem 3 shows that LMC converges to first-order stationary points of GNNs.
We provide detailed proofs of the theorems in Appendix D.
In the theoretical analysis, we suppose that the following assumptions hold in this paper.
Assumption 1.
Assume that (1) at the k -th iteration, a batch of nodes V k B is uniformly sampled from
V and the corresponding labeled node set V kL B = V k B ∩V
L is uniformly sampled from V L , (2) functions f θ l , ϕ θ l , ∇ w L , ∇ θ l L , ∇ w ℓ w , and ∇ θ l u θ l are γ -Lipschitz with γ > 1 , ∀ l ∈
[ L ] , (3) norms ∥
H l,k ∥ F , ∥ H l,k ∥ F , ∥ (cid:98) H l,k ∥ F , ∥ (cid:101) H l,k ∥ F , ∥ V l,k ∥ F , ∥ V l,k ∥ F , ∥ (cid:98) V l,k ∥ F , ∥ (cid:101) V l,k ∥ F , ∥∇ w L∥ 2 , ∥∇ θ l L∥ 2 , ∥ (cid:101) g θ l ∥ 2 , and ∥ (cid:101) g w ∥ 2 are bounded by G > 1 ,
∀ l ∈ [ L ] , k ∈ N ∗ .
Theorem 2.
Suppose that Assumption 1 holds, then with η = O ( ε 2 ) and β i
= O ( ε 2 ) , i ∈ [ n ] , there exists C > 0 and ρ ∈ (0 , 1) such that
[Equation]
Theorem 3.
Suppose that Assumption 1 holds.
Besides, assume that the optimal value L ∗ = inf w, Θ L ( w, Θ) is bounded by G .
Then, with η = O ( ε 4 ) , β i =
O ( ε 4 ) , i ∈ [ n ] , and N =
O ( ε − 6 ) ,
LMC ensures to find an ε -stationary solution such that E [ ∥∇ w, Θ L ( w R , Θ R ) ∥ 2 ] ≤ ε after running for N iterations, where R is uniformly selected from [ N ] and Θ R =
( θ l,R ) Ll =1 .
E XPERIMENTS
We introduce experimental settings in Section 7.1.
We then evaluate the convergence and efficiency of LMC in Sections 7.2 and 7.3.
Finally, we conduct ablation studies about the proposed compensations in Section 7.4.
We run all experiments on a single GeForce RTX 2080 Ti (11 GB).
7. E XPERIMENTAL S ETTINGS
Datasets.
Some recent works (Hu et al., 2020) have indicated that many frequently-used graph datasets are too small compared with graphs in real-world applications.
Therefore, we evaluate LMC on four large datasets, PPI, REDDIT , FLICKR(Hamilton et al., 2017), and Ogbn-arxiv (Hu et al., 2020).
These datasets contain thousands or millions of nodes/edges and have been widely used in previous works (Fey et al., 2021; Zeng et al., 2020; Hamilton et al., 2017; Chiang et al., 2019; Chen et al., 2018a;b).
For more details, please refer to Appendix A.1.
Baselines and Implementation Details.
In terms of prediction performance, our baselines include node-wise sampling methods (GraphSAGE (Hamilton et al., 2017) and VR-GCN (Chen et al., 2018a)), layer-wise sampling method (FASTGCN (Chen et al., 2018b) and LADIES (Zou et al., 2019)), subgraph-wise sampling methods (CLUSTER-GCN (Chiang et al., 2019), GRAPHSAINT (Zeng et al., 2020), FM (Yu et al., 2022), and GAS (Fey et al., 2021)), and a precomputing method (SIGN (Rossi et al., 2020)).
By noticing that GAS and FM achieve the state-of-the-art prediction performance (Table 1) among the baselines, we further compare the efficiency of LMC with GAS,
FM, and CLUSTER-GCN, another subgraph-wise sampling method using METIS partition.
We implement LMC, FM, and CLUSTER-GCN based on the codes and toolkits of GAS (Fey et al., 2021) to ensure a fair comparison.
For other implementation details, please refer to Appendix A.3.
Hyperparameters.
To ensure a fair comparison, we follow the data splits, training pipeline, and most hyperparameters in (Fey et al., 2021) except for the additional hyperparameters in LMC such as β i .
We use the grid search to find the best β i (see Appendix A.4 for more details).
7. LMC IS F AST WITHOUT S ACRIFICING A CCURACY
Table 1 reports the prediction performance of LMC and the baselines.
We report the mean and the standard deviation by running each experiment five times for GAS, FM, and LMC.
LMC, FM, and
GAS all resemble full-batch performance on all datasets while other baselines may fail, especially on the FLICKR dataset.
Moreover, LMC, FM, and GAS with deep GNNs, i.e., GCNII (Chen et al., 2020) outperform other baselines on all datasets.
[Table]
As LMC, FM, and GAS share the similar prediction performance, we additionally compare the convergence speed of LMC, FM, GAS, and
CLUSTER-GCN, another subgraph-wise sampling method using METIS partition, in Figure 2 and Table 2.
We use a sliding window to smooth the convergence curve in Figure 2 as the accuracy on test data is unstable.
The solid curves correspond to the mean, and the shaded regions correspond to values within plus or minus one standard deviation of the mean.
Table 2 reports the number of epochs, the runtime to reach the full-batch accuracy in Table 1, and
the
GPU memory.
As shown in Table 2 and Figure 2a, LMC is significantly faster than GAS, especially with a speed-up of 2x on REDDIT.
Notably, the test accuracy of LMC is more stable than GAS, and thus the smooth test accuracy of
LMC outperforms GAS in Figure 2b.
Although
GAS finally resembles full-batch performance in Table 1 by selecting the best performance on the valid data, it may fail to resemble under small batch sizes due to its unstable process (see Section 7.3).
Another appealing feature of LMC is that
LMC shares comparable GPU memory costs with GAS, and thus LMC avoids the neighbor explosion problem.
FM is slower than other methods, as they additionally update historical embeddings in the storage for the nodes outside the mini-batches.
Please see Appendix E.2 for the comparison in terms of training time per epoch.
[Table]
Ogbn-arxiv & GCNII — 234.8 373.6 197.
—   —  568 FLICKR & GCNII — 400.
To further illustrate the convergence of LMC, we compare the errors of mini-batch gradients computed by CLUSTER, GAS, and LMC.
At epoch training step, we record the relative errors ∥ (cid:101) g θ l − ∇ θ l L∥ 2 / ∥∇ θ l L∥ 2 , where ∇ θ l L is the full-batch gradient for the parameters θ l at the l -th MP layer and the (cid:101) g
θ l is a mini-batch gradient.
To avoid the randomness of the full-batch gradient ∇ θ l L , we set the dropout rate as zero.
We report average relative errors during training in Figure 3.
LMC enjoys the smallest estimated errors in the experiments.
[Figure]
Layer Layer
Layer 2 0.
R e l a t i v e e rr o r s o f g r a d
i e n t s
CLUSTERGASLMC (a) Ogbn-arxiv
Layer
Layer 1 0.
R e l a t i v e e rr o r s o f g r a d
i e n t s
CLUSTERGASLMC (b) FLICKR
Layer
Layer 1 0.
R e l a t i v e e rr o r s o f g r a d
i e n t s
CLUSTERGASLMC (
[Figure]
REDDIT
7. LMC IS R OBUST IN TERMS OF B ATCH S IZES
[Table]
An appealing feature of mini-batch training methods is that they can avoid the out-of-memory issue by decreasing the batch size.
Thus, we evaluate the prediction performance of LMC on Ogbn-arxiv datasets with different batch sizes (numbers of clusters).
We conduct experiments under different sizes of sampled clusters per mini-batch.
We run each experiment with the same epoch and search learning rates in the same set.
We report the best prediction accuracy in Table 3.
LMC outperforms GAS under small batch sizes (batch size = 1 or 2 ) and achieve comparable performance with GAS (batch size = 5 or 10 ).
7. A BLATION
The improvement of LMC is due to two parts: the compensation in forward passes C lf and the compensation in back passes C lb .
Compared with GAS, the compensation in forward passes C lf additionally combines the incomplete up-to-date messages.
Figure 4 shows the convergence curves of LMC using both C lf and C lb (denoted by C f & C b ), LMC using only C lf (denoted by C f ), and
GAS on the Ogbn-arxiv dataset.
Under small batch sizes, the improvement mainly is due to C lb and the incomplete up-to-date messages in forward passes may hurt the performance.
This is because the mini-batch and the union of their neighbors are hard to contain most neighbors of out-of-batch nodes when the batch size is small.
Thus, the compensation in back passes C lb is the most important component by correcting the bias of the mini-batch gradients.
Under large batch sizes, the improvement is due to C lf , as the large batch sizes decrease the discarded messages and improve the accuracy of the mini-batch gradients (see Table 7 in Appendix).
Notably, C lb still slightly improves the performance.
We provide more ablation studies about β i in Appendix E.4.
[Figure]
C ONCLUSION
In this paper, we propose a novel subgraph-wise sampling method with a convergence guarantee, namely Local Message Compensation (LMC).
LMC uses efficient and effective compensations to correct the biases of mini-batch gradients and thus accelerates convergence.
We show that LMC converges to first-order stationary points of GNNs.
To the best of our knowledge, LMC is the first subgraph-wise sampling method for GNNs with provable convergence.
Experiments on large-scale benchmark tasks demonstrate that LMC significantly outperforms state-of-the-art subgraph-wise sampling methods in terms of efficiency.
A CKNOWLEDGEMENT
The authors would like to thank all the anonymous reviewers for their insightful comments.
This work was supported in part by National Nature Science Foundations of China grants U19B2026,
U19B2044, 61836011, and 61836006, and the Fundamental Research Funds for the
Central Universities grant WK
A M ORE D ETAILS ABOUT E XPERIMENTS
In this section, we introduce more details about our experiments, including datasets, training and evaluation protocols, and implementations.
A. D ATASETS
We evaluate LMC on four large datasets, PPI, REDDIT, FLICKR (Hamilton et al., 2017), and Ogbnarxiv (Hu et al., 2020).
All of the datasets do not contain personally identifiable information or offensive content.
Table 4 shows the summary statistics of the datasets.
Details about the datasets are as follows.
• PPI contains 24 protein-protein interaction graphs.
Each graph corresponds to a human tissue.
Each node indicates a protein with positional gene sets, motif gene sets and immunological signatures as node features.
Edges represent interactions between proteins.
The task is to classify protein functions.
• REDDIT is a post-to-post graph constructed from REDDIT.
Each node indicates a post and each edge between posts indicates that the same user comments on both.
The task is to classify REDDIT posts into different communities based on (1) the GloVe CommonCrawl word vectors (Pennington et al., 2014) of the post titles and comments, (2) the post’s scores, and (3) the number of comments made on the posts.
• Ogbn-arxiv is a directed citation network between all Computer Science (CS) arXiv papers indexed by MAG (Wang et al., 2020).
Each node is an arXiv paper and each directed edge indicates that one paper cites another one.
The task is to classify unlabeled arXiv papers into different primary categories based on labeled papers and node features, which are computed by averaging word2vec (Mikolov et al., 2013) embeddings of words in papers’ title and abstract.
• FLICKR categorizes types of images based on their descriptions and properties (Fey et al., 2021; Zeng et al., 2020).
[Table]
REDDIT  232, 11,606,
Ogbn-arxiv  169, 1,157,
FLICKR  89, 449,
A. T RAINING AND E VALUATION P ROTOCOLS
We run all the experiments on a single GeForce RTX 2080 Ti (11 GB).
All the models are implemented in Pytorch (Paszke et al., 2019) and PyTorch Geometric (Fey & Lenssen, 2019) based on the official implementation of (Fey et al., 2021) 1 .
The code of LMC is available on GitHub at https://github.com/MIRALab-USTC/GNN-LMC .
Data Splitting.
We use the data splitting strategies following previous works (Fey et al., 2021; Gu et al., 2020).
A. I MPLEMENTATION D ETAILS A.3. N ORMALIZATION T ECHNIQUE
In Section 4.2 in the main text, we assume that the subgraph V B is uniformly sampled from V and the corresponding set of labeled nodes V L B = V B ∩ V L is uniformly sampled from V L .
To enforce the assumption, we use the normalization technique to reweight Equations (6) and (7) in the main text.
Suppose we partition the whole graph V into b parts {V B i } bi =
and then uniformly sample c clusters without replacement to construct subgraph V B .
By the normalization technique, Equation (6) becomes
[Equation]
where b |V L B | c |V L | is the corresponding weight.
Similarly, Equation (7) becomes
[Equation]
where b |V B | c |V| is the corresponding weight.
A.3. I NCORPORATING B ATCH N ORMALIZATION
We uniformly sample a mini-batch of nodes V B and generate the induced subgraph of N ( V B ) .
If we directly feed the H ( l ) N ( V B ) to a batch normalization layer, the learned mean and standard deviation of the batch normalization layer may be biased.
Thus, LMC first feeds the embeddings of the minibatch H ( l ) V B to a batch normalization layer and then feeds the embeddings outside the mini-batch
H ( l ) N ( V B ) \V B to another batch normalization layer.
A. S ELECTION OF β i
We select β i = score ( i ) α for each node v i , where α ∈ [0 , 1] is a hyperparameter and score is a function to measure the quality of the incomplete up-to-date messages.
We search score in a { f ( x ) = x 2 , f ( x ) = 2 x − x 2 , f ( x ) = x, f ( x ) = 1; x = deg local ( i ) /deg
global ( i ) } , where deg global ( i ) is the degree of node i in the whole graph and deg local ( i ) is the degree of node i in the subgraph induced by N ( V B ) .
B C OMPUTATIONAL C OMPLEXITY
We summarize the computational complexity in Table 5, where n max is the maximum of neighborhoods, L is the number of message passing layers, V B is a set of nodes in a sampled mini-batch, d is the embedding dimension, V is the set of nodes in the whole graph, and E is the set of edges in the whole graph.
As GD, backward SGD, CLUSTER, GAS, and LMC share the same memory complexity of parameters θ ( l ) , we omit it in Table
[Table]
[Table]
C A DDITIONAL R ELATED W ORK C. M AIN DIFFERENCES BETWEEN LMC AND G RAPH FM
• First, LMC focuses on the convergence of subgraph-wise sampling methods, which is orthogonal to the idea of GraphFM-OB to alleviate the staleness problem of historical values.
The
advanced approach to alleviating the staleness problem of historical values can further improve the performance of LMC and it is easy to establish provable convergence by the extension of LMC.
• Second, LMC uses nodes in both mini-batches and their 1 -hop neighbors to compute incomplete up-to-date messages.
In contrast, GraphFM-OB only uses nodes in the mini-batches.
For the nodes whose neighbors are contained in the union of the nodes in mini-batches and their 1 -hop neighbors, the aggregation results of LMC are exact, while those of GraphFM-OB are not.
• Third, by noticing that aggregation results are biased and the I/O bottleneck for the history access, LMC does not update the historical values in the storage for nodes outside the mini-batches.
However, GraphFM-OB updates them based on the aggregation results.
D D ETAILED PROOFS
Notations.
Unless otherwise specified, C and C ′ with any superscript or subscript denotes constants.
We denote the learning rate by η .
In this section, we suppose that Assumption 1 holds.
D. P ROOF OF T HEOREM 1: UNBIASED MINI - BATCH GRADIENTS OF BACKWARD SGD
In this subsection, we give the proof of Theorem 1, which shows that the mini-batch gradients computed by backward SGD are unbiased.
Proof.
As V L B = V B ∩ V L is uniformly sampled from V L , the expectation of g w ( V B ) is
[Equation]
As the subgraph V B is uniformly sampled from V , the expectation of g θ l ( V B ) is
[Equation]
D. D IFFERENCES BETWEEN EXACT VALUES AT ADJACENT ITERATIONS
We first show that the differences between the exact values of the same layer in two adjacent iterations can be bounded by setting a proper learning rate.
Lemma 1.
Suppose that Assumption 1 holds.
Given an L -layer GNN, for any ε > 0 , by letting
[Equation]
we have ∥ H l,k +1 − H l,k ∥
F < ε, ∀ l ∈ [ L ] , k ∈ N ∗ .
Proof.
Since η ≤ ε (2 γ
[Equation]
L G <
[Equation]
1 G , we have
[Equation]
Then, because η ≤ ε (2 γ ) L G < ε (2 γ )
L − 1 G , we have
[Equation]
And so on, we have
[Equation]
Since (2 γ )
L − l > 1 , we have
[Equation]
Lemma 2.
Suppose that Assumption 1 holds.
Given an L -layer GNN, for any ε > 0 , by letting
[Equation]
, we have
[Equation]
∥ F < ε, ∀ l ∈ [ L ] , k ∈ N ∗ .
Proof.
Since η ≤
[Equation]
we have
[Equation]
Then, because η ≤ ε (2 γ )
[Equation]
<
[Equation]
2 G , we have
[Equation]
And so on, we have
[Equation]
Since (2 γ ) l − 1 > 1 , we have
[Equation]
D. H ISTORICAL VALUES AND TEMPORARY VALUES
Suppose that we uniformly sample a mini-batch V k B ⊂ V at the k -th iteration and |V k B | = S .
For the simplicity of notations, we denote the temporary node embeddings and auxiliary variables in the l -th layer by (cid:98) H l,k and (cid:98) V l,k , respectively, where
[Equation]
and
[Equation]
We abbreviate the process that LMC updates the node embeddings and auxiliary variables of V k B in the l -th layer at the k -th iteration as
[Equation]
For each v i ∈ V k B , the update process of v i in the l -th layer at the k -th iteration can be expressed by
[Equation]
where f θ l,k ,i and ϕ θ l +1 ,k ,i are the components for node v i of f θ l,k and ϕ θ l +1 ,k , respectively.
D.3. C ONVEX COMBINATION COEFFICIENTS
We first focus on convex combination coefficients
β i , i ∈ [ n ] .
For the simplicity of analysis, we assume β
i
=
β for i ∈
[ i ] .
The analysis of the case where ( β i ) ni
=1 are different from each other is the same.
Lemma 3.
Suppose that Assumption 1 holds.
For any ε > 0 , by letting
[Equation]
we have
[Equation]
Proof.
Since (cid:98) H l,k = (1 − β ) H l,k + β (cid:101) H l,k , we have
[Equation]
Hence letting β ≤ ε 2 G leads to
[Equation]
Lemma 4.
Suppose that Assumption 1 holds.
For any ε > 0 , by letting
[Equation]
we have
[Equation]
Proof.
Since (cid:98) H l,k = (1 − β ) H l,k + β (cid:101) H l,k , we have ∥
[Equation]
Hence letting β ≤ ε 2 G leads to
[Equation]
D.3. A PPROXIMATION ERRORS OF HISTORICAL VALUES
Next, we focus on the approximation errors of historical node embeddings and auxiliary variables
[Equation]
Lemma 5.
For an L -layer GNN
[Equation]
suppose that Assumption 1 holds.
Besides
[Equation]
we suppose that
( d l, 1 h ) 2 is bounded by G > 1 , ∀ l ∈
[ L ] , 2.
there exists N ∈ N ∗ such that
[Equation]
then there exist constants C ′∗ , 1 , C ′∗ , 2 , and C ′∗ , 3 that do not depend on k, l, N , and η , such that
[Equation]
where ρ = n − Sn < 1 , n = |V| , and S is number of sampled nodes at each iteration.
Proof.
We have
[Equation]
About the first term, for l ≥ 1 , we have
[Equation]
For l = 0 , we have
[Equation]
=
[Equation]
About the second term, for l ≥ 1 , we have
E
[Equation] [Equation]
+
[Equation]
f
[Equation] [Equation]
F ]
[Equation]
E
[Equation]
i
[Equation]
2 F
[Equation]
i
[Equation]
f
[Equation]
l +
[Equation] [Equation]
l
[Equation] [Equation]
F ]
[Equation]
i
[Equation]
h
[Equation] [Equation]
l
[Equation]
i − f θ l +1 ,k +1 ,i (
[Equation]
l,k +1 ) ⟩ ] ≤ E
[Equation]
l
[Equation]
i − h l +
[Equation]
k i ∥ 2 F ] +
[Equation] [Equation]
h l
[Equation] [Equation]
k i − f θ l +1 ,k +1 ,i (
[Equation]
l,k )
[Equation]
f θ l +1 ,k
+1 ,i ( H l,k ) − f θ l +1 ,k
+1 ,
i ( H l,k +1 )
∥ 2 F ] +
[Equation] [Equation]
h l +
[Equation]
k
i
[Equation]
h l +
[Equation]
k i , h l +
[Equation]
k i − f θ l +1 ,k +1 ,i ( H l,k +1 ) ⟩ ] ≤ E
[Equation]
i
[Equation]
l
[Equation] [Equation]
i
[Equation]
F ]
[Equation]
l
[Equation]
k i − f θ l +1 ,k +1 ,
[Equation]
( H l,
[Equation] [Equation]
2 F ]
[Equation] [Equation]
f θ l +1 ,k
+1 ,i ( H l,k ) − f θ l +1 ,k
+1 ,i (
[Equation]
l,k +1 )
∥ 2 F ]
[Equation] [Equation]
h l
[Equation] [Equation]
k i − f θ l +1 ,k +1 ,i ( H l,k +1 )
∥ F ] ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k
i
[Equation]
F ] +
[Equation]
γ
[Equation]
E [ ∥ θ l +1 ,k − θ l +1 ,k +1 ∥ 2 ] +
[Equation]
γ
[Equation]
E [ ∥ H l,k
[Equation]
H l,k +1 ∥ 2 F ]
[Equation] [Equation]
E [
[Equation]
l +1 ,k
[Equation]
l +1 ,k +
[Equation]
∥ H l,k
[Equation]
l,k
[Equation]
F ] ≤
[Equation] [Equation]
h l +1 ,k
i
[Equation] [Equation]
i
[Equation] [Equation]
G
[Equation]
N
[Equation]
N≤ E [
[Equation]
h l
[Equation]
1 ,k i − h l
+
,k
i
[Equation]
F ]
[Equation]
G
[Equation]
[Equation] [Equation]
G )
N.
For l = 0 , we have
E [ ∥ h l +1 ,k
i
[Equation]
f θ l +1 ,k
+1 ,
i ( H l,k +1 )
[Equation]
2 F ] ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k
i + h l +1 ,k i − f θ l +1 ,k
+1 ,
i ( H l,k +1 )
∥ 2 F ] ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k i ∥
[Equation]
F ]
[Equation] [Equation]
h l +1 ,k
i − f θ l +1 ,k +1 ,i ( H l,k +1 )
∥ 2 F ] + 2 E [ ⟨ h l +1 ,k
i − h l
+1 ,k
i , h l +1 ,k
i − f θ l +1 ,k +1 ,i ( H l,k +1 ) ⟩ ] ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k i ∥ 2 F ] + E [ ∥ h l +1 ,k
i − f θ l +1 ,k +1 ,i ( H l,k ) + f θ l +1 ,k +1 ,i ( H l,k ) − f θ l +1 ,k
+1 ,
i ( H l,k +1 )
∥ 2 F ] + 2 E [ ⟨ h l +1 ,k
i − h l
+1 ,k
i , h l +1 ,k
i − f θ l +1 ,k +1 ,i ( H l,k +1 ) ⟩ ] ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k i ∥ 2 F ] + 2 E
[ ∥ h l +1 ,k i − f θ l +1 ,k +1 ,
i ( H l,k ) ∥ 2 F ] + 2 E [ ∥ f θ l +1 ,k
+1 ,i ( H l,k ) − f θ l +1 ,k
+1 ,
i ( H l,k +1 )
∥ 2 F ] + 4 G E
[ ∥ h l +1 ,k i − f θ l +1 ,k +1 ,i ( H l,k +1 )
∥ F ] ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k i ∥ 2 F ] + 2 γ 2 E
[ ∥ θ l +1 ,k − θ l +1 ,k +1 ∥ 2 ] + 4 Gγ E
[ ∥ θ l +1 ,k − θ l +1 ,k +1 ∥ ] ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k i ∥ 2 F ] + 2 γ 2 G 2 η 2 + 4 G 2 γη, ≤ E [ ∥ h l +1 ,k
i − h l
+1 ,k i ∥ 2 F ] + 2 G 2 γ ( γ + 2) η
+ 4 G 2 γη.
Hence we have ( d l +1 ,k +1 h ) 2 ≤ ( n − S ) n ( d l +1 ,k h ) 2 + 2( n − S ) γ ( γ + 2) G 2 η + (cid:40) 0 , l = 0 , 2 γ 2 S ( d l,k +1 h ) 2 + 4 nγ ( γ + G )
N 2, l ≥ 1 .
Let ρ = n − Sn < 1 .
For l = 0 , we have
[Equation]
which leads to
[Equation]
Then, for l = 1 we have
[Equation]
where C 2 , 1 , C 2 , 2 , and C 2 , 3 are all constants.
Hence we have
[Equation]
which leads to
[Equation]
And so on, there exist constants C ′∗ , 1 , C ′∗ , 2 , and C ′∗ , 3 that are independent with η, k, l, N such that
[Equation]
Lemma 6.
For an L -layer GNN, suppose that Assumption 1 holds.
Besides, we suppose that
( d l, 1 h ) 2 is bounded by G > 1 , ∀ l ∈
[ L ] , 2.
there exists N ∈ N ∗ such that
[Equation]
then there exist constants C ′∗ , 1 , C ′∗ , 2 , and C ′∗ , 3 that are independent with k, l, ε ∗ , and η , such that
[Equation]
where ρ = n − Sn < 1 , n = |V| , and S is number of sampled nodes at each iteration.
Proof.
Similar to the proof of Lemma 5.
D. P ROOF OF T HEOREM 2: APPROXIMATION ERRORS OF MINI - BATCH GRADIENTS
In this subsection, we focus on the mini-batch gradients computed by LMC, i.e.,
[Equation]
and
[Equation]
where V k B is the sampled mini-batch and V kL B is the corresponding labeled node set at the k -th iteration.
We denote the mini-batch gradients computed by backward SGD by
[Equation]
and
[Equation]
Below, we omit the sampled subgraph V k B and simply write the mini-batch gradients as (cid:101) g w ( w k ) , (cid:101) g θ l ( θ l,k ) , g w ( w k ) , and g θ l ( θ l,k ) .
The approximation errors of gradients are denoted by
[Equation]
and
[Equation]
We restate Theorem 2 as follows.
Theorem 4.
For any k ∈ N ∗ and l ∈
[ L ] , the expectations of ∥ ∆ kw ∥ 22 ≜ ∥ (cid:101) g w ( w k ) − ∇ w L ( w k ) ∥and ∥
∆ kθ l ∥ 22 ≜ ∥ (cid:101) g θ l ( θ l,k ) − ∇ θ l L ( θ l,k ) ∥ 22 have the bias-variance decomposition
[Equation]
where
[Equation]
Suppose that Assumption 1 holds, then with η = O ( ε 2 ) and β i
= O ( ε 2 ) , i ∈ [ n ] , there exist C >and ρ ∈ (0 , 1) such that for any k ∈ N ∗ and l ∈
[ L ] , the bias terms can be bounded as
[Equation]
Hence we have
[Equation]
Lemma 7.
Suppose that Assumption 1 holds.
For any k ∈ N ∗ , the difference between (cid:101) g w ( w k ) and g w ( w k ) can be bounded as
[Equation]
Proof.
We have
[Equation]
Lemma 8.
Suppose that Assumption 1 holds.
For any k ∈ N ∗ and l ∈
[ L ] , the difference between (cid:101) g θ l ( θ l,k ) and g θ l ( θ l,k ) can be bounded as
[Equation]
Proof.
As ∥ Aa − Bb ∥ 2 ≤ ∥ A ∥ F ∥ a − b ∥ 2 + ∥
A − B ∥
F ∥ b ∥ 2 , we can bound ∥ (cid:101) g θ l ( θ l,k ) − g θ l ( θ l,k ) ∥by
[Equation]
≤
[Equation]
∇ θ l u θ l,k ( h l − 1 ,k j , m l − 1 ,k N ( v j ) , x j ) (cid:17) V l,kj − (cid:16) ∇ θ l u θ l,k
( h l − 1 ,k j , m l − 1 ,k N ( v j ) , x j ) (cid:17) V l,kj ∥≤ |V| max v i ∈V k B ∥ (cid:16)
[Equation]
θ l u θ l
[Equation]
k
[Equation]
l
[Equation] [Equation]
k
[Equation]
m l
[Equation] [Equation]
k
[Equation]
v
[Equation]
x
[Equation] [Equation]
l
[Equation]
k
[Equation]
θ l u θ l
[Equation]
k
[Equation]
l
[Equation]
k
[Equation]
, m l − 1 ,k N ( v j ) , x j ) (cid:17) V l,kj ∥≤ |V| max v i ∈V k
[Equation]
∥∇ θ l
[Equation]
θ
[Equation]
m l
[Equation]
x
[Equation]
l
[Equation]
∥
[Equation]
l
[Equation]
l
[Equation]
} ≤ |V| G ∥ V l
[Equation]
k − V l,k ∥ F + |V| Gγ ∥
H l,k − H l,k ∥ F .
Lemma 9.
For an L -layer ConvGNN, suppose that Assumption 1 holds.
For any N ∈ N ∗ , by letting
[Equation]
and β
[Equation] [Equation]
[Equation]
] , there exists G 2 , ∗ > 0 and ρ ∈ (0 , 1) such that for any k ∈ N ∗ we have
[Equation]
where
[Equation]
and
[Equation]
Proof.
By Lemmas 1 and 2 we know that
[Equation]
By Lemmas 3 and 4 we know that for any k ∈ N ∗ and l ∈
[ L ] we have
[Equation]
and
[Equation]
Thus, by Lemmas 5 and 6 we know that there exist C ′∗ , 1 , C ′∗ , 2 , and C ′∗ , 3 that do not depend on k, l, η, N such that for ∀ l ∈ [ L ] and k ∈ N ∗ hold
[Equation]
and
[Equation]
We can decompose ∥ ∆ kw ∥ 22 as
[Equation]
We take expectation of both sides of the above expression, leading to
[Equation]
16) where
[Equation]
Var
[Equation]
as
E
[Equation]
By Lemma 7, we can bound the bias term as
[Equation]
where G 2 , 1 = γ max { (cid:113) C ′∗ , 1 , (cid:113) C ′∗ , 2 , (cid:113) C ′∗ , 3 } .
Similar to Eq.
equation 16, we can decompose E [ ∥
∆ kθ l ∥ 22 ] as
[Equation]
where
Bias
[Equation]
By Lemma
[Equation]
we can bound the bias term as
[Equation]
18) where G 2 , 2 = √ 2 |V| G (1 + γ ) max { (cid:113) C ′∗ , 1 , (cid:113) C ′∗ , 2 , (cid:113) C ′∗ , 3 } .
Let G 2 , ∗ = max { G 2 , 1 , G 2 , 2 } , then we have
[Equation]
By letting ε = 1 N 13 and C = 2 G 2 , ∗ , we have
[Equation]
which leads to
[Equation]
and
[Equation]
Theorem 2 and Theorem 4 follow immediately.
D. P ROOF OF T HEOREM 3: CONVERGENCE GUARANTEES
In this subsection, we give the convergence guarantees of LMC.
We first give sufficient conditions for convergence.
Lemma 10.
Suppose that function f : R n → R is continuously differentiable.
Consider an optimization algorithm with any bounded initialization x 1 and an update rule in the form of
[Equation]
where η > 0 is the learning rate and d ( x k ) is the estimated gradient that can be seen as a stochastic vector depending on x k .
Let the estimation error of the gradient be ∆ k
=
d ( x k ) −∇ f ( x k ) .
Suppose that
the optimal value f ∗ = inf x f ( x ) is bounded; 2.
the gradient of f is γ -Lipschitz, i.e.,
[Equation]
there exists G 0 > 0 that does not depend on η such that
[Equation]
there exists N ∈ N ∗ and ρ ∈ (0 , 1) that do not depend on η such that
[Equation]
where G 0 is the same constant as that in Condition 3, then by letting η = min { 1 γ , 1 N 23 } , we have
[Equation]
where R is chosen uniformly from [ N ]
[Equation]
Proof.
As the gradient of f is γ -Lipschitz
[Equation]
we have
[Equation]
Then, we have
[Equation]
By taking expectation of both sides
[Equation]
we have
[Equation]
By summing up the above inequalities for k ∈ [ N ] and dividing both sides by Nη (1 − ηγ 2 )
[Equation]
we have
[Equation]
where the second inequality comes from ηγ > 0 and f ( x k ) ≥ f ∗ .
According to the above conditions, we have
[Equation]
Notice that
[Equation]
where R is uniformly chosen from [ N ] , hence we have
[Equation]
By letting η = min { 1 γ , 1 N 23 }
[Equation]
we have
[Equation]
Given an L -layer GNN, forllowing (Chen et al., 2018a)
[Equation]
we directly assume that:
the optimal value
[Equation]
is bounded by G > 1 ; 2.
the gradients of L with respect to parameters w and θ l , i.e., ∇ w L , ∇ θ l L are γ -Lipschitz for ∀ l ∈
[ L ] .
To show the convergence of LMC by Lemma
[Equation]
it suffices to show that
there exists
G
>
that does
not depend on η such that
[Equation]
for any N ∈ N ∗ , there exist G 2 > 0 and ρ ∈ (0 , 1) such that
[Equation]
by letting
[Equation]
and
[Equation]
Lemma 11.
Suppose that Assumption 1 holds, then
[Equation]
Proof.
We have
[Equation]
and
E
[Equation]
Lemma 12.
Suppose that Assumption 1 holds.
For any N ∈ N ∗ , there exist G 2 > 0 and ρ ∈ (0 , 1) such that
[Equation]
by letting
[Equation]
and
[Equation]
Proof.
By Eqs.
(17) and (18) we know that there exists G 2 , ∗ such that for any k ∈ N ∗ we have
[Equation]
and
[Equation]
where ρ = n − Sn < 1 is a constant.
Hence
[Equation]
and
[Equation]
g θ l ( θ l
[Equation]
k ) − g θ l (
[Equation]
l
[Equation]
k
[Equation] [Equation]
where G 2 = GG 2 , ∗ .
According to Lemmas 11 and 12, the conditions in Lemma 10 hold.
By letting
[Equation]
Theorem 3 follows immediately.
E M ORE E XPERIMENTS E. P ERFORMANCE ON S MALL D ATASETS
Figure 5 reports the convergence curves GD, GAS, and LMC for GCN on three small datasets, i.e.,
Cora, Citeseer, and PubMed from Planetoid (Yang et al., 2016).
LMC is faster than GAS, especially on the CiteSeer and PubMed datasets.
Notably, the key bottleneck on the small datasets is graph sampling rather than forward and backward passes.
Thus, GD is faster than GAS and LMC, as it avoids graph sampling by directly using the whole graph.
[Figure]
E. C OMPARISON IN TERMS OF T RAINING T IME PER E POCH
We evaluate the training time per epoch of CLUSTER, GAS, FM, and LMC in Table 6.
Compared with GAS, LMC additionally accesses historical auxiliary variables.
Inspired by GAS (Fey et al., 2021), we use the concurrent mini-batch execution to asynchronously access historical auxiliary variables.
Moreover, from the convergence analysis of LMC, we can sample clusters to construct fixed subgraphs at preprocessing step (Line 2 in Algorithm 1) rather than sample clusters to construct various subgraphs at each training step 2 .
This further avoids sampling costs.
Finally, the training time per epoch of LMC is comparable with GAS.
CLUSTER is slower than GAS and LMC, as it prunes edges in forward passes, introducing additional normalization operation for the adjacency matrix of the sampled subgraph by [ A V B ] i,j / (cid:112) deg V B ( i ) deg V B ( j ) , where deg V B ( i ) is the degree in the sampled subgraph rather than the whole graph.
The normalized adjacency matrix is difficult to store and reuse, as the sampled subgraph may be different.
FM is slower than other methods, as they additionally update historical embeddings in the storage for the nodes outside the mini-batches.
CLUSTER, GAS, FM, and LMC.
[Table] [Table]
E. C OMPARISON IN TERMS OF M EMORY UNDER DIFFERENT B ATCH S IZES
In Table 7, we report the GPU memory consumption, and the proportion of reserved messages (cid:80) bk
=1 ∥ (cid:101)
A alg V bk ∥ 0 / ∥ (cid:101)
A ∥ 0 in forward and backward passes for GCN, where (cid:101)
A is the adjacency matrix of full-batch GCN, (cid:101)
A alg V bk is the adjacency matrix used in a subgraph-wise method alg (e
[Table]
g., CLUSTER,
GAS, and LMC), and ∥ · ∥ 0 denotes the ℓ 0 -norm.
As shown in Table 7, LMC makes full use of all sampled nodes in both forward and backward passes, which is the same as full-batch GD.
Default indicates the default batch size used in the codes and toolkits of GAS (Fey et al., 2021).
[Table]
E. A BLATION ABOUT β i
As shown in Section A.4, β i = score ( i ) α in LMC.
We report the prediction performance under α ∈ { 0 . 0 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 . 0 } and score ∈ { f ( x ) = x 2 , f ( x ) =
2 x − x 2 , f ( x ) =
x, f ( x ) = 1; x = deg local ( i ) /deg
global ( i ) } in Tables 8 and 9 respectively.
When exploring the effect of a specific hyper-parameter, we fix the other hyper-parameters as their best values.
Notably, α =implies that LMC directly uses the historical values as affordable without alleviating their staleness, which is the same as that in GAS.
Under large batch sizes, LMC achieves the best performance with large β i = 1 , as large batch sizes improve the quality of the incomplete up-to-date messages.
Under small batch sizes, LMC achieves the best performance with small β i = 0 .
4 score 2 x − x 2 ( i ) , as small learning rates alleviate the staleness of the historical values.
[Table]
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Table]
[Table] [Table]
[Table] [Table]
[Table]
F P OTENTIAL S OCIETAL I MPACTS
In this paper, we propose a novel and efficient subgraph-wise sampling method for the training of GNNs, i.e., LMC.
This work is promising in many practical and important scenarios such as search engine, recommendation systems, biological networks, and molecular property prediction.
Nonetheless, this work may have some potential risks.
For example, using this work in search engine and recommendation systems to over-mine the behavior of users may cause undesirable privacy disclosure.