Impact of Representation Learning in Linear Bandits

√ kN + √ dkNT ) regret , where N is the number of rounds we play for each bandit . When T is sufficiently large , our algorithm significantly outperforms the naive algorithm ( playing T bandits independently ) that achieves Õ ( T √ dN ) regret . We also provide an Ω ( T √ kN + √ dkNT ) regret lower bound , showing that our algorithm is minimax-optimal up to poly-logarithmic factors . Furthermore , we extend our algorithm to the infinite-action setting and obtain a corresponding regret bound which demonstrates the benefit of representation learning in certain regimes . We also present experiments on synthetic and realworld data to illustrate our theoretical findings and demonstrate the effectiveness of our proposed algorithms . 1 INTRODUCTION . This paper investigates the benefit of using representation learning for sequential decision-making problems . Representation learning learns a joint low-dimensional embedding ( feature extractor ) from different but related tasks and then uses a simple function ( often a linear one ) on top of the embedding ( Baxter , 2000 ; Caruana , 1997 ; Li et al. , 2010 ) The mechanism behind is that since the tasks are related , we can extract the common information more efficiently than treating each task independently . Empirically , representation learning has become a popular approach for improving sample efficiency across various machine learning tasks ( Bengio et al. , 2013 ) . In particular , recently , representation learning has become increasingly more popular in sequential decision-making problems ( Teh et al. , 2017 ; Taylor & Stone , 2009 ; Lazaric & Restelli , 2011 ; Rusu et al. , 2015 ; Liu et al. , 2016 ; Parisotto et al. , 2015 ; Higgins et al. , 2017 ; Hessel et al. , 2019 ; Arora et al. , 2020 ; D ’ Eramo et al. , 2020 ) . For example , many sequential decision-making tasks share the same environment but have different reward functions . Thus a natural approach is to learn a succinct representation that describes the environment and then make decisions for different tasks on top of the learned representation . While representation learning is already widely applied in sequential decision-making problems empirically , its theoretical foundation is still limited . One important problem remains open : When does representation learning provably improve efficiency of sequential decision-making problems ? We take a step to characterize the benefit of representation learning in sequential decision-making problems . We tackle the above problem in the linear bandits setting , one of the most fundamental and popular settings in sequential decision-making problems . This model is widely used in applications as such clinical treatment , manufacturing process , job scheduling , recommendation systems , etc ( Dani et al. , 2008 ; Chu et al. , 2011 ) . We study the multi-task version of linear bandits , which naturally models the scenario where one needs to deal with multiple different but closely related sequential decision-making problems concurrently . We will mostly focus on the finite-action setting . Specifically , we have T tasks , each of which is governed by an unknown linear coefficient θt ∈ Rd . At the n-th round , for each task t ∈ [ T ] , the player chooses an action an , t that belongs to a finite set , and receive a reward rn , t with expectation E rn , t = 〈θt , xn , t , an , t〉 where xn , t , an , t represents the context of action an , t . For this problem , a straightforward approach is to treat each task independently , which leads to Õ ( T √ dN ) 1 total regret . Can we do better ? Clearly , if the tasks are independent , then by the classical Ω ( √ dN ) per task lower bound for linear bandit , it is impossible to do better . We investigate how representation learning can help if the tasks are related . Our main assumption is the existence of an unknown linear feature extractorB ∈ Rd×k with k d and a set of linear coefficients { wt } Tt=1 such that θt = Bwt . Under this assumption , the tasks are closely related as B is a shared linear feature extractor that maps the raw contexts xn , t , a ∈ Rd to a low-dimensional embedding B > xn , t , a ∈ Rk . In this paper , we focus on the regime where k d , N , T . This regime is common in real-world problems , e.g. , computer vision , where the input dimension is high , the number of data is large , many task are related , and there exists a low-dimension representation among these tasks that we can utilize . Problems with similar assumptions have been studied in the supervised learning setting ( Ando & Zhang , 2005 ) . However , to our knowledge , this formulation has not been studied in the bandit setting . Our Contributions We give the first rigorous characterization on the benefit of representation learning for multi-task linear bandits . Our contributions are summarized below . • We design a new algorithm for the aforementioned problem . Theoretically , we show our algorithm incurs Õ ( √ dkTN + T √ kN ) total regret in N rounds for all T tasks . Therefore , our algorithm outperforms the naive approach with O ( T √ dN ) regret . To our knowledge , this is the first theoretical result demonstrating the benefit of representation learning for bandits problems . • To complement our upper bound , we also provide an Ω ( √ dkTN+T √ kN ) lower bound , showing our regret bound is tight up to polylogarithmic factors . • We further design a new algorithm for the infinite-action setting , which has a regret Õ ( d1.5k √ TN + kT √ N ) , which outperforms the naive approach with O ( Td √ N ) regret in the regime where T = Ω̃ ( dk2 ) . • We provide simulations and an experiment on MNIST dataset to illustrate the effectiveness of our algorithms and the benefits of representation learning . Organization This paper is organized as follows . In Section 2 , we discuss related work . In Section 3 , we introduce necessary notation , formally set up our problem , and describe our assumptions . In Section 4 , we present our main algorithm for the finite-action setting and its performance guarantee . In Section 5 , we describe our algorithm and its theoretical guarantee for the infinite-action setting . In Section 6 , we provide simulation studies and real-world experiments to validate the effectiveness of our approach . We conclude in Section 7 and defer all proofs to the Appendix . 2 RELATED WORK . Here we mainly focus on related theoretical results . We refer readers to Bengio et al . ( 2013 ) for empirical results of using representation learning . For supervised learning , there is a long line of works on multi-task learning and representation learning with various assumptions ( Baxter , 2000 ; Ando & Zhang , 2005 ; Ben-David & Schuller , 2003 ; Maurer , 2006 ; Cavallanti et al. , 2010 ; Maurer et al. , 2016 ; Du et al. , 2020 ; Tripuraneni et al. , 2020 ) . 1Õ ( · ) omits logarithmic factors . All these results assumed the existence of a common representation shared among all tasks . However , this assumption alone is not sufficient . For example , Maurer et al . ( 2016 ) further assumed every task is i.i.d . drawn from an underlying distribution . Recently , Du et al . ( 2020 ) replaced the i.i.d . assumption with a deterministic assumption on the input distribution . Finally , it is worth mentioning that Tripuraneni et al . ( 2020 ) gave the method-of-moments estimator and built the confidence ball for the feature extractor , which inspired our algorithm for the infinite-action setting . The benefit of representation learning has been studied in sequential decision-making problems , especially in reinforcement learning domains . D ’ Eramo et al . ( 2020 ) showed that representation learning can improve the rate of approximate value iteration algorithm . Arora et al . ( 2020 ) proved that representation learning can reduce the sample complexity of imitation learning . Both works require a probabilistic assumption similar to that in ( Maurer et al. , 2016 ) and the statistical rates are of similar forms as those in ( Maurer et al. , 2016 ) . We remark that representation learning is also closely connected to meta-learning ( Schaul & Schmidhuber , 2010 ) . Raghu et al . ( 2019 ) empirically suggested that the effectiveness of metalearning is due to its ability to learn a useful representation . There is a line of works that analyzed the theoretical properties of meta-learning ( Denevi et al. , 2019 ; Finn et al. , 2019 ; Khodak et al. , 2019 ; Lee et al. , 2019 ; Bertinetto et al. , 2018 ) . We also note that there are analyses for other representation learning schemes ( Arora et al. , 2019 ; McNamara & Balcan , 2017 ; Galanti et al. , 2016 ; Alquier et al. , 2016 ; Denevi et al. , 2018 ) . Linear bandits ( stochastic linear bandits / linearly parameterized bandits / contextual linear bandits ) have been studied in recent years ( Auer , 2002 ; Dani et al. , 2008 ; Rusmevichientong & Tsitsiklis , 2010 ; Abbasi-Yadkori et al. , 2011 ; Chu et al. , 2011 ; Li et al. , 2019a ; b ) . The studies are divided into two branches according to whether the action set is finite or infinite . For the finite-action setting , Θ̃ ( √ dN ) has been shown to be the near-optimal regret bound ( Chu et al. , 2011 ; Li et al. , 2019a ) , and for the infinite-action setting , Θ̃ ( d √ N ) regret bound has been shown to be near-optimal ( Dani et al. , 2008 ; Rusmevichientong & Tsitsiklis , 2010 ; Li et al. , 2019b ) . Some previous work studied the impact of low-rank structure in linear bandit . Lale et al . ( 2019 ) studied a setting where the context vectors share a low-rank structure . Specifically , in their setting , the context vectors consist of two parts , i.e . x̂ = x+ψ , so that x is from a hidden low-rank subspace and ψ is i.i.d . drawn from an isotropic distribution . Jun et al . ( 2019 ) and Lu et al . ( 2020 ) studied the bilinear bandits with low-rank structure . In their setting , the player chooses two actions x , y and receives the stochastic reward with mean x > Θy , where Θ is an unknown low-rank bilinear form . The algorithms proposed in the aforementioned papers share some similarities with our Algorithm 2 for our infinite-action setting , in that both used Davis-Kahan theorem to recover and exploit the low-rank structure . Some previous work proposed multi-task bandits with different settings . Deshmukh et al . ( 2017 ) proposed a setting under the contextual bandit framework . They assumed similarities among arms . Bastani et al . ( 2019 ) studied a setting where the coefficients of the tasks were drawn from a gaussian distribution fixed across tasks and proposed an algorithm based on Thompson sampling . Soare et al . ( 2018 ) proposed a setting where tasks were played one by one sequentially and the coefficients of the tasks were near in ` 2 distance . In our setting , the tasks are played simultaneously and the coefficients share a common linear feature extractor .

Summary: This paper studies the benefits of learning a low-rank feature extractor in multi-task linear bandits. Specifically, the paper studies the setting where an unknown common linear feature extractor $B \in R^{d \times k}$ maps the original $d$-dimensional contexts $x$ to a $k$-dimensional representation. Essentially, for multi-task linear bandit problem $r_t = \theta_t^T x$, this paper assumes the matrix of model parameters $\Theta \in R^{d\times T}$ is low-rank and be factorized as $\Theta = BW$ with rank k. The paper proposed algorithms to estimate both $B$ and $W$ in finite actions setting and infinite actions setting. In finite action setting, the proposed solution is a greedy algorithm while in infinite actions setting, the proposed solution is a explore-the-commit method. Theoretical result shows that the regret is $O(T\sqrt{kN} + \sqrt{dkTN})$ and matches the lower bound. Simulation result shows that the algorithm outperforms baselines running independent linear bandit algorithms.