Neural Thompson Sampling

1 INTRODUCTION . The stochastic multi-armed bandit ( Bubeck & Cesa-Bianchi , 2012 ; Lattimore & Szepesvári , 2020 ) has been extensively studied , as an important model to optimize the trade-off between exploration and exploitation in sequential decision making . Among its many variants , the contextual bandit is widely used in real-world applications such as recommendation ( Li et al. , 2010 ) , advertising ( Graepel et al. , 2010 ) , robotic control ( Mahler et al. , 2016 ) , and healthcare ( Greenewald et al. , 2017 ) . In each round of a contextual bandit , the agent observes a feature vector ( the “ context ” ) for each of the K arms , pulls one of them , and in return receives a scalar reward . The goal is to maximize the cumulative reward , or minimize the regret ( to be defined later ) , in a total of T rounds . To do so , the agent must find a trade-off between exploration and exploitation . One of the most effective and widely used techniques is Thompson Sampling , or TS ( Thompson , 1933 ) . The basic idea is to compute the posterior distribution of each arm being optimal for the present context , and sample an arm from this distribution . TS is often easy to implement , and has found great success in practice ( Chapelle & Li , 2011 ; Graepel et al. , 2010 ; Kawale et al. , 2015 ; Russo et al. , 2017 ) . Recently , a series of work has applied TS or its variants to explore in contextual bandits with neural network models ( Blundell et al. , 2015 ; Kveton et al. , 2020 ; Lu & Van Roy , 2017 ; Riquelme et al. , 2018 ) . Riquelme et al . ( 2018 ) proposed NeuralLinear , which maintains a neural network and chooses the best arm in each round according to a Bayesian linear regression on top of the last network layer . Kveton et al . ( 2020 ) proposed DeepFPL , which trains a neural network based on perturbed training data and chooses the best arm in each round based on the neural network output . Similar approaches have also been used in more general reinforcement learning problem ( e.g. , Azizzadenesheli et al. , 2018 ; Fortunato et al. , 2018 ; Lipton et al. , 2018 ; Osband et al. , 2016a ) . Despite the reported empirical success , strong regret guarantees for TS remain limited to relatively simple models , under fairly restrictive assumptions on the reward function . Examples are linear functions ( Abeille & Lazaric , 2017 ; Agrawal & Goyal , 2013 ; Kocák et al. , 2014 ; Russo & Van Roy , 2014 ) , generalized linear functions ( Kveton et al. , 2020 ; Russo & Van Roy , 2014 ) , or functions with small RKHS norm induced by a properly selected kernel ( Chowdhury & Gopalan , 2017 ) . In this paper , we provide , to the best of our knowledge , the first near-optimal regret bound for neural network-based Thompson Sampling . Our contributions are threefold . First , we propose a new algorithm , Neural Thompson Sampling ( NeuralTS ) , to incorporate TS exploration with neural networks . It differs from NeuralLinear ( Riquelme et al. , 2018 ) by considering weight uncertainty in all layers , and from other neural network-based TS implementations ( Blundell et al. , 2015 ; Kveton et al. , 2020 ) by sampling the estimated reward from the posterior ( as opposed to sampling parameters ) . Second , we give a regret analysis for the algorithm , and obtain an Õ ( d̃ √ T ) regret , where d̃ is the effective dimension and T is the number of rounds . This result is comparable to previous bounds when specialized to the simpler , linear setting where the effective dimension coincides with the feature dimension ( Agrawal & Goyal , 2013 ; Chowdhury & Gopalan , 2017 ) . Finally , we corroborate the analysis with an empirical evaluation of the algorithm on several benchmarks . Experiments show that NeuralTS yields competitive performance , in comparison with stateof-the-art baselines , thus suggest its practical value in addition to strong theoretical guarantees . Notation : Scalars and constants are denoted by lower and upper case letters , respectively . Vectors are denoted by lower case bold face letters x , and matrices by upper case bold face letters A . We denote by [ k ] the set { 1 , 2 , · · · , k } for positive integers k. For two non-negative sequence { an } , { bn } , an = O ( bn ) means that there exists a positive constant C such that an ≤ Cbn , and we use Õ ( · ) to hide the log factor inO ( · ) . We denote by ‖ · ‖2 the Euclidean norm of vectors and the spectral norm of matrices , and by ‖ · ‖F the Frobenius norm of a matrix . 2 PROBLEM SETTING AND PROPOSED ALGORITHM . In this work , we consider contextualK-armed bandits , where the total number of rounds T is known . At round t ∈ [ T ] , the agent observes K contextual vectors { xt , k ∈ Rd | k ∈ [ K ] } . Then the agent selects an arm at and receives a reward rt , at . Our goal is to minimize the following pseudo regret : RT = E [ T∑ t=1 ( rt , a∗t − rt , at ) ] , ( 2.1 ) where a∗t is the optimal arm at round t that has the maximum expected reward : a ∗ t = argmaxa∈ [ K ] E [ rt , a ] . To estimate the unknown reward given a contextual vector x , we use a fully connected neural network f ( x ; θ ) of depth L ≥ 2 , defined recursively by f1 = W1 x , fl = Wl ReLU ( fl−1 ) , 2 ≤ l ≤ L , f ( x ; θ ) = √ mfL , ( 2.2 ) where ReLU ( x ) : = max { x , 0 } , m is the width of neural network , W1 ∈ Rm×d , Wl ∈ Rm×m , 2 ≤ l < L , WL ∈ R1×m , θ = ( vec ( W1 ) ; · · · ; vec ( WL ) ) ∈ Rp is the collection of parameters of the neural network , p = dm + m2 ( L − 2 ) + m , and g ( x ; θ ) = ∇θf ( x ; θ ) is the gradient of f ( x ; θ ) w.r.t . θ . Our Neural Thompson Sampling is given in Algorithm 1 . It maintains a Gaussian distribution for each arm ’ s reward . When selecting an arm , it samples the reward of each arm from the reward ’ s posterior distribution , and then pulls the greedy arm ( lines 4–8 ) . Once the reward is observed , it updates the posterior ( lines 9 & 10 ) . The mean of the posterior distribution is set to the output of the neural network , whose parameter is the solution to the following minimization problem : min θ L ( θ ) = t∑ i=1 [ f ( xi , ai ; θ ) − ri , ai ] 2/2 +mλ‖θ − θ0‖22/2 . ( 2.3 ) We can see that ( 2.3 ) is an ` 2-regularized square loss minimization problem , where the regularization term centers at the randomly initialized network parameter θ0 . We adapt gradient descent to solve ( 2.3 ) with step size η and total number of iterations J . Algorithm 1 Neural Thompson Sampling ( NeuralTS ) Input : Number of rounds T , exploration variance ν , network width m , regularization parameter λ . 1 : Set U0 = λI 2 : Initialize θ0 = ( vec ( W1 ) ; · · · ; vec ( WL ) ) ∈ Rp , where for each 1 ≤ l ≤ L − 1 , Wl = ( W,0 ; 0 , W ) , each entry of W is generated independently from N ( 0 , 4/m ) ; WL = ( w > , −w > ) , each entry of w is generated independently from N ( 0 , 2/m ) . 3 : for t = 1 , · · · , T do 4 : for k = 1 , · · · , K do 5 : σ2t , k = λg > ( xt , k ; θt−1 ) U −1 t−1 g ( xt , k ; θt−1 ) /m 6 : Sample estimated reward r̃t , k ∼ N ( f ( xt , k ; θt−1 ) , ν2σ2t , k ) 7 : end for 8 : Pull arm at and receive reward rt , at , where at = argmaxa r̃t , a 9 : Set θt to be the output of gradient descent for solving ( 2.3 ) 10 : Ut = Ut−1 + g ( xt , at ; θt ) g ( xt , at ; θt ) > /m 11 : end for A few observations about our algorithm are in place . First , compared to typical ways of implementing Thompson Sampling with neural networks , NeuralTS samples from the posterior distribution of the scalar reward , instead of the network parameters . It is therefore simpler and more efficient , as the number of parameters in practice can be large . Second , the algorithm maintains the posterior distributions related to parameters of all layers of the network , as opposed to the last layer only ( Riquelme et al. , 2018 ) . This difference is crucial in our regret analysis . It allows us to build a connection between Algorithm 1 and recent work about deep learning theory ( Allen-Zhu et al. , 2018 ; Cao & Gu , 2019 ) , in order to obtain theoretical guarantees as will be shown in the next section . Third , different from linear or kernelized TS ( Agrawal & Goyal , 2013 ; Chowdhury & Gopalan , 2017 ) , whose posterior can be computed in closed forms , NeuralTS solves a non-convex optimization problem ( 2.3 ) by gradient descent . This difference requires additional techniques in the regret analysis . Moreover , stochastic gradient descent can be used to solve the optimization problem with a similar theoretical guarantee ( Allen-Zhu et al. , 2018 ; Du et al. , 2018 ; Zou et al. , 2019 ) . For simplicity of exposition , we will focus on the exact gradient descent approach . 3 REGRET ANALYSIS . In this section , we provide a regret analysis of NeuralTS . We assume that there exists an unknown reward function h such that for any 1 ≤ t ≤ T and 1 ≤ k ≤ K , rt , k = h ( xt , k ) + ξt , k , with |h ( xt , k ) | ≤ 1 where { ξt , k } forms an R-sub-Gaussian martingale difference sequence with constant R > 0 , i.e. , E [ exp ( λξt , k ) |ξ1 : t−1 , k , x1 : t , k ] ≤ exp ( λ2R2 ) for all λ ∈ R. Such an assumption on the noise sequence is widely adapted in contextual bandit literature ( Agrawal & Goyal , 2013 ; Bubeck & CesaBianchi , 2012 ; Chowdhury & Gopalan , 2017 ; Chu et al. , 2011 ; Lattimore & Szepesvári , 2020 ; Valko et al. , 2013 ) . Next , we provide necessary background on the neural tangent kernel ( NTK ) theory ( Jacot et al. , 2018 ) , which plays a crucial role in our analysis . In the analysis , we denote by { xi } TKi=1 the set of observed contexts of all arms and all rounds : { xt , k } 1≤t≤T,1≤k≤K where i = K ( t− 1 ) + k. Definition 3.1 ( Jacot et al . ( 2018 ) ) . Define H̃ ( 1 ) i , j = Σ ( 1 ) i , j = 〈x i , xj〉 , A ( l ) i , j = ( Σ ( l ) i , i Σ ( l ) i , j Σ ( l ) i , j Σ ( l ) j , j ) , Σ ( l+1 ) i , j = 2E ( u , v ) ∼N ( 0 , A ( l ) i , j ) max { u , 0 } max { v , 0 } , H̃ ( l+1 ) i , j = 2H̃ ( l ) i , jE ( u , v ) ∼N ( 0 , A ( l ) i , j ) 1 ( u ≥ 0 ) 1 ( v ≥ 0 ) + Σ ( l+1 ) i , j . Then , H = ( H̃ ( L ) + Σ ( L ) ) /2 is called the neural tangent kernel matrix on the context set . The NTK technique builds a connection between deep neural networks and kernel methods . It enables us to adapt some complexity measures for kernel methods to describe the complexity of the neural network , as given by the following definition . Definition 3.2 . The effective dimension d̃ of matrix H with regularization parameter λ is defined as d̃ = log det ( I + H/λ ) log ( 1 + TK/λ ) . Remark 3.3 . The effective dimension is a metric to describe the actual underlying dimension in the set of observed contexts , and has been used by Valko et al . ( 2013 ) for the analysis of kernel UCB . Our definition here is adapted from Yang & Wang ( 2019 ) , which also considers UCB-based exploration . Compared with the maximum information gain γt used in Chowdhury & Gopalan ( 2017 ) , one can verify that their Lemma 3 shows that γt ≥ log det ( I + H/λ ) /2 . Therefore , γt and d̃ are of the same order up to a ratio of 1/ ( 2 log ( 1 + TK/λ ) ) . Furthermore , d̃ can be upper bounded if all contexts xi are nearly on some low-dimensional subspace of the RKHS space spanned by NTK ( Appendix D ) . We will make a regularity assumption on the contexts and the corresponding NTK matrix H. Assumption 3.4 . Let H be defined in Definition 3.1 . There exists λ0 > 0 , such that H λ0I . In addition , for any t ∈ [ T ] , k ∈ [ K ] , ‖xt , k‖2 = 1 and [ xt , k ] j = [ xt , k ] j+d/2 . The assumption that the NTK matrix is positive definite has been considered in prior work on NTK ( Arora et al. , 2019 ; Du et al. , 2018 ) . The assumption on context xt , a ensures that the initial output of neural network f ( x ; θ0 ) is 0 with the random initialization suggested in Algorithm 1 . The condition on x is easy to satisfy , since for any context x , one can always construct a new context x̃ as [ x/ ( √ 2‖x‖2 ) , x/ ( √ 2‖x‖2 ) ] > . We are now ready to present the main result of the paper : Theorem 3.5 . Under Assumption 3.4 , set the parameters in Algorithm 1 as λ = 1 + 1/T , ν = B + R √ d̃ log ( 1 + TK/λ ) + 2 + 2 log ( 1/δ ) where B = max { 1/ ( 22e √ π ) , √ 2h > H−1h } with h = ( h ( x1 ) , . . . , h ( xTK ) ) > , and R is the sub-Gaussian parameter . In line 9 of Algorithm 1 , set η = C1 ( mλ+mLT ) −1 and J = ( 1 +LT/λ ) ( C2 + log ( T 3Lλ−1 log ( 1/δ ) ) ) /C1 for some positive constants C1 , C2 . If the network width m satisfies : m ≥ poly ( λ , T , K , L , log ( 1/δ ) , λ−10 ) , then , with probability at least 1− δ , the regret of Algorithm 1 is bounded as RT ≤ C3 ( 1 + cT ) ν √ 2λL ( d̃ log ( 1 + TK ) + 1 ) T + ( 4 + C4 ( 1 + cT ) νL ) √ 2 log ( 3/δ ) T + 5 , where C3 , C4 are some positive absolute constants , and cT = √ 4 log T + 2 logK . Remark 3.6 . The definitionB in Theorem 3.5 is inspired by the RKHS norm of the reward function defined in Chowdhury & Gopalan ( 2017 ) . It can be verified that when the reward function h belongs to the function space induced by NTK , i.e. , ‖h‖H < ∞ , we have √ h > H−1h ≤ ‖h‖H according to Zhou et al . ( 2019 ) , which suggests that B ≤ max { 1/ ( 22e √ π ) , √ 2‖h‖H } . Remark 3.7 . Theorem 3.5 implies the regret of NeuralTS is on the order of Õ ( d̃T 1/2 ) . This result matches the state-of-the-art regret bound in Chowdhury & Gopalan ( 2017 ) ; Agrawal & Goyal ( 2013 ) ; Zhou et al . ( 2019 ) ; Kveton et al . ( 2020 ) . Remark 3.8 . In Theorem 3.5 , the requirement of m is specified in Condition 4.1 and the proof of Theorem 3.5 , which is a high-degree polynomial in the time horizon T , number of layers L and number of actions K. However , in our experiments , we can choose reasonably small m ( e.g. , m = 100 ) to obtain good performance of NeuralTS . See Appendix A.1 for more details . This discrepancy between theory and practice is due to the limitation of current NTK theory ( Du et al. , 2018 ; Allen-Zhu et al. , 2018 ; Zou et al. , 2019 ) . Closing the gap is a venue for future work . Remark 3.9 . Theorem 3.5 suggests that we need to know T before we run the algorithm in order to set m. When T is unknown , we can use the standard doubling trick ( See e.g. , Cesa-Bianchi & Lugosi ( 2006 ) ) to set m adaptively . In detail , we decompose the time interval ( 0 , +∞ ) as a union of non-overlapping intervals [ 2s , 2s+1 ) . When 2s ≤ t < 2s+1 , we restart NeuralTS with the input T = 2s+1 . It can be verified that similar Õ ( d̃ √ T ) regret still holds .

Summary: The paper proposes a novel Thompson sampling algorithm for neural networks which can be applied to any arbitrary, bounded reward function. While existing works apply Thompson sampling (TS) to neural networks in a heuristic way (e.g., sampling parameters in the last layer only), this algorithm considers the posterior distribution of all the parameters and is the first to provide a theoretical, tight regret upper bound. The work builds on the neural tangent kernel theory, which enables the use of techniques developed for linear reward functions (Agrawal and Goyal, 2013). Actually, the paper is analogous to the work of Zhou et al. (2020) which proposed the NeuralUCB algorithm by combining the neural tangent kernel theory and the Linear UCB methodology (Abbasi-Yadkori et al., 2011).