A Biased Graph Neural Network Sampler with Near-Optimal Regret

1 Introduction . Graph convolution networks ( GCN ) and Graph neural networks ( GNN ) in general [ 21 , 17 ] have recently become a powerful tool for representation learning for graph structured data [ 6 , 2 , 33 ] . These neural networks iteratively update the representation of a node using a graph convolution operator or message passing operator which aggregate the embeddings of the neighbors of the node , followed by a non-linear transformation . After stacking multiple graph convolution layers , these models can learn node representations which can capture information from both immediate and distant neighbors . GCNs and variants [ 32 ] have demonstrated the start-of-art performance in a diverse range of graph learning prolems [ 21 , 17 , 3 , 30 , 13 , 15 , 23 ] . However , they face significant computational challenges given the increasing sizes of modern industrial datasets . The multilayers of graph convolutions is equivalent to recursively unfold the neighbor aggregation in a top-down manner which will lead to an exponentially growing neighborhood size with respect to the number of layers . If the graph is dense and scale-free , the computation of embeddings will involve a large portion of the graph even with a few layers , which is intractable for large-scale graph [ 21 , 34 ] . Several sampling methods have been proposed to alleviate the exponentially growing neighborhood sizes , including node-wise sampling [ 17 , 9 , 24 ] , layer-wise sampling [ 8 , 37 , 20 ] and subgraph sampling [ 10 , 35 , 19 ] . However , the optimal sampler with minimum variance is a function of the neighbors ’ embeddings unknown apriori before the sampling and only partially observable for those sampled neighbors . Most previous methods approximate the optimal sampler with a static distribution which can not reduce variance properly . And most of existing approaches [ 8 , 37 , 20 , 10 , 35 , 19 ] ∗Corresponding author . Work done during the internship at Amazon Shanghai AI Lab . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . do not provide any asymptotic convergence guarantee on the sampling variance . We are therefore less likely to be confident of their behavior as GNN models are applied to larger and larger graphs . Recently , Liu et al . [ 24 ] propose a novel formulation of neighbor sampling as a multi-armed bandit problem ( MAB ) and apply bandit algorithms to update sampler and reduce variance . Theoretically , they provide an asymptotic regret analysis on sampling variance . Empirically , this dynamic sampler named as BanditSampler is more flexible to capture the underlying dynamics of embeddings and exhibits promising performance in a variety of datasets . However , we will show in Section 2.3 that there are several critical issues related to the numerical stability and theoretical limitations of the BanditSampler [ 24 ] . First , the reward function designed is numerically unstable . Second , the bounded regret still can be regarded as a linear function of training horizon T . Third , their analysis relies on two strong implicit assumptions , and does not account for the unavoidable dependency between embedding-dependent rewards and GNN training dynamics . In this paper , we build upon the bandit formulation for GNN sampling and propose a newly-designed reward function that trades bias with variance . In Section 3.1 , we highlight that the proposed reward has the following crucial advantages : ( i ) It is numerically stable . ( ii ) It leads to a more meaningful notion of regret directly connected to sampling approximation error , the expected error between aggregation from sampling and that from full neighborhood . ( iii ) Its variation can be formulated by GNN training dynamics . Then in Section 3.2 , we clarify how the induced regret is connected to sampling approximation error and emphasize that the bounded variation of rewards is essential to derive a meaningful sublinear regret , i.e. , a per-iteration regret that decays to zero as T becomes large . In that sense , we are the first to explicitly account for GNN training dynamic due to stochastic gradient descent ( SGD ) so as to establish a bounded variation of embedding-dependent rewards , which we present in Section 3.3 . Based on that , in Section 4 , we prove our main result , namely , that the regret of the proposed algorithm as the order of ( T √ lnT ) 2/3 , which is near-optimal and manifest that the sampling approximation error of our algorithm asymptotically converges to that of the optimal oracle with the near-fastest rate . Hence we name our algorithm as Thanos from `` Thanos Has A Near-Optimal Sampler '' . Finally , empirical results in Section 5 demonstrate the improvement of Thanos over BanditSampler and others in terms of variance reduction and generalization performance . 2 Background . 2.1 Graph Neural Networks and Neighbor Sampling . Graph Neural Networks . Given a graph G = ( V , E ) , where V and E are node and edge sets respectively , the forward propagation of a GNN is formulated as h ( l+1 ) v , t = σ ( ∑ i∈Nv avih ( l ) i , tW ( l ) t ) for the node v ∈ V at training iteration t. Here h ( l ) i , t ∈ Rd is the hidden embedding of node i at the layer l , h ( 0 ) i , t = xi is the node feature , and σ ( · ) is the activation function . Additionally , Nv is the neighbor set of node v , Dv = |Nv| is the degree of node v , and avi > 0 is the edge weight between node v and i . And W ( l ) t ∈ Rd×d is the GNN weight matrix , learned by minimizing the stochastic loss L̂ with SGD . Finally , we denote z ( l ) i , t = avih ( l ) i , t as the weighted embedding , [ Dv ] = { i|1 ≤ i ≤ Dv } , and for a vector x ∈ Rd0 , we refer to its 2-norm as ‖x‖ ; for matrix W , its spectral norm is ‖W‖ . Neighbor Sampling . Recursive neighborhood expansion will cover a large portion of the graph if the graph is dense or scale-free even within a few layers . Therefore , we consider to neighbor sampling methods which samples k neighbors under the distribution p ( l ) v , t to approximate ∑ i∈Nv z ( l ) i , t with this subset St. We also call p ( l ) v , t the policy . For ease of notation , we simplify p ( l ) v , t as pt = { pi , t|i ∈ Nv } ; pi , t is the probability of neighbor i to be sampled . We can then approximate µ ( l ) v , t = ∑ i∈Nv z ( l ) i , t with an unbiased estimator µ̂ ( l ) v , t = 1 k ∑ i∈St z ( l ) i , t/pi , t . As it is unbiased , only the variance term Vpt ( µ̂ ( l ) v , t ) need to be considered when optimizing the policy pt . Define the variance term when k = 1 as Vpt . Then following [ 29 ] , Vpt ( µ̂ ( l ) v , t ) = Vpt/k with Vpt decomposes as Vpt = Ve −Vc . with Ve = ∑ i∈Nv ‖z ( l ) i , t‖2/pi , t , which is dependent on pt and thus refereed as the effective variance . And Vc = ‖ ∑ j∈Nv z ( l ) j , t‖2 is independent on the policy and therefore referred to as constant variance . 2.2 Formulate Neighbor Sampling as Multi-Armed Bandit . The optimal policy in terms of reducing the variance Vpt is given by p ∗ i , t = ‖zi , t‖∑ j∈Nv‖zj , t‖ [ 29 ] . However , this expression is intractable to compute for the following reasons : ( i ) It is only after sampling and forward propagation that we can observe z ( l ) i , t , and z ( l ) i , t changes with time along an optimization trajectory with unknown dynamics . ( ii ) z ( l ) i , t is only partially observable in that we can not see the embeddings of the nodes we do not sample . While static policies [ 17 , 8 , 37 ] are capable of dealing with ( ii ) , they are not equipped to handle ( i ) as required to approximate p∗t and reduce the sampling variance effectively . In contrast , adversarial MAB frameworks are capabable of addressing environments with unknown , non-stationary dynamics and partial observations alike ( See Fig.1 ) . The basic idea is that a hypothetical gambler must choose which of K slot machines to play ( See Fig . 1a ) . For neighbor sampling , K is equal to the degree Dv of root node v. At each time step , the gambler takes an action , meaning pulling an arm It ∈ [ K ] according to his policy pt , and then receives a reward rIt . To maximize cumulative rewards , an algorithm is applied to update the policy based on the observed reward history { rIτ : τ = 1 . . . , t } . Liu et al . [ 24 ] formulate node-wise neighbor sampling as a MAB problem . Following the general strategy from Salehi et al . [ 29 ] designed to reduce the variance of stochastic gradient descent , they apply an adversarial MAB to GNN neighbor sampling using the reward ri , t = −∇pi , tVe ( pt ) = ‖z ( l ) i , t‖ / p 2 i , t , ( 1 ) which is the negative gradient of the effective variance w.r.t . the policy . Since Ve ( pt ) −Ve ( p∗t ) ≤ 〈pt − p∗t , ∇ptVe ( pt ) 〉 , maximizing this reward over a sequence of arm pulls , i.e. , ∑T t=1 rIt , t , is more-or-less equivalent to minimizing an upper bound on ∑T t=1 Ve ( pt ) − ∑T t=1 Ve ( p ∗ t ) . The actual policy is then updated using one of two existing algorithms designed for adversarial bandits , namely Exp3 [ 1 ] and Exp3.M [ 31 ] . Please see Appendix C for details . Finally , Liu et al . [ 24 ] prove that the resulting BanditSampler can asymptotically approach the optimal variance with a factor of three : ∑T t=1 Ve ( pt ) ≤ 3 ∑T t=1 Ve ( p ∗ t ) + 10 √ TD4v ln ( Dv/k ) /k 3 . ( 2 ) Critically however , this result relies on strong implicit assumptions , and does not account for the unavoidable dependency between the reward distribution and GNN model training dynamics . We elaborate on this and other weaknesses of the BanditSampler next . 2.3 Limitation of BanditSampler . Updated by Exp3 , BanditSampler as described is sufficiently flexible to capture the embeddings ’ dynamics and give higher probability to ones with larger norm . And the dynamic policy endows it with promising performance on large datasets . Moreover , it can be applied not only to GCN but GAT models [ 32 ] , where avi change with time as well . It is an advantage over previous sampling approaches . Even so , we still found several crucial drawbacks of the BanditSampler . Numerical Instability Due to the pi , t in the denominator of ( 1 ) , the reward of BanditSampler suffers from numerical instability especially when the neighbors with small pi , t are sampled . From Fig . 5a ( in Appendix ) , we can observe that the rewards ( 1 ) of BanditSampler range between a large scale . Even though the mean of received rewards is around 2.5 , the max of received rewards can attain 1800 . This extremely heavy tail distribution forces us to choose a quite small temperature hyperparameter η ( Algorithm 3 and 5 in Appendix C ) , resulting in dramatic slowdown of the policy optimization . By contrast , the reward proposed by us in the following section is more numerically stable ( See Fig . 5b in Appendix ) and possesses better practical interpretation ( Fig . 2c ) . Limitation of Existing Regret and Rewards There are two types of regret analyses for bandit algorithms [ 1 , 4 ] : ( i ) the weak regret with a static oracle given by R̂ ( T ) = maxj∈ [ Dv ] ( ∑T t=1 E [ rj , t ] ) −∑T t=1 E [ rIt , t ] , which measures performance relative pulling the single best arm ; and ( ii ) the worstcase regret with a dynamic oracle given byR ( T ) = ∑T t=1 maxj∈ [ Dv ] E [ rj , t ] − ∑T t=1 E [ rIt , t ] , where the oracle can pull the best arm at each t. When the growth of the regret as a function of T is sublinear , the policy is long-run average optimal , meaning the long-run average performance converges to that of the oracle . But from this perspective , the bound from ( 2 ) can actually function more like worst-case regret . To see this , note that the scale factor on the oracle variance is 3 , which implies that once we subtract Ve ( p∗t ) from the upper bound , the effective regret satisfiesR ( T ) ≤ 2 ∑T t=1 Ve ( p ∗ t ) +O ( √ T ) . By substituting p∗t into Ve , we obtain Ve ( p ∗ t ) = ∑ i∈Nv ∑ j∈Nv‖z ( l ) i , t‖‖z ( l ) j , t‖ , which can be regarded as a constant lower bound given the converged variation of zi , t ( Lemma 1 ) . Consequently , the regret is still linear about T . And linear worst-case regret can not confirm the effectiveness of policy since uniform random guessing will also achieve linear regret . Crucial Implicit Assumptions There are two types of adversaries : if the current reward distribution is independent with the previous actions of the player , it is an oblivious adversary ; otherwise , it is a non-oblivious adversary [ 7 ] . GNN neighbor sampling is apparently non-oblivious setting but it is theoretically impossible to provide any meaningful guarantees on the worst-case regret in the non-oblivious setting ( beyond what can be achieved by random guessing ) unless explicit assumptions are made on reward variation [ 4 ] . BanditSampler [ 24 ] circumvents this issue by implicitly assuming bounded variation and oblivious setting ( See Appendix H ) , but this can not possibly be true since embedding-dependent rewards must depend on training trajectory and previous sampling . In contrast , we are the first to explicitly account for training dynamic in deriving reward variation and further regret bound in non-oblivious setting , and without this consideration no meaningful bound can possibly exist .

Summary: This paper tackles the high variance issues of the existing sampling methods, which approximate the full-graph training of GNNs. The proposed method is built upon existing work that treats GNN neighbor sampling as a multi-armed bandit problem. The authors design a novel biased reward function to reduce variance and avoid numerical instability. The proposed reward function leads to a more meaningful notion of regret directly connected to sampling approximation error. Based on that, the authors prove that the regret of the proposed algorithm is near-optimal and demonstrate the fast convergence of the sampling approximation error. Empirical results verify that the proposed algorithm, Thanos, has improved variance estimates and performance across benchmarks over BanditSampler and others.