Optimal Gradient-based Algorithms for Non-concave Bandit Optimization

Abstract:
Bandit problems with linear or concave reward have been extensively studied, but relatively few works have studied bandits with non-concave reward. This work considers a large family of bandit problems where the unknown underlying reward function is non-concave, including the low-rank generalized linear bandit problems and two-layer neural network with polynomial activation bandit problem. For the low-rank generalized linear bandit problem, we provide a minimax-optimal algorithm in the dimension, refuting both conjectures in [55, 43]. Our algorithms are based on a unified zeroth-order optimization paradigm that applies in great generality and attains optimal rates in several structured polynomial settings (in the dimension). We further demonstrate the applicability of our algorithms in RL in the generative model setting, resulting in improved sample complexity over prior approaches. Finally, we show that the standard optimistic algorithms (e.g., UCB) are sub-optimal by dimension factors. In the neural net setting (with polynomial activation functions) with noiseless reward, we provide a bandit algorithm with sample complexity equal to the intrinsic algebraic dimension. Again, we show that optimistic approaches have worse sample complexity, polynomial in the extrinsic dimension (which could be exponentially worse in the polynomial degree).

heading: 1 Introduction
text: Bandits [51] are a class of online decision-making problems where an agent interacts with the environment, only receives a scalar reward, and aims to maximize the reward. In many real-world applications, bandit and RL problems are characterized by large or continuous action space. To encode the reward information associated with the action, function approximation for the reward function is typically used, such as linear bandits [20]. Stochastic linear bandits assume the mean reward to be the inner product between the unknown model parameter and the feature vector associated with the action. This setting has been extensively studied, and algorithms with optimal regret are known [20, 51, 2, 13].
However, linear bandits suffer from limited representation power unless the feature dimension is prohibitively large. A comprehensive empirical study [60] found that real-world problems required non-linear models and thus non-concave rewards to attain good performance on a testbed of bandit problems. To take a step beyond the linear setting, it becomes more challenging to design optimal algorithms. Unlike linear bandits, more sophisticated algorithms beyond optimism are necessary. For instance, a natural first step is to look at quadratic [44] and higher-order polynomial [35] reward. In
∗Alphabetical order. Correspondence to: Qi Lei, qilei@princeton.edu, Jason D. Lee, jasonlee@ princeton.edu.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
the context of phase retrieval, which is a special case for the quadratic bandit, people have derived algorithms that achieve minimax risks in the statistical learning setting [9, 52, 15]. However, the straightforward adaptation of these algorithms results in sub-optimal dimension dependency.
In the bandit domain, existing analysis on the nonlinear setting includes eluder dimension [62], subspace elimination [55, 43], etc. Their results also suffer from a larger dimension dependency than the best known lower bound in many settings. (See Table 1 and Section 3 for a detailed discussion of these results.) Therefore in this paper, we are interested in investigating the following question:
What is the optimal regret for non-concave bandit problems, including structured polynomials (low-rank etc.)? Can we design algorithms with optimal dimension dependency?
Contributions: In this paper, we answer the questions and close the gap (in problem dimension) for various non-linear bandit problems.
1. First, we design stochastic zeroth-order gradient-like2 ascent algorithms to attain minimax regret for a large class of structured polynomials. The class of structured polynomials contains bilinear and low-rank linear bandits and symmetric and asymmetric higher-order homogeneous polynomial bandits with action dimension d. Though the reward is non-concave, we combine techniques from two bodies of work, non-convex optimization and numerical linear algebra, to design robust gradient-based algorithms that converge to global maxima. Our algorithms are also computationally efficient, practical, and easily implementable. In all cases, our algorithms attain the optimal dependence on dimension d, which was not previously attainable using existing optimism techniques. As a byproduct, our algorithm refutes3 the conjecture from [43] on the bilinear bandit, and the conjecture from [55] on low-rank linear bandit by giving an algorithm that attains the optimal dimension dependence. 2. We demonstrate that our techniques for non-concave bandits extend to RL in the generative setting, improving upon existing optimism techniques. 3. When the reward is a general polynomial without noise, we prove that solving polynomial equations achieves regret equal to the intrinsic algebraic dimension of the underlying polynomial class, which is often linear in d for interesting cases. In general, this complexity cannot be further improved. 4. Furthermore, we provide a lower bound showing that all UCB algorithms have a sample complexity of Ω(dp), where p is the degree of the polynomial. The dimension of all homogeneous polynomials of degree p in dimension d is dp, showing that UCB is oblivious to the polynomial class and highly sub-optimal even in the noiseless setting.

heading: 1.1 Related Work
text: Linear Bandits. Linear bandit problems and their variants are studied in [20, 51, 1, 13, 2, 20, 61, 34, 33]. The matching upper bound and minimax lower bound achievesO(d √ T ) regret. Structured linear bandits, including sparse linear bandit [3] which developed an online-to-confidence-set technique. This technique yields the optimal O( √ sdT ) rate for sparse linear bandit. However [55] employed
the same technique for low-rank linear bandits giving an algorithm with regret O( √ d3poly(k)T )
which we improve to O( √ d2poly(k)T ), which meets the lower bound given in [55]. Eluder Dimension. [62] proposed the eluder dimension as a general complexity measure for nonlinear bandits. However, the eluder dimension is only known to give non-trivial bounds for linear function classes and monotone functions of linear function classes. For structured polynomial classes, the eluder dimension simply embeds into an ambient linear space of dimension dp, where d is the dimension and p is the degree. This parallels the linearization/NTK line in supervised learning [71, 30, 7] which show that linearization also incurs a similarly large penalty of dp sample complexity, and more advanced algorithm design is need to circumvent linearization [10, 17, 25, 72, 27, 59, 29, 57, 36, 68, 19]. Non-concave Bandits. To our knowledge, there is no general study of non-concave bandits, likely due to the difficulty of globally maximizing non-concave functions. A natural starting point of
2Our algorithm estimates the gradient, but with some irreducible bias for the tensor case. Note that our algorithms converge linearly despite the bias.
3Both papers conjectured regret of the form O( √
d3poly(k)T ) based on convincing but potentially misleading heuristics.
studying the non-concave setting are quadratic rewards such as the Rayleigh quotient, or namely bandit PCA [47, 28]. In the bilinear [43] and low-rank linear setting [55] with rank k parameter matrices achieved Õ( √ d3poly(k)T ) regret. Other literature considered related but different settings that are not comparable to our results [44, 42, 32, 48]. We note that the regret of all previous work is at least O( √ d3T ). This includes the subspace exploration and refinement algorithms from [43, 55], or from eluder dimension [62]. Recently [35] considers online problem with a low-rank tensor associated with axis-aligned set of arms, which corresponds to finding the largest entry of the tensor. Finally, [47] study the bandit PCA problem in the adversarial setting, attaining regret of O( √ d3T ). We leave adapting our results to the adversarial setting as an open problem. Due to space limitation we defer the discussion of some equally important literature to the appendix.

heading: 2 Preliminaries

heading: 2.1 Setup: Structured Polynomial Bandit
text: We study structured polynomial bandit problems where the reward function is from a class of structured polynomials (see below). A player plays actions for T rounds; at each round t ∈ [T ], the player chooses one action at from the feasible action set A and receives the reward rt afterward. We consider both the stochastic case where rt = fθ(at) + ηt where ηt is the random noise, and the noiseless case rt = fθ(at). Specifically the function fθ is unknown to the player, but lies in a known function class F . We use the notation vec(M) to denote the vectorization of a matrix or a tensor M , and v⊗p to denote the p-order tensor product of a vector v. For vectors we use ∥ · ∥2 or ∥ · ∥ to denote its ℓ2 norm. For matrices ∥ · ∥2 or ∥ · ∥ stands for its spectral norm, and ∥ · ∥F is Frobenius norm. For integer n, [n] denotes set {1, 2, · · ·n}. For cleaner presentation, in the main paper we use Õ, Θ̃ or Ω̃ to hide universal constants, polynomial factors in p and polylog factors in dimension d, error ε, eigengap ∆, total round number T or failure rate δ.
We now present the outline with the settings considered in the paper:
The stochastic bandit eigenvector case, FEV, considers action set A = {a ∈ Rd : ∥a∥2 ≤ 1} as shown in Section 3.1, and
FEV = { fθ(a) = a TMa,M = ∑k j=1 λjvjv ⊤ j , for orthonormal vj
M ∈ Rd×d, 1 ≥ λ1 ≥ |λ2| ≥ · · · ≥ |λk|
} .
The stochastic low-rank linear reward case, FLR considers action sets on bounded matrices A = {A ∈ Rd×d : ∥A∥F ≤ 1} in Section 3.2, and
FLR = { fθ(A) = ⟨M ,A⟩ = vec(M)⊤vec(A), M ∈ Rd×d, rank(M) = k,M = M⊤, ∥M∥F ≤ 1 } .
We illustrate how to apply the established bandit oracles to attain a better sample complexity for RL problems with the simulator in Section 3.2.1.
The stochastic homogeneous polynomial reward case is presented in Section 3.3. For the symmetric case in Section 3.3.1, the action sets are A = {a ∈ Rd : ∥a∥2 ≤ 1}, and
FSYM = { fθ(a) = ∑k j=1 λj(v ⊤ j a)
p for orthonormal vj , 1 ≥ r∗ = λ1 > |λ2| ≥ · · · ≥ |λk|
} ;
in the asymmetric case in Section D.3.1, the action sets are A = {a = a(1)⊗ a(2)⊗ · · · ⊗ a(p) ∈ Rdp : ∀q ∈ [p], ∥a(q)∥2 ≤ 1}, and
FASYM = { fθ(a) = ∑k j=1 λj ∏p q=1(vj(q)
⊤a(q)) for orthonormal vj(q) for each q, 1 ≥ r∗ = |λ1| ≥ |λ2| ≥ · · · ≥ |λk|
} .
In the above settings, there is stochastic noise on the observed rewards. We also consider noiseless settings as below: The noiseless polynomial reward case is presented in Section 3.4. The action sets A are subsets of Rd, and FP = {fθ(a) = ⟨θ, ã⊗p⟩ : θ ∈ V, ã = [1,a⊤]⊤,V ⊆ (Rd+1)⊗p is an algebraic variety}.
Additionally, F needs to be admissible (Definition 3.20). This class includes two-layer neural networks with polynomial activations (i.e. structured polynomials). We study the fundamental limits of all UCB algorithms in Section G.0.1 as they are Ω(dp−1) worse than our algorithm presented in Section 3.4.1.
In the above settings, we are concerned with the cumulative regret R(T ) for T rounds. Let f∗θ =
sup a∈A
fθ(a), Rθ(T ) := T∑
t=1 (f∗θ − fθ(at)). Since the parameters can be chosen adversarially, we are
bounding R(T ) = supθ Rθ(T ) in this paper. In all the stochastic settings above, we make the standard assumption on stochasticity that ηt is conditionally zero-mean 1-sub-Gaussian random variable regarding the randomness before t.

heading: 2.2 Warm-up: Adapting Existing Algorithms
text: In all of the above settings, the function class can be viewed as a generalized linear function. Namely, there is fixed feature maps ψ, ϕ so that fθ(a) = ψ(θ)Tϕ(a). Thus it is straightforward to adapt linear bandit algorithms like the renowned LinUCB [53] to our settings. Furthermore, another baseline is given by the eluder dimension argument [62][67] which gives explicit upper bounds for general function classes. We present the best upper bound by adapting these methods as a baseline in Table 1, together with our newly-derived lower bound and upper bound in this paper. The best-known statistical rates are based on the following result. Theorem 2.1 (Proposition 4 in [62]). With α = O(T−2) appropriately small, given the α-coveringnumber N (under ∥·∥∞) and the α-eluder-dimension dE of the function class F , Eluder UCB (Algorithm 2) achieves regret Õ( √ dET logN).
In the first row of Table 1, we further elaborate on the best results obtained from Theorem 2.1 in individual settings. More details can be found in Appendix B.

heading: 3 Main results
text: We now present our main results. We consider 4 different stochastic settings (see Table 1) and one noiseless setting with structured polynomials. In the cases of stochastic reward, all our algorithms can be unified as gradient-based optimization. At each stage with a candidate action a, we define the estimatorGn(a) := 1n ∑n i=1(fθ((1−ζ)a+ζzi)+ ηi)zi, with zi ∼ N (0, σ2Id) and proper step-size ζ [26]. Therefore Ez[G(a)] = σ2ζ∇fθ((1 − ζ)a)+O(ζ2) = ζ(1−ζ)p−1σ2∇fθ(a)+O(ζ2) for p-th order homogeneous polynomials. Therefore with enough samples, we are able to implement noisy gradient ascent with bias. In the noiseless setting, our algorithm solves for the parameter θ with randomly sampled actions {at} and the noiseless reward {fθ(at)}, and then determines the optimal action by computing argmaxa fθ(a).
3.1 Stochastic Eigenvalue Reward (FEV)
Now consider bandits with stochastic reward r(a) = a⊤Ma+η with action setA = {a|∥a∥2 ≤ 1}. M = ∑r i=1 λiviv ⊤ i is symmetric and satisfies r
∗ = λ1 > |λ2| ≥ |λ3| · · · ≥ |λr|, η ∼ N (0, 1). Denote by a∗ the optimal action (±v1), the leading eigenvector of M , (a∗)⊤Ma∗ = λ1. Let ∆ = λ1 − |λ2| > 0 be the eigengap and κ := λ1/∆ be the condition number. Remark 3.1 (Negative leading eigenvalue). For a symmetric matrix M , we will conduct noisy power method to recover its leading eigenvector, and therefore we require its leading eigenvalue λ1
to be positive. It is straightforward to extend to the setting where the nonzero eigenvalues satisfy: r∗ ≡ λ1 > λ2 ≥ λl > 0 > λl+1 · · · ≥ λk, and |λk| > λ1. For this problem, we can shift M to get M + |λk|I and the eigen-spectrum now becomes λ1 + |λk|, λ2 + |λk|, · · · 0; therefore, we can still recover the optimal action with dependence on the new condition number (λ1 + |λk|)/(λ1 − λ2). Remark 3.2 (Asymmetric matrix). Our algorithm naturally extends to the asymmetric setting: f(a1,a2) = a ⊤ 1 M̃a2, where M̃ = UΣV
⊤. This setting can be reduced to the symmetric case via defining
M =
[ 0 M̃⊤
M̃ 0
] = 1
2
[ V V U −U ] [ Σ 0 0 −Σ ] [ V V U −U ]⊤ ,
which is a symmetric matrix, and its eigenvalues are ±σi(M̃), the singular values of M̃ . Therefore our analysis on symmetric matrices also applies to the asymmetric setting and will equivalently depend on the gap between the top singular values of M̃ . A formal asymmetric to symmetric conversion algorithm is presented in Algorithm 1 in [28].
Algorithm. We note that by conducting zeroth-order gradient estimate 1/n ∑n
i=1(f(a/2 + zi/2) + ηi)zi with step-size 1/2 [26] and sample size n, we get an estimate for Eη,z[(f(a/2+z/2)+η)z] = σ2
2 Ma when z ∼ N (0, σ 2). Therefore we are able to use noisy power method to recover the top eigenvector. We present the complete algorithm in Algorithm 1 and attain a gap-dependent risk bound:
Algorithm 1 Noisy power method for bandit eigenvalue problem. 1: Input: Quadratic function f : A → R with noisy reward, failure probability δ, error ε. 2: Initialization: Initial action a0 ∈ Rd randomly sampled on the unit sphere Sd−1. We set α = |λ2/λ1|, sample size per iteration n = Cnd2 log(d/δ)(λ1α)−2ε−2, sample variance m = Cmd log(n/δ), total iteration L = ⌊CL log(d/ε)⌋+ 1.
3: for Iteration l from 1 to L do 4: Sample zi ∼ N (0, 1/mId), i = 1, 2, · · ·n. (Re-sample the whole batch if exists zi with
norm greater than 1.) 5: Noisy power method: 6: Take actions ãi = al−1+zi 2 and observe ri = f(ai) + ηi,∀i ∈ [n]
7: Update normalized action al ← mn ∑n
i=1 rizi, and normalize al ← al/∥al∥2. 8: Output: aL.
Theorem 3.3 (Regret bound for noisy power method (NPM)). In Algorithm 1, we set ε ∈ (0, 1/2), δ = 0.1/(L0S) and let CL, CS , Cn, Cm be large enough universal constants. Then with high probability 0.9 we have: the output aL satisfies tan θ(a∗,aL) ≤ ε and yields r∗ε2-optimal reward; and the total number of actions we take is Õ( κd 2
∆2ε2 ). By explore-then-commit (ETC) the cumulative regret is at most Õ( √ κ3d2T ).
All proofs in this subsection are in Appendix C.1.
Remark 3.4 (Intuition of [43] and how to overcome the conjectured lower bound via the design of adaptive algorithms). Let us consider the rank 1 case of r(a) = (aTθ∗)2 + η. A random action a ∼ Unif(Sd−1) has f(a) ≍ 1/d2, and the noise has standard deviation O(1). Thus the signalto-noise-ratio is O(1/d2) and the optimal action θ∗ requires d bits to encode. If we were to play non-adaptively, this would require O(d3) queries and result in regret √ d3T which matches the result of [43].
To go beyond this, we must design algorithms that are adaptive, meaning the information in f(a) + η is strictly larger than 1d2 . As an illustration of why this is possible, consider batching the timesteps into d stages so that each stage decode 1 bit of θ∗. At the first stage, random exploration a ∼ Unif(Sd−1) gives signal-to-noise-ratio O(1/d2). Suppose k bits of θ are decoded at k-th stage by θ̂, adaptive algorithms can boost the signal-to-noise-ratio to O(k/d2) by using θ̂ as bootstrap (e.g. exploring with θ̂ ± a where a is random exploration in the unexplored subspace). In this way
adaptive algorithms only need d2/k queries in (k + 1)-th stage and so the total number of queries sums up to ∑d k=1 d
2/k ≈ d2 log d. Gradient descent and power method offer a computationally efficient and seamless way to implement the above intuition. For every iterate action a, we estimate Ma from noisy observations and take it as our next action a+. With d2/(∆2ε2) samples, noisy power method enjoys linear progress tan θ(a+,a∗) ≤ max{c tan(a,a∗), ε}, where c < 1 is a constant that depends on λ1, λ2, and ε. Therefore even though every step costs d2 samples, overall we only need logarithmic (in d, ε, λ1, λ2) iterations to find an ε-optimal action. Remark 3.5 (Connection to phase retrieval and eluder dimension). For rank-1 case M = xx⊤, the bilinear bandits can be viewed as phase retrieval, where one observes yr = (a⊤r x)
2 = a⊤r Mar plus some noise ηr ∼ N (0, σ2). The optimal (among non-adaptive algorithms) sample complexity to recover x is σ2d/ϵ2 [16, 15] where they play a from random Gaussian N (0, I). However, in bandit, we need to set the variance of a to at most 1/d to ensure ∥a∥ ≤ 1. Our problem is equivalent to observing yr/d where their a/ √ d ∼ N (0, 1/dI) and noise level ηr/d ∼ N (0, 1), i.e., σ2 = d2. Therefore one gets d3/ϵ2 even for the rank-1 problems. On the other hand, for all rank-1 M the condition number κ = 1; and thus our results match the lower bound (see Section 3.3.2) up to logarithmic factors, and also have fundamental improvements for phase retrieval problems by leveraging adaptivity.
For UCB algorithms based on eluder dimension, the regret upper bound is O( √ dE log(N)T ) = O( √ d3T ) as presented in Theorem 2.1, where the dependence on d is consistent with [43, 55] and is non-optimal.
The previous result depends on the eigen-gap. When the matrix is ill-conditioned, i.e., λ1 is very close to λ2, we can obtain gap-free versions with a modification: The first idea stems from finding higher reward instead of recovering the optimal action. Therefore when λ1 and λ2 are very close (gap being smaller than desired accuracy ε), it is acceptable to find any direction in the span of (v1,v2). More formally, we care about the convergence speed of identifying any action in the space spanned by any top eigenvectors (whose associated eigenvalues are higher than λ1 − ε). Therefore the convergence speed will depend on ε instead of λ1 − λ2. Corollary 3.6 (Gap-free regret bound). For positive semi-definite matrix M , by setting α = 1 − ε2/2 in Algorithm 1 and performing ETC afterwards, one can obtain cumulative regret of Õ(λ
3/5 1 d 2/5T 4/5).
Again, the PSD assumption is not essential. For general symmetric matrices with λ1 ≥ λ2 ≥ · · · > 0 > · · ·λk. We can still conduct shifted power method on M − λkI , yielding a cumulative regret of (λ1 + |λk|)3/5d2/5T 4/5. Another novel gap-free algorithm requires to identify any top eigenspace V1:l, l ∈ [k]: V1:l is the column span of {v1,v2, · · ·vl}. Notice that in traditional subspace iteration, the convergence rate of recovering Vl depends on the eigengap ∆l := |λl|−|λl+1|. Meanwhile, since ∑k l=1 ∆l = λ1, at least one eigengap is larger or equal to λ1/k. Suppose ∆l∗ ≥ λ1/k; we can, therefore, set α = 1− 1/k and recover the top l∗ subspace up to λ1ϵ error, which will give an ϵ-optimal reward in the end. We don’t know l∗ beforehand and will try recovering the top subspace V1,V2, · · ·Vk respectively, which will only lose a k factor. With the existence of l∗, at least one trial (on recovering Vl∗ ) will be successful with the parameters of our choice, and we simply output the best action among all trials. Theorem 3.7 (Informal statement: (gap-free) subspace iteration). By running subspace iteration (Algorithm 3) with proper choices of parameters, we attain a cumulative regret of Õ(λ1/31 k
4/3(dT )2/3). Algorithm 3) with another set of parameters can also recover the whole eigenspace, and achieve cumulative regret of Õ((λ1k)1/3(κ̃dT )2/3), where κ̃ = λ1/|λk|.
3.2 Stochastic Low-rank Linear Bandits (FLR)
In the low-rank linear bandit, the reward function is f(A) = ⟨A,M⟩, with noisy observations rt = f(A) + ηt, and the action space is {A ∈ Rd×d : ∥A∥F ≤ 1}. Without loss of generality we assume k, the rank of M satisfies k ≤ d2 , since when k is of the same order as d, the known upper and lower bound are both √ d2k2T = Θ( √ d4T ) [55] and there is no room for improvement. We
write r∗ = ∥M∥F ≤ 1, therefore the optimal action is A∗ = M/r∗. In this section, we write X(s) to be the s-th column of any matrix X .
As presented in Algorithm 4, we conduct noisy subspace iteration to estimate the right eigenspace of M . Subspace iteration requires calculating MXt at every step. This can be done by considering a change of variable of g(X) := f(XX⊤) = ⟨XX⊤,M⟩ whose gradient4 is ∇g(X) = 2MX . The zeroth-order gradient estimator can then be employed to stochastically estimate MX . We instantiate the analysis with symmetric M while extending to asymmetric setting is straightforward since the problem can be reduced to symmetric setting (suggested in Remark 3.2).
With stochastic observations and randomly sampled actions, we achieve the next iterate Yl that satisfies MXl−1 ≡ E[Yl] in Algorithm 4. Let M = V ΣV ⊤. With proper concentration bounds presented in the appendix, we can apply the analysis of noisy power method [37] and get: Theorem 3.8 (Informal statement; sample complexity for low-rank linear reward). With Algorithm 4, XL satisfies ∥(I −XLX⊤L )V ∥ ≤ ε/4, and output Â satisfies ∥Â−A∗∥F ≤ ε∥M∥F with sample size Õ(d2kλ−2k ε −2).
We defer the proofs for low-rank linear reward in Appendix C.2. Corollary 3.9 (Regret bound for low-rank linear reward). We first call Algorithm 4 with ε4 = Θ̃(d2kλ−2k T −1) to obtain Â; we then play Â for the remaining steps. The cumulative regret satisfies
R(T ) ≤ Õ( √ d2k(r∗)2λ−2k T ),
with high probability 0.9.
To be more precise, we need T ≥ Θ̃(d2k/λ2k) for Algorithm 4 to take sufficient actions; however, the conclusion still holds for smaller T . Since simply playing 0 for all T actions will give a sharper bound of R(T ) ≤ r∗T ≤ Õ(r∗ √ d2kλ−2k T )). For cleaner presentation, we won’t stress this for every statement.
Notice that k ≤ ∥M∥2F /λ2k ≤ kκ̃2 is order k. Thus for well-conditioned matrices M , our upper bound of Õ( √ d2k∥M∥2F /λ2kT ) matches the lower bound √ d2k2T except for logarithmic factors.
In the previous setting with the bandit eigenvalue problem, estimating M up to an ϵ-error (measured by operator norm) gives us an ϵ-optimal reward. Therefore the sample complexity for eigenvalue reward with similar subspace iteration is ∥M∥22d2k/(λ2kϵ2). In this section, on the other hand, we need Frobenius norm bound ∥AL −A∗∥F ≤ ϵ; naturally the complexity becomes ∥M∥2F d2k/(λ2kϵ2). Theorem 3.10 (Regret bound for low-rank linear reward: gap-free case). Set ε6 = Θ( d 2k2
(r∗)2T ),
n = Θ̃( d 2k2 (r∗)2ε4 ), L = Θ(log(d/ε)) and k ′ = 2k in Algorithm 4 and get Â. Then we play it for the remaining steps, the cumulative regret satisfies:
R(T ) ≤ Õ((dkT )2/3(r∗)1/3). 4Directly performing projected gradient descent on f(A) would not work, since this is not an adaptive
algorithm as the gradient of a linear function is constant. This would incur regret √ d3poly(k)T .
We summarize all our results and prior work for quadratic reward in Table 2.
3.2.1 RL with Simulator: Q-function is Quadratic and Bellman Complete
In this section we demonstrate how our results for non-concave bandits also apply to reinforcement learning. Let Th be the Bellman operator applied to the Q-function Qh+1 defined as:
Th(Qh+1)(s, a) = rh(s, a) + Es′∼P(·|s,a)[max a′ Qh+1(s ′, a′)].
Definition 3.11 (Bellman complete). Given MDPM = (S,A,P, r,H), function classFh : S×A 7→ R, h ∈ [H] is called Bellman complete if for all h ∈ [H] and Qh+1 ∈ Fh+1, Th(Qh+1) ∈ Fh. Assumption 3.12 (Bellman complete for low-rank quadratic reward). We assume the function class Fh = {fM : fM (s, a) = ϕ(s, a)⊤Mϕ(s, a), rank(M) ≤ k, and 0 < λ1(M)/λmin(M) ≤ κ̃.} is a class of quadratic function and the MDP is Bellman complete. Here fM (ϕ(s, a)) =∑k
j=1 λj(v ⊤ j ϕ(s, a))
2 when M = ∑k
j=1 λjvjv ⊤ j . Write Q ∗ h = fM∗h ∈ Fh.
Observation: When querying sh−1, ah−1, we observe s′h ∼ P(·|sh−1, ah−1) and reward rh−1(sh−1, ah−1). Oracle to recover parameter M̂ : Given n ≥ Θ̃(d2k2κ̃2(M)/ε2), if one can play ≥ n samples ai and observe yi ∼ a⊤i Mai + ηi, i ∈ [n] with 1-sub-gaussian and mean-zero noise η, we can recover M̂ = M̂({(ai, yi)}) such that ∥M̂ −M∥2 ≤ ε. This oracle is implemented via our analysis from the bandit setting.
With the oracle, at time step H , we can estimate M̂H that is ϵ/H close to M∗H in spectral norm through noisy observations from the reward function with Õ(κ̃2d2H2/ϵ2) samples. Next, for each time step h = H − 1, H − 1, · · · , 1, sample s′i ∼ P(·|s, a), we define ηi = maxa′ fM̂h+1(s ′ i, a
′)− Es′∼P(·|s,a) maxa′ fM̂h+1(s
′, a′). ηi is mean-zero andO(1)-sub-gaussian since it is bounded. Denote Mh as the matrix that satisfies fMh := T fM̂h+1 , which is well-defined due to Bellman completeness. We estimate M̂h from the noisy observations yi = rh(s, a) + maxa′ fM̂h+1(s ′ i, a ′) = T f M̂h+1 + ηi =: fMh + ηi. Therefore with the oracle, we can estimate M̂h such that ∥M̂h −Mh∥2 ≤ ϵ/H with Θ(κ̃2d2k2H2/ϵ2) bandits. More details are deferred to Algorithm 5 and the Appendix. We state the theorem on sample complexity of finding ϵ-optimal policy here:
Theorem 3.13. Suppose F is Bellman complete associated with parameter κ̃. With probability 1− δ, Algorithm 5 learns an ϵ-optimal policy π with Θ̃(d2k2κ̃2H3/ϵ2) samples.
Existing approaches require O(d3H2/ϵ2) trajectories, or equivalently O(d3H3/ϵ2) samples, though they operate in the online RL setting [75, 23, 41], which is worse by a factor of d.
It is an open problem to attain d2 sample complexity in the online RL setting. The quadratic Bellman complete setting can also be easily extended to any of the polynomial settings of Section 3.3.

heading: 3.3 Stochastic High-order Homogeneous Polynomial Reward
text: Next we move on to homogeneous high-order polynomials.

heading: 3.3.1 The symmetric setting
text: Let reward function be a p-th order stochastic polynomial function f : A → R, where the action set A := {Bd1 = {a ∈ Rd, ∥a∥ ≤ 1.}. f(a) = T (a⊗p), rt = f(a) + ηt, where T = ∑k j=1 λjv ⊗p j is an orthogonally decomposable rank-k tensor. {v1, · · ·vk} form an orthonormal basis. Optimal reward r∗ satisfies 1 ≥ r∗ = λ1 ≥ |λ2| · · · ≥ |λk|. Noise ηt ∼ N (0, 1). In this setting, the problem is fundamentally more challenging than quadratic reward functions. On the one hand, it has a higher noise-to-signal ratio with larger p. One can tell from the rank-1 setting where T = λ1v ⊗p 1 . For a randomly generated action a on the unit ball, E[∥a⊤v1∥2] = 1/d. Therefore on average the signal strength is only (a⊤v1)p ∼ d−p/2, much smaller than the noise level 1. Intuitively this demonstrates why higher complexity is needed for high-order polynomials. On the
other hand, it is also technically more challenging. Unlike the matrix case, the expected zeroth-order update is no longer equal to any tensor product. Therefore existing tensor decomposition arguments do not apply. Fortunately, we prove that zeroth-order optimization still pushes the iterated actions toward the optimal action with linear convergence, given a good initialization. We show the bandit optimization procedure in Algorithm 6 and present the result in Theorem 3.14: Theorem 3.14 (Staged progress). For each stage s, with high probability As is not empty; and at least one action a ∈ As satisfies: tan θ(a,a∗) ≤ ε̃s = 2−s.
We defer the proofs together with formal statements to Appendix D.1.1 and D.1. Remark 3.15 (Choice of step-size). Here we choose step-size ζ = 1/2p. Note that E[y] = ζ(1− ζ)pσ2∇af(a). The scaling (1− ζ)p ≥ 1/ √ e ensures the signal to noise ratio not too small. The choice of weighted action is a delicate balance between making progress in optimization and controlling the noise to signal ratio.
Corollary 3.16 (Regret bound for tensors). Algorithm 6 yields an regret of: R(T ) ≤ Õ (√ kdpT )
, with high probability. Corollary 3.17 (Regret bound with burn-in period). In Algorithm 6, we first set ε = 1/p, ns = Θ̃(dp/λ21). We can first estimate the action a such that v ⊤ 1 a ≥ 1 − 1/p with Õ(kdp/λ21) samples. Next we change ε = Θ̃(k1/4d1/2λ−1/21 T −1/4), and ns = Θ̃(d2ε−2λ−21 ) in Algorithm 6. This procedure suffices to find λ1ε2-optimal reward with Õ(kd2/λ21ε 2) samples in the candidate set with size at most Õ(k). Finally with the UCB algorithm altogether we have a regret bound of: Õ(kd p
λ1 + √ kd2T ).
In Section 3.3.2 we also demonstrate the necessity of the sample complexity in this burn-in period.

heading: 3.3.2 Lower Bounds for Stochastic Polynomial Bandits
text: In this section we show lower bound for stochastic polynomial bandits.
f(a) = p∏ i=1 (θ⊤i a) + η, where η ∼ N(0, 1), ∥a∥2 ≤ 1 and f(a) ≤ 1. (1)
Theorem 3.18. Define minimax regret as follow
R(d, p, T ) = inf π sup (θ1,...,θp) E(θ1,...,θp)
[ T max
a p∏ i=1 (θ⊤i a)− T∑ t=1 p∏ i=1 (θ⊤i a (t))
] .
For all algorithmsA that adaptively interact with bandit (Eq (1)) for T rounds, we have R(d, p, T ) ≥ Ω( √ dpT/pp).
From the theorem, we can see even when the problem is rank-1, any algorithm incurs at least Ω( √ dpT/pp) regret. This further implies algorithm requires sample complexity of Ω((d/p2)p/ϵ2) to attain ϵ-optimal reward. This means our regret upper bound obtained in Corollary 3.16 is optimal in terms of dependence on d. In Appendix G we also show a Ω( √ dpT ) lower bound for asymmetric actions setting, which our upper bound up to poly-logarithmic factors. We note that with burn-in period, our algorithm also obtains a cumulative regret of kdp/λ1+ √ kd2T , as shown in Corollary 3.17. Here for a fixed λ1 and very large T , this is better result than the previous upper bound. We note that there is no contradiction with the lower bound above, since this worst case is achieved with a specific relation between T and λ1.
The burn-in period requires kdp/λ21 samples to get a constant of r ∗. We want to investigate whether our dependence on λ1 ≡ r∗ is optimal. Next we show a gap-dependent lower bound for finding an arm that is close to the optimal arm by a constant factor. Theorem 3.19. For all algorithms A that adaptively interact with bandit (Eq (1)) for T rounds and output a vector a(T ) ∈ Rd, it requires at least T = Ω(dp/∥θ∥2p) rounds to find an arm a(T ) ∈ Rd such that ∏p i=1(θ ⊤ i a (T )) ≥ 34 ·maxa ∏p i=1(θ ⊤ i a).
In the rank-1 setting r∗ = ∥θ∥p and the lower bound for achieving a constant approximation for the optimal reward is Ω(dp/(r∗)2). Therefore our burn-in sample complexity is also optimal in the dependence on d and r∗.

heading: 3.4 Noiseless Polynomial Reward
text: In this subsection, we study the regret bounds for learning bandits with noiseless polynomial rewards. First we present the definition of admissible polynomial families. Definition 3.20 (Admissible Polynomial Family). For a ∈ Rd, define ã = [1,a⊤]⊤. For a algebraic variety V ⊆ (Rd+1)⊗p, define RV := {rθ(a) = ⟨θ, ã⊗p⟩ : θ ∈ V} as the polynomial family with parameters in V . We define the dimension of the familyRV as the algebraic dimension of V . Next, define X := { ã⊗p : a ∈ Rd } . An polynomial familyRV is said to be admissible5 w.r.t. X if for any θ ∈ V , dim(X ∩ {X ∈ X : ⟨X,θ⟩ = 0⟩}) < dim(X ) = d.

heading: 3.4.1 Upper Bounds via Solving Polynomial Equations
text: We show that if the action set A is of positive measure with respect to the Lebesgue measure µ, then by playing actions randomly, we can uniquely solve for the ground-truth reward function rθ(a) almost surely with samples of size that scales with the intrinsic algebraic dimension of V , provided that V is an admissible algebraic variety. Theorem 3.21. Assume that the reward function class is an admissible polynomial familyRV , and the maximum reward is upper bounded by 1. If µ(A) > 0, where µ is the Lebesgue meaure, then by randomly sample actions a1, . . . ,aT from Pa∼N (0,Id)(·|a ∈ A), when T ≥ 2dim(V), we can uniquely solve for the ground-truth θ and thus determine the optimal action almost surely. Therefore, the cumulative regret at round T can be bounded as
R(T ) ≤ min{T, 2dim(V)}.
We state two important examples of admissible polynomial families with O(d) dimensions. Example 3.22 (low-rank polynomials). The function classRV of possibly inhomogeneous degree-p polynomials with k summandsRV = {r(a) = ∑k i=1 λi⟨vi,a⟩pi | λi ∈ R,vi ∈ Rd} is admissible with dim(RV) ≤ dk, where p = max{pi}. Neural network with monomial/polynomial activation functions are low-rank polynomials. Example 3.23 ([18]). The function class RV = {r(a) = q(Ua) | U ∈ Rk×d,deg q(·) ≤ p} is admissible with dim(V) ≤ dk + (k + 1)p.

heading: 4 Conclusion
text: In this paper, we design minimax-optimal algorithms for a broad class of bandit problems with non-concave rewards. For the stochastic setting, our algorithms and analysis cover the low-rank linear reward setting, bandit eigenvector problem, and homogeneous polynomial reward functions. We improve the best-known regret from prior work and attain the optimal dependence on problem dimension d. Our techniques naturally extend to RL in the generative model settings. Furthermore, we obtain the optimal regret, dependent on the intrinsic algebraic dimension, for general polynomial reward without noise. Our regret bound demonstrates the fundamental limits of UCB algorithms, being Ω(dp−1) worse than our result for cases of interest.
We leave to future work several directions. First, the gap-free algorithms for the low-rank linear and bandit eigenvector problem do not attain √ T regret. We believe √ d2T regret without any dependence on gap is impossible, but do not have a lower bound. Secondly, our study of degree p polynomials only covers the noiseless setting or orthogonal tensors with noise. We believe entirely new algorithms are needed for general polynomial bandits to surpass eluder dimension. As a first step, designing optimal algorithms would require understanding the stability of algebraic varieties under noise, so we leave this as a difficult future problem. Finally, we conjecture our techniques can be used to design optimal algorithms for representation learning in bandits and MDPs [74, 39].

heading: Acknowledgment
text: JDL acknowledges support of the ARO under MURI Award W911NF-11-1-0303, the Sloan Research Fellowship, NSF CCF 2002272, and an ONR Young Investigator Award. QL is supported by NSF
5Intuitively, admissibility means the dimension of X decreases by one when there is an additional linear constraint ⟨θ, X⟩ = 0
2030859 and the Computing Research Association for the CIFellows Project. SK acknowledges funding from the NSF Award CCF-1703574 and the ONR award N00014-18-1-2247. The authors would like to thank Qian Yu for numerous conversations regarding the lower bound in Theorem 3.23. JDL would like to thank Yuxin Chen and Anru Zhang for several conversations on tensor power iteration, Simon S. Du and Yangyi Lu for explaining to him to the papers of [43, 55], and Max Simchowitz and Chao Gao for several conversations regarding adaptive lower bounds.

Review summaries:
The paper studies nonlinear bandit optimization with structured polynomial reward functions, with a focus on the pure exploration setting. The paper proposes some stochastic zeroth-order gradient ascent algorithms with phase elimination and show that they attain minimax sample complexity (in terms of $1/\epsilon$ and the dimension $d$) for a class of structured polynomial reward functions (i.e., stochastic low-rank linear reward function, stochastic bilinear reward function, stochastic bilinear reward reward function). The sample complexity bounds usually improve upon the bounds obtained via existing techniques (i.e., LinUCB and eluder dimension-based arguments) by a $O(d)$ factor. Combined with matching lower bounds, these results help to reveal the right role of $d$ in the minimax complexity of the considered problems. The paper also shows the suboptimality of UCB algorithms (in terms of $d$) by proving a $\Omega(d^p)$-type lower bound for UCB in the noise-less polynomial reward case (for which solving polynomial equations usually leads to $O(d)$-type sample complexity).

This paper studied a fundamental question in non-linear bandit: suppose the mean reward function f_\theta in the model r_t = f_\theta(a_t) + noise is non-linear or only linear after some high-dimensional reparametrization of \theta, is there a good bandit algorithm to find near-optimal policies with sample complexities dictated by the intrinsic dimension? In this paper, the authors investigated the following examples:   1. low-rank linear reward;  2. bilinear reward (or RL with quadratic Bellman-complete Q-function);  3. homogeneous polynomial reward;  4. noiseless polynomial reward.   In first three examples, the zeroth-order gradient ascent algorithm achieves a PAC sample complexity with the optimal dependence on the dimension d, strictly better than the upper bound given by LinUCB when viewed as a generalized linear model, as well as the eluder dimension argument. In the last example, there is an instance where solving polynomial equations after random sampling only has sample complexity O(d), but any UCB algorithm incurs a sample complexity O(d^p).

This paper studies the zeroth-order optimization methods for the nonlinear bandit problem. It studies PAC bounds for bandit problems where the reward function is structured polynomials, including a low-rank generalized linear bandit problem and a two-layer neural network with polynomial activations bandit problem. In the low-rank generalized linear bandit setting, the authors provide a minimax optimal algorithm improving the results in [28, 35]. They further apply their algorithms in the RL with the generative model setting, achieving improved sample complexity. In the neural net setting, they propose a bandit algorithm that has sample complexity equal to the intrinsic algebraic dimension.  They also show that the standard optimistic algorithms are sub-optimal for both settings.