Title: Multi-task Representation Learning for Pure Exploration in Bilinear Bandits
Venue: NeurIPS 2023
Accepted: yes
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

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 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.

1 Introduction

Bilinear bandits (Jun et al., 2019; Lu et al., 2021; Kang et al., 2022) are an important class of sequential decision-making problems. In bilinear bandits (as opposed to the standard linear bandit setting) we are given a pair of arms x t ∈ R d 1 and z t ∈ R d 2 at every round t and the interaction of this pair of arms with a low-rank hidden parameter, Θ ∗ ∈ R d 1 × d 2 generates the noisy feedback (reward) r t = x ⊤ t Θ ∗ z t + η t . The η t is random 1 -subGaussian noise.

Alot 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 a drug discovery application, scientists may want to determine whether a particular (drug, protein) pair interacts in the desired way (Luo et al., 2017). Likewise, an online dating service might match a pair of people and gather feedback about their compatibility (Shen et al., 2023). A clothing website's recommendation system may suggest a pair of items (top, bottom) for a customer based on their likelihood of matching (Reyes et al., 2021). In all of these scenarios, the two items are considered as a single unit, and the system must utilize available feature vectors ( x t , z t ) 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 et al., 2019; Lu et al., 2021; Kang et al., 2022) exclusively

focused on maximizing the number of pairs with desired interactions discovered over time (regret minimization). However, 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 (Fiez et al., 2019; Katz-Samuels et al., 2020) rather than regret minimization.

Moreover, in various decision-making 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 a 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 low-dimensional representation common to multiple diseases. 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 Θ m, ∗ ∈ R d 1 × d 2 . Let each Θ m, ∗ has a decomposition of Θ m, ∗ = B 1 S m, ∗ B ⊤ 2 , where B 1 ∈ R d 1 × k 1 and B 2 ∈ R d 2 × k 2 are shared across tasks, but S m, ∗ ∈ R k 1 × k 2 is specific for task m . We assume that k 1 , k 2 ≪ d 1 , d 2 and M ≫ d 1 , d 2 . Thus, B 1 and B 2 provide a means of dimensionality reduction. Furthermore, we assume that each S m, ∗ has rank r ≪ min { k 1 , k 2 } . In the terminology of multi-task representation learning B 1 , B 2 are called feature extractors and x m,t , z m,t are called rich observations (Yang et al., 2020, 2022; Du et al., 2023). The reward for the task m ∈ { 1 , 2 , . . . , M } at round t is

Observe that similar to the learning procedure in Yang et al. (2020, 2022), at each round t = 1 , 2 , · · · , for each task m ∈ [ M ] , the learner selects a left and right action x m,t ∈ X and z m,t ∈ Z . After the player commits the batch of actions for each task { x m,t , z m,t : m ∈ [ M ] } , it receives the batch of rewards { r m,t : m ∈ [ M ] } . Also note that in (1) we define the ˜ g m,t ∈ R k 1 , ˜ v m,t ∈ R k 2 as the latent features, and both ˜ g m,t , ˜ v m,t 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 x m, ∗ and right arm z m, ∗ for each task m with a minimum number of samples (fixed confidence setting). First, consider a single-task setting and let Θ ∗ have low rank r . Let the SVD of the Θ ∗ = UDV ⊤ . Prima-facie, if U and V are known then one might want to project all the left and right arms in the r × r subspace of U and V and reduce the bilinear bandit problem into a r 2 dimension linear bandit setting. Then one can apply one of the algorithms from Soare et al. (2014); Fiez et al. (2019); Katz-Samuels et al. (2020) to solve this r 2 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 ˜ O ( r 2 / ∆ 2 ) is possible where ∆ is the minimum reward gap and ˜ O ( · ) 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 r 2 dimensional linear bandit problems with shared representation and achieve a sample complexity bound of ˜ O ( Mr 2 / ∆ 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 r 2 dimensional linear bandit sample complexity analysis. Following the regret analysis of bilinear bandit setting by Jun et al. (2019); Lu et al. (2021); Kang et al. (2022) we know that the effective dimension is actually ( d 1 + d 2 ) r . Similarly for the multi-task representation learning the effective dimension should scale with the learned latent features ( k 1 + k 2 ) r . Hence the natural questions to ask are these:

1) Can we design a single-task pure exploration bilinear bandit algorithm whose sample complexity scales as ˜ O (( d 1 + d 2 ) r/ ∆ 2 ) ?

2) Can we design an algorithm for multi-task pure exploration bilinear bandit problem that can learn the latent features and has sample complexity that scales as ˜ O ( M ( k 1 + k 2 ) r/ ∆ 2 ) ?

In this paper, we answer both these questions affirmatively. In doing so, we make the following novel contributions to the growing literature of multi-task representation learning in online settings:

We formulate the multi-task bilinear representation learning problem. To our knowledge, this is the first work that explores pure exploration in a multi-task bilinear representation learning setting.

We proposed the algorithm GOBLIN for a single-task pure exploration bilinear bandit setting whose sample complexity scales as ˜ O (( d 1 + d 2 ) r/ ∆ 2 ) . This improves over RAGE (Fiez et al., 2019) whose sample complexity scales as ˜ O (( d 1 d 2 ) / ∆ 2 ) .

Our algorithm GOBLIN for multi-task pure exploration bilinear bandit problem learns the latent features and has sample complexity that scales as ˜ O ( M ( k 1 + k 2 ) r/ ∆ 2 ) . This improves over DouExpDes (Du et al., 2023) whose samples complexity scales as ˜ O ( M ( k 1 k 2 ) / ∆ 2 ) .

̸

̸

Preliminaries: We assume that ∥ x ∥ 2 ≤ 1 , ∥ z ∥ 2 ≤ 1 , ∥ Θ ∗ ∥ F ≤ S 0 and the r -th largest singular value of Θ ∗ ∈ R d 1 × d 2 is S r . Let p := d 1 d 2 denote the ambient dimension, and k = ( d 1 + d 2 ) r denote the effective dimension. Let [ n ] := { 1 , 2 , . . . , n } . Let x ∗ , z ∗ := arg max x , z x ⊤ Θ ∗ z . For any x , z define the gap ∆( x , z ) := x ⊤ ∗ Θ ∗ z ∗ -x ⊤ Θ ∗ z and furthermore ∆ = min x = x ∗ , z = z ∗ ∆( x , z ) . Similarly, for any arbitrary vector w ∈ W define the gap of w ∈ R p as ∆( w ) := ( w ∗ -w ) ⊤ θ ∗ , for some θ ∗ ∈ R p and furthermore, ∆ = min w = w ∗ ∆( w ) . If A ∈ R d × d ≥ 0 is a positive semidefinite matrix, and w ∈ R p is a vector, let ∥ w ∥ 2 A := w ⊤ Aw denote the induced semi-norm. Given any vector b ∈ R |W| we denote the w -th component as b w . Let ∆ W := { b ∈ R |W| : b w ≥ 0 , ∑ w ∈W b w = 1 } denote the set of probability distributions on W . We define Y ( W ) = { w -w ′ : ∀ w , w ′ ∈ W , w = w ′ } as the directions obtained from the differences between each pair of arms and Y ∗ ( W ) = { w ∗ -w : ∀ w ∈ W\ 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 a 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 t = 1 , 2 , . . . the learner observes the reward r t = x ⊤ t Θ ∗ z t + η t where the low rank hidden parameter Θ ∗ ∈ R d 1 × d 2 is unknown to the learner, x t ∈ R d 1 , z t ∈ R d 2 are visible to the learner, and η t is a 1 -sub-Gaussian noise. We assume that the matrix Θ ∗ has a low rank r which is known to the learner and d 1 , d 2 ≫ r . Finally recall that the goal is to identify the optimal left and right arms x ∗ , z ∗ with a minimum number of samples.

Wepropose a phase-based, two-stage arm elimination algorithm called G -O ptimal Design for B i lin ear Bandits (abbreviated as GOBLIN). GOBLIN proceeds in phases indexed by ℓ = 1 , 2 , . . . As this is a pure-exploration 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 τ E ℓ rounds, and pure exploration in rotated arms stage that runs for τ G ℓ rounds. We will define τ E ℓ in Section 2.1, while rotated arms and τ G ℓ 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 are remaining. Now we discuss the individual stages that occur at every phase ℓ of GOBLIN.

2.1 Estimating Subspaces of Θ ∗ (Stage 1 of the ℓ -th phase)

In the first stage of phase ℓ , GOBLIN estimates the row and column sub-spaces Θ ∗ . Then GOBLIN uses these estimates to reduce the bilinear bandit problem in the original ambient dimension p := d 1 d 2 to a lower effective dimension k := ( d 1 + d 2 ) r . To do this, GOBLIN first vectorizes the x ∈ R d 1 , z ∈ R d 2 into a new vector w ∈ R p and then solves the E -optimal design in Step 3 of

̸

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 ℓ -th phase be denoted by b E ℓ . Then GOBLIN samples each w for ⌈ τ E ℓ b E ℓ, w ⌉ times, where τ E ℓ = ˜ O ( √ d 1 d 2 r/S r ) (step 7 of Algorithm 1). In this paper, we sample an arm ⌈ τ E ℓ b E ℓ, w ⌉ number of times. However, this may lead to over-sampling of an arm than what the design ( G or E -optimal) 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 Θ ∗ in stage 1 of phase ℓ . GOBLIN estimates this by solving the following well-defined regularized minimization problem with nuclear norm penalty:

where Q ( · ) , ˜ ψ ν ( · ) , are appropriate functions stated in Definition 1, 3 respectively in Appendix A.3. The Q ( · ) function takes as input the rank-one matrix x s z ⊤ s which is obtained after reshaping w s . Note that x s , and z s are the observed vectors in d 1 and d 2 dimension and ̂ Θ ℓ ∈ R d 1 × d 2 Finally, set the regularization parameter γ ℓ := 4 √ 2 ( 4+ S 2 0 ) Cd 1 d 2 log(2( d 1 + d 2 ) /δ ) τ E ℓ . This is in step 8 of Algorithm 1.

2.2 Optimal Design for Rotated Arms (Stage 2 of ℓ -th phase)

In stage 2 of phase ℓ , GOBLIN leverages the information about the learned sub-space of Θ ∗ to rotate the arm set and then run the optimal design on the rotated arm set. Once we 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 d 1 d 2 . In contrast GOBLIN uses the information about the learned sub-space of Θ ∗ to reduce the problem from ambient dimension d 1 d 2 to effective dimension ( d 1 + d 2 ) r . This reduction is done as follows: Let ̂ Θ ℓ = ̂ U ℓ ̂ D ℓ ̂ V ⊤ ℓ be the SVD of ̂ Θ ℓ in the ℓ -th phase. Let ̂ U ℓ ⊥ and ̂ V ℓ ⊥ be orthonormal bases of the complementary subspaces of ̂ U ℓ and ̂ V ℓ respectively. Let X ℓ and Z ℓ be the active set of arms in the stage 2 of phase ℓ . Then rotate the arm sets such that new rotated arm sets are as follows:

Let ̂ H ℓ = [ ̂ U ℓ ̂ U ⊥ ℓ ] ⊤ ̂ Θ ℓ [ ̂ V ℓ ̂ V ⊥ ℓ ] . Then define vectorized arm set so that the last ( d 1 -r ) · ( d 2 -r ) components are from the complementary subspaces as follows:

which implies ∥ ̂ θ k +1: p ∥ 2 = O ( d 1 d 2 r/τ E ℓ ) by Lemma 3 in Appendix A.1. So the last p -k components of ̂ θ ℓ are very small compared to the first k components. Hence, GOBLIN has now reduced the d 1 d 2 dimensional linear bandit to ( d 1 + d 2 ) r dimensional linear bandit using (3), (4). This is shown in step 10 of Algorithm 1.

Now in stage 2 of phase ℓ , GOBLIN implements G -optimal design (Pukelsheim, 2006; Fiez et al., 2019) in the rotated arm set X ℓ , Z ℓ defined in (3). To do this, first GOBLIN defines the rotated vector w = [ x 1: d 1 ; z 1: d 2 ] ∈ R p that belong to the set W ℓ . Then GOBLIN solves the G -optimal design (Pukelsheim, 2006) as follows:

This is shown in step 11 of Algorithm 1 and Λ ℓ is defined in (6). It can be shown that sampling according to ̂ b G ℓ 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 supp ( ̂ b G ℓ ) ≤ ˜ O ( 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 d 1 d 2 dimension then the support of ̂ b G ℓ will scale with d 1 d 2 which will lead to higher sample complexity. Then GOBLIN samples each w ∈ W ℓ for ⌈ τ G ℓ b G ℓ, w ⌉ times, where τ G ℓ := ⌈ 8 B ℓ ∗ ρ G ( Y ( W ℓ )) log(4 ℓ 2 |W| /δ ) ϵ 2 ℓ ⌉ . Note that the total length of phase ℓ , combining stages 1 and 2 is ( τ E ℓ + τ G ℓ ) rounds. Observe that the stage 1 design is on the whole arm set W whereas stage 2 design is on the refined active set W .

ℓ

Let the observed features in stage 2 of phase ℓ be denoted by W ℓ ∈ R τ G ℓ × p , and r ℓ ∈ R τ G ℓ be the observed rewards. Define the diagonal matrix Λ ℓ as

where, λ ⊥ ℓ := τ G ℓ -1 / 8 k log(1 + τ G ℓ -1 /λ ) ≫ λ . Deviating from Soare et al. (2014); Fiez et al. (2019) GOBLIN constructs a regularized least square estimator at phase ℓ as follows

This regularized least square estimator in (7) forces the last p -k components of ̂ θ ℓ to be very small compared to the first k components. Then GOBLIN builds the estimate ̂ θ ℓ from (7) only from the observations from this phase (step 13 in Algorithm 1) and eliminates sub-optimal actions in step 14 in Algorithm 1 using the estimator ̂ θ ℓ . Finally GOBLIN eliminates sub-optimal arms to build the next phase active set W ℓ and stops when |W ℓ | = 1 . GOBLIN outputs the arm in W ℓ and reshapes it to get the ̂ x ∗ and ̂ z ∗ . 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 δ , rank r of Θ ∗ , spectral bound S r of Θ ∗ , S, S ⊥ ℓ := 8 d 1 d 2 r τ E ℓ S 2 r log ( d 1 + d 2 δ ℓ ) , λ, λ ⊥ ℓ := τ G ℓ -1 / 8( d 1 + d 2 ) r log(1+ τ G ℓ -1 λ ) . Let p := d 1 d 2 , k := ( d 1 + d 2 ) r .

3: Define a vectorized arm w := [ x 1: d 1 ; z 1: d 2 ] and w ∈ W . Let τ E ℓ := 8 d 1 d 2 r log(4 ℓ 2 |W| /δ ℓ ) S r . Let the E -optimal design be b E ℓ := arg min b ∈△ W ∥ ∥ ( ∑ w ∈W b w ww ⊤ ) -1 ∥ ∥ .

2: Let W 1 ←W , ℓ ← 1 , τ G 0 := log(4 ℓ 2 |X| /δ ) . Define Λ ℓ as in (6), B ℓ ∗ := (8 √ λS + √ λ ⊥ ℓ S ⊥ ℓ ) . √

4: while |W ℓ | > 1 do

5: ϵ ℓ = 2 -ℓ , δ ℓ = δ/ℓ 2 .

6: (Stage 1:) Explore the Low-Rank Subspace

7: Pull arm w ∈ W exactly ⌈ ̂ b E ℓ, w τ E ℓ ⌉ times and observe rewards r t , for t = 1 , . . . , τ E ℓ .

8: Compute ̂ Θ ℓ using (2).

9: (Stage 2:) Reduction to low dimensional linear bandits

10: Let the SVD of ̂ Θ ℓ = ̂ U ℓ ̂ D ℓ ̂ V ⊤ ℓ . Rotate arms in active set W ℓ -1 to build W ℓ following (4).

and construct the least squares estimator ̂ θ ℓ using only the observations of this phase where ̂ θ ℓ is defined in (7). Note that ̂ θ ℓ is also rotated following (4).

14: Eliminate arms such that W ℓ +1 ←W ℓ \{ w ∈ W ℓ : max w ′ ∈W ℓ ⟨ w ′ -w , ̂ θ ℓ ⟩ > 2 ϵ ℓ }

15: ℓ ← ℓ +1

16: Output the arm in W ℓ and reshape to get the ̂ x ∗ and ̂ z ∗

2.3 Sample Complexity Analysis of Single-Task GOBLIN

We now analyze the sample complexity of GOBLIN in the single-task setting through the following theorem.

Theorem 1. (informal) With probability at least 1 -δ , GOBLIN returns the best arms x ∗ , z ∗ , and the number of samples used is bounded by ˜ O ( ( d 1 + d 2 ) r ∆ 2 + √ d 1 d 2 r S r ) .

Discussion 1. In Theorem 1 the first quantity is the number of samples needed to identify the best arms x ∗ , z ∗ while the second quantity is the number of samples to learn Θ ∗ (which is required to find the best arms). Note that the magnitude of S r would be free of d 1 , d 2 since Θ ∗ contains only r nonzero singular values and ∥ Θ ∗ ∥ ≤ 1 , and hence we assume that S r = Θ(1 / √ r ) (Kang et al., 2022). So the sample complexity of single-task GOBLIN scales as ˜ O ( ( d 1 + d 2 ) r ∆ 2 ) . However, if one runs RAGE (Fiez et al., 2019) on the arms in X , Z then the sample complexity will scale as ˜ O ( d 1 d 2 ∆ 2 ) .

Proof (Overview) of Theorem 1: Step 1 (Subspace estimation in high dimension): We denote the vectorized arms in high dimension as w ∈ 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 et al. (2022) which enables us to apply the Lemma 3 in Appendix A.1. This leads to ∥ ̂ Θ ℓ -Θ ∗ ∥ 2 F ≤ C 1 d 1 d 2 r log(2( d 1 + d 2 ) /δ ) τ E ℓ for some C 1 > 0 . Also, note that in the first stage of the ℓ -th phase by setting τ E ℓ = √ 8 d 1 d 2 r log(4 ℓ 2 |W| /δ ℓ ) S r and sampling each arm w ∈ W exactly ⌈ ̂ b E ℓ, w τ E ℓ ⌉ times we are guaranteed that ∥ θ ∗ k +1: p ∥ 2 = O ( d 1 d 2 r/τ E ℓ ) . Summing up over ℓ = 1 to ⌈ log 2 ( 4∆ -1 )⌉ we get that the total sample complexity of the first stage is bounded by ˜ O ( √ d 1 d 2 r/S r ) .

Step 2 (Effective dimension for rotated arms): We rotate the arms w ∈ W in high dimension to get the rotated arms w ∈ W ℓ in step 10 of Algorithm 1. Then we show that the effective dimension of w scales 8 k log ( 1 + τ G ℓ -1 /λ ) when λ ⊥ ℓ = τ G ℓ -1 8 k log ( 1+ τ ℓ G -1 /λ ) in Lemma 7 of Appendix A.4. Note that this requires a different proof technique than Valko et al. (2014) where the budget n 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-Samuels et al. (2020) as there is no parameter λ ⊥ ℓ 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 1 ≤ ρ G ( Y ( W ℓ )) ≤ p/γ 2 Y where, γ Y = max { c > 0 : c Y ⊂ conv( W ∪ -W ) } is the gauge norm of Y (Rockafellar, 2015). Note that this is a worst-case dependence when ρ G ( Y ( W ℓ )) scales with p . Substituting this value of ρ G ( Y ( W ℓ )) in the definition of λ ⊥ ℓ we can show that Λ ℓ 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 ̂ b G ℓ is equivalent to D -optimal design ̂ b D ℓ = arg max b log ∣ ∣ ∣ ∑ w ∈W ℓ b w w w ⊤ + Λ ℓ ∣ ∣ ∣ | Λ ℓ | . Then using Frank-Wolfe algorithm (Jamieson and Jain, 2022) we can show the support ̂ b G ℓ or equivalently ̂ b D ℓ is bounded by at most 8 k log(1+ τ G ℓ -1 /λ )(8 k log(1+ τ G ℓ -1 /λ )+1) 2 . This is shown in Lemma 9 (Appendix A.4).

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 τ G ℓ = ⌈ 8 B ℓ ∗ ρ ( Y ( W ℓ )) log(2 |W| /δ ) ( x ⊤ ( ̂ θ ℓ -θ ∗ )) 2 ⌉ . This is discussed in 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 ⌈ ̂ b ℓ, w τ G ℓ ⌉ 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 ( k γ 2 Y log( k log 2 (∆ -1 ) |W| δ ) ⌈ log 2 (∆ -1 ) ⌉ ) where the effective dimension k = ( d 1 + d 2 ) r . Finally, we combine the total samples of phase ℓ as ( τ E ℓ + τ G ℓ ) . The final sample complexity is given by summing over all phases from ℓ = 1 to ⌈ log 2 ( 4∆ -1 )⌉ . The claim of the theorem follows by noting ˜ O ( k/γ 2 Y ) ≤ ˜ O ( k/ ∆ 2 ) .

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 a reward model stated

in (1). The learning proceeds as follows: At each round t = 1 , 2 , · · · , for each task m ∈ [ M ] , the learner selects a left and right action x m,t ∈ X and z m,t ∈ Z . After the player commits the batch of actions for each task { x m,t , z m,t : m ∈ [ M ] } , it receives the batch of rewards { r m,t : m ∈ [ M ] } . Finally recall that the goal is to identify the optimal left and right arms x m, ∗ , z m, ∗ for each task m with a minimum number of samples. We now state the following assumptions to enable representation learning across tasks.

Assumption 1. (Low-rank Tasks) We assume that the hidden parameter Θ m, ∗ for all the m ∈ [ M ] have a decomposition Θ m, ∗ = B 1 S m, ∗ B ⊤ 2 and each S m, ∗ 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 σ min ( 1 M ∑ M m =1 Θ m, ∗ ) ≥ c 0 S r , for some c 0 > 0 , S r is the r -th largest singular value of Θ m, ∗ and σ min ( A ) denotes the minimum eigenvalue of matrix A .

This assumption is similar to the diverse tasks assumption of Yang et al. (2020, 2022); Tripuraneni et al. (2021); Du et al. (2023) and ensures the possibility of recovering the feature extractors B 1 and B 2 shared across tasks.

Our extension of GOBLIN to the multi-task setting is now a phase-based, three-stage arm elimination algorithm. In GOBLIN each phase ℓ = 1 , 2 , . . . consists of three stages; the stage for estimation of feature extractors B 1 , B 2 , which runs for τ E ℓ rounds, the stage for estimation of S m, ∗ which runs for ∑ m ˜ τ E m,ℓ rounds, and a stage of pure exploration with rotated arms that runs for ∑ m τ G m,ℓ rounds. We will define τ E m,ℓ in Section 3.1, ˜ τ E m,ℓ in Section 3.2, while the rotated arms and τ G m,ℓ are defined in Section 3.3. 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 are remaining. Now we discuss the individual stages that occur at every phase ℓ = 1 , 2 , . . . for multi-task GOBLIN.

3.1 Estimating Feature Extractors B 1 and B 2 (Stage 1 of Phase ℓ )

In the first stage of phase ℓ , GOBLIN leverages the batch of rewards { r m,t : m ∈ [ M ] } at every round t from M tasks to learn the feature extractors B 1 and B 2 . To do this, GOBLIN first vectorizes the x ∈ X , z ∈ Z into a new vector w = [ x 1: d 1 ; z 1: d 2 ] ∈ W m and then solves the E -optimal design in step 3 of Algorithm 2. Similar to the single-task setting (Section 2) GOBLIN samples each w ∈ W m for ⌈ τ E ℓ b E ℓ, w ⌉ times for each task m , where τ E ℓ = ˜ O ( √ d 1 d 2 r/S r ) and b E ℓ, w is the solution to E -optimal design on w . Let the sampled arms for each task m at round s be denoted by x m,s , z m,s which is obtained after reshaping w s . Then it builds the estimator ̂ Z ℓ as follows:

Then it performs SVD decomposition on ̂ Z ℓ , and let ̂ B 1 , ̂ B 2 be the topk 1 and topk 2 left and right singular vectors of ̂ Z ℓ respectively. These are the estimation of the feature extractors B 1 and B 2 .

3.2 Estimating Hidden Parameter S m, ∗ per Task (Stage 2 of phase ℓ )

In the second stage of phase ℓ , the goal is to recover the hidden parameter S m, ∗ for each task m . GOBLIN proceeds as follows: First, let ˜ g m = x ⊤ ̂ B 1 ,ℓ and ˜ v m = z ⊤ ̂ B 2 ,ℓ be the latent left and right arm respectively for each m . Then GOBLIN defines the vector ˜ w = [ ˜ g m ; ˜ v m ] ∈ ˜ W m and then solves the E -optimal design in step 11 of Algorithm 2. It then samples for each task m , the latent arm ˜ w ∈ ˜ W m for ⌈ ˜ τ E m,ℓ ˜ b E m,ℓ, ˜ w ⌉ times, where ˜ τ E m,ℓ := ˜ O ( √ k 1 k 2 r/S r ) and ˜ b E m,ℓ, ˜ w is the solution to

E -optimal design on ˜ w . Then it builds estimator ̂ S m,ℓ for each task m in step 12 as follows:

Once GOBLIN recovers the ̂ S m,ℓ for each task m it has reduced the d 1 d 2 bilinear bandit to a k 1 k 2 dimension bilinear bandit where the left and right arms are ˜ g m ∈ G m , ˜ v m ∈ V m respectively.

3.3 Optimal Design for Rotated Arms per Task (Stage 3 of phase ℓ )

In the third stage of phase ℓ , similar to Algorithm 1, the multi-task GOBLIN defines the rotated arm set G m , V m for each task m for these k 1 k 2 dimensional bilinear bandits. Let the SVD of ̂ S m,ℓ = ̂ U m,ℓ ̂ D m,ℓ ̂ V ⊤ m,ℓ . Define ̂ H m,ℓ = [ ̂ U m,ℓ ̂ U ⊥ m,ℓ ] ⊤ ̂ S m,ℓ [ ̂ V m,ℓ ̂ V ⊥ m,ℓ ] . Then define the vectorized arm set so that the last ( k 1 -r ) · ( k 2 -r ) components are from the complementary subspaces as follows:

˜

g

m,r

+1:

k

1

˜

v

⊤

m,r

+1:

k

2

This is shown in step 14 of Algorithm 2. Now we proceed similarly to Section 2.2. We construct a per-task optimal design for the rotated arm set V m , G m and define the w = [ ˜ g m, 1: d 1 ; ˜ v m, 1: d 2 ] and ˜ w ∈ ˜ W m where ˜ g m ∈ G m and ˜ v m ∈ V m respectively. Following (5) we know that to minimize the sample complexity for the m -th bilinear bandit we need to sample according to G -optimal design

Then GOBLIN runs G -optimal design on the arm set W ℓ following the (11) and then samples each w ∈ W m,ℓ for ⌈ τ G m,ℓ ̂ b G m,ℓ, w ⌉ times where ̂ b G m,ℓ, w is the solution to the G -optimal design, and τ G ℓ is defined in step 17 of Algorithm 2. So the total length of phase ℓ , combining stages 1 , 2 and 3 is ( τ E ℓ + ∑ m ˜ τ E m,ℓ + ∑ m τ G m,ℓ ) rounds. Observe that the stage 1 and 2 design is on the whole arm set W , ˜ W m whereas the stage 3 design is on the refined active set W m,ℓ . Let at the stage 3 of ℓ -th phase the actions sampled be denoted by the matrix W m,ℓ ∈ R τ G m,ℓ × k 1 k 2 and observed rewards r m ∈ R τ G m,ℓ × k 1 k 2 . Define the positive diagonal matrix Λ m,ℓ according to (6) but set p = k 1 k 2 and k = ( k 1 + k 2 ) r . Then similar to Section 2.2 we can build for each task m only from the observations from this phase

Finally GOBLIN eliminates the sub-optimal arms using the estimator ̂ θ m,ℓ to build the next phase active set W m,ℓ and stops when |W m,ℓ | = 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 2. (informal) With probability at least 1 -δ , GOBLIN returns the best arms x m, ∗ , z m, ∗ for each task m , and the total number of samples is bounded by ˜ O ( M ( k 1 + k 2 ) r ∆ 2 + M √ k 1 k 2 r S r + √ d 1 d 2 r S r ) .

Discussion 2. In Theorem 2 the first quantity is the sample complexity to identify the best arms x m, ∗ , z m, ∗ and the second quantity is the number of samples to learn S m, ∗ for each task m . This is required to rotate the arms to reach the effective dimension of ( k 1 + k 2 ) r . Finally, the third quantity is the number of samples needed to learn Θ m, ∗ (which in turn is used to estimate the feature extractors B 1 and B 2 to learn the S m, ∗ ). Again we assume that S r = Θ(1 / √ r ) (Kang et al., 2022). So the sample complexity of multi-task GOBLIN scales as ˜ O ( M ( k 1 + k 2 ) r/ ∆ 2 ) . However, if one runs DouExpDes (Du et al., 2023) then the sample complexity will scale as ˜ O ( M ( k 1 k 2 ) / ∆ 2 ) which is worse than GOBLINwhen r ≪ k 1 or k 2 .

(

)]}

Algorithm 2 G-Optimal Design for Bilinear Bandits (GOBLIN) for multi-task setting

1: Input: arm set X , Z , confidence δ , rank r of Θ ∗ , spectral bound S r of Θ ∗ , S, S ⊥ m,ℓ = 8 k 1 k 2 r ˜ τ E m,ℓ S 2 r log( k 1 + k 2 δ ℓ ) , λ, λ ⊥ m,ℓ = τ G m,ℓ -1 (8( k 1 + k 2 ) r log(1+ τ G m,ℓ -1 /λ )) . Let p = k 1 k 2 , k = ( k 1 + k 2 ) r . √

3: Define arm w = [ x 1: d 1 ; z 1: d 2 ] and w ∈ W m . Let τ E ℓ = 8 d 1 d 2 r log(4 ℓ 2 |W| /δ ℓ ) S r . Let E -optimal design be b E ℓ =argmin b ∈△ W ∥ ∥ ( ∑ w ∈W b w ww ⊤ ) -1 ∥ ∥ .

2: Let W m, 1 ← W m , ℓ ← 1 , τ G 0 = log(4 ℓ 2 |X| /δ ) . Define Λ m,ℓ as in (6), B ℓ m, ∗ := (8 λS + √ λ ⊥ m,ℓ S ⊥ m,ℓ ) √

4: while ∃ m ∈ [ M ] , ∣ ∣ W m,ℓ ∣ ∣ > 1 do

5: ϵ ℓ = 2 -ℓ , δ ℓ = δ/ℓ 2 .

6: (Stage 1:) Explore the Low-Rank Subspace

7: Pull arm w ∈ W exactly ⌈ ̂ b E ℓ, w τ E ℓ ⌉ times for each task m and observe rewards { r m,t } τ E ℓ t =1 .

8: Compute ̂ Z ℓ using (8).

9: (Stage 2:) Build ̂ S m,ℓ for each task m

10: Let ̂ B 1 ,ℓ , ̂ B 2 ,ℓ be the topk 1 left and topk 2 right singular vectors of ̂ Z ℓ respectively. Build ˜ g m = x ⊤ ̂ B 1 ,ℓ and ˜ v m = z ⊤ ̂ B 2 ,ℓ for all x ∈ X and z ∈ Z for each m .

12: Pull arm ˜ w ∈ ˜ W m exactly ⌈ ˜ b E m,ℓ, ˜ w ˜ τ E m,ℓ ⌉ times and observe rewards r m,t , for t = 1 , . . . , ˜ τ E m,ℓ , for each task m . Then compute ̂ S m,ℓ using (9) for each m .

11: Define a vectorized arm ˜ w = [ ˜ g m, 1: k 1 ; ˜ v m, 1: k 2 ] and ˜ w ∈ ˜ W m for each m . Let ˜ τ E m,ℓ = √ 8 k 1 k 2 r log(4 ℓ 2 |W| /δ ℓ ) S r , and ˜ b E m,ℓ = arg min b m ∈△ ˜ W m ∥ ∥ ( ∑ ˜ w ∈ ˜ W m b m, ˜ w ˜ w ˜ w ⊤ ) -1 ∥ ∥ .

13: (Stage 3:) Reduction to low dimensional linear bandits for each task m

14: SVD of ̂ S m,ℓ = ̂ U m,ℓ ̂ D m,ℓ ̂ V ⊤ m,ℓ . Rotate arms in active set W m,ℓ -1 to build W m,ℓ using (10).

15: Let ̂ b G m,ℓ =arg min b m, w max w , w ′ ∈W m,ℓ ∥ w -w ′ ∥ 2 ( ∑ w m ∈W m b m, w w m w ⊤ m + Λ m,ℓ /n ) -1 .

16: Define ρ G ( Y ( W m,ℓ ))=min b m, w max w , w ′ ∈W m,ℓ ∥ w -w ′ ∥ 2 ( ∑ w ∈W m b m, w w w ⊤ + Λ m,ℓ n ) -1 .

17: Set τ G m,ℓ = 64 B ℓ m, ∗ ρ G ( Y ( W m,ℓ )) log(4 ℓ 2 |W m | /δ ℓ ) ϵ 2 ℓ . Then pull arm w ∈ W m for each task m exactly ⌈ ̂ b m,ℓ, w τ G m,ℓ ⌉ times and construct the least squares estimator ̂ θ m,ℓ using only the observations of this phase where ̂ θ m,ℓ is defined in (12).

18: Eliminate arms such that W m,ℓ +1 ←W m,ℓ \ { w m ∈ W m,ℓ : max w ′ m ∈W m,ℓ 〈 w ′ m -w m , ̂ θ m,ℓ 〉 > 2 ϵ m,ℓ } 19: ℓ ← ℓ +1

20: Output the arm in W m,ℓ and reshape to get the ̂ x m, ∗ and ̂ z m, ∗ for each task m .

Proof (Overview) of Theorem 2: Step 1 (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 ∗ = 1 M ∑ M m =1 Θ ∗ ,m using (8). This requires us to modify the Lemma 3 in Appendix A.1 and apply Stein's lemma (Lemma 1) to get a bound of ∥ ̂ Z ℓ -Z ∗ ∥ 2 F ≤ C 1 d 1 d 2 r log(2( d 1 + d 2 ) /δ ) τ E ℓ for some C 1 > 0 . This is shown in Lemma 12 in Appendix A.6. Summing up over ℓ = 1 to ⌈ log 2 ( 4∆ -1 )⌉ we get that the total samples complexity of the first stage is bounded by ˜ O ( √ d 1 d 2 r/S r ) .

Step 2 (Estimation of left and right feature extractors): Now using the estimator in (8) we get a good estimation of the feature extractors B 1 and B 2 . Let ̂ B 1 ,ℓ , ̂ B 2 ,ℓ be the topk 1 left and topk 2 right singular vectors of ̂ Z ℓ respectively. Then using the Davis-Kahan sin θ Theorem (Bhatia, 2013) in Lemma 14, 15 (Appendix A.6) we have ∥ ( ̂ B ⊥ 1 ,ℓ ) ⊤ B 1 ∥ , ∥ ( ̂ B ⊥ 2 ,ℓ ) ⊤ B 2 ∥ ≤ ˜ O ( √ ( d 1 + d 2 ) r/Mτ E ℓ ) .

Step 3 (Estimation of ̂ S m,ℓ in low dimension): Now we estimate the quantity ̂ S m,ℓ ∈ R k 1 × k 2 for each task m . To do this we first build the latent arms ˜ g m = x ⊤ ̂ U ℓ and ˜ v m = z ⊤ ̂ V ℓ for all x ∈ X

and z ∈ 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 σ min ( ∑ ˜ w ∈ ˜ W b ˜ w ˜ w ˜ w ⊤ ) > 0 which enables us to sample following E -optimal design. Then use the estimator in (9). Then in Lemma 19 we show that ∥ ̂ S m,ℓ -µ ∗ S m, ∗ ∥ 2 F ≤ C 1 k 1 k 2 r log ( 2( k 1 + k 2 ) δ ℓ ) /τ E m,ℓ holds with probability greater than (1 -δ ) . Also, note that in the second phase by setting ˜ τ E m,ℓ = √ 8 k 1 k 2 r log(4 ℓ 2 |W| /δ ℓ ) /S r and sampling each arm w ∈ W exactly ⌈ ̂ b E ℓ, w ˜ τ E m,ℓ ⌉ times we are guaranteed that ∥ θ ∗ k +1: p ∥ 2 = O ( k 1 k 2 r/ ˜ τ E m,ℓ ) in the ℓ -th phase. Summing up over ℓ = 1 to ⌈ log 2 ( 4∆ -1 )⌉ across each task M we get that the total samples complexity of the second stage is bounded by ˜ O ( M √ k 1 k 2 r/S r ) .

Step 4 (Convert to k 1 k 2 bilinear bandits): Once GOBLIN recovers ̂ S m,τ E ℓ it rotates the arm set following (10) to build W m to get the k 1 k 2 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 ( k γ 2 Y log( k log 2 (∆ -1 ) |W| δ ) ⌈ log 2 (∆ -1 ) ⌉ ) where the effective dimension k = ( k 1 + k 2 ) r . The total samples of phase ℓ is given by τ E ℓ + ∑ m ( ˜ τ E m,ℓ + τ G m,ℓ ) . Finally, we get the total sample complexity by summing over all phases from ℓ =1 to ⌈ log 2 ( 4∆ -1 ) ⌉ . The claim of the theorem follows by noting ˜ O ( k/γ 2 Y ) ≤ ˜ O ( k/ ∆ 2 ) .

4 Experiments

In this section, we conduct proof-of-concept experiments on both single and multi-task bilinear bandits. In the single-task experiment, we compare against the state-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. In 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 a central contribution, we defer a fuller description of the experimental set-up to Appendix A.8.

5

le3

BanditSetting(LogPlot)

1e4

BanditSetting(Log Plot)

GOBLIN(ours)

1.6

GOBLIN(ours)

4

RAGE

1.4

DouExpDes

Samples

3

ples

S

1.2

1.0

w

0.8

S

2

S

0.6

1

0.4

0.2

6

10

14

5

10

1520

25

30

Arms

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 ˜ O (( d 1 + d 2 ) r/ ∆ 2 ) which improves upon the ˜ O (( d 1 d 2 ) / ∆ 2 ) sample complexity of RAGE (Fiez et al., 2019) in a singletask setting. We then extend GOBLIN for multi-task pure exploration bilinear bandit problems by learning latent features which enables sample complexity that scales as ˜ O ( M ( k 1 + k 2 ) r/ ∆ 2 ) which improves over the ˜ O ( M ( k 1 k 2 ) / ∆ 2 ) sample complexity of DouExpDes (Du et al., 2023). Our analysis opens an exciting opportunity to analyze representation learning in 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-linear representations.

References

Agarwal, D., Chen, B.-C., and Elango, P. (2009). Explore/exploit schemes for web content optimization. In 2009 Ninth IEEE International Conference on Data Mining , pages 1-10. IEEE.

Bengio, Y., Courville, A., and Vincent, P. (2013). Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence , 35(8):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, J. R., and Cardoso, M. J. (2018). Uncertainty in multitask learning: joint representations for probabilistic mr-only radiotherapy planning. In Medical Image Computing and Computer Assisted Intervention-MICCAI 2018: 21st International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part 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 , 35:19971-19983.

Katz-Samuels, J., Jain, L., Jamieson, K. G., et al. (2020). An empirical process approach to the union bound: Practical algorithms for combinatorial and linear bandits. Advances in Neural Information Processing Systems , 33: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, J., Zhang, H., Zhang, L., Huang, X., and Zhang, L. (2014). Joint collaborative representation with multitask learning for hyperspectral image classification. IEEE Transactions on Geoscience and Remote Sensing , 52(9):5923-5936.

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). A network integration approach for drug-target interaction prediction and computational drug repositioning from heterogeneous information. Nature communications , 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.01768 .

Maurer, A., Pontil, M., and Romera-Paredes, B. (2016). The benefit of multitask representation learning. Journal of Machine Learning Research , 17(81):1-32.

Minsker, S. (2018). Sub-gaussian estimators of the mean of a 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. In 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:2301.12357 .

Pukelsheim, F. (2006). Optimal design of experiments . SIAM.

Reyes, L. J. P., Oviedo, N. B., Camacho, E. C., and Calderon, J. M. (2021). Adaptable recommendation system for outfit selection with deep learning approach. IFAC-PapersOnLine , 54(13):605-610.

Rockafellar, R. (2015). Convex analysis. princeton landmarks in mathematics and physics.

Shamir, O. (2011). A variant of azuma's inequality for martingales with subgaussian tails. arXiv preprint arXiv:1110.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):e0283896.

Soare, M., Lazaric, A., and Munos, R. (2014). Best-arm identification in 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, T. (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. S. (2020). Impact of representation learning in linear bandits. arXiv preprint arXiv:2010.06531 .

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). Multi-modal multi-task learning for joint prediction of multiple regression and classification variables in alzheimer's disease. NeuroImage , 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):3294-3315.

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:11618-11630.

Peer reviews:
comments: 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.

STRENGTHS:
1.	This paper is well-written. The considered problem is interesting and well-motivated.
2.	The theoretical analysis looks sound.

WEAKNESSES:
1.	What is the main technical novelty of this work, beyond the existing SVD, G-optimal design, phased elimination methods in the bilinear/linear bandit literature?
2.	Could 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?

QUESTIONS:
Please see the weaknesses above.

LIMITATIONS:
Please see the weaknesses above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1) 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\nWEAKNESSES:\n1) 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\nQUESTIONS:\n1) 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\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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}$.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}