Pure Exploration in Kernel and Neural Bandits
Year: 2022

Abstract:
We study pure exploration in bandits, where the dimension of the feature representation can be much larger than the number of arms. To overcome the curse of dimensionality, we propose to adaptively embed the feature representation of each arm into a lower-dimensional space and carefully deal with the induced model misspecifications. Our approach is conceptually very different from existing works that can either only handle low-dimensional linear bandits or passively deal with model misspecifications. We showcase the application of our approach to two pure exploration settings that were previously under-studied: (1) the reward function belongs to a possibly infinite-dimensional Reproducing Kernel Hilbert Space, and (2) the reward function is nonlinear and can be approximated by neural networks. Our main results provide sample complexity guarantees that only depend on the effective dimension of the feature spaces in the kernel or neural representations. Extensive experiments conducted on both synthetic and real-world datasets demonstrate the efficacy of our methods.

heading: 1 Introduction
text: Pure exploration in bandits [11, 12, 6] has been extensively studied in machine learning. Consider a set of arms, where each arm is associated with an unknown reward distribution. The goal is to ∗Equal contribution
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
approximately identify the optimal arm using as few samples as possible. Applications of bandit pure exploration range from medical domains [3] to online content recommendation [40].
Despite the popularity of bandit pure exploration, it was previously mainly studied in two relatively restrictive settings: (1) the standard multi-armed bandit setting [22, 19, 9, 18], where the expected rewards among arms are completely unrelated to each other, and (2) the (generalized) linear bandit setting [38, 13, 10, 26], where the expected rewards are assumed to be linearly parameterized by some unknown weight vector. The standard multi-armed bandit setting fails to deal with large arm sets, and the linear bandit setting suffers from both model misspecifications (due to its simplified linear form) and the curse of dimensionality in the high-dimensional setting. Pure exploration is also studied in continuous spaces. However, guarantees therein scale exponentially with dimension [31, 4].
In this paper, we generalize bandit pure exploration to the nonlinear and high-dimensional settings. More specifically, we study the following two settings: (1) the rewards of arms are parameterized by a function belonging to a Reproducing Kernel Hilbert Space (RKHS), and (2) the rewards of arms are nonlinear functions that can be approximated by an overparameterized neural network. Problems in these two settings are often high-dimensional in nature. To overcome the curse of dimensionality, we propose to adaptively embed each arm’s feature representation in a lower-dimensional space and carefully deal with the induced misspecifications. Note that our approach is conceptually very different from all existing work dealing with model misspecifications: they assume the existence of misspecifications and address it in the original space (thus dealing with model misspecifications in a passive way) [28, 7]. On the other hand, we deliberately induce (acceptable) misspecifications to embed arms into lower-dimensional spaces and thus overcome the curse of dimensionality.

heading: 1.1 Contribution and Outline
text: We make the following main contributions:
• In Section 3, we introduce the idea of adaptive embedding to avoid the curse of dimensionality. The induced model misspecifications are carefully handled, which is novel in the bandit pure exploration setting. The sample complexity is theoretically analyzed and we relate the instancedependent sample complexity to the complexity of a closely-related linear bandit problem without model misspecification. As a by-product, our algorithm can also be applied to constrained highdimensional linear bandit pure exploration to reduce sample complexity.
• In Section 4, we specialize the adaptive embedding scheme to pure exploration in an RKHS. We construct feature mappings from eigenfunctions and eigenvalues of the associated kernel. The effective dimension of the kernel is analyzed, and we provide sample complexity guarantees in terms of the eigenvalue decay of the associated kernel. We rely on a known kernel in this setting.
• In Section 5, we further extend our adaptive embedding scheme to pure exploration with a general nonlinear reward function and model the reward function with an over-parameterized neural network. Sample complexity guarantees are provided with respect to the eigenvalue decay of the associated Neural Tangent Kernel. To the best of our knowledge, this provides the first theoretically founded pure exploration algorithm with a neural network approximation.
• In Section 6, we conduct extensive experiments on both synthetic and real-world datesets to confirm the efficacy of our proposed algorithms. We conclude our paper in Section 7 with open problems.

heading: 1.2 Related Work
text: The bandit pure exploration problem has a long history, dating back to the seminal work by Bechhofer [5], Paulson et al. [32]. One classical objective of pure exploration is Best Arm Identification (BAI), where the goal is to identify the best arm using as few samples as possible [22, 19, 9, 15]. To make it applicable to a large action space, the BAI problem is also extensively studied as the good arm identification problem, where the goal is to identify an -optimal arm [11, 12, 20, 21, 34, 23, 30].
The pure exploration problem in linear bandits is initially analyzed in Soare et al. [38], where optimal experimental design [27] is applied to guide the allocation of samples. Other approaches dealing with linear bandits, with various sample complexity guarantees, include adaptive sampling [45] and an approach called track-and-stop [10]. Constrained linear bandit pure exploration is also commonly studied with additional assumptions on the reward parameters [41, 10]. We note that the
track-and-stop approach only achieves optimal instance-dependent sample complexity in the regime where the confidence parameter approaches 0, but fails to do so in the moderate confidence regime. Fiez et al. [13] propose an elimination-based algorithm (with optimal design) that achieves (nearly) instance-dependent sample complexity. such algorithm is further generalized to the combinatorial bandit setting [24].
Learning with model misspecifications was recently introduced in bandit learning, with the primary emphasis placed on the regret minimization problem [16, 28, 14]. A very recent independent work studies pure exploration in kernel bandits with misspecifications [7]; both their and our algorithms follow the framework of RAGE [13] and draw inspiration from [28]. Camilleri et al. [7] propose a robust estimator that works in high-dimensional spaces and also explore the project-then-round idea through regularized least squares. Our algorithms adaptively embed actions into lower dimensional spaces according to some error tolerances (different embeddings from round to round); our rounding and elimination steps are thus computed only with respect to lower-dimensional embeddings. We additionally study the pure exploration problem with an overparameterized neural network. As mentioned before, our approach is also conceptually different from existing ones: rather than passively dealing with model misspecifications in its original representation, we deliberately and adaptively embed arms into a lower-dimensional space to avoid the curse of dimensionality; the induced model misspecifications are also carefully dealt with in our algorithms.

heading: 2 Problem setting
text: We introduce the general setting and notations for pure exploration in bandits. Consider a set of arms X ⊆ RD where the number of arms |X | = K is possibly very large. We use an unknown function h : X → [−1, 1] to represent the true reward of each arm. A noisy feedback h(x) + ξ is observed after each sample arm x, where the noise ξ is assumed to be 1-sub-Gaussian. The learner is allowed to allocate her samples based on previously collected information, and the goal is to approximately identify an approximately optimal arm using as few samples as possible. Let x? = arg maxx∈X h(x) denote the optimal arm among X . We aim at developing ( , δ)-PAC guarantees: for any δ ∈ (0, 1), with probability at least 1− δ, the algorithm outputs an -optimal arm x̂ such that h(x̂) ≥ h(x?)− using a finite number of samples. The performance of the algorithm is measured by its sample complexity, i.e., the number of samples pulled before it stops and recommends a candidate arm.
Notations. We define ∆x = h(x?)− h(x) as the sub-optimality gap of arm x. We use the notations Sk := {x ∈ X : ∆x < 4 · 2−k} (with S1 = X ). We consider feature mappings of the form ψd(·) : X → Rd, and define ψd(X ) = {ψd(x) : x ∈ X}. We use ΛX = {λ ∈ R|X | : ∑ x∈X λx = 1, λx ≥ 0} to denote the (|X | − 1)-dimensional probability simplex over arms in X ; and set Aψd(λ) = ∑ x∈X λxψd(x)ψd(x) >.2 We use ‖x‖A = √ x>Ax to represent the Mahalanobis norm. We also define Y(V) = {v − v′ : v,v′ ∈ V} for any set V . For a matrixH ∈ R|X |×|X|, we useH(x,x′) to denote the entry ofH which locates at row x and column x′.

heading: 3 Bandit pure exploration with adaptive embedding
text: We introduce the idea of bandit pure exploration with adaptive embedding, which can be viewed as an approach that actively trades off sample complexity with accuracy guarantees: we adaptively embed the feature representation into lower-dimensional spaces to avoid the curse of dimensionality, and conduct pure exploration with misspecified linear bandits. The embedding dimensions are carefully selected so that we can identify an -optimal arm.
We formalize the idea as follows. For any d ∈ N, we assume the existence of a feature mapping ψd : X → Rd and a unknown reward vector θd ∈ Rd such that, for any x ∈ X ,
h(x) = 〈ψd(x),θd〉+ ηd(x), where ηd(x) represents the induced approximation error on arm xwith respect to the low-dimensional embedding ψd(·). Without loss of generality, we assume that the action set X is rich enough so that ψd(X ) spans Rd for d considered in this paper. Otherwise, one can always project feature
2A generalized inversion is used for singular matrices. We refer to Appendix A.1 for detailed discussion.
representations ψd(X ) into an even lower-dimensional space without losing information in the linear component.
We use γ̃ : N → R to represent the misspecification level: an upper bound of the induced approximation error across all arms, i.e., maxx∈X |ηd(x)| ≤ γ̃(d). We define g(d, ζ) := (1 + ζ) infλ∈ΛX supy∈Y(ψd(X ))‖y‖2Aψd (λ)−1 , which represents the optimal value of a transductive design among embeddings in Rd. We define γ(d) := (16 + 8 √ g(d, ζ)) γ̃(d), which quantifies the sub-optimality gap of the identified arm in the worst case. One can easily show γ(d) ≤ O(γ̃(d) √ d) through Kiefer-Wolfowitz theorem [27]. Remark 1. We believe such optimality guarantees are un-improvable in general. In fact, a hard instance is constructed in [28] showing that, even with deterministic feedback, identifying a o(γ̃(d) √ d)optimal arm requires sample complexity exponential on d. On the other side, identifying a Ω(γ̃(d) √ d)optimal only requires sample complexity polynomially in d. Such a sharp trade-off between optimality and sample complexity motivates our definition of γ(d) (and our sample complexities are polynomially in d).
We assume the knowledge of both the feature mapping ψd(·) and the error function γ̃(·). This assumption is mild since one can explicitly construct/analyze ψd(·) and γ̃(·) in many cases (as discussed in Section 3.2, Section 4 and Section 5). We further assume that γ(d) can be made arbitrarily small for large enough d. Such assumption trivially holds if the rewards are perfectly explained for d large enough, i.e., γ̃(d) = 0. We now define the effective dimension with respect to γ(d) (induced from feature mapping ψd(·)) as follows. Definition 1. For any > 0, we define the effective dimension as deff( ) := min{d ≥ 1 : γ(d) ≤ }. In general, the effective dimension deff( ) captures the smallest dimension one needs to explore in order to identify an -optimal arm. Similar notions have been previously used in regret minimization settings [42, 43]. One can easily see that deff( 1) ≤ deff( 2) as long as 1 ≥ 2.

heading: 3.1 Algorithm and analysis
text: Algorithm 1 follows the framework of RAGE [13] to eliminate arms with sub-optimality gap≥ O(2−k) at the k-th iteration. It runs for n = O(log(1/ )) iterations and identifies an -optimal arm. We use optimal experimental design to select arms for faster elimination. For any fixed design λ ∈ ΛX , with N ≥ rd(ζ) samples and an approximation factor ζ (with default value ζ ∈ [1/10, 1/4]), the rounding procedure in Rd, i.e., ROUND(λ,N, d, ζ), outputs a discrete allocation {x1,x2, . . . ,xN} satisfying
max y∈Y(ψd(X )) ‖y‖2(∑Ni=1 ψd(xi)ψd(xi)>)−1 ≤ (1 + ζ) maxy∈Y(ψd(X )) ‖y‖2Aψd (λ)−1/N. (1)
Efficient rounding procedures exist with rd(ζ) = d 2+d+2 ζ [33] or rd(ζ) = 180d ζ2 [1, 13]. We refer reads to [13, 33, 1] for detailed rounding algorithms and the associated computational complexities.
Unlike RAGE that directly works in the original high-dimensional space, Algorithm 1 adaptively embeds arms into lower-dimensional spaces and carefully deals with the induced misspecifications. More specifically, the embedding dimension dk is selected as the smallest dimension such that the induced error term k is well controlled, i.e., k ≤ O(2−k). The embedding is more aggressive at initial iterations due to larger error tolerance; The embedding dimension selected at the last iteration is (roughly) deff( ) to identify an -optimal arm. The number of samples required for each iteration Nk is with respect to an experimental design in the lower-dimensional space after embedding. The ROUND procedure also becomes more efficient due to the embedding. Before stating our main theorem, we introduce the following complexity measure [38, 13, 10], which quantifies the hardness of the pure exploration problem (with respect to mapping ψd(·)).
ρ?d( ) := inf λ∈ΛX sup x∈X\{x?} ‖ψd(x?)−ψd(x)‖2Aψd (λ)−1 max{h(x?)− h(x), }2 .
Theorem 1. With probability of at least 1− δ, Algorithm 1 correctly outputs an -optimal arm with sample complexity upper bounded by
640
dlog2(2/ )e∑
k=1
(( k ρ?dk(2 2−k) log(k2|X |2/δ) ) + (rdk(ζ) + 1) ) ≤ Õ ( deff( ) ·max{∆min, }−2 ) ,
Algorithm 1 Arm Elimination with Adaptive Embedding and Induced Misspecification Input: Action set X , confidence parameter δ, accuracy parameter and rounding approximation
factor ζ. 1: Set n = dlog2(2/ )e and Ŝ1 = X . 2: for k = 1, 2, . . . , n do 3: Set δk = δ/k2, dk = deff(4 · 2−k). 4: Select feature representation ψdk(·), and calculate the induced misspecification level γ̃(dk).
Set rdk(ζ) = O(dk/ζ 2) as the number of samples needed for ROUND in Rdk .
5: Set λk and τk be the design and the value of the following optimization problem inf λ∈ΛX
sup y∈Y(ψdk (Ŝk)) ‖y‖2Aψdk (λ)−1 .
6: Set k = 2γ̃(dk) + γ̃(dk) √
(1 + ζ) τk, and Nk = max{d(2−k − k)−22(1 + ζ) τk log(|Ŝk|2/δk)e, rdk(ζ)}.
7: Get {x1,x2, . . . ,xNk} = ROUND(λk, Nk, dk, ζ). 8: Pull arms {x1,x2, . . . ,xNk} and receive rewards {y1, . . . , yNk}. 9: Set θ̂k = A−1k bk, whereAk = ∑Nk i=1ψdk(xi)ψdk(xi) > and bk = ∑Nk i=1ψdk(xi)yi.
10: Eliminate arms with respect to criteria
Ŝk+1 = Ŝk \ {x ∈ Ŝk : ∃x′ such that (ψdk(x′)−ψdk(x))>θ̂k ≥ ωk(ψdk(x′)−ψdk(x))},
where ωk(y) = k + ‖y‖A−1k √
2 log(|Ŝk|2/δk). 11: end for Output: Output any arm in Ŝn+1.
where dk = deff(4 · 2−k) ≤ deff( ) since 4 · 2−k ≥ when k ≤ dlog2(2/ )e.
The rounding term rd(ζ) commonly appears in the sample complexity of linear bandits [13, 24]; and our rounding term is with respect to the lower-dimensional space after embedding, which only scales with dk rather than the ambient dimension. To further interpret the complexity, we define another complexity measure of a closely related linear bandit problem in the low-dimensional space and without model misspecifications.
ρ̃?d( ) := inf λ∈ΛX sup x∈X\{x?} ‖ψd(x?)−ψd(x)‖2Aψd (λ)−1 max{〈ψd(x?)−ψd(x),θd〉, }2 ,
where 〈ψd(x?)−ψd(x),θd〉 on the denominator represents the sub-optimality gap characterized by the linear component rather than the true sub-optimality gap h(x?)− h(x). The relation between ρ?( ) and ρ̃?( ) is discussed as follows.
Proposition 1. Suppose maxx∈X |h(x) − 〈ψd(x),θd〉| ≤ γ̃(d). For any ≥ γ̃(d), we have ρ?d( ) ≤ 9ρ̃?d( ). When γ̃(d) < ∆min/2, ρ̃?d(0) represents the sample complexity of a closely-related linear bandit problem without model misspecifications, i.e., h̃(x) = 〈ψd(x),θd〉. Remark 2. When γ̃(d) < ∆min/2, our sample complexity upper bound is relevant to the sample complexity of closely-related linear bandit problems without model misspecifications in lowerdimensional spaces. In fact, ρ̃?d(0) log(1/2.4δ) is the lower bound of the corresponding linear bandit problem in Rd [38, 13, 10]. Remark 3. Although the misspecification levels are generally known for situations considered in this paper, we also provide an algorithm that deals with unknown misspecification levels in Appendix D. Similar sample complexity guarantees are provided, but only in an unverifiably way (due to unknown misspecification levels): the algorithm starts to output -good arms after N samples, yet it doesn’t know when to stop. We refer readers to [23] for details on the unverifiable sample complexity.

heading: 3.2 Application to high-dimensional linear bandits
text: We apply the idea of adaptive embedding to high-dimensional linear bandits. We consider linear bandit problem of the form h = Xθ? where X ∈ RK×D and the i-th row of X represents the
feature vector of arm xi. We assume that ‖θ?‖2 ≤ C, which is commonly studied as the constrained linear bandit problem [41, 10].
LetX = UΣV > be the singular value decomposition (SVD) ofX , with singular values σ1 ≥ σ2 ≥ · · · ≥ σr > 0 for some r ≤ min{K,D}. Let ui,j denote the (i, j)-th entry of matrix U and u:,i denote the i-th column of U (similar notations for V ). We have
h = Xθ? = UΣV >θ? =
d∑
i=1
σiu:,iv > :,iθ? +
D∑
i=d+1
σiu:,iv > :,iθ? =:
d∑
i=1
σiu:,iv > :,iθ? + η,
where ‖η‖∞ ≤ C ∑D i=d+1 σi. As a result, for any d ≤ r, we can construct the feature mapping ψd(xi) = [σ1ui,1, . . . , σdui,d] > ∈ Rd such that h(xi) = 〈ψd(xi), θ̃?〉 + η(xi), where θ̃? = [ V >θ ] [1:d]
∈ Rd is the associated reward parameter.3 The upper bound of the induced misspecification can be expressed as γ̃(d) = C ∑D i=d+1 σi, which allows us to calculate γ(d). We can then apply Algorithm 1 to identify an -optimal arm. A high-dimensional linear bandit instance is provided in Appendix B.3 showing that: Algorithm 1 takes Õ(1/ 2) samples to identify an -optimal arm, while the sample complexity upper bound of RAGE scales as Õ(D/ 2).

heading: 4 Pure exploration in RKHS
text: We consider a kernel function K : Z ×Z → R over a compact set Z; we assume the kernel function satisfies condition stated in the Mercer’s Theorem (see Appendix E.1) and has eigenvalues decay fast enough (see Assumption 1). LetH be the Reproducing Kernel Hilbert Space (RKHS) induced from K. We assume X ⊆ Z and the true reward of any arm x ∈ X is given by an unknown function h ∈ H such that ‖h‖H ≤ 1. Let {φj}∞j=1 and {µj}∞j=1 be sequences of eigenfunctions and non-negative eigenvalues associated with kernel K.4 A corollary of Mercer’s theorem shows that any h ∈ H can be written in the form of h(·) = ∑∞j=1 θjφj(·) for some {θj}∞j=1 ∈ `2(N) such that ∑∞ j=1 θ 2 j/µj < ∞. We also have ‖h‖2H = ∑∞ j=1 θ 2 j/µj . Although functions in RKHS are non-linear in nature, we now can represent them in terms of an infinite-dimensional linear function. We construct feature mappings for the embedding next. For any x ∈ X , we have h(x) = ∑∞j=1 θjφj(x) = ∑∞ j=1 θj√ µj √ µjφj(x). Let Cφ := supx∈X̃ ,j≥1|φj(x)|. Since ∑∞ j=1 θ 2 j/µj = ‖h‖2H ≤ 1 is bounded, for any d ∈ N, we define feature mapping ψd(x) = [ √ µ1φ1(x), . . . , √ µdφd(x)] > ∈ Rd such that
h(x) = 〈θd,ψd(x)〉+ ηd(x),
where θd = [θ1/ √ µ1, . . . , θd/ √ µd] > ∈ Rd and |ηd(x)| ≤ γ̃(d) := Cφ √∑ j>d µj . We remark here that the constant Cφ is calculable and usually mild, e.g., Cφ = 1 for φj(x) = sin((2j − 1)πx/2). We can then construct γ(d) and deff( ) as in Section 3 and specialize Algorithm 1 to the kernel setting. Both γ(d) and deff( ) depend on eigenvalues of the associated kernel. Fortunately, fast eigenvalue decay are satisfied by most kernel functions, e.g., Gaussian kernel. We quantify such properties through the following assumption.
Assumption 1. We consider kernels with the following eigenvalue decay with some absolute constants Ck and β.
1. Kernel K is said to have (Ck, β)-polynomial eigenvalue decay (with β > 3/2) if µj ≤ Ckj−β for all j ≥ 1. 3We note that the embeddings and associated quantities can also be constructed on the fly with respect to the set of uneliminated arms. 4With a known kernel, the sequence of eigenfunctions and eigenvalues can be analytically calculated or numeriaclly approximated [36, 35]. We assume the knowledge of eigenfunctions and eigenvalues in this paper.
2. Kernel K is said to have (Ck, β)-exponential eigenvalue decay (with β > 0) if µj ≤ Cke−βj for all j ≥ 1.
Theorem 2. Suppose Assumption 1 holds. For any > 0, the following statements hold when we specialize Algorithm 1 to the kernel setting.
1. Suppose K has (Ck, β)-polynomial eigenvalue decay. We have deff( ) ≤ O( −2/(2β−3)), and the sample complexity of identifying an -optimal arm is upper bounded by Õ( −2/(2β−3) max{∆min, }−2).
2. Suppose K has (Ck, β)-exponential eigenvalue decay. We have deff( ) ≤ O(log(1/ )), and the sample complexity of identifying an -optimal arm is upper bounded by Õ(max{∆min, }−2).
Remark 4. Our sample complexity guarantees are directly related to the eigenvalue decay of the underlying kernel function, rather than the empirical kernel matrix as studied in previous works [7, 42]. Although one can also provide an instance dependent bound as in Theorem 1, the worst-case sample complexity bound in Theorem 2 provides insightful characterizations of the sample complexity in terms of eigenvalue decay. One should notice that with exponential eigenvalue decay, the sample complexity Õ( −2) essentially matches, up to logarithmic factors, the complexity of distinguishing a two-armed bandit up to accuracy [25].

heading: 5 Pure exploration with neural networks
text: In this section we present a neural network-based pure exploration algorithm in Algorithm 2. Our algorithm is inspired by the recently proposed neural bandits algorithms for regret minimization [47, 46]. At the core of our algorithm is to use a neural network f(x;θ) to learn the unknown reward function h. Specifically, following [8, 47], we consider a fully connected neural network f(x;θ) with depth L ≥ 2
f(x;θ) = √ mWLσ ( WL−1σ ( · · ·σ(W1x) )) , (2)
where σ(x) := max(x, 0) is the ReLU activation function, W1 ∈ Rm×d,WL ∈ R1×m, and for 2 ≤ i ≤ L − 1, Wi ∈ Rm×m. Moreover, we denote θ = [vec(W1)>, . . . , vec(WL)>]> ∈ Rp, where p = m+md+m2(L− 2) is the number of all the network parameters. We use g(x;θ) = ∇θf(x;θ) to denote the gradient of the neural network output with respect to the weights. In detail, at k-th iteration, Algorithm 2 firstly applies its current gradient mapping g(x;θk−1) over the whole action set X , and obtains the collection of gradientsG ∈ R|X |×p. Then Algorithm 2 does SVD overG and constructs a dk-dimensional feature mappingψdk , which can be regarded as the projection of the gradient feature mapping g(x;θk−1) to the most informative dk-dimensional eigenspace. Here we choose dk such that the summation of the eigenvalues of the remaining eigenspace be upper bounded by some error ̄. Algorithm 2 then computes the optimal design λk over ψk(X ) and pulls arms {x1, . . . ,xNk} based on both the design λk and the total number of allocations Nk. Finally, Algorithm 2 trains a new neural network f(x;θk) using gradient descent starting from the initial parameter θ0 (details are deferred to Appendix F), then eliminates the arms x in the current arm set Ŝk which are sub-optimal with respect to the neural network function value f(x;θk). The main difference between Algorithm 2 and its RKHS counterpart is as follows. Unlike Algorithm 1 which works on known feature mappings ψd (derived from a known kernel K), Algorithm 2 does not have information about the feature mapping, and thus it constructs the feature mapping from the raw high-dimensional up-to-date gradient mapping g(x;θk−1). The feature mapping is constructed with respect to a trained neural work, which leverages the great representation power of neural networks. This makes Algorithm 2 a more general and flexible algorithm than Algorithm 1.
Now we present the main theorem of Algorithm 2. Let H |X |×|X| be the Neural Tangent Kernel (NTK)[17] gram matrix over all arms X (the detailed definition ofH is deferred to Appendix F). We define the effective dimension for the neural version as below. The definition is similar to Definition 1.
Algorithm 2 Neural Arm Elimination Input: Action set X , initial parameter θ0, neural network f(x;θ), gradient mapping g(x,θ),
width of the matrix m, parameter of the number of allocations A, approximation parameter ζ, regularization parameter α, error parameter ̄, , confidence level δk = δ/(8k2) 1: Set Ŝ1 = X . 2: for k = 1, 2, . . . , n do 3: Construct the truncated feature representation ψk(X ) based on gradient mapping g(x;θk−1).
In detail, letG ∈ R|X |×p be the collection of gradients such that G = [g(x1;θk−1) >; . . . ; g(x|X |;θk−1) >]/ √ m ∈ R|X |×p (3)
Let [U ,Σ,V ] be the SVD of G, where U = (ui,j) ∈ R|X |×|X|, Σ = [diag(e1, . . . , e|X |), 0] ∈ R|X |×p, V ∈ Rp×p. Get dk = min{d ∈ [|X |] : ∑|X | i=d+1 ei ≤ ̄},
and set ψdk(xi) = (e1ui,1, . . . , edkui,dk) ∈ Rdk . 4: Set λk and τk be the experimental design and the value of the following optimization problem
inf λ∈ΛX sup y∈Y(ψdk (Ŝk)) ‖y‖2Aψdk (λ)−1 . (4)
5: Set Nk = max { 22k A(1 + ζ) log ( |X |2/δk ) , rdk(ζ)}. 6: Get {x1,x2, . . . ,xNk} = ROUND(λk, Nk, dk, ζ). 7: Pull arms {x1,x2, . . . ,xNk} and receive rewards {y1, . . . , yNk}. 8: Using Jk step ηk-step size gradient descent to optimize the following loss function to obtain
θk, θk = arg minL(θ) := Nk∑
j=1
(f(xj ;θ)− yj)2 + mα
2 ‖θ − θ0‖22. (5)
9: SetAk = αI + ∑Nk i=1ψdk(x)ψdk(x)
> and eliminate arms with respect to criteria Ŝk+1 = Ŝk \ {x ∈ Ŝk : ∃x′ such that f(x′;θk)− f(x;θk) ≥ 2−k/8 + 3 /8}.
10: end for Output: Output any arm in Ŝn+1.
Definition 2 (Neural version). For any > 0, we define the effective dimension as deff( ) := min{d ∈ [|X |] : ∑|X |i=d+1 λi(H) ≤ }, where λi(H) is the i-th eigenvalue ofH .
Next, we make standard assumptions for the initialization of neural networks and the arms x ∈ X . Assumption 2 ([47]). There exists λ0 > 0 such thatH λ0I . For any x ∈ X , the arm x satisfies ‖x‖2 = 1 and that its j-th coordinate is identical to its j + d/2-th coordinate. Meanwhile, the initial parameter θ0 = [vec(W1)>, . . . , vec(WL)>]> is initialized as follows: for 1 ≤ l ≤ L− 1, Wl is set to ( W 0 0 W ) , where each entry ofW is generated independently from N (0, 4/m);WL is set to (w>,−w>), where each entry of w is generated independently from N (0, 2/m).
We now present our main theorem for pure exploration with neural network approximation. The formal version of our theorem is deferred to Appendix F.
Theorem 3 (Informal). Under Assumption 2, with proper selection of parameters α, n, ̄, A, ηk, Jk, then when m = poly(|X |, L, λ−10 , log(|X |/δk), Nk, α, ̄−1), with probability at least 1 − δ, ŜK+1 only includes arm x satisfying ∆x ≤ , and the total sample complexity of Algorithm 2 is bounded by
N = Õ ( (1 + ζ)deff(̄ 2/|X |)/ 2 + rdeff (̄2/|X |)(ζ) ) = Õ ( deff(̄ 2/|X |) −2 ) .
Remark 5. For the case where the effective dimension can be well bounded, e.g., deff(̄2/|X |) = O(log(|X |/̄2)), Theorem 7 suggests that Algorithm 2 is able to identify an -optimal arm within Õ( −2) samples. That suggests that our neural network-based algorithm is efficient without constructing a low-dimensional linear approximation of h in prior, like the previous linear or RKHS approaches.

heading: 6 Experiments
text: We conduct four experiments on synthetic and real-world datasets. We specialize our embedding idea to the neural, kernel, and linear regimes, and denote the algorithms as NeuralEmbedding (Algorithm 2), KernelEmbedding (Algorithm 1 with Gaussian kernel), and LinearEmbedding (Algorithm 1 with linear representation), respectively. We compare our algorithms with two baselines: RAGE and ActionElim. RAGE [13] conducts pure exploration in linear bandits and ActionElim [18, 11] ignores all feature representations. The (empirical) sample complexity of each algorithm is calculated as the number of samples needed so that the uneliminated set contains only -optimal arms. Unsuccessful runs, i.e., those terminate with non- -optimal arms, are reported as failures. In our experiments, we set = 0.1 and δ = 0.05. All results are averaged over 50 runs.5
Synthetic datasets. We first generate the feature matrix X̃ = X + E ∈ RK×D where X is constructed as a rank-2 matrix andE is a perturbation matrix with tiny spectral norm (See Appendix G for details). Each row of X̃ represents the feature representation of an arm, and those features can be grouped into two clusters with equal size. In the case with linear rewards, we let θ? equal to the first row of X̃ . For nonlinear rewards, the reward of each arm is set as the 2-norm of its feature representation. We vary the number of arms K and the ambient dimension D in our experiments.
Fig. 1 shows experimental results on synthetic datasets. All algorithms successfully identify - optimal arms with zero empirical failure probability (due to the simplicity of the datasets). In terms of sample complexity, NeuralEmbedding outperforms all other algorithms in most cases, and KernelEmbedding and LinearEmbedding significantly outperform RAGE and Action Elim. The sample complexities of NeuralEmbedding, KernelEmbedding and LinearEmbedding are not affected when increasing number of arms or dimensions since they first identify the important subspace and then conduct elimination. On the other side, the sample complexity of ActionElim gets larger with increasing number of arms and the sample complexity of RAGE gets larger with increasing dimensions.
MNIST dataset. The MNIST dataset [29] contains hand-written digits from 0 to 9. We view each digit as an arm, and set their rewards according to the confusion matrix of a trained classifier. Digit 7 is chosen as the optimal arm with reward 1; the reward of digits 1, 2 and 9 are set to be 0.8, and all other digits have reward 0.5. In each experiment, we randomly draw 200 samples (20 samples each digit) from the dataset. We project the raw feature matrixX ∈ R200×784 into a lower-dimensional space X̃ ∈ R200×200 so that it becomes full rank (but without losing any information) and feasible for RAGE. Our goal is to correctly identify a digit 7.
Yahoo dataset. The Yahoo! User Click Log Dataset R6A6 contains users’ click-through records. Each record consists of a 36-dimensional feature representation (obtained from an outer product), and a binary outcome stating whether or not a user clicked on the article. We view each record as
5All algorithms are elimination-styled for fair comparison. Other implementation details are deferred to Appendix G.
6https://webscope.sandbox.yahoo.com
an arm, and set the reward as 0.8 (if clicked) or 0.3 (if not clicked) to makes the problem harder. In each experiment, we randomly draw 200 arms from the dataset, where 5 of them having rewards 0.8 (proportional to true click-through ratio), Our goal is to identify an arm with rewards 0.8. Our experimental setup is similar to the one used in Fiez et al. [13]. However, their true rewards are obtained from a least square regression. We do not enforce linearity in rewards in our experiment.
Box plots in Fig. 2 show the sample complexity of each algorithm on real-world datasets. NeuralEmbedding significantly outperforms all other algorithms thanks to (1) the representation power of neural networks and (2) efficient exploration in low-dimensional spaces. KernelEmbedding and LinearEmbedding have competitive performance on the MNIST dataset. Table 1 shows the success rate of each algorithm. Linear methods such as RAGE and LinearEmbedding have relatively low success rates on the Yahoo dataset (with nonlinear rewards). Our NeuralEmbedding and KernelEmbedding methods have high success rates since they are designed for nonlinear setting.

heading: 7 Conclusion
text: We introduce the idea of adaptive embedding in bandit pure exploration. Unlike existing works that passively deal with model misspecifications, we adaptively embed high-dimensional feature representations into lower-dimensional spaces to avoid the curse of dimensionality. The induced misspecifications are carefully dealt with. We further apply our approach to two under-studied settings with the nonlinearity: (1) pure exploration in an RKHS and (2) pure exploration with neural networks. Our sample complexity guarantees depend on the effective dimension of the feature spaces in the kernel or neural representations. We conduct extensive experiments on both synthetic and real-world datasets, and our algorithms greatly outperform existing ones.
Our current analysis with neural representations is in the NTK regime, which can only describe a part of the representation of the neural networks. We leave extending our algorithm to more general settings (beyond the NTK regime) as a future direction.
Acknowledgments and Disclosure of Funding
We thank the anonymous reviewers and area chair for their helpful comments. YZ and RN are partially supported by AFOSR grant FA9550-18-1-0166. DZ and QG are partially supported by the National Science Foundation CAREER Award 1906169, IIS-1904183 and AWS Machine Learning Research Award. RJ and RW are partially supported by AFOSR FA9550-18-1-0166, NSF OAC-1934637, NSF DMS-2023109, and NSF DGE-2022023. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.

Review summaries:
This paper introduces a new approach to tackle bandit pure exploration in high dimensional setting. Recall that the literature generally deals with a finite number of arms, or scale exponentially with the dimension of the feature space. The idea, developed in Algorithm 1, is as follows. At each iteration of a RAGE-like algorithm, the arms are embedded into a lower dimensional space. The dimension is chosen as the smallest integer such that the induced error does not harm the subsequent computations and guarantees. Guarantees for Algorithm 1 are derived in Theorem 1, and expressed in terms of the effective dimension. An example of sequence of feature maps is given for linear bandits, based on the SVD decomposition of the design matrix. Under boundedness of parameter $\theta_\star$, it is possible to control the induced error by the successive approximations. When the reward function is assumed to be in a RKHS, using the eigendecomposition of the associated integral operator allows to design similarly a series of feature maps with controllable error. The effective dimension can be bounded in terms of the eigendecay of the kernel. The reward function can also be approximated using a neural network. The algorithm (Algorithm 2) then also needs to learn the weights of the network jointly. Guarantees for Algorithm 2 are stated in Theorem 3. Finally, Section 6 presents some experiments on synthetic and real world datasets showing that the proposed Algorithms exhibit smaller sample complexities than Action Elim.

This paper considers the best arm identification problem for kernel and neural bandits, where the reward function belongs to an RKHS and is approximated by a neural network, respectively. The main contribution is that the authors propose to consider a variant of the Randomized Adaptive Gap Elimination (RAGE) algorithm with adaptive embedding (so that the problem can be viewed as misspecified linear bandits) and then address the corresponding misspecification issue. The sample complexity results are established under known feature mappings and misspecification errors and then adapted to the cases of RKHS and neural bandits. Experimental results on both synthetic datasets and real datasets are provided to validate the proposed algorithms. 

Summary: This paper studies \delta-PAC \epsilon-best arm identification when a kernel or neural representation is used to embed the arms. The paper targets regimes where the number of arms K is really large and the similarity between arms can be exploited to reduce the sample complexity of the tests needed to find the best arm.  General remark: The paper is well written and addresses a hard problem. The algorithms are smart and the empirical results are impressive. Despite a few minor remarks listed below, I recommend accepting this paper.  Disclaimer: Unfortunately, my level of expertise does not allow me to verify Appendix F.