Year: 2021

Abstract:
We study how representation learning can improve the efficiency of bandit problems. We study the setting where we play T linear bandits with dimension d concurrently, and these T bandit tasks share a common k( d) dimensional linear representation. For the finite-action setting, we present a new algorithm which achieves Õ(T √ 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 Õ(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.

text: √ 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARIES
text: Notation. We use bold lowercases for vectors and bold uppercases for matrices. For any positive integer n, we use [n] to denote the set of integers {1, 2, . . . , n}. For any vector x, we use ‖x‖ to denote its `2 norm. For a matrix A, we use ‖A‖ to denote the 2-norm of A, ‖A‖F to denote the Frobenius norm, and ‖A‖max = maxi,j |Aij | to denote the max-norm. For two expressions α, β > 0, we denote α . β if there is a numerical constant c > 0 such that α ≤ cβ. We denote α & β if β . α.
Problem Setup. Let d be the ambient dimension and k(≤ d) be the representation dimension. In total, we have T tasks and we play each task concurrently for N rounds. Each task t ∈ [T ] has an
unknown vector θt ∈ Rd. At each round n ∈ [N ], the player chooses action an,t ∈ An,t for each task t ∈ [T ] where An,t is the action set at round n for the task t. After the player commits to a batch of actions {an,t}t∈[T ], it receives a batch of rewards {rn,t}t∈[T ], where we assume rn,t = 〈an,t,θt〉 + εn,t. Here we assume the noise εn,t are independent 1-subGaussian random variables, which is a standard assumption in the literature.
We use the total expected regret to measure the performance of our algorithm. When we have N rounds and T tasks, it is defined as RN,T = ∑N n=1 ∑T t=1 maxa∈An,t〈a,θt〉− 〈an,t,θt〉. When the action set is finite, we assume that all An,t have the same size K, i.e. |An,t| ≡ K. Furthermore, we writeAn,t = {xn,t,1, . . . ,xn,t,K}. Besides, we interchangeably use the number an,t ∈ [K] and the vector an,t = xn,t,an,t ∈ Rd to refer to the same action. Assumptions. Our main assumption is the existence of a common linear feature extractor. Assumption 1 (Common Feature Extractor). There exists a linear feature extractor B ∈ Rd×k and a set of linear coefficients {wt}Tt=1 such that the expected reward of the t-th task at the n-th round satisfies E[rt,n] = 〈wt,B>xn,t,an,t〉.
For simplicity, we let W = [w1, . . . ,wT ]. Assumption 1 implies that Θ , [θ1, . . . ,θT ] = BW . Note this assumption is in a sense necessary to guarantee the effectiveness of representation learning because without it one cannot hope that representation learning helps.
In this paper, we mostly focus on the finite-action setting. We put the following assumption on action sets. Assumption 2. Marginally, for every n ∈ [N ], t ∈ [T ], a ∈ [K], the contexts satisfy xn,t,a ∼ N (0,Σt) such that λmax(Σt) ≤ O(1/d) and λmin(Σt) ≥ Ω(1/d).
With this assumption, we have an unknown covariance matrix Σt for each task. At each round, the actions of the t-th task are sampled from a Gaussian distribution with covariance Σt. This is a prototypical setting for theoretical development on linear bandits with finite actions (See e.g., Han et al. (2020)). At a population level, each one of the K actions is equally good but being able to select different actions based on the realized contexts allows the player to gain more reward.
We will also study the infinite-action setting. We first state our assumption about the action sets. Assumption 3 (Ellipsoid Action Set). We assume An,t = At = {x>Q−1t x ≤ 1 : x ∈ Rd} is an ellipsoid with λmin(Qt) ≥ λ0 = Ω(1).
The first assumption states that each action set is an ellipsoid that covers all directions. This is a standard assumption, e.g., see Rusmevichientong & Tsitsiklis (2010).
In this setting, we will also need to put some additional assumptions on the underlying parameters B andW . Assumption 4 (Diverse Source Tasks). We assume that λmin( 1TWW >) ≥ νk , where ν = Ω(1).
This assumption roughly states that the underlying linear coefficients {wt}Tt=1 equally spans all directions in Rk. This is a common assumption in representation learning literature that enables us to learn the linear feature extractor (Du et al., 2020; Tripuraneni et al., 2020). For example, the assumption holds with high probability when wi is uniformly chosen from the sphere Sk−1. Assumption 5. We assume ‖wt‖ ≥ ω = Ω(1).
This is a normalization assumption on the linear coefficients.

heading: 4 MAIN RESULTS FOR FINITE-ACTION SETTING
text: In this section focus on the finite-action setting. The pseudo-code is listed in Algorithm 1. Our algorithm uses a doubling schedule rule (Gao et al., 2019; Simchi-Levi & Xu, 2019; Han et al., 2020; Ruan et al., 2020). We only update our estimation of θ after an epoch is finished, and we only use samples collected within the epoch. In Line 5, we solve an empirical `2-risk minimization problem on the data collected in the last epoch to estimate the feature extractor B and linear predictors
Algorithm 1: MLinGreedy: Multi-task Linear Bandit with Finite Actions
1: Let M = dlog2 log2Ne,G0 = 0,GM = N,Gm = N1−2 −m
for 1 ≤ m ≤M − 1, let θ̂0,t ← 0; 2: for m← 1, . . . ,M do 3: for n← Gm−1 + 1, . . . ,Gm do 4: For each task t ∈ [T ]: choose action an,t = arg maxa∈[K] x>n,t,aθ̂m−1,t; 5: Compute B̂Ŵ ← arg min
B∈Rd×k,W∈Rk×T
∑Gm n=Gm−1+1 ∑T t=1[x > n,t,an,tBwt − rn,t] 2;
6: For each task t ∈ [T ]: let θ̂m,t = B̂ŵt;
W, similar to Du et al. (2020). Given estimated feature extractor B̂ and linear predictors Ŵ , we compute our estimated linear coefficients of task t as θ̂t , B̂ŵt in Line 6. For choosing actions, for each task, we use a greedy rule, i.e., we choose the action that maximizes the inner product with our estimated θ (cf. Line 4).
The following theorem gives an upper bound on the regret of Algorithm 1. Theorem 1 (Regret of Algorithm 1). Suppose K,T ≤ poly(d) and N ≥ d2. Under Assumption 1 and Assumption 2, the expected regret of Algorithm 1 is upper bounded by
E[RN,T ] = Õ(T √ kN + √ dkNT ).
There are two terms in Theorem 1, and we interpret them separately. The first term Õ(T √ kN) represents the regret for playing T independent linear bandits with dimension k for N rounds. This is the regret we need to pay even if we know the optimal feature extractor, with which we can reduce the original problem to playing T independent linear bandits with dimension k (recallwt are different for different tasks). The second term Õ( √ dkNT ) represents the price we need to pay to learn the feature extractorB. Notably, this term shows we are using data across all tasks to learnB as this term scales with √ NT .
Now comparing with the naive strategy that plays T independent d-dimensional linear bandits with regret Õ(T √ dN), our upper bound is smaller as long as T = Ω(k). Furthermore, when T is large, our bound is significantly stronger than Õ(T √ dN), especially when k d. To our knowledge, this is the first formal theoretical result showing the advantage of representation learning for bandit problems. We remark that requiring T = Ω(k) is necessary. One needs at least k tasks to recover the span ofW , so only in this regime representation learning can help.
Our result also puts a technical requirement on the scaling K,T ≤ poly(d) and N ≥ d2. These are conditions that are often required in linear bandits literature. The first condition ensures that K and T are not too large, so we need not characterize log(KT ) factors in regret bound. If they are too large, e.g. K,T ≥ Ω(ed), then we would have logKT = O(d) and we could no longer omit the logarithmic factors in regret bounds. The second condition ensures one can at least learn the linear coefficients up to a constant error. See more discussions in (Han et al., 2020, Section 2.5).
While Algorithm 1 is a straightforward algorithm, the proof of Theorem 1 requires a combination of representation learning and linear bandit techniques. First, we prove the in-sample guarantee of representation learning, as done in Lemma 2. Second, we exploit Assumption 2 to show that the learned parameters could extrapolate well on new contexts, as shown in Lemma 4. The regret analysis then follows naturally. We defer the proof of Theorem 1 to Appendix A.
The following theorem shows that Algorithm 1 is minimax optimal up to logarithmic factors. Theorem 2 (Lower Bound for Finite-Action Setting). Let A denote an algorithm and I denote a finite-actioned multi-task linear bandit instance that satisfies Assumption 1 and Assumption 2. Then for any N,T, d, k ∈ Z+ with k ≤ d, k ≤ T , we have
inf A sup I
E[RN,TA,I ] = Ω ( T √ kN + √ dkNT ) . (1)
Theorem 2 has the exactly same two terms as in Theorem 1. This confirms our intuition that the two prices to pay are real: 1) playing T independent k-dimensional linear bandits and 2) learning
Algorithm 2: E2TC: Explore-Explore-Then-Commit Input: N : total number of rounds , N1: number of rounds for stage 1 , N2: number of rounds
for stage 2 1: Stage 1: Estimating Linear Feature Extractor with Method-of-Moments; 2: for ∀t ∈ [T ], n ∈ [N1] do 3: Play xn,t ∼ Unif(λ0 · Sd−1) and receive reward rn,t; 4: Compute M̂← 1N1T ∑N1 n=1 ∑T t=1 r 2 n,txn,tx > n,t; 5: Let B̂DB̂> ← top-k singular value decomposition of M̂ . Denote b̂i the i-th column of B̂; 6: Stage 2: Estimating Optimal Actions on Low-dimensional Space ; 7: for ∀t ∈ [T ], i ∈ [k] do 8: Play vt,i , √ λ0b̂i for N2/k times and receive rewards {rn,t}N1+iN2/kn=N1+(i−1)N2/k+1;
9: for ∀t ∈ [T ] do 10: Estimate ŵt ← arg minw∈Rk 12N2 ∑N1+N2 n=N1+1
[〈xn,t, B̂w〉 − rn,t]2; 11: Let θ̂t = B̂ŵt; 12: Stage 3: Committing to Near-optimal Actions; 13: for ∀t ∈ [T ], n = N1 +N2 + 1, . . . , N do 14: Play an,t ← arg maxa∈At〈a, θ̂t〉 and receive reward rn,t;
the d× k-dimensional feature extractor. We defer the proof of Theorem 2 to Appendix B. At a high level, we separately prove the two terms in the lower bound. The first term is established by the straightforward observation that our multi-task linear bandit problem is at least as hard as solving T independent k-dimensional linear bandits. The second term is established by the observation that multi-task linear bandit can be seen as solving k independent d-dimensional linear bandits, each has NT/k rounds. Note that the observation would directly imply the regret lower bound k √ d(NT/k) = √ dkNT , which is exactly the second term. To our knowledge, this lower bound is also the first one of its kind for multi-task sequential decision-making problems. We believe our proof framework can be used in proving lower bounds for other related problems.

heading: 5 EXTENSION TO INFINITE-ACTION SETTING
text: In this section we present and analyze an algorithm for the infinite-action setting. Pseudo-code is listed in Algorithm 2. Our algorithm has three stages, and we explain each step below.
Stage 1: Estimating Linear Feature Extractor with Method-of-Moments. The goal of the first stage is to estimate the linear feature extractor B. Our main idea to view this problem as a lowrank estimation problem for which we use a method-of-moments estimator. In more detail, we first sample each xn,t ∼ Unif[λ0 · Sd] for N1 times. We can use this sampling scheme because the action set is an ellipsoid. Note this sampling scheme has a sufficient coverage on all directions, which help us estimate B. Next, we compute the empirical weighted covariance matrix M̂ =
1 N1T ∑N1 n=1 ∑T t=1 r 2 n,txn,tx > n,t. To proceed, we compute the singular value decomposition of M̂ and keep its top-k column space as our estimated linear feature extractor B̂, which is a sufficiently accurate estimator (cf. Theorem 5).
Stage 2: Estimating Optimal Actions on Low-Dimensional Space. In the second stage, we use our estimated linear feature extractor to refine our search space for the optimal actions. Specifically, we denote B̂ = [b̂1, . . . , b̂k] and for i ∈ [k] and t ∈ [T ], we let vt,i = √ λ0b̂i. Under Assumption 3, we know vt,i ∈ At for all i ∈ [k]. Therefore, we can choose vt,i to explore. Technically, our choice of actions vt,i also guarantees a sufficient cover in the sense that λmin (∑k i=1 B̂ >vt,iv > t,iB̂ ) ≥ λ0. In particular, this coverage is on a low-dimensional space instead of the original ambient space.
The second stage has N2 rounds and on the t-th task, we just play each vt,i for N2/k rounds. After that, we use linear regression to estimate wt for each task. Given the estimation ŵt, we can obtain an estimation to the true linear coefficient θ̂t , B̂ŵt.
Stage 3: Committing to Near-Optimal Actions. After the second stage, we have an estimation θ̂t for each task. For the remaining (N − N1 − N2) rounds, we just commit to the optimal action indicated by our estimations. Specifically, we play the action an,t ← arg maxa∈At〈at, θ̂t〉 for round n = N1 +N2 + 1, . . . , N .
The following theorem characterizes the regret of our algorithm. Theorem 3 (Regret of Algorithm 2). If we choose N1 = c1d1.5k √ N T and N2 = c2k √ N for some constants c1, c2 > 0. The regret of Algorithm 2 is upper bounded by
E[RN,T ] = Õ(d1.5k √ TN + kT √ N).
The first term represents the regret incurred by estimating the linear feature extractor. Notably, the term scales with √ TN , which means we utilize all TN data here. The first term scales with d1.5k, which we conjecture is unavoidable at least by our algorithm. The second term represents playing T independent k-dimensional infinite-action linear bandits.
Notice that if one uses a standard algorithm, e.g. the PEGE algorithm (Rusmevichientong & Tsitsiklis, 2010), to play T tasks independently, one can achieve an Õ(dT √ N) regret. Comparing with
this bound, our bound’s second term is always smaller and the first is smaller when T = Ω̃(dk2). This demonstrates that more tasks indeed help us learn the representation and reduce the regret.
We complement our upper bound with a lower bound below. This theorem suggests our second term is tight but there is still a gap in the first term. We leave it as an open problem to design new algorithm to match the lower bound or to prove a stronger lower bound.
Theorem 4 (Lower Bound for Infinite-Action Setting). LetA denote an algorithm and I denote an infinite-action multi-task linear bandit instance that satisfies Assumption 1, Assumption 3, Assumption 4, Assumption 5. Then for any N,T, d, k ∈ Z+ with k ≤ d, k ≤ T , we have
inf A sup I
E[RN,TA,I ] = Ω ( d √ kNT + kT √ N ) . (2)

heading: 6 EXPERIMENTS
text: In this section, we use synthetic data and MNIST data to illustrate our theoretical findings and demonstrate the effectiveness of our Algorithm for the finite-action setting. We also have simulation studies for the infinite-action setting, which we defer to Appendix F. The baseline is the naive algorithm which plays T tasks independently, and for each task, thie algorithm uses linear regression to estimate θt and choose the action greedily according to the estimated θt.

heading: 6.1 SYNTHETIC DATA
text: Setup. The linear feature extractor B is uniformly drawn from the set of d × k matrices with orthonormal columns.2 Each linear coefficient wt is uniformly chosen from the k-dimensional sphere. The noises are i.i.d. Gaussian: εn,t,a = N (0, 1) for every n ∈ [N ], t ∈ [T ], a ∈ [K]. We fix K = 5 and N = 10000 for all simulations on finite-action setting. We vary k, d and T to compare Algorithm 1 and the naive algorithm.
Results and Discussions. We present the simulation results in Figure 1 and Figure 2. We emphasize that the y-axis in our figures corresponds to the regret per task, which is defined as RN,T /T . We fix K = 5, N = 10000. These simulations verify our theoretical findings. First, as the number of tasks increases, the advantage of our algorithm increases compared to the naive algorithm. Secondly, we notice that as k becomes larger (relative to d), the advantage of our algorithm becomes smaller. This can be explained by our theorem that as k increases, our algorithm pays more regret, whereas the naive algorithm’s regret does not depend on k.

heading: 6.2 FINITE-ACTION LINEAR BANDITS FOR MNIST
text: Setup. We create a linear bandits problem on MNIST data (LeCun et al., 2010) to illustrate the effectiveness of our algorithm on real-world data. We fix K = 2 and create T = ( 10 2 ) tasks and each task is parameterized by a pair (i, j), where 0 ≤ i < j ≤ 9. We use Di to denote the set of MNIST images with digit i. At each round n ∈ [N ], for each task (i, j), we randomly choose one picture
2We uniformly (under the Haar measure) choose a random element from the orthogonal group O(d) and uniformly choose k of its columns to generate B.
fromDi and one fromDj , then we present those two pictures to the algorithm and assign the picture with larger digit with reward 1 and the other with reward 0. The algorithm is now required to select an image (action). We again compare our algorithm with the naive algorithm.
Results and Discussions. The experimental results are displayed in Figure 3 for T = 10 (done by constructing tasks with first five digits) and Figure 4 for T = 45. We observe for both T = 10 and T = 45, our algorithm significantly outperforms the naive algorithm for all k. Interestingly, unlike our simulations, we find the advantage of our algorithm does not decrease as we increase k. We believe the reason is the optimal predictor is not exactly linear, and we need to develop an agnostic theory to explain this phenomenon, which we leave as a future work.

heading: 7 CONCLUSION
text: We initiate the study on the benefits of representation learning in bandits. We proposed new algorithms and demonstrated that in the multi-task linear bandits, if all tasks share a common linear feature extractor, then representation learning provably reduces the regret. We demonstrated empirical results to corroborate our theory. An interesting future direction is to generalize our results to general reward function classes (Li et al., 2017; Agrawal et al., 2019). In the following, we discuss some future directions.
Adversarial Contexts For the finite-action setting, we assumed the context are i.i.d. sampled from a Gaussian distribution. In the bandit literature, there is a large body on developing lowregret algorithms for the adversarial contexts setting. We leave it as an open problem to develop a algorithm with an Õ(T √ kN + √ dkNT ) upper bound or show this bound is not possible in the adversarial contexts setting. One central challenge for the upper bound is that existing analyses for multi-task representation learning requires i.i.d. inputs even in the supervised learning setting. Another challenge is how to develop a confidence interval for an unseen input in the multi-task linear bandits setting. This confidence interval should utilize the common feature extractor and is tighter than the standard confidence interval for linear bandits, e.g. LinUCB.
Robust Algorithm The current approach is tailored to the assumption that there exists a common feature extractor. One interesting direction is to develop a robust algorithm. For example, consider the scenario where whether there exists a common feature extractor is unknown. We want to develop an algorithm with regret bound as in this paper when the common feature extractor exists and gracefully degrades to the regret of T independent linear bandits when the common feature extractor does not exist.
General Function Approximation In this paper, we focus on linear bandits. In the bandits literature, sublinear regret guarantees have been proved for more general reward function classes beyond the linear one (Li et al., 2017; Agrawal et al., 2019). Similarly, in the supervised representation learning literature, general representation function classes have also been studied (Maurer et al., 2016; Du et al., 2020). An interesting future direction is to merge these two lines of research by developing provably efficient algorithms for multi-task bandits problems where general function class is used for representation.

heading: ACKNOWLEDGMENTS
text: The authors would like to thank the anonymous reviewers for their comments and suggestions on our paper. WH is supported by NSF, ONR, Simons Foundation, Schmidt Foundation, Amazon Research, DARPA and SRC. JDL acknowledges support of the ARO under MURI Award W911NF-11-1-0303, the Sloan Research Fellowship, and NSF CCF 2002272.

heading: A PROOF OF THEOREM 1
text: Lemma 1 (General Hoeffding’s inequality, Vershynin (2018), Theorem 2.6.2). Let X1, . . . , Xn be independent random variables such that E[Xi] = 0 and Xi is σi-sub-Gaussian. Then there exists a constant c > 0, such that for any δ > 0, we have
Pr ∣∣∣∣∣ n∑ i=1 Xi ∣∣∣∣∣ ≥ c √√√√ n∑ i=1 σ2i log(1/δ)  ≤ δ. (3) Lemma 2. Let T be the number of tasks and N0 be the number of samples. For every (n, t) ∈ [N0] × [T ], let xn,t ∈ Rd be fixed vectors and let yn,t = x>n,tθt + εn,t, where θt ∈ Rd is a vector and εn,t is an independent 1-sub-Gaussian variable. Let
B̂Ŵ = arg min B̂∈Rd×k,Ŵ∈Rk×T N0∑ n=1 T∑ t=1 [x>n,tB̂ŵt − yn,t]2, (4)
where Ŵ = (ŵ1 · · · ŵT ). Then with probability 1− δ, we have N0∑ n=1 T∑ t=1 [x>n,t(B̂ŵt −Bwt)]2 . (dk + kT ) log(NTdk) + log(1/δ).
Proof. By (4), we have N0∑ n=1 T∑ t=1 [x>n,tB̂ŵt − yn,t]2 ≤ N0∑ n=1 T∑ t=1 [x>n,tBwt − yn,t]2.
Since yn,t = x>n,tBwt + εn,t, we have
N0∑ n=1 T∑ t=1 [x>n,t(B̂ŵt −Bwt) + εn,t]2 ≤ N0∑ n=1 T∑ t=1 ε2n,t,
which implies N0∑ n=1 T∑ t=1 [x>n,t(B̂ŵt −Bwt)]2 ≤ N0∑ n=1 T∑ t=1 2εn,tx > n,t(B̂ŵt −Bwt). (5)
Next we bound the right-hand side of (5) via a uniform concentration argument. We let
B = {B′ ∈ Rd×k : ‖B′‖max ≤ 1}, W = {W ′ ∈ Rk×T : ‖W ′‖max ≤ 1}.
For any fixed matricesB′ ∈ B,W ′ ∈ W , we writeW ′ = (w′1 · · · w′T ) and define
ηn,t(B ′,W ′) = 2εn,tx > n,t(B ′w′t −Bwt).
Note that ηn,t(B′,W ′) is an independent sub-Gaussian variable with sub-Gaussian norm 2x>n,t(B
′w′t−Bwt). By the general Hoeffding’s inequality (Lemma 1), with probability 1− δ, we have
N0∑ n=1 T∑ t=1 ηn,t(B ′,W ′) . √ f(B′,W ′) log(1/δ), (6)
where f(B′,W ′) = N∑ n=1 T∑ t=1 [x>n,t(B ′w′t −Bwt)]2.
Next we apply the -net. Let B′ = N (B, ‖·‖max, ),W ′ = N (W, ‖·‖max, ). Applying an union bound over B′ ×W ′ for (6), we have
Pr [ ∀(B′,W ′) ∈ B′ ×W ′ :
N0∑ n=1 T∑ t=1 ηn,t . √ f(B′,W ′) log(1/δ)
] ≥ 1− δ|B′ ×W ′|.
Since f(B′,W ′) is (3NTdk)-Lipschitz with respect to ‖·‖max, we have
Pr [ ∀(B′,W ′) ∈ B ×W :
N0∑ n=1 T∑ t=1 ηn,t . √ f(B′,W ′) log(1/δ) + 24NTdk
] ≥ 1− δ|B′ ×W ′|.
Note that |B′| = O((1/ )dk) and |W ′| = O((1/ )kT ). Let δ = δ0/|B′ ×W ′| and = (24NTdk)−1. We have
Pr [ ∀(B′,W ′) ∈ B ×W :
N0∑ n=1 T∑ t=1 ηn,t . √ f(B′,W ′)[(dk + kT ) log(NTdk) + log(1/δ0)]
] ≥ 1− δ0.
(7)
Now we assume that the above event holds. Combining (5) and (7), we have
Pr [√ f(B̂, Ŵ ) . √ (dk + kT ) log(NTdk) + log(1/δ0) ] = Pr [ f(B̂, Ŵ ) . √ f(B̂, Ŵ )[(dk + kT ) log(NTdk) + log(1/δ0)]
] ≥ 1− δ0,
which proves our lemma.
Lemma 3. With probability 1−O((NT )−2), for every m ∈ [M ], t ∈ [T ], we have
λmin  Gm∑ n=Gm−1+1 xn,t,an,tx > n,t,an,t  & Gm − Gm−1 d .
Proof. The proof can be done by following the proof of Lemma 4 in Han et al. (2020).
Lemma 4. For each epoch m ∈ [M ], with probability 1−O((NT−2)), we have∥∥∥B̂Ŵ −BW∥∥∥2 F . (dk + kT ) log(NTdk) + log(1/δ) (Gm − Gm−1)/d ,
where B̂, Ŵ are computed at Line 5 in Algorithm 1.
Proof. Placing N0 = Gm − Gm−1 in Lemma 2, with probability 1−O((NT )−2), we have Gm∑
n=Gm−1
T∑ t=1 [x>n,t,an,t(B̂ŵt −Bwt)] 2 . (dk + kT ) log(NTdk). (8)
By Lemma 3, with probability 1−O((NT )−2), we have Gm∑
n=Gm−1
T∑ t=1 [x>n,t,an,t(B̂ŵt −Bwt)] 2
= T∑ t=1 (B̂ŵt −Bwt)>  Gm∑ n=Gm−1 xn,t,an,tx > n,t,an,t  (B̂ŵt −Bwt) &
T∑ t=1 (B̂ŵt −Bwt)> Gm − Gm−1 d (B̂ŵt −Bwt)
= Gm − Gm−1
d
∥∥∥B̂Ŵ −BW∥∥∥2 F . (9)
We conclude by combining (9) with (8).
Let
Rm = Gm∑ n=Gm−1+1 T∑ t=1 max a∈[K] 〈xn,t,a,θt〉 − 〈xn,t,an,t ,θt〉
be the regret incurred in the m-th epoch. We have the following lemma. Lemma 5. We have
E[Rm] . ( √ NTdk + T √ kN) √ log(NTdk) log(NKT ).
Proof. At round n that belongs to epoch m, for task t, we have
max a∈[K] θ>t (xn,t,a − xn,t,an,t) ≤ max a∈[K] {θ>t (xn,t,a − xn,t,an,t) + θ̂>m−1,t(xn,t,an,t − xn,t,a)}
= max a∈[K]
(θt − θ̂m−1,t)>xn,t,a + (θ̂m−1,t − θt)>xn,t,an,t
≤ 2 max a∈[K] ∣∣∣(θt − θ̂m−1,t)>xn,t,a∣∣∣. By Assumption 2 and a union bound over K actions, T tasks, and (Gm − Gm−1) rounds, we have with probability 1− (NKT )−2, for every n ∈ (Gm−1,Gm] and t ∈ [T ],
max a∈[K] ∣∣∣(θt − θ̂m−1,t)>xn,t,a∣∣∣ . ∥∥∥θt − θ̂m−1,t∥∥∥√ log(NKT ) d . (10)
Using a union bound over (10) and Lemma 4, with probability 1−O((NT )−2), the regret incurred in the m-th epoch is
E[Rm] . Gm∑
n=Gm−1+1
T∑ t=1 max a∈[K] θ>t (xn,t,a − xn,t,an,t)
. (Gm − Gm−1) T∑ t=1 ∥∥∥θt − θ̂m−1,t∥∥∥√ log(NKT ) d ≤ Gm √ T ∥∥∥Θ̂−Θ∥∥∥ F √ log(NKT ) d
(11)
≤ Gm √ T
√ (dk + kT ) log(NTdk)
(Gm−1 − Gm−2)/d
√ log(NKT )
d (12)
. √ NT (dk + kT ) log(NTdk) log(NKT ) (13) . ( √ NTdk + T √ kN) √ log(NTdk) log(NKT ),
where we denote Θ̂ = ( θ̂m−1,1 · · · θ̂m−1,T ) in (11) and the second inequality in (11) uses Cauchy, (12) uses Lemma 4, (13) uses Gm−1 . Gm−1 − Gm−2 and Gm/Gm−1 . √ N . Since that the above bound holds with probability 1− O((NT )−2) and that the regret is bounded by NT , we prove the lemma.
Proof of Theorem 1. The regret is bounded by
E[RN,T ] = M∑ m=1 E[Rm] .M( √ NTdk + T √ kN) √ log(NTdk) log(NKT )
= ( √ NTdk + T √ kN) √ log(NTdk) log(NKT ) log logN.

heading: B PROOF OF THEOREM 2
text: In this appendix, we assume that Σt = I in Assumption 2 and that all noises are gaussian, i.e. εn,t ∼ N (0, 1) for all n ∈ [N ], t ∈ [T ].
For each task t ∈ [T ], we denote the regret incurred on task t as
RN,(t) = N∑ n=1 max a∈[K] 〈xn,t,a,θt〉 − 〈xn,t,an,t ,θt〉.
We divide Theorem 2 into the following two lemmas.
Lemma 6. Under the setting of Theorem 2, we have infA supI E[R N,T A,I ] ≥ Ω(T
√ kN).
Lemma 7. Under the setting of Theorem 2, we have infA supI E[R N,T A,I ] ≥ Ω(
√ dkNT ).
Proof of Theorem 2. We combine Lemma 6 and Lemma 7.
Our proofs to the lemmas will be based on the lower bounds for the (single-task) linear bandit setting, which corresponds to the T = 1 case in our multi-task setting. For this single-task setting, we assume k = d and B = Id. We write the regret as RN = RN,1 and call algorithms for the single-task setting as single-task algorithms. Lemma 8 (Han et al. (2020), Theorem 2). Assume N ≥ d2 and d ≥ 2. Let N (µ,Σ) be the multivariate normal distribution with mean µ and covariance matrix Σ. There is a constant C > 0, such that for any single-task algorithm S, we have
sup ‖w‖≤1
E[RNS,I ] ≥ C √ dN,
where I is the instance with hidden linear coefficients w.
Next we use it to prove Lemma 6 and Lemma 7. The main idea to prove Lemma 6 is to note that we can treat our setting as T independent k-dimensional linear bandits.
Proof of Lemma 6. Suppose there is an algorithmA that achieves supI E[R N,T A,I ] ≤ CT
√ kN . Then
we have
sup ‖wt‖≤1 E [ T∑ t=1 R N,(t) A,I ] ≤ sup I E[RN,TA,I ] ≤ CT √ kN.
Therefore, there exists t ∈ [T ] such that
sup I
E[RN,(t)A,I ] ≤ 1 T CT √ kN = C √ kN,
which contradicts Lemma 8.
Proof of Lemma 7. Suppose there is an algorithm A that achieves supI E[R N,T A,I ] ≤ C
√ dkNT . We
complete the proof separately, based on whether k ≤ d2 or not. Note that when k > d 2 , the lower
bound in Lemma 6 becomes Ω(T √ kN) = Ω(T √ dN). Since T ≥ k, we have √ dkTN . T √ dN . Thus we conclude by Lemma 6.
In the remaining, we assume k ≤ d2 . Without loss of generality, we assume that d is even and that 2k divides T . For i = 1, . . . , k, we denote the regret of group i as
R N,((i)) A,I = iT/k∑ t=(i−1)T/k+1 R N,(t) A,I .
We consider instances such that tasks from the same group share the same hidden linear coefficients θt. Since there are k groups, we have
sup ‖θt‖≤1 E [ k∑ i=1 R N,((i)) A,I ] ≤ sup I E[RN,TA,I ] ≤ C √ dkNT . (14)
Therefore, there exists a group i ∈ [k] such that
sup Ii
E[RN,((i))A,I ] ≤ 1 k C √ dkNT = C
√ dNT
k , (15)
which means that the regret incurred in group i is less than C √ dNT/k. Since the tasks in group i share the same hidden linear coefficients, they could be regarded as one large linear bandit problem. Since there are T/k tasks in group i, the large linear bandit is played for N · T/k rounds. By Lemma 8, the algorithm A must have incurred C √ dNT/k regret on group i, which contradicts with (15).

heading: C METHOD-OF-MOMENTS ESTIMATOR UNDER BANDIT SETTING
text: The following theorem shows the guarantee of the method-of-moments (MoM) estimator we used to find the linear feature extractorB. In this appendix, for a matrixB with orthogonal columns, we write B⊥ to denote its orthogonal complement matrix (a matrix whose columns are the orthogonal complement of those ofB).
Theorem 5 (MoM Estimator). Assume N1T & polylog(N1, T ) · d 1.5k λ0ν
. We have with probability at least 1− (N1T )−100, ∥∥∥B̂>⊥B∥∥∥ . Õ( d1.5kλ0ν√N1T ). (16) The theorem guarantees that our estimated B̂ is close to the underlying B in the operator norm so long as the values N1 and T are sufficiently large. We add a remark that our theorem is similar to Theorem 3 in Tripuraneni et al. (2020). The key differences are: (i) we use a uniform distribution to find the feature extractor, while they assumed the input distribution is standard d-dimensional Gaussian; (ii) the SNR (signal-to-noise ratio) in our linear bandit setting is worse than that in their supervised learning setting, and thus we get an extra d factor in our theorem.
In the sequel, we prove the theorem.
Lemma 9 (Hoeffding). Let ε1, . . . , εN be i.i.d. 1-sub-Gaussian random variables. We have
Pr [∣∣∣∣∣ 1N N∑ i=1 (εi − E εi) ∣∣∣∣∣ ≥ t ] ≤ 2e− t 2 2N .
Lemma 10 (Matrix Bernstein’s inequality, Vershynin (2018), Theorem 5.4.1). Let X1, . . . ,Xm ∈ Rd×d be independent, mean zero, symmetric random matrices that ‖Xi‖ ≤M almost surely for all i ∈ [m]. Let σ2 = ∥∥∑m i=1 EX2i ∥∥. We have Pr
[∥∥∥∥∥ m∑ i=1 Xi ∥∥∥∥∥ ≥ δ ] ≤ 2d exp ( − δ 2/2 σ2 +Mδ/3 ) .
Equivalently, with probability at least 1− δ, we have∥∥∥∥∥ m∑ i=1 Xi ∥∥∥∥∥ .√σ2 log(d/δ) +M log(d/δ). Lemma 11 (Moments of Uniform Distribution on Sphere). Let x ∼ Unif(Sd−1) be a uniformly chosen unit vector. We have Ex61 = 15d(d+2)(d+4) , Ex 4 1 = 3 d(d+2) , Ex 2 1 = 1 d . Moreover, we have Ex41x22 = 3d(d+2)(d+4) .
Proof. We need to recall the fact that when x ∼ Unif(Sd−1), its coordinate xi follows the Beta distribution: x1+12 ∼ Beta( d−1 2 , d−1 2 ). Then we prove the lemma by noting the moments of the Beta distribution (Fang, 2018).
Corollary 6 (Uniform Distribution on Sphere). Let x ∼ Unif(Sd−1) be a uniformly chosen unit vector. We have the following statements.
(a) E 〈x,θ〉2xx> = 2θθ >+I
d(d+2) .
(b) E 〈x,θ〉4xx> = 12θθ >+3I
d(d+2)(d+4) .
Proof. Let e1 = (1, 0, . . . , 0) ∈ Rd be the unit vector. For (a), note that
(E 〈x, e1〉2xx>)ij = Ex21xixj =  0, i, j 6= 1 and i 6= j, 0, i = 1 6= j or j = 1 6= i, 1 d(d+2) , i = j 6= 1,
3 d(d+2) , i = j = 1.
Therefore, we have E[〈x, e1〉2xx>] = 1d(d+2) (2e1e > 1 +I). By the isotropy (rotation invariance) of uniform distribution, we have
E 〈x,θ〉2xx> = 1
d(d+ 2) (2θθ> + I).
For (b), note that
(E 〈x, e1〉4xx>)ij = Ex41xixj =  0, i, j 6= 1 and i 6= j, 0, i = 1 6= j or j = 1 6= i, 3 d(d+2)(d+4) , i = j 6= 1,
15 d(d+2)(d+4) , i = j = 1.
Therefore, we have E 〈x, e1〉4xx> = 1d(d+2)(d+4) (12e1e > 1 + 3I). By the isotropy, we have
E 〈x,θ〉4xx> = 1
d(d+ 2)(d+ 4) (12θθ> + 3I).
Let An,t = r2n,txn,tx > n,t,M = 1 N1T ∑N1 n=1 ∑T t=1An,t. We decompose M into three terms M = M1 +M2 +M3, where
M1 = 1
N1T N1∑ n=1 T∑ t=1 〈xn,t,θt〉2xn,tx>n,t,
M2 = 1
N1T N1∑ n=1 T∑ t=1 2εn,t〈xn,t,θt〉xn,tx>n,t,
M3 = 1
N1T N1∑ n=1 T∑ t=1 ε2n,txn,tx > n,t.
LetA1nt = 1N1T 〈xn,t,θt〉 2xn,tx > n,t. Next we analyze each error ‖Mi − EMi‖. Lemma 12. With probability at least 1− 1 N31T 3 , we have
‖M1 − EM1‖ .
√ log(dN1T )
d3N1T +
log(dN1T )
N1T .
Proof. We have
EA21nt = 1
N21T 2 E 〈xn,t,θt〉 4xn,tx > n,t =
12θtθ > t + 3I
d(d+ 2)(d+ 4)N21T 2 .
Using Lemma 10 with m = N1T,M = 1N1T , σ 2 . m · 1 d3N21T 2 = 1 d3N1T , we have with probability at least 1− 1 N31T 3 ,
‖M1 − EM1‖ .
√ log(dN1T )
d3N1T +
log(dN1T )
N1T .
Next we analyze the errors ofM2 andM3. Since these errors contain the unbounded sub-Gaussian terms εn,t, we need to cut their tails before applying the matrix Bernstein inequality. Define ε′n,t = εn,t I{|εn,t| ≤ R}. We have∣∣E εn,t − E ε′n,t∣∣ ≤ E |εn,t| I{|εn,t| > R}
= R · Pr[|εn,t| > R] + ∫ +∞ R Pr[|εn,t| > x] dx
≤ 2Re−R 2 2 + ∫ +∞ R 2e− x2 2 dx ≤ 2(R+ 1 R )e− R2 2 .
Let ε′′n,t = ε 2 n,t I{|εn,t| ≤ R}. We have∣∣E ε2n,t − E ε′′n,t∣∣ ≤ E ε2n,t I{ε2n,t > R2}
= R2 · Pr [ ε2n,t > R 2 ] + ∫ +∞ R2 Pr ( ε2n,t > x ) dd x
≤ 2R2e−R 2 2 + ∫ +∞ R2 2e− x 2 dd x ≤ (2R2 + 4)e−R 2 2 .
Define
M ′2 = 1
N1T N1∑ n=1 T∑ t=1 2ε′n,t〈xn,t,θt〉xn,tx>n,t,
M ′′3 = 1
N1T N1∑ n=1 T∑ t=1 ε′′n,txn,tx > n,t.
Lemma 13. With probability at least 1− (N1T )−3, we have
‖M ′2 − EM ′2‖ .
√ log(dN1T )
d2N1T + R log(dN1T ) N1T .
Proof. LetA2nt = 1N1T 2ε ′ n,t〈xn,t,θt〉xn,tx>n,t. We find that
EA22nt = 4E (ε′n,t)2
N21T 2 E 〈xn,t,θt〉 2xn,tx > n,t =
4E (ε′n,t)2
N21T 2
2θtθ > t + I
d(d+ 2) .
We conclude by using Lemma 10 with m = N1T,M = RN1T , σ 2 . m · 1 d2N21T 2 = 1 d2N1T .
Lemma 14. With probability at least 1− (N1T )−3, we have
‖M ′′3 − EM ′′3 ‖ .
√ log(dN1T )
dN1T + R2 log(dN1T ) N1T .
Proof. LetA3nt = 1N1T ε ′′ n,txn,tx > n,t. We find that
EA23nt = E (ε′′n,t)2
N21T 2 Exn,tx
> n,t =
E (ε′′n,t)2
N21T 2
I d .
We conclude by using Lemma 10 with m = N1T,M = R 2 N1T , σ2 . m · 1 dN21T 2 = 1 dN1T .
Lemma 15. With probability at least 1− (N1T )−2, we have
‖M − EM‖ .
√ log(dN1T )
dN1T +
log2(dN1T )
N1T .
Proof. LetR = √
8 log(N1T ). With probability at least 1−(N1T )−3, we have |εn,t| ≤ R for every n ∈ [N1] and t ∈ [T ]. Note that in this case, we have M2 = M ′2 and M3 = M ′′3 . Using a union bound over Lemma 12, Lemma 13, and Lemma 14, we have with probability at least 1− (N1T )−2,
‖M − EM‖ ≤ ‖M1 − EM1‖+ ‖M2 − EM2‖+ ‖M3 − EM3‖ = ‖M1 − EM1|+ ‖M ′2 − EM2‖+ ‖M ′′3 − EM3‖ ≤ ‖M1 − EM1‖+ ‖M ′2 − EM ′2‖+ ‖EM ′2 − EM2‖
+ ‖M ′′3 − EM ′′3 ‖+ ‖EM ′′3 − EM3‖
.
√ log(dN1T )
dN1T +
log2(dN1T )
N1T .
Proof of Theorem 5. We note that σk+1(M)− σk+1(EM) ≤ ‖E‖. Under Assumption 3, we have
EM1 = 2λ0
d(d+ 2)T θθ> + c1I, EM2 = 0, EM3 = c3I,
where c1, c3 ∈ R are constants. Since
σk( 1
T θθ>)− σk+1(
1 T θθ>) = σk( 1 T WW>) = ν k ,
we have
σk(EM)− σk+1(EM) λ0ν
d2k .
AssumeN1T & polylog(N1, T )·dk so that ‖E‖ ≤ λ0νd2k . Together with Davis-Kahan sin θ theorem, we have with probability at least 1− (N1T )−100,
‖B̂>⊥B‖ . ‖B̂>⊥EB‖
σk(EM)− σk+1(EM)− ‖E‖ (17)
≤ ‖E‖ σk(EM)− σk+1(EM)− ‖E‖ . ‖E‖ λ0ν/d2k
. d2k
λ0ν (
√ log(dN1T )
dN1T +
log2(dN1T )
N1T ) . d1.5k√ N1T · polylog(d,N, T ),
where (17) uses the Davis-Kahan sin θ theorem (Bhatia, 2013, Section VII.3).

heading: D PROOF OF THEOREM 3
text: Our proof is similar to the proof of Theorem 3.1 of Rusmevichientong & Tsitsiklis (2010).
Lemma 16 (Rusmevichientong & Tsitsiklis (2010), Lemma 3.5). For two vectors u,v ∈ Rd, we have ∥∥∥∥ u‖u‖ − v‖v‖
∥∥∥∥ ≤ 2‖u− v‖max{‖u‖, ‖v‖} . Lemma 17. Let xt = arg maxx∈At〈x, θ̂t〉. We have
max x∈At 〈x,θt〉 − 〈xt,θt〉 ≤ J
‖θt − θ̂t‖2
‖θt‖ .
Proof. For θ ∈ Rd, we define ft(θ) = maxa∈At〈a,θ〉. Let x∗t = arg maxx∈At〈x,θt〉. Then we have
max x∈At
〈x,θt〉 − 〈xt,θt〉 = 〈x∗t − xt,θt〉 = 〈x∗t ,θt − θ̂t〉+ 〈x∗t − xt, θ̂t〉+ 〈xt, θ̂t − θt〉
≤ 〈x∗t ,θt − θ̂t〉+ 〈xt, θ̂t − θt〉 = 〈x∗t − xt,θt − θ̂t〉
= 〈ft(θt)− ft(θ̂t),θt − θ̂t〉 = 〈ft( θt ‖θt‖ )− ft( θ̂t ‖θ̂t‖ ),θt − θ̂t〉
≤ ‖ft( θt ‖θt‖ )− ft( θ̂t ‖θ̂t‖ )‖ · ‖θt − θ̂t‖ ≤ J‖ θt ‖θt‖ − θ̂t ‖θ̂t‖ ‖ · ‖θt − θ̂t‖
(18)
≤ 2J ‖θt − θ̂t‖ 2
‖θt‖ , (19)
where the first inequality in (18) uses Cauchy and (19) uses Lemma 16.
Lemma 18. For each task t ∈ [T ], we have E‖θ̂t − θt‖2 . k 2 λ21N2 + ‖B̂>⊥B‖2.
Proof. We define θ′t = B̂B̂ >θt. Note that θt = θ′t + B̂⊥B̂ > ⊥ . We have
θ̂t − θt = B̂ŵt − (B̂B̂>θt + B̂⊥B̂>⊥)θt = B̂(ŵt − B̂>θt)− B̂⊥B̂>⊥Bwt.
Note that B̂ is perpendicular to B̂⊥ and that ‖B̂‖ = ‖B̂⊥‖ = 1. We have
‖θ̂t − θt‖2 = ‖B̂(ŵt − B̂>θt)‖2 + ‖B̂⊥B̂>⊥Bwt‖2 ≤ ‖ŵt − B̂>θt‖2 + ‖B̂>⊥B‖2.
Let vi = √ λ0b̂i. The OLS estimator is given by
ŵt = ( N1+N2+1∑ n=N1+1 B̂>viv > i B̂ )−1 N1+N2∑ n=N1+1 B̂>xn,trn,t
= ( N1+N2+1∑ n=N1+1 B̂>xn,tx > n,tB̂ )−1 N1+N2∑ n=N1+1 B̂>xn,t(x > n,tB̂w ′ t + εn,t)
= w′t + ( N1+N2+1∑ n=N1+1 B̂>xn,tx > n,tB̂ )−1 N1+N2∑ n=N1+1 B̂>xn,tεn,t.
WriteA = ∑N1+N2+1 n=N1+1 B̂>xn,tx > n,tB̂.
E‖ŵt −w′t‖2 = N1+N2∑ n=N1+1 x>n,tB̂A −2B̂>xn,t E ε2n,t
≤ N1+N2∑ n=N1+1 x>n,tB̂A −2B̂>xn,t ≤ N2‖A−2‖ ≤ N2( k λ0N2 )2 = k2 λ20N2 .
Putting together, we have
E‖θ̂t − θt‖2 ≤ k2
λ20N2 + ‖B̂>⊥B‖2.
Proof of Theorem 3. Note that J = O(1) under our assumptions, as indicated by Rusmevichientong & Tsitsiklis (2010). So we have
E[RN,T ] ≤ TN1 + TN2 + T (N −N1 −N2)J E‖θ̂t − θt‖2
‖θt‖
≤ TN1 + TN2 + TNJ ‖B̂>⊥B‖2 + k2/(λ21N2)
ω
. TN1 + TN2 + TN · d3k2
N1T log3(NT ) + TN
k2
N2
≤ TN1 + Nd3k2
N1 log3(NT ) + TN2 + TN
k2
N2
≤ d1.5k √ NT log3(NT ) + kT √ N.

heading: E PROOF OF THEOREM 4
text: In this appendix, we assume all action sets are spherical, i.e. An,t ≡ Sd−1 for n ∈ [N ], t ∈ [T ]. Note that these action sets meet Assumption 3.
For each task t ∈ [T ], we useRN,(t) = ∑N n=1[maxx∈An,t〈x,θt〉−E 〈xn,t,θt〉] to denote the regret incurred on task t.
Lemma 19. Under the setting of Theorem 4, we have infA supI E[R N,T A,I ] ≥ CkT
√ N, where
C = 0.0001.
Lemma 20. Under the setting of Theorem 4, we have infA supI E[R N,T A,I ] ≥ Cd
√ kTN, where
C = 0.0001.
Proof of Theorem 4. We combine Lemma 19 and Lemma 20.
The proofs in this appendix would largely follow the proofs in Appendix B, yet this appendix is much longer, because we need to construct instances that satisfies Assumption 3, Assumption 4, and Assumption 5.
Our proofs to the lemmas are based on the lower bounds for the (single-task) linearly parameterized bandit setting, which corresponds to the T = 1 case in our setting. For this single-task setting, we assume k = d andB = Id. (This setting need not meet Assumption 4 and Assumption 5). We write the regret as RN = RN,1 and call algorithms for the single-task setting as single-task algorithms.
We invoke the following lower bound for the single-task setting.
Lemma 21. Assume N ≥ d2 and d ≥ 2. Let N (µ,Σ) be the multivariate normal distribution with mean µ and covariance matrix Σ. For any single-task algorithm, we have
E w∼N (0,Id/d)
[RN · I{0.09 ≤ ‖w‖ ≤ 3}] ≥ 0.006d √ N. (20)
The lemma can be proved by following the proof of Theorem 2.1 of Rusmevichientong & Tsitsiklis (2010). Next we use it to prove Lemma 19 and Lemma 20. The main idea to prove Lemma 19 is to note that we can treat our setting as T independent k-dimensional linear bandits.
Let µd be the conditional probability measure whose density function is given by
f(x) =
{ g(x)
PrX [0.09≤‖X‖≤3] , 0.09 ≤ ‖x‖ ≤ 3, 0, otherwise,
where X ∼ N (0, Id/d) is the multivariate gaussian vector and g(x) is the probability density function of X . Then (20) implies
E w∼µd
[RN ] ≥ 0.006d √ N. (21)
Proof of Lemma 19. Without loss of generality, we assume 2k divides T . Suppose, for contradiction, that there is an algorithm A′, such that for every instance I, it incurs regret E[RN,TA′,I ] ≤ CTk √ N . We replace the condition ‖wt‖ ≤ 1 by ‖wt‖ ≤ 3 in our setting. Note that A′ implies an algorithm A that incurs regret E[RN,TA,I ] ≤ 3CTk √ N .
We construct the following instances I = (B,W ), where B = (Ik 0)> and W = (w1 · · · wT ) as follows. Let
w1 = · · ·wT/2k = e1,wT/2k+1 = · · · = wT/k = e2, . . . ,w(k−1)T/2k+1 = · · · = wT/2 = ek,
where e1, . . . , ek ∈ Rk is the standard basis. Let
wT/2+1, . . . ,wT ∼ µd
be i.i.d. drawn. Thanks to the first T2 tasks, the instance I always satisfies Assumption 4. Note that ‖wt‖ ≥ 0.09, so the instance also satisfies Assumption 5. Then we have
E wT/2+1,...,wT∼µd
[RN,TA,I ] ≤ 3CTk √ N.
Thus
T∑ t=T/2+1 E wT/2+1,...,wT∼µd [R N,(t) A,I ] ≤ E wT/2+1,...,wT∼µd [RN,TA,I ] ≤ 3CTk √ N.
Therefore, we can find t ∈ [T/2 + 1, T ] such that
E wT/2+1,...,wT∼µd
[R N,(t) A,I ] ≤
3CT T/2 k √ N = 6Ck √ N.
We note that the expectation operator EwT/2+1,...,wT∼µd is over all randomness on tasks τ 6= t and its parameter wτ . So there is a realization of the randomness on other tasks τ 6= t that satisfies
E wt∼µd
[R N,(t) A,I | wτ , ετ , τ 6= t] ≤ 6Ck
√ N.
Based on the realization wτ , ετ , we design a single-task algorithm S, which plays task t and simulates other tasks τ 6= t with wτ , ετ . The algorithm achieves
E w∼µd [RNS,I ] = E wt∼µd [R N,(t) A,I | wτ , ετ , τ 6= t] ≤ 6Ck
√ N,
which contradicts to (21) because C = 0.0001.
Proof of Lemma 20. Without loss of generality, we assume that d is even and that 2k divides T . Suppose, for contradiction, that there is an algorithm A′, such that for every instance I, it incurs regret E[RN,TA′,I ] ≤ Cd √ kTN . We replace the condition ‖wt‖ ≤ 1 by ‖wt‖ ≤ 3 in our setting. Note that A′ implies an algorithm A that incurs regret E[RN,TA,I ] ≤ 3Cd √ kTN .
We prove the lemma separately, based on whether k ≥ d2 or not.
1. Consider k ≤ d2 . We generate k vectors ψ1, . . . , ψk so that ψi ∼ µd−i+1. For every dimension i, we consider a map
ci : Rd×i → Rd×(d−i+1), (x1, . . . , xi) 7→ (y1, . . . , yd−i+1),
so that when {x1, . . . , xi} are orthogonal, the set {y1, . . . , yd−i+1} is the orthonormal basis of the orthogonal complement span{x1, . . . , xi}⊥ of Rd. Note that ci can be computed efficiently, e.g. by Gram-Schmidt process.
We define φ1 = ψ1 and φi = ci−1(φ1, . . . , φi−1) · ψi for i ≥ 2. We observe that {φ1, . . . , φk} are orthogonal. Next we define our instance I = (B,W), where B = (b1 · · · bk) and bi = φi‖φi‖ . ForW = (w1 · · · wT ), we let
w1 = · · · = wT/k = ‖φ1‖ · e1, wT/k+1 = · · · = w2T/k = ‖φ2‖ · e2,
... w(k−1)T/k+1 = · · · = wT = ‖φk‖ · ek,
where {e1, . . . , ek} ⊆ Rk is the standard basis. Note that the instance I always satisfies Assumption 4 and Assumption 5 by our choice of wt. Note that we have divided the tasks into k groups, so that each group share the same vector wt. For i = 1, . . . , k, we denote the regret of group i as
R N,((i)) A,I = iT/k∑ t=(i−1)T/k+1 R N,(t) A,I .
For the algorithm A, it incurs regret
E φ1,...,φk
[R N,((i)) A,I ] ≤ 3Cd
√ kNT .
So there is a group i ∈ [k], such that
E φ1,...,φk
[R N,((i)) A,I ] ≤
1 k 3Cd √ kNT = 3Cd
√ NT
k .
Similar to the proof of Lemma 19, we can fix the randomness for groups j 6= i to obtain a realization, such that
E φi|φj(j 6=i)
[R N,((i)) A,I | φj , εj , j 6= i] ≤
1 k 3Cd √ kNT = 3Cd
√ NT
k .
Here we note that φi could depend on φι for ι ≥ i+ 1. Now we let ψ′i ∼ µd−k+1 and let φ′i = Aψ ′ i, A = ck−1(φ1, . . . , φi−1, φi+1, . . . , φk). (22)
Note that φ′i and φi | φj(j 6=i) are identical, so we have
E φi=φ′i
[R N,((i)) A,I | φj , εj , j 6= i] ≤
1 k 3Cd √ kNT = 3Cd
√ NT
k . (23)
We complete the proof by showing that (23) implies a single-task algorithm S that plays a (d−k+1)dimensional linear bandit for NTk times. Let w = ψ ′ i. Then w is independently drawn from µd−k+1. Next we design the algorithm S, which runs A by playing the tasks t ∈ T = {(i − 1)T/k + 1, . . . , iT/k} and simulates other tasks τ /∈ T . Note that playing the task t ∈ T is the same as playing the single-task bandit defined by w, because we have θt = Bwt = φi = Aψ′i and the matrix A is known (as in (22)) after we fix the randomness on other tasks. Since |T | = Tk and A is played for N times, S can play the single-task bandit specified by w for NTk times. As a result, we have
E w∼µd−k+1 [RNS,I(I)] = E φi=φ′i [R N,((i)) A,I | φj , εj , j 6= i]
≤ 3Cd √ NT
k ≤ 6C(d− k + 1) √ NT
k ,
which contradicts to (21) because C = 0.0001. 2. Consider k > d2 . In this case, the lower bound in Lemma 19 becomes Ω(Tk √ N) = Ω(Td √ N).
Since T ≥ k, we have d √ kTN . Td √ N . Thus we conclude by Lemma 19.

heading: F EXPERIMENTS FOR INFINITE-ARM SETTING
text: Setup In all experiments, we set d = 10, N = 104 and the action At = Sd−1. The parameters are generated as follows. We consider k = 2, 3 in our experiments. The noise εn,t ∼ N (0, 1) are i.i.d. Gaussian random variables. To verify our theoretical results, we consider a hyper-parameter
c ∈ {0.5, 1, 1.5, 2}. For each c, we run E2TC with N1 = dck √ N T and N2 = k √ N .
Results and Discussions We present the simulation results in Figure 5 and Figure 6. We emphasize that the y-axis in our figures corresponds to the regret per task, which is defined as RegN,TT .
Our main observation is that only when the number of tasks T is large and we choose the right scaling N1 = d1.5k √ N T , our method can outperform the PEGE algorithm. We discuss several implications of our results. First, representation learning does help, especially when there are many tasks available for us to learn the representation, as we see in all figures that the regret per task of E2TC decreases as T increases. Second, the help of representation learning is bounded. In the figures, we see that the curves of E2TC bends to a horizontal line as T increases, which suggests a lower bound on the regret per task. Meanwhile, we also proved an Ω(k √ N) lower bound on the regret per task in Theorem 4. Third, representation learning may have adverse effect without enough task. In our figures, this was established by noting that our algorithm cannot outperform PEGE when T is small. This corresponds to the “negative transfer” phenomenon observed in previous work (Wang et al., 2019). Fourth, the correct hyper-parameter c = 1.5 is crucial. For hyperparameter other than c = 1.5, the figures show that our algorithm would require much more tasks to outperform PEGE. Lastly, by comparing the two figures, we notice that our algorithm has bigger advantage when k is smaller, which corroborates the scaling with respect to k in our regret upper bound. In contrast, PEGE does not benefit from a smaller k.

Review summaries:
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. 

This paper theoretically studies the benefits of representation learning in linear bandit problems. The key assumption is the existence of a common linear feature extractor. Two different setting are studied. In the finite-action setting, the authors provide the MLinGreedy algorithm that achieves matching upper and lower bounds (up to polylog factors). In the infinite-action setting, the authors provide the $E^2TC$ algorithm that can achieve lower regret than the naive method when the number of tasks is large. Experiments on both synthetic and real-world data are conducted, which confirm the theoretical results.