SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where,  $\lambda_{\ell}^{\perp}\coloneqq \tau_{\ell -1}^{G} / 8k\log (1 + \tau_{\ell -1}^{G} / \lambda)\gg \lambda$  . Deviating from Soare et al. (2014); Fiez et al. (2019) GOBLIN constructs a regularized least square estimator at phase  $\ell$  as follows
$$
\widehat {\boldsymbol {\theta}} _ {\ell} = \underset {\boldsymbol {\theta} \in \mathbb {R} ^ {p}} {\arg \min } \frac {1}{2} \| \underline {{\mathbf {W}}} _ {\ell} \boldsymbol {\theta} - \mathbf {r} _ {\ell} \| _ {2} ^ {2} + \frac {1}{2} \| \boldsymbol {\theta} \| _ {\boldsymbol {\Lambda} _ {\ell}} ^ {2}. \tag {7}
$$
This regularized least square estimator in (7) forces the last  $p - k$  components of  $\widehat{\theta}_{\ell}$  to be very small compared to the first  $k$  components. Then GOBLIN builds the estimate  $\widehat{\theta}_{\ell}$  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  $\widehat{\theta}_{\ell}$ . Finally GOBLIN eliminates sub-optimal arms to build the next phase active set  $\underline{W}_{\ell}$  and stops when  $|\underline{W}_{\ell}| = 1$ . GOBLIN outputs the arm in  $\underline{W}_{\ell}$  and reshapes it to get the  $\widehat{\mathbf{x}}_{*}$  and  $\widehat{\mathbf{z}}_{*}$ . The full pseudocode is presented in Algorithm 1.
Algorithm 1 G-Optimal Design for Bilinear Bandits (GOBLIN) for single-task setting
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 - \delta$ , GOBLIN returns the best arms  $\mathbf{x}_*$ ,  $\mathbf{z}_*$ , and the number of samples used is bounded by  $\widetilde{O}\left(\frac{(d_1 + d_2)r}{\Delta^2} + \frac{\sqrt{d_1d_2r}}{S_r}\right)$ .
Discussion 1. In Theorem 1 the first quantity is the number of samples needed to identify the best arms  $\mathbf{x}_{*},\mathbf{z}_{*}$  while the second quantity is the number of samples to learn  $\Theta_{*}$  (which is required to find the best arms). Note that the magnitude of  $S_{r}$  would be free of  $d_{1},d_{2}$  since  $\Theta_{*}$  contains only  $r$  nonzero singular values and  $\| \Theta_{*}\| \leq 1$ , and hence we assume that  $S_{r} = \Theta (1 / \sqrt{r})$  (Kang et al., 2022). So the sample complexity of single-task GOBLIN scales as  $\widetilde{O} (\frac{(d_1 + d_2)r}{\Delta^2})$ . However, if one runs RAGE (Fiez et al., 2019) on the arms in  $\mathcal{X}$ ,  $\mathcal{Z}$  then the sample complexity will scale as  $\widetilde{O} (\frac{d_1d_2}{\Delta^2})$ .
Proof (Overview) of Theorem 1: Step 1 (Subspace estimation in high dimension): We denote the vectorized arms in high dimension as  $\overline{\mathbf{w}}\in \overline{\mathcal{W}}$ . We run the  $E$ -optimal design to sample the arms in  $\overline{\mathcal{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  $\| \widehat{\Theta}_{\ell} - \Theta_{*}\|_{F}^{2}\leq$ $\frac{C_1d_1d_2r\log(2(d_1 + d_2) / \delta)}{\tau_\ell^E}$  for some  $C_1 > 0$ . Also, note that in the first stage of the  $\ell$ -th phase by setting  $\tau_{\ell}^{E} = \frac{\sqrt{8d_{1}d_{2}r\log(4\ell^{2}|\mathcal{W}| / \delta_{\ell})}}{S_{r}}$  and sampling each arm  $\overline{\mathbf{w}}\in \overline{\mathcal{W}}$  exactly  $\lceil \widehat{\mathbf{b}}_{\ell ,\overline{\mathbf{w}}}^{E}\tau_{\ell}^{E}\rceil$  times we are guaranteed that  $\| \pmb{\theta}_{k + 1:p}^{*}\|_{2} = O(d_{1}d_{2}r / \tau_{\ell}^{E})$ . Summing up over  $\ell = 1$  to  $\lceil \log_2(4\Delta^{-1})\rceil$  we get that the total sample complexity of the first stage is bounded by  $\widetilde{O} (\sqrt{d_1d_2r} /S_r)$ .
Step 2 (Effective dimension for rotated arms): We rotate the arms  $\overline{\mathbf{w}}\in \overline{\mathcal{W}}$  in high dimension to get the rotated arms  $\underline{\mathbf{w}}\in \underline{\mathcal{W}}_{\ell}$  in step 10 of Algorithm 1. Then we show that the effective dimension of  $\underline{\mathbf{w}}$  scales  $8k\log \left(1 + \tau_{\ell -1}^{G} / \lambda\right)$  when  $\lambda_{\ell}^{\perp} = \frac{\tau_{\ell - 1}^{G}}{8k\log\left(1 + \tau_{\ell - 1}^{G} / \lambda\right)}$  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  $\lambda_{\ell}^{\perp}$  to control during phase  $\ell$ , and the effective dimensions in those papers do not depend on phase length.
Step 3 (Bounded Support): For any phase  $\ell$ , we can show that  $1 \leq \rho^{G}(\mathcal{V}(\underline{\mathcal{W}}_{\ell})) \leq p / \gamma_{\mathcal{Y}}^{2}$  where,  $\gamma_{\mathcal{Y}} = \max \{c > 0 : c\mathcal{Y} \subset \mathrm{conv}(\underline{\mathcal{W}} \cup -\underline{\mathcal{W}})\}$  is the gauge norm of  $\mathcal{Y}$  (Rockafellar, 2015). Note that this is a worst-case dependence when  $\rho^{G}(\mathcal{Y}(\underline{\mathcal{W}}_{\ell}))$  scales with  $p$ . Substituting this value of  $\rho^{G}(\mathcal{Y}(\underline{\mathcal{W}}_{\ell}))$  in the definition of  $\lambda_{\ell}^{\perp}$  we can show that  $\Lambda_{\ell}$  does not depend on  $\underline{\mathbf{w}}$  or  $\mathbf{y} = \underline{\mathbf{w}} - \underline{\mathbf{w}}'$ . Then following Theorem 21.1 in Lattimore and Szepesváris (2020) we can show that the  $G$ -optimal design  $\widehat{\mathbf{b}}_{\ell}^{G}$  is equivalent to  $D$ -optimal design  $\widehat{\mathbf{b}}_{\ell}^{D} = \arg \max_{\mathbf{b}} \log \frac{\left|\sum_{\underline{\mathbf{w}} \in \mathcal{W}_{\ell}} \mathbf{b}_{\underline{\mathbf{w}}} \underline{\mathbf{w}} \underline{\mathbf{w}}^{\top} + \boldsymbol{\Lambda}_{\ell}\right|}{|\Lambda_{\ell}|}$ . Then using Frank-Wolfe algorithm (Jamieson and Jain, 2022) we can show the support  $\widehat{\mathbf{b}}_{\ell}^{G}$  or equivalently  $\widehat{\mathbf{b}}_{\ell}^{D}$  is bounded by at most  $\frac{8k \log(1 + \tau_{\ell-1}^{G} / \lambda)(8k \log(1 + \tau_{\ell-1}^{G} / \lambda) + 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  $\tau_{\ell}^{G} = \lceil \frac{8B_{*}^{\ell}\rho(\mathcal{Y}(\underline{W}_{\ell}))\log(2|\mathcal{W}| / \delta)}{(\mathbf{x}^{\top}(\widehat{\theta}_{\ell} - \theta^{*}))^{2}}\rceil$ . 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  $\underline{W}_{\ell}$  exactly  $\lceil \widehat{\mathbf{b}}_{\ell,\underline{\mathbf{w}}} \tau_{\ell}^{G} \rceil$  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\left(\frac{k}{\gamma_{\mathcal{Y}}^{2}} \log \left( \frac{k \log_{2}(\Delta^{-1}) |\underline{\mathcal{W}}|}{\delta} \right) \lceil \log_{2}(\Delta^{-1}) \rceil \right)$  where the effective dimension  $k = (d_1 + d_2)r$ . Finally, we combine the total samples of phase  $\ell$  as  $(\tau_{\ell}^{E} + \tau_{\ell}^{G})$ . The final sample complexity is given by summing over all phases from  $\ell = 1$  to  $\lceil \log_{2}(4\Delta^{-1}) \rceil$ . The claim of the theorem follows by noting  $\widetilde{O}(k / \gamma_{\mathcal{Y}}^2) \leq \widetilde{O}(k / \Delta^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 \in [M]$  has a reward model stated
in (1). The learning proceeds as follows: 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]\}$ . Finally recall that the goal is to identify the optimal left and right arms  $\mathbf{x}_{m,*},\mathbf{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  $\Theta_{m,*}$  for all the  $m\in [M]$  have a decomposition  $\Theta_{m,*} = \mathbf{B}_1\mathbf{S}_{m,*}\mathbf{B}_2^\top$  and each  $\mathbf{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  $\sigma_{\mathrm{min}}\left(\frac{1}{M}\sum_{m=1}^{M}\Theta_{m,*}\right) \geq \frac{c_0}{S_r}$ , for some  $c_0 > 0$ ,  $S_r$  is the  $r$ -th largest singular value of  $\Theta_{m,*}$  and  $\sigma_{\mathrm{min}}(\mathbf{A})$  denotes the minimum eigenvalue of matrix  $\mathbf{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  $\mathbf{B}_1$  and  $\mathbf{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  $\ell = 1,2,\ldots$  consists of three stages; the stage for estimation of feature extractors  $\mathbf{B}_1,\mathbf{B}_2$ , which runs for  $\tau_{\ell}^{E}$  rounds, the stage for estimation of  $\mathbf{S}_{m,*}$  which runs for  $\sum_{m}\widetilde{\tau}_{m,\ell}^{E}$  rounds, and a stage of pure exploration with rotated arms that runs for  $\sum_{m}\tau_{m,\ell}^{G}$  rounds. We will define  $\tau_{m,\ell}^{E}$  in Section 3.1,  $\widetilde{\tau}_{m,\ell}^{E}$  in Section 3.2, while the rotated arms and  $\tau_{m,\ell}^{G}$  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  $\ell = 1,2,\ldots$  for multi-task GOBLIN.
3.1 Estimating Feature Extractors  $\mathbf{B}_1$  and  $\mathbf{B}_2$  (Stage 1 of Phase  $\ell$ )
In the first stage of phase  $\ell$ , GOBLIN leverages the batch of rewards  $\{r_{m,t}: m \in [M]\}$  at every round  $t$  from  $M$  tasks to learn the feature extractors  $\mathbf{B}_1$  and  $\mathbf{B}_2$ . To do this, GOBLIN first vectorizes the  $\mathbf{x} \in \mathcal{X}, \mathbf{z} \in \mathcal{Z}$  into a new vector  $\overline{\mathbf{w}} = [\mathbf{x}_{1:d_1}; \mathbf{z}_{1:d_2}] \in \overline{\mathcal{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  $\overline{\mathbf{w}} \in \overline{\mathcal{W}}_m$  for  $\lceil \tau_\ell^E \mathbf{b}_{\ell,\overline{\mathbf{w}}}^E \rceil$  times for each task  $m$ , where  $\tau_\ell^E = \widetilde{O}(\sqrt{d_1 d_2 r} / S_r)$  and  $\mathbf{b}_{\ell,\overline{\mathbf{w}}}^E$  is the solution to  $E$ -optimal design on  $\overline{\mathbf{w}}$ . Let the sampled arms for each task  $m$  at round  $s$  be denoted by  $\mathbf{x}_{m,s}, \mathbf{z}_{m,s}$  which is obtained after reshaping  $\overline{\mathbf{w}}_s$ . Then it builds the estimator  $\widehat{\mathbf{Z}}_\ell$  as follows:
$$
\widehat{\mathbf{Z}}_{\ell} = \operatorname *{arg  min}_{\boldsymbol {\Theta}\in \mathbb{R}^{d_{1}\times d_{2}}}L_{\ell}(\boldsymbol {\Theta}) + \gamma_{\ell}\| \boldsymbol {\Theta}\|_{\mathrm{nuc}},
$$
$$
L _ {\ell} (\boldsymbol {\Theta}) = \langle \boldsymbol {\Theta}, \boldsymbol {\Theta} \rangle - \frac {2}{M \tau_ {\ell} ^ {E}} \sum_ {m = 1} ^ {M} \sum_ {s = 1} ^ {\tau_ {\ell} ^ {E}} \langle \widetilde {\psi} _ {\nu} \left(r _ {m, s} \cdot Q \left(\mathbf {x} _ {m, s} \mathbf {z} _ {m, s} ^ {\top}\right)\right), \boldsymbol {\Theta} \rangle \tag {8}
$$
Then it performs SVD decomposition on  $\widehat{\mathbf{Z}}_{\ell}$ , and let  $\widehat{\mathbf{B}}_1, \widehat{\mathbf{B}}_2$  be the top- $k_1$  and top- $k_2$  left and right singular vectors of  $\widehat{\mathbf{Z}}_{\ell}$  respectively. These are the estimation of the feature extractors  $\mathbf{B}_1$  and  $\mathbf{B}_2$ .
3.2 Estimating Hidden Parameter  $\mathbf{S}_{m,*}$  per Task (Stage 2 of phase  $\ell$ )
In the second stage of phase  $\ell$ , the goal is to recover the hidden parameter  $\mathbf{S}_{m,*}$  for each task  $m$ . GOBLIN proceeds as follows: First, let  $\widetilde{\mathbf{g}}_m = \mathbf{x}^\top \widehat{\mathbf{B}}_{1,\ell}$  and  $\widetilde{\mathbf{v}}_m = \mathbf{z}^\top \widehat{\mathbf{B}}_{2,\ell}$  be the latent left and right arm respectively for each  $m$ . Then GOBLIN defines the vector  $\widetilde{\mathbf{w}} = [\widetilde{\mathbf{g}}_m; \widetilde{\mathbf{v}}_m] \in \widetilde{\mathcal{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  $\widetilde{\mathbf{w}} \in \widetilde{\mathcal{W}}_m$  for  $[\widetilde{\tau}_{m,\ell}^E \widetilde{\mathbf{b}}_{m,\ell,\widetilde{\mathbf{w}}}^E]$  times, where  $\widetilde{\tau}_{m,\ell}^E \coloneqq \widetilde{O}(\sqrt{k_1 k_2 r} / S_r)$  and  $\widetilde{\mathbf{b}}_{m,\ell,\widetilde{\mathbf{w}}}^E$  is the solution to
$E$ -optimal design on  $\widetilde{\mathbf{w}}$ . Then it builds estimator  $\widehat{\mathbf{S}}_{m,\ell}$  for each task  $m$  in step 12 as follows:
$$
\widehat{\mathbf{S}}_{m,\ell} = \operatorname *{arg  min}_{\boldsymbol {\Theta}\in \mathbb{R}^{k_{1}\times k_{2}}}L^{\prime}_{\ell}(\boldsymbol {\Theta}) + \gamma_{\ell}\| \boldsymbol {\Theta}\|_{\mathrm{nuc}},
$$
$$
L _ {\ell} ^ {\prime} (\boldsymbol {\Theta}) = \langle \boldsymbol {\Theta}, \boldsymbol {\Theta} \rangle - \frac {2}{\widetilde {\tau} _ {m , \ell} ^ {E}} \sum_ {s = 1} ^ {\widetilde {\tau} _ {m, \ell} ^ {E}} \left\langle \widetilde {\psi} _ {\nu} \left(r _ {m, s} \cdot Q \left(\widetilde {\mathbf {g}} _ {m, s} \widetilde {\mathbf {v}} _ {m, s} ^ {\top}\right)\right), \boldsymbol {\Theta} \right\rangle \tag {9}
$$
Once GOBLIN recovers the  $\widehat{\mathbf{S}}_{m,\ell}$  for each task  $m$  it has reduced the  $d_1d_2$  bilinear bandit to a  $k_{1}k_{2}$  dimension bilinear bandit where the left and right arms are  $\widetilde{\mathbf{g}}_m\in \mathcal{G}_m,\widetilde{\mathbf{v}}_m\in \mathcal{V}_m$  respectively.
3.3 Optimal Design for Rotated Arms per Task (Stage 3 of phase  $\ell$ )
In the third stage of phase  $\ell$ , similar to Algorithm 1, the multi-task GOBLIN defines the rotated arm set  $\underline{\mathcal{G}}_m, \underline{\mathcal{V}}_m$  for each task  $m$  for these  $k_{1}k_{2}$  dimensional bilinear bandits. Let the SVD of  $\widehat{\mathbf{S}}_{m,\ell} = \widehat{\mathbf{U}}_{m,\ell}\widehat{\mathbf{D}}_{m,\ell}\widehat{\mathbf{V}}_{m,\ell}^{\top}$ . Define  $\widehat{\mathbf{H}}_{m,\ell} = [\widehat{\mathbf{U}}_{m,\ell}\widehat{\mathbf{U}}_{m,\ell}^{\perp}]^{\top}\widehat{\mathbf{S}}_{m,\ell}[\widehat{\mathbf{V}}_{m,\ell}\widehat{\mathbf{V}}_{m,\ell}^{\perp}]$ . Then define the vectorized arm set so that the last  $(k_{1} - r)\cdot (k_{2} - r)$  components are from the complementary subspaces as follows:
$$
\underline {{\mathcal {W}}} _ {m, \ell} = \left\{\left[ \mathbf {v e c} \left(\widetilde {\mathbf {g}} _ {m, 1: r} \widetilde {\mathbf {v}} _ {m, 1: r} ^ {\top}\right); \mathbf {v e c} \left(\widetilde {\mathbf {g}} _ {m, r + 1: k _ {1}} \widetilde {\mathbf {v}} _ {m, 1: r} ^ {\top}\right); \mathbf {v e c} \left(\widetilde {\mathbf {g}} _ {m, 1: r} \widetilde {\mathbf {v}} _ {m, r + 1: k _ {2}} ^ {\top}\right); \right. \right.
$$
$$
\left. \left. \mathbf {v e c} \left(\widetilde {\mathbf {g}} _ {m, r + 1: k _ {1}} \widetilde {\mathbf {v}} _ {m, r + 1: k _ {2}} ^ {\top}\right) \right] \right\}
$$
$$
\widehat {\boldsymbol {\theta}} _ {m, \ell , 1: k} = \left[ \operatorname {v e c} \left(\widehat {\mathbf {H}} _ {m, \ell , 1: r, 1: r}\right); \operatorname {v e c} \left(\widehat {\mathbf {H}} _ {m, \ell , r + 1: k _ {1}, 1: r}\right); \operatorname {v e c} \left(\widehat {\mathbf {H}} _ {m, \ell , 1: r, r + 1: k _ {2}}\right) \right],
$$
$$
\boldsymbol {\theta} _ {\ell , k + 1: p} = \operatorname {v e c} \left(\widehat {\mathbf {H}} _ {m, \ell , r + 1: k _ {1}, r + 1: k _ {2}}\right). \tag {10}
$$
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  $\underline{\nu}_m, \underline{G}_m$  and define the  $\underline{\mathbf{w}} = [\widetilde{\mathbf{g}}_{m,1:d_1};\widetilde{\mathbf{v}}_{m,1:d_2}]$  and  $\widetilde{\mathbf{w}} \in \widetilde{\mathcal{W}}_m$  where  $\widetilde{\mathbf{g}}_m \in \underline{\mathcal{G}}_m$  and  $\widetilde{\mathbf{v}}_m \in \underline{\nu}_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
$$
\widehat {\mathbf {b}} _ {m, \ell} ^ {G} = \arg \min  _ {\mathbf {b} _ {m, \underline {{\mathbf {w}}}}} \max  _ {\underline {{\mathbf {w}}}, \underline {{\mathbf {w}}} ^ {\prime} \in \underline {{\mathcal {W}}} _ {m, \ell}} \| \underline {{\mathbf {w}}} - \underline {{\mathbf {w}}} ^ {\prime} \| _ {(\sum_ {\underline {{\mathbf {w}}} \in \underline {{\mathcal {W}}} _ {m}} \mathbf {b} _ {m, \underline {{\mathbf {w}}}} \underline {{\mathbf {w}}} \underline {{\mathbf {w}}} ^ {\top} + \boldsymbol {\Lambda} _ {m, \ell / n}) ^ {- 1}} \tag {11}
$$
Then GOBLIN runs  $G$ -optimal design on the arm set  $\underline{\mathcal{W}}_{\ell}$  following the (11) and then samples each  $\underline{\mathbf{w}} \in \underline{\mathcal{W}}_{m,\ell}$  for  $\lceil \tau_{m,\ell}^{G} \widehat{\mathbf{b}}_{m,\ell,\underline{\mathbf{w}}}^{G} \rceil$  times where  $\widehat{\mathbf{b}}_{m,\ell,\underline{\mathbf{w}}}^{G}$  is the solution to the  $G$ -optimal design, and  $\tau_{\ell}^{G}$  is defined in step 17 of Algorithm 2. So the total length of phase  $\ell$ , combining stages 1, 2 and 3 is  $(\tau_{\ell}^{E} + \sum_{m} \widetilde{\tau}_{m,\ell}^{E} + \sum_{m} \tau_{m,\ell}^{G})$  rounds. Observe that the stage 1 and 2 design is on the whole arm set  $\overline{\mathcal{W}}, \widetilde{\mathcal{W}}_m$  whereas the stage 3 design is on the refined active set  $\underline{\mathcal{W}}_{m,\ell}$ . Let at the stage 3 of  $\ell$ -th phase the actions sampled be denoted by the matrix  $\underline{\mathbf{W}}_{m,\ell} \in \mathbb{R}^{\tau_{m,\ell}^{G} \times k_1k_2}$  and observed rewards  $\mathbf{r}_m \in \mathbb{R}^{\tau_{m,\ell}^{G} \times k_1k_2}$ . Define the positive diagonal matrix  $\Lambda_{m,\ell}$  according to (6) but set  $p = k_1k_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
$$
\widehat {\boldsymbol {\theta}} _ {m, \ell} = \arg \min  _ {\boldsymbol {\theta}} \frac {1}{2} \| \underline {{\mathbf {W}}} _ {m, \ell} \boldsymbol {\theta} - \mathbf {r} _ {m} \| _ {2} ^ {2} + \frac {1}{2} \| \boldsymbol {\theta} \| _ {\boldsymbol {\Lambda} _ {m, \ell}} ^ {2} \tag {12}
$$
Finally GOBLIN eliminates the sub-optimal arms using the estimator  $\widehat{\theta}_{m,\ell}$  to build the next phase active set  $\underline{W}_{m,\ell}$  and stops when  $|\underline{W}_{m,\ell}| = 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 - \delta$ , GOBLIN returns the best arms  $\mathbf{x}_{m,*}$ ,  $\mathbf{z}_{m,*}$  for each task  $m$ , and the total number of samples is bounded by  $\widetilde{O}\left(\frac{M(k_1 + k_2)r}{\Delta^2} + \frac{M\sqrt{k_1k_2r}}{S_r} + \frac{\sqrt{d_1d_2r}}{S_r}\right)$ .
Discussion 2. In Theorem 2 the first quantity is the sample complexity to identify the best arms  $\mathbf{x}_{m,*}$ ,  $\mathbf{z}_{m,*}$  and the second quantity is the number of samples to learn  $\mathbf{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  $\Theta_{m,*}$  (which in turn is used to estimate the feature extractors  $\mathbf{B}_1$  and  $\mathbf{B}_2$  to learn the  $\mathbf{S}_{m,*}$ ). Again we assume that  $S_r = \Theta(1/\sqrt{r})$  (Kang et al., 2022). So the sample complexity of multi-task GOBLIN scales as  $\widetilde{O}(M(k_1 + k_2)r/\Delta^2)$ . However, if one runs DouExpDes (Du et al., 2023) then the sample complexity will scale as  $\widetilde{O}(M(k_1k_2)/\Delta^2)$  which is worse than GOBLIN when  $r \ll k_1$  or  $k_2$ .
Algorithm 2 G-Optimal Design for Bilinear Bandits (GOBLIN) for multi-task setting
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  $\widehat{\mathbf{Z}}_{\ell}$  to estimate the quantity  $\mathbf{Z}_{*} = \frac{1}{M}\sum_{m = 1}^{M}\Theta_{*,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  $\| \widehat{\mathbf{Z}}_{\ell} - \mathbf{Z}_{*}\|_{F}^{2}\leq \frac{C_{1}d_{1}d_{2}r\log(2(d_{1} + d_{2}) / \delta)}{\tau_{\ell}^{E}}$  for some  $C_1 > 0$ . This is shown in Lemma 12 in Appendix A.6. Summing up over  $\ell = 1$  to  $\left[\log_2\left(4\Delta^{-1}\right)\right]$  we get that the total samples complexity of the first stage is bounded by  $\widetilde{O} (\sqrt{d_1d_2r} /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  $\mathbf{B}_1$  and  $\mathbf{B}_2$ . Let  $\widehat{\mathbf{B}}_{1,\ell}$ ,  $\widehat{\mathbf{B}}_{2,\ell}$  be the top- $k_1$  left and top- $k_2$  right singular vectors of  $\widehat{\mathbf{Z}}_\ell$  respectively. Then using the Davis-Kahan sin  $\theta$  Theorem (Bhatia, 2013) in Lemma 14, 15 (Appendix A.6) we have  $\| (\widehat{\mathbf{B}}_{1,\ell}^{\perp})^{\top}\mathbf{B}_1\|$ ,  $\| (\widehat{\mathbf{B}}_{2,\ell}^{\perp})^{\top}\mathbf{B}_2\| \leq \widetilde{O} (\sqrt{(d_1 + d_2)r / M\tau_\ell^E})$ .
Step 3 (Estimation of  $\widehat{\mathbf{S}}_{m,\ell}$  in low dimension): Now we estimate the quantity  $\widehat{\mathbf{S}}_{m,\ell} \in \mathbb{R}^{k_1 \times k_2}$  for each task  $m$ . To do this we first build the latent arms  $\widetilde{\mathbf{g}}_m = \mathbf{x}^\top \widehat{\mathbf{U}}_\ell$  and  $\widetilde{\mathbf{v}}_m = \mathbf{z}^\top \widehat{\mathbf{V}}_\ell$  for all  $\mathbf{x} \in \mathcal{X}$