Graph Neural Network Bandits
Abstract
We consider the bandit optimization problem with the reward function deﬁned over graph-structured data.
This problem has important applications in molecule design and drug discovery, where the reward is naturally invariant to graph permutations.
The key challenges in this setting are scaling to large domains, and to graphs with many nodes.
We resolve these challenges by embedding the permutation invariance into our model.
In particular, we show that graph neural networks (GNNs) can be used to estimate the reward function, assuming it resides in the
Reproducing Kernel Hilbert Space of a permutation-invariant additive kernel.
By establishing a novel connection between such kernels and the graph neural tangent kernel (GNTK), we introduce the ﬁrst GNN conﬁdence bound and use it to design a phased-elimination algorithm with sublinear regret.
Our regret bound depends on the GNTK’s maximum information gain, which we also provide a bound for.
While the reward function depends on all N node features, our guarantees are independent of the number of graph nodes N .
Empirically, our approach exhibits competitive performance and scales well on graph-structured domains.
Introduction
Contemporary bandit optimization approaches consider problems on large or continuous domains and have successfully been applied to a signiﬁcant number of machine learning and real-world applications, e.g., in mobile health, environmental monitoring, economics and hyperparameter tuning, to name a few.
The main idea behind them is to exploit the correlations between the rewards of “similar” actions.
This in turn, has resulted in increasingly rich models of reward functions (e g., in linear and kernelized bandits [37, 12]), including several recent attempts to harness deep neural networks for bandit tasks (see, e.g, [45, 29]).
A vast majority of previous works only focus on standard input domains and obtaining theoretical regret bounds.
Learning on graph-structured data, such as molecules or biological graph representations, requires designing sequential methods that can effectively exploit the structure of graphs.
Consequently, graph neural networks (GNNs) have received attention as a rapidly expanding class of machine learning models.
They deem remarkably well-suited for prediction tasks in applications such as designing novel materials [20], drug discovery [22], structure-based protein function prediction [16], etc.
This rises the question of how to bridge the gap, and design bandit optimization algorithms on graphstructured data that can exploit the power of graph neural networks to approximate a graph reward.
In this paper, we consider bandit optimization over graphs and propose to employ graph neural networks with one convolutional layer to estimate the unknown reward.
To scale to large graph domains (both in the number of graphs and number of nodes), we propose practical structural assumptions to model the reward function.
In particular, we propose to use permutation invariant additive kernels.
We show a novel connection between such kernels and the graph neural tangent kernel (GNTK) that we deﬁne in Section 3.
Our main result are GNN conﬁdence bounds that can be readily used in sequential decision-making algorithms to achieve sublinear regret bounds (see Section 4.3).
Related Work.
Our work extends the rich toolbox of methods for kernelized bandits and Bayesian optimization (BO) that work under the norm bounded Reproducing Kernel Hilbert Space assumption [37, 14, 42, 12]).
The majority of these methods are designed for general Euclidean domains and rely on kernelized conﬁdence sets to select which action to query next.
The exception is [43], that consider the spectral setting in which the reward function is a linear combination of the eigenvectors of the graph Laplacian and the bandit problem is deﬁned over nodes of a single graph.
In contrast, our focus is on optimizing over graph domains (i.e., set of graphs), and constructing conﬁdence sets that can quantify the uncertainty of graph neural networks estimates.
This work contributes to the literature on neural bandits , in which a fully-connected [45, 44, 19], or a single hidden layer convolutional network [25] is used to estimate the reward function.
These works provide sublinear cumulative regret bounds in their respective settings, however, when applied directly to graph features (as we demonstrate in Section 4.3), these approaches do not scale well with the number of graph nodes.
Due to its important applications in molecule design, sequential optimization on graphs has recently received considerable attention.
For example, in [28], the authors propose a kernel to capture similarities between graphs, and at every step, select the next graph through a kernelized random walk.
Other works (e.g., [17, 18, 23, 38]) encode graph representations to the (continuous) latent space of a variational autoencoder and perform BO in the latent space.
While practically relevant for discovering novel molecules with optimized properties, these approaches lack theoretical guarantees and deem computationally demanding.
A primary focus in our work is on embedding the natural structure of the data, i.e., permutation invariance, into the reward model.
This is inspired by the works of [6, 32] that consider invariances in kernel-based supervised learning.
Consequently, the graph neural tangent kernel plays an integral role in our theoretical analysis.
Du et al.
[15] provide a recursive expression for the tangent kernel of a GNN, without showing that the obtained expression is the limiting tangent kernel as deﬁned in Jacot et al.
[21] (i.e., as in Eq.
In contrast, we analyze the learning dynamics of the GNN and properties of the GNTK by exploiting the connection between the structure of a graph neural network and that of a neural network (in Section 3).
We recover that the graph neural tangent kernel also encodes additivity.
Additive models for bandit optimization have been previously studied in [24] and [35], however, these works only focus on Euclidean domains and standard base kernels.
Finally, we build upon the recent literature on elimination-based algorithms that make use of maximum variance reduction sampling [13, 8, 7, 9, 40, 30].
One of our proposed algorithms, GNN-PE, employs a phased elimination strategy together with our GNN conﬁdence sets.
Main Contributions.
We introduce a bandit problem over graphs and propose to capture prior knowledge by modeling the unknown reward function using a permutation invariant additive kernel.
We establish a key connection between such kernel assumptions and the graph neural tangent kernel (Proposition 3.2).
By exploiting this connection, we provide novel statistical conﬁdence bounds for the graph neural network estimator (Theorem 4.2).
We further prove that a phased elimination algorithm that uses our GNN-conﬁdence bounds (GNN-PE) achieves sublinear regret (Theorem 4.3).
Importantly, our regret bound scales favorably with the number of graphs and is independent of the number of graph nodes (see Table 1).
Finally, we empirically demonstrate that our algorithm consistently outperforms baselines across a range of problem instances.
Problem Statement
We consider a bandit problem where the learner aims to optimize an unknown reward function via sequential interactions with a stochastic environment.
At every time step t ∈ { 1 , . . .
, T } , the learner selects a graph G t from a graph domain G and observes a noisy reward y t
= f ∗ ( G t )
+ (cid:15) t , where f ∗ :
G → R is the reward function and (cid:15) t is i.i.d.
zero-mean sub-Gaussian noise with known variance proxy σ 2 .
Over a time horizon T , the learner seeks a small cumulative regret R T =
[Equation]
where G ∗ ∈ arg max G ∈G f ∗ ( G ) .
The aim is to attain regret that is sublinear in T , meaning that
R T /T
→ 0 as T → ∞ , which implies convergence to the optimal graph.
As an example application, consider drug or material design, where molecules may be represented with graph structures (e.g., from S MILES representations [1]) and the reward f ∗ ( G ) can correspond to an unknown molecular property of interest, e.g., atomization energy.
Evaluating such properties typically requires running costly simulations or experiments with noisy outcomes.
To identify the most promising candidate, e.g., the molecule with the highest atomization energy, molecules are sequentially recommended for testing and the goal is to ﬁnd the optimal molecule with the least number of evaluations.
Graph Domain.
We assume that the domain G is a ﬁnite set of undirected graphs with N nodes.
Without exploiting structure, standard bandit algorithms (e.g., [3]) cannot generalize across graphs, and their regret linearly depends on |G| .
To capture the structure, we consider reward functions de- pending on features associated with the graph nodes.
Similar to Du et al.
[15], we associate each node j ∈ [ N ] with a feature vector h G,j ∈ R d , for every graph G ∈ G .
We use h G =
( h G,j ) Nj =1 ∈ R Nd to denote the concatenated vector of all node features, and N ( j ) as the neighborhood of node j , including itself.
We deﬁne the aggregated node feature ¯
[Equation]
the normalized sum of the neighboring nodes’ features.
Similarly, ¯ h G ∈ R Nd denotes the aggregated features, stacked across all nodes.
Lastly, we let P N be the group of all permutations of length N , and use c · G to denote a permuted graph, where a permutation c ∈ P N is a bijective mapping from { 1 , . . .
, N } onto itself.
Permuting the nodes of a graph c · G produces a permuted feature vector h c · G :=
( h G,c ( j ) )
Nj =1 , and the same holds for the aggregated features ¯ h c · G .
Reward Model.
Practical graph optimization problems, such as drug discovery and materials optimization often do not depend on how the graphs’ nodes in the dataset are ordered.
We incorporate this structural prior into modeling the reward function, and consider functions that are invariant to node permutations .
We assume that f ∗ depends on the graph only through the aggregated node features and gives the same reward for all permutations of a graph, i.e., f ∗ ( c · G ) =
f ∗ ( G ) , for any
G ∈ G and c ∈ P N .
To guarantee such an invariance, we assume that the reward belongs to the reproducing kernel Hilbert space (RKHS) corresponding to a permutation invariant kernel
[Equation]
where k can be any kernel deﬁned on graph representations ¯ h G .
This assumption further restricts the hypothesis space to permutation invariant functions deﬁned on Nd –dimensional vector representations of graphs.
This is due to the reproducing property of the RKHS which allows us to write f ( G ) =
[Equation]
f,
[Equation]
) (cid:105) =
f ( c · G ) .
To make progress when optimizing over graphs with a large number of nodes N , we assume that k decomposes additively over node features, i.e.,
[Equation]
Thereby, we obtain an additive graph kernel that is invariant to node permutations:
[Equation]
[Equation]
For an arbitrary choice of k ( j ) , calculating ¯ k requires a costly sum over ( N !)
2 operands, since | P N | = N ! .
In Section 3, we select a base kernel for which the sum can be reduced to N 2 terms.
We are now in a position to state our main assumption.
We assume that f ∗ belongs to the RKHS of ¯ k and has a B -bounded RKHS norm.
The norm-bounded RKHS regularity assumption is typical in the kernelized and neural bandits literature [37, 12, 45, 25].
Note that Eq.
(1) only puts a structural prior on the kernel function, i.e., it describes the generic form of an additive permutation invariant graph kernel.
Specifying the base kernels k ( j ) determines the representation power of ¯ k .
The smoother the base kernels are, the less complex the RKHS of ¯ k will be.
In Section 3, we set the base kernels k ( j ) such that ¯ k becomes the expressive graph neural tangent kernel.
Graph Neural Networks
Graph neural networks are effective models for learning complex functions deﬁned on graphs.
As in Du et al.
[15], we consider graph networks that have a single graph convolutional layer and L fully-connected ReLU layers of equal width m .
Such a network f GNN ( G ; θ ) : G → R may be recursively deﬁned as follows:
[Equation]
where θ := ( W ( i ) )
i ≤ L +1 is initialized randomly with standard normal i.i.d.
entries, and σ relu ( x ) := max( 0 , x ) .
The network operates on aggregated node features ¯ h G,j as typical in Graph
Convolutional Networks [27].
For convenience, we assume that at initialization f GNN ( G ; θ 0 ) = 0 , for all G ∈ G , similar to [25, 45].
This assumption can be fulﬁlled without loss of generality, with a similar treatment as in [25, Appendix B.2].
Embedded Invariances.
In this work, we use graph neural networks to estimate the unknown reward function f ∗ .
This choice is motivated by the expressiveness of the GNN, the fact that it scales well with graph size, and particularly due to the invariances embedded in its structure.
We observe that the graph neural network f GNN is invariant to node permutations, i.e., for all G ∈ G and c ∈ P N ,
[Equation]
The key step to show this property is proving that f GNN can be expressed as an additive model of
L -layer fully-connected ReLU networks,
[Equation]
where f NN has a similar recursive deﬁnition as f GNN (see Equation A.1).
The above properties are formalized in Lemma A.1 and Lemma A.2.
Lazy (NTK) Regime.
We initialize and train f GNN in the well-known lazy regime [11].
In this initialization regime, when the width m is large, training with gradient descent using a small learning rate causes little change in the network’s parameters.
Let g GNN ( G, θ ) =
∇ θ f GNN ( G, θ ) denote the gradient of the network.
It can be shown that during training, for all G ∈ G , the network remains close to f GNN ( G, θ 0 )
+ g T GNN ( G, θ 0 )( θ − θ 0 ) , that is, its ﬁrst order approximation around initialization parameters θ 0 .
Training this linearized model with a squared error loss is equivalent to kernel regression with a tangent kernel
˜ k GNN ( G, G (cid:48) ) :
=
g T GNN ( G ; θ 0 )
g GNN ( G (cid:48) ; θ 0 ) .
For networks of ﬁnite width, this kernel function is random since it depends on the random network parameters at initialization.
We show in Proposition 3.1, that in the inﬁnite width limit, the tangent kernel converges to a deterministic kernel.
This proposition introduces the Graph Neural Tangent Kernel as the limiting kernel, and links it to the Neural Tangent Kernel ([21], also deﬁned in Appendix A).
Proposition 3.1.
Consider any two graphs G and G (cid:48) with N nodes and d -dimensional node features.
In the inﬁnite width limit, the tangent kernel ˜ k GNN ( G, G (cid:48) ) /m converges to a deterministic kernel,
[Equation]
which we refer to as the Graph Neural Tangent Kernel (GNTK).
Moreover,
[Equation]
where k NN : S d − 1 × S d − 1 → R is the Neural Tangent Kernel.
The proof is given in Appendix A.1.
We note that ¯ h G,j lies on the d -dimensional sphere, since the aggregated node features are normalized.
The NTK is bounded by 1 for any two points on the sphere [5].
Therefore, Proposition 3.1 implies that the GNTK is also bounded, i.e., k GNN ( G, G (cid:48) ) ≤for any G, G (cid:48) ∈ G .
This proposition yields a kernel which captures the behaviour of the lazy GNN.
While deﬁned on graphs with Nd dimensional representations, the effective input domain of this kernel is d -dimensional.
This advantage directly stems from the additive construction of the GNTK.
The next proposition uncovers the embedded structure of the GNTK by showing a novel connection between the GNTK and ¯ k , the permutation invariant additive kernel from Eq.
The proof is presented in Appendix A.1.
Proposition 3.2.
Consider ¯ k from Eq.
(1) , where for every 1 ≤ j ≤ N the base kernel k ( j ) is set to be equal to k NN ,
[Equation]
Then the permutation invariant additive kernel and the GNTK are identical, i.e., for all G, G (cid:48) ∈ G ,
[Equation]
This result implies that k GNN inherits the favorable properties of the permutation invariant additive kernel class.
Hence, functions residing in H GNN , the RKHS of k GNN , are additive, invariant to node permutations, and act on G through its aggregated node features.
While we use the GNTK as an analytical tool, this kernel can be of independent interest in kernel methods over graph domains.
In particular, calculating k GNN requires signiﬁcantly fewer operations compared to a kernel ¯ k with an arbitrary choice of k ( j ) , for which calculating ¯ k ( G, G (cid:48) ) requires super-exponentially many operations in N (See Eq.
In contrast, due to the decomposition in Eq.
(4), calculating k GNN only costs a quadratic number of summations.
GNN Bandits
The bandit literature is rich with algorithms that effectively balance exploration and exploitation to achieve sublinear regret.
Two components are common in kernelized bandit optimization algorithms.
The maximum information gain , for characterizing the worst-case complexity of the learning problem [37, 24, 12, 40]; and conﬁdence sets , for quantifying the learner’s uncertainty [4, 39, 36, 12, 31].
Our ﬁrst main result is an upper bound for the maximum information gain when the hypothesis space is H GNN (Theorem 4.1).
We then propose valid conﬁdence sets that utilize GNNs in Theorem 4.2.
These theorems may be of independent interest, as they can be used towards bounding the regret for a variety of bandit algorithms on graphs.
Lastly, we introduce the GNN-PE algorithm, together with its regret guarantee.
4.1 Information Gain
In bandit tasks, the learner seeks actions that give a large reward while, at the same time, provide information about the unknown reward function.
The speed of learning about f ∗ is commonly quantiﬁed via the maximum information gain.
Assume that the learner chooses a sequence of actions ( G 1 , . .
, G T ) and observes noisy rewards, where the noise is i.i.d.
and drawn from a zero-mean sub-Gaussian distribution with a variance proxy λ .
The information gain of this sequence calculated via the GNTK is
[Equation]
with the kernel matrix K GNN ,T =
[ k GNN ( G i , G j )] i,j ≤ T .
The maximum information gain (MIG) [37] is then deﬁned as:
[Equation]
In Section 4.3, we express regret bounds in terms of this quantity, as common in kernelized and neural bandits.
In Theorem 4.1 we obtain a data-independent bound on the MIG.
The proof is given in Appendix B.
Theorem 4.1 (GNTK Information Gain Bound) .
Suppose the observation noise is i.i.d., and drawn from a zero-mean sub-Gaussian distribution, and the input domain is G .
Then the maximum information gain associated with k GNN is bounded by
[Equation]
We observe that the obtained MIG bound does not depend on N the number of nodes in the graphs.
To highlight this advantage, we compare Theorem 4.1 to the equivalent bound for the vanilla neural tangent kernel which ignores the graph structure.
We consider the neural tangent kernel that operates on graphs through the Nd -dimensional vector of aggregated node features ¯ h G ,
[Equation]
For κ NN the maximum information gain scales as γ NN ,T = O ( T ( Nd − 1) /Nd log 1 /Nd T ) , where
N appears in the exponent [25].
This results in poor scalability with graph size in the bandit optimization task, as we further demonstrate in Section 4.3.
Table 1 summarizes this comparison.
4.2 Conﬁdence Sets
Quantifying the uncertainty over the reward helps the learner to guide exploration and balance it against exploitation.
Conﬁdence sets are an integral tool for uncertainty quantiﬁcation.
Conditioned on the history H t − 1 =
( G i , y i )
i<t , for any G ∈ G , the set C t − 1 ( G, δ ) deﬁnes an interval to which f ∗ ( G ) belongs with a high probability such that,
[Equation]
An approach common to the kernelized bandit literature is to construct sets of the form C t − 1 ( G, δ ) =
[ µ t − 1 ( G ) ± β t σ t − 1 ( G )] where β t depends on the conﬁdence level δ .
The center of the interval, characterized by µ t − 1 ( · ) , is the estimate of the reward, and the width β t σ t − 1 ( · ) , reﬂects the uncertainty.
In this work, we utilize GNNs for construction of such sets.
To this end, we train a graph neural network to estimate the reward.
We use the gradient of this network at initialization to approximate the uncertainty over the reward, as in [45].
Let f GNN ( G ; θ ( J ) t − 1 ) be the GNN trained with gradient descent for J steps and by using learning rate η on the loss
[Equation]
where λ is the regularization coefﬁcient, and θ 0 the network parameters at initialization.
We propose conﬁdence sets of the form
[Equation]
where the center and width of the set are calculated via,
[Equation]
Here g GNN ( G ; θ 0 )
=
∇ θ f GNN ( G ; θ 0 ) denotes the gradient at initialization.
Moreover, we use λ 0 :=
λ min ( K GNN ) > 0 to denote the minimum eigenvalue of the kernel matrix calculated for the entire domain, i.e., K GNN =
[ k GNN ( G, G (cid:48) )] G,G (cid:48) ∈G .
Theorem 4.2 shows that this construction gives valid conﬁdence intervals, i.e., it satisﬁes Eq.
(7), when the reward function lies in H GNN and has a bounded RKHS norm.
Theorem 4.2 (GNN Conﬁdence Bound) .
Set δ ∈ (0 , 1) .
Suppose f ∗ ∈
H k GNN with a bounded norm (cid:107) f ∗ (cid:107)
k
GNN ≤ B .
Assume that the random sequences ( G i ) i<t and ( (cid:15) i )
i<t are statistically independent.
Let the width m = poly
[Equation]
, learning rate η = C ( Lm + mλ ) − 1 with some universal constant C , and J ≥ 1 .
Then for all graphs G ∈ G , with probability of at least 1 − δ ,
[Equation]
where ˆ µ t − 1 and ˆ σ t − 1 are deﬁned in Eq.
(8) and
[Equation]
The " ≈ " notation in Theorem 4.2 omits the terms that vanish with t , i.e., are o (1) .
An exact version of the theorem without the aforementioned approximations is given in Appendix C.1.
[Table]
4.3 Bandit Optimization with Graph Neural Networks
The developed conﬁdence sets can be used to assist the learner with controlling the growth of regret.
In this section, we give a concrete example on how our GNN conﬁdence sets (Equation 8) can be used by an algorithm to solve bandit optimization tasks on graphs.
We introduce GNN-Phased Elimination (GNN-PE; see Algorithm 1) that consists of episodes of pure exploration over a set of plausible maximizer graphs, similar to [7, 30].
Each episode is followed by an elimination step, that makes use of GNN conﬁdence bounds to shrink the set of plausible maximizers.
More formally, at step t during an episode e , the learner selects actions via G e,t = arg max G ∈G e ˆ σ e,t − 1 ( G ) , where G e is the set of graphs that might maximize f ∗ according to the learner’s current knowledge.
Once the episode is over after T e steps, the set
G e is updated to contain graphs that still have a chance of being a maximizer according to the conﬁdence bounds [ˆ µ e,T e ( G ) ±
β T e ˆ σ e,T e ( G )] where ˆ µ e,T e and ˆ σ e,T e are only computed based on the points within episode e .
Theorem 4.3 shows that GNN-PE incurs a sublinear control over the cumulative regret.
We provide the proof in Appendix C.
We use ˜ O ( · ) notation to hide polylog( T ) factors.
Theorem 4.3.
Set δ ∈ (0 , 1) .
Suppose f ∗ ∈ H
f GNN with a bounded norm (cid:107) f ∗ (cid:107)
k GNN ≤ B .
Let the width m = poly
[Equation]
, learning rate η = C ( Lm + mλ ) − 1 with some universal constant C , and J ≥ 1 .
Then with probability at least 1 − δ , GNN-PE satisﬁes
[Equation]
We can observe the beneﬁt of working with a graph neural network by comparing the bound in
Theorem 4.3 with the regret for a structure-agnostic algorithm.
Recall κ NN the vanilla NTK, deﬁned over the concatenated feature vectors (Equation 6).
For the sake of this comparison, we ignore the geometric structure and assume that f ∗ ∈ H NN .
Swapping out k GNN for κ NN , and respectively the
GNN with an NN as deﬁned in Eq.
(A.1), we obtain NN-PE, the neural network counterpart of
GNN-PE.
This algorithm accepts Nd -dimensional input vectors as actions.
Similar to Theorem 4.3, we can show that NN-PE can satisfy a guarantee of O
[Equation]
for the regret.
This bound suggests that as N grows, ﬁnding the optimal graph can become more challenging for the learner.
Working with k GNN to encode the structure of the bandit problem, and consequently using the
GNN to solve it, removes the dependency on N in the exponent.
This result is summarized in Table 1.
We provide some intuition on why working with a permutation invariant model is beneﬁcial for bandit optimization on graphs.
Conﬁdence sets which are constructed for member of H NN are larger, and result in sub-optimal action selection.
Further, training the neural network is a more challenging task, since permutation invariance is not hard coded in the network architecture and has to be learned from the data.
This results in less accurate reward estimates.
We refer the reader to Appendix A.3 for a more rigorous discussion.
There we compare H GNN and H NN , the hypothesis spaces corresponding to the two models, through the Mercer decomposition of their kernels.
Experiments
We create synthetic datasets which may be of independent interest and can be used for evaluating and benchmarking machine learning algorithms on graph domains.
Each dataset is constructed from a ﬁnite graph domain together with a reward function.
The domains are generated randomly and differ in properties of the member graphs that inﬂuence the problem complexity, e.g., number of nodes and edge density.
Each domain G p,N consists of Erd˝os-Rényi random graphs, where each graph has N nodes, and between each two nodes there exists an edge with probability p .
The node features are i.i.d.
d = 10 dimensional standard Gaussian vectors.
We choose N ∈ { 5 , 20 , 100 } , p ∈ { 0 . 05 , 0 . 2 , 0 . 95 } , and thereby sample a total of 9 different domains each containing 10000 graphs.
For instance,
G 0 .
05 , 5 denotes the domain with sparse and small graphs, while G 0 .
95 , 100 is the domain of dense graphs with many nodes.
For every domain, we sample a random reward function f : G p,N → R that is invariant to node permutations.
We use GP(0 , k GNN ) as a prior, and sample f from its posterior GP.
The posterior is calculated using a small random dataset ( G i , y i ) i ≤ 5 , where y i are drawn independently from N (0 , 1) and G i are randomly chosen from G p,N .
The corresponding dataset is then D p,N = { ( G i , f ( G i ))
| G i ∈ G p,N } .
Experiment Setup.
Every performance curve in the paper shows an average over 20 runs of the corresponding bandit problem, each with a different action set sampled from G p,N .
The shaded areas in all ﬁgures show the standard error across runs.
In all experiments, the reward is observed with a zero-mean Gaussian noise of variance σ = 10 − 2 .
We always set width m = 2048 and layers L = 2 , for every type of network architecture.
Four algorithms appear in our experiments.
In addition to our main algorithm GNN-PE, we introduce GNN-UCB, which selects actions via
G t = arg max ˆ µ t − 1 ( G ) + β t ˆ σ t − 1 ( G ) , the classic UCB policy based on the GNN conﬁdence sets.
The pseudo-code is given in Appendix D.3.
NN-UCB, introduced by [45], is the neural counterpart of GNN-UCB, and NN-PE as discussed in Section 4.3.
To conﬁgure these algorithms, we only tune λ and β = β t , and we do so by using the simplest dataset D 0 .
We ﬁnd that the algorithms are not sensitive to domain conﬁgurations and work for all D p,N out of the box.
Therefore, the same values for λ and β are used across all experiments.
We include the complete result of our hyperparameter search in Figure 5.
Lazy training.
We initialize the graph neural networks (and the NNs) in the lazy regime as described in Eq.
(2) (and Eq.
A.1).
Training a network in this regime with gradient descent causes little change in the weights.
Consequently, it is challenging to effectively train a lazy network in practice.
Therefore, the stopping criterion for gradient descent plays a crucial role in achieving sublinear regret.
Inaccurate estimation of the reward function disturbs the balance of exploration and exploitation, and leads the learner to poor optima.
To prevent this issue, we devise a stopping criterion that depends on the history
H t − 1 , such that, as t grows, the network is often trained for more gradient descent steps J .
This criterion can be employed by any neural bandit algorithm and may be of independent practical interest.
The details of training with gradient descent, stopping and batching are given in Appendix D.2.
Regret Experiments.
We assess the performance of the algorithms on bandit optimization tasks over different domains.
In Figure 1, we show the inference cumulative regret ˆ R T =
(cid:80)
[Equation]
G ) , for which we select graph domains with N = 20 nodes and edge probability p = 0 .
Figure 6 shows the regret for all dataset conﬁgurations.
To verify scalability with |G| , we run the algorithms on action sets of increasing size |G| ∈ { 200 , 500 , 1000 } .
Figure 1 presents the results: GNN-PE consistently outperforms the other methods.
It is evident that the algorithms built with GNN conﬁdence sets ﬁnd the optimal graph, regardless of the size of the domain.
The GNN algorithms exhibit competitive performance, and attain sublinear regret for all dataset conﬁgurations.
The neural methods however, may fail to scale and ﬁnd the optima in limited time.
Scalability with Graph Size.
In Section 4.3, we argue that using a neural network which takes ¯ h G ∈ R Nd as the input, causes the regret to grow with O ( T (2 Nd − 1) / 2 Nd ) .
The additive structure of the GNN, however, allows the learner to work on a d -dimensional domain, independent of graph size.
Figure 2 reﬂects this behaviour.
Fixing p = 0 .
2 , and |G| = 200 , we run the algorithm over domains with two graph sizes N ∈ { 20 , 100 } .
GNN-PE achieves sublinear regret in both cases, and manages to ﬁnd a global maxima within roughly the same number of steps.
This is in contrast to
NN-PE, which is more affected by increasing graph size.
A similar comparison for all conﬁgurations and algorithms is plotted in Figure 7, and the same behaviour is observed across all settings: the performance of GNN methods scales well with N , while this is not the case for NN methods.
Effect of Graph Density.
As a ﬁnal observation, we discuss the effect of edge density.
Consider a complete graph G with (cid:0)
N(cid:1) edges.
The neighborhoods are symmetric and the aggregated node features ¯ h G,i are identical for all i ≤ N .
Permutations on this graph will not change the output of either f GNN or f NN .
Therefore, we expect that for dense graphs, i.e., large values of p , using a permutation invariant model comes with fewer beneﬁts for the learner.
This is opposed to when the graph is sparse
[Figure] [Figure]
[Figure]
and the neighborhoods are asymmetric.
To verify this conjecture, we ﬁx N = 20 , |G| = 200 , and run the algorithm over domains with graphs of different edge probability p ∈ { 0 . 05 , 0 . 2 , 0 . 95 } .
Figure 3 shows that while GNN-UCB always achieves sublinear regret, it takes longer to ﬁnd the optima when the graphs are more dense.
NN-UCB however, improves as the edge probability p grows, since, roughly put, the graphs in the domain are becoming invariant to permutations.
Therefore
Figure 3 conﬁrms that the performance gap between the two method is reduced for graphs that are more dense.
In Figure 8, we plot the effect of graph density for other dataset conﬁgurations and bandit algorithms.
This behaviour is observed predominantly for the UCB algorithms.
Conclusion
We analyze the use of graph neural networks in bandit optimization tasks over large graph domains.
The main takeaway is that encoding the natural structure of the environment into the model, reduces the complexity of the task for the learner.
By selecting a kernel that embeds invariances, we introduce structure into the algorithm in a principled manner.
Importantly, we propose key structural assumptions on the graph reward function and establish a novel connection between additive permutation invariant kernels and the GNTK.
We construct valid graph neural network conﬁdence sets, and use it to build a GNN bandit algorithm that achieves sublinear regret.
While all node features contribute to the graph’s reward, our bounds are independent of the number of nodes.
This result holds for GNNs with a single convolutional layer and graphs with node feature representation.
An immediate next step is to generalize this approach to other more complex graph neural network architectures and representations (e.g., by including information about graph edges) and investigating their effectiveness for bandit optimization.
Our analysis opens up two avenues of future research.
The proposed kernel and the graph conﬁdence sets may be used in other algorithms for sequential decision-making tasks on graphs.
Additionally, our approach of embedding the environment’s permutation invariant structure into the algorithm may inspire further work on structured bandit optimization in presence of invariances.
Acknowledgments and Disclosure of Funding
We thank Jonas Rothfuss for his valuable suggestions regarding the experiments.
We acknowledge
Deepak Narayanan’s effort on an earlier version of the code.
We thank Nicolas Emmenegger and
Scott Sussex for their thorough feedback and lastly, we thank Alex Hägele for fruitful discussions regarding the writing.
This research was supported by the European Research Council (ERC) under the European Union’s Horizon 2020 research and Innovation Program Grant agreement no. 815943.