SAMPLING-BASED GNN TRAINING
Year: 2021

Abstract:
Neighbor sampling is a commonly used technique for training Graph Neural Networks (GNNs) on large graphs. Previous work has shown that samplingbased GNN training can be considered as Stochastic Compositional Optimization (SCO) problems and can be better solved by SCO algorithms. However, we find that SCO algorithms are impractical for training GNNs on large graphs because they need to store the moving averages of the aggregated features of all nodes in the graph. The moving averages can easily exceed the GPU memory limit and even the CPU memory limit. In this work, we propose a variant of SCO algorithms with sparse moving averages for GNN training. By storing the moving averages in the most recent iterations, our algorithm only requires a fixed size buffer, regardless of the graph size. We show that our algorithm preserves the convergence rate of the original SCO algorithm when the buffer size satisfies certain conditions. Our experiments validate our theoretical results and show that our algorithm outperforms the traditional Adam SGD for GNN training with a small memory overhead.

1 INTRODUCTION
Graph Neural Networks (GNNs) have become the state-of-the-art models for machine learning tasks on graph-structured data. By recursively aggregating the features of neighboring nodes, GNNs learn an embedding of the nodes and use the embedding for downstream tasks such as node classification (Kipf & Welling, 2017; Duran & Niepert, 2017) or link prediction (Zhang & Chen, 2017; 2018).
Due to the recursive neighbor aggregation, training GNNs on large graphs is computationally challenging. To alleviate the computation burden, various neighbor sampling methods have been proposed (Hamilton et al., 2017; Ying et al., 2018; Chen et al., 2018; Zou et al., 2019; Li et al., 2018; Chiang et al., 2019; Zeng et al., 2020). The idea is to compute an unbiased estimation of the aggregation result in each layer based on a sampled subset of neighbors. These sampling techniques enable GNN training on large graphs. However, due to the composition of the aggregation functions in multiple layers, the stochastic gradient obtained with sampled neighbor aggregation is not an unbiased estimation of the true gradient, which undermines the convergence property of SGD-based training algorithms.
Previous work has shown that sampling-based GNN training is actually a Stochastic Compositional Optimization (SCO) problem (Cong et al., 2020; 2021). Cong et al. (2021) show that SCO algorithms can achieve faster convergence than the commonly used Adam SGD for GNN training on small graphs. Despite their good convergence property, SCO algorithms are not widely adopted for GNN training due to two reasons. First, although SCO algorithms achieve smaller training losses, the obtained GNN models usually have poor generalization – the validation and test accuracy are lower than the models trained by Adam SGD. Second, SCO algorithms need to maintain the moving averages of aggregation results of all nodes in the graph. For large graphs, the moving averages may exceed the memory capacity of the GPU. While it is possible to store the moving averages in CPU memory, copying the data from CPU to GPU in each iteration is expensive, which may negate the benefits of the faster convergence of SCO algorithms.
To address the above issues, we propose a Sparse Stochastic Compositional (SpSC) gradient method in this work. Our main idea is to store the moving averages for nodes sampled in the most recent
iterations instead of all nodes. As only a small number of nodes are stored, our algorithm has small memory consumption even for large graphs. We provide a convergence analysis on SpSC and show that, when the number of stored iterations satisfies certain constraints, SpSC can preserve the asymptotic convergence rate of the original SCO algorithm. In practice, the sparse moving averaging slightly slows down the convergence of SCO algorithm, but it surprisingly overcomes its poor generalization problem and achieves higher accuracy for sampling-based GNN training. Compared with Adam SGD, our algorithm incurs a small overhead for updating the moving averages in each iteration, but the overhead can be easily justified by the faster convergence of our algorithm.
Our experiments with two GNN models on different input graphs validate our theoretical results and show that our algorithm achieves higher accuracy than Adam SGD with the same or less amount of training time.
2 BACKGROUND AND MOTIVATION
To facilitate our discussion, we first give background on sampling-based GNN training and its relation to stochastic compositional optimization.
2.1 GNN COMPUTATION
The computation at each layer of a GNN is conducted in two steps: aggregate and update. For each node v in the graph, the aggregate function gathers data from its neighboring nodes and returns the aggregation result as
zv = Aggv(hne[v], xv, xne[v]). (1)
Here, hne[v] is the intermediate features of v’s neighbors from the previous layer, xv is the input feature of v, and xne[v] is the input features of v’s neighbors. The update function uses the aggregated value to produce the intermediate features of v as
hv = Updv(zv, xv). (2)
By stacking the intermediate features and the input features of all nodes, the computation at layer l can be written as
Z(l) = Agg(H(l−1), X), H(l) = Upd(Z(l), X,W (l)). (3)
Here, H(l−1) = [h(l−1)1 , . . . , h (l) N ] denotes the intermediate features of all nodes at layer l−1, Z(l) ∈ RN×dl is the aggregated features of all nodes at layer l, X = [x1, . . . , xN ] is the input features of all nodes, andW (l) is the learnable weights. As an example, Graph Convolutional Network (GCN) (Kipf & Welling, 2017) has Agg(H(l−1), X) = PH(l−1) where P is the normalized Laplacian matrix of the graph, and Upd(Z(l), X,W (l)) = σ(Z(l)W (l)) where σ is a non-linear activation function. Many other GNNs can be expressed in this form with different definitions of Agg and Upd (Zhou et al., 2018).
2.2 SAMPLING-BASED GNN TRAINING AS STOCHASTIC COMPOSITIONAL OPTIMIZATION
When the graph is large, the neigbhbor aggregation operation Agg incurs a large overhead, making the training of GNNs computationally challenging. Therefore, prior work has proposed to replace the Agg function with a sampled neighbor aggregation operation Ãgg. By sampling the neighboring nodes, an unbiased estimate of Z(l) is computed at each layer, i.e.,
Z̃(l) = Ãgg(H(l−1), X) (4)
with E[Z̃(l)] = Z(l). If we define the computation at layer l of the original GNN as a function
f (l)(Z(l−1),W (l−1), ...,W (T )) = [Z(l),W (l), ...,W (T )] (5)
=[Agg(Upd(Z(l−1), X,W (l−1))),W (l), ...,W (T )]],
the computation with sampled neighbor aggregation can be written as a stochastic function
f (l) ξl (Z̃(l−1),W (l−1), ...,W (T )) = [Z̃(l),W (l), ...,W (T )] (6)
=[Ãgg(Upd(Z̃(l−1), X,W (l−1))),W (l), ...,W (T )]]
where ξl represents the sampled neighbors at layer l. Since Z̃(l) is an unbiased estimate of Z(l), we have E[f (l)ξl ] = f (l), and the computation of a T -layer GNN can be written as
F (θ) = EξT+1 [ f (T+1) ξT+1 ( EξT [ f (T ) ξT ( ...Eξ1 [f (1)(θ)]... )])]
(7)
where θ = [X,W (1), ...,W (T )], f (T+1) is the loss function, and f (T+1)ξT+1 corresponds to the estimated loss with mini-batch sampling. Note that we put all the learnable weights in θ to formulate the computation as a stochastic compositional function. Our goal is to minimize F (θ), which is exactly a multi-level SCO problem.
2.3 LARGE MEMORY CONSUMPTION ISSUE WITH A NAIVE IMPLEMENTATION
SCO has been well studied in the past few years, and many algorithms with guaranteed convergence have been proposed (Zhang & Xiao, 2019; Yang et al., 2019; Chen et al., 2020; Yang et al., 2019; Balasubramanian et al., 2020; Chen et al., 2020; Lian et al., 2017; Wang et al., 2017b; Ghadimi et al., 2020). It seems straightforward to adopt these SCO algorithms to achieve faster training of GNNs. However, these algorithms have large memory consumption when applied to GNN training and cannot run on GPUs for large graphs.
To see the problem, let us consider the implementation of the SCGD algorithm (Yang et al., 2019) for GNN training. Formally, the algorithm is written as
y (1) k+1 = (1− βk)y (1) k + βkf (1) ξ1,k (θk), (8)
y (l) k+1 = (1− βk)y (l) k + βkf (l) ξl,k (y (l−1) k+1 ), 2 ≤ l ≤ T, (9)
θk+1 = θk − αk∇f (1)ξ1,k(θk)∇f (2) ξ2,k (y (1) k ) . . .∇f (T+1) ξT+1,k (y (T ) k ). (10)
The key idea is to store an auxiliary variable y(i) to maintain the moving average of each composite function. Since f (l)ξl returns the exact values of W (l), ...,W (T ), we only need to maintain a moving average of Z̃(l) for each layer. The computation is shown in Figure 1a. The moving average of the aggregated features is stored in Y (l) with each row for one node. In each iteration, some nodes (rows) are sampled, and the estimated aggregation results Z̃(l) are merged into Y (l) based on Formula (8) and (9). For the nodes that are not sampled, we simply multiply the corresponding rows of Z̄(l) by (1 − βk). Since the number of rows in Y (l) is the number of nodes in the graph, Y (l) takes a lot of memory when the graph is large. For example, for training a 3-layer GCN on a graph with two million nodes, suppose the hidden state dimension dl = 512 and a floating point has 4 bytes, Y takes 3× 2M× 512× 4 = 12GB of memory. All of the existing SCO algorithms need to maintain this moving average, which impedes their application to large-scale GNN training.
3 SPARSE STOCHASTIC COMPOSITIONAL GRADIENT DESCENT
To reduce the memory consumption of SCO algorithms for GNN training, we propose a Sparse Stochastic Compositional (SpSC) gradient method. Instead of storing the moving averages of all nodes in the graph, we only store the moving averages of nodes that are sampled in the most recent iterations.
As shown in Figure 1b, we maintain a fixed size buffer of the moving averages. The buffer is divided into t chunks with each chunk for the Z̃(l) of one iteration. The size of each chunk is ml · dl where ml is the maximum number of the nodes that can be sampled at layer l and dl is the hidden state dimension. Initially, the buffer is empty. In every iteration, we first check if the sampled nodes are in the buffer. For the nodes that are found in the buffer, we collect the corresponding rows of the buffer and add them to Z̃(l)k based on Formula (8) and (9). For the nodes that are not found in the buffer, we multiply the corresponding rows of Z̃(l)k by βk. The updated Z̃ (l) k is then written to chunk-(k mod t). All the other chunks are multiplied by (1 − βk). Since the sampled nodes found in the buffer are updated to chunk-(k mod t), the original values in chunk-0 and chunk-1 are invalidated, as shown by the shadowed rows in Figure 1b. As the buffer size is a constant (T · t ·ml · dl) regardless of the graph size, our algorithm can be employed to train GNN on very large graphs.
Our algorithm overwrites chunk-(k mod t) in iteration k. The information of the overwritten nodes is lost. The update of the moving averages can be written as
y (l) k+1 = (1− βk)y (l) k + βkf (l) ξl,k (y (l−1) k+1 )− k−1∏ j=k−t+1 (1− βj)u(l)k (11)
where u
(l) k = P (ξl,k−t/(ξl,k−t+1 ∪ · · · ∪ ξl,k))y (l) k−t+1. (12)
P (ξl,k−t/(ξl,k−t+1 ∪ · · · ∪ ξl,k)) is a matrix with binary values representing the overwritten nodes in the current iteration k, i.e., the nodes that are sampled in iteration k − t and are not sampled in the following t iterations. Since we only store the moving averages of nodes that are sampled in the most recent t iterations, the information of the overwritten nodes is lost. These nodes are multiplied by (1− βj) in every iteration after iteration k − t. The values of the overwritten rows are∏k−1 j=k−t+1(1− βj)u (l) k . Our algorithm simply replaces Formula (8) and (9) in the SCGD algorithm with Formula (11).
To study the convergence property of SpSC, we make the following assumptions that are commonly used in the analysis of SCO algorithms (Yang et al., 2019; Balasubramanian et al., 2020).
Assumption 1. The composite functions f (l) are Ll-smooth. That is, for any y and y′, we have ‖∇f (l)ξl (y)−∇f (l) ξl (y′)‖ ≤ Ll‖y − y′‖.
Assumption 2. The stochastic gradients of the composite functions f (l) are bounded in expectation, i.e., E[‖∇f (l)ξl (y)‖
2] ≤ C2l . Assumption 3. The estimated aggregation results obtained by sampled neighbor aggregation is unbiased, i.e., E[f (l)ξl (y)] = f (l)(y), and the stochastic gradient of f (l) is unbiased, i.e., E[∇f (l)ξl (y)] = ∇f (l)(y).
Following the single-timescale analysis of the algorithm (Balasubramanian et al., 2020), we use large batches for estimating the composite functions and assume that the estimation variance is small.
Assumption 4. The estimated aggregation results have small bounded variance, i.e., E[‖f (l)ξl,k(y)− f (l)(y)‖] ≤ βkV 2.
This is a reasonable assumption for GNN training on GPUs as we always sample a batch of nodes for neighbor aggregation to achieve better utilization of the GPU parallelism.
In additional to the conventional assumptions, we make an assumption on the moving averages.
Assumption 5. The moving average of the aggregated features are bounded, i.e., E[‖y(l)‖2] ≤ D2.
The convergence rate of our algorithm is summarized in the following theorem.
Theorem 1. Under Assumptions 1-5, if we set β ≤ 1 − 2(−1/(2T−1)) and α = β ·
min
( 1
2C1T
√ 2(1−β)2T−1−1
Cmax , 2∑T
l=1 A 2 l+16CmaxC 2 1T
) , the model parameters {θk} of our training
algorithm with (11) for updating sparse moving averages satisfy
1
K K−1∑ k=0 E [‖∇F (θk)‖]2 ≤O(β) + 4CmaxT 2 ( (1 + Cmax)(1− β)2(t−1) αβ2 + (1− β)2t+1 Tαβ2 ) D2
(13)
where Cmax = max(C22C 2 3 · · ·C2l , C23C24 · · ·C2l , · · · , C2l , 1) for l = 2 . . . T , t is the number of buffer chunks used in our algorithm.
The last term on the RHS of (13) reveals how the convergence of the SCGD algorithm is affected by the sparse moving averages. The larger t we use (i.e., the more chunks we have in the buffer), the smaller (1− β)2(t−1) and (1− β)2t+1we have, and the faster convergence we achieve. In theory, if we can make (1− β)2(t−1) = O(β4), the algorithm will achieve O( √ 1/K) convergence rate. As β → 0, we need larger t to maintain the convergence rate, and eventually, we will need to store the moving average of all nodes in the graph.
Applying Sparse Moving Averages to SCSC. Our sparse moving average can also be applied to other SCO algorithms. For example, the Stochastically Corrected Stochastic Compositional gradient method (SCSC) (Chen et al., 2020) has a correction term in the update of the moving averages. Because of the correction term, SCSC needs a relaxed assumption on the estimation error of the composite functions. More specifically, if we change (11) to
y (l) k+1 = (1− βk)y (l) k + f (l) ξl,k (y (l−1) k+1 )− (1− βk)f (l) ξl,k (y (l−1) k )− k−1∏ j=k−t+1 (1− βj)u(l)k (14)
and replace (8) and (9) with (14), we can relax Assumption 4 as
Assumption 6. The estimated aggregation results have bounded variance, i.e., E[‖f (l)ξl,k(y) − f (l)(y)‖] ≤ V 2.
The convergence rate of SCSC with sparse moving averages is summarized as follows.
Theorem 2. Under Assumptions 1-3 and 5-6, if we choose αk = α = cα√K and βk = β = cβ√ K , the model parameters {θk} of SCSC with (14) for updating sparse moving averages satisfy
1
K K−1∑ k=0 E[‖∇F (θk)‖2] ≤O (β) + 2(1 + β) (1− β)2(t−1) αβ T∑ l=1 D2l
+ 6 (1− β)2(t−1)
α Cu T∑ l=1 D2l (15)
where Cu = max ( 4C2l + γl ) for l = 1 . . . T .
The result suggests that, if we can set t such that (1− β)2(t−1) = O(β3), the algorithm will achieve O( √ 1/K) convergence rate.
4 IMPLEMENTATION DETAILS
Algorithm 1 describes an implementation of Formula (11) in our algorithm. For each layer, we allocate a buffer (buf ) of size tml × dl for the moving averages. The buffered nodes are maintained in a list node list ∈ Rtml where node list[i] is the index of the node stored at buf [i]. If buf [i] is empty, node list[i] is set to -1. In every iteration k, we first get the location of chunk-(k mod t). Then, we look up each of the sampled nodes in the buffer (line 3). The LookUp function computes
Algorithm 1: Updating sparse moving average of aggregated features at layer l in iteration k Input: Sampled nodes S, Buffered nodes node list ∈ Rtml , buf (l) ∈ Rtml×dl , Z̃(l)k ∈ R |S|×dl
// Get the location of chunk-(k mod t) 1 start = (k mod t) ∗ml; 2 end = chunk start+ |S|; // Look up the sampled nodes in the buffer 3 idx in buf, idx in z = LookUp(S, node list); // Update the moving average for the sampled nodes
4 Z̃ (l) k = βk ∗ Z̃ (l) k ; Z̃ (l) k [idx in z] = (1− βk) ∗ buf (l)[idx in buf ] + Z̃ (l) k [idx in z];
// Update the moving average for all buffered nodes
5 buf (l) = (1− βk) ∗ buf (l); buf (l)[start : end] = Z̃(l)k ; // Invalidate the old buffer for the sampled nodes 6 node list[idx in buf ] = −1; // Add the sampled nodes to node list 7 node list[start : end] = S;
the intersection of S and node list and returns the indices of the overlapping nodes in the two arrays. If a sampled node is not found in the buffer, we multiply the corresponding row of Z̃(l)k by βk. If a sampled node is in the buffer, we read in its current moving average and update the corresponding row of Z̃(l)k (line 4). For buffered nodes that are not sampled, we simply multiply their moving averages by (1−βk) (line 5). For buffered nodes that are sampled, we invalidate their original buffer by setting node list[idx in buf ] to −1 (line 6). Last, we add the sampled nodes to the node list. Most of the operations in Algorithm 1 are simple vector operations, and they incur little overhead. The performance bottleneck is the LookUp function. With a naive implementation, it has O(tml log |S|) time complexity, assuming S is sorted. In our implementation, we use an auxiliary array node loc ∈ RN to store the locations of all nodes in the buffer and accelerate the LookUp function. Specifically, if node-i is in the buffer, we store its location in buffer in node loc[i]; otherwise, node loc[i] is set to -1. With the auxiliary array, the idx in z can be obtained by comparing node loc[S] with zero, and the idx in buf is simply node loc[S][idx in z]. Before updating the node list at line 7 of Algorithm 1, we remove the overwritten nodes from node loc by setting node loc[node list[start : end]] to -1. Finally, we store the locations of the newly sampled nodes to node loc by setting node loc[S] to [start, start+ 1, . . . , end]. It is easy to see that all these operations have O(|S|) time complexity.
5 EVALUATION 5.1 EXPERIMENTAL SETUP
We conduct our experiments on a workstation with an Nvidia RTX 3090 GPU, an Intel Xeon Gold 6226R CPU, and 512GB RAM. Our code is implemented with PyTorch 1.8.0 and PyTorch Geometric 1.7.0.
We evaluate our algorithm on five graphs as listed in Table 1. The reddit and yelp graph are adopted from GraphSAINT (Zeng et al., 2020), and the arxiv, proteins, products are from the Open Graph Benchmark (Hu et al., 2020).
We apply our algorithm to two GNN models: GCN (Kipf & Welling, 2017) and GraphSAGE (Hamilton et al., 2017). Both models have three convolutional layers. We use Formula (11) for updating the moving average instead of Formula (14). This is because Formula (14) requires two forward passes which incurs extra overheads. The algorithm is run for 50 epochs. We set β to 0.2 initially, and decrease it to 0.1 at epoch 20, and further decrease it to 0.05 at epoch 40. The number of buffered chunks (t) is set to 8. We adopt the layer-wise sampling method in Zou et al. (2019) for neighbor sampling. The batch size is set to 4096, and the number of sampled neighbors in each layer is set to 8192.
5.2 TRAINING RESULTS
Validation Accuracy. Figure 2a shows the validation accuracy of different algorithms for training a GCN on arxiv. We can see that full neighbor aggregation (Adam Full) achieves the highest
accuracy. This is reasonable because full neighbor aggregation returns unbiased estimates of gradients. If neighbor sampling is used, the accuracy clearly drops with the same training algorithm (Adam Sample). Our algorithm (Sparse SCO) is able to improve the convergence of sampling-based GNN training and achieves almost the same accuracy as full neighbor aggregation. The results on reddit and yelp are similar. On products graph, Adam Full runs out of memory, so we only show the results of sampling-based training in Figure 2d. Interestingly, we find that our algorithm with neighbor sampling achieves even higher accuracy than Adam Full on proteins graph, as shown in Figure 2c. The original SCO algorithm (which stores the moving averages for all nodes in the graph) also has lower accuracy than Sparse SCO. This is probably due to the overfitting of models by Adam Full and SCO.
Training Loss. Figure 3 shows the training loss of different algorithms on different graphs. For reddit and yelp, we are able to run full neighbor aggregation on our GPU. As expected, Adam Full achieves the smallest training loss. Adam Sample, however, has the slowest convergence. SCO achieves training loss close to Adam Full. We run our algorithm with different t’s. The larger t we use, the smaller training loss we obtain. The results are consistent with our theoretical analysis and also suggesting that the poor accuracy of the original SCO algorithm is probably due to overfitting. For products graph, we are not able to run full neighbor aggregation. The results show a clearly faster convergence of our algorithm than Adam SGD for sampling-based training.
Test Accuracy. Table 2 lists the test accuracy of the models trained by different algorithms. For reddit and yelp, we follow the GraphSAINT paper (Zeng et al., 2020) and report the F1-micro score. For proteins, we follow the OGB (Hu et al., 2020) and report the ROC-AUC. We can see that our training algorithm achieves the highest test accuracy for both GCN and GraphSAGE on almost all the graphs. We do not include the results for GCN on yelp graph because its accuracy is apparently lower than GraphSAGE with all training algorithms, probably due to the limited expressiveness of the GCN model. While it is hard to draw a direct comparison with GraphSAINT because the sampling methods and the model architectures are different, our test accuracy on reddit and yelp matches the best accuracy reported by GraphSAINT (Zeng et al., 2020). It is worth noting that our algorithm achieves higher accuracy than Adam SGD with both full neighbor aggregation and neighbor sampling on proteins and products graph. The numbers are higher than the accuracy
Execution Time. Table 3 lists the time per epoch of different training algorithms on different graphs. Although full neighbor aggregation achieves good convergence in many cases, it incurs a much large computation overhead than sampled neighbor aggregation. The execution time of Adam Full is 7x to 55x longer than that of Adam Sample. Compared with Adam Sample, the SCO algorithm incurs a small overhead for updating the moving average of aggregation results. Our algorithm runs slightly slower than SCO because the LookUp operation in Algorithm 1 incurs an extra overhead. However, the execution time is still much smaller than full neighbor aggregation.
Memory Consumption. Figure 4 shows the memory consumption of different algorithms. We collect the numbers by calling the max memory allocated function in PyTorch at the end of the first epoch. For arxiv, reddit, proteins, and yelp, full neighbor aggregation takes much larger memory space than sampled aggregation. SCO and Sparse SCO require additional memory for storing the moving averages. Because we only store the moving average of nodes sampled in recent iterations, Sparse SCO uses less memory than SCO. On products graph, SCO requires extremely large memory due to the massive number of nodes, while our Sparse SCO can run with only 2GB of GPU memory.
6 RELATED WORK
To overcome the scalability limitation of GNN training, various neighbor sampling methods have been proposed, including node-wise sampling (Hamilton et al., 2017; Ying et al., 2018), layer-wise sampling (Chen et al., 2018; Huang et al., 2018; Zou et al., 2019), and subgraph sampling (Zeng et al., 2020; Chiang et al., 2019). These sampling-based training methods achieve good accuracy in practice (particularly on small graphs), but they lack theoretical justification.
Some recent works (Cong et al., 2020; 2021) point out that sampling-based GNN training is actually multi-level Stochastic Compositional Optimization (SCO). However, they either use this connection to justify their sampling techniques and fall back to Adam SGD for training (Cong et al., 2020), or they directly adopt an SCO algorithm without considering the large memory consumption issue (Cong et al., 2021).
The research on SCO traces back to (Ermoliev, 1976) where a two-timescale stochastic approximation scheme was proposed for two-level problems. Various SCO algorithms and convergence analyses have been proposed ever since (Wang et al., 2017a;b; Lian et al., 2017; Yang et al., 2019; Zhang & Xiao, 2019; Chen et al., 2020; Balasubramanian et al., 2020; Ghadimi et al., 2020; Hu et al., 2019; Lin et al., 2018). Despite a substantial volume of work on SCO in recent years, none of the existing work has considered the large memory consumption issue and the data movement overhead of the algorithms. Our work is the first to establish a convergence analysis for SCO algorithms with sparse moving averages.
7 CONCLUSION
In this work, we propose a new variant of SCO algorithm for training graph neural networks on large graphs. Our main idea is to maintain a sparse moving average of the aggregation results in each convolutional layer. We study the convergence property of our algorithm and show that the algorithm can achieve O( √ 1/K) convergence rate when a sufficient amount of moving averages are maintained. Our experiments with two GNN models on different graphs validate our theoretical results and show a clear advantage of our algorithm against Adam SGD for GNN training.
REFERENCES Ogb leaderboards for node classification, 2021. URL https://ogb.stanford.edu/docs/ leader_nodeprop/.
Krishnakumar Balasubramanian, Saeed Ghadimi, and Anthony Nguyen. Stochastic multi-level composition optimization algorithms with level-independent convergence rates. arXiv preprint arXiv:2008.10526, 2020.
Jie Chen, Tengfei Ma, and Cao Xiao. FastGCN: Fast learning with graph convolutional networks via importance sampling. In International Conference on Learning Representations, 2018.
Tianyi Chen, Yuejiao Sun, and Wotao Yin. Solving stochastic compositional optimization is nearly as easy as solving stochastic optimization. arXiv preprint arXiv:2008.10847, 2020.
Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. Cluster-gcn: An efficient algorithm for training deep and large graph convolutional networks. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 257–266, 2019.
Weilin Cong, Rana Forsati, Mahmut Kandemir, and Mehrdad Mahdavi. Minimal variance sampling with provable guarantees for fast training of graph neural networks. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1393–1403, 2020.
Weilin Cong, Morteza Ramezani, and Mehrdad Mahdavi. On the importance of sampling in learning graph convolutional networks. arXiv preprint arXiv:2103.02696, 2021.
Alberto Garcia Duran and Mathias Niepert. Learning graph representations with embedding propagation. In Advances in neural information processing systems, pp. 5119–5130, 2017.
Yu. Ermoliev. Methods of stochastic programming. Monographs in Optimization and OR, 1976. In Russian.
Saeed Ghadimi, Andrzej Ruszczynski, and Mengdi Wang. A single timescale stochastic approximation method for nested stochastic optimization. SIAM Journal on Optimization, 30(1):960–979, 2020.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in neural information processing systems, pp. 1024–1034, 2017.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint arXiv:2005.00687, 2020.
Wenqing Hu, Chris Junchi Li, Xiangru Lian, Ji Liu, and Huizhuo Yuan. Efficient smooth non-convex stochastic compositional optimization via stochastic recursive gradient descent. In Advances in Neural Information Processing Systems, pp. 6929–6937, 2019.
Wenbing Huang, Tong Zhang, Yu Rong, and Junzhou Huang. Adaptive sampling towards fast graph representation learning. In Advances in neural information processing systems, pp. 4558–4567, 2018.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (ICLR), 2017.
Ruoyu Li, Sheng Wang, Feiyun Zhu, and Junzhou Huang. Adaptive graph convolutional neural networks. In AAAI Conference on Artificial Intelligence, 2018.
Xiangru Lian, Mengdi Wang, and Ji Liu. Finite-sum composition optimization via variance reduced gradient descent. In Artificial Intelligence and Statistics, pp. 1159–1167, 2017.
Tianyi Lin, Chenyou Fan, and Mengdi Wang. Improved oracle complexity of variance reduced methods for nonsmooth convex stochastic composition optimization. arXiv, pp. arXiv–1802, 2018.
Mengdi Wang, Ethan X Fang, and Han Liu. Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. Mathematical Programming, 161(1-2): 419–449, 2017a.
Mengdi Wang, Ji Liu, and Ethan X Fang. Accelerating stochastic composition optimization. The Journal of Machine Learning Research, 18(1):3721–3743, 2017b.
Shuoguang Yang, Mengdi Wang, and Ethan X Fang. Multilevel stochastic gradient methods for nested composition optimization. SIAM Journal on Optimization, 29(1):616–659, 2019.
Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L Hamilton, and Jure Leskovec. Graph convolutional neural networks for web-scale recommender systems. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 974–983, 2018.
Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna. Graphsaint: Graph sampling based inductive learning method. In International Conference on Learning Representations, 2020.
Junyu Zhang and Lin Xiao. Multi-level composite stochastic optimization via nested variance reduction. arXiv preprint arXiv:1908.11468, 2019.
Muhan Zhang and Yixin Chen. Weisfeiler-lehman neural machine for link prediction. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 575–583, 2017.
Muhan Zhang and Yixin Chen. Link prediction based on graph neural networks. In Advances in Neural Information Processing Systems, pp. 5165–5175, 2018.
Jie Zhou, Ganqu Cui, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. Graph neural networks: A review of methods and applications. arXiv preprint arXiv:1812.08434, 2018.
Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layer-dependent importance sampling for training deep and large graph convolutional networks. In Advances in Neural Information Processing Systems, pp. 11249–11259, 2019.
8 APPENDIX
Link to our source code: https://anonymous.4open.science/r/iclr2022_ artifact-0FE6/
8.1 PROOF TO THEOREM 1
Based on the discussion in Section 3, our proposed algorithm can be written as
y (1) k+1 = (1− βk)y (1) k + βkf (1) ξ1,k
(θk)− k−1∏
j=k−t+1
(1− βj)u(1)k , (16)
y (l) k+1 = (1− βk)y (l) k + βkf (l) ξl,k (y (l−1) k+1 )− k−1∏ j=k−t+1 (1− βj)u(l)k , 2 ≤ l ≤ T, (17)
θk+1 = θk − αk∇k, (18)
with u
(l) k = P (ξl,k−t/(ξl,k−t+1 ∪ · · · ∪ ξl,k))y (l) k+t−1 (19)
and ∇k = ∇f (1)ξ1,k(θk) · · · ∇f (T ) ξT,k (y (T−1) k+1 )∇f (T+1) ξT+1,k (y (T ) k+1). (20)
8.1.1 SUPPORTING LEMMA
Lemma 3. The change of y(l)k in every iteration is bounded, i.e., E [∥∥∥y(1)k+1 − y(1)k ∥∥∥2] ≤ β2k1− βkE [∥∥∥f (1) (θk)− y(1)k+1∥∥∥2]+ 1βk k−1∏
j=k−t+1
(1− βj)2 E [∥∥∥u(1)k ∥∥∥2]+ β3kV 2
(l = 1) (21)
and E [∥∥∥y(l)k+1 − y(l)k ∥∥∥2] ≤ β2k1− βkE [∥∥∥f (l) (y(l−1)k+1 )− y(l)k+1∥∥∥2]+ 1βk k−1∏
j=k−t+1
(1− βj)2 E [∥∥∥u(l)k ∥∥∥2]+ β3kV 2
(l ≤ 2) (22)
Proof. According to (17), we have( y
(l) k+1 − y (l) k ) =βk ( f (l) ξl,k ( y (l−1) k+1 ) − y(l)k ) − k−1∏ j=k−t+1 (1− βj)u(l)k
=βk
( f (l) ( y
(l−1) k+1 ) − y(l)k ) + βk ( f (l) ξl,k (y (l−1) k+1 )− f (l)(y (l−1) k+1 ) ) − k−1∏ j=k−t+1 (1− βj)u(l)k
(23) Taking the square of both sides and using Assumption 4, we have E [∥∥∥y(l)k+1 − y(l)k ∥∥∥2] = E  ∥∥∥∥∥∥βk ( f (l) ( y (l−1) k+1 ) − y(l)k ) − k−1∏ j=k−t+1 (1− βj)u(l)k ∥∥∥∥∥∥ 2 + β3kV 2
≤ β 2 k 1− βk E [∥∥∥f (l) (y(l−1)k+1 )− y(l)k+1∥∥∥2]+ 1βk k−1∏ j=k−t+1 (1− βj)2 E [∥∥∥u(l)k ∥∥∥2]+ β3kV 2
(24)
Lemma 4. The difference between the stochastic gradient and the true gradient is bounded, i.e.,
‖E [∇k|Fk]−∇F (θk)‖ ≤ T+1∑ l=2 l−1∑ m=1 Am,lE [∥∥∥y(m)k+1 − f (m) (y(m−1)k+1 )∥∥∥ |Fk]
:= T∑ l=1 AlE [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥ |Fk]
(25)
where Al := ∑T m=l+1Al,m.
This lemma is commonly used in the analysis of SCO algorithms Chen et al. (2020); Yang et al. (2019). Proof can be found in Chen et al. (2020), page 20-23.
Lemma 5. The difference between the composite function and its moving average satisfies
K−1∑ k=0 T∑ l=1 E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2]
≤4CmaxC 2 1Tα 2
β2
K−1∑ k=0 E [ ‖∇F (θk)‖2 ] + 2CmaxT 2K ( (1− β)2t−1
β3 + Cmax(1− β)2(t−1) β3 + (1− β)2t+1 Tβ3
) D2
+ 2CmaxT 2K ( βCmax +
β2(1− β)(1 + T ) T
) V 2, (26)
where Cmax = max(C22C 2 3 · · ·C2l , C23C24 · · ·C2l , · · · , C2l , 1) for l = 2 . . . T .
Proof. According to (17), we have
y (l) k+1 − f
(l) ( y
(l−1) k+1 ) = (1− βk) ( y (l) k − f (l) ( y (l−1) k+1 )) + βk ( f (l) ξl,k ( y (l−1) k+1 ) − f (l) ( y (l−1) k+1 )) − k−1∏ j=k−t+1 (1− βj)u(l)k
= (1− βk) ( y (l) k − f (l) ( y (l−1) k )) + (1− βk) ( f (l) ( y (l−1) k ) − f (l) ( y (l−1) k+1 )) + βk ( f
(l) ξl,k
( y
(l−1) k+1
) − f (l) ( y
(l−1) k+1
)) − k−1∏ j=k−t+1 (1− βj)u(l)k
= (1− βk) ( y (l) k − f (l) ( y (l−1) k )) + f (l) ( y (l−1) k ) − f (l) ( y (l−1) k+1 ) ︸ ︷︷ ︸
:=Y (l) k
+ βk
( f
(l) ξl,k
( y
(l−1) k+1
) − f (l) ( y
(l−1) k
)) − k−1∏ j=k−t+1 (1− βj)u(l)k .
(27)
Taking the square of both sides of (27) and taking the expectation conditioned on Fk,n :={ Fk, y(1)k+1, · · · , y (l−1) k+1 } , we have:
E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2]
=E  ∥∥∥∥∥∥(1− βk) ( y (l) k − f (l) ( y (l−1) k )) + Y (l) k − k−1∏ j=k−t+1 (1− βj)u(l)k ∥∥∥∥∥∥ 2 
+ β2kE [∥∥∥(f (l)ξl,k (y(l−1)k+1 )− f (l) (y(l−1)k ))∥∥∥2]
(Assumption4) ≤ (1− βk)E [∥∥∥y(l)k − f (l) (y(l−1)k )∥∥∥2]+ 1βk · 11− βkE [∥∥∥Y (l)k ∥∥∥2] + 1
β2k k∏ j=k−t+1 (1− βj)2 E [∥∥∥u(l)k ∥∥∥2]+ β3kV 2.
(28)
Let us define
S (l) k+1 = E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2] , (29) Formula (28) can be rewritten as
S (l) k+1 ≤(1− βk)S (l) k +
1 βk · 1 1− βk E [∥∥∥Y (l)k ∥∥∥2]
+ 1
β2k k∏ j=k−t+1 (1− βj)2 E [∥∥∥u(l)k ∥∥∥2]+ β3kV 2. (30)
Because
E [∥∥∥Y (l)k ∥∥∥2] = E [∥∥∥f (l)(y(l−1)k+1 )− f (l)(y(l−1)k )∥∥∥2]
(Assumption2) ≤ C2l E [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥2] , (31)
when l = 1,
E [∥∥∥Y (1)k ∥∥∥2] ≤ C21E [‖θk+1 − θk‖2]
≤ C21E [ ‖αk∇k + αk∇F (θk)− αk∇F (θk)‖2 ] ≤ C21 · α2k · (2E [ ‖∇F (θk)‖2 ] + 2T
T∑ l=1 A2l S (l) k+1).
(32)
plugging (32) into (30), we have
S (1) k+1 ≤(1− βk)S (1) k +
1
β2k k−1∏ j=k−t+1 (1− βj)2E [∥∥∥u(1)k ∥∥∥2]+ β3kV 2
+ 1 βk · 1 1− βk · E [∥∥∥Y (1)k ∥∥∥2]
≤(1− βk)S(1)k + 1
β2k k−1∏ j=k−t+1 (1− βj)2E [∥∥∥u(1)k ∥∥∥2]+ β3kV 2
+ 2 βk · 1 1− βk · C21α2kE ‖∇F (θk)‖ 2 + 2 βk · 1 1− βk · C21α2kT T∑ l=1 A2l S (l) k+1.
(33)
When k = 0, ||u10||2 = 0, and
S (1) 1 ≤ (1− β0)S (1) 0 + β0 (β 2 0V
2 + 2 β20 · 1 1− β0 C21α 2 0 ‖∇F (θ0)‖ 2 + 2 β20 · 1 1− β0 C21α 2 0T T∑ l=1
S (l) 1 )︸ ︷︷ ︸
:=∆ (1) 0
= (1− β0)S(1)0 + β0∆ (1) 0 .
(34)
Similarly, we have
S (1) 2 ≤ (1− β1)S (1) 1 + β1∆ (1) 1
...
S (1) k+1 ≤ (1− βk)S (1) k + βk∆ (1) k .
(35)
Here, for i > t,
∆ (1) i =
1
β3i i−1∏ j=i−t−1 (1− βj)2E [∥∥∥u(1)i ∥∥∥2]+ β2i V 2
+ 2 β2i · 1 1− βi C21α
2 iE [ ‖∇F (θi)‖2 ] + 2
β2i · 1 1− βi C21α 2 iT T∑ l=1 S (l) i+1;
(36)
for 0 < i ≤ t,
∆ (1) i =β 2 i V
2 + 2 β2i · 1 1− βi C21α
2 iE [ ‖∇F (θi)‖2 ] + 2
β2i · 1 1− βi C21α 2 iT T∑ l=1 S (l) i+1.
(37)
Because S(1)0 = 0, we can recursively write:
S (1) 1 ≤β0∆ (1) 0
S (1) 2 ≤(1− β1)β0∆ (1) 0 + β1∆ (1) 1
S (1) 3 ≤(1− β2)(1− β1)β0∆ (1) 0 + (1− β2)β1∆ (1) 1 + β2∆ (1) 2
...
S (1) k+1 ≤(1− βk) · · · (1− β1)β0∆ (1) 0 + (1− βk) · · · (1− β2)β1∆ (1) 1 + · · ·
+ (1− βk)βk−1∆(1)k−1 + βk∆ (1) k
= k∑ i=1 wi,k∆ (1) i ,
(38)
where wi,k = (1− βk) · · · (1− βi+1)βi. For l > 1, (30) can be rewritten as
S (l) k+1 ≤ (1− βk)S (l) k + C2l βk (1− βk) E [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥2]+ 1β2k k−1∏ j=k−t+1 (1− βj)2 E [∥∥∥u(l)k ∥∥∥2]+ β3kV 2
(Lemma 3) ≤ (1− βk)S(l)k + βkC 2 l
(1− βk)2 S
(l−1) k+1 + βkZ (l) k ,
(39)
where
Z (l) k =
1
β3k k−1∏ j=k−t+1 (1− βj)2 E [∥∥∥u(l)k ∥∥∥2]+ C2lβ3k (1− βk) k−1∏ j=k−t+1 (1− βj)2 E [∥∥∥u(l−1)k ∥∥∥2]+ β2kV 2 + βkC2l1− βk V 2.
(40)
It follows that
S (2) k+1 ≤(1− βk)S (2) k +
βkC 2 2
(1− βk)2 S
(1) k+1 + βkZ (2) k
Accordingto(38)
≤ (1− βk)S(2)k + βkC
2 2 (1− βk)2 k∑ i=0 wi,k∆ (1) i + βkZ (2) k
=(1− βk)S(2)k + βk
( C22
(1− βk)2 k∑ i=0 wi,k∆ (1) i + Z (2) k
) (41)
Following the same steps as (38), by recursively plugging S(2)k into S (2) k+1, we can obtain
S (2) k+1 ≤ k∑ j=0 wj,k
( C22
(1− βj)2 j∑ i=0 wi,k∆ (1) i + Z (2) j
)
= k∑ j=0 C22wj,k (1− βj)2 j∑ i=0 wi,k∆ (1) i + k∑ j=0 wj,kZ (2) j
= k∑ i=0  k∑ j=i C22wj,k (1− βj)2 wi,k∆(1)i + k∑ i=0 wi,kZ (2) i
(42)
According to (38), it is easy to find ∑k j=i wj,k < 1. Thus, we have
k∑ j=i C22wj,k (1− βj)2 ≤ k∑ j=i C22wj,k (1− βi)2 ≤ C 2 2 (1− βi)2 k∑ j=i wj,k ≤ C22 (1− βi)2 . (43)
Combining (43) and (42), we have
S (2) k+1 ≤ k∑ i=0 ( C22 (1− βi)2 wi,k∆ (1) i + wi,kZ (2) i )
≤ k∑ i=0 wi,k ( C22 (1− βi)2 ∆ (1) i + Z (2) i )
= k∑ i=0 wi,k∆ (2) i
(44)
Similarly, for any l > 1, we have
S (l) k+1 ≤ k∑ i=0 wi,k∆ (l) i (45)
where
∆ (2) i = C22 (1− βi)2 ∆ (1) i + Z (2) i
∆ (3) i = C23 (1− βi)2 ∆ (2) i + Z (3) i
= C22C 2 3
(1− βi)4 ∆
(1) i + C23 (1− βi)2 Z (2) i + Z (3) i
∆ (l) i =
C22C 2 3 · · ·C2l
(1− βi)2(l−1) ∆
(1) i + C23 · · ·C2l (1− βi)2(l−2) Z (2) i + C24 · · ·C2l (1− βi)2(l−3) Z (3) i + · · ·+ Z (l) i
(46)
LetCmax = max(C22C 2 3 · · ·C2l , C23C24 · · ·C2l , · · · , C2l , 1),∀l ∈ [1, T ] and Zmaxi defined as follows:
Zmaxi = 1
β3i i−1∏ j=i−t+1 (1− βj)2D2 + Cmax β3i (1− βi) i−1∏ j=i−t+1 (1− βj)2D2 + β2i V 2 + βiCmax 1− βi V 2
(47)
Plugging (46) into (45), we can obtain
S (l) k+1 ≤ k∑ i=0 Cmax (1− βi)2(T−1) wi,k∆ (1) i + k∑ i=0 wi,k l∑ j=2 Cmax (1− βi)2(T−1) Z (l) i . (48)
Summing (48) from l = 1 to T , setting βk = β and αk = α, then substituting ∆ (1) k we have: K−1∑ k=0 T∑ l=1 S (l) k+1 ≤ K−1∑ k=0 k∑ i=0 T∑ l=1 Cmax (1− β)2(T−1) wi,k∆ (1) i + K−1∑ k=0 k∑ i=0 wi,k T∑ l=1 l∑ j=2 Cmax (1− β)2(T−1) Z (l) i
≤ K−1∑ k=0 K−1∑ i=k
CmaxT
(1− β)2(T−1) wi,K−1∆
(1) k + K−1∑ k=0 K−1∑ i=k CmaxT 2 (1− β)2(T−1) wi,K−1Z max i
≤ K−1∑ k=0
CmaxT
(1− β)2(T−1) ∆
(1) k + K−1∑ k=0 K−1∑ i=k CmaxT 2 (1− β)2(T−1) wi,K−1Z max i
≤ CmaxT (1− β)2(T−1) K−1∑ k=0
( β2V 2 + 2C21α 2 β2(1− β) E [ ‖∇F (θk)‖2 ] + 2C21α 2T β2(1− β) T∑ l=1 S (l) k+1 )
+ CmaxTD
2K
β3(1− β)2(T−1−t) + CmaxT
2KZmax
(1− β)2(T−1) .
(49)
It follows that
K−1∑ k=0 T∑ l=1 S (l) k+1 ≤
CmaxTβ 2(1− β) β2(1− β)2T−1 − 2C21Cmaxα2T 2 K−1∑ k=0 ( β2V 2 + 2C21α 2 β2(1− β) E [ ‖∇F (θk)‖2 ]) + CmaxTD 2K(1− β)2t+1
β3(1− β)2T−1 − 2C21CmaxT 2α2β +
CmaxT 2β2(1− β)KZmax
β2(1− β)2T−1 − 2C21Cmaxα2T 2 (50)
If we set β ≤ 1 − 2T−1 √
1 2 and α ≤ β 2C1T
√ 2(1−β)2T−1−1
Cmax , we can ensure β2(1 − β)2T−1 −
2C21Cmaxα 2T 2 ≥ β2/2, and (50) can be simplified as
K−1∑ k=0 T∑ l=1 S (l) k+1 ≤2CmaxTKβ 2(1− β)2V 2 + 4CmaxC 2 1Tα 2 β2 K−1∑ k=0 E [ ‖∇F (θk)‖2 ] + 2CmaxTK(1− β)2t+1
β3 D2 + 2CmaxT 2K(1− β)Zmax
≤2CmaxTKβ2(1− β)2V 2 + 4CmaxC
2 1Tα 2
β2
K−1∑ k=0 E [ ‖∇F (θk)‖2 ] + 2CmaxTK(1− β)2t+1
β3 D2 + 2CmaxT
2K(1− β) ( β2 +
βCmax 1− β
) V 2
+ 2CmaxT
2K(1− β)2t−1
β3 (1 + Cmax 1− β )D2
≤4CmaxC 2 1Tα 2
β2
K−1∑ k=0 E [ ‖∇F (θk)‖2 ] + 2CmaxT 2K ( (1− β)2t−1
β3 + Cmax(1− β)2(t−1) β3 + (1− β)2t+1 Tβ3
) D2
+ 2CmaxT 2K ( βCmax +
β2(1− β)(1 + T ) T
) V 2
(51)
8.1.2 REMAINING STEPS TOWARDS THEOREM 1.
Using the smoothness of F (θk), we have
F (θk+1) ≤ F (θk) + 〈∇F (θk), θk+1 − θk〉+ L
2 ‖θk+1 − θk‖2
= F (θk)− αk 〈∇F (θk), gk+1〉+ L
2 α2k ‖gk+1‖
2
= F (θk)− αkE [ ‖∇F (θk)‖2 ] + Lα2k
2 ‖gk+1‖2 + αk 〈∇F (θk),∇F (θk)− gk+1〉
= F (θk)− αkE [ ‖∇F (θk)‖2 ] + Lα2k
2 ‖gk+1‖2 + αk 〈∇F (θk),∇F (θk)−∇k〉+ αk 〈∇F (θk),∇k − gk+1〉 .
(52)
Conditioning on Fk and taking expectation of both sides w.r.t. ξk, we have
E [F (θk+1) |Fk] ≤F (θk)− αkE [ ‖∇F (θk)‖2 ] + L 2 E [ ‖θk+1 − θk‖2 | Fk ] + αk 〈∇F (θk),E [∇F (θk)−∇k|Fk]〉
≤F (θk)− αkE [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2Tα2k + αkE [‖∇F (θk)‖] ‖E [∇k|Fk]−∇F (θk)‖
(Lemma 4) ≤ F (θk)− αkE [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2Tα2k + αk T∑ l=1 (Al ‖∇F (θk)‖E [∥∥∥y(l)k+1 − f (l)(y(l−1)k+1 )|Fk∥∥∥]
≤F (θk)− αk ( 1− αk
4βk T∑ l=1 A2l
) E [ ‖∇F (θk)‖2 ] + βk
T∑ l=1 E [∥∥∥y(l)k+1 − f (l)(y(l−1)k+1 )∥∥∥2 |Fk]+ L2C21 · · ·C2Tα2k.
(53)
Summing from k = 0 to K − 1, we have K−1∑ k=0 E [F (θk+1|Fk)] ≤ K−1∑ k=0 F (θk)− K−1∑ k=0 αk(1− αk 4βk T∑ l=1 A2l )E [‖∇F (θk)‖]2
+ K−1∑ k=0 βk T∑ l=1 S (l) k+1 + K−1∑ k=0 L 2 C21 · · ·C2Tα2k.
(54)
Substitute (50) into the above equation, we have K−1∑ k=0 E [F (θk+1|Fk)] ≤ K−1∑ k=0 F (θk)− α(1− α 4β T∑ l=1 A2l ) K−1∑ k=0 E [‖∇F (θk)‖]2
+ 4CmaxC
2 1Tα 2
β
K−1∑ k=0 E [ ‖∇F (θk)‖2 ] + LK 2 C21 · · ·C2Tα2
+ 2CmaxT 2K
( (1− β)2t−1
β2 + Cmax(1− β)2(t−1) β2 + (1− β)2t+1 Tβ2
) D2
+ 2CmaxT 2K ( β2Cmax +
β3(1− β)(1 + T ) T
) V 2.
(55)
If we set α ≤ 2β∑T l=1 A 2 l+16CmaxC 2 1T , we can ensure 1− α4β ∑T l=1A 2 l − 4CmaxC 2 1Tα β ≥ 1 2 , and thus, K−1∑ k=0 E [F (θk+1|Fk)] ≤ K−1∑ k=0 F (θk) + α 2 K−1∑ k=0 E [‖∇F (θk)‖]2 + LK 2 C21 · · ·C2Tα2
+ 2CmaxT 2K
( (1− β)2t−1
β2 + Cmax(1− β)2(t−1) β2 + (1− β)2t+1 Tβ2
) D2
+ 2CmaxT 2K ( β2Cmax +
β3(1− β)(1 + T ) T
) V 2.
(56)
Combining with the condition for α and β in (50), if we set
β ≤ 1− 2(−1/(2T−1))
α = β ·min  1 2C1T √ 2(1− β)2T−1 − 1 Cmax , 2∑T l=1A 2 l + 16CmaxC 2 1T  , (57) we have
1
K K−1∑ k=0 E [‖∇F (θk)‖]2 ≤ 2(F (θ0)− F (θ∗)) Kα + LC21 · · ·C2Tα
+ 4CmaxT 2
( (1− β)2t−1
αβ2 + Cmax(1− β)2(t−1) αβ2 + (1− β)2t+1 Tαβ2
) D2
+ 4CmaxT 2
( β2Cmax
α + β3(1− β)(1 + T ) αT
) V 2
≤O(β) + 4CmaxT 2 ( 1 + Cmax)(1− β)2(t−1)
αβ2 +
(1− β)2t+1
Tαβ2
) D2
(58)
This completes the proof of Theorem 1.If we can set t such that (1−β)2(t−1) = O(β4), the algorithm achieves O(1/ √ K) convergence rate.
8.2 PROOF TO THEOREM 2
The SCSC algorithm with our sparse moving average can be written as
y (1) k+1 = (1− βk)y (1) k + f (1) ξ1,k (θk)− (1− βk)f (1)ξ1,k(θk−1)− k−1∏ j=k−t (1− βj)u(1)k , (59)
y (l) k+1 = (1−βk)y (l) k + f (l) ξl,k (y (l−1) k+1 )− (1−βk)f (l) ξl,k (y (l−1) k )− k−1∏ j=k−t (1−βj)u(l)k , 2 ≤ l ≤ T, (60)
θk+1 = θk − αk∇k, (61)
with u
(l) k = P (ξ1,k−t/(ξ1,k−t+1 ∪ · · · ∪ ξ1,k))y (l) k+t−1 (62)
and ∇k = ∇f (1)ξ1,k(θk) · · · ∇f (T ) ξN,k (y (T−1) k+1 )∇f (T+1) ξT+1,k (y (T ) k+1). (63)
8.2.1 SUPPORTING LEMMA
Lemma 6. The change of y(l)k in every iteration is bounded, i.e.,
E [∥∥∥y(l)k+1 − y(l)k ∥∥∥2] ≤3( βk1− βk )2 E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2]
+
( βk
1− βk
)2 V 2 + 3C2l−1E [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥2] + 3
k−1∏ j=k−t (1− βj)2E [∥∥∥u(l)k ∥∥∥2]
(64)
Proof. According to (60), we have (1− βk) ( y (l) k+1 − y (l) k ) =βk ( f (l) ξl,k ( y (l−1) k+1 ) − y(l)k+1 ) + (1− βk) ( f
(l) ξl,k
( y
(l−1) k+1 ) − f (l)ξl,k ( y (l−1) k )) − k−1∏ j=k−t (1− βj)u(l)k
=βk
( f (l) ( y
(l−1) k+1 ) − y(l)k+1 ) + βk ( f (l) ξl,k ( y (l−1) k+1 ) − f (l) ( y (l−1) k+1 )) + (1− βk) ( f
(l) ξl,k
( y
(l−1) k+1 ) − f (l)ξl,k ( y (l−1) k )) − k−1∏ j=k−t (1− βj)u(l)k
(65)
Taking the square of both sides of (65) gives us (64).
Lemma 7. The difference between the composite function and its moving average is bounded, i.e., E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2] ≤ (1− βk)E [∥∥∥y(l)k − f (l) (y(l−1)k )∥∥∥2]
+ 4 (1− βk)2 C2l E [∥∥∥y(l−1)k − y(l−1)k+1 ∥∥∥2]+ 2β2kV 2
+ 1
βk k−1∏ j=k−t (1− βj)2 E [ ‖u(l)k ‖ 2 ] (66)
Proof.
y (l) k+1 − f
(l) ( y
(l−1) k+1 ) = (1− βk) ( y (l) k − f (l) ( y (l−1) k )) + (1− βk) ( f (l) ( y (l−1) k ) − f (l) ( y (l−1) k+1 )) ︸ ︷︷ ︸
:=T1
+ βk
( f
(l) ξl,k
( y
(l−1) k+1
) − f (l) ( y
(l−1) k+1 )) ︸ ︷︷ ︸
:=T2 + (1− βk) ( f
(l) ξl,k
( y
(l−1) k+1 ) − f (l)ξl,k ( y (l−1) k )) ︸ ︷︷ ︸
:=T3
+ k−1∏ j=k−t (1− βj)u(l)k
(67)
Conditioned on Fk, taking expectation over ξl,k, we have:
E [(1− βk)T1 + βkT2 + (1− βk)T3] = E [ (1− βk) f (l) ( y (l−1) k ) − f (l) ( y (l−1) k+1 ) + f (l) ξl,k ( y (l) k+1 ) − (1− βk) f (l)ξl,k ( y (l) k )] = 0
Taking the square of both sides of (67) and taking the expectation conditioned on Fk,n :={ Fk, y(1)k+1, · · · , y (l−1) k+1 } , we have:
E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2]
= E  ∥∥∥∥∥∥(1− βk) ( y (l) k − f (l) ( y (l−1) k )) − k−1∏ j=k−t (1− βj)u(l)k ∥∥∥∥∥∥ 2 
+ E [ ‖(1− βk)T1 + βkT2 + (1− βk)T3‖2 ] ≤ (1− βk)E
[∥∥∥y(l)k − f (l) (y(l−1)k )∥∥∥2]+ 1βk k−1∏ j=k−t (1− βj)2 E [ ‖u(l)k ‖ 2 ]
+ 2E [ ‖(1− βk)T1 + βkT2‖2 |Fk,n ] + 2 (1− βk)2 E [ ‖T3‖2 |Fk,n ] ≤ (1− βk)E
[∥∥∥y(l)k − f (l) (y(l−1)k )∥∥∥2]+ 2 (1− βk)2 E [‖T1‖2 |Fk,n]+ 2β2kE [‖T2‖2 |Fk,n] + 2 (1− βk)2 E [ ‖T3‖2 |Fk,n ] + 1
βk k−1∏ j=k−t (1− βj)2 E [ ‖u(l)k ‖ 2 ]
≤ (1− βk)E [∥∥∥y(l)k − f (l) (y(l−1)k )∥∥∥2]+ 2 (1− βk)2 E [∥∥∥f (l) (y(l−1)k )− f (l) (y(l−1)k+1 )∥∥∥2]
+ 2 (1− βk)2 E [∥∥∥f (l)ξl,k (y(l−1)k )− f (l)ξl,k (y(l−1)k+1 )∥∥∥2]+ 2β2kV 2
+ 1
βk k−1∏ j=k−t (1− βj)2 E [ ‖u(l)k ‖ 2 ]
(68)
≤ (1− βk)E [∥∥∥y(l)k − f (l) (y(l−1)k )∥∥∥2]
+ 4 (1− βk)2 C2l E [∥∥∥y(l−1)k − y(l−1)k+1 ∥∥∥2]+ 2β2kV 2
+ 1
βk k−1∏ j=k−t (1− βj)2 E [ ‖u(l)k ‖ 2 ] (69)
8.2.2 REMAINING STEPS TOWARDS THEOREM 2.
Using the smoothness of F (θk), we have
F (θk+1) ≤ F (θk) + 〈∇F (θk), θk+1 − θk〉+ L
2 ‖θk+1 − θk‖2
= F (θk)− αk 〈∇F (θk),∇k〉+ L
2 ‖θk+1 − θk‖2
= F (θk)− αkE[‖∇F (θk)‖2] + L
2 ‖θk+1 − θk‖2 + αk 〈∇F (θk),∇F (θk)−∇k〉 .
(70)
Conditioning on Fk and taking expectation of both sides w.r.t. ξk, we have
E [F (θk+1) |Fk] ≤F (θk)− αkE [ ‖∇F (θk)‖2 ] + L 2 E [ ‖θk+1 − θk‖2 | Fk ] + αk 〈∇F (θk),E [∇F (θk)−∇k|Fk]〉
≤F (θk)− αkE [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2T+1α2k + αkE [‖∇F (θk)‖] ‖E [∇k|Fk]−∇F (θk)‖
(Lemma 4) ≤ F (θk)− αkE [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2T+1α2k + αk T∑ l=1 AlE [‖∇F (θk)‖]E [∥∥∥y(T )k+1 − f (T )(y(T−1)k+1 )∥∥∥ |Fk]
≤F (θk)− αk ( 1− αk
2βk T∑ l=1 A2l
) E [ ‖∇F (θk)‖2 ] + βk 2 T∑ l=1 E [∥∥∥y(T )k+1 − f (T )(y(T−1)k+1 )∥∥∥2 |Fk]
+ L
2 C21 · · ·C2T+1α2k.
(71)
The subsequent analysis builds on the following Lyapunov function:
Vk := F (θk)− F (θ∗) + T∑ l=1 E [∥∥∥y(l)k − f (l)(y(l−1)k )∥∥∥2] (72)
where θ∗ is the optimal solution of the problem.
Conditioning on Fk and taking expectation of Vk+1 w.r.t. ξk, we have
E [Vk+1|Fk] ≤Vk − αk ( 1− αk
2βk T∑ l=1 A2l
) E [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2T+1α2k
+ (1 + βk) T∑ l=1 E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2 |Fk]− T∑ l=1 E [∥∥∥y(l)k − f (l) (y(l−1)k )∥∥∥2 |Fk]
− βk 2 T∑ l=1 E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2 |Fk]
(Lemma 8) ≤ Vk − αk ( 1− αk
2βk T∑ l=1 A2l
) E [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2T+1α2k
+ ((1 + βk) (1− βk)− 1) T∑ l=1 E [∥∥∥y(l)k − f (l)(y(l−1)k )∥∥∥2 |Fk]
+ 4 (1 + βk) (1− βk)2C21E [ ‖θk − θk−1‖2 |Fk ] +
T∑ l=2 [ 4 (1 + βk) (1− βk)2C2l + γl ] E [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥ |Fk]
+ 2(1 + βk)β 2 kNV 2 + 1 + βk βk k−1∏ j=k−t (1− βj)2 T∑ l=1 E [ ‖u(l)k ‖ 2 ] − T∑ l=2 γlE [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥ |Fk]− βk2 T∑ l=1 E [∥∥∥y(l)k+1 − f (l)(y(l−1)k+1 )∥∥∥2]
≤Vk − αk ( 1− αk
2βk T∑ l=1 A2l
) ‖∇F (θk)‖2 + L
2 C21 · · ·C2T+1α2k + 2(1 + βk)β2kNV 2
+ 4C21E [ ‖θk − θk−1‖2 |Fk ] +
T∑ l=2 ( 4C2l + γl ) E [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥2 |Fk]
− T∑ l=2 γlE [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥ |Fk]− βk2 T∑ l=1 E [∥∥∥y(l)k+1 − f (l)(y(l−1)k+1 )∥∥∥2]
+ 1 + βk βk k−1∏ j=k−t (1− βj)2 T∑ l=1 E [ ‖u(l)k ‖ 2 ]
(Lemma 7) ≤ Vk − αk ( 1− αk
2βk T∑ l=1 A2l
) E [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2T+1α2k
+ 4C21E [ ‖θk − θk−1‖2 |Fk ] +3 ( βk
1− βk )2 T∑ l=2 ( 4C2l + γl ) E [∥∥∥y(l−1)k+1 − f (l) (y(l−2)k+1 )∥∥∥2]︸ ︷︷ ︸
:=Ik1
−βk 2 T∑ l=1 E [∥∥∥y(l)k+1 − f (l) (y(l−1)k+1 )∥∥∥2]︸ ︷︷ ︸
:=Ik2
+3 T∑ l=2 ( 4C2l + γl ) C2l−1E [∥∥∥y(l−2)k+1 − y(l−2)k ∥∥∥2 |Fk]︸ ︷︷ ︸ :=Ik3
− T∑ l=2 γlE [∥∥∥y(l−1)k+1 − y(l−1)k ∥∥∥ |Fk]︸ ︷︷ ︸
:=Ik4
+ ( 2 (1 + βk)β 2 kN + ( βk
1− βk )2 T∑ l=2 ( 4C2l + γl )) V 2
+ 1 + βk βk k−1∏ j=k−t (1− βj)2 T∑ l=1 E [ ‖u(l)k ‖ 2 ] + 3 k−1∏ j=k−t (1− βj)2 T∑ l=2 ( 4C2l + γl ) E [∥∥∥u(l−1)k ∥∥∥2]
(73)
If we define γl such that:
3
( βk
1− βk )2 T∑ l=2 ( 4C2l + γl ) < βk 2
(74)
and
3 ( 4C2l + γl ) C2l−1 < γl−1, (75)
then Ik1 + Ik2 < 0 (76)
Ik3 + Ik4 < 0. (77)
It follows that
E [Vk+1|Fk] ≤Vk − αk ( 1− αk
2βk T∑ l=1 A2l
) E [ ‖∇F (θk)‖2 ] + L
2 C21 · · ·C2T+1α2k
+ (3(4C22 + γ2) + 4)C 2 1E [ ‖θk − θk−1‖2 |Fk ] + ( 2 (1 + βk)β 2 kN + ( βk
1− βk )2 T∑ l=2 ( 4C2l + γl )) V 2
+ 1 + βk βk k−1∏ j=k−t (1− βj)2 T∑ l=1 E [ ‖u(l)k ‖ 2 ] + 3 k−1∏ j=k−t (1− βj)2 T∑ l=2 ( 4C2l + γl ) E [∥∥∥u(l−1)k ∥∥∥2]
(78)
Setting βk = αk ∑T−1 l=1 A 2 l , we have
E [Vk+1|Fk] ≤Vk − αk 2 E[‖∇F (θk)‖2] + L 2 C21 · · ·C2T+1α2k + (3(4C22 + γ2) + 4)C21E
[ ‖θk − θk−1‖2 |Fk ] + ( 2 (1 + βk)β 2 kN + ( βk
1− βk )2 T∑ l=2 ( 4C2l + γl )) V 2
+ 1 + βk βk k−1∏ j=k−t (1− βj)2 T∑ l=1 E [∥∥∥u(l)k ∥∥∥2]+ 3 k−1∏ j=k−t (1− βj)2 T∑ l=2 ( 4C2l + γl ) E [∥∥∥u(l−1)k ∥∥∥2]
(79)
Setting αk = α = cα√K and βk = β = cβ√ K and summing up both sides of (79) from k = 0 to K − 1, we have
1
K K−1∑ k=0 ‖∇F (θk)‖2 ≤ 2V0 Kα + LC21 · · ·C2T+1α+ 2(3(4C22 + γ2) + 4)C41C22 · · ·C2T+1α
+ 2
α
( 2 (1 + β)β2N + ( β
1− β )2 T∑ l=2 ( 4C2l + γl )) V 2
+ 2(1 + β) (1− β)2(t−1)
αβ
1
K K−1∑ k=0 T∑ l=1 E [∥∥∥u(l)k ∥∥∥2]
+ 6 (1− β)2(t−1)
α
1
K K−1∑ k=0 T∑ l=2 ( 4C2l + γl ) E [∥∥∥u(l−1)k ∥∥∥2]
(Assumption 5) ≤ O ( 1√ K ) + 2(1 + β) (1− β)2(t−1)
αβ
T∑ l=1 D2l
+ 6 (1− β)2(t−1)
α Cu T∑ l=1 D2l
(80)
where Cu = max ( 4C2l + γl ) for l = 1 . . . T .
This completes the proof of Theorem 2. If we can set t such that (1− β)2(t−1) = O(β3), the algorithm achieves O(1/ √ K) convergence rate.

Review summaries:
This paper formulated the graph neural network as a stochastic compositional optimization problem and then developed a new optimization algorithm to train GNN. However, there are some errors. It is not ready for publication.

This paper proposed an improved variant of the Stochastic Compositional Optimization (SCO) framework to train GNNs, replacing all nodes' moving averages with a sparse representation. The proposed algorithm only requires a fixed-size buffer, regardless of the graph size, solving the memory issue of SCO algorithms and making it practically applicable to large graphs. The paper showed that the proposed algorithm preserves the convergence rate of the original SCO algorithm and experimentally validated that the algorithm could outperform the traditional Adam SGD for GNN training with a small memory overhead.

This paper studies neighbor sampling techniques for training GNNs. Previous work has observed that sampling-based GNN training can be formulated as a Stochastic Compositional Optimization (SCO) problem. The authors argue that naive implementation of existing SCO algorithms incurs huge memory cost for training GNNs. This paper propose a Sparse Stochastic Compositional (SpSC) gradient method, which only stores the data for nodes sampled in the past few iterations. Convergence analysis on SpSC is provided and empirical results show that SpSC has better performance than naive SCO. 

This paper studies the problem of training GNNs on large-scale graphs. Specifically, it proposes to use sparse moving average for sampling-based GNN training strategies. The convergence rate of the proposed approach has been proved under several assumptions. The authors conduct experiments on several large-scale datasets to show the effectiveness and efficiency of the proposal.