summaries: The paper presents a new neural-bandit algorithm with shallow exploration and provides a regret bound for the proposed method. The existing approaches have introduced deep neural networks based bandit algorithms to learn reward functions, in which exploration takes place over the entire network parameter space, which can be inefficient for large-size networks which are typical in NTK based approaches. The authors address this by taking an existing approach that decouples the deep neural network feature representation learning from most of the exploration of the network parameters by only exploring over the final layer of the network.   Despite the fact that this idea of shallow exploration has been proposed previously, there has not been a theoretical analysis with a regret bound. The authors analyze a UCB version of this approach, then build from techniques from both deep neural contextual bandits and linear contextual bandits to prove an O(\sqrt(T)) regret bound. Finally, the authors present experimental results to show that their algorithm work well in practice. This paper studies neural contextual bandits and proposes an algorithm that transforms the raw feature vector using the last hidden layer of a deep ReLU neural network, and uses an UCB approach to explore in the last linear layer. Compared with existing neural contextual bandit algorithms, the proposed algorithm attains computation efficiency. Regret guarantees and empirical results are provided to demonstrate the effectiveness of the proposed algorithms This paper study a novel contextual bandit algorithm: Neural-LinUCB. As in (Riquelme et al 2019), the idea of this algorithm is based on decoupling deep representation learning and exploration. A deep neural network learns the mapping between the context $x_{t,a_t}$, while a linear bandit, OFUL (Abbasi-Yadkori 2011), chooses the arm to play. In contrast to (Riquelme et al 2019), a regret upper bound of the algorithm a regret upper bound of the algorithm is stated in Corollary 4.6. The proposed algorithm is also an improvement over NeuralUCB (Zhou et al 2020) for two reasons: the computational cost of the exploration is lesser, since the exploration is only done in the last layer of weights, and the regret upper bound is tighter. Indeed, in contrast to (Zhou et al 2020) it does not depend on the dimension of the tangent kernel matrix, which can be in O(KT). Experiments, done on four contextual bandit problems, show that Neura-lLinUCB outperforms LinUCB and performs as well as NeuralUCB and NeuralTS.  Authors tackle the setting of contextual bandits, using deep representation learning combined with an upper confidence bound algorithm. The main contribution of this work is to provide a regret bound for the setup which decouples the representation learning from the UCB search, by searching only over the last layer of the network. The setting had been studied before, but only empirically, and using Thompson sampling rather than UCB. Authors validate their results empirically on several domains from the UCI data repo, as well as on MNIST, comparing against state of the art baselines. 
abstractText: We study neural contextual bandits, a general class of contextual bandits, where each context-action pair is associated with a raw feature vector, but the specific reward generating function is unknown. We propose a novel learning algorithm that transforms the raw feature vector using the last hidden layer of a deep ReLU neural network (deep representation learning), and uses an upper confidence bound (UCB) approach to explore in the last linear layer (shallow exploration). We prove that under standard assumptions, our proposed algorithm achieves Õ( √ T ) finitetime regret, where T is the learning time horizon. Compared with existing neural contextual bandit algorithms, our approach is computationally much more efficient since it only needs to explore in the last layer of the deep neural network.
authors.name: DEEP REPRE
authors.name: SHALLOW EXPLORATION
authors.name: Pan Xu
authors.name: Zheng Wen
authors.name: Handong Zhao
references.authors: Yasin Abbasi-Yadkori Dávid Pál Csaba Szepesvári
references.title: Improved algorithms for linear stochastic bandits
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Alekh Agarwal Daniel Hsu Satyen Kale John Langford Lihong Li Robert Schapire
references.title: Taming the monster: A fast and simple algorithm for contextual bandits
references.venue: In International Conference on Machine Learning,
references.authors: Zeyuan Allen-Zhu Yuanzhi Li Yingyu Liang
references.title: Learning and generalization in overparameterized neural networks, going beyond two layers
references.venue: In Advances in neural information processing systems,
references.authors: Zeyuan Allen-Zhu Yuanzhi Li Zhao Song
references.title: A convergence theory for deep learning via overparameterization
references.venue: In International Conference on Machine Learning,
references.authors: Robin Allesiardo Raphaël Féraud Djallel Bouneffouf
references.title: A neural networks committee for the contextual bandit problem
references.venue: In International Conference on Neural Information Processing,
references.authors: Sanjeev Arora Simon Du Wei Hu Zhiyuan Li Ruosong Wang
references.title: Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks
references.venue: In International Conference on Machine Learning,
references.authors: Sanjeev Arora Simon S Du Wei Hu Zhiyuan Li Russ R Salakhutdinov Ruosong Wang
references.title: On exact computation with an infinitely wide neural net
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Jean-Yves Audibert Rémi Munos Csaba Szepesvári
references.title: Exploration–exploitation tradeoff using variance estimates in multi-armed bandits
references.venue: Theoretical Computer Science,
references.authors: Peter Auer Nicolo Cesa-Bianchi Paul Fischer
references.title: Finite-time analysis of the multiarmed bandit problem
references.venue: Machine learning,
references.authors: Sivaraman Balakrishnan Martin J Wainwright Bin Yu
references.title: Statistical guarantees for the em algorithm: From population to sample-based analysis
references.venue: The Annals of Statistics,
references.authors: Qi Cai Zhuoran Yang Jason D Lee Zhaoran Wang
references.title: Neural temporal-difference learning converges to global optima
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Yuan Cao Quanquan Gu
references.title: A generalization theory of gradient descent for learning overparameterized deep relu networks
references.venue: arXiv preprint arXiv:1902.01384,
references.authors: Yuan Cao Quanquan Gu
references.title: Generalization bounds of stochastic gradient descent for wide and deep neural networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Olivier Chapelle Lihong Li
references.title: An empirical evaluation of thompson sampling
references.venue: In Advances in neural information processing systems,
references.authors: Sayak Ray Chowdhury Aditya Gopalan
references.title: On kernelized multi-armed bandits
references.venue: In International Conference on Machine Learning,
references.authors: Wei Chu Lihong Li Lev Reyzin Robert Schapire
references.title: Contextual bandits with linear payoff functions
references.venue: In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics,
references.authors: Mark Collier Hector Urdiales Llorens
references.title: Deep contextual multi-armed bandits
references.venue: arXiv preprint arXiv:1807.09809,
references.authors: George Cybenko
references.title: Approximation by superpositions of a sigmoidal function
references.venue: Mathematics of control, signals and systems,
references.authors: Varsha Dani Thomas P Hayes Sham M Kakade
references.title: Stochastic linear optimization under bandit feedback
references.venue: In Conference on Learning Theory,
references.authors: Aniket Anand Deshmukh Abhimanu Kumar Levi Boyles Denis Charles Eren Manavoglu Urun Dogan
references.title: Self-supervised contextual bandits in computer vision
references.venue: arXiv preprint arXiv:2003.08485,
references.authors: Simon Du Jason Lee Haochuan Li Liwei Wang Xiyu Zhai
references.title: Gradient descent finds global minima of deep neural networks
references.venue: In International Conference on Machine Learning,
references.authors: Simon S. Du Xiyu Zhai Barnabas Poczos Aarti Singh
references.title: Gradient descent provably optimizes over-parameterized neural networks
references.venue: In International Conference on Learning Representations,
references.authors: Dheeru Dua Casey Graff
references.title: UCI machine learning repository, 2017
references.venue: URL http://archive. ics.uci.edu/ml
references.authors: Sarah Filippi Olivier Cappe Aurélien Garivier Csaba Szepesvári
references.title: Parametric bandits: The generalized linear case
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Geoffrey E Hinton Ruslan R Salakhutdinov
references.title: Reducing the dimensionality of data with neural networks
references.authors: Arthur Jacot Franck Gabriel Clément Hongler
references.title: Neural tangent kernel: Convergence and generalization in neural networks
references.venue: In Advances in neural information processing systems,
references.authors: Branislav Kveton Manzil Zaheer Csaba Szepesvari Lihong Li Mohammad Ghavamzadeh Craig Boutilier
references.title: Randomized exploration in generalized linear bandits
references.venue: In International Conference on Artificial Intelligence and Statistics,
references.authors: John Langford Tong Zhang
references.title: The epoch-greedy algorithm for multi-armed bandits with side information
references.venue: In Advances in neural information processing systems,
references.authors: Tor Lattimore Csaba Szepesvári
references.title: Bandit algorithms
references.authors: Yann LeCun Léon Bottou Yoshua Bengio Patrick Haffner
references.title: Gradient-based learning applied to document recognition
references.venue: Proceedings of the IEEE,
references.authors: Jaehoon Lee Lechao Xiao Samuel Schoenholz Yasaman Bahri Roman Novak Jascha SohlDickstein Jeffrey Pennington
references.title: Wide neural networks of any depth evolve as linear models under gradient descent
references.venue: In Advances in neural information processing systems,
references.authors: Lihong Li Wei Chu John Langford Robert E Schapire
references.title: A contextual-bandit approach to personalized news article recommendation
references.venue: In Proceedings of the 19th international conference on World wide web,
references.authors: Lihong Li Yu Lu Dengyong Zhou
references.title: Provably optimal algorithms for generalized linear contextual bandits
references.venue: In International Conference on Machine Learning,
references.authors: Boyi Liu Qi Cai Zhuoran Yang Zhaoran Wang
references.title: Neural trust region/proximal policy optimization attains globally optimal policy
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Samet Oymak Mahdi Soltanolkotabi
references.title: Towards moderate overparameterization: global convergence guarantees for training shallow neural networks
references.venue: IEEE Journal on Selected Areas in Information Theory,
references.authors: Carlos Riquelme George Tucker Jasper Snoek
references.title: Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling
references.venue: In International Conference on Learning Representations,
references.authors: Paat Rusmevichientong John N Tsitsiklis
references.title: Linearly parameterized bandits
references.venue: Mathematics of Operations Research,
references.authors: Jyotirmoy Sarkar
references.title: One-armed bandit problems with covariates
references.venue: The Annals of Statistics,
references.authors: Jasper Snoek Oren Rippel Kevin Swersky Ryan Kiros Nadathur Satish Narayanan Sundaram Mostofa Patwary Mr Prabhat Ryan Adams
references.title: Scalable bayesian optimization using deep neural networks
references.venue: In International conference on machine learning,
references.authors: William R Thompson
references.title: On the likelihood that one unknown probability exceeds another in view of the evidence of two samples
references.authors: Michal Valko Nathan Korda Rémi Munos Ilias Flaounas Nello Cristianini
references.title: Finite-time analysis of kernelised contextual bandits
references.venue: In Proceedings of the Twenty-Ninth Conference on Uncertainty in Artificial Intelligence,
references.authors: Lingxiao Wang Qi Cai Zhuoran Yang Zhaoran Wang
references.title: Neural policy gradient methods: Global optimality and rates of convergence
references.venue: In International Conference on Learning Representations,
references.authors: Zhaoran Wang Han Liu Tong Zhang
references.title: Optimal computational and statistical rates of convergence for sparse nonconvex learning problems
references.venue: Annals of statistics,
references.authors: Pan Xu Quanquan Gu
references.title: A finite-time analysis of q-learning with neural network function approximation
references.venue: In International Conference on Machine Learning,
references.authors: Pan Xu Jian Ma Quanquan Gu
references.title: Speeding up latent variable gaussian graphical model estimation via nonconvex optimization
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Tom Zahavy Shie Mannor
references.title: Deep neural linear bandits: Overcoming catastrophic forgetting through likelihood matching
references.venue: arXiv preprint arXiv:1901.08612,
references.authors: Chiyuan Zhang Samy Bengio Moritz Hardt Benjamin Recht Oriol Vinyals
references.title: Understanding deep learning requires rethinking generalization
references.venue: In International Conference on Learning Representations,
references.authors: Dongruo Zhou Lihong Li Quanquan Gu
references.title: Neural contextual bandits with ucb-based exploration
references.venue: In International Conference on Machine Learning,
references.authors: Difan Zou Quanquan Gu
references.title: An improved analysis of training over-parameterized deep neural networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Difan Zou Yuan Cao Dongruo Zhou Quanquan Gu
references.title: Stochastic gradient descent optimizes over-parameterized deep relu networks
references.venue: arXiv preprint arXiv:1811.08888,
references.authors: Zhang
references.title: 4,m = 10000, and thus the hidden layer width is (10000
references.authors: Abbasi-Yadkori
references.title: The next lemma provides some standard bounds on the feature matrix At, which is a combination of Lemma 10 and Lemma
references.authors: Abbasi-Yadkori
references.title: Published as a conference paper at ICLR 2022 D.3 PROOF OF LEMMA C.4 The following lemma characterizes the concentration property of self-normalized martingales
sections.text: √ T ) finite-
  time regret, where T is the learning time horizon. Compared with existing neural contextual bandit algorithms, our approach is computationally much more efficient since it only needs to explore in the last layer of the deep neural network.
sections.heading: 1 INTRODUCTION
sections.text: Multi-armed bandits (MAB) (Auer et al., 2002; Audibert et al., 2009; Lattimore & Szepesvári, 2020) are a class of online decision-making problems where an agent needs to learn to maximize its expected cumulative reward by repeatedly interacting with a partially known environment. Following a bandit algorithm (also called a strategy or policy), in each round, the agent adaptively chooses an arm, and then receives a reward associated with that arm. Since only the reward of the chosen arm will be observed (bandit information feedback), a good bandit algorithm has to deal with the exploration-exploitation dilemma: trade-off between pulling the best arm based on existing knowledge/history data (exploitation) and trying the arms that have not been fully explored (exploration).
  In many real-world applications, the agent will also be able to access detailed contexts associated with the arms. For example, when a company wants to choose an advertisement to present to a user, the recommendation will be much more accurate if the company takes into consideration the contents, specifications, and other features of the advertisements in the arm set as well as the profile of the user. To encode the contextual information, contextual bandit models and algorithms have been developed, and widely studied both in theory and in practice (Dani et al., 2008; Rusmevichientong & Tsitsiklis, 2010; Li et al., 2010; Chu et al., 2011; Abbasi-Yadkori et al., 2011). Most existing contextual bandit algorithms assume that the expected reward of an arm at a context is a linear function in a known context-action feature vector, which leads to many useful algorithms such as LinUCB (Chu et al., 2011), OFUL (Abbasi-Yadkori et al., 2011), etc. The representation power of the linear model can be limited in applications such as marketing, social networking, clinical studies, etc., where the rewards are usually counts or binary variables. The linear contextual bandit problem has also been extended to richer classes of parametric bandits such as the generalized linear bandits (Filippi et al., 2010; Li et al., 2017) and kernelised bandits (Valko et al., 2013; Chowdhury & Gopalan, 2017).
  With the prevalence of deep neural networks (DNNs) and their phenomenal performances in many machine learning tasks (LeCun et al., 2015; Goodfellow et al., 2016), there has emerged a line of work that employs DNNs to increase the representation power of contextual bandit algorithms (Allesiardo et al., 2014; Riquelme et al., 2018; Collier & Llorens, 2018; Zahavy & Mannor, 2019;
  Zhou et al., 2020; Deshmukh et al., 2020; Zhang et al., 2020). The problems they solve are usually referred to as neural contextual bandits. For example, Zhou et al. (2020) developed the NeuralUCB algorithm, which can be viewed as a natural extension of LinUCB (Chu et al., 2011; Abbasi-Yadkori et al., 2011), where they use the output of a deep neural network with the feature vector as input to approximate the reward. Zhang et al. (2020) adapted neural networks in Thompson Sampling (Thompson, 1933; Chapelle & Li, 2011; Russo et al., 2018) for both exploration and exploitation and proposed NeuralTS . For a fixed time horizon T , it has been proved that both NeuralUCB and NeuralTS achieve a O(d̃ √ T ) regret bound, where d̃ is the effective dimension of a neural tangent kernel matrix which can potentially scale with O(TK) for K-armed bandits. This high complexity is mainly due to that the exploration is performed over the entire huge neural network parameter space, which is inefficient and even infeasible when the number of neurons is large. A more realistic and efficient way of learning neural contextual bandits may be to just explore different arms using the last layer as the exploration parameter. More specifically, Riquelme et al. (2018) provided an extensive empirical study of benchmark algorithms for contextual-bandits through the lens of Thompson Sampling, which suggests decoupling representation learning and uncertainty estimation improves performance.
  In this paper, we show that the decoupling of representation learning and the exploration can be theoretically validated. We study a new neural contextual bandit algorithm, which learns a mapping to transform the raw features associated with each context-action pair using a deep neural network (deep representation), and then performs an upper confidence bound (UCB)-type exploration over the linear output layer of the network (shallow exploration). We prove a sublinear regret of the proposed algorithm by exploiting the UCB exploration techniques in linear contextual bandits (Abbasi-Yadkori et al., 2011) and the analysis of deep overparameterized neural networks using neural tangent kernels (Jacot et al., 2018). Our theory confirms the empirically observed effectiveness of decoupling the deep representation learning and the UCB exploration in contextual bandits (Riquelme et al., 2018; Zahavy & Mannor, 2019).
  Contributions we summarize the main contributions of this paper as follows.
  • We propose a contextual bandit algorithm, Neural-LinUCB, for solving a general class of contextual bandit problems without knowing the specific reward generating function. The proposed algorithm learns a deep representation to transform the raw feature vectors and performs UCBtype exploration in the last layer of the neural network, which we refer to as deep representation and shallow exploration. Compared with LinUCB (Li et al., 2010; Chu et al., 2011) and neural bandits such as NeuralUCB (Zhou et al., 2020) and NeuralTS (Zhang et al., 2020), our algorithm enjoys the best of two worlds: strong expressiveness due to the deep representation and computational efficiency due to the shallow exploration. • Despite the usage of a DNN as the feature mapping, we prove a Õ( √ T ) regret for the proposed
  Neural-LinUCB algorithm, which matches the regret bound of linear contextual bandits (Chu et al., 2011; Abbasi-Yadkori et al., 2011). To the best of our knowledge, this is the first work that theoretically shows the convergence of bandits algorithms under the scheme of deep representation and shallow exploration. It is notable that a similar scheme called Neural-Linear was proposed by Riquelme et al. (2018) for Thompson sampling algorithms, and they empirically showed that decoupling representation learning and uncertainty estimation improves the performance. Our work confirms this observation from a theoretical perspective. • We conduct experiments on contextual bandit problems based on real-world datasets, demonstrating a better performance and computational efficiency of Neural-LinUCB over LinUCB and existing neural bandits algorithms such as NeuralUCB, which well aligns with our theory.
sections.heading: 1.1 ADDITIONAL RELATED WORK
sections.text: There is a line of related work to ours on the recent advance in the optimization and generalization analysis of deep neural networks. In particular, Jacot et al. (2018) first introduced the neural tangent kernel (NTK) to characterize the training dynamics of network outputs in the infinite width limit. From the notion of NTK, a fruitful line of research emerged and showed that loss functions of deep neural networks trained by (stochastic) gradient descent can converge to the global minimum (Du et al., 2019b; Allen-Zhu et al., 2019b; Du et al., 2019a; Zou et al., 2018; Zou & Gu, 2019). The generalization bounds for overparameterized deep neural networks are also established in Arora et al. (2019a;b); Allen-Zhu et al. (2019a); Cao & Gu (2019a;b). Recently, the NTK based analysis
  is also extended to the study of sequential decision problems including bandits (Zhou et al., 2020; Zhang et al., 2020), and reinforcement learning algorithms (Cai et al., 2019; Liu et al., 2019; Wang et al., 2020; Xu & Gu, 2020).
  Our algorithm is also different from Langford & Zhang (2008); Agarwal et al. (2014) which reduce the bandit problem to supervised learning. Moreover, their algorithms need to access an oracle that returns the optimal policy in a policy class given a sequence of context and reward vectors, whose regret depends on the VC-dimension of the policy class. Notation We use [k] to denote a set {1, . . . , k}, k ∈ N+. ‖x‖2 = √
  x>x is the Euclidean norm of a vector x ∈ Rd. For a matrix W ∈ Rm×n, we denote by ‖W‖2 and ‖W‖F its operator norm and Frobenius norm respectively. For a semi-definite matrix A ∈ Rd×d and a vector x ∈ Rd, we denote the Mahalanobis norm as ‖x‖A = √ x>Ax. Throughout this paper, we reserve the notations {Ci}i=0,1,... to represent absolute positive constants that are independent of problem parameters such as dimension, sample size, iteration number, step size, network length and so on. The specific values of {Ci}i=0,1,... can be different in different context. For a parameter of interest T and a function f(T ), we use notations such asO(f(T )) and Ω(f(T )) to hide constant factors and Õ(f(T )) to hide constant and logarithmic dependence of T .
sections.heading: 2 PRELIMINARIES
sections.text: In this section, we provide the background of contextual bandits and deep neural networks.
sections.heading: 2.1 LINEAR CONTEXTUAL BANDITS
sections.text: A contextual bandit is characterized by a tuple (S,A, r), where S is the context (state) space, A is the arm (action) space, and r encodes the unknown reward generating function at all context-arm pairs. A learning agent, who knows S and A but does not know the true reward r (values bounded in (0, 1) for simplicity), needs to interact with the contextual bandit for T rounds. At each round t = 1, . . . , T , the agent first observes a context st ∈ S chosen by the environment; then it needs to adaptively select an arm at ∈ A based on its past observations; finally it receives a reward
  r̂t(xs,at) = r(xs,at) + ξt, (2.1)
  where xs,a ∈ Rd is a known feature vector for context-arm pair (s, a) ∈ S ×A, and ξt is a random noise with zero mean. The agent’s objective is to maximize its expected total reward over these T rounds, which is equivalent to minimizing the pseudo regret (Audibert et al., 2009):
  RT = E [ T∑ t=1 ( r̂(xst,a∗t )− r̂(xst,at) )] , (2.2)
  where a∗t ∈ argmaxa∈A{r(xst,a) = E[r̂(xst,a)]}. To simplify the exposition, we use xt,a to denote xst,a since it only depends on the round index t in most bandit problems, and we assume A = [K]. In linear contextual bandits, the reward function in (2.1) is assumed to have a linear structure r(xs,a) = x > s,aθ
  ∗ for some unknown weight vector θ∗ ∈ Rd. One provably sample efficient algorithm for linear contextual bandits is Linear Upper Confidence Bound (LinUCB) (Chu et al., 2011) or Optimism in the Face of Uncertainty Linear bandit algorithm (OFUL) (Abbasi-Yadkori et al., 2011). Specifically, at each round t, LinUCB chooses the action at = argmaxa∈[K]{x>t,aθt + αt‖xt,a‖A−1t }, where θt is a point estimate of θ ∗, At = λI + ∑t i=1 xi,aix > i,ai
  with some λ > 0 is a matrix defined based on the historical context-arm pairs, and αt > 0 is a tuning parameter that controls the exploration rate in LinUCB.
sections.heading: 2.2 DEEP NEURAL NETWORKS
sections.text: In this paper, we use f(x) to denote a neural network with input data x ∈ Rd. Let L be the number of hidden layers and Wl ∈ Rml×ml−1 be the weight matrices in the l-th layer, where l = 1, . . . , L, m1 = . . . = mL−1 = m and m0 = mL = d. Then a L-hidden layer neural network is defined as
  f(x) = √ mθ∗>σL(WLσL−1(WL−1 · · ·σ1(W1x) · · · )), (2.3)
  where σl is an activation function and θ∗ ∈ Rd is the weight of the output layer. To simplify the presentation, we will assume σ1 = σ2 = . . . = σL = σ is the ReLU activation function, i.e., σ(x) = max{0, x} for x ∈ R. We denote w = (vec(W1)>, . . . , vec(WL)>)>, which is the concatenation of the vectorized weight parameters of all hidden layers of the neural network. We also write f(x;θ∗,w) = f(x) in order to explicitly specify the weight parameters of neural network f . It is easy to show that the dimension p of vector w satisfies p = (L− 2)m2 + 2md. To simplify the notation, we define φ(x; w) as the output of the L-th hidden layer of neural network f .
  φ(x; w) = √ mσ(WLσ(WL−1 · · ·σ(W1x) · · · )). (2.4)
  Note that φ(x; w) itself can also be viewed as a neural network with vector-valued outputs.
sections.heading: 3 DEEP REPRESENTATION AND SHALLOW EXPLORATION
sections.heading: 3.1 HIGH-LEVEL IDEA OF THE PROPOSED ALGORITHM
sections.text: The linear parametric form in linear contextual bandits might produce biased estimates of the reward due to the lack of representation power (Snoek et al., 2015; Riquelme et al., 2018). In contrast, it is well known that deep neural networks are powerful enough to approximate an arbitrary function (Cybenko, 1989). Therefore, a natural extension of linear contextual bandits is to use a deep neural network to approximate the reward generating function r(·). Nonetheless, DNNs usually have a prohibitively large dimension for weight parameters, which makes the exploration in neural networks based UCB algorithm inefficient (Kveton et al., 2020; Zhou et al., 2020).
  In this work, we study a neural contextual bandit algorithm, where the hidden layers of a deep neural network are used to represent the features and the exploration is only performed in the last layer of the neural network. In particular, for any arm feature vector x, we use 〈θ∗,φ(x; w)〉 to approximate the unknown reward function r(x), where φ(x; w) defined as in (2.4) is a neural network with weight w, and θ∗ is a unknown weight parameter. Note that we can also view 〈θ∗,φ(x; w)〉 as a neural network with φ(x; w) being the output of the last hidden layer and θ∗ the weight parameter of the last (linear) layer. Different from existing neural bandit algorithms, we only add a UCB bonus term involving the last layer instead of all the weight parameter of this large neural network.
  This decoupling of the representation and the exploration achieves the best of both worlds: efficient exploration of shallow (linear) models and high expressive power of deep models. In what follows, we will describe a neural contextual bandit algorithm that uses the output of the last hidden layer of a neural network to transform the raw feature vectors (deep representation) and performs UCB-type exploration in the last layer of the neural network (shallow exploration). Since the exploration is performed only in the last linear layer, we call this procedure Neural-LinUCB, which is displayed in Algorithm 1.
sections.heading: 3.2 DETAILED IMPLEMENTATION OF DEEP REPRESENTATION AND SHALLOW EXPLORATION
sections.text: Now we describe the details of Algorithm 1. In round t, the agent receives an action set with raw features Xt = {xt,1, . . . ,xt,K}. Then the agent chooses an arm at that maximizes the following upper confidence bound:
  at = argmax k∈[K]
  { 〈φ(xt,k; wt−1),θt−1〉+ αt‖φ(xt,k; wt−1)‖A−1t−1 } , (3.1)
  where θt−1 is a point estimate of the unknown weight in the last layer, φ(x; w) is defined as in (2.4), wt−1 is an estimate of all the weight parameters in the hidden layers of the neural network, αt > 0 is the algorithmic parameter controlling the exploration, and At defined as follows.
  At = λI + t∑ i=1 φ(xi,ai ; wi−1)φ(xi,ai ; wi−1) >, (3.2)
  and λ > 0. After pulling arm at, the agent will observe a noisy reward r̂t := r̂(xt,at) = r(xt,k)+ξt, where ξt is an independent ν-subGaussian random noise for some ν > 0 and r(·) is an unknown reward function. In this paper, we will interchangeably use notation r̂t to denote the reward received at the t-th step and an equivalent notation r̂(x) to express its dependence on the feature vector x.
  Upon receiving the reward r̂t, the agent updates its estimate θt of the output layer weight by using the same `2-regularized least-squares estimate in linear contextual bandits (Abbasi-Yadkori et al., 2011). In particular, we have θt = A−1t bt, where bt = ∑t i=1 r̂iφ(xi,ai ; wi−1).
  To save the computation, the neural networkφ(·; wt) will be updated once everyH steps. Therefore, we have w(q−1)H+1 = . . . = wqH for q = 1, 2, . . .. We call the time steps {(q−1)H+1, . . . , qH} an epoch with length H . At time step t = Hq, we will retrain the neural network based on all the historical data via Algorithm 2, which minimizes the following empirical loss function:
  Lq(w) = qH∑ i=1 ( θ>i φ(xi,ai ; w)− r̂i )2 . (3.3)
  In practice, one can further save computational cost by only feeding data {xi,ai , r̂i,θi} qH i=(q−1)H+1 from the q-th epoch into Algorithm 2 to update the parameter wt, which does not hurt the performance since the historical information has been encoded into the estimate of θi. In this paper, we will perform the following gradient descent step w(s)q = w (s−1) q − ηq∇wLq(w(s−1)), for s = 1, . . . , n, where w(0)q = w(0) is chosen as the same random initialization point. We will discuss more about the initial point w(0) in the next paragraph. Then Algorithm 2 outputs w(n)q and we set it as the updated weight parameter wHq+1 in Algorithm 1. In the next round, the agent will receive another action set Xt+1 with raw feature vectors and repeat the above steps to choose the sub-optimal arm and update estimation for contextual parameters.
  Initialization: Recall that w is the collection of all hidden layer weight parameters of the neural network. We will follow the same initialization scheme as used in Zhou et al. (2020), where each entry of the weight matrices follows some Gaussian distribution. Specifically, for any l ∈ {1, . . . , L− 1},
  we set Wl = [ W 0 0 W ] , where each entry of W follows distributionN(0, 4/m) independently; for
  WL, we set it as [V −V], where each entry of V follows distribution N(0, 2/m) independently. Comparison with LinUCB and NeuralUCB: Compared with linear contextual bandits in Section 2.1, Algorithm 1 has a distinct feature that it learns a deep neural network to obtain a deep representation of the raw data vectors and then performs UCB exploration. This deep representation allows our algorithm to characterize more intrinsic and latent information about the raw data {xt,k}t∈[T ],k∈[K] ⊂ Rd. However, the increased complexity of the feature mapping φ(·; w) also introduces great hardness in training. For instance, a recent work by Zhou et al. (2020) also studied the neural contextual bandit problem, but different from (3.1), their algorithm (NeuralUCB) performs the UCB exploration on the entire network parameter space, which is Rp̃+d, where p̃ = m + md + (L − 1)m2. Note that in Zhou et al. (2020), they need to compute the inverse of a matrix Zt ∈ R(p̃+d)×(p̃+d), which is defined in a similar way to the matrix At in our paper except that Zt is defined based on the gradient of the network instead of the output of the last hidden layer as in (3.2). In sharp contrast, At in our paper is only of size d × d and thus is much more efficient and practical in implementation, which will be seen from our experiments in later sections.
  We note that there is also a similar algorithm to our Neural-LinUCB presented in Deshmukh et al. (2020), where they studied the self-supervised learning loss in contextual bandits with neural network representation for computer vision problems. However, no regret analysis has been provided. When the feature mapping φ(·; w) is an identity function, the problem reduces to linear contextual bandits where we directly use xt as the feature vector. In this case, it is easy to see that Algorithm 1 reduces to LinUCB (Chu et al., 2011) since we do not need to learn the representation parameter w.
  Comparison with Neural-Linear: The high-level idea of decoupling the representation and exploration in our algorithm is also similar to that of the Neural-Linear algorithm (Riquelme et al., 2018; Zahavy & Mannor, 2019), which trains a deep neural network to learn a representation of the raw feature vectors, and then uses a Bayesian linear regression to estimate the uncertainty in the bandit problem. However, these two algorithms are significantly different since Neural-Linear (Riquelme et al., 2018) is a Thompson sampling based algorithm that uses posterior sampling to estimate the weight parameter θ∗ via Bayesian linear regression, whereas Neural-LinUCB adopts upper confidence bound based techniques to estimate the weight θ∗. Nevertheless, both algorithms share the same idea of deep representation and shallow exploration, and we view our Neural-LinUCB algorithm as one instantiation of the Neural-Linear scheme.
  Algorithm 1 Deep Representation and Shallow Exploration (Neural-LinUCB) 1: Input: regularization parameter λ > 0, number of total steps T , episode length H , exploration
  parameters {αt > 0}t∈[T ] 2: Initialization: A0 = λI, b0 = 0; entries of θ0 follow N(0, 1/d), and w(0) is initialized as
  described in Section 3; q = 1; w0 = w(0) 3: for t = 1, . . . , T do 4: receive feature vectors {xt,1, . . . ,xt,K} 5: choose arm at = argmaxk∈[K] θ
  > t−1φ(xt,k; wt−1) +αt‖φ(xt,k; wt−1)‖A−1t−1 , and obtain
  reward r̂t 6: update At and bt as follows:
  At = At−1 + φ(xt,at ; wt−1)φ(xt,at ; wt−1) >, bt = bt−1 + r̂tφ(xt,at ; wt−1),
  7: update θt = A−1t bt 8: if mod(t,H) = 0 then 9: wt ← output of Algorithm 2
  10: q = q + 1 11: else 12: wt = wt−1 13: end if 14: end for 15: Output wT
  Algorithm 2 Update Weight Parameters with Gradient Descent
  1: Input: initial point w(0)q = w(0), maximum iteration number n, step size ηq , and loss function defined in (3.3). 2: for s = 1, . . . , n do 3: w(s)q = w (s−1) q − ηq∇wLq(w(s−1)q ). 4: end for 5: Output w(n)q
sections.heading: 4 MAIN RESULTS
sections.text: To analyze the regret bound of Algorithm 1, we first lay down some important assumptions on the neural contextual bandit model. Assumption 4.1. For all i ≥ 1 and k ∈ [K], we assume that ‖xi,k‖2 = 1 and its entries satisfy [xi,k]j = [xi,k]j+d/2.
  The assumption that ‖xi,k‖2 = 1 is not essential and is only imposed for simplicity, which is also used in Zou & Gu (2019); Zhou et al. (2020). The condition on the entries of xi,k is also mild since otherwise we could always construct x′i,k = [x > i,k,x > i,k] >/ √
  2 to replace it. An implication of Assumption 4.1 is that the initialization scheme in Algorithm 1 results in φ(xi,k; w(0)) = 0 for all i ∈ [T ] and k ∈ [K]. We assume the following stability condition on the spectral norm of the neural network gradient:
  Assumption 4.2. There is a constant `Lip > 0 such that it holds ∥∥ ∂φ ∂w (x; w0) − ∂φ ∂w (x ′; w0) ∥∥ 2 ≤ `Lip‖x− x′‖2 for all x,x′ ∈ {xi,k}i∈[T ],k∈[K].
  The inequality in Assumption 4.2 resembles the Lipschitz condition on the gradient of the neural network. However, it is essentially different from the smoothness condition since here the gradient is taken with respect to the neural network weights while the Lipschitz condition is imposed on the feature parameter x. Similar conditions are widely made in nonconvex optimization (Wang et al., 2014; Balakrishnan et al., 2017; Xu et al., 2017), in the name of first-order stability, which is essential to derive the convergence of alternating optimization algorithms. Furthermore, Assumption 4.2 is only required on the TK training data points and a specific weight parameter w0. Therefore, the condition will hold if the raw feature data lie in a certain subspace of Rd. We provided some further discussions in the supplementary material about this assumption for interested readers.
  In order to analyze the regret bound of Algorithm 1, we need to characterize the properties of the deep neural network in (2.3) that is used to represent the feature vectors. Following a recent line of research (Jacot et al., 2018; Cao & Gu, 2019a; Arora et al., 2019b; Zhou et al., 2020), we define the covariance between two data point x,y ∈ Rd as follows.
  Σ̃(0)(x,y) = Σ(0)(x,y) = x>y,
  Λ(l)(x,y) = [ Σl−1(x,x) Σl−1(x,y) Σl−1(y,x) Σl−1(y,y) ] ,
  Σ(l)(x,y) = 2E(u,v)∼N(0,Λ(l−1)(x,y))[σ(u)σ(v)],
  Σ̃(l)(x,y) = 2Σ̃(l−1)(x,y)Eu,v[σ̇(u)σ̇(v)] + Σ(l)(x,y), (4.1)
  where (u, v) ∼ N(0,Λ(l−1)(x,y)), and σ̇(·) is the derivative of activation function. We denote the neural tangent kernel (NTK) matrix H ∈ RTK×TK based on all feature vectors {xt,k}t∈[T ],k∈[K]. Renumbering {xt,k}t∈[T ],k∈[K] as {xi}i=1,...,TK , then each entry Hij is defined as
  Hij = 1
  2
  ( Σ̃(L)(xi,xj) + Σ (L)(xi,xj) ) , (4.2)
  for all i, j ∈ [TK]. Based on the above definition, we impose the following assumption on H. Assumption 4.3. The neural tangent kernel defined in (4.2) is positive definite, i.e., λmin(H) ≥ λ0 for some constant λ0 > 0.
  Assumption 4.3 essentially requires the neural tangent kernel matrix H to be non-singular, which is a mild condition and also imposed in other related work (Du et al., 2019a; Arora et al., 2019b; Cao & Gu, 2019a; Zhou et al., 2020). Moreover, it is shown that Assumption 4.3 can be easily derived from Assumption 4.1 for two-layer ReLU networks (Oymak & Soltanolkotabi, 2020; Zou & Gu, 2019). Therefore, Assumption 4.3 is mild or even negligible given the non-degeneration assumption on the feature vectors. Also note that matrix H is only defined based on layers l = 1, . . . , L of the neural network, and does not depend on the output layer θ. It is easy to extend the definition of H to the NTK matrix defined on all layers including the output layer θ, which would also be positive definite by Assumption 4.3 and the recursion in (4.2).
  Before we present the regret analysis of the neural contextual bandit, we need to modify the regret defined in (2.2) to account for the randomness of the neural network initialization. For a fixed time horizon T , we define the regret of Algorithm 1 as follows.
  RT = E [ T∑ t=1 ( r̂(xt,a∗t )− r̂(xt,at) ) |w(0) ] , (4.3)
  where the expectation is taken over the randomness of the reward noise. Note that RT defined in (4.3) is still a random variable since the initialization of Algorithm 2 is randomly generated.
  Now we are going to present the regret bound of the proposed algorithm.
  Theorem 4.4. Suppose Assumptions 4.1, 4.2 and 4.3 hold. Assume that ‖θ∗‖2 ≤ M for some positive constant M > 0. For any δ ∈ (0, 1), let us choose αt in Neural-LinUCB as
  αt = ν √ 2 ( d log(1 + t log(HK)/λ) + log(1/δ) ) + λ1/2M.
  We choose the step size ηq of Algorithm 2 as ηq ≤ C0 ( d2mnT 5.5L6 log(TK/δ) )−1 and the width of the neural network satisfies m = poly(L, d, 1/δ,H, log(TK/δ)). With probability at least 1− δ over the randomness of the initialization of the neural network, it holds that
  RT ≤ C1αT √ Td log ( 1 + TG2
  λd
  ) + C2`LipL
  3d5/2T √
  logm log( 1 δ ) log(TK δ )‖r− r̃‖H−1
  m1/6 ,
  where constants {Ci}i=0,1,2 are independent of the problem, r = (r(x1), r(x2), . . . , r(xTK))> ∈ RTK and r̃ = (f(x1;θ0,w0), . . . , f(xTK ;θT−1,wT−1))> ∈ RTK , and ‖r‖A = √ r>Ar.
  Remark 4.5. Theorem 4.4 shows that the regret of Algorithm 1 can be bounded by two parts: the first part is of order Õ( √ T ), which resembles the regret bound of linear contextual bandits (Abbasi-
  Yadkori et al., 2011); the second part is of order Õ(m−1/6T √
  (r− r̃)>H−1(r− r̃)), which depends on the estimation error of the neural network f for the reward generating function r and the neural tangent kernel H.
  It is worth noting that our theoretical analysis depends on the reward structure assumption that r(·) = 〈θ∗,ψ(·)〉. However, the linear structure between θ∗ and ψ(·) is not essential. As long as the deep representation of the feature vector and the uncertainty weight parameter can be decoupled, Algorithm 1 can be easily extended to settings with milder assumptions on the reward structure such as generalized linear models (Sarkar, 1991; Filippi et al., 2010; Li et al., 2017; Kveton et al., 2020). For more general bandit models where no assumption is imposed to the reward generating function, it is still unclear whether the decoupled deep representation and shallow exploration would work especially in cases a thorough exploration may be needed.
  Based on the result in Theorem 4.4, we can easily verify the following conclusion:
  Corollary 4.6. Under the same conditions of Theorem 4.4, if we choose a sufficiently overparameterized neural network mapping φ(·) such that m ≥ T 3, then the regret of Algorithm 1 is RT = Õ( √ T √
  (r− r̃)>H−1(r− r̃)). Remark 4.7. For the ease of presentation, let us denote E := ‖r− r̃‖H−1 . If we have E = O(1), the total regret in Theorem 4.4 becomes Õ( √ T ) which matches the regret of linear contextual bandits (Abbasi-Yadkori et al., 2011). We remark that there is a similar assumption in Zhou et al. (2020) where they assume that r>H−1r can be upper bounded by a constant. They show that this term can be bounded by the RKHS norm of r if it belongs to the RKHS induced by the neural tangent kernel (Arora et al., 2019a;b; Lee et al., 2019). In addition, E here is the difference between the true reward function and the neural network function, which can also be small if the deep neural network function well approximates the reward generating function r(·).
sections.heading: 5 EXPERIMENTS
sections.text: In this section, we provide empirical evaluations of Neural-LinUCB on real-world datasets. As we have discussed in Section 3, Neural-LinUCB could be viewed as an instantiation of the NeuralLinear scheme studied in Riquelme et al. (2018) except that we use the UCB exploration instead of the posterior sampling exploration therein. Note that there has been an extensive comparison (Riquelme et al., 2018) of the Neural-Linear methods with many other baselines such as greedy algorithms, Variational Inference, Expectation-Propagation, Bayesian Non-parametrics and so on. Therefore, we do not seek a thorough empirical comparison of Neural-LinUCB with all existing bandits algorithms. In this experiment, we only aim to validate the advantages of our algorithm over the following baselines: (1) Neural-Linear (Riquelme et al., 2018); (2) LinUCB (Chu et al., 2011), which does not have a deep representation of the feature vectors; (3) NeuralUCB (Zhou et al., 2020), and (4) NeuralTS (Zhang et al., 2020) which perform UCB/TS exploration on all the parameters of the neural network. All numerical experiments were run on a workstation with Intel(R) Xeon(R) CPU E5-2637 v4 @ 3.50GHz.
  Datasets: we evaluate the performances of all algorithms on bandit problems created from realworld data. Specifically, following the experimental setting in Zhou et al. (2020),we use datasets (Shuttle) Statlog, Magic and Covertype from UCI machine learning repository (Dua & Graff, 2017), and the MINST dataset from LeCun et al. (1998). The details of these datasets are presented in Table 1. In Table 1, each instance represents a feature vector x ∈ Rd that is associated with one of the K arms, and dimension d is the number of attributes in each instance.
  Implementations: for LinUCB, we follow the setting in Li et al. (2010) to use disjoint models for different arms. For neural network based algorithms, we use a ReLU neural network defined as in (2.3) with L = 2 and m = 100 for the UCI datasets (Statlog, Magic, Covertype). Thus the neural network weights are W1 ∈ Rm×d, W2 ∈ Rk×m, and θ ∈ Rk respectively, where k = 100, m = 100, and d is the dimension of features in the corresponding task. Since the problem size of the MNIST dataset is larger, inspired by Hinton & Salakhutdinov (2006), we use a deeper NN and set L = 3, k = 100 and m = 100, with weights W1 ∈ Rm×d, W2 ∈ Rm×m, W3 ∈ Rk×m, and θ ∈ Rk. We set the time horizon T = 15, 000, which is the total number of rounds for each algorithm on each dataset. We use stochastic gradient decent to optimize the network weights, with a step size ηq =1e-5 and maximum iteration number n = 1, 000. To speed up the training process, the network parameter w is updated everyH = 100 rounds starting from round 2000. We also apply early stopping when the loss difference of two consecutive iterations is smaller than a threshold of 1e-6. We set λ = 1 and αt = 0.02 for all algorithms, t ∈ [T ]. For NeuralUCB and NeuralTS, since it is computationally unaffordable to perform the original UCB exploration as displayed in Zhou et al. (2020), we follow their experimental setting to replace the matrix Zt ∈ R(d+p̃)×(d+p̃) in their papers with its diagonal matrix.
  Results: we plot the cumulative regret of all algorithms versus round in Figures 1(a), 1(b) and 1(c) for UCI datasets and in Figure 1(d) for MNIST. The results are reported based on the average of 10 repetitions over different random shuffles of the datasets. It can be seen that algorithms based on neural network representations (NeuralUCB, NeuralTS, Neural-Linear and Neural-LinUCB) consistently outperform the linear contextual bandit method LinUCB, which shows that linear models may lack representation power and find biased estimates for the underlying reward generating function. Furthermore, our proposed Neural-LinUCB achieves a comparable regret with NeuralUCB in all experiments despite the fact that our algorithm only explores in the output layer of the neural network, which is more computationally efficient as we will show in the sequel.The results in our experiment are well aligned with our theory that deep representation and shallow exploration are sufficient to guarantee a good performance of neural contextual bandit algorithms, which is also consistent with the findings in existing literature (Riquelme et al., 2018) that decoupling the representation learning and uncertainty estimation improves the performance.
  We also conducted experiments to study the effects of different widths of deep neural networks on the regret performance and to show the computational efficiency of Neural-LinUCB compared with existing neural bandit algorithms. Due to the space limit, we defer the results to Appendix A.
sections.heading: 6 CONCLUSIONS
sections.text: In this paper, we propose a new neural contextual bandit algorithm called Neural-LinUCB, which uses the hidden layers of a ReLU neural network as a deep representation of the raw feature vectors and performs UCB type exploration on the last layer of the neural network. By incorporating techniques in liner contextual bandits and neural tangent kernels, we prove that the proposed algorithm achieves a sublinear regret when the width of the network is sufficiently large. This is the first regret analysis of neural contextual bandit algorithms with deep representation and shallow exploration, which have been observed in practice to work well on many benchmark bandit problems (Riquelme et al., 2018). We also conducted experiments on real-world datasets to demonstrate the advantage of the proposed algorithm over LinUCB and existing neural contextual bandit algorithms.
sections.heading: A ADDITIONAL EXPERIMENTAL RESULTS
sections.text: In this section, we provide more experimental results that are omitted in Section 5 due to space limit.
  A.1 COMPUTATIONAL EFFICIENCY OF NEURAL-LINUCB
  Throughout the experiments, our Neural-LinUCB algorithm is much more computationally efficient than NeuralUCB since we only perform the UCB exploration on the last layer of the neural network, where the dimension is much lower. In specific, on the Statlog dataset, it takes on average 1.11 seconds for NeuralUCB to finish 100 rounds (one epoch in Algorithm 1) and achieve the regret in Figure 1(a), while it only takes 0.58 seconds for Neural-LinUCB to finish 100 rounds and achieve the comparable or even better regret in Figure 1(a). On the Magic dataset, the average runtimes for 100 rounds of NeuralUCB and Neural-LinUCB are 1.32 seconds and 0.81 seconds respectively. On the Covertype dataset, the runtimes of NeuralUCB and Neural-LinUCB are 1.02 seconds and 0.66 seconds respectively. And on the MNIST dataset, the average runtimes for 100 rounds of NeuralUCB and Neural-LinUCB are 4.67 seconds and 1.29 seconds respectively. For practical applications in the real-world with larger problem sizes, we believe that the improvement of our algorithm in terms of the computational efficiency will be more pronounced.
  As we discussed in Section 5 and in the above paragraph as well, the computational efficiency of Neural-LinUCB mainly stems from the design of shallow exploration. This is because in UCB based bandit algorithms we need to compute the inverse of matrix A at every time step for arm selection (Line 5 of Algorithm 1). Due to the large width of the neural network used in practice, the arm selection operation could be rather time consuming. However, the neural network weight can be updated periodically (i.e., in our paper it is only updated every H steps). To validate our analysis on computational efficiency, we further studied the time profiling of the experiments conducted on MNIST to compared our proposed algorithm with NeuralUCB in more details.
  In particular, the setting is the same as that in Section 5 for MNIST experiments. We record the time cost of the most expensive two subroutines: (1) the operation of arm selection (Line 5 in Algorithm 1); and (2) the operation of updating the neural network weights (Line 9 in Algorithm 1), for H = 100 rounds. The time cost is presented in Table 2. For Neural-LinUCB, the arm selection operation takes about 0.28 seconds (this is 21.71% of the total time cost by the algorithm in these H = 100 rounds), among which the matrix inverse step only takes 0.17 seconds. For NeuralUCB, the arm selection operation takes 3.60 seconds (this is 77.19% of the total time time cost by NeuralUCB for H = 100 rounds). Therefore, the operation of arm selection in NeuralUCB is much (almost 13 times) more time consuming than that in Neural-LinUCB. Moreover, since the UCB matrix Zt in NeuralUCB is defined as ∇f(x; w)∇f(x; w)>, it needs to compute the gradients via back-propagation (0.93 seconds) and compute the matrix inverse (1.54 seconds), while our Neural-LinUCB algorithm only needs to compute the matrix inverse of a small matrix (0.17 seconds). To summarize, our method is much more computationally efficient.
  A.2 IMPACT OF LARGE WIDTHS
  Note that the requirement of width m in our Theorem 4.4 is extremely high. On one hand, our theory may be too conservative since the current understanding of deep learning is still very limited in the field. We believe our work is a good starting point towards understanding the behavior of deep bandits algorithms. On the other hand, we would also like to investigate the impact of mild overparameterization on the regret performance of Neural-LinUCB in practice. Therefore, we conducted additional experiments on the Statlog dataset with wider neural networks. In particular, the neural
  network parameters are listed as follows
  W1 ∈ Rm×d,W2 ∈ Rm×m, . . . ,WL ∈ Rk×m,θ ∈ Rk, where L is the depth, k = 100, d is the feature dimensions, and m is the width. We conducted experiments for the following settings: (1) L = 2, m = 100, and thus the hidden layer width is (100); (2) L = 2, m = 1000, and thus the hidden layer width is (1000); (3) L = 2, m = 10000, and thus the hidden layer width is (10000); (4) L = 2, and the hidden layer width is (50000); (5) L = 3, m = 10000, and thus the hidden layer width is (10000, 10000); and (6)L = 4,m = 10000, and thus the hidden layer width is (10000, 10000, 10000). The results are plotted in Figure A.2. We observe that the performance of our Neural-LinUCB algorithm is not negatively impacted by the width of the neural network. In fact, Figure A.2 shows improved performance of Neural-LinUCB when the total number of hidden nodes increases. This is consistent to the observations in Zhang et al. (2017) that an overparameterized neural network trained by gradient descent does not necessarily lead to overfitting and also aligns with our Theorem 4.4 that the regret bound of Neural-LinUCB decreases as the width m increases.
sections.heading: B MORE DISCUSSIONS ON ASSUMPTION 4.2
sections.text: In this section, we are going to show that Assumption 4.2 could be satisfied as long as the feature vectors {x} lie in a begin subspace of Rd. Let us start with the case that φ : Rd → Rm is a two-layer ReLU neural network with vector output. In particular, we define φ(x; w) as follows φ(x; w) = σ(W2σ(W1x)), where w = (vec(W1), vec(W2))>, W1 ∈ Rm×d, W2 ∈ Rd×m, and σ is the ReLU activation function applied elementwise. We use u>i to denote the i-th row of W1 and thus W1 = (u1, . . . ,um)>, where ui ∈ Rd, ∀i ∈ [m]. Similarly, we have W2 = (v1, . . . ,vd)>, where vj ∈ Rm is the j-th row of W2, ∀j ∈ [d]. Let us denote h as the vector σ(W1x). We thus obtain
  h = 1{u > 1 x ≥ 0}u>1 x
  ... 1{u>mx ≥ 0}u>mx
   , φ(x; w) = 1{v > 1 h ≥ 0}v>1 h
  ... 1{v>d h ≥ 0}v>d h  . We use φl(x; w) to denote the l-th entry of vector φ(x; w), for any l ∈ [d]. Then it holds that
  ∂φl(x; w) ∂vec(W1) = 1{v>1 h ≥ 0}
  [ v11 1{u>1 x ≥ 0}x>, . . . , vm1 1{u>mx ≥ 0}x> ] ,
  for all l ∈ [d], where vi1 is the i-th element in v1, i ∈ [m]. This further implies that
  ∂φ(x; w) ∂vec(W1) =
  1{v > 1 h ≥ 0} ( v11 1{u>1 x ≥ 0}x>, . . . , vm1 1{u>mx ≥ 0}x> ) ...
  1{v>d h ≥ 0} ( v1d 1{u>1 x ≥ 0}x>, . . . , vmd 1{u>mx ≥ 0}x>
  )  ∈ Rd×md.
  Similarly, we can compute the gradient of φ(x; w) with respect to W2. In particular, we have
  ∂φ1(x; w)
  ∂vec(v1) = 1{v>1 h ≥ 0}
  [ 1{u>1 x ≥ 0}u>1 x, . . . ,1{u>mx ≥ 0}u>mx ] ,
  ∂φ1(x; w)
  ∂vec(vj) = [0, . . . , 0], j 6= 1.
  Therefore, the gradient of φ(x; w) with respect to W2 is
  ∂φ(x; w) ∂vec(W2) =
   ∂φ1(x;w) ∂vec(v1) . . . 0>
  . . . 0> . . . ∂φd(x;w)∂vec(vd)  ∈ Rd×md. Lastly, we have
  ∂φ(x; w) ∂w = [ ∂φ(x;w) ∂vec(W1) ∂φ(x;w) ∂vec(W2) ] ∈ Rd×(md+md).
  Therefore, for any two feature vectors x and x′ from {xi,k}i∈[T ],k∈[K], if many nodes in the initial neural network φ(x; w0) are activated or deactivated at the same time for both x and x′, then the spectral norm of the matrix ∂φ(x;w0)∂w − ∂φ(x′;w0) ∂w would satisfy the condition in Assumption 4.2. A more thorough study of this stability condition is out of the scope of this paper, though it would be an interesting open direction in the theory of deep neural networks.
sections.heading: C PROOF OF THE MAIN RESULTS
sections.text: In this section, we provide the proof of the regret bound for Neural-LinUCB. Recall that in neural contextual bandits, we do not assume a specific formulation of the underlying reward generating function r(·). Instead, we use deep neural networks defined in Section 2.2 to approximate r(·). We will first show that the reward generating function r(·) can be approximated by the local linearization of the overparameterized neural network near the initialization weight w(0). In particular, we denote the gradient of φ(x; w) with respect to w by g(x; w), namely,
  g(x; w) = ∇wφ(x; w), (C.1)
  which is a matrix in Rd×p. We define φj(x; w) to be the j-th entry of vector φ(x; w), for any j ∈ [d]. Then, we can prove the following lemma. Lemma C.1. Suppose Assumptions 4.3 hold. Then there exists w∗ ∈ Rp such that ‖w∗−w(0)‖2 ≤ 1/ √ m √ (r− r̃)>H−1(r− r̃) and it holds that
  r(xt,k) = θ ∗>φ(xt,k; wt−1) + θ > 0 g(xt,k; w
  (0)) ( w∗ −w(0) ) ,
  for all k ∈ [K] and t = 1, . . . , T .
  Lemma C.1 implies that the reward generating function r(·) at points {xi,k}i∈[T ],k∈[K] can be approximated by a linear function around the initial point w(0). Note that a similar lemma is also proved in Zhou et al. (2020) for NeuralUCB.
  The next lemma shows the upper bounds of the output of the neural network φ and its gradient. Lemma C.2. Suppose Assumptions 4.1 and 4.3 hold. For any round index t ∈ [T ], suppose it is in the q-th epoch of Algorithm 2, i.e., t = (q − 1)H + i for some i ∈ [H]. If the step size ηq in Algorithm 2 satisfies
  η ≤ C0 d2mnT 5.5L6 log(TK/δ) ,
  and the width of the neural network satisfies
  m ≥ max{L log(TK/δ), dL2 log(m/δ), δ−6H18L16 log3(TK)}, (C.2)
  then, with probability at least 1− δ we have
  ‖wt −w(0)‖2 ≤ δ3/2
  m1/2Tn9/2L6 log3(m) ,
  ‖g(xt,k; w(0))‖F ≤ C1 √ dLm,
  ‖φ(x; wt)‖2 ≤ √ d log(n) log(TK/δ),
  for all t ∈ [T ], k ∈ [K], where the neural network φ is defined in (2.4) and its gradient is defined in (C.1).
  The next lemma shows that the neural network φ(x; w) is close to a linear function in terms of the weight w parameter around a small neighborhood of the initialization point w(0).
  Lemma C.3 (Theorems 5 in Cao & Gu (2019b)). Let w,w′ be in the neighborhood of w0, i.e., w,w′ ∈ B(w0, ω) for some ω > 0. Consider the neural network defined in (2.4), if the width m and the radius ω of the neighborhood satisfy
  m ≥ C0 max{dL2 log(m/δ), ω−4/3L−8/3 log(TK) log(m/(ωδ))}, ω ≤ C1L−5(logm)−3/2,
  then for all x ∈ {xt,k}t∈[T ],k∈[K], with probability at least 1− δ it holds that |φj(x; w)− φ̂j(x; w)| ≤ C2ω4/3L3d−1/2 √ m logm,
  where φ̂j(x; w) is the linearization of φj(x; w) at w′ defined as follow:
  φ̂j(x; w) = φj(x; w ′) + 〈∇wφj(x; w′),w −w′〉. (C.3)
  Similar results on the local linearization of an overparameterized neural network are also presented in Allen-Zhu et al. (2019b); Cao & Gu (2019b).
  For the output layer θ∗, we perform a UCB type exploration and thus we need to characterize the uncertainty of the estimation. The next lemma shows the confidence bound of the estimate θt in Algorithm 1.
  Lemma C.4. Suppose Assumption and 4.3 hold. For any δ ∈ (0, 1), with probability at least 1− δ, the distance between the estimated weight vector θt by Algorithm 1 and θ∗ can be bounded as follows: ∥∥∥∥θt − θ∗ −A−1t t∑
  s=1
  φ(xs,as ; ws−1)θ > 0 g(xs,as ; w (0))(w∗ −w(0)) ∥∥∥∥
  At ≤ ν √ 2 ( d log(1 + t(logHK)/λ) + log 1/δ ) + λ1/2M,
  for any t ∈ [T ].
  Note that the confidence bound in Lemma C.4 is different from the standard result for linear contextual bandits in Abbasi-Yadkori et al. (2011). The additional term on the left hand side of the confidence bound is due to the bias caused by the representation learning using a deep neural network. To deal with this extra term, we need the following technical lemma.
  Lemma C.5. Assume that At = λI + ∑t s=1 φsφ > s , where φt ∈ Rd and ‖φt‖2 ≤ G for all t ≥ 1 and some constants λ,G > 0. Let {ζt}t=1,... be a real-value sequence such that |ζt| ≤ U for some constant U > 0. Then we have∥∥∥∥A−1t t∑
  s=1
  φsζs ∥∥∥∥ 2 ≤ 2Ud, ∀t = 1, 2, . . .
  The next lemma provides some standard bounds on the feature matrix At, which is a combination of Lemma 10 and Lemma 11 in Abbasi-Yadkori et al. (2011).
  Lemma C.6. Let {xt}∞t=1 be a sequence in Rd and λ > 0. Suppose ‖xt‖2 ≤ G and λ ≥ max{1, G2} for some G > 0. Let At = λI + ∑t s=1 xtx > t . Then we have
  det(At) ≤ (λ+ tG2/d)d, and T∑ t=1 ‖xt‖2A−1t−1 ≤ 2 log det(AT ) det(λI) ≤ 2d log(1 + TG2/(λd)).
  Now we are ready to prove the regret bound of Algorithm 1.
  Proof of Theorem 4.4. For a time horizon T , without loss of generality, we assume T = QH for some epoch number Q. By the definition of regret in (4.3), we have
  RT = E [ T∑ t=1 (r̂(xt,a∗t )− r̂(xt,at)) ] = E [ Q∑ q=1 H∑ i=1 (r̂(xqH+i,a∗qH+i)− r̂(xqH+i,aqH+i)) ] .
  Note that for the simplicity of presentation, we omit the conditional expectation notation of w(0) in the rest of the proof when the context is clear. In the second equation, we rewrite the time index t = qH + i as the i-th iteration in the q-th epoch.
  By the definition in (2.1), we have E[r̂(xt,k)|xt,k] = r(xt,k) for all t ∈ [T ] and k ∈ K. Based on the linearization of reward generating function, we can decompose the instaneous regret into different parts and upper bound them individually. In particular, by Lemma C.1, there exists a vector w∗ ∈ Rp such that we can write the expectation of the reward generating function as a linear function. Then it holds that
  r(xt,a∗t )− r(xt,at) = θ > 0 [ g ( xt,a∗t ; w (0) ) − g ( xt,at ; w (0) )]( w∗ −w(0) )
  + θ∗> [ φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 )] = θ>0 [ g ( xt,a∗t ; w (0) ) − g ( xt,at ; w (0) )]( w∗ −w(0) )
  + θ>t−1 [ φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 )] − (θt−1 − θ∗)> [ φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 )] . (C.4)
  The first term in (C.4) can be easily bounded using the first order stability in Assumption 4.2 and the distance between w∗ and w(0) in Lemma C.1. The second term in (C.4) is related to the optimistic rule of choosing arms in Line 5 of Algorithm 1, which can be bounded using the same technique for LinUCB (Abbasi-Yadkori et al., 2011). For the last term in (C.4), we need to prove that the estimate of weight parameter θt−1 lies in a confidence ball centered at θ∗. For the ease of notation, we define
  Mt = A −1 t t∑ s=1 φ(xs,as ; ws−1)θ > 0 g(xs,as ; w (0))(w∗ −w(0)). (C.5)
  Then the second term in (C.4) can be bounded in the following way:
  − (θt−1 − θ∗)> [ φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 )] = − ( θt−1 − θ∗ −Mt−1 )> φ ( xt,a∗t ; wt−1 ) + ( θt−1 − θ∗ −Mt−1 )> φ ( xt,at ; wt−1
  ) −M>t−1 [ φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1
  )] ≤ ‖θt−1 − θ∗ −Mt−1‖At−1 · ‖φ(xt,a∗t ; wt−1)‖A−1t−1
  + ‖θt−1 − θ∗ −Mt−1‖At−1 · ‖φ(xt,at ; wt−1)‖A−1t−1 + ∥∥M>t−1[φ(xt,a∗t ; wt−1)− φ(xt,at ; wt−1)]∥∥2
  ≤ αt‖φ(xt,a∗t ; wt−1)‖A−1t−1 + αt‖φ(xt,at ; wt−1)‖A−1t−1 + ‖Mt−1‖2 · ‖φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 ) ‖2. (C.6)
  where the last inequality is due to Lemma C.4 and the choice of αt. Plugging (C.6) back into (C.4) yields
  r(xt,a∗t )− r(xt,at) ≤ αt‖φ(xt,at ; wt−1)‖A−1t−1 − αt‖φ(xt,a∗t ; wt−1)‖A−1t−1 + αt‖φ(xt,a∗t ; wt−1)‖A−1t−1 + αt‖φ(xt,at ; wt−1)‖A−1t−1 + ‖Mt−1‖2 · ‖φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 ) ‖2
  + ‖θ0‖2 · ‖g(xt,a∗t ; w (0))− g(xt,at ; w(0))‖F · ‖w∗ −w(0)‖2 ≤ 2αt‖φ(xt,at ; wt−1)‖A−1t−1 + ‖Mt−1‖2 · ‖φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 ) ‖2
  + `Lip‖θ0‖2 · ‖xt,a∗t − xt,at‖2 · ‖w ∗ −w(0)‖2, (C.7)
  where in the first inequality we used the definition of upper confidence bound in Algorithm 1 and the second inequality is due to Assumption 4.2. Recall the linearization of φj in Lemma C.3, we have
  φ̂(x; wt−1) = φ(x; w0) + g(x; w0)(wt−1 −w0).
  Note that by the initialization, we have φ(x; w0) = 0 for any x ∈ Rd. Thus, it holds that φ ( xt,a∗t ; wt−1 ) − φ ( xt,at ; wt−1 ) = φ ( xt,a∗t ; wt−1 ) − φ ( xt,a∗t ; w0 ) + φ ( xt,at ; w0 ) − φ ( xt,at ; wt−1
  ) = φ ( xt,a∗t ; wt−1 ) − φ̂ ( xt,a∗t ; wt−1 ) + g(xt,a∗t ; w0)(wt−1 −w0)
  + φ ( xt,at ; wt−1 ) − φ̂ ( xt,at ; wt−1 ) − g(xt,at ; w0)(wt−1 −w0), (C.8)
  which immediately implies that∥∥φ(xt,a∗t ; wt−1)− φ(xt,at ; wt−1)∥∥2 ≤ ∥∥φ(xt,a∗t ; wt−1)− φ̂(xt,a∗t ; wt−1)∥∥2 + ∥∥φ(xt,at ; wt−1)− φ̂(xt,at ; wt−1)∥∥2
  + ∥∥(g(xt,a∗t ; w0)− g(xt,at ; w0))(wt−1 −w0)∥∥2
  ≤ C0ω4/3L3d1/2 √ m logm+ `Lip‖xt,a∗t − xt,at‖2‖wt−1 −w (0)‖2, (C.9)
  where the second inequality is due to Lemma C.3 and Assumption 4.2. Therefore, the instaneous regret can be further upper bounded as follows.
  r(xt,a∗t )− r(xt,at) ≤ 2αt‖φ(xt,at ; wt−1)‖A−1t−1 + `Lip‖θ0‖2 · ‖xt,a∗t − xt,at‖2 · ‖w ∗ −w(0)‖2
  + ‖Mt−1‖2 · ( C0ω 4/3L3d1/2 √ m logm+ `Lip‖xt,a∗t − xt,at‖2‖wt−1 −w (0)‖2 ) . (C.10)
  By Assumption 4.1 we have ‖xt,a∗t − xt,at‖2 ≤ 2. By Lemma C.1 and Lemma C.2, we have ‖w∗ −w(0)‖2 ≤ √ 1/m(r− r̃)>H−1(r− r̃),
  ‖wt −w(0)‖2 ≤ δ3/2
  m1/2Tn9/2L6 log3(m) .
  (C.11)
  In addition, since the entries of θ0 are i.i.d. generated from N(0, 1/d), we have ‖θ0‖2 ≤ 2(2 + √ d−1 log(1/δ)) with probability at least 1 − δ for any δ > 0. By Lemma C.2, we have ‖g(xt,at ; w(0))‖F ≤ C1 √ dm. Therefore,∣∣θ>0 g(xs,as ; w(0))(w∗ −w(0))∣∣ ≤ C2d√log(1/δ)(r− r̃)>H−1(r− r̃).
  Then, by the definition of Mt in (C.5) and Lemma C.5, we have ‖Mt−1‖2 ≤ C3d2 √ log(1/δ)(r− r̃)>H−1(r− r̃). (C.12)
  Substituting (C.12) and the above results on ‖xt,at − xt,a∗t ‖2, ‖θ0‖2, ‖w ∗ −w(0)‖2 and ‖wt−1 − w(0)‖2 back into (C.10) further yields
  r(xt,a∗t )− r(xt,at)
  ≤ 2αt‖φ(xt,at ; wt−1)‖A−1t−1 + C4`Lipm −1/2
  √ log(1/δ)(r− r̃)>H−1(r− r̃)
  + ( C0ω 4/3L3d1/2 √ m logm+
  2`Lipδ 3/2
  m1/2Tn9/2L6 log3(m)
  ) C3d 2 √ log(1/δ)(r− r̃)>H−1(r− r̃).
  Note that we have ω = O(m−1/2‖r − r̃‖H−1) by Lemma C.1. Therefore, the regret of the Neural-LinUCB is
  RT ≤ √√√√QH max t∈[T ] α2t Q∑ q=1 H∑ i=1 ‖φ(xi,ai ; wqH+i)‖2A−1i + C4`Lipm −1/2T √ log(1/δ)‖r− r̃‖H−1
  +
  ( C0TL 3d1/2 √
  logm‖r− r̃‖4/3H−1 m1/6 + 2`Lipδ 3/2 m1/2n9/2L6 log3(m)
  ) C3d 2 √ log(1/δ)‖r− r̃‖H−1
  ≤ C5 √ Td log(1 + TG2/(λd)) ( ν √ d log(1 + T (log TK)/λ) + log 1/δ + λ1/2M ) + C6`LipL 3d5/2m−1/6T √ logm log(1/δ) log(TK/δ)‖r− r̃‖H−1 ,
  where the first inequality is due to Cauchy’s inequality, the second inequality comes from the upper bound ofαt in Lemma C.4 and Lemma C.6. {Cj}j=0,...,6 are absolute constants that are independent of problem parameters.
  Proof of Corollary 4.6. It directly follows the result in Theorem 4.4.
sections.heading: D PROOF OF TECHNICAL LEMMAS
sections.text: In this section, we provide the proof of technical lemmas used in the regret analysis of Algorithm 1.
  D.1 PROOF OF LEMMA C.1
  Before we prove the lemma, we first present some notations and a supporting lemma for simplification. Let β = (θ>,w>)> ∈ Rd+p be the concatenation of the exploration parameter and the hidden layer parameter of the neural network f(x;β) = θ>φ(x; w). Note that for any input data vector x ∈ Rd, we have
  ∂
  ∂β f(x;β) =
  ( φ(x; w)>,θ> ∂
  ∂w φ(x; w)
  )> = ( φ(x; w)>,θ>g(x; w) )> , (D.1)
  where g(x; w) is the partial gradient of φ(x; w) with respect to w defined in (C.1), which is a matrix in Rd×p. Similar to (4.2), we define HL+1 to be the neural tangent kernel matrix based on all L+ 1 layers of the neural network f(x;β). Note that by the definition of H in (4.2), we must have HL+1 = H + B for some positive definite matrix B ∈ RTK×TK . The following lemma shows that the NTK matrix is close to the matrix defined based on the gradients of the neural network on TK data points. Lemma D.1 (Theorem 3.1 in Arora et al. (2019b)). Let > 0 and δ ∈ (0, 1). Suppose the activation function in (2.3) is ReLU, i.e., σl(x) = max(0, x), and the width of the neural network satisfies
  m ≥ Ω ( L14
  4 log
  ( L
  δ
  )) . (D.2)
  Then for any x,x′ ∈ Rd with ‖x‖2 = ‖x′‖2 = 1, with probability at least 1−δ over the randomness of the initialization of the network weight w it holds that∣∣∣∣〈 1√m ∂f(β,x)∂β , 1√m ∂f(β,x′)∂β 〉 −HL+1(x,x′) ∣∣∣∣ ≤ .
  Note that in the above lemma, there is a factor 1/ √ m before the gradient. This is due to the additional √ m factor in the definition of the neural network in (2.3), which ensures the value of the neural network function evaluated at the initialization is of the order O(1).
  Proof of Lemma C.1. Recall that we renumbered the feature vectors {xt,k}t∈[T ],k∈[K] for all arms from round 1 to round T as {xi}i=1,...,TK . By concatenating the gradients at different inputs and the gradient in (D.1), we define Ψ ∈ RTK×(d+p) as follows.
  Ψ = 1√ m
   ∂ ∂βθ >φ(x1; w) ...
  ∂ ∂βθ >φ(xTK ; w)
   = 1√ m  φ(x1; w (0))> θ>0 g(x1; w (0)) ... ... φ(xi; w (0))> θ>0 g(xi; w (0)) ...
  ... φ(xTK ; w (0))> θ>0 g(xTK ; w (0))
   .
  By Applying Lemma D.1, we know with probability at least 1− δ it holds that |〈Ψj∗,Ψl∗〉 −HL+1(xj ,xl)| ≤
  for any > 0 as long as the width m satisfies the condition in (D.2). By applying union bound over all data points {x1, . . . ,xt, . . . ,xTK}, we further have
  ‖ΨΨ> −HL+1‖F ≤ TK . Note that H is the neural tangent kernel (NTK) matrix defined in (4.2) and HL+1 is the NTK matrix defined based on all L + 1 layers. By Assumption 4.3, H has a minimum eigenvalue λ0 > 0, which is defined based on the first L layers of f . Furthermore, by the definition of NTK matrix in (4.2), we know that HL+1 = H + B for some semi-positive definite matrix B. Therefore, the NTK matrix HL+1 defined based on all L + 1 layers is also positive definite and its minimum eigenvalue is lower bounded by λ0. Let = λ0/(2TK). By triangle equality we have ΨΨ> HL+1 − ‖ΨΨ> − HL+1‖2I HL+1 − ‖ΨΨ> − HL+1‖F I HL+1 − λ0/2I 1/2HL+1, which means that Ψ is semi-definite positive and thus rank(Ψ) = TK since m > TK.
  We assume that Ψ can be decomposed as Ψ = PDQ>, where P ∈ RTK×TK is the eigenvectors of ΨΨ> and thus PP> = ITK , D ∈ RTK×TK is a diagonal matrix with the square root of eigenvalues of ΨΨ>, and Q> ∈ RTK×(d+p) is the eigenvectors of Ψ>Ψ and thus Q>Q = ITK . We use Q1 ∈ Rd×TK and Q2 ∈ Rp×TK to denote the two blocks of Q such that Q> = [Q>1 ,Q>2 ]. By definition, we have
  Q>Q = [Q>1 ,Q > 2 ] [ Q1 Q2 ] = Q>1 Q1 + Q > 2 Q2 = ITK .
  Note that the minimum singular value of Q1 ∈ Rd×TK is zero since d is a fixed number and TK > d. Therefore, it must hold that rank(Q2) = TK and thus Q>2 Q2 is positive definite. Let r = (r(x1), . . . , r(xi), . . . , r(xTK))
  > ∈ RTK denote the vector of all possible rewards. We further define G ∈ RTKd×p and Φ ∈ RTKd as follows
  G = 1√ m  g(x1; w (0)) ... g(xi; w (0))
  ... g(xTK ; w (0))
   , Φ =  φ(x1,1; w0) ... φ(xt,k; wt−1)
  ... φ(xT,K ; wT−1)
   . (D.3)
  and Θ,Θ0 ∈ RTK×TKd as follows
  Θ∗ =  θ∗> . . . θ∗>
  . . . θ∗>
   , Θ0 =  θ>0 . . . θ>0
  . . . θ>0
   , (D.4)
  It can be verified that Ψ = PD[Q>1 ,Q > 2 ] and PDQ > 2 = Θ0G. Note that we have Q > 2 Q2 is positive definite by Assumption 4.3, which corresponds to the neural tangent kernel matrix defined on the first L layers. Then we can define w∗ as follows
  w∗ = w(0) + 1/ √ mQ2(Q > 2 Q2)
  −1D−1P>(r−Θ∗Φ). (D.5) We can verify that
  Θ∗Φ + √ mPDQ>2 (w
  ∗ −w(0)) = r. On the other hand, we have
  ‖w∗ −w(0)‖22 ≤ 1/m(r−Θ∗Φ)>PD−1(Q>2 Q2)−1D−1P>(r−Θ∗Φ) ≤ 1/m(r−Θ∗Φ)>H−1(r−Θ∗Φ),
  which completes the proof.
  D.2 PROOF OF LEMMA C.2
  Note that we can view the output of the last hidden layer φ(x; w) defined in (2.4) as a vector-output neural network with weight parameter w. The following lemma shows that the output of the neural network φ is bounded at the initialization. Lemma D.2 (Lemma 4.4 in Cao & Gu (2019b)). Let δ ∈ (0, 1), and the width of the neural network satisfy m ≥ C0L log(TKL/δ). Then for all t ∈ [T ], k ∈ [K] and j ∈ [d], we have |φj(xt,k; w(0))| ≤ C1 √ log(TK/δ) with probability at least 1 − δ, where w(0) is the initialization of the neural network.
  In addition, in a smaller neighborhood of the initialization, the gradient of the neural network φ is uniformly bounded. Lemma D.3 (Lemma B.3 in Cao & Gu (2019b)). Let ω ≤ C0L−6(logm)−3 and w ∈ B(w0, ω). Then for all t ∈ [T ], k ∈ [K] and j ∈ [d], the gradient of the neural network φ defined in (2.4) satisfies ‖∇wφj(xt,k; w)‖2 ≤ C1 √ Lm with probability at least 1− TKL2 exp(−C2mω2/3L).
  The next lemma provides an upper bound on the gradient of the squared loss function defined in (3.3). Note that our definition of the loss function is slightly different from that in Allen-Zhu et al. (2019b) due to the output layer θi and thus there is an additional term on the upper bound of ‖θi‖2 for all i ∈ [T ]. Lemma D.4 (Theorem 3 in Allen-Zhu et al. (2019b)). Let ω ≤ C0δ3/2/(T 9/2L6 log3m). For all w ∈ B(w(0), ω), with probability at least 1 − exp(−C1mω2/3L) over the randomness of w(0), it holds that
  ‖∇L(w)‖22 ≤ C2TmL(w) supi=1,...,H ‖θi‖22
  d .
  Proof of Lemma C.2. Fix the epoch number q and we omit it in the subscripts in the rest of the proof when no confusion arises. Recall that w(s) is the s-th iterate in Algorithm 2. Let δ > 0 be any constant. Let ω be defined as follows.
  ω = δ3/2m−1/2T−9/2L−6 log−3(m). (D.6)
  We will prove by induction that with probability at least 1− δ the following statement holds for all s = 0, 1, . . . , n
  φj(x; w (s)) ≤ C0 s∑ h=0
  √ log(TK/δ)
  h+ 1 , for ∀j ∈ [d]; and ‖w(s)q −w(0)‖ ≤ ω. (D.7)
  First note that (D.7) holds trivially when s = 0 due to Lemma D.2. Now we assume that (D.7) holds for all j = 0, . . . , s. The loss function in (3.3) can be bounded as follows.
  L(w(j)) = qH∑ i=1 (θ>i φ(xi; w (j))− r̂i)2 ≤ qH∑ i=1 2(‖θi‖22 · ‖φ(xi; w(j))‖22 + 1).
  By the update rule of θt, we have ‖θt‖2 = ∥∥∥∥(λI + t∑
  i=1
  φ(xi; wi−1)φ(xi; wi−1) > )−1 t∑
  i=1
  φ(xi; wi−1)r̂ ∥∥∥∥ 2 ≤ 2d, (D.8)
  where the inequality is due to Lemma C.5, which combined with (D.7) immediately implies
  L(w(j)) ≤ C1Td3 log(TK/δ) ( j∑ h=0 1 h+ 1 )2 ≤ C1Td3 log(TK/δ) log2 n. (D.9)
  Substituting (D.8) and (D.9) into the inequality in Lemma D.4, we also have∥∥∇L(w(j))∥∥ 2 ≤ C2 √ dTmL(w(j)) ≤ C3d2T log(n) √ m log(TK/δ). (D.10)
  Now we consider w(s+1). By triangle inequality we have∥∥w(s+1) −w(0)∥∥ 2 ≤ s∑ j=0 ∥∥w(j+1) −w(j)∥∥ 2
  = s∑ j=0 η ∥∥∇L(w(j))∥∥ 2
  ≤ s∑ j=0 ηd2T log(n) √ m log(TK/δ), (D.11)
  where the last inequality is due to (D.10). If we choose the step size ηq in the q-th epoch such that
  η ≤ ω d2Tn log(n) √ m log(TK/δ) , (D.12)
  then we have ‖w(s+1)q − w(0)‖2 ≤ ω. Note that the choice of m,ω satisfies the condition in Lemma C.3. Thus we know φj(x; w) is almost linear in w, which leads to
  |φj(x; w(s+1))| ≤ |φj(x; w(s)) + 〈∇φj(x; w(s)),w(s+1) −w(s)〉|+ C5ω4/3L3d−1/2 √ m logm
  ≤ s∑
  h=0
  C √ log(TK/δ)
  h+ 1 + η √ dm‖∇L(w(s))‖2 + 2C5ω4/3L3d−1/2
  √ m logm
  ≤ s∑
  h=0
  C0 √ log(TK/δ)
  h+ 1 + C3η
  √ dm √ CT 2d4m log(TK/δ) log n
  + 2C5ω 4/3L3d−1/2 √ m logm
  = s∑ h=0
  C0 √ log(TK/δ)
  h+ 1 + ω √ dm n + 2C5ω
  4/3L3d−1/2 √ m logm, (D.13)
  where in the second inequality we used the induction hypothesis (D.7), Cauchy-Schwarz inequality and Lemma D.3, and the third inequality is due to (D.10). Note that the definition of ω in (D.6) ensures that ω √ dm < 1/2 and ω4/3L3d−1/2 √ m logm ≤ m−1/6T−6L−5d−1/2 √ logm ≤ 1/n as long as m ≥ n6. Plugging these two upper bounds back into (D.13) finishes the proof of (D.7).
  Note that for any t ∈ [T ], we have wt = w(n)q for some q = 1, 2, . . .. Since we have wt ∈ B(w, ω), the gradient g(x; w(0)) can be directly bounded by Lemma D.3, which implies ‖g(x; w(0))‖F ≤ C6 √ dLm. Applying (D.7) with s = n, we have the following bound of the neural network function φ(x; w (n) q ) = φ(x; wt) for all t in the q-th epoch
  ‖φ(x; wt)‖2 ≤ C0 √ d log(n) log(TK/δ),
  which completes the proof. In this proof, {Cj > 0}j=0,...,6 are constants independent of problem parameters.
  D.3 PROOF OF LEMMA C.4
  The following lemma characterizes the concentration property of self-normalized martingales. Lemma D.5 (Theorem 1 in Abbasi-Yadkori et al. (2011)). Let {ξ}∞t=1 be a real-valued stochastic process and {xt}∞t=1 be a stochastic process in Rd. Let Ft = σ(x1, . . . ,xt+1, ξ − 1, . . . , ξt) be a σ-algebra such that xt and ξt are Ft−1-measurable. Let At = λI + ∑t s=1 xsx > s for some constant
  λ > 0 and St = ∑t s=1 ξsxi. If we assume ξt is ν-subGaussian conditional on Ft−1, then for any η ∈ (0, 1), with probability at least 1− δ, we have
  ‖St‖2A−1t ≤ 2ν 2 log
  ( det(At) 1/2 det(λI)−1/2
  δ
  ) .
  Proof of Lemma C.4. Let Φt = [φ(x1,a1 ; w0), . . . ,φ(xt,at ; wt−1)] ∈ Rd×t be the collection of feature vectors of the chosen arms up to time t and r̂t = (r̂1, . . . , r̂t)> be the concatenation of all received rewards. According to Algorithm 1, we have At = λI + ΦtΦ>t and thus
  θt = A −1 t bt = (λI + ΦtΦ > t ) −1Φtr̂t.
  By Lemma C.1, the underlying reward generating function rt = r(xt,at) = E[r̂(xt,at)|xt,at ] can be rewritten as
  rt = 〈θ∗,φ(xt,at ; wt−1)〉+ θ>0 g(xt,at ; w(0))(w∗ −w(0)).
  By the definition of the reward in (2.1) we have r̂t = rt + ξt. Therefore, it holds that
  θt = A −1 t ΦtΦ > t θ ∗ + A−1t t∑ s=1 φ(xs,as ; ws−1)(θ > 0 g(xs,as ; w (0))(w∗ −w(0)) + ξs)
  = θ∗ − λA−1t θ∗ + A−1t t∑
  s=1
  φ(xs,as ; ws−1)(θ > 0 g(xs,as ; w (0))(w∗ −w(0)) + ξs).
  Note that At is positive definite as long as λ > 0. Therefore ‖ · ‖At and ‖ · ‖At are well defined norms. Then for any δ ∈ (0, 1) by triangle inequality we have
  ‖θt − θ∗ −A−1t ΦtΘtGt(w∗ −w(0))‖At ≤ λ‖θ∗‖A−1t + ‖Φtξt‖A−1t
  ≤ ν √ 2 log ( det(At)1/2 det(λI)−1/2
  δ
  ) + λ1/2M
  holds with probability at least 1 − δ, where in the last inequality we used Lemma D.5 and the fact that ‖θ∗‖A−1t ≤ λ
  −1/2‖θ∗‖2 ≤ λ−1/2M by Lemma C.1. Plugging the definition of Φt,Θt and Gt and apply Lemma C.6, we further have∥∥∥∥θt − θ∗ −A−1t t∑
  s=1
  φ(xs,as ; ws−1)θ > 0 g(xs,as ; w (0))(w∗ −w(0)) ∥∥∥∥
  At ≤ ν √ 2 ( d log(1 + t(logHK)/λ) + log 1/δ ) + λ1/2M,
  where we used the fact that ‖φ(x; w)‖2 ≤ C √ d logHK by Lemma C.2.
  D.4 PROOF OF LEMMA C.5 We now prove the technical lemma that upper bounds ‖A−1t ∑t s=1 φsζs‖2.
  Proof of Lemma C.5. We first construct auxiliary vectors φ̃t ∈ Rd+1 and matrices Bt ∈ R(d+1)×(d+1) for all t = 1, . . . in the following way:
  φ̃t =
  [ G−1φt√
  1−G−2‖φt‖22
  ] , Bt = [ A−1t 0d 0>d 0 ] , (D.14)
  where 0d ∈ Rd is an all-zero vector. Then by definition we immediately have∥∥∥∥A−1t t∑ s=1 φsζs ∥∥∥∥ 2 = ∥∥∥∥Bt t∑ s=1 φ̃sζs ∥∥∥∥ 2 . (D.15)
  For all s = 1, 2 . . ., let {βs,j}d+1j=1 be the coefficients of the decomposition ofU−1ζsφ̃s on the natural basis. Specifically, let {e1, . . . , ed+1} be the natural basis of Rd+1 such that the entries of ej are all zero except the j-th entry which equals 1. Then we have
  U−1ζsφ̃s = d∑ j=1 βs,jej , ∀s = 1, 2, . . . (D.16)
  We can conclude that |βs,j | ≤ 1 since |ζs| ≤ U and ‖φ̃s‖2 ≤ 1. Moreover, it is easy to verify that ‖φ̃t‖2 = 1 for all t ≥ 1. Therefore, we have∥∥∥∥Bt t∑
  s=1
  φ̃sζs ∥∥∥∥ 2 = ∥∥∥∥Bt t∑ s=1 φ̃sφ̃ > s φ̃sζs ∥∥∥∥ 2
  = ∥∥∥∥Bt t∑ s=1 φ̃sφ̃ > s U d∑ j=1 βs,jej ∥∥∥∥ 2
  = U ∥∥∥∥ d∑ j=1 Bt t∑ s=1 φ̃sφ̃ > s βs,jej ∥∥∥∥ 2 ≤ U d∑ j=1 ∥∥∥∥Bt t∑ s=1 φ̃sφ̃ > s βs,j ∥∥∥∥ 2
  = U d∑ j=1 ∥∥∥∥A−1t t∑ s=1 φsφ > s βs,j ∥∥∥∥ 2 , (D.17)
  where the inequality is due to triangle inequality and the last equation is due to the definition of φ̃t and Bt in (D.14). For each j = 1, . . . , d+ 1, we have∥∥∥∥A−1t t∑
  s=1
  φsφ > s βs,j ∥∥∥∥ 2 = ∥∥∥∥A−1t ∑ s∈[t]:βs,j≥0 φsφ > s βs,j + A −1 t ∑ s∈[t]:βs,j<0 φsφ > s βs,j ∥∥∥∥ 2
  ≤ ∥∥∥∥A−1t ∑
  s∈[t]:βs,j≥0
  φsφ > s βs,j ∥∥∥∥ 2 + ∥∥∥∥A−1t ∑ s∈[t]:βs,j<0 φsφ > s (−βs,j) ∥∥∥∥ 2 .
  (D.18) Since we have |βs,j | ≤ 1, it immediately implies
  At = λI + t∑ s=1 φsφ > s ∑ s∈[t]:βs,j≥0 φsφ > s βs,j ,
  At = λI + t∑ s=1 φsφ > s ∑ s∈[t]:βs,j<0 φsφ > s (−βs,j).
  Further by the fact that ‖A−1B‖2 ≤ 1 for any A B 0, combining the above results with (D.18) yields ∥∥∥∥A−1t t∑
  s=1
  φsφ > s βs,j ∥∥∥∥ 2 ≤ 2.
  Finally, substituting the above results into (D.17) and (D.15) we have∥∥∥∥A−1t t∑ s=1 φsζs ∥∥∥∥ 2 ≤ 2Ud, which completes the proof.