SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Subhojyoti Mukherjee
Wisconsin, Madison
smukherjee27@wisc.edu
Qiaomin Xie
ISyE Department
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.
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  $\mathbf{x}_t \in \mathbb{R}^{d_1}$  and  $\mathbf{z}_t \in \mathbb{R}^{d_2}$  at every round  $t$  and the interaction of this pair of arms with a low-rank hidden parameter,  $\Theta_* \in \mathbb{R}^{d_1 \times d_2}$  generates the noisy feedback (reward)  $r_t = \mathbf{x}_t^\top \Theta_* \mathbf{z}_t + \eta_t$ . The  $\eta_t$  is random 1-subGaussian noise.
A 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 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  $(\mathbf{x}_t,\mathbf{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,\dots ,M$  with each task having its own hidden parameter  $\Theta_{m,*}\in \mathbb{R}^{d_1\times d_2}$ . Let each  $\Theta_{m,*}$  has a decomposition of  $\Theta_{m,*} = \mathbf{B}_1\mathbf{S}_{m,*}\mathbf{B}_2^\top$ , where  $\mathbf{B}_1\in \mathbb{R}^{d_1\times k_1}$  and  $\mathbf{B}_2\in \mathbb{R}^{d_2\times k_2}$  are shared across tasks, but  $\mathbf{S}_{m,*}\in \mathbb{R}^{k_1\times k_2}$  is specific for task  $m$ . We assume that  $k_{1},k_{2}\ll d_{1},d_{2}$  and  $M\gg d_{1},d_{2}$ . Thus,  $\mathbf{B}_1$  and  $\mathbf{B}_2$  provide a means of dimensionality reduction. Furthermore, we assume that each  $\mathbf{S}_{m,*}$  has rank  $r\ll \min \{k_1,k_2\}$ . In the terminology of multi-task representation learning  $\mathbf{B}_1,\mathbf{B}_2$  are called feature extractors and  $\mathbf{x}_{m,t},\mathbf{z}_{m,t}$  are called rich observations (Yang et al., 2020, 2022; Du et al., 2023). The reward for the task  $m\in \{1,2,\dots ,M\}$  at round  $t$  is
$$
r _ {m, t} = \mathbf {x} _ {m, t} ^ {\top} \boldsymbol {\Theta} _ {m, *} \mathbf {z} _ {m, t} + \eta_ {m, t} = \underbrace {\mathbf {x} _ {m , t} ^ {\top} \mathbf {B} _ {1}} _ {\mathbf {g} _ {m, t} ^ {\top}} \mathbf {S} _ {m, *} \underbrace {\mathbf {B} _ {2} ^ {\top} \mathbf {z} _ {m , t}} _ {\mathbf {v} _ {m, t}} + \eta_ {m, t} = \mathbf {g} _ {m, t} ^ {\top} \mathbf {S} _ {m, *} \mathbf {v} _ {m, t} + \eta_ {m, t}. \tag {1}
$$
Observe that similar to the learning procedure in Yang et al. (2020, 2022), at each round  $t = 1, 2, \dots$ , for each task  $m \in [M]$ , the learner selects a left and right action  $\mathbf{x}_{m,t} \in \mathcal{X}$  and  $\mathbf{z}_{m,t} \in \mathcal{Z}$ . After the player commits the batch of actions for each task  $\{\mathbf{x}_{m,t}, \mathbf{z}_{m,t} : m \in [M]\}$ , it receives the batch of rewards  $\{r_{m,t} : m \in [M]\}$ . Also note that in (1) we define the  $\widetilde{\mathbf{g}}_{m,t} \in \mathbb{R}^{k_1}$ ,  $\widetilde{\mathbf{v}}_{m,t} \in \mathbb{R}^{k_2}$  as the latent features, and both  $\widetilde{\mathbf{g}}_{m,t}, \widetilde{\mathbf{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  $\mathbf{x}_{m,*}$  and right arm  $\mathbf{z}_{m,*}$  for each task  $m$  with a minimum number of samples (fixed confidence setting). First, consider a single-task setting and let  $\Theta_{*}$  have low rank  $r$ . Let the SVD of the  $\Theta_{*} = \mathbf{UDV}^{\top}$ . Prima-facie, if  $\mathbf{U}$  and  $\mathbf{V}$  are known then one might want to project all the left and right arms in the  $r\times r$  subspace of  $\mathbf{U}$  and  $\mathbf{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  $\widetilde{O}(r^2/\Delta^2)$  is possible where  $\Delta$  is the minimum reward gap and  $\widetilde{O}(\cdot)$  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  $\widetilde{O}(Mr^2/\Delta^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  $\widetilde{O}((d_1 + d_2)r / \Delta^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  $\widetilde{O}(M(k_1 + k_2)r / \Delta^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:
1) 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.
2) We proposed the algorithm GOBLIN for a single-task pure exploration bilinear bandit setting whose sample complexity scales as  $\widetilde{O}((d_1 + d_2)r/\Delta^2)$ . This improves over RAGE (Fiez et al., 2019) whose sample complexity scales as  $\widetilde{O}((d_1d_2)/\Delta^2)$ .
3) Our algorithm GOBLIN for multi-task pure exploration bilinear bandit problem learns the latent features and has sample complexity that scales as  $\widetilde{O}(M(k_1 + k_2)r / \Delta^2)$ . This improves over DouExpDes (Du et al., 2023) whose samples complexity scales as  $\widetilde{O}(M(k_1k_2) / \Delta^2)$ .
Preliminaries: We assume that  $\| \mathbf{x}\| _2\leq 1$ $\| \mathbf{z}\| _2\leq 1$ $\| \Theta_{*}\|_{F}\leq S_{0}$  and the  $r$ -th largest singular value of  $\Theta_{*}\in \mathbb{R}^{d_{1}\times d_{2}}$  is  $S_{r}$ . Let  $p\coloneqq d_1d_2$  denote the ambient dimension, and  $k = (d_{1} + d_{2})r$  denote the effective dimension. Let  $[n]\coloneqq \{1,2,\ldots ,n\}$ . Let  $\mathbf{x}_{*},\mathbf{z}_{*}\coloneqq$  arg  $\max_{\mathbf{x},\mathbf{z}}\mathbf{x}^{\top}\Theta_{*}\mathbf{z}$ . For any  $\mathbf{x},\mathbf{z}$  define the gap  $\Delta (\mathbf{x},\mathbf{z})\coloneqq \mathbf{x}_*^\top \Theta_*\mathbf{z}_* - \mathbf{x}^\top \Theta_*\mathbf{z}$  and furthermore  $\Delta = \min_{\mathbf{x}\neq \mathbf{x}_{*},\mathbf{z}\neq \mathbf{z}_{*}}\Delta (\mathbf{x},\mathbf{z})$ . Similarly, for any arbitrary vector  $\mathbf{w}\in \mathcal{W}$  define the gap of  $\mathbf{w}\in \mathbb{R}^p$  as  $\Delta (\mathbf{w})\coloneqq (\mathbf{w}_{*} - \mathbf{w})^{\top}\boldsymbol{\theta}_{*}$ , for some  $\boldsymbol{\theta}_{*}\in \mathbb{R}^{p}$  and furthermore,  $\Delta = \min_{\mathbf{w}\neq \mathbf{w}_{*}}\Delta (\mathbf{w})$ . If  $\mathbf{A}\in \mathbb{R}_{\geq 0}^{d\times d}$  is a positive semidefinite matrix, and  $\mathbf{w}\in \mathbb{R}^p$  is a vector, let  $\| \mathbf{w}\|_{\mathbf{A}}^{2}\coloneqq \mathbf{w}^{\top}\mathbf{A}\mathbf{w}$  denote the induced semi-norm. Given any vector  $\mathbf{b}\in \mathbb{R}^{|\mathcal{W}|}$  we denote the w-th component as  $\mathbf{b_w}$ . Let  $\Delta_{\mathcal{W}}\coloneqq \left\{\mathbf{b}\in \mathbb{R}^{|\mathcal{W}|}:\mathbf{b_w}\geq 0,\sum_{\mathbf{w}\in \mathcal{W}}\mathbf{b_w} = 1\right\}$  denote the set of probability distributions on  $\mathcal{W}$ . We define  $\mathcal{V}(\mathcal{W}) = \{\mathbf{w} - \mathbf{w}':\forall \mathbf{w},\mathbf{w}'\in \mathcal{W},\mathbf{w}\neq \mathbf{w}'\}$  as the directions obtained from the differences between each pair of arms and  $\mathcal{V}^{*}(\mathcal{W}) = \{\mathbf{w}_{*} - \mathbf{w}:\forall \mathbf{w}\in \mathcal{W}\backslash \mathbf{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,\ldots$  the learner observes the reward  $r_t = \mathbf{x}_t^\top \boldsymbol{\Theta}_*\mathbf{z}_t + \eta_t$  where the low rank hidden parameter  $\boldsymbol{\Theta}_* \in \mathbb{R}^{d_1 \times d_2}$  is unknown to the learner,  $\mathbf{x}_t \in \mathbb{R}^{d_1}$ ,  $\mathbf{z}_t \in \mathbb{R}^{d_2}$  are visible to the learner, and  $\eta_t$  is a 1-sub-Gaussian noise. We assume that the matrix  $\boldsymbol{\Theta}_*$  has a low rank  $r$  which is known to the learner and  $d_1, d_2 \gg r$ . Finally recall that the goal is to identify the optimal left and right arms  $\mathbf{x}_*$ ,  $\mathbf{z}_*$  with a minimum number of samples.
We propose a phase-based, two-stage arm elimination algorithm called G-Optimal Design for Bilinear Bandits (abbreviated as GOBLIN). GOBLIN proceeds in phases indexed by  $\ell = 1,2,\ldots$ . 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  $\ell$  of GOBLIN consists of two stages; the estimation of  $\Theta_{*}$  stage, which runs for  $\tau_{\ell}^{E}$  rounds, and pure exploration in rotated arms stage that runs for  $\tau_{\ell}^{G}$  rounds. We will define  $\tau_{\ell}^{E}$  in Section 2.1, while rotated arms and  $\tau_{\ell}^{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  $\ell$  of GOBLIN.
2.1 Estimating Subspaces of  $\Theta_{*}$  (Stage 1 of the  $\ell$ -th phase)
In the first stage of phase  $\ell$ , GOBLIN estimates the row and column sub-spaces  $\Theta_{*}$ . Then GOBLIN uses these estimates to reduce the bilinear bandit problem in the original ambient dimension  $p \coloneqq d_1d_2$  to a lower effective dimension  $k \coloneqq (d_1 + d_2)r$ . To do this, GOBLIN first vectorizes the  $\mathbf{x} \in \mathbb{R}^{d_1}$ ,  $\mathbf{z} \in \mathbb{R}^{d_2}$  into a new vector  $\overline{\mathbf{w}} \in \mathbb{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  $\ell$ -th phase be denoted by  $\mathbf{b}_{\ell}^{E}$ . Then GOBLIN samples each  $\overline{\mathbf{w}}$  for  $\left[\tau_{\ell}^{E}\mathbf{b}_{\ell,\overline{\mathbf{w}}}^{E}\right]$  times, where  $\tau_{\ell}^{E} = \widetilde{O}(\sqrt{d_{1}d_{2}r}/S_{r})$  (step 7 of Algorithm 1). In this paper, we sample an arm  $\left[\tau_{\ell}^{E}\mathbf{b}_{\ell,\overline{\mathbf{w}}}^{E}\right]$  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 efficient Rounding Procedures (see Pukelsheim (2006); Fiez et al. (2019)). Let  $\widehat{\Theta}_{\ell}$  be estimate of  $\Theta_{*}$  in stage 1 of phase  $\ell$ . GOBLIN estimates this by solving the following well-defined regularized minimization problem with nuclear norm penalty:
$$
\widehat {\boldsymbol {\Theta}} _ {\ell} = \underset {\boldsymbol {\Theta} \in \mathbb {R} ^ {d _ {1} \times d _ {2}}} {\arg \min } L _ {\ell} (\boldsymbol {\Theta}) + \gamma_ {\ell} \| \boldsymbol {\Theta} \| _ {\mathrm {n u c}}, \quad L _ {\ell} (\boldsymbol {\Theta}) = \langle \boldsymbol {\Theta}, \boldsymbol {\Theta} \rangle - \frac {2}{\tau_ {\ell} ^ {E}} \sum_ {s = 1} ^ {\tau_ {\ell} ^ {E}} \langle \widetilde {\psi} _ {\nu} (r _ {s} \cdot Q (\mathbf {x} _ {s} \mathbf {z} _ {s} ^ {\top})), \boldsymbol {\Theta} \rangle \tag {2}
$$
where  $Q(\cdot), \widetilde{\psi}_{\nu}(\cdot)$ , are appropriate functions stated in Definition 1, 3 respectively in Appendix A.3. The  $Q(\cdot)$  function takes as input the rank-one matrix  $\mathbf{x}_s\mathbf{z}_s^\top$  which is obtained after reshaping  $\overline{\mathbf{w}}_s$ . Note that  $\mathbf{x}_s$ , and  $\mathbf{z}_s$  are the observed vectors in  $d_{1}$  and  $d_{2}$  dimension and  $\widehat{\Theta}_{\ell} \in \mathbb{R}^{d_{1} \times d_{2}}$ . Finally, set the regularization parameter  $\gamma_{\ell} := 4\sqrt{\frac{2(4 + S_{0}^{2})Cd_{1}d_{2}\log(2(d_{1} + d_{2}) / \delta)}{\tau_{\ell}^{E}}}$ . This is in step 8 of Algorithm 1.
2.2 Optimal Design for Rotated Arms (Stage 2 of  $\ell$ -th phase)
In stage 2 of phase  $\ell$ , GOBLIN leverages the information about the learned sub-space of  $\Theta_{*}$  to rotate the arm set and then run the optimal design on the rotated arm set. Once we recover  $\widehat{\Theta}_{\ell}$ , 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  $\mathbf{x}_{*}$  and  $\mathbf{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  $\Theta_{*}$  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  $\widehat{\Theta}_{\ell} = \widehat{\mathbf{U}}_{\ell}\widehat{\mathbf{D}}_{\ell}\widehat{\mathbf{V}}_{\ell}^{\top}$  be the SVD of  $\widehat{\Theta}_{\ell}$  in the  $\ell$ -th phase. Let  $\widehat{\mathbf{U}}_{\perp}^{\ell}$  and  $\widehat{\mathbf{V}}_{\perp}^{\ell}$  be orthonormal bases of the complementary subspaces of  $\widehat{\mathbf{U}}_{\ell}$  and  $\widehat{\mathbf{V}}_{\ell}$  respectively. Let  $\mathcal{X}_{\ell}$  and  $\mathcal{Z}_{\ell}$  be the active set of arms in the stage 2 of phase  $\ell$ . Then rotate the arm sets such that new rotated arm sets are as follows:
$$
\underline {{\mathcal {X}}} _ {\ell} = \left\{\underline {{\mathbf {x}}} = \left[ \widehat {\mathbf {U}} _ {\ell} \widehat {\mathbf {U}} _ {\ell} ^ {\perp} \right] ^ {\top} \mathbf {x} \mid \mathbf {x} \in \mathcal {X} _ {\ell} \right\}, \underline {{\mathcal {Z}}} _ {\ell} = \left\{\underline {{\mathbf {z}}} = \left[ \widehat {\mathbf {V}} _ {\ell} \widehat {\mathbf {V}} _ {\ell} ^ {\perp} \right] ^ {\top} \mathbf {z} \mid \mathbf {z} \in \mathcal {Z} _ {\ell} \right\}. \tag {3}
$$
Let  $\widehat{\mathbf{H}}_{\ell} = [\widehat{\mathbf{U}}_{\ell}\widehat{\mathbf{U}}_{\ell}^{\perp}]^{\top}\widehat{\boldsymbol{\Theta}}_{\ell}[\widehat{\mathbf{V}}_{\ell}\widehat{\mathbf{V}}_{\ell}^{\perp}]$ . Then define vectorized arm set so that the last  $(d_1 - r)\cdot (d_2 - r)$  components are from the complementary subspaces as follows:
$$
\underline {{\mathcal {W}}} _ {\ell} = \left\{\left[ \operatorname {v e c} \left(\underline {{\mathbf {x}}} _ {1: r} \underline {{\mathbf {z}}} _ {1: r} ^ {\top}\right); \operatorname {v e c} \left(\underline {{\mathbf {x}}} _ {r + 1: d _ {1}} \underline {{\mathbf {z}}} _ {1: r} ^ {\top}\right); \operatorname {v e c} \left(\underline {{\mathbf {x}}} _ {1: r} \underline {{\mathbf {z}}} _ {r + 1: d _ {2}} ^ {\top}\right); \right. \right.
$$
$$
\left. \mathbf {v e c} \left(\underline {{\mathbf {x}}} _ {r + 1: d _ {1}} \underline {{\mathbf {z}}} _ {r + 1: d _ {2}} ^ {\top}\right) \right] \in \mathbb {R} ^ {d _ {1} d _ {2}}: \underline {{\mathbf {x}}} \in \mathcal {X} _ {\ell}, \underline {{\mathbf {z}}} \in \mathcal {Z} _ {\ell} \bigg \}
$$
$$
\widehat {\boldsymbol {\theta}} _ {\ell , 1: k} = \left[ \mathbf {v e c} \left(\widehat {\mathbf {H}} _ {\ell , 1: r, 1: r}\right); \mathbf {v e c} \left(\widehat {\mathbf {H}} _ {\ell , r + 1: d _ {1}, 1: r}\right); \mathbf {v e c} \left(\widehat {\mathbf {H}} _ {\ell , 1: r, r + 1: d _ {2}}\right) \right],
$$
$$
\widehat {\boldsymbol {\theta}} _ {\ell , k + 1: p} = \operatorname {v e c} \left(\widehat {\mathbf {H}} _ {\ell , r + 1: d _ {1}, r + 1: d _ {2}}\right). \tag {4}
$$
which implies  $\| \widehat{\pmb{\theta}}_{k + 1:p}\| _2 = O\left(d_1d_2r / \tau_\ell^E\right)$  by Lemma 3 in Appendix A.1. So the last  $p - k$  components of  $\widehat{\pmb{\theta}}_{\ell}$  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  $\ell$ , GOBLIN implements  $G$ -optimal design (Pukelsheim, 2006; Fiez et al., 2019) in the rotated arm set  $\underline{\mathcal{X}}_{\ell}, \underline{\mathcal{Z}}_{\ell}$  defined in (3). To do this, first GOBLIN defines the rotated vector  $\underline{\mathbf{w}} = [\underline{\mathbf{x}}_{1:d_1}; \underline{\mathbf{z}}_{1:d_2}] \in \mathbb{R}^p$  that belong to the set  $\underline{\mathcal{W}}_{\ell}$ . Then GOBLIN solves the  $G$ -optimal design (Pukelsheim, 2006) as follows:
$$
\widehat {\mathbf {b}} _ {\ell} ^ {G} = \underset {\mathbf {b} _ {\underline {{\mathbf {w}}}}} {\arg \min } \underset {\underline {{\mathbf {w}}}, \underline {{\mathbf {w}}} ^ {\prime} \in \underline {{\mathcal {W}}} _ {\ell}} {\max } \left\| \underline {{\mathbf {w}}} - \underline {{\mathbf {w}}} ^ {\prime} \right\| _ {\left(\sum_ {\underline {{\mathbf {w}}} \in \underline {{\mathcal {W}}}} \mathbf {b} _ {\underline {{\mathbf {w}}}} \underline {{\mathbf {w}}} \underline {{\mathbf {w}}} ^ {\top} + \boldsymbol {\Lambda} _ {\ell} / n\right) ^ {- 1}} ^ {2}. \tag {5}
$$
This is shown in step 11 of Algorithm 1 and  $\Lambda_{\ell}$  is defined in (6). It can be shown that sampling according to  $\widehat{\mathbf{b}}_{\ell}^{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  $\underline{W}_{\ell}$  we are guaranteed that the support of  $\mathrm{supp}(\widehat{\mathbf{b}}_{\ell}^{G}) \leq \widetilde{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  $\widetilde{\mathbf{b}}_{\ell}^{G}$  will scale with  $d_{1}d_{2}$  which will lead to higher sample complexity. Then GOBLIN samples each  $\underline{\mathbf{w}}\in \underline{\mathcal{W}}_{\ell}$  for  $\lceil \tau_{\ell}^{G}\mathbf{b}_{\ell ,\underline{\mathbf{w}}}^{G}\rceil$  times, where  $\tau_{\ell}^{G}\coloneqq \lceil \frac{8B_{*}^{\ell}\rho^{G}(\mathcal{Y}(\mathcal{W}_{\ell}))\log(4\ell^{2}|\mathcal{W}| / \delta)}{\epsilon_{\ell}^{2}}\rceil$  Note that the total length of phase  $\ell$  combining stages 1 and 2 is  $(\tau_{\ell}^{E} + \tau_{\ell}^{G})$  rounds. Observe that the stage 1 design is on the whole arm set  $\overline{\mathcal{W}}$  whereas stage 2 design is on the refined active set  $\underline{\mathcal{W}}_{\ell}$
Let the observed features in stage 2 of phase  $\ell$  be denoted by  $\mathbf{W}_{\ell} \in \mathbb{R}^{\tau_{\ell}^{G} \times p}$ , and  $\mathbf{r}_{\ell} \in \mathbb{R}^{\tau_{\ell}^{G}}$  be the observed rewards. Define the diagonal matrix  $\Lambda_{\ell}$  as
$$
\boldsymbol {\Lambda} _ {\ell} = \operatorname {d i a g} [ \underbrace {\lambda , \dots , \lambda} _ {k}, \underbrace {\lambda_ {\ell} ^ {\perp} , \dots , \lambda_ {\ell} ^ {\perp}} _ {p - k} ] \tag {6}
$$