A Biased Graph Neural Network Sampler with Near-Optimal Regret
Year: 2021

Abstract:
Graph neural networks (GNN) have recently emerged as a vehicle for applying deep network architectures to graph and relational data. However, given the increasing size of industrial datasets, in many practical situations the message passing computations required for sharing information across GNN layers are no longer scalable. Although various sampling methods have been introduced to approximate full-graph training within a tractable budget, there remain unresolved complications such as high variances and limited theoretical guarantees. To address these issues, we build upon existing work and treat GNN neighbor sampling as a multi-armed bandit problem but with a newly-designed reward function that introduces some degree of bias designed to reduce variance and avoid unstable, possibly-unbounded pay outs. And unlike prior bandit-GNN use cases, the resulting policy leads to near-optimal regret while accounting for the GNN training dynamics introduced by SGD. From a practical standpoint, this translates into lower variance estimates and competitive or superior test accuracy across several benchmarks.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 Background

heading: 2.1 Graph Neural Networks and Neighbor Sampling
text: 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.

heading: 2.2 Formulate Neighbor Sampling as Multi-Armed Bandit
text: 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 cannot 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.

heading: 2.3 Limitation of BanditSampler
text: 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 cannot 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 cannot 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.

heading: 3 Towards a More Meaningful Notion of Regret
text: To address the limitations of the BanditSampler, we need a new notion of regret and the corresponding reward upon which it is based. In this section we motivate a new biased reward function, interpret the resulting regret that emerges, and then conclude by linking with the GCN training dynamics.

heading: 3.1 Rethinking the Reward
text: Consider the following bias-variance decomposition of approximation error:
E[‖µ̂(l)v,t − µ (l) v,t‖2] =‖µ (l) v,t − E[µ̂ (l) v,t]‖2 + Vpt(µ̂ (l) v,t) , Bias(µ̂ (l) v,t) + Vpt(µ̂ (l) v,t).
Prior work has emphasized the enforcement of zero bias as the starting point when constructing samplers; however, we will now argue that broader estimators that do introduce bias should be reconsidered for the following reasons: (i) Zero bias itself may not be especially necessary given that even an unbiased µ̂(l)v,t will become biased for approximating h (l+1) v,t once it is passed through the non-linear activation function. (ii) BanditSampler only tackles the variance reduction after enforcing zero bias in the bias-variance trade-off. However, it is not clear that the optimal approximation error must always be achieved via a zero bias estimator, i.e., designing the reward to minimize the approximation error in aggregate could potentially perform better, even if the estimator involved is biased. (iii) Enforcing a unbiased estimator induces other additional complications: the reward can become numerically unstable and hard to bound in the case of a non-oblivious adversary. And as previously argued, meaningful theoretical analysis must account for optimization dynamics that fall under the non-oblivious setting. Consequently, to address these drawbacks, we propose to trade variance with bias by adopting the biased estimator: µ̂(l)v,t = Dv k ∑ i∈St z (l) i,t and redefine the reward:
ri,t = 2z (l)> i,t z̄ (l) v,t − ∥∥∥z(l)i,t∥∥∥2, with z̄(l)v,t = 1Dvµ(l)v,t = 1Dv ∑ j∈Nv z (l) j,t. (3)
Equation (3) is derived by weighting the gradient of bias and variance w.r.t. pi,t equally, which we delegate to Appendix. Additionally, because of partial observability, we approximate z̄(l) with 1 k ∑ i∈St z (l) i,t . We also noticed, due to the exponential function from the Exp3 algorithm (see line
6, Algorithm 3), the negative rewards of some neighbors will shrink wi,t considerably, which can adversely diminish their sampling probability making it hard to sample these neighbors again. Consequently, to encourage the exploration on the neighbors with negative rewards, we add ReLU function over rewards (note that our theory from Section 4 will account for this change). The practical reward is then formulated as
r̃i,t = ReLU ( 2z(l)> ∑
j∈St
1 k z (l) j,t − ‖z (l) i,t‖
2 ) . (4)
The intuition of (3) and by extension (4) is that the neighbors whose weighted embeddings z(l)i,t are closer to z̄(l)v,t will be assigned larger rewards (See Fig. 2c). Namely, our reward will bias the policy towards neighbors that having contributed to the accurate approximation instead of ones with large norm as favored by BanditSampler. And in the case of large but rare weighted embeddings far from µ
(l) v,t, BanditSampler tends to frequently sample these large and rare embeddings, causing significant
deviations. The empirical evidence is shown in Section 5.3.
The reward (3) possesses following practical and theoretical advantages, which will be expanded more in next sections:
• Since it is well bounded by ri,t = ‖z̄(l)v,t‖2 − ‖z (l) i,t − z̄ (l) v,t‖2 ≤ ‖z̄ (l) v,t‖2, the proposed reward is
more numerical stable as we show in Fig. 5 (See Appendix). • It will incur a more meaningful notion of regret, meaning the regret defined by (3) is equivalent to
the gap between the policy and the oracle w.r.t. approximation error. • The variation of reward (3) is tractable to bound as a function of training dynamics of GCN in non-
oblivious setting, leading to a provable sublinear regret as the order of (Dv lnDv)1/3(T √
lnT )2/3, which means the approximation error of policy asymptotically converges to the optimal oracle with a factor of one rather than three.

heading: 3.2 Interpreting the Resulting Regret
text: We focus on the worst-case regret in the following analysis. The regret defined by reward (3) is directly connected to approximation error. More specifically, we notice ri,t = −‖z(l)i,t − z̄ (l) v,t‖2 + ‖z̄(l)v,t‖2. Since ‖z̄ (l) v,t‖2 will be canceled out in R(T ), we haveR(T ) = ∑T t=1(E‖z (l) It,t − z̄(l)v,t‖2 − maxj∈[Dv] E‖z (l) j,t− z̄ (l) v,t‖2), where the former term is the expected approximation error of the policy and the latter is that of the optimal oracle. Consequently, the regret defined by (3) is the gap between the policy and the optimal oracle w.r.t. the approximation error.
Then we clarify how to bound this regret. The worst-case regret is a more solid guarantee of optimality than the weak regret. Even though some policies can establish the best achievable weak regretO( √ T ), their worst-case regret still be linear. This is because the gap between static and dynamic oracles can be a linear function of T if there is no constraint on rewards. For example, consider the following worst-case scenario. Given three arms {i1, i2, i3}, at every iteration, one of them will be assigned a reward of 3 while the others receive only 1. In that case, consistently pulling any arm will match the static oracle and any static oracle will have a linear gap with the dynamic oracle. Hence it is impossible to establish a sublinear worst-case regret unless additional assumptions are introduced on the variation of the rewards to bound the gap between static and dynamic oracles [4]. Besbes et al. [4] claim that the worst-case regret can be bounded as a function of the variation budget:
T−1∑ t=1 sup i∈[Dv ] ∣∣∣E[r̃i,t+1]− E[r̃i,t]∣∣∣ ≤ VT (5) where VT is called the variation budget. Then, Besbes et al. [4] derived the regret bound asR(T ) = O(K lnK · V 1/3T T 2/3) for Rexp3. Hence, if the variation budget is a sublinear function of T in the given environment, the worst-case regret will be sublinear as well.
To fix the theoretical drawbacks of BanditSampler, we first drop the assumption of oblivious adversary, i.e. considering the dependence between rewards and previous sampling along the training horizon of GCN. Then to bound the variation budget, we account for GCN training dynamic in practicallymeaningful setting (i.e. no unrealistic assumptions) as described next.

heading: 3.3 Accounting for the Training Dynamic of GCN
text: One of our theoretical contributions is to study the dynamics of embeddings in the context of GNN training optimized by SGD. We present our assumptions as follows:
• Lipschitz Continuous Activation Function: ∀x,y, ‖σ(x)−σ(y)‖ ≤ Cσ‖x−y‖ and σ(0) = 0. • Bounded Parameters: For any 1 ≤ t ≤ T and 0 ≤ l ≤ L− 1, ‖W (l)t ‖ ≤ Cθ. • Bounded Gradients: For 1 ≤ t ≤ T , ∃Cg , such that ∑L−1 l=0 ‖∇W (l)t L̂‖ ≤ Cg .
Besides, given the graph G and its feature X , since avi is fixed in GCN, define Cx = maxv∈V‖ ∑ i∈Nv avixi‖. Define D̄ = maxv∈V Dv, Ā = maxv,i avi, G = CσCθD̄Ā, and ∆zt,l = maxi∈V‖z (l) i,t+1−z (l) i,t‖. For SGD, we apply the learning rate schedule as αt = 1/t. The above assumptions are reasonable. The bounded gradient is generally assumed in the non-convex/convex convergence analysis of SGD [25, 28]. And the learning rate schedule is necessary for the analysis of SGD to decay its constant gradient variance [16]. Then we will bound ∆zt,l as a function of gradient norm and step size by recursively unfolding the neighbor aggregation. Lemma 1 (Dynamic of Embedding). Based on our assumptions on GCN, for any i ∈ V at the layer l, we have: ∥∥∥z(l)i,t∥∥∥ ≤ Cz, ∣∣∣r̃i,t∣∣∣ ≤ Cr, ∣∣∣ri,t∣∣∣ ≤ Cr, (6) where Cz = Gl−1ĀCσCθCx and Cr = 3C2z . Then, consider the training dynamics of GCN optimized by SGD. For any node i ∈ V at the layer l, we have
∆zt,l = max i∈V ∥∥∥z(l)i,t+1 − z(l)i,t∥∥∥ ≤ αtGl−1ĀCσCxCg. (7) Lemma 1 is obtained by recursively unfolding neighbor aggregations and training steps, and can be generally applied to any GCN in practical settings. Based on it, we can derive the variation budget of reward (3) and (4) as a function of ∆zt,l in the non-oblivious setup.
Lemma 2 (Variation Budget). Given the learning rate schedule of SGD as αt = 1/t and our assumptions on the GCN training, for any T ≥ 2, any v ∈ V , the variation of the expected reward in (3) and (4) can be bounded as:
T∑ t=1 ∣∣∣E[ri,t+1]− E[ri,t]∣∣∣ ≤ VT = C̄v lnT, T∑ t=1 ∣∣∣E[r̃i,t+1]− E[r̃i,t]∣∣∣ ≤ VT = C̄v lnT (8) where C̄v = 12G2(l−1)C2σC 2 xĀ
2CθCg . The derivation of Lemma 2 is attributed to that our reward variation can be explicitly formulated as a function of embeddings’ variation. In contrast, pi,t emerging in the denominator of (1) incurs not only the numerically unstable reward but hardship to bound its variation. More specifically, pi,t is proportional to the summation of observed reward history of neighbor i, which is hard to bound due to the complication to explicitly keep track of overall sampling trajectory as well as its bilateral dependency with pi,t. It is potentially why BanditSampler’s regret (2) ignores the dependency between rewards and previous training/sampling steps. On the contrary, our rewards are tractable to bound as a function of embeddings’ dynamic in practical non-oblivious setting, leading to a sublinear variation budget (8), and further a solid near-optimal worst-case regret as presented next.

heading: 4 Main Result: Thanos and Near-Optimal Regret
text: Algorithm 1 Thanos
1: Input: η > 0, γ ∈ (0, 1), k, T,∆T ,G, X, {αt}Tt=1. 2: Initialize: For any v ∈ V , any i ∈ Nv , set pi,1 = 1/Dv . 3: for t = 1, 2, . . . , T do 4: Reinitialize the policy every ∆T steps: ∀v ∈ V,∀i ∈ Nv , set pi,t = 1/Dv . 5: Sample k neighbors with pt and estimate µ (l) v,t with the estimator µ̂ (l) v,t = Dv k ∑ i∈St z (l) i,t . 6: Forward GNN model and calculate the reward ri,t according to (4). 7: Update the policy and optimize the model following [24] using η, γ, and {αt}Tt=1. 8: end for
Algorithm 1 presents the condensed version of our proposed algorithm. See Algorithm 2 in Appendix B for the detailed version. Besides the trade-off between bias and variance, and exploration and exploitation, our proposed algorithm also accounts for a third trade-off between remembering and forgetting: given the non-stationary reward distribution, while keeping track of more observations
can decrease the variance of reward estimation, the non-stationary environment implies that “old” information is potentially less relevant due to possible changes in the underlying rewards. The changing rewards give incentive to dismiss old information, which in turn encourages exploration. Therefore, we apply Rexp3 algorithm [4] to tackle the trade-off between remembering and forgetting by reinitializing the policy every ∆T steps (line 4 in Algorithm 1).
Then, we present our main result: bounding the worst-case regret of the proposed algorithm:
R(T ) = T∑ t=1 ∑ i∈N∗k E[ri,t]− T∑ t=1 ∑ It∈St Eπ[rIt,t]. (9)
where N ∗k = argmaxNk⊂Nv ∑ i∈Nk E[ri,t], |N ∗ k | = k. Because we consider the non-oblivious adversary, E[ri,t] is taken over the randomness of rewards caused by the previous history of randomized arm pulling. Eπ[rIt,t] is taken over the joint distribution π of the action sequence (S1,S2, . . . ,ST ). Theorem 3 (Regret Bound). Consider Algorithm 1 as the neighbor sampling algorithm for training GCN. Given either (3) or (4) as reward function, we can bound its regret as
follows. Let ∆T = (C̄v lnT )− 2 3 (Dv lnDv) 1 3T 2 3 , η =
√ 2k ln(Dv/k)
Cr(exp(Cr)−1)DvT , and γ = min{1, √
(exp(Cr)−1)Dv ln(Dv/k) 2kCrT }. Given the variation budget in (8), for every T ≥ Dv ≥ 2, we have the regret bound for either (3) or (4) as
R(T ) ≤ C̄(Dv lnDv) 1 3 · (T √ lnT ) 2 3 . (10)
where C̄ is a absolute constant independent with Dv and T . The obtained regret is as the order of (Dv lnDv)1/3(T √
lnT )2/3. According to Theorem 1 in [4], the worst-case regret of any policy is lower bounded by O((DvVT )1/3T 2/3), suggesting our algorithm is near-optimal (with a modest lnT factor from optimal). In that sense, we name our algorithm as Thanos from “Thanos Has A Near-Optimal Sampler.”
The near-optimal regret from Theorem 3 can be obtained due to the following reasons: (i) Our proposed reward leads to a more meaningful notion of regret which is directly connected to approximation error. (ii) Its variation budget is tractable to be formulated by the dynamic of embeddings. (iii) We explicitly study training dynamic of GCN to bound embeddings’ dynamic by recursively unfolding the neighbor aggregation and training steps in the practical setting.
As mentioned in Section 3.2, the regret based on rewards (3) is equivalent to approximation error. The result of Theorem 3 says the approximation error of Thanos asymptotically converges to that of the optimal oracle with the near-fastest convergence rate. In the case of enforcing zero bias like BanditSampler, sampling variance is the exact approximation error. However, even if we ignore other previously-mentioned limitations, its regret (2) suggests the approximation error of their policy asymptotically converges to three (as opposed to one) times of the oracle’s approximation error, so the regret is still linear. We compare the existing theoretical convergence guaratees in Table 1.

heading: 5 Experiments
text: We describe the experiments to verify the effectiveness of Thanos and its improvement over BanditSampler in term of sampling approximation error and final practical performance.

heading: 5.1 Illustrating Policy Differences via Synthetic Stochastic Block Model Data
text: As mentioned in Section 3.1, our reward will bias to sample the neighbors having contributed to accurate approximation. Fig. 2c is the visualization of this intuition: after setting z̄v,t = (1, 1)>,
the reward inside the dashed circle is positive; otherwise negative. And the embeddings closer to z̄v,t will have larger rewards. In order to understand how this bias differentiates the policy of two samplers given different distribution of features and edges, we propose to use cSBM[14, 11] to generate synthetic graphs. We consider a cSBM [14] with two classes, whose node set V1 and V2 have 500 nodes. The node features are sampled from class-specific Gaussians N1, N2. We set feature size to 100, average degree 2d̄ = 20, k = 10, and µ = 1, and we note that µ controls the difference between two Gaussian’s mean [14]. The average number of inter-class and intra-class edges per node is d̄ − λd̄1/2 and d̄ + λd̄1/2 respectively. Then, we scale down the node features of V1 by 0.1 to differentiate the distribution of feature norm and test the sampler’s sensitivity to it. The configuration of training and samplers is same as Section 5.4 and listed in Appendix.
In the case of cSBMs, an ideal sampler should sample more intra-class neighbors than inter-class neighbors to get linear-separable embeddings and better classification. Thus, we inspect for each v the k neighbors having the top-k highest sampling probability, and compute the ratio of intra-class neighbors among them, i.e. pintra = ∑ i∈Nv 1{(yi = yv) ∩ (pi,t is top-k)}/k. We report the average of pintra for V1 ∩ Vtrain versus t in Fig. 2a. For the scaled community V1, Thanos will be biased to sample more intra-class neighbors due to the intuition explained by Fig. 2c, leading to more accurate approximation and improvement on test accuracy over BanditSampler as shown in Fig. 3a and 2b. This claim holds true under different edge distributions (λ ∈ {0.5, 1, 1.5}). We additionally report the results on unscaled V2 for comparison in Appendix.

heading: 5.2 Evaluating the Sampling Approximation Error
text: We numerically compare the approximation error between two samplers in the training of GCN and GAT on Cora dataset from Kipf and Welling [21] as well as cSBM synthetic data in Section 5.1. At each iteration, given a batch of nodes VL at the top layer, we perform sampling with BanditSampler and Thanos respectively, getting two subgraphs Gbs and Gour. For Cora, we perform forward propagation on the original graph G as well as Gbs and Gour respectively with the same model parameters {W (l)t }l∈[L], and we get the accurate µ (1) v,t of the first layer aggregation as well as its estimated values µ̂(bs)v,t and µ̂ (our) v,t from both samplers. We compute distbs = ∑ v∈VL‖µ̂ (bs) v,t − µ (1) v,t‖
and distour = ∑ v∈VL‖µ̂ (our) v,t − µ (1) v,t‖. We set k = 2, γ = 0.1, η = 0.1 for Thanos, η = 0.01 for BanditSampler (since its unstable rewards require smaller η), ∆T = 200, αt = 0.001, L = 2 and the dimension of hidden embeddings d = 16. Fig. 3 plots the mean and the standard deviation of ∆dist = ∑T t=1 distour − ∑T t=1 distbs with 10 trials. The mean curves of both GCN and GAT are
below zero, suggesting Thanos establishes lower approximation error in practice. For cSBM synthetic graphs, we follow the setting as Section 5.1, compare two samplers under different edge distributions (λ ∈ {0.5, 1, 1.5}) and directly plot ∑ t distbs (blue) and ∑ t distour (red). From Fig. 3a, we know Thanos achieves quite lower approximation error and higher test accuracy (Fig. 2b) in the setting of less inter-edges (e.g. λ = 1 or 1.5) due to the intuition manifested by Fig. 2c, whereas BanditSampler is biased to sample large-norm neighbors, resulting in high approximation error and degenerated performance. For small λ = 0.5, the almost-equal number of inter/intra edges will shift µv,t to the unscaled community V2. Hence two samplers’ approximation error are close.

heading: 5.3 Sensitivity to Embedding Norms
text: Previously, we claim that BanditSampler will bias policy to the neighbors with large norm, potentially leading to severe deviation from µv,t as well as a drop of performance. In this section, we present the evidence on CoraFull [5] with corrupted features and demonstrate that our algorithm resolves this issue. For CoraFull, we randomly corrupt 800 (roughly 5% of) training nodes by multiplying their features by 40. We run both samplers 300 epochs with the corrupted CoraFull and count the total number of times that these corrupted nodes were sampled per epoch. We set k = 3, η = 1, γ = 0.2 and the other hyperparameters the same as Section 5.4. We repeat 5 trials for each algorithm and report the average over epochs and trials. We also record the test accuracy with the best validation accuracy in each trial and report its mean across trials. From Fig. 4a, we can tell BanditSampler biases to corrupted nodes, degenerating its performance more as shown in Fig. 4b.

heading: 5.4 Accuracy Comparisons across Real-World Benchmark Datasets
text: We conduct node classification experiments on several benchmark datasets with large graphs: ogbnarxiv, ogbn-products [18], CoraFull, Chameleon [11] and Squirrel [27]. The models include GCN and GAT. For GCN, we compare the test accuracy among Thanos, BanditSampler, GraphSage[17], LADIES[37], GraphSaint[35], ClusterGCN[10] and vanilla GCN. For GAT, we compare test accuracy among Thanos, BanditSampler and vanilla GAT. The experimental setting is similar with Liu et al. [24]. The dimension of hidden embedding d is 16 for Chameleon and Squirrel, 256 for the others. The number of layer is fixed as 2. We set k = 3 for CoraFull; k = 5 for ogbn-arxiv, Chameleon, Squirrel; k = 10 for ogbn-products. We searched the learning rate among {0.001, 0.002, 0.005, 0.01} and found 0.001 optimal. And we set the penalty weight of l2 regularization 0.0005 and dropout rate 0.1. We do grid search for sampling hyperparameters: η, γ, and ∆T and choose optimal ones for each. Their detailed settings and dataset split are listed in Appendix. Also we apply neighbor sampling for test nodes for all methods, which is consistent with prior LADIES and GraphSaint experiments, and is standard for scalability in practical setting. From Table 2, we can tell our algorithm achieves superior performance over BanditSampler for training GAT, and competitive or superior performance for training GCN.

heading: 5.5 Sample Size Ablation
text: To verify the sensitivity of Thanos w.r.t. sample size k, we compare the test accuracy between Thanos and BanditSampler as sample size k increases on Ogbn-arxiv. The other hyperparameter setting is the same as Section 5.4. We compare two samplers with k = 3, k = 5, k = 10, k = 15. The result from Fig. 4c suggests Thanos still exhibits a mainfest improvement over BanditSampler as k increases.

heading: 6 Related Work
text: Hamilton et al. [17] initially proposed to uniformly sample subset for each root node. Many other methods extend this strategy, either by reducing variance [9], by redefining neighborhoods [34] [36] [22], or by reweighting the policy with MAB [24] and reinforcement learning [26]. Layer-wise sampling further reduces the memory footprint by sampling a fixed number of nodes for each layer. Recent layer-wise sampling approaches include [8] and [37] that use importance sampling according to graph topology, as well as [20] and [12] that also consider node features. Moreover, training GNNs with subgraph sampling involves taking random subgraphs from the original graph and apply them for each step. Chiang et al. [10] partitions the original graph into smaller subgraphs before training. Zeng et al. [35] and Hu et al. [19] samples subgraphs in an online fashion. However, most of them do not provide any 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.

heading: 7 Conclusion
text: In this paper, we build upon bandit formulation for GNN sampling and propose a newly-designed reward function that introduce some degree of bias to reduce variance and avoid numerical instability. Then, we study the dynamic of embeddings introduced by SGD so that bounding the variation of our rewards. Based on that, we prove our algorithm incurs a new-optimal regret. Besides, our algorithm named Thanos addresses another trade-off between remembering and forgetting caused by the non-stationary rewards by employing Rexp3 algorithm. The empirical results demonstrate the improvement of Thanos over BanditSampler in term of approximation error and test accuracy.

heading: Acknowledgements
text: We would like to thank Amazon Web Service for supporting the computational resources, Hanjun Dai for the extremely helpful discussion, and the anonymous reviewers for providing constructive feedback on our manuscript.

Review summaries:
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.

The authors propose Thanos, a bandit-based approach for subsampling graph neighbourhoods, to be used as context for graph neural network predictions. It improves on the prior art (BanditSampler) by proposing trading off increased reward bias for less variance. This allows the authors to reduce numerical optimisation issues, achieve competitive benchmarking performance, as well as explicitly take into account the GNN training dynamics in their theoretical analysis, demonstrating near-optimal regret (with an added ln T factor).

The authors propose a novel bandit sampler, addressing the limitations of the previous work, for graph neural networks. The previous work utilized EXP3 which as known issues with low probabilities (variance can be really high) and leads to small learning rate. The regret bound was given in terms of a static oracle instead of a dynamic one. Also, the plays are considered in the oblivious adversary setting.  These issues are addressed in the current work by rethinking the reward function. Experimental results are shown on a wide variety of datasets as well as GNN architectures to validate the proposed algorithm.  

This paper proposes an adaptive neighbor sampling method for graph neural networks applied to large graphs. This method is based on Rexp3, an adversarial bandit algorithm. Central to it is a novel reward, which unlike previous work induces a biased sampler, which is shown to have a near optimal regret. Even though it is biased, this sampler makes up for it by having better numerical stability, requiring less unrealistic assumption for the bound, and having provably faster convergence to variance reduction of the estimator.  Empirically, the paper confirms that the behavior of the proposed method matches theoretical predictions and performs better than a more naive bandit approach on a toy graph problem and Cora. Large scale experiments are then conducted on several standard benchmarks where the proposed methods appears to outperform baselines.