SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yifang Chen $^{1}$ , Yingbing Huang $^{2}$ , Simon S. Du $^{1*}$ , Kevin Jamieson $^{1*}$ , Guanya Shi $^{3*}$
<sup>1</sup> Paul G. Allen School of Computer Science & Engineering University of Washington, Seattle, WA
{yifangc, ssdu, jamieson, guanyas}@cs.washington.edu
<sup>2</sup> University of Illinois Urbana-Champaign, Champaign, IL {yh21}@illinois.edu
$^{3}$  Robotics Institute, Carnegie Mellon University, Pittsburgh, PA {{guanyas}  $@$  andrew.cmu.edu * Equal advising
Representation learning based on multi-task pretraining has become a powerful approach in many domains. In particular, task-aware representation learning aims to learn an optimal representation for a specific target task by sampling data from a set of source tasks, while task-agnostic representation learning seeks to learn a universal representation for a class of tasks. In this paper, we propose a general and versatile algorithmic and theoretic framework for active representation learning, where the learner optimally chooses which source tasks to sample from. This framework, along with a tractable meta algorithm, allows most arbitrary target and source task spaces (from discrete to continuous), covers both task-aware and task-agnostic settings, and is compatible with deep representation learning practices. We provide several instantiations under this framework, from bilinear and feature-based nonlinear to general nonlinear cases. In the bilinear case, by leveraging the non-uniform spectrum of the task representation and the calibrated source-target relevance, we prove that the sample complexity to achieve  $\varepsilon$ -excess risk on target scales with  $(k^{*})^{2}\| v^{*}\|_{2}^{2}\varepsilon^{-2}$  where  $k^{*}$  is the effective dimension of the target and  $\| v^{*}\|_{2}^{2} \in (0,1]$  represents the connection between source and target space. Compared to the passive one, this can save up to  $\frac{1}{d_W}$  of sample complexity, where  $d_W$  is the task space dimension. Finally, we demonstrate different instantiations of our meta algorithm in synthetic datasets and robotics problems, from pendulum simulations to real-world drone flight datasets. On average, our algorithms outperform baselines by  $20\% - 70\%$ .
Recently, few-shot machine learning has enjoyed significant attention and has become increasingly critical due to its ability to derive meaningful insights for target tasks that have minimal data, a scenario commonly encountered in real-world applications. This issue is especially prevalent in robotics where data collection and training data is prohibitive to collect or even non-reproducible (e.g., drone flying with complex aerodynamics [1] or legged robots on challenging terrains [2]). One
$^{1}$ Code in https://github.com/cloudwaysX/ALMultiTask_Robotics
promising approach to leveraging the copious amount of data from a variety of other sources is multi-task learning, which is based on a key observation that different tasks may share a common low-dimensional representation. This process starts by pretraining a representation on source tasks and then fine-tuning the learned representation using a limited amount of target data ([3-7]).
In conventional supervised learning tasks, accessing a large amount of source data for multi-task representation learning may be easy, but processing and training on all that data can be costly. In real-world physical systems like robotics, this challenge is further amplified by two factors: (1) switching between different tasks or environments is often significantly more expensive (e.g., reset giant wind tunnels for drones [7]); (2) there are infinitely many environments to select from (i.e., environmental conditions are continuous physical parameters like wind speed). Therefore, it is crucial to minimize not only the number of samples, but the number of sampled source tasks, while still achieving the desired performance on the target task. Intuitively, not all source tasks are equally informative for learning a universally good representation or a target-specific representation. This is because source tasks can have a large degree of redundancy or be scarce in other parts of the task space. In line with this observation, Chen et al. [8] provided the first provable active representation learning method that improves training efficiency and reduces the cost of processing source data by prioritizing certain tasks during training with theoretical guarantees. On the other hand, many existing works [9-13] prove that it is statistically possible to learn a universally good representation by randomly sampling source tasks (i.e., the passive learning setting).
The previous theoretical work of [8] on active multi-task representation learning has three main limitations. First, it only focuses on a finite number of discrete tasks, treating each source independently, and therefore fails to leverage the connection between each task. This could be sub-optimal in many real-world systems like robotics for two reasons: (1) there are often infinitely many sources to sample from (e.g., wind speed for drones); (2) task spaces are often highly correlated (e.g., perturbing the wind speed will not drastically change the aerodynamics). In our paper, by considering a more general setting where tasks are parameterized in a vector space  $\mathcal{W}$ , we can more effectively leverage similarities between tasks compared to treating them as simply discrete and different. Secondly, the previous work only considers a single target, while we propose an algorithm that works for an arbitrary target space and distribution. This is particularly useful when the testing scenario is time-variant. Thirdly, we also consider the task-agnostic setting by selecting  $\mathcal{O}(k)$  representative tasks among the  $d_W$  high dimension task space, where  $k \ll d_W$  is the dimension of the shared representation. Although this result does not improve the total source sample complexity compared to the passive learning result in the bilinear setting [12], it reduces the number of tasks used in the training and therefore implicitly facilitates the training process.
In addition to those theoretical contributions, we extend our proposed algorithmic framework beyond a pure bilinear representation function, including the known nonlinear feature operator with unknown linear representation (e.g., random features with unknown coefficients), and the totally unknown nonlinear representation (e.g., deep neural network representation). While some prior works have considered nonlinear representations [9, 10, 14, 13] in passive learning, the studies in active learning are relatively limited [8]. All of these works only consider non-linearity regarding the input, rather than the task parameter. In this paper, we model task-parameter-wise non-linearity and show its effectiveness in experiments. Note that it particularly matters for task selections because the mapping from the representation space to task parameters to is no longer linear.
See more related works and how our problem scope is different from theirs in Appendix A.
1.1 Summary of contributions
- We propose the first generic active representation learning framework that admits any arbitrary source and target task space. This result greatly generalizes previous works where tasks lie in the discrete space and only a single target is allowed. To show its flexibility, we also provide discussions on how our framework can accommodate various supervised training oracles and optimal design oracles. (Section 3)
- We provide theoretical guarantees under a benign setting, where inputs are i.i.d. and a unit ball is contained in the overall task space, as a compliment to the previous work where tasks lie on the vertices of the whole space. In the target-aware setting, to identify an  $\varepsilon$ -good model our method requires a sample complexity of  $\widetilde{\mathcal{O}}(kd_X(k^*)^2 \| v^* \|_2^2 \min \{k^*, \kappa^2\} \varepsilon^{-2})$  where  $k^*$  is the effective dimension of the target,  $\kappa$  is the conditional number of representation matrix, and
$\| v^{*}\|_{2}^{2}\in (0,1]$  represents the connection between source and target space that will be specified in the main paper. Compared to passive learning, our result saves up to a factor of  $\frac{k^2}{d_W}$  in the sample complexity when targets are uniformly spread over the  $k$  -dim space and up to  $\frac{1}{d_W}$  when targets are highly concentrated. Our results further indicate the necessity of considering the continuous space by showing that directly applying the previous algorithm onto some discretized sources in the continuous space (e.g., orthonormal basis) can lead to worse result. Finally, ignoring the tasks used in the warm-up phases, in which only a few samples are required, both the target-aware and the target-agnostic cases can save up to  $\widetilde{\mathcal{O}} (k^{*} + k)$  number of tasks compared to the passive one which usually requires  $d_W$  number of tasks. (Section 4)
- We provide comprehensive experimental results under different instantiations beyond the benign theoretical setting, studying synthetic and real-world scenarios: 1) For the synthetic data setting in a continuous space, we provide results for pure linear, known nonlinear feature operator  $\psi_{X}$  and unknown nonlinear representation  $\phi_{X}$ . Our target-aware active learning (AL) approach shows up to a significant budget saving (up to  $68\%$ ) compared to the passive approach and the target-agnostic AL approach also shows an advantage in the first two cases. 2) In a pendulum simulation with continuous task space, we provide the results for known nonlinear feature operator  $\psi_{X}$  and  $\psi_{W}$  and show that our target-aware AL approach has up to  $20\%$  loss reduction compared to the passive one, which also translates to better nonlinear control performance. 3) Finally, in the real-world drone dataset with a discrete task space, we provide results for unknown linear and nonlinear representation  $\phi_{X}$  and show that our target-aware AL approach converges much faster than the passive one. (Section 5)
2 Preliminary
Multi-task (or multi-environments). Each task or environment is parameterized by a known vector  $w \in \mathbb{R}^{d_W}$ . We denote the source and target task parameter space as  $\mathcal{W}_{\mathrm{source}} \subset \mathbb{R}^{d_W}$ ,  $\mathcal{W}_{\mathrm{target}} \subset \mathbb{R}^{d_W}$ . These spaces need not be the same (e.g., they could be different sub-spaces). In the discrete case, we set  $w$  as a one-hot encoded vector and therefore we have in total  $d_W$  number of candidate tasks while in the continuous space, there exist infinitely many tasks. For convenience, we also use  $w$  as the subscript to index certain tasks. In addition, we use  $\nu_{\mathrm{source}} \in \Delta(\mathcal{W}_{\mathrm{source}})$ ,  $\nu_{\mathrm{target}} \in \Delta(\mathcal{W}_{\mathrm{target}})$  to denote the task distribution for the sources and targets.
Data generation. Let  $\mathcal{X} \in \mathbb{R}^{d_X}$  be the input space. We first assume there exists some known feature/augmentation operator  $\psi_X: \mathcal{X} \to \mathbb{R}^{d_{\psi_X} \geq d_W}$ ,  $\psi_W: \mathcal{W} \to \mathbb{R}^{d_{\psi_W} \geq d_W}$ , that can be some nonlinear operator that lifts  $w, x$  to some higher dimensional space (e.g., random Fourier features [15]). Notice that the existence of non-identical  $\psi$  indicates the features are not pairwise independent and the design space of  $\mathcal{W}_{\mathrm{source}}$  is not benign (e.g., non-convex), which adds extra difficulty to this problem.
Then we assume there exists some unknown underlying representation function  $\phi_X: \psi(\mathcal{X}) \to \mathcal{R}$  which maps the augmented input space  $\psi(\mathcal{X})$  to a shared representation space  $\mathcal{R} \in \mathbb{R}^k$  where  $k \ll d_{\psi_X}, k \leq d_{\psi_W}$ , and its task counterparts  $\phi_W: \psi(\mathcal{W}) \to \mathcal{R}$  which maps parameterized task space to the feature space. Here the representation functions are restricted to be in some function classes  $\Phi$ , e.g., linear functions, deep neural networks, etc.
In this paper, we further assume that  $\phi_W$  is a linear function  $B_W \in \mathbb{R}^{k \times d_{\psi_W}}$ . To be more specific, for any fixed task  $w$ , we assume each sample  $(x,y) \sim \nu_w$  satisfies
$$
y = \phi_ {X} \left(\psi_ {X} (x)\right) ^ {\top} B _ {W} \psi_ {W} (w) + \xi , \quad \xi \sim \mathcal {N} (0, \sigma^ {2}) \tag {1}
$$
For convenience, we denote  $Z_{w}$  as the collection of  $n_w$  sampled data  $(x_w^1, y_w^1), \ldots, (x_w^{n_w}, y_w^{n_w}) \sim \mu_w$ . We note that when  $\psi_X, \psi_W$  is identity and  $\phi_X$  is linear, this is reduced to standard linear setting in many previous papers [9, 11, 12, 8].
The task diversity assumption. There exists some distribution  $p \in \Delta(\mathcal{W}_{\mathrm{source}})$  that  $\mathbb{E}_{w \sim p} \lambda_{\min}(B_W \psi_W(w) \psi_W(w)^\top B_W^\top) > 0$ , which suggests the source tasks are diverse enough to learn the representation.
Data collection protocol. We assume there exists some i.i.d. data sampling oracle given the environment and the budget. To learn a proper representation, we are allowed access to an unlimited  $n_{\mathrm{source}}$  number of data from source tasks during the learning process by using such an oracle. Then at
the end of the algorithm, we are given a few-shot of mix target data  $Z_{\mathrm{target}} = \{Z_w\}_{w \sim \nu_{\mathrm{target}}}$  which is used for fine-tuning based on learned representation  $\hat{\phi}_X$ . Denote  $n_{\mathrm{target}}$  as the number of data points in  $Z_{\mathrm{target}}$ .
Data collection protocol for target-aware setting. When the target task is not a singleton, we additionally assume a few-shot of known environment target data  $\dot{Z}_{\mathrm{target}} := \{Z_w, w\}_{w \in \dot{W}_{\mathrm{target}}}$ , where  $|\dot{W}_{\mathrm{target}}| = \dim(\mathcal{W}_{\mathrm{target}})$  and  $\dot{W}_{\mathrm{target}} = \{\arg \max_{W \in \mathcal{W}_{\mathrm{target}}} \lambda_{\min}(WW^{\top})\}$ . Again denote  $\dot{n}_{\mathrm{target}}$  as the number of data points in  $\dot{Z}_{\mathrm{target}}$ , we have  $\dot{n}_{\mathrm{target}} \approx n_{\mathrm{target}}^{2/3} \ll n_{\mathrm{source}}$ .
Remark 2.1. Here  $|\dot{W}_{\mathrm{target}}|$  represents vectors that can cover every directions of  $\mathcal{W}_{\mathrm{target}}$  space. This extra  $\dot{Z}_{\mathrm{target}}$  requirement comes from the non-linearity of  $l_{2}$  loss and the need to learn the relationship between sources and targets. We want to emphasize that such an assumption implicitly exists in previous active representation learning [8] since  $\dot{Z}_{\mathrm{target}} = Z_{\mathrm{target}}$  in their single target setting. Nevertheless, in a passive learning setting, only mixed  $\dot{Z}_{\mathrm{target}}$  is required since no source selection process involves. Whether such a requirement is necessary for target-aware active learning remains an open problem.
Other notations. Let  $e_i$  to be one-hot vector with 1 at  $i$ -th coordinates and let  $\epsilon_i = 2^{-i}$ .
2.1 Goals
Expected excess risk. For any target task space  $\mathcal{W}_{\mathrm{target}}$  and its distribution  $\nu_{\mathrm{target}}$  over the space, as well as a few-shot examples as stated in section 2, our goal is to minimize the expected excess risk with our estimated  $\hat{\phi}_X$
$$
\operatorname {E R} \left(\hat {\phi} _ {X}, \nu_ {\text {t a r g e t}}\right) = \mathbb {E} _ {w _ {0} \sim \nu_ {\text {t a r g e t}}} \mathbb {E} _ {(x, y) \sim \nu_ {w _ {0}}} \| \hat {\phi} _ {X} \left(\psi_ {X} (x)\right) ^ {\top} \hat {w} _ {\text {a v g}} - y \| _ {2}
$$
where  $\hat{w}_{\mathrm{avg}} = \arg \min_{w}\sum_{(x,y)\in Z_{\mathrm{target}}}\| \hat{\phi}_X(\psi_X(x))w - y\| _2$  , which average model estimation that captures the data behavior under the expected target distribution. Note that the  $\mathcal{W}_{\mathrm{target}},\nu_{\mathrm{target}}$  are given in advance in the target-aware setting.
The number of tasks. Another side goal is to save the number of long-term tasks we are going to sample during the learning process. Since a uniform exploration over  $d_W^{\mathrm{source}}$ -dimension is unavoidable during the warm-up stage, we define long-term task number as
$$
\left| \left\{w \in \mathcal {W} _ {\text {s o u r c e}} \mid n _ {w} \geq \tilde {\Omega} (\varepsilon^ {- \alpha}) \right\} \right|
$$
where  $\alpha$  is some arbitrary exponent and  $\varepsilon$  is the target accuracy and  $n_w$  is number of samples sampled from task  $w$  as defined above.
3 A general framework
Our algorithm 1 iteratively estimates the shared representation  $\hat{\phi}_X, \hat{B}_W$  and the next target relevant source tasks which the learner should sample from by solving several optimal design oracles
$$
g (f, A) = \min  _ {q \in \Delta (\mathcal {W} _ {\text {s o u r c e}})} \lambda_ {\max } \left(\left(\int q (w) f (w) f (w) ^ {\top}\right) ^ {- 1} A\right) \tag {2}
$$
This exploration and exploitation (target-aware exploration here) trade-off is inspired by the classical  $\epsilon$ -greedy strategy, but the key difficulty in our work is to combine that with multi-task representation learning and different optimal design problems. The algorithm can be generally divided into three parts, and some parts can be skipped depending on the structure and the goal of the problem.
- Coarse exploration: The learner uniformly explores all the directions of the  $\mathcal{W}_{\mathrm{source}}$  (denoted by distribution  $q_{0}$ ) in order to find an initial  $k$ -dimension subspace  $V$  that well spans over the representation space (i.e.,  $\frac{1}{c} B_W B_W^\top \leq B_W V V^\top B_W^\top \leq c B_W B_W^\top$  for some arbitrary constant  $c \leq \frac{d_{\psi W}}{k}$ ). To give an intuitive example, suppose  $B_W \in \mathbb{R}^{2 \times d_{\psi W}^{\mathrm{source}} + 1}$  has the first half column equals  $e_1$  and the second half equals  $e_2$ . Then instead of uniformly choosing  $\{e_i\}_{i \in [d_{\psi W}^{\mathrm{source}}]$  task, we only need explore over two tasks  $V[1] = \sqrt{\frac{2}{d_{\psi W}^{\mathrm{source}}}} [1, 1, \ldots, 0, 0, \ldots], V[2] = \sqrt{\frac{2}{d_W^{\mathrm{source}}}} [0, 0, \ldots, 1, 1, \ldots]$ .
Algorithm 1 Active multi-task representation learning (general templates)
1: Inputs: Candidate source set  $\mathcal{W}_{\mathrm{source}}$ . Classes of candidate representation function  $\Phi_X, \Phi_W$  and the known feature operator  $\psi_X, \psi_W$ .
2: [Target-aware only] Inputs: Target set  $\mathcal{W}_{\mathrm{target}}$  and distribution  $\nu_{\mathrm{target}}$ . Few-shot sample  $\dot{Z}_{\mathrm{target}}$  as defined in the preliminary.
3: Stage 1: Coarse exploration. (Warm-up stage)
4: Set initial sampling distribution  $q_{0} = g(\psi_{W}, I_{d_{\psi_{W}}})$  where  $g$  is defined in Eqn. 2
5: Set  $n_0 \approx \mathrm{poly}(d_{\psi_X}, k) + \mathrm{poly}(d_{\psi_W}, k)$ . Collect  $n_0q_0(w)$  data for each task denoted as  $\{Z_w\}_{w|q_0(w) \neq 0}$  and update  $\hat{\phi}_X \gets \mathcal{O}_{\mathrm{offline}}^X(\{Z_w\}_{w|q_0(w) \neq 0}, \psi_X)$  and  $\hat{B}_W \gets \mathcal{O}_{\mathrm{offline}}^W(\{Z_w\}_{w|q_0(w) \neq 0}, \hat{\phi}_X)$
6: for  $j = 1,2,3,\ldots$  do
7: Stage 2: Fine target-agnostic exploration (Directly choose  $q_1^j = q_0$  when  $k = \Theta(d_W)$ )
8: Compute the exploration sampling distribution  $q_1^j = g(\hat{B}_W \circ \psi_W, I_k)$
9:  $n_1^j \approx \mathrm{poly}(d_{\psi_X}, k) \epsilon_j^{-\frac{4}{3}}$ . Collect  $n_1^j q_1^j(w)$  data for each task denoted as  $\{Z_w\}_{w|q(w) \neq 0}$  and update  $\hat{\phi}_X \gets \mathcal{O}_{\mathrm{offline}1}^X(\{Z_w\}_{w|q_1^j(w) \neq 0}, \psi_X)$  and  $\hat{B}_W \gets \mathcal{O}_{\mathrm{offline}}^W(\{Z_w\}_{w|q_1^j(w) \neq 0}, \hat{Z}_{\mathrm{target}}, \hat{\phi}_X)$ .
0: [Target-aware only] Stage 3: Fine target-aware exploration
11: Compute the exploitation sampling distribution  $q_2^j = g(\hat{B}_W \circ \psi_W, \Sigma_{\mathrm{regu}})$  where  $\Sigma_{\mathrm{regu}}$  is the regularized version of  $\hat{B}_W(\mathbb{E}_{w_0 \sim \nu_0} w_0 w_0^\top) \hat{B}_W^\top$  after clipping out insignificant eigenvalues.
12: Set  $n_2^j \approx \mathrm{poly}(d_{\psi_X}, k) \epsilon_j^{-2}$ . Collect  $n_2^j q_2^j(w)$  data for each task denoted as  $\{Z_w\}_{w|q_2^j(w) \neq 0}$  and update  $\hat{\phi}_X \gets \mathcal{O}_{\text{offline 3}}^X(\{Z_w\}_{w|q_1^j(w) \neq 0 \text{ and } q_2^j(w) \neq 0}, \psi_X)$ .
13: end for
14: Return  $\hat{\phi}_X$
We want to highlight that the sample complexity of this warm-up stage only scales with  $d_{\psi_X}$ ,  $k$  and the spectrum-related parameters of  $B_W$  (i.e.,  $\kappa(B_W)$ ,  $\sigma_{\min}(B_X)$ ), not the desired accuracy  $\varepsilon$ .
- Fine target-agnostic exploration: The learner iteratively updates the estimation of  $V$  and uniformly explore for  $\widetilde{\mathcal{O}}(\epsilon_j^{-\frac{4}{3}})$  times on this  $k$ , instead of  $d_{\psi_W}$  subspace, denoted by distribution  $q_1$ . (Note this  $\epsilon_j^{-\frac{4}{3}}$  comes from the exploration part in  $\epsilon$ -greedy, which is  $(n_2^j)^{\frac{2}{3}}$ ) Such reduction not only saves the cost of maintaining a large amount of physical environment in real-world experiments but also simplifies the non-convex multi-task optimization problem. Of course, when  $k = \Theta(d_{\psi_W})$ , we can always uniformly explore the whole  $(d_{\psi_W}$  space as denoted in the algorithm. Note that theoretically,  $q_1$  only needs to be computed once as shown in 4. In practice, to further improve the accuracy while saving the task number, the  $q_1$  can be updated only when a significant change from the previous one happens, which is adopted in our experiments as shown in appendix E.1.
- Fine target-aware exploration. In the task-awareness setting, the learner estimates the most-target-related sources parameterized by  $\{w\}$  based on the current representation estimation and allocates more budget on those, denoted by distribution  $q_{2}$ . By definition,  $q_{2}$  should be more sparse than  $q_{1}$  and thus allowing the final sample complexity only scales with  $k^{*}$ , which measures the effective dimension in the source space that is target-relevant.
Computational oracle for optimal design problem. Depending on the geometry of  $\{\psi_W(w)\}_{w\in \mathcal{W}_{\mathrm{source}}}$ , the learner should choose proper offline optimal design algorithms to solve  $g(f,A)$ . Here we propose several common choices. 1). When  $\mathcal{W}_{\mathrm{source}}$  contains a ball, we can approximate the solution via an eigendecomposition-based closed-form solution with an efficient projection as detailed in Section 4. 2) When  $\mathcal{W}_{\mathrm{source}}$  is some other convex geometry, we can approximate the result via the Frank-Wolfe type algorithms [16], which avoids explicitly looping over the infinite task space. 3) For other even harder geometry, we can use discretization or adaptive sampling-based approximation [17]. In our experiments, we adopt the latter one and found out that its running time cost is almost negligible in our pendulum simulator experiment in Section 5, where the  $\psi_W$  is a polynomial augmentation.
Offline optimization oracle  $\mathcal{O}_{\mathrm{offline}}^{X}$ . Although we are in the continuous setting, the sampling distribution  $q_{0}, q_{1}, q_{2}$  is sparse. Therefore, our algorithm allows any proper passive multi-task
learning algorithm, either theoretical or heuristic one, to plugin the  $\mathcal{O}_{\mathrm{offline}}^{X}$ . Some common choices include gradient-based joint training approaches[18-21], the general non-convex ERM [9] and other more carefully designed algorithms [12, 22]. We implement the first one in our experiments (Section 5) to tackle the nonlinear  $\psi_{X}, \phi_{X}$  and give more detailed descriptions of the latter two in Section 4 and Appendix B.1 to tackle the bilinear model.
4 A theoretical analysis under the benign  $\mathcal{W}_{\mathrm{source}}$  setting
4.1 Assumptions