1: Input: arm set  $\mathcal{X},\mathcal{Z}$ , confidence  $\delta$ , rank  $r$  of  $\Theta_{*}$ , spectral bound  $S_{r}$  of  $\Theta_{*}$ ,  $S,S_{m,\ell}^{\perp} = \frac{8k_1k_2r}{\widetilde{\tau}_{m,\ell}^{E}S_r^2}\log \left(\frac{k_1 + k_2}{\delta_\ell}\right)$ ,  $\lambda, \lambda_{m,\ell}^{\perp} = \frac{\tau_{m,\ell - 1}^{G}}{(8(k_1 + k_2)r\log(1 + \tau_{m,\ell - 1}^{G} / \lambda))}$ . Let  $p = k_1k_2$ ,  $k = (k_1 + k_2)r$ .
2: Let  $\underline{\mathcal{W}}_{m,1} \gets \underline{\mathcal{W}}_m, \ell \gets 1$ ,  $\tau_0^G = \log (4\ell^2 |\mathcal{X}| / \delta)$ . Define  $\Lambda_{m,\ell}$  as in (6),  $B_{m,*}^{\ell} \coloneqq (8\sqrt{\lambda} S + \sqrt{\lambda_{m,\ell}^{\perp}} S_{m,\ell}^{\perp})$
3: Define arm  $\overline{\mathbf{w}} = [\mathbf{x}_{1:d_1};\mathbf{z}_{1:d_2}]$  and  $\overline{\mathbf{w}}\in \overline{\mathcal{W}}_m$  . Let  $\tau_{\ell}^{E} = \frac{\sqrt{8d_{1}d_{2}r\log(4\ell^{2}|\mathcal{W}| / \delta_{\ell})}}{S_{r}}$  . Let  $E$  -optimal design be  $\mathbf{b}_{\ell}^{E} = \arg \min_{\mathbf{b}\in \triangle_{\overline{\mathcal{W}}}}\| (\sum_{\overline{\mathbf{w}}\in \overline{\mathcal{W}}}\mathbf{b}_{\overline{\mathbf{w}}}\overline{\mathbf{w}}\overline{\mathbf{w}}^{\top})^{-1}\|$
4: while  $\exists m\in [M],\left|\underline{\mathcal{W}}_{m,\ell}\right| > 1$  do
5:  $\epsilon_{\ell} = 2^{-\ell}, \delta_{\ell} = \delta / \ell^{2}$ .
6: (Stage 1:) Explore the Low-Rank Subspace
7: Pull arm  $\overline{\mathbf{w}}\in \overline{\mathcal{W}}$  exactly  $\lceil \widehat{\mathbf{b}}_{\ell ,\overline{\mathbf{w}}}^{E}\tau_{\ell}^{E}\rceil$  times for each task  $m$  and observe rewards  $\{r_{m,t}\}_{t = 1}^{\tau_\ell^E}$
8: Compute  $\widehat{\mathbf{Z}}_{\ell}$  using (8).
9: (Stage 2:) Build  $\mathbf{S}_{m,\ell}$  for each task  $m$
10: 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. Build  $\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}$  for all  $\mathbf{x}\in \mathcal{X}$  and  $\mathbf{z}\in \mathcal{Z}$  for each  $m$
11: Define a vectorized arm  $\widetilde{\mathbf{w}} = [\widetilde{\mathbf{g}}_{m,1:k_1};\widetilde{\mathbf{v}}_{m,1:k_2}]$  and  $\widetilde{\mathbf{w}} \in \widetilde{\mathcal{W}}_m$  for each  $m$ . Let  $\widetilde{\tau}_{m,\ell}^{E} = \frac{\sqrt{8k_{1}k_{2}r\log(4\ell^{2}|\mathcal{W}| / \delta_{\ell})}}{S_{r}}$ , and  $\widetilde{\mathbf{b}}_{m,\ell}^{E} = \arg \min_{\mathbf{b}_{m} \in \triangle_{\widetilde{\mathcal{W}}_{m}}} \left\| \left( \sum_{\widetilde{\mathbf{w}} \in \widetilde{\mathcal{W}}_{m}} \mathbf{b}_{m,\widetilde{\mathbf{w}}} \widetilde{\mathbf{w}} \widetilde{\mathbf{w}}^{\top} \right)^{-1} \right\|$ .
12: Pull arm  $\widetilde{\mathbf{w}}\in \widetilde{\mathcal{W}}_m$  exactly  $\left|\widetilde{\mathbf{b}}_{m,\ell ,\widetilde{\mathbf{w}}}^{E}\widetilde{\tau}_{m,\ell}^{E}\right|$  times and observe rewards  $r_{m,t}$ , for  $t = 1,\ldots ,\widetilde{\tau}_{m,\ell}^{E}$ , for each task  $m$ . Then compute  $\widehat{\mathbf{S}}_{m,\ell}$  using (9) for each  $m$ .
13: (Stage 3:) Reduction to low dimensional linear bandits for each task  $m$
14: SVD of  $\widehat{\mathbf{S}}_{m,\ell} = \widehat{\mathbf{U}}_{m,\ell}\widehat{\mathbf{D}}_{m,\ell}\widehat{\mathbf{V}}_{m,\ell}^{\top}$ . Rotate arms in active set  $\underline{\mathcal{W}}_{m,\ell-1}$  to build  $\underline{\mathcal{W}}_{m,\ell}$  using (10).
15: Let  $\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}\|^{2}(\sum_{\underline{\mathbf{w}}_{m}\in \underline{\mathcal{W}}_{m}}\mathbf{b}_{m,\underline{\mathbf{w}}}\underline{\mathbf{w}}_{m}\underline{\mathbf{w}}_{m}^{\top} + \boldsymbol{\Lambda}_{m,\ell /n})^{-1}$ .
16: Define  $\rho^G (\mathcal{V}(\underline{\mathcal{W}}_{m,\ell})) = \min_{\mathbf{b}_{m,\underline{\mathbf{w}}}}\max_{\underline{\mathbf{w}},\underline{\mathbf{w}} '\in \underline{\mathcal{W}}_{m,\ell}}\| \underline{\mathbf{w}} -\underline{\mathbf{w}} '\| ^2$ $(\sum_{\underline{\mathbf{w}}\in \underline{\mathcal{W}}_m}\mathbf{b}_{m,\underline{\mathbf{w}}\underline{\mathbf{w}}}\underline{\mathbf{w}}^\top +\frac{\boldsymbol{\Lambda}_{m,\ell}}{n})^{-1}$
17: Set  $\tau_{m,\ell}^{G} = \frac{64B_{m,*}^{\ell}\rho^{G}(\mathcal{Y}(\mathcal{W}_{m,\ell}))\log(4\ell^{2}|\mathcal{W}_{m}| / \delta_{\ell})}{\epsilon_{\ell}^{2}}$ . Then pull arm  $\underline{\mathbf{w}} \in \underline{\mathcal{W}}_m$  for each task  $m$  exactly  $[\widehat{\mathbf{b}}_{m,\ell,\underline{\mathbf{w}}} \tau_{m,\ell}^{G}]$  times and construct the least squares estimator  $\widehat{\pmb{\theta}}_{m,\ell}$  using only the observations of this phase where  $\widehat{\pmb{\theta}}_{m,\ell}$  is defined in (12).
18: Eliminate arms such that  $\underline{\mathcal{W}}_{m,\ell +1}\gets \underline{\mathcal{W}}_{m,\ell}\backslash \left\{\underline{\mathbf{w}}_m\in \underline{\mathcal{W}}_{m,\ell}:\max_{\underline{\mathbf{w}}'_m\in \underline{\mathcal{W}}_{m,\ell}}\left\langle \underline{\mathbf{w}}'_m - \underline{\mathbf{w}}_m,\widehat{\pmb{\theta}}_{m,\ell}\right\rangle >2\epsilon_{m,\ell}\right\}$
19:  $\ell \gets \ell + 1$
20: Output the arm in  $\underline{\mathcal{W}}_{m,\ell}$  and reshape to get the  $\widehat{\mathbf{x}}_{m,*}$  and  $\widehat{\mathbf{z}}_{m,*}$  for each task  $m$ .