Multi-task Representation Learning for Pure Exploration in Bilinear Bandits
Authors: Subhojyoti Mukherjee
Qiaomin Xie
Josiah P. Hanna
Robert D Nowak
Venue: NeurIPS 2023 poster (2023)
Keywords: Linear Bandits
Experimental design
Pure Exploration
Representation Learning

Abstract: We study multi-task representation learning for the problem of pure exploration in bilinear bandits. In bilinear bandits, an action takes the
form of a pair of arms from two different entity types and the reward is a bilinear function of the known feature vectors of the arms. In the \textit{multi-task bilinear bandit problem}, we aim to find optimal actions for multiple tasks that share a common low-dimensional linear representation. The objective is to leverage this characteristic to expedite the process of identifying the best pair of arms for all tasks. We propose the algorithm GOBLIN that uses an experimental design approach to optimize sample allocations for learning the global representation as well as minimize the number of samples needed to identify the optimal pair of arms in individual tasks. To the best of our knowledge, this is the first study to give sample complexity analysis for pure exploration in bilinear bandits with shared representation. Our results demonstrate that by learning the shared representation across tasks, we achieve significantly improved sample complexity compared to the traditional approach of solving tasks independently.

Multi-task Representation Learning for Pure Exploration in Bilinear Bandits
Subhojyoti Mukherjee
Qiaomin Xie ISyE Department UW-Maddisn Wisconsin, Madison
Josiah P. Hanna CS Department UW-Madison Wisconsin, Madison
ECE Department
UW- UW-Maddisn Wisconsin, Madison smukherjee2 2wwiee7wwwee sc. edu
Robert Nowak ECE Department UUW-Madison Wisconsin, Madison
Abstract
We study multi-task representation learning for the problem of pure exploration in bilinear bandits. In bilinear bandits, an action takes the form of pair of arms from two different entity types and the reward is a bilinear function of the known feature vectors of the arms. In the multi-task bilinear bandit problem, we aim to find optimal actions for multiple tasks that share common low-dimensional linear representation. The objective is to leverage this characteristic to expedite the process of identifying the best pair of arms for all tasks. We propose the algorithm GOBLIN that uses an experimental design approach to optimize sample allocations for learning the global representation needed to identify the optimal pair of arms in individual tasks. To the best knowledge, this is the first study give sample complexity analysis for pure exploration bilinear bandits with shared representation. Our results demonstrate that by learning the shared representation across tasks, we achieve significantly improved sample complexity compared to the traditional approach of solving tasks independently.
well minimize the number of samples
Introduction
Bilinear bandits (Jun et al., 2019; Lu et al., 2021; Kang al., 2022) are important class of sequential decision making problems. bilinear bandits (as opposed to the standard linear bandit setting) given pair of arms x; E dd and E dd2 at every round t and the interaction of this pair of arms with low-rank hidden parameter, O E Rd xd2 generates the noisy feedback (reward) x. O.z, Mt- The 7/ is random 1-subGaussian noise.
lot of real-world applications exhibit the above bilinear feedback structure, particularly applications that involve selecting pairs of items and evaluating their compatibility. For example, in drug discovery application, scientists may want to determine whether particular (drug, protein) pair interacts in the desired way (Luo et al., 2017). Likewise, online dating service might match a pair of people and gather feedback about their compatibility (Shen al., 2023). clothing website's 's recommendation system may suggest pair of items (top, bottom) for customer based on their likelihood matching (Reyes et al., 2021), In all these scenarios, the two items are considered as single unit, and the system must utilize available feature vectors (X,, z;) to learn which features of the pairs are most indicative of positive feedback in order to make effective recommendations. All the previous works in this setting (Jun al., 2019; Lu et al., 2021; Kang al., 2022) exclusively
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

focused on maximizing the number of pairs with desired interactions discovered over time (regret minimization). However, in many real-world applications where obtaining sample expensive and time-consuming. e.g., clinical trials (Zhao et al., 2009; Zhang et al., 2012), it is often desirable to identify the optimal option using as few samples as possible, i.e., we face the pure exploration scenario (Fiez. et al., 2019; Katz-Samuels et al., 2020) rather than regret minimization.
Moreover, in various decision- masiing scenarios, we may encounter multiple interrelated tasks such as treatment planning for different diseases (Bragman et al., 2018) and content optimization for multiple websites (Agarwal et al., 2009). Often, there exists a shared representation among these tasks, such as the features of drugs or the representations of website items. Therefore, we can leverage this shared representation to accelerate learning. This area of research is called multi-task representation learning and has recently generated lot of attention in machine learning (Bengio et al., 2013; Li et al., 2014; Maurer et al., 2016; Du et al., 2020; Tripuraneni et al., 2021). There are many applications of this multi-task representation learning in real-world settings. For instance, in clinical treatment planning, we seek to determine the optimal treatments for multiple diseases, and there may exist a lowddimensional representation common to multiple discases. To avoid the time-consuming process of conducting clinical trials for individual tasks and collecting samples, we utilize the shared representation and decrease the number of required samples.
The above multi-task representation learning naturally shows up in bilinear bandit setting as follows: Let there be M tasks indexed as m = 1,2,.... .M with each task having its own hidden parameter O mm.. E Rd xd2, Let each Om., has a decomposition of O
= B SmS... B,, where B, E R44 x and B2 E da are shared across tasks, but S Sm,= E R ki x ka is specific for task m. We assume that 11, l:2 < d1. d2 and M > d1, d2. Thus, B1 and B2 provide means of dimensionality reduction. Furthermore, we assume that each S,.. has rank r < min k2). In the terminology of multi-task representation learning B,, B2 are called feature extractors and Xm.t. Zm, are called rich observations (Yang et al., 2020, 2022; Du et al., 2023). The reward for the task m E {1, 2,,...... M} at round is
Tm,t X,n, 4 m I + 7m, x B1 S B. zm, time B Vm,t + 7m,
(1)
Bm,
Ven.t
Observe that similar to the learning procedure in Yang et al. (2020, 2022), each round t 1, 1,2, 2,2,.... for each task m E [M], the learner selects a left and right action Xm.t E X and Zm, E 2. After the player commits the batch of actions for each task {xm,tz Zm.t : m E [M]] it receives the batch of rewards (Fm,t= : E [M]] Also note that in (1) we define the gm.t E Re Vi,,t RR the latent features, and both gm.t; a are unknown to the learner and needs to be learned for each task m
(hence the name multi-task representation learning).
In this paper, we focus on pure exploration for multi-task representation learning in bilinear bandits where the goal is to find the optimal left arm Xm,,, and right arm Z,,, + for each task with minimum number of samples (fixed confidence setting). First, consider a single-task setting and let O, have low rank r. Let the SVD of the O, UDV T_ Prima-facie if U and V are known then one might want to project all the left and right arms in the r x subspace of U and V and reduce the bilinear bandit problem into 22 dimension linear bandit setting. Then from Soare et al. (2014); Fiez et al. (2019); Katz-Samuels et al. (2020) to solve this r2 dimensional linear bandit pure exploration problem. Following the analysis of this line of work (in linear bandits) (Mason et al., 2021; Mukherjee et al., 2022, 2023) one might conjecture that a sample complexity bound of 0 / is possible where A is the minimum reward gap and Ō(-) hides log factors. Similarly, for the multi-task setting one might be tempted to use the linear bandit analysis of Du et al. (2023) to convert this problem into M concurrent p2 dimensional linear bandit problems with shared representation and achieve a sample complexity bound of Ö(M)² 2). However, these matrices (subspaces) are unknown and so there is a model mismatch as noted in the regret analysis of bilinear bandits (Jun et al., 2019; Lu et al., 2021; Kang et al., 2022). Thus it is difficult to apply the ?² dimensional linear bandit sample complexity analysis. Following the regret analysis of bilinear bandit setting by Jun et al. (2019); Lu al. (2021); Kang et al. (2022) we know that the effective dimension is actually (d, + dr)r. Similarly for the multi-task representation learning the effective dimension should scale with the learned latent features (k1 + k2)r. Hence the natural questions to ask are these:
can apply one of the algorithms
1) Can we design single-task pure exploration bilinear bandit algorithm whose sample complexity scales as Ö((d) + d

---

2) Can design algorithm for multi-task pure exploration bilinear ban- dit problem that can learn the latent features and has sample complexity that scales as (M //)
In this paper, we answer both these questions affirmatively. In doing so, we make the following novel contributions the growing literature of multi-task representation learning in online settings:
I) We formulate the multi-task bilinear representation learning problem. To our knowledge, this the first work that explores pure exploration in multi-task bilinear representation learning setting.
2) We proposed the algorithm GOBLIN for a single-task pure exploration bilinear bandit setting whose sample complexity scales as Ö(d, + d2)r/A²). This improves over RAGE (Fiez et al., 2019) whose sample complexity scales as ((d1d2))/A²)
3) Our algorithm GOBLIN for multi-task pure exploration bilinear bandit problem learns the latent features and has sample complexity that scales as ((M(k) /2)r/A²). This improves DouExpDes (Du et al., 2023) whose samples complexity scales as O(M (k: A²).
Preliminaries: We assume that ||x||2 < 1, || < 1, |@.11 F < So and the r-th largest singular value of θ E Rd Xd2 is S,.. Let p = d1d2 denote the ambient dimension, and k (d1 + d2)r denote the effective dimension. Let [n]
arg x O.z. For any x, define the gap A(x, == x. O. 21. n). Let x., x.,2. O.z and further- more A = 27 A(x, z). Similarly, for any arbitrary vector E W define the gap of E R A(w) = (w. w) 0., for some 0, E R2 and furthermore, A min A(w). If A E Raxd is a positive semidefinite matrix, and w E RP is a vector, let ||w]2 == w T Aw denote the induced semi-norm. Given any vector b R WI denote the w-th component as b Let Aw = {b E RIW bw 2 0, E Ewew b = }} denote the set of probability distributions on W. We define ((W) (w - w' Vw. w' E W, w + w'} as the directions obtained from the differences between each pair of arms and (WV))) = [w. w V/w E W.,., as the directions obtained from the differences between the optimal arm and each suboptimal arm.
2 Pure Exploration in Single-Task Bilinear Bandits
In this section, we consider pure exploration in single-task bilinear bandit setting as a warm-up to the main goal of learning representations for the multi-task bilinear bandit. To our knowledge, this is the first study of pure exploration in single-task bilinear bandits. We first recall the single-task bilinear bandit setting as follows: At every round
2.... the learner observes the reward x, Oz + 77 where the low rank hidden parameter θ E IRd xd2 is unknown to the learner, X2 E Rd4, 2, E IR are visible to the learner, and n/ a sub--auussian noise. We assume that the matrix O has low rank which is known to the learner and d1, d2 > r. Finally recall that the goal is to identify the optimal left and right arms x., Z. with minimum number of samples.
We propose phase-based, two-stage arm elimination algorithm called G-Optimal Design for Bilinear Bandits (abbreviated as GOBLIN). GOBLIN proceeds in phases indexed by l = 1, 2.... As this pure-cxploration problem, the total number of samples is controlled by the total phases which depends on the intrinsic problem complexity. Each phase of GOBLIN consists of two stages; the estimation of θ. stage, which runs for 7f rounds, and pure exploration in rotated arms stage that runs for 7 rounds. We will define 7F in Section 2.1, while rotated arms and T are defined in Section 2.2. At the end of every phase, GOBLIN eliminates sub-optimal arms to build the active set for the next phase and stops when only the optimal left and right arms remaining. Now discuss the individual stages that occur at every phase l of GOBLIN.
2.1 Estimating Subspaces of e (Stage of the 6-th phase)
In the first stage of phase f, GOBLIN estimates the row and column sub-spaces θ,, Then GOBLIN uses these estimates to reduce the bilinear bandit problem in the original ambient dimension == dyda to a lower effective dimension h (dl + dr)r. To do this, GOBLIN first vectorizes the Rd 1d into new vector E R2 and then solves the E-optimal design in Step 3 of
3

---

Algorithm 1 (Pukelsheim, 2006; Jun et al., 2019; Du et al., 2023). Let the solution to the E-optimal design problem at the stage 1 of f-th phase be denoted by b,". Then GOBLIN samples each for Treb times, where T = O S+) (step 7 of Algorithm 1). In this paper, we sample an arm b number of times. However, this may lead to over-sampling of design (G or -ooptimal) is actually suggesting. However, we can match the number of allocations of an arm to the design using an efficcient Rounding Procedures (see Pukelsheim (2006); Fiez et al. (2019)). Let θ be estimate of e in stage 1 of phase f. GOBLIN estimates this by solving the following well-defined regularized minimization problem with nuclear norm penalty:
than what the
- arg min (2( + 7el1 O muc: G Xdg
Le( = ( θ
Wolrs (), O)
where Q(-), (-), are appropriate functions stated in Definition 1, respectively Appendix A.3. The Q--) function takes as input the rank-one matrix x,z, which is obtained after reshaping w.. Note that Xs, and 2g are the observed vectors in d1 and d2 dimension and θ E IR xd2 Finally, set
e (+)) da log(2(d) +d2)/8) This is in step 8 of Algorithm 1.
the regularization parameter " == AV
of
2.2 Optimal Design for Rotated Arms (Stage of f-th phase)
In stage 2 of phase E, GOBLIN leverages the information about the learned sub-space of O to rotate the arm set and then run the optimal design the rotated arm set. Once recover + one might be tempted to run a pure exploration algorithm (Soare et al., 2014; Fiez et al., 2019; Katz-Samuels et al., 2020; Zhu et al., 2021) to identify x, and z.. However, then the sample complexity will scale with di d2. In contrast GOBLIN uses the information about the learned sub-space of O to reduce the problem from ambient dimension dı d2 to effective dimension (d1 + d2)r. This reduction is done as follows: Let θ UT be the SVD of θ in the f-th phase. Let U1 and V be orthonormal bases of the complementary subspaces of U, and V, respectively. Let X₂ and 2g be the active set of arms in the stage of phase f. Then rotate the arm sets such that new rotated arm sets are as follows:
X, = fx [ū,0+1xx x x X). 2, {z = ( V= z Ze}.
(3)
Let H - | V V components are from the complementary subspaces as follows:
Then define vectorized arm set so that the last (d1 r) (d2 -r)
W₁ { [[ee (xr.Zi:r)) vec (x++1:d, zT... vec (Znn
vec (X,+1ad 2 E Rdida x E Ne, E Zt}
0 1.. 1.k
[vec(H, H f, 1or 1or) vec H h 11 lod, 1sr) vec (1r3
O = vec (He Iad +1dg))
(4)
which implies 10 1l2 O (d/d2r/Tf) by Lemma 3 in Appendix A.1. So the last p components of O2 are very small compared to the first k components. Hence, GOBLIN has now reduced the di d2 dimensional linear bandit (d1 + de)r dimensional linear bandit using (3), (4). This is shown in step 10 of Algorithm 1.
Now in stage 2 of phase f, GOBLIN implements G-optimal design (Pukelsheim, 2006; Fiez et al., 2019) in the rotated arm set &x,, 2, defined in (3). To do this, first GOBLIN defines the rotated vector w [ ; [.22)] IRP that belong to the set W,- Then GOBLIN solves the G-optimal design (Pukelsheim, 2006) as follows:
by arg min
mAx W.W'S W.W.
W - w/ 1 Lew bw T
(5)
b
This is shown in step 11 of Algorithm I and A, is defined in (6). It can be shown that sampling ac- cording to by leads to the optimal sample complexity. This is discussed in Remark 1 in Appendix A.2. The key point to note from (5) is that due to the estimation in the rotated arm space W; we are guaranteed that the support of (b)) < Ō(k(k
1)/2) (Pukelsheim, 2006). On the other hand,

---

if the G-optimal design of Fiez et al. (2019); Katz-Samuels et al. (2020) are run in d1dg dimension then the support of by will scale with d1dg which will lead to higher sample complexity. Then GOBLIN samples each E We for [y times, where 7 = ( (O((W)) log(4e [W1/0]
Note that the total length of phase f, combining stages and 2 is (T + T rounds. Observe that the stage design is on the whole arm W whereas stage 2 design is on the refined active set W,-
Let the observed features in stage 2 of phase be denoted by We IRT xp, and rg E R be the observed rewards. Define the diagonal matrix A₃ as
A diag X
Xe
(6)
p-k
where, X - - log(1 / >> A. Deviating from Soare al. (2014); Fiez et al. (2019)
GOBLIN constructs a regularized least square estimator at phase as follows
θ
= arg min BERP
|1 rell2 + 2
111122
This regularized least square estimator (7) forces the last p - k components of 0g to be very small compared to the first k components. Then GOBLIN builds the estimate 0₂ from (7) only from the observations from this phase (step 13 in Algorithm 1) and eliminates sub-optimal actions in step 14 in Algorithm using the estimator 0g. Finally GOBLIN eliminates sub-optimal arms build the next phase active set W, and stops when [Wyl = 1. GOBLIN outputs the arm in W. and reshapes it to get the X, and 2, The full pseudocode is presented in Algorithm 1.
Algorithm 1 G-Optimal Design for Bilinear Bandits (GOBLIN) for single-task setting
1: Input: arm set X, Z, confidence 6, rank r of O., spectral bound S, of O., S, = dd. log dd A, x = --- /8(4 (1 + d2)r log(1 + - Letp d1d2, k == (d1 + d2)r.
7, S.
1, T 1ol ( /8). Define A as (6), B (8VAS / \ /8d, der log (19(11))))))))) Let 8,
2: -,
3: Define vectorized arm w == [ i:d 2 [dd and w W. Let of
the E-optimal design be b = arg min 11( Ewew b w -1.
4: while [Wel do
5:
6 2- -2- 5g 5/6².
6: (Stage 1:) Explore the Low-Rank Subspace
7: Pull arm w W exactly D Te times and observe rewards r,, for t. 1,...,
8: Compute θ using (2).
9:
(Stage 2:) Reduction to low dimensional linear bandits
Let the SVD of θ U.D.V. Rotate arms in active set W build W, following (4).
11: Let Dî arg minl, max w ew, ||w w
new bw W +T+e//n)-1
12: Define , (C)(WW)) minp max W.W,
[wew bw w +Ae/n)-1--
Set T - (((()) (og(1)" Then pull arm w E W exactly [ times
13:
of
and construct the least squares estimator O₂ using only the observations of this phase where
O, is defined in (7). Note that O₂ is also rotated following (4).
14:
Eliminate arms such that W <- W, (w E W, maxw'ew, (w' w. 0.) 201}
15:
L +- {+}
16:
Output the arm in W. and reshape to get the x. and 2,
2.3
Sample Complexity Analysis of Single-Task GOBLIN
We now analyze the sample complexity of GOBLIN the single-task setting through the following thcorem.
5

---

Theorem 1. (informal) With probability at least 1 - 6, GOBLIN returns the best arms X., Z,, and the number of samples used is bounded by Ö (d. (ditd
dddider
+ A²
5.
Discussion 1. In Theorem I the first quantity is the number of samples needed to identify the best arms Xx, Z, while the second quantity is the number of samples to learn O (which is required to find the best arms). Note that the magnitude of S, would be free of d1, d2 since θ contains only r nonzero singular values and |@o1l I 1, and hence we assume that S, = (11/VF) (Kang et al.,
2022). So the sample complexity of single-task GOBLIN scales as Ö (d) (d)+)))r) ) However, if one runs RAGE (Fiez et al., 2019) on the arms in X, 2 then the sample complexity will scale as Ö ((dd)))
Proof (Overview) of Theorem 1: Step 1 (Subspace estimation in high dimension): We denote the vectorized arms in high dimension as w E W. We run the E-optimal design to sample the arms in W. Note that this E-optimal design satisfies the distribution assumption of Kang al. (2022) which enables us to apply the Lemma in Appendix A.I. This leads ||O, -0.1 I Cadidar (og()(d) //4)/6) for some C1 > 0. Also, note that in the first stage of the f-th phase by
(8d (o1())) and sampling each arm w E W exactly ( times we are S. guaranteed that 110 11ppll Odd1d2//T) Summing up over & = 111 to [log2 (4A "1] we get that
setting of
the total sample complexity of the first stage is bounded by ((vd,d2r/S.))
Step 2 (Effective dimension for rotated arms): We rotate the arms w E W in high dimension to get the rotated arms w E We in step 10 of Algorithm 1. Then we show that the effective dimension of w
scales 8k log (1 + / when X 8A log( (og(1++f_./A) in Lemma 7 of Appendix A.4. Note that
this requires different proof technique than Valko et al. (2014) where the budget is given apriori and effective dimension scales with log(n). This step also diverges from the pure exploration proof technique of Fiez et al. (2019); Katz-Samuel et al. (2020) as there is no parameter A to control during phase É, and the effective dimensions in those papers do not depend on phase length.
Step 3 (Bounded Support): For any phase É, we can show that S (>(WW, S p// where, yy = max > 0 <)/ C conv (W U -W) is the gauge norm of y (Rockafellar, 2013). Note that this is a worst-case dependence when p ()(())) scales with p. Substituting this value of o ()WW)) in the definition of X we can show that A₃ does not depend on w or y w w'. Then following Theorem 21.1 in Lattimore and Szepesvári (2020) we can show that the G-optimal
design D is equivalent to D-optimal design bl = arg max, log Ewew, bw w w +AAc Then using
Trank-Wolfe algorithm (Jamieson and Jain, 2022) we can show the support by or equivalently 6P is 8k (oog(+++/-//// (A) log(1 + (1)+1) This is shown in Lemma (Appendix A.4).
bounded by at most
2
Step 4 (Phase length and Elimination): Using the Lemma 9, concentration Lemma 5, and using the log determinant inequality in Lemma 7 and Proposition 1 (Appendix A.4) we show that the phase length in the second stage is given by Th = [ {BB ((W(,)) (og(2)/41//))) This is discussed in (x -0.--')) Discussion 3 (Appendix A.4). We show in Lemma 10 (Appendix A.4) that setting this phase length and sampling each active arm in W. exactly [f times results in the elimination of sub-optimal actions with high probability.
Step 5 (Total Samples): We first show that the total samples in the second phase are bounded by
& O 75)
log( & loga(4 - [logg [1]] where the effective dimension h = (d1 + de)r. Finally,
we combine the total samples of phase as (7F + T.)) The final sample complexity is given by summing over all phases from 1 to [log2 --)- The claim of the theorem follows by noting
Ō(k/73)< < (kk/A²).
3 Multi-task Representation Learning
In this section, we extend GOBLIN to multi-task representation learning for the bilinear bandit setting. In the multi-task setting, we now have M tasks, where each task m [M] has reward model stated

---

.......foo for each task m E [M.] the learner selects a left and right action Xm.s E X' and Zm.t E 2. After the player commits the batch of actions for each task { (Xm.t, 2m,tt m E [M]] it receives the batch of rewards {rm.: m E [M]] Finally recall that the goal to identify the optimal left and right arms Xm,. Zm,. for each task m with minimum number of samples. We now state the following assumptions to enable representation
in (I). The learning proceeds as follows: At each round = 1.2.....
learning across tasks.
Assumption 1. (Low-rank Tasks) We assume that the hidden parameter m, all the m E [M] have a decomposition O m, m, B,2 and each Sm,= has rank r.
This is similar to the assumptions in Yang et al. (2020, 2022); Du et al. (2023) ensuring the feature extractors are shared across tasks in the bilinear bandit setting.
Assumption 2. (Diverse Tasks) We assume that Fmin M Em ② m, for some co > 0, S, is the r-th largest singular value of m,,, and a min (A) denotes the minimum eigenvalue of matrix A.
1
This assumption is similar to the diverse tasks assumption of Yang et al. (2020, 2022); Tripuraneni al. (2021); Du al. (2023) and ensures the possibility of recovering the feature extractors B, and B2 shared across tasks.
Our extension of GOBLIN to the multi-task setting is now phase-based, three-stage arm elimination algorithm. In GOBLIN each phase = 1,2.... consists of three stages; the stage for estimation of feature extractors B1, B2, which runs for Tf rounds, the stage for estimation of S,,+ which runs for Em -m rounds, and stage of pure exploration with rotated arms that runs for Em rounds. We will define - in Section 3.1, F in Section 3.2, while the rotated arms and - are defined in Section 3.3. At the end of every phase, GOBLIN eliminates sub-optimal arms build the active set for the next phase and stops when only the optimal left and right arms are remaining. Now we discuss the individual stages that occur every phase f 1.2, ... for multi-task GOBLIN.
3.1 Estimating Feature Extractors BJ and B₂ (Stage of Phase !)
In the first stage of phase L, GOBLIN leverages the batch of rewards {rimt : m E [M] at every round from M tasks to learn the feature extractors B, and B. To do this, GOBLIN first vectorizes the E X, E Z into a new vector w [X1:d, : Z1:d,] W and then solves the E-optimal design in step of Algorithm 2. Similar the single-task setting (Section 2) GOBLIN samples each w E Wm, for by. times for each task m, where Tf = //d/dd///// and bf is the solution to E-optimal design on w. Let the sampled arms for each task m at round s be denoted by Xm Zm which is obtained after reshaping w.. Then it builds the estimator Z, as follows:
Z, arg min (e(O) + O buer @ER4 Xdg
M - L(( = (,, θ) MTE
(Tm.s ((,mm,ZZT,.)) O)
(8)
m=
Then it performs SVD) decomposition on Z,, and let B,, B2 be the top-k; and top-kz left and right singular vectors of Z, respectively. These are the estimation of the feature extractors B, and Bg.
3.2 Estimating Hidden Parameter Sm.. per Task (Stage 2 of phase !
In the second stage of phase f, the goal is to recover the hidden parameter S,,. for each task m. GOBLIN proceeds as follows: First, let gm = x BB, and Vn = z B2,t be the latent left and right arm respectively for each m. Then GOBLIN defines the vector solves the E-optimal design in step 11 of Algorithm 2. It then samples for each task m, the latent arm E Wm for E Ber.cbee.. m. times, where 7 = O(VKrkzr/// (S.) and E w is the solution to
[gm; vm] E Wm and then
7

---

E-optimal design on w. Then it builds estimator Sm,e for each task m in step 12 as follows:
S m arg min (2( + vel O |nnue E
= 2
()) (0,
(r (B.m.Vm.ss))
Once GOBLIN recovers the Sm,e for each task m it has reduced the didg bilinear bandit to k2K2 dimension bilinear bandit where the left and right arms are gm E 9m. Vo Vm respectively.
3.3 Optimal Design for Rotated Arms per Task (Stage 3 of phase E)
In the third stage of phase similar Algorithm 1, the multi-task GOBLIN defines the rotated arm set gm K. for each task m for these ky hi dimensional bilinear bandits. Let the SVD of Sm,t = U Um, Dm,e VT, m.t- Define H mu, Ú m.e) S (V e V Then define the vectorized arm set so that the last (k) r)- (k2 r) components are from the complementary subspaces follows:
Wm.t {vvec (gm,1:r V mm...) vec (8m,r+1:d) T mmr vec (Bm,1.r vm.
vec (m..++ V m
0 (vec (, vec vec(H L. H 02 = vc
vee (,
(10)
This is shown in step 14 of Algorithm 2. Now we proceed similarly Section 2.2. We construct per-task optimal design for the rotated arm set L 2m and define the w = 1:d, Vm,1.d2] and
w E Wm where gm E and Vm E Vm respectively. Following (5) we know that minimize the sample complexity for the m-th bilinear bandit we need to sample according to G-optimal design
m arg min max w. w.W'EW...
|W
w "'liiEeew bm w +Am,t/n)->
(11)
bm
Then GOBLIN runs G-optimal design on the arm set W, following (11) and then samples each P
w E Wm, for [ m. times where D m,,,w is the solution to the G-optimal design, and T is defined in step 17 of Algorithm 2. So the total length of phase f, combining stages 1, 2 and is (Tf + +m FF + Em - rounds. Observe that the stage and design is on the whole arm W, W... whereas the stage 3 design is the refined active W.mee e- Let at the stage 3 of L-th phase the actions sampled be denoted by the matrix W m, RRT.. X and observed rewards Im R xkkika Define the positive diagonal matrix Am. according (6) but set ką k2 and = (k1 + k2)r. Then similar to Section 2.2 we can build for each task m only from the observations
from this phase
i mm.t = arg min
W O 100mmll 112 1
(12)
0
Finally GOBLIN eliminates the sub-optimal arms using the estimator mm, to build the next phase active set Wm,t and stops when [Wm,,l = 1. The full pseudo-code is given in Algorithm 2.
3.4 Sample Complexity analysis of Multi-task GOBLIN
We now present the sample complexity of GOBLIN for the multi-task setting. Theorem (informal) With probability at least GOBLIN returns the best arms X Z for each task m. and the total number of samples bounded by Ō [M(k+kk)rr +
MvEike + S.
ydider Se
2²
Discussion 2. In Theorem 2 the first quantity is the sample complexity to identify the best arms Xm, Zm,- and the second quantity is the number of samples to learn Sn,+ for each task m. This is required to rotate the arms to reach the effective dimension of (ki + )r. Finally, the third quantity is the number of samples needed to learn + m. (which in turn used estimate the feature extractors B, and Bg to learn the Sm,,). Again we assume that S, = 0(1/v7) (Kang et al., 2022). So the sample complexity of multi-task GOBLIN scales as Õ(M + k2)r/A?). However, if one runs DouExpDes (Du et al., 2023) then the sample complexity will scale as (M k2) which is worse than GOBLINwhen < kj or 62.
8

---

Algorithm 2 G-Optimal Design for Bilinear Bandits (GOBLIN) for multi-task setting
1: Input: arm set X, 2, confidence 6, rank r of .., spectral bound S, of ①.. S,
Skyker log( ( ), A, Am, (B(krtka)r log(1 + /A)) Let = k2k2, k (k1 k2)r.
m
2: Let Wm.1 + W., l + 1, TO log(482 (2100)) Define Am,t as in (6), B == (8VAS
3: Define arm w [x1:d Z1:d2] and E Wm m. Let Tp design be by = arg mim Aw 122 bw w w 1-
dd log( S,
Let E-optimal
4: while 3m [M]] [ do
5: C ===2 2 62 = 5/62.
6: (Stage 1:) Explore the Low-Rank Subspace
7: Pull arm w E W exactly [b F1 times for each task m and observe rewards {rm,ali
8:
Compute 2, using (8).
9:
(Stage 2:) Build S,,E for each task m
Let B1,G, Bg, be the top-k1 left and top-k2 right singular vectors of Ze respectively. Build
gm = x BI.t and V. = z Bz, for all x E X' and z E Z for each m. w 1:k, Vm,1 and w E Wm for each m. Let -,.
11:
Define vectorized
[8k har (01(1))))/5)) and b arg min. b...EA em || (Eweww. w w -
Sr
12:
Pull arm w E Wm exactly
[D bb m,,| times and observe rewards Tm,t, for t
1,..., Fm.e. for each task m. Then compute Sm,t using (9) for each m. (Stage 3:) Reduction to low dimensional linear bandits for each task m
13:
14:
SVD of S Sm,=ÜU, --mmmm Dm,ee VT, m, Rotate arms in active set Wmme-1 to build Wm,e using (10).
15:
Let b m. aar min max w WWW,
|ww W1 we.. EWW... bm W Am. e/n)
16:
Define 0 CD'(W e)) e))-mimn min max
I|w w/ll I12
w.w. E.W...
(2 wew. bm Amm.
Set - 64m S0 (J(W () log log( (M. 1/66) Then pull
17:
Wm for each task m
of
exactly [bm,t.w. G times and construct the least squares estimator θm, using only the observations of this phase where Om,e is defined (12).
Eliminate arms such that Wm, Wm Wm.e' (wm W Wm,e. max EW... m.c (''m W Ome) > 2mml)
19:
+-1+1
20:
W.,r Output the arm in W m and reshape to get the Xm,, and 2m,- for each task m.
Proof (Overview) of Theorem 2: Step (Subspace estimation in high dimension): The first steps diverge from the proof technique of Theorem 1. We now build the average estimator Z; to estimate the quantity Z, and apply Stein's 's lemma (Lemma to get bound of 2, 11 < C,d, +da))() for some
MM
O..m using (8). This requires us modify the Lemma in Appendix A.1
T/
C1 > 0. This is shown in Lemma 12 in Appendix A.6. Summing up over 8 = to [loga (4A 4A)1)]
we get that the total samples complexity of the first stage is bounded by Ö(vd,dgr/S.)
Step 2 (Estimation of left and right feature extractors): Now using the estimator in (8) we get good estimation of the feature extractors B and B2. Let Bi,z. B2, be the top-k; left and top-k2 right singular vectors of Z₂ respectively. Then using the Bavis-Kahan sin 0 Theorem (Bhatia, 2013) in Lemma 14, 15 (Appendix A.6) we have 1 (B4,) B2ll S 0lv(d, (di + d2)r/M
Step 3 (Estimation of Sm.e e in low dimension): Now we estimate the quantity Sm, E Rhixks for each task m. To do this we first build the latent arms gm x U and vn = z Ve for all X

---

and z E Z for each m, and sample them following the E-optimal design in step 12 of Algorithm 2. We also show in Lemma 16 (Appendix A.6) that m ( by ww > 0 which enables to sample following F-optimal design. Then use the estimator in (9). Then in Lemma 19 we show that
||mm / m. 3 & C kj ker log (2 22kk++ /Tm.e holds with probability greater than (1 - 6).
Also, note that in the second phase by setting mm."" V8kikgr (og( 101/ and sampling each arm w E W exactly [b times we are guaranteed that 1 l|2 Okkikrr/T in the é-th phase. Summing up over = 1 to |log2 (42-1) across each task M we get that the total samples complexity of the second stage is bounded by k(rSs..
Step 4 (Convert to k l2 bilinear bandits): Once GOBLIN recovers S =,,,"" it rotates the arm set following (10) to build W.m to get the k k2 bilinear bandits. The rest of the steps follow the same way as in steps 2,3 and 4 of proof of Theorem 1.
Step 5 (Total Samples): We show the total samples in the third phase are bounded by O R. log( k (ll,,, ) [logg(
(A A 1)1) where the effective dimension = (ki + k2)r. The to-
tal samples phase is given by T." Em (FR. +mm,)) Finally, we get the total sample complexity by summing over all phases from L = 1 to [logz (4A 1)1. The claim of the theorem follows by noting Ō(k/73) v Ō(k/A²).
Experiments
In this section, we conduct proof-of-concer concett experiments on both single and multi-task bilinear bandits. the single-task experiment, we compare against the ofate-of-the-art RAGE algorithm (Fiez et al., 2019). We show in Figure 1 (left) that GOBLIN requires fewer samples than the RAGE with an increasing number of arms. the multi-task experiment, we compare against the state-of-the-art DouExpDes algorithm (Du et al., 2023). We show in Figure 1 (right) that GOBLIN requires fewer samples than DouExpDes with an increasing number of tasks. As experiments are not central contribution, we defer fuller description of the experimental set-up Appendix A.8.
5
30
Setting (Log
Setting (Log
GOBLIN(ours)
1.6
GOBLIN(ours)
RAGE
1.4
+
DouExpDes
Samples
1.2
3-
1.0
0.8
2
@
0.6
1
0.4
0.2
6
10 Arms
14
5
10
15
20 Tasks
25
30
Figure 1: (Left) Single-task experiment: results show the number of samples required to identify the optimal action pair for differing numbers of actions. (Right) Multi-task experiment: results show the number of samples required identify the optimal action pair for varying numbers of tasks.
5 Conclusions and Future Directions
In this paper, we formulated the first pure exploration multi-task representation learning problem. We introduce an algorithm, GOBLIN that achieves a sample complexity bound of Ö(dh + )2) which improves upon the ((d, task setting. We then extend GOBLIN for multi-task pure exploration bilinear bandit problems by learning latent features which enables sample complexity that scales as (M + k2)r/A?) which improves over the (M k2) ?) sample complexity of DouExpDes (Du et al., 2023). Our analysis opens an exciting opportunity to analyze representation learning the kernel and neural bandits (Zhu et al., 2021; Mason et al., 2021). We can leverage the fact that this type of optimal design does not require the arm set to be an ellipsoid (Du et al., 2023) which enables us to extend our analysis to non-lincar representations.
sample complexity of RAGE (Fiez et al., 2019) in single-
10

---

References
Agarwal, D., Chen, B.-C., and Elango, (2009). Explore/exploit schemes for web content optimiza- tion. In 2009 Ninth IEEE International Conference on Data Mining, pages 1-10. IEEE.
Bengio, Y., Courville, A., and Vincent, (2013). Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 3588):1798-1828
Bhatia, R. (2013). Matrix analysis, volume 169. Springer Science & Business Media.
Bragman, F. J., Tanno, R., Eaton-Rosen, Z., Li, W., Hawkes, D. J., Ourselin, S., Alexander, D. C., McClelland, R., and Cardoso, M. J. (2018). Uncertainty in multitask learning: joint representa- tions for probabilistic mr-only radiotherapy planning. In Medical Image Computing and Computer Assisted InterventoonMMICAA 2018: 21st International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Purt IV 11, pages 3-11. Springer.
Du, S. S., Hu, W., Kakade, S. M., Lee, J. D., and Lei, Q. (2020). Few-shot learning via learning the representation, provably. arXiv preprint arXiv:2002.09434
Du, Y., Huang, L.,, and Sun, W. (2023). Multi-task representation learning for pure exploration in linear bandits. arXiv preprint arXiv:2302. 04441.
Fiez, T. Jain, L., Jamieson, K. G., and Ratliff, L. (2019). Sequential experimental design for transductive linear bandits. Advances in neural information processing systems, 32.
Jamieson, K. and Jain, L. (2022). Interactive machine learning.
Jun, K.-S., Willett, R., Wright, S., and Nowak, R. (2019). Bilinear bandits with low-rank structure. In International Conference on Machine Learning, pages 3163-3172. PMLR.
Kang, Y., Hsieh, C.-J,, and Lee, T. C. M. (2022). Efficient frameworks for generalized low-rank matrix bandit problems. Advances in Neural Information Processing Systems, 15:19971-19983
Katz-Samuels, J., Jain, L., Jamieson, K. G., al. (2020). An empirical process approach the union bound: Practical algorithms for combinatorial and linear bandits. Advances in Neural Information Processing Systems, 13:10371-10382
Kiefer, J. and Wolfowitz, J. (1960). The equivalence of two extremum problems. Canadian Journal of Mathematics, 12:363-366.
Lattimore, T. and Szepesvári, C. (2020). Bandit algorithms. Cambridge University Press.
Li, Zhang. H., Zhang, I.,, Huang, X., and Zhang, L. (2014). Joint collaborative representation with multitask learning for hyperspectral image classification. JEEE Transactions on Geoscience and Remote Sensing, 5299):593355936.
Lu, Y., Meisami, A., and Tewari, A. (2021). Low-rank generalized linear bandit problems. In International Conference on Artificial Intelligence and Statistics, pages 460-468, PMLR.
Luo, Y., Zhao, X., Zhou, J., Yang. J., Zhang, Y., Kuang, W., Peng. J., Chen, L., and Zeng. J. (2017). network integration approach for drug-target interaction prediction and computational drug repositioning from heterogeneous information. Nature communications, 8(1 8(1):573.
Mason, B., Camilleri, R., Mukherjee, S., Jamieson, K., Nowak, R., and Jain, L. (2021). Nearly optimal algorithms for level set estimation. arXiv preprint arXiv:2111.01763
Maurer, A., Pontil, M., and Comera-Pareders B. (2016). The benefit of multitask representation learning. Journal of Machine Learning Research, 17(81):1-32.
Minsker, (2018). Sub- Sub-gaussian estimators of the mean of random matrix with heavy-tailed entries. The Annals of Statistics, 46(6A);2871-2903.
Mukherjee, S., Tripathy, A. S., and Nowak, R. (2022). Chernoff sampling for active testing and extension to active regression. International Conference on Artificial Intelligence and Statistics, pages 7384-7432. PMLR.

---

Mukherjee, S., Xie. Q., Hanna, J., and Nowak, R. (2023). Speed: Experimental design for policy
evaluation in linear heteroscedastic bandits. arXiv preprint arXiv:330.123357.
Pukelsheim, F. (2006). Optimal design of experiments. SIAM.
Reyes, L. P., Oviedo, N. B., Camacho, E. C., and Calderon, J. M. (2021). Adaptable recommenda- tion system for outfit selection with deep learning approach. IPAC-PPppesSOnLLine, 54(13):605-611
Rockafellar, R. (2015). Convex analysis. princeton landmarks in mathematics and physics.
Shamir, O. (2011). variant of azuma's inequality for martingales with subgaussian tails. arXiv preprint arXiv: 10.2392.
Shen, Q. Han, S., Han, Y., and Chen, X. (2023). User review analysis of dating apps based on text mining, Plos one, 18(4):c0283896.
Soare, M., Lazaric, A., and Munos, R. (2014). Best-arm identification linear bandits. Advances in Neural Information Processing Systems, 27.
Stein, C., Diaconis, P., Holmes, S., and Reinert, G. (2004). Use of exchangeable pairs in the analysis of simulations. Lecture Notes-Monograph Series, pages 1-26.
Tripuraneni. N., Jin, C., and Jordan, M. (2021). Provable meta-learning of linear representations. In International Conference on Machine Learning, pages 10434-10443, PMLR.
Valko, M., Munos, R., Kveton, B., and Kocák, (2014). Spectral bandits for smooth graph functions. In International Conference on Machine Learning, pages 46-54. PMLR.
Yang. J., Hu, W., Lee, J. D., and Du, S. (2020). Impact of representation learning in linear bandits. arXiv preprint X11::2000000631.
Yang, J., Lei, Q., Lee, J. D., and Du, S. S. (2022). Nearly minimax algorithms for linear bandits with shared representation. arXiv preprint arXiv:2203. 15664.
Zhang, D., Shen, D., Initiative, A. D. N., et al. (2012). Mudti-modal multi-task learning for joint prediction of multiple regression and classification variables in alzheimer's ' disease. Neurolmage, 59(2):895-907.
Zhao, Y., Kosorok, M. R., and Zeng, D. (2009). Reinforcement learning design for cancer clinical trials. Statistics in medicine, 28(26):3244-3115
Zhu, Y., Zhou, D., Jiang, R., Gu, Q. Willett, R., and Nowak. R. (2021). Pure exploration in kernel and neural bandits. Advances in neural information processing systems, 34:1 -1630.
12


Reviews:
[{"review_id": "svcTSe3opK", "reviewer": "Reviewer_8J9B", "summary": "This paper studies pure exploration of bilinear bandits with multi-task representation learning, i.e., different tasks share a common low-dimensional linear representation.\nAs an intermediate step, the authors propose the first single-task bilinear bandit pure exploration algorithm, based on G-optimal design. Each phase of the proposed algorithm first estimates the unknown parameter $\\hat{\\Theta}$ (stage 1), and then use it to rotate the active arm set, and convert the problem to a $(d_1 + d_2)r$-dimensional linear bandit problem, on which G-optimal design is applied (stage 2).\nWhen extending to multi-task setting, stage 1 is modified to first estimate the feature extractors that are common to all tasks, and then estimate hidden parameter unique to each task. Then G-optimal design is applied as before.", "questions": "1. What is the lower bound for sample complexity in (single-task) bilinear bandit pure exploration? Does Theorem 1 already attains optimality in terms of $(d_1 + d_2)r, \\Delta, S_r$?\n\n2. In multi-task setting, as the estimation of task-specific parameter $\\hat{S}$ depends on the estimation of the common feature extractor, how is this reflected in the theoretical guarantee on $||\\hat{S}-S||^2_F$? Intuitively, larger number of tasks means more samples for the estimation of feature extractors, which further helps the estimation of $\\hat{S}$.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The problem setting studied in this paper, i.e., different tasks share common feature extractors while each having unique hidden parameter, is natural and well-motivated.\n\nTo the best of my knowledge, this paper proposes the first solution to pure exploration of bilinear bandits for both single-task and multi-task settings.", "weaknesses": "The authors mention improvement over RAGE (Fiez et al., 2019), but there lacks a more rigorous discussion about lower bound of the sample complexity in either single-task or multi-task setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating the soundness and contribution of our work. We answer the questions raised below.\n\n*Weakness*:\n\n1) Thank you for raising this point. There is no lower bound for the single task setting for pure exploration in bi-linear bandits with unknown underlying structures. In the paper [Rizk et al. (2021)]() they conjecture that under certain graph structure for pure exploration the lower bound of [Soare et al (2014)]() or [Fiez et al. (2019)]() can be reached. But currently, without assuming any additional assumption (like underlying graph structure) there is no lower bound. We leave this line of research for future work and will discuss this point in the camera-ready version.\n\n- G Rizk, A Thomas, I Colin, R Laraki, Y Chevaleyre, Best arm identification in graphical bilinear bandits, International Conference on Machine Learning, 2021\n- M Soare, A Lazaric, R Munos, Best-arm identification in linear bandits, Advances in Neural Information Processing Systems 27 (NIPS 2014)\n- T Fiez, L Jain, K G. Jamieson, L Ratliff, Sequential Experimental Design for Transductive Linear Bandits, Advances in Neural Information Processing Systems 32 (NeurIPS 2019)\n\n2) That's a great observation and thanks for reading the proof in detail. Note that the $M$ tasks have different unknown task-specific parameters $\\mathbf{S}\\_{m, \\star}$. The estimation of each $\\mathbf{S}\\_{m, \\star}$ requires using samples from the corresponding task $m$, thus resulting in a linear scaling with the number of tasks $M$ for the total sample complexity. However, this scaling is with $O\\left(\\frac{M(k_1+k_2)r\\log(\\delta^{-1})}{\\Delta^2}\\right)$ which depends in $k_1$, $k_2$ and rank $r$ instead of $d_1d_2$ and scales linear with $M$. This is where the dominating term of our sample complexity arises in Theorem 2. Also note that the estimation of $\\mathbf{B}_1$ and $\\mathbf{B}_2$ cannot be further improved as we used the standard Davis-Kahan theorem for bounding the error in estimation similar to [Yang et al 2021](), [Yang et al 2022](), [Du et al. 2023](). Again note that these works address linear bandit settings whereas we address bilinear bandit setting. \n\n- Jiaqi Yang, Wei Hu, Jason D Lee, and Simon S Du. Impact of representation learning in linear bandits. In International Conference on Learning Representations, 2021.\n- J Yang, Q Lei, JD Lee, SS Du, Nearly minimax algorithms for linear bandits with shared representation, arXiv preprint arXiv:2203.15664, 2022\n- Y Du, L Huang, W Sun, Multi-task Representation Learning for Pure Exploration in Linear Bandits, arXiv preprint arXiv:2302.04441, 2023\n\nPlease reach out to us during the discussion period if you have further questions and we hope that you will consider raising our score given these clarifications.\n"}, {"review_id": "gVsvWI7OA3", "reviewer": "Reviewer_9og1", "summary": "The authors study the problem of pure exploration in bilinear bandits, where the reward is a bilinear function of the feature vectors of two types of arms. They consider both single-task and multi-task settings, where the latter involves learning a shared low-dimensional representation across multiple tasks. They propose a phase-based algorithm called GOBLIN, which uses optimal experimental design techniques to estimate the hidden parameters and identify the optimal pair of arms for each task. They provide sample complexity analysis for both settings and show that their algorithm achieves significant improvement over existing methods. They also conduct experiments to demonstrate the empirical performance of their algorithm.", "questions": "1) What is the motivation and application of pure exploration in bilinear bandits? Authors gave motivation for bilinear bandits but it was not clear why pure exploration. \n2) What are the limitations of the proposed algorithms and analysis?\n3) How does GOBLIN handle the cases when the rank r of Θ∗ is unknown or when the spectral bound Sr of Θ∗ is unknown or inaccurate?\n4) How does GOBLIN handle the cases when the feature extractors B1 and B2 are not shared across tasks or when they are noisy or incomplete?\n5) How does GOBLIN handle the cases when the tasks have different reward distributions or different reward gaps?\n6) How does GOBLIN handle the cases when the arm sets X and Z are infinite or continuous or dynamic?\n7) How does GOBLIN handle the cases when the rewards are not sub-Gaussian or when there are outliers or adversarial noise?\n8) What are the computational and memory complexities of GOBLIN and how can they be improved or optimized?\n\n", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "1) Originality: The paper presents a novel approach to the problem, combining existing ideas in a creative way and formulating the problem in a new and insightful manner.\n2) Quality: The methodology used in the paper is rigorous and well-executed, with clear and thorough explanations of the techniques employed.\n3) Clarity: The writing is clear and concise, making it easy to follow the arguments and understand the main contributions of the paper.\n4) Significance: The results presented in the paper have important implications for the field, removing limitations from prior results and opening up new avenues for research.\n", "weaknesses": "1) The paper does not provide a comprehensive literature review of the existing work on the topic, and does not compare or contrast its approach with the previous methods or results. (Related Work) The paper only cites a few papers on bilinear bandits, linear bandits, and representation learning, but does not discuss how they are related to or different from its own approach. It also does not mention any existing work on pure exploration in bilinear bandits or multi-task representation learning in online settings. For example some missing papers are: [1-2]\n[1] Soare, Marta, Ouais Alsharif, Alessandro Lazaric, and Joelle Pineau. \"Multi-task linear bandits.\" In NIPS2014 workshop on transfer and multi-task learning: theory meets practice. 2014.\n[2] Deshmukh, Aniket Anand, Urun Dogan, and Clay Scott. \"Multi-task learning for contextual bandits.\" Advances in neural information processing systems 30 (2017).\n2) The paper does not explain the rationale or motivation behind the choice of the data set, the experimental design, the evaluation metrics, or the baselines. (Methodology and Experiments) The paper only briefly describes the data set and the experimental setup, but does not justify why they are suitable or relevant for the problem. It also does not explain how it chooses the evaluation metrics or the baselines, or how they reflect the performance of its algorithm.\n3) The paper does not report the details of the data preprocessing, feature engineering, model architecture, hyperparameter tuning, or implementation. (Methodology and Experiments) The paper does not provide any information on how it preprocesses the data, extracts features, designs the model architecture, tunes the hyperparameters, or implements its algorithm. It also does not provide any pseudocode or algorithmic description of its method.\n4) \"We assume that the matrix Θ∗ has a low rank r which is known to the learner\" This seems like a big assumption too. Authors should include this in their list of assumptions. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating the soundness and contribution of our work. We answer the questions raised below.\n\n*Weakness*\n\n1) Thanks for pointing out these papers. We will discuss the papers pointed out in the revision. Here we provide a brief comparison of our work with [1][2]. [1] studies regret minimization for the multi-task linear bandit, whereas we study pure exploration for multi-task bilinear bandit, which requires a different algorithmic design and analysis. Similarly [2] studies regret minimization for multi-task contextual bandit, which is a different setup than the pure exploration studied in our work. We remark that there is only one work on bilinear bandit pure exploration [Geovani et al. (2021)](), which focuses on graph structure for a single task setting. We will add more detailed discussion and comparison to the revision.\n2) We are not sure how some of these comments are related to our experimental setting. We use a standard single and multi-task bilinear bandit experimental setting similar to [Jun et al., 2019](); [Lu et al. 2021](),  [Kang et al., 2022](). However note that these papers only study single task regret minimization setup whereas we study multi-task pure exploration setting. Our metric of performance is sample complexity against arms or tasks (similar to the setting of [Fiez et al. (2019)]() and [Du et al. 2023]()). \n3) We are not sure how this comment is related to our paper. We use a standard bilinear bandit setting similar to [Jun et al., 2019](); [Lu et al. 2021](),  [Kang et al., 2022](), [Du et al. 2023](). It will be very helpful if the reviewer reaches us to us during the discussion phase to point out any specific questions on the experimental setup. \n4) We will make it clear that rank $r$ knowledge is required and will state this as a separate assumption.  Similar assumptions have been commonly considered in prior works like [Jun et al., 2019](); [Lu et al. 2021](),  [Kang et al., 2022](). In practice, one can leverage domain knowledge or cross-validation to choose the rank.\n\n*Questions*\n\n1) In many real-world applications where obtaining a sample is expensive and time-consuming, e.g., clinical trials (Zhao et al., 2009; Zhang et al., 2012), it is often desirable to identify the optimal option using as few samples as possible, i.e., we face the pure exploration scenario rather than regret minimization. This is a well-studied setting in linear bandits but has not been studied in the bilinear setting. To avoid the time-consuming process of conducting clinical trials for individual tasks and collecting samples, we utilize the shared representation and decrease the number of required samples. We mention the reason for pure exploration in bilinear bandits in multi-task setting in lines 37-48. \n2) In the current analysis we explicitly rely on the linear structure analysis. It will be very interesting to extend this line of research to generalized linear bandits. Also, note that currently there is no lower bound for multi-task pure exploration in the bilinear bandit setting with no assumption on the underlying structure. It will be interesting for future works to give lower bounds for this setting.\n3) Thanks for pointing out this new direction. The current analysis requires knowledge of the rank and $S_r$ similar to the works of [Jun et al., 2019](); [Lu et al., 2021](); [Kang et al., 2022](). We leave this line of research for future works.\n4) The key idea of a multi-task setting is that the feature extractors are shared across tasks so that there is some gain in accelerating learning and reducing the rate of regret (See [Yang et al., 2020](), [Yang et al., 2022](); [Du et al., 2023]()). If tasks do not share the common feature extractors $B_1$ and $B_2$ then the problem is reduced to $M$ independent subproblems, with a regret scaling as $O\\left(M(d_1d_2 r)/\\Delta^2\\right)$.\n5) Currently we assume each task has different distributions (but all are parametric sub-Gaussian distributions) with different reward gaps. The regret scales with minimum gap across all tasks. This is an interesting direction for future work.\n6) For continuous action space currently there are no existing works in bilinear bandits. It is not clear how to extend our current analysis to continuous action space.\n7) It is an interesting future work to extend our analysis to non-parametric or even adversarial noise. We have taken the first step towards understanding pure exploration for bilinear bandit setting in the case of parametric sub-Gaussian reward distribution setting.\n8) The computational and memory complexity of GOBLIN is same as that of [Jun et al., 2019](); [Lu et al. 2021](),  [Kang et al., 2022]() and scales as $\\tilde{O}(d T)$.\n\nThank you for taking the time to read our response. We hope we have answered your questions and clarified your concerns sufficiently for you to consider raising your final score. Please reach out to us during the discussion period if you have further questions."}, {"review_id": "hrPzNrlUO2", "reviewer": "Reviewer_vf5o", "summary": "This paper studies multi-task representation learning for the problem of pure exploration in bilinear bandits. The authors design algorithm GOBLIN that uses an experimental design approach to optimize sample allocations for learning the global representation as well as minimize the number of samples needed to identify the optimal pair of arms in individual tasks. Their results corroborate that by learning the shared representation across tasks, they achieve significantly improved sample complexity compared to the traditional approach of solving tasks independently.", "questions": "Please see the weaknesses above.", "limitations": "Please see the weaknesses above.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThis paper is well-written. The considered problem is interesting and well-motivated.\n2.\tThe theoretical analysis looks sound.\n", "weaknesses": "1.\tWhat is the main technical novelty of this work, beyond the existing SVD, G-optimal design, phased elimination methods in the bilinear/linear bandit literature?\n2.\tCould you elaborate more on the difference of formulation, algorithms and results between this work and [Du et al. 2023]? Can the considered multi-task representation learning for bilinear bandit problem somehow be transformed to the multi-task representation learning for linear bandit problem in [Du et al. 2023]? Could you explain more about the claim “if one runs DouExpDes [Du et al. 2023] then its sample complexity will scale as $\\tilde{O}(M k_1 k_2 / \\Delta^2)$” in Discussion 2?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper. Below, we clarify the key novelty in our approach and compare our work with [Du et al. 2023].\n\n*Weakness* \n\n1) Even though we use existing ideas in SVD, G-optimal design phased elimination in linear bandit setting, it is highly non-trivial to extend that analysis to multi-task bilinear bandit setting where the regret must scale with the intrinsic dimensions $k_1$ and $k_2$ instead of $d_1$ and $d_2$. We list these technical novelties below:\n    - We use E-optimal design in an initial exploration to learn two feature extractors $\\mathbf{B}\\_1$ and $\\mathbf{B}\\_2$ jointly across all the $M$ tasks. This requires a new estimator (see eq (8)) and concentration techniques not studied in [Jun et al., 2019](); [Lu et al. 2021](),  [Kang et al., 2022](), [Du et al. 2023](). The sample complexity of this step scales as $\\frac{\\sqrt{d_1 d_2 r}}{S_r}$ where $S_r$ is the spectral bound of $\\mathbf{\\Theta}_\\star$ and scalses as $S_r= \\Theta(1/\\sqrt{r})$.\n    - We then estimate the individual task-specific parameter matrix $\\mathbf{S}_{m,*}\\in\\mathbb{R}^{k_1 \\times k_2}$ using a new estimator in eq (9). In particular, we use a regularized minimization optimization problem with a nuclear norm penalty to construct the estimator. Note that [Du et al. 2023]() does not involve this type of estimator as it only needs a good estimate of the vector $w_m,$ which is obtained by a simple OLS estimation. Our sample complexity of this step scales as $\\frac{M \\sqrt{k_1 k_2 r}}{S_r}$.\n    - Then we rotate the arm sets (similar to [Jun et al., 2019]()) to reduce the $M$ bilinear bandits to $(k_1 + k_2)r$ dimensional linear bandits. We then use G-optimal design to learn the optimal arm for these $(k_1 + k_2)r$ dimensional linear bandits such that the sample complexity scales as $\\frac{M\\left(k_1+k_2\\right) r}{\\Delta^2}$. We are inspired by the proofs of G-optimal design from [Fiez et al. (2019)]()) to derive a guarantee for this step. Finally, note that this step is also not there in [Du et al. 2023](). So the final sample complexity scales as $O\\left(\\frac{M\\left(k_1+k_2\\right) r}{\\Delta^2}\\right)$ as $S_r = \\Theta(1/\\sqrt{r})$.\n    - All of these steps have to be done carefully and it is highly non-trivial to combine them so that the desirable sample complexity scaling appears. \n    - We mention these technical novelties/approaches of combining these methods in the proof overview of Theorem 1 in lines 192-225 or Theorem 2 in lines 300-327. Please reach out to us if you have any specific questions about our novelty. \n2) Thanks for asking this clarification question.  [Du et al. 2023]() critically consider the linear bandit setting which is a simpler setting than the bilinear bandit setting. Crucially there is a single feature extractor in the linear bandit setting of [Du et al. 2023]() whereas we have two feature extractors. [Du et al. 2023]() do not require to rotate the arms and learn in the rotated arm space as we do (and explain) in section 2.2 and section 3.3. Learning in the rotated arm space leads to a regret of $M(k_1 + k_2)r/\\Delta^2$ instead of $M(k_1k_2)/\\Delta^2$ (as done in [Du et al. 2023]()). Hence, we make the statement that DouExpDes [Du et al. 2023]() has a sample complexity of $\\tilde{O}\\left(M k_1 k_2 / \\Delta^2\\right)$ in Discussion 2. We will add the clarification in the camera-ready version. \n\n- KS Jun, R Willett, S Wright, R Nowak, Bilinear bandits with low-rank structure, ICML 2019\n- T Fiez, L Jain, K G. Jamieson, L Ratliff, Sequential Experimental Design for Transductive Linear Bandits, NeurIPS 2019\n- Y Du, L Huang, W Sun, Multi-task Representation Learning for Pure Exploration in Linear Bandits, arXiv preprint arXiv:2302.04441, 2023\n- Y Lu, A Meisami, A Tewari, Low-rank generalized linear bandit problems,  International Conference on Artificial Intelligence and Statistics, 2021\n- Y Kang, CJ Hsieh, TCM Lee, Efficient frameworks for generalized low-rank matrix bandit problems, NeurIPS, 2022\n\nWe hope we have answered your questions and clarified your concerns sufficiently for you to consider raising your final score. Please reach out to us during the discussion period if you have further questions."}]

Author rebuttal:
References:

- KS Jun, R Willett, S Wright, R Nowak, Bilinear bandits with low-rank structure, ICML 2019
- T Fiez, L Jain, K G. Jamieson, L Ratliff, Sequential Experimental Design for Transductive Linear Bandits, NeurIPS 2019
- Y Du, L Huang, W Sun, Multi-task Representation Learning for Pure Exploration in Linear Bandits, arXiv preprint arXiv:2302.04441, 2023
- Y Lu, A Meisami, A Tewari, Low-rank generalized linear bandit problems,  International Conference on Artificial Intelligence and Statistics, 2021
- Y Kang, CJ Hsieh, TCM Lee, Efficient frameworks for generalized low-rank matrix bandit problems, NeurIPS, 2022
- G Rizk, A Thomas, I Colin, R Laraki, Y Chevaleyre, Best arm identification in graphical bilinear bandits, International Conference on Machine Learning, 2021
- M Soare, A Lazaric, R Munos, Best-arm identification in linear bandits, Advances in Neural Information Processing Systems 27 (NIPS 2014)
- Jiaqi Yang, Wei Hu, Jason D Lee, and Simon S Du. Impact of representation learning in linear bandits. In International Conference on Learning Representations, 2021.
- J Yang, Q Lei, JD Lee, SS Du, Nearly minimax algorithms for linear bandits with shared representation, arXiv preprint arXiv:2203.15664, 2022

Meta-review: The paper studies multi-task bilinear bandits, where the $d_1 \times d_2$ reward matrix $\Theta_m$ for the task $m$ can be decomposed as $B_1 S_mB_2$. The task-specific matrix $S_m$ is of dimension $k_1\times k_2$ with $k_1,k_2\ll d_1,d_2$, and is additionally of rank $r\ll k_1,k_2$. This model seems complicated, and would deserve a proper practical justification. The authors present a pure exploration algorithm GOBLIN that achieves a sample complexity scaling as $r(d_1+d_2)/\Delta^2$ in the single-task case, and when $M\gg d_1,d_2$ as $Mr(k_1+k_2)/\Delta^2$ in the $M$-task case.

The reviewers are rather positive about this work, and the sample complexity guarantees are state-of-the-art. The claimed guarantees are only better when the gap $\Delta$ is very close to 0. See e.g., Theorem 1 where the sample complexity also scales as $\sqrt{d_1d_2r}$. The same observation holds for Theorem 2 (by the way, the definition of $\Delta$ in the multi-task case is not given). 

For a technical and algorithmic perspective, the paper does not seem to introduce very novel ideas. The authors combine the idea of June et al. (2019) to reduce the problem to a classical linear bandit problem (see e.g. the first lines of Stage 2 in Algorithm 1) and the algorithm from Fiez et al. 2019. 

Some parts and claims in the paper remained obscure for the revewiers and needed clarification – the authors provided sound explanations in the rebuttal and during the discussion phase. To improve the paper, the authors could account for the recommendations made by the reviewers as summarized below.

-	Some of the relevant literature is missing (as mentioned by Reviewer 9og1), but this can be addressed. 

-	The assumption $M\gg d_1,d_2$ should be discussed in more detail.

-	It is hard to understand why the authors compare their algorithm to RAGE and to that presented in (du et al. 2023), because these are made for linear bandits. 

-	The authors could discuss the existence of lower bounds.

-	It would be nice to have a thorough practical justification of the model. 

-	The presentation could be significantly improved. See e.g., Section 4 .
Decision: Accept (poster)