SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Assumption 4.1 (Geometry of the task space). We assume the source task space  $\mathcal{W}_{\text{source}}$  is a unit ball  $\mathbb{B}^{d_W^{\text{source}}}$  (1) that span over the first  $d_W^{\text{source}} \geq \frac{1}{2} d_W$  without loss of generality, while the target task space  $\mathcal{W}_{\text{target}} \subset \mathbb{R}^{d_W}$  can be any arbitrary  $\mathbb{B}^{d_W^{\text{target}}}$  (1).
Under this assumption, we let  $B_W^{\mathrm{source}}$  denote the first  $d_W^{\mathrm{source}}$  columns of  $B_W$ , which stands for the source-related part of  $B_W$ . And  $B_W^{\mathrm{target}}$
Then we assume the bilinear model where  $\phi_X = B_X \in \mathbb{B}^{d_X \times k}$  and  $\psi_X, \psi_W = I$ . Therefore,  $d_{\psi_X} = d_X, d_{\psi_W} = d_W$ . Moreover the model satisfies the following assumptions
Assumption 4.2 (Benign  $B_X, B_W$ ).  $B_X$  is an orthonormal matrix. Each column of  $B_W$  has magnitude  $\Theta(1)$  and  $\sigma_{\min}(B_W^{source}) > 1$ . Suppose we know  $\bar{\kappa} \geq \kappa(B_W^{source})$ ,  $\sigma_{\max}(B_W^{target})$  and  $\underline{\sigma} \leq \sigma_{\min}(B_W^{source})$ ,  $\sigma_{\min}(B_W^{target})$ . Trivially,  $\bar{\kappa} = \sqrt{d_W}$ ,  $\underline{\sigma} = 1$ .
Finally, the following assumption is required since we are using a training algorithm in [12] and might be able to relax to sub-gaussian by using other suboptimal oracles.
Assumption 4.3 (Isotropic Gaussian Input). For each task  $w$ , its input  $i$  satisfies  $x_{i,w} \sim \mathcal{N}(0, I_d)$ .
4.2 Algorithm
Here we provide the target-aware theory and postpone the target-agnostic in the Appendix. C since its analysis is covered by the target-aware setting.
This target-aware algorithm 2 follows the 3-stage which corresponds to sampling distribution  $q_{0}, q_{1}, q_{2}$  with explicit solutions. Notice that calculating  $q_{1}$  once is enough for theoretical guarantees.
We use existing passive multi-task training algorithms as oracles for  $\mathcal{O}_{\mathrm{offline}1}^{X},\mathcal{O}_{\mathrm{offline}2}^{X}$  and use the simple ERM methods for  $\mathcal{O}_{\mathrm{offline}}^{W}$  based on the learned  $\hat{B}$ . For the coarse exploration and fine target-agnostic exploration stage, the main purpose is to have a universal good estimation in all directions of  $B_{X}$ . (i.e., upper bound the  $\sin (\hat{B}_X,B_X)$ ) Therefore we choose the alternating minimization (MLLAM) proposed in [12]. On the other hand, for the fine target-aware exploration, we mainly care about final transfer learning performance on learned representation. Therefore, we use a non-convex ERM from [9]. We defer the details and its theoretical guarantees for  $\mathcal{O}_{\mathrm{offline}}$  into Appendix B.1.
Note the major disadvantage from [9] comes from its sample complexity scaling with a number of training source tasks, which will not be a problem here since in  $\mathcal{O}_{\mathrm{offline},3}^{X}$  since only  $k + k^{*} \ll d_{W}$  number of tasks are used. The major benefit of using non-convex ERM comes from its generality that it works even for the non-linear setting and is not tied with a specific algorithm. That is to say, as long as there exists other theoretical or heuristic oracles  $\mathcal{O}_{\mathrm{offline},1}^{X}, \mathcal{O}_{\mathrm{offline},2}^{X}$  giving a similar guarantee, stage 3 always works.
4.3 Results
Theorem 4.1 (Informal). By running Algo. 2, in order to let  $ER(\hat{\phi}_X, \nu_{\text{target}}) \leq \varepsilon^2$  with probability  $1 - \delta$ , the number of source samples  $n_{\text{source}}$  is at most
$$
\widetilde {\mathcal {O}} \left(\left(k d _ {X} + \log (1 / \delta)\right) \left(k ^ {*}\right) ^ {2} \min  \left\{k ^ {*}, \kappa^ {2} \left(B _ {W}\right) \right\} \max  _ {i} \| W _ {i} ^ {*} \| _ {2} ^ {2} \varepsilon^ {- 2} + l o w - o r d e r\right)
$$
Here  $k^{*} = \mathrm{rank}(\mathbb{E}_{w_{0}\sim \nu_{target}}B_{W}w_{0}w_{0}^{\top}B_{W}^{\top})$  represents the effective dimension of target and
$$
W _ {i} ^ {*} = \underset {w \in \mathcal {W} _ {\text {s o u r c e}}} {\arg \min } \| w \| _ {2} \quad s. t \quad B _ {W} ^ {\text {s o u r c e}} w = u _ {i} \sqrt {\lambda_ {i}} \text {w h e r e} U, \Lambda \leftarrow E i g (\mathbb {E} _ {w _ {0} \sim \nu_ {\text {t a r g e t}}} B _ {W} w _ {0} w _ {0} ^ {\top} B _ {W} ^ {\top}).
$$
Algorithm 2 Target-aware algorithm for benign source space
1: Inputs: Target probability  $\delta$ ,  $\bar{\kappa}$ ,  $\underline{\sigma}$ . Some constant  $\beta_{1}, \beta_{2}, \beta_{3}$ . Others same as Algo. 1.
2: Set  $q_0$  as  $q_0(e_t) = \frac{1}{d_W}, \forall t \in d_W, \text{and} q_0(w) = 0$  otherwise
3: Set  $n_0 = \beta_1\bar{\kappa}^2\left(k^3 d_X\bar{\kappa}^2 +d_W^{\frac{3}{2}}\underline{\sigma}^{-2}\sqrt{k + \log(1 / \delta)}\right)$ . Collect  $n_0q_0(w)$  data for each task denoted as  $\{Z_w\}_{w|q(w)\neq 0}$
4: Update  $\hat{B}_X \gets \mathcal{O}_{\text{offline 1}}^X(\{Z_w\}_{w|q_0(w) \neq 0})$  and  $\hat{B}_W^{\text{source}} \gets \mathcal{O}_{\text{offline}}^W(\{Z_w\}_{w|q_0(w) \neq 0}, \hat{B}_X)$
5: Compute  $q_{1}$  as  $q_{1}(v_{i}) = \frac{1}{k}, \forall i \in k$ , and  $q_{0}(w) = 0$  otherwise. Here  $v_{i}$  is the  $i$ -th vector of  $V$ , where  $U, D, V \gets \mathrm{SVD}(\hat{B}_{W}^{\mathrm{source}})$
6: for  $j = 1,2,3,\ldots$  do
7: Set  $n_1^j = \beta_2\epsilon_j^{-\frac{4}{3}}k^{\frac{5}{3}}d_W^{\frac{2}{3}}d_X^{\frac{1}{3}}\left(k^{\frac{2}{3}}d_W^{\frac{1}{3}}\underline{\sigma}^{-\frac{4}{3}} + \bar{\kappa}^2\underline{\sigma}^{-\frac{2}{3}}\right)$ . Collect  $n_1^j q_1(w)$  data for each task denoted as  $\{Z_w\}_{w|q_1(w)\neq 0}$ .
8: Update  $\hat{B}_X \gets \mathcal{O}_{\text{offline 2}}^X(\{Z_w\}_{w|q_1(w) \neq 0}), \hat{B}_W^{\text{source}} \gets \mathcal{O}_{\text{offline}}^W(\{Z_w\}_{w|q_1(w) \neq 0}, \hat{B}_X)$  and  $\hat{B}_W^{\text{target}} \gets \mathcal{O}_{\text{offline}}^W(\dot{Z}_{\text{target}}, \hat{B}_X)$
9: Find a set of target-aware tasks parameterized by  $\tilde{W}_j$  with each column  $i$  as
$$
\tilde {W} _ {j} (i) = \operatorname {P r o j} _ {\mathcal {W} _ {\text {s o u r c e}}} w _ {i} ^ {\prime} = \frac {w _ {i} ^ {\prime}}{\| w _ {i} ^ {\prime} \| _ {2}}
$$
where  $w_{i}^{\prime} = \underset {w}{\arg \min}\| w\|_{2}$  s.t.  $\hat{B}_{W,j}^{\mathrm{source}}w = u_i\sqrt{\lambda_i}\quad \forall \Lambda_i\geq 8(kd_W)^{\frac{3}{2}}\sqrt{\frac{d_X}{n_1}}$
where  $U,\Lambda \gets \operatorname {Eig}\left(\mathbb{E}_{w_0\sim \nu_{\mathrm{target}}}\left[\hat{B}_{W,j}^{\mathrm{target}}w_0(\hat{B}_{W,j}^{\mathrm{target}}w_0)^\top \right]\right)$
10: Compute  $q_2^j$  as  $q_2^j(w) = \frac{1}{\# \operatorname{col}(\tilde{W}_j)}$ ,  $\forall w \in \operatorname{col}(\tilde{W}_j)$  and  $q_2^j(w) = 0$  otherwise
11: Assign  $n_2^j$  total sampling budget as # col  $(\tilde{W}_j) \beta_3 \max_i \| W_j'(i) \|_2^2 \epsilon_j^{-2}$
12: Collect  $n_2^j(w) = n_2^j q_2^j(w)$  data for each task denoted as  $\{Z_w\}_{w|q_2(w) \neq 0}$ .
13: Update the model, note that both data collected from stage 2 and stage 3 are used.
$$
\tilde {B} _ {X} \leftarrow \mathcal {O} _ {\text {o f f l i n e 3}} ^ {X} \left(\left\{Z _ {w} \right\} _ {w | q _ {1} (w) \neq 0 \text {a n d} q _ {2} (w) \neq 0}\right)
$$
14: end for
15: Return  $\tilde{B}_X$
As long as the number of target samples satisfies
$$
n _ {t a r g e t} \geq \tilde {\Omega} ((k + \log (1 / \delta)) \varepsilon^ {- 2}), \quad \dot {n} _ {t a r g e t} \gtrsim \tilde {\Omega} \left(\varepsilon^ {- \frac {4}{3}} (k ^ {*}) ^ {\frac {2}{3}} \sqrt {k} \left(d _ {W} ^ {\frac {1}{2}} \underline {{\sigma}} ^ {- \frac {4}{3}} + k ^ {- \frac {2}{3}} d _ {W} ^ {\frac {1}{6}} \bar {\kappa} ^ {2} \underline {{\sigma}} ^ {- \frac {1}{3}}\right)\right)
$$
Comparison with passive learning. By choosing  $\{e_i\}_{i\in [d_W^{\mathrm{source}}]}$  as a fixed source set, we reduce the problem to a discrete setting and compare it with the passive learning. In [9], the authors get  $N_{\mathrm{total}}$  as most  $\frac{k d_X d_W\|\mathbb{E}_{w_0\sim\nu_{\mathrm{target}}}B_W w_0 w_0^\top B_W^\top\|}{\sigma_{\min}^2(B_W^{\mathrm{source}})}\varepsilon^{-2}$ . We first consider the cases in their paper that the target task is uniformly spread  $\|\mathbb{E}_{w_0\sim\nu_{\mathrm{target}}}B_W w_0 w_0^\top B_W^\top\| = \frac{1}{k}$ .
- When the task representation is well-conditioned  $\sigma_{\mathrm{min}}^2 (B_W^{\mathrm{source}}) = \frac{d_W}{k}$ . We have a passive one as  $\widetilde{\mathcal{O}}(kd_X\varepsilon^{-2})$  while the active one  $\widetilde{\mathcal{O}}(kd_X\frac{k^2}{d_W}\varepsilon^{-2})$  (See Lemma B.8 for details), which suggests as long as  $d_W \gg k^2$ , our active learning algorithm gain advantage even in a relatively uniform spread data and representation conditions.
- Otherwise, we consider the extreme case that  $\sigma_{\mathrm{min}}^2(B_W^{\mathrm{source}}) = 1$ . We have passive one  $\widetilde{\mathcal{O}}(d_X d_W \varepsilon^{-2})$  while the active one  $\widetilde{\mathcal{O}}(k^3 d_X \varepsilon^{-2})$ . Notice here we require  $d_W \gg k^3$ .
Both of them indicate the necessity of considering the continuous case with large  $d_W$  even if everything is uniformly spread. On the other hand, whether we can achieve the same result as the passive one when  $d_W \leq k^3$  remains to be explored in the future.
We then consider the single target  $w_{0}$  case.
- With well-conditioned  $B_W$ , the passive one now has sample complexity  $\mathcal{O}(k^2 d_X \varepsilon^{-2})$  while the active gives a strictly improvement  $\mathcal{O}\left(\frac{k^2 d_X}{d_W} \varepsilon^{-2}\right)$ .
- With ill-conditioned  $B_W$  where  $\sigma_{\min}(B_W) = 1$  and  $\max_i \| W_i^* \| = 1$ , that is, only a particular direction in source space contributes to the target. The Passive one now has sample complexity  $\mathcal{O}(kd_Xd_W\varepsilon^{-2})$  while our active one only has  $kd_X\varepsilon^{-2}$ , which demonstrates the benefits of our algorithm in unevenly distributed source space.
Comparison with previous active learning. By using the same discrete reduction and set single target  $w_0$ , we compare our result with the current state-of-art active representation algorithm in [23]. They achieve  $\widetilde{\mathcal{O}}(kd_X\|\nu\|_1^2\varepsilon^{-2})$ , where  $\nu = \arg\min_{\nu}\|\nu\|_1$  s.t  $B_W\nu = B_Ww_0$ . On the other hand, our active one gives  $\widetilde{\mathcal{O}}(kd_X\|w^*\|_2^2\varepsilon^{-2})$ , where  $w^* = \arg\min_{\nu}\|\nu\|_2$  s.t  $B_W\nu = B_Ww_0$ , which is strictly better than the discrete one. This again indicates the separation between continuous and discrete cases where in fixed discrete sets, the  $L_1$  norm regularization is strictly better than  $L_2$ .
Furthermore, when a fixed discrete set is given, which is exactly the setting in [23]. Their algorithm can be seen as a computationally efficient reduction under ours.(Appendix B.5.)
Save task number. When ignoring the short-term initial warm-up stage, we only require maintaining  $\widetilde{\mathcal{O}}(k + \log(N_{\mathrm{total}}k^{*}))$  number of source tasks, where the first term comes from  $q_{1}$  in the target-agnostic stage and the second term comes from  $q_{2}$  in the target-aware stage.
5 Experiment
In this section, we provide experimental results under different instantiations of the Algorithm 1, and all of them show the effectiveness of our strategy both in target-aware and target-agnostic settings.
5.1 Settings
Datasets and problem definition. Our results cover the different combinations of  $\psi_{X},\phi_{X},\psi_{W}$  as shown in Table 1. Here we provide a brief introduction for the three datasets and postpone the details into Appendix E.2
identity ψW	nonlinear ψW
identity ψX and linear φX	synthetic, drone	NA
nonlinear ψX and linear φX	synthetic	pendulum simulator
identity ψX and nonlinear φX	synthetic, drone	NA
Table 1: Summary of different instantiations
- Synthetic data. We generate data that strictly adhere to our data-generating assumptions and use the same architecture for learning and predicting. When  $\phi_X$  is nonlinear, we use a neural network  $\phi_X$  to generate data and use a slightly larger neural net for learning. The goal for synthetic data is to better illustrate our algorithm as well as serve as the first step to extend our algorithm on various existing datasets.
- Pendulum simulator. To demonstrate our algorithm in the continuous space, we adopt the multi-environment pendulum model in [24] and the goal is to learn a  $w$ -dependent residual dynamics model  $f(x, w) \in \mathbb{R}$  where  $x$  is the pendulum state and  $w \in \mathbb{R}^5$  including external wind, gravity and damping coefficients.  $f(x, w)$  is highly nonlinear with respective to  $x$  and  $w$ . Therefore we use known non-linear feature operators  $\psi_X, \psi_W$ . In other words, this setting can be regarded as a misspecified linear model. It is also worth noting that due to the non-invertibility of  $\psi_W$ , the explicit selection of a source via a closed form is challenging. Instead, we resort to an adaptive sampling-based method discussed in Section 3. Specifically, we uniformly sample  $w$  from the source space, select the best  $w'$ , and then uniformly sample around this  $w'$  at a finer grain. Our findings indicate that about 5 iterations are sufficient to approximate the most relevant source.
- Real-world drone flight dataset [7]. The Neural-Fly dataset [7] includes real flight trajectories using two different drones in various wind conditions. The objective is to learn the residual
$^{2}$ Github Link: https://github.com/cloudwaysX/ALMultiTask_Robotics
aerodynamics model  $f(x, w) \in \mathbb{R}^3$  where  $x \in \mathbb{R}^{11}$  is the drone state (including velocity, attitude, and motor speed) and  $w$  is the environment condition (including drone types and wind conditions). We collect 6 different  $w$  and treat each dimension of  $f(x, w)$  as a separate task. Therefore  $w$  is reformulated as a one-hot encoded vector in  $\mathbb{R}^{18}$ .
For each dataset/problem, we can choose different targets. For simplicity, in the following subsection, we present results for one target task for each problem with 10 random seeds regarding random data generation and training, and put more results in Appendix E. In all the experiments, we use a gradient-descent joint training oracle, which is a standard approach in representation learning.
5.2 Results
Those results encapsulate the effectiveness of active learning in terms of budget utilization and test loss reduction. In the drone dataset, we further demonstrate its ability in identifying relevant source tasks (see Figure 2). We note that in two robotics problems (pendulum simulation and real-world drone dataset), the active learning objective is to learn a better dynamics model. However, in the pendulum simulation, we deploy a model-based nonlinear controller which translates better dynamics modeling to enhanced control performance (see Figure 1 and Appendix E.2).
Target-aware AL	Target-agnostic AL
identity ψX and linear φX	38.7%	51.6%
nonlinear ψX and linear φX	38.7%	45.2%
identity ψX and non-linear φX	32.0%	68.0%
Table 2: Results on synthetic data. Using the test loss of the final output model from passive learning as a baseline, we show the ratio between the budget required by target-aware/target-agnostic active learning to achieve a similar loss and the budget required by passive learning.
Figure 1: Results on pendulum simulator for a specific target. Left: The test loss of the estimated model  $\hat{f}$ . The passive strategy suffers from negative transfer while the active strategy steadily decreases. Right: The control error using final output  $\hat{f}$ . Here we use a model-based nonlinear policy  $\pi(x, \hat{f})$ . The model learned from active strategy leads to better control performance.
Figure 2: Results on the real drone dataset [7] with target drone_type_A_30_z. Source data includes two drone types A and B, six wind speeds from 0 to 50, and three directions x-y-z. We present results for linear  $\phi_{X}$  here and postpone the non-linear  $\phi_{X}$  case in Appendix E.3. Left: The test loss of the estimated bilinear model  $\tilde{f}$ . The passive strategy converges slower than the active strategy. Right: Top 10 the most similar source tasks. Given the target environment, the algorithm successfully finds the other drone_type_A environments as relevant sources. See more explanations in Appendix E.3.