SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alon Albalak
University of California,
alon_albalak@ucsb.edu
Colin Raffel
craffel@gmail.com
William Yang Wang
william@cs.ucsb.edu
Few-shot learning is valuable in many real-world applications, but learning a generalizable model without overfitting to the few labeled datapoints is challenging. In this work, we focus on Few-shot Learning with Auxiliary Data (FLAD), a training paradigm that assumes access to auxiliary data during few-shot learning in hopes of improving generalization. Previous works have proposed automated methods for mixing auxiliary and target data, but these methods typically scale linearly (or worse) with the number of auxiliary datasets, limiting their practicality. In this work we relate FLAD to the explore-exploit dilemma that is central to the multi-armed bandit setting and derive algorithms whose computational complexity is independent of the number of auxiliary datasets, allowing us to scale to  $100 \times$  more auxiliary datasets than prior methods. We propose two algorithms - EXP3-FLAD and UCB1-FLAD - and compare them with prior FLAD methods that either explore or exploit, finding that the combination of exploration and exploitation is crucial. Through extensive experimentation we find that our methods outperform all pre-existing FLAD methods by  $4\%$  and lead to the first 3 billion parameter language models that outperform the 175 billion parameter GPT-3. Overall, our work suggests that the discovery of better, more efficient mixing strategies for FLAD may provide a viable path towards substantially improving generalization in few-shot learning. All of our code is available at github.com/alon-albalak/FLAD.
Few-shot learning is an attractive learning setting for many reasons: it promises efficiency in cost and time, and in some scenarios data is simply not available due to privacy concerns or the nature of the problem. However, few-shot learning is also a challenging setting that requires a delicate balance between learning the structure of the feature and label spaces while preventing overfitting to the limited training samples [1, 2, 3]. One approach to improving the generalizability of models in the few-shot setting is Few-shot Learning with Auxiliary Data (FLAD), where additional auxiliary datasets are used to improve generalization on the target few-shot task [4, 5, 6, 7].
However, FLAD methods introduce their own challenges, including increased algorithmic and computational complexity. Specifically, incorporating auxiliary data during training introduces a large space of design choices (e.g. how and when to train on auxiliary data). Manually designing the curriculum for training on large quantities of auxiliary data is not feasible due to the combinatorially large search space, and hand-picking which auxiliary data to use based on heuristics (e.g. from the same domain or task as the target few-shot dataset) can lead to sub-optimal results [8]. Delegating such choices to an algorithm can lead to better solutions, as demonstrated in the transfer learning [8, 9, 10], meta-learning [11, 12], multi-task learning [13, 14, 15, 16], and auxiliary learning literature [4, 17]. However, prior auxiliary learning algorithms often assume that only 1-3 related auxiliary datasets are
Figure 1: Overview of few-shot learning with auxiliary data (FLAD) as a multi-armed bandit problem. On the left is the learner which defines a policy  $\pi$  that determines which auxiliary dataset to sample from. On the right is the environment that includes the set of auxiliary datasets  $\mathcal{D}_A$ , target dataset  $\mathcal{D}_{\mathcal{T}}$ , and the model  $f_{\theta}$ . At each turn  $t$ , the following five steps take place, further described in Section 3.1: 1. The learner selects an auxiliary dataset  $\mathcal{D}_a$  according to its policy  $\pi$ . 2. The environment samples a batch  $\{\mathbf{x},\mathbf{y}\} \sim \mathcal{D}_a$ . 3. The model  $f_{\theta}$  calculates gradients for the sampled batch  $(\nabla_a)$  and the target dataset  $(\nabla_{\mathcal{T}})$ , then updates the parameters  $\theta$ . 4. A reward  $\mathcal{R}_{a,t}$  is calculated based on  $\nabla_a$  and  $\nabla_{\mathcal{T}}$ . 5. The learner updates  $\pi$  based on  $\mathcal{R}_{a,t}$ .
available and design algorithms whose computational complexity grows linearly (or worse) with the number of auxiliary datasets [18, 8], motivating the search for more efficient methods as the number of auxiliary datasets grows.
To overcome the challenges of prior works, we desire a FLAD algorithm that (1) makes no assumptions on available auxiliary data a-priori (in-domain, on-task, quality, quantity, etc.), (2) scales well with the number of auxiliary datasets, and (3) adds minimal memory and computational overhead. We design algorithms that satisfy our desiderata by drawing inspiration from the central problem in multi-armed bandit (MAB) settings: the exploration-exploitation trade-off [19, 20]. We relate the set of auxiliary datasets to the arms of a MAB and tailor the classic EXP3 [21] and UCB1 [22] algorithms to fit the FLAD framework by designing three efficient gradient-based reward signals. The combination of our MAB-based algorithms and efficient gradient-based rewards allows us to scale to  $100 \times$  more auxiliary datasets than previous methods. Figure 1 provides a basic illustration of how we formulate FLAD as a MAB problem.
To empirically validate our approaches, we focus on few-shot training of language models and utilize P3 [23], a readily available resource with hundreds of auxiliary language datasets. We evaluate our methods on the same held-out tasks as the T0 language model [16] and show that, when using the same collection of auxiliary datasets, our algorithms outperform a directly fine-tuned T0 by  $5.6\%$  (EXP3-FLAD) and  $5.7\%$  (UCB1-FLAD) absolute. Furthermore, incorporating all available datasets in P3 (i.e. not just those used to train T0) increases the improvement to  $9.1\%$  and  $9.2\%$ . Finally, we compare models trained with our methods against state-of-the-art few-shot methods, finding that our methods improve performance by  $>3\%$ , even though one model utilizes a large collection of unlabeled target dataset samples. Furthermore, to the best of our knowledge, our methods lead to the first 3 billion parameter model that improves over 175B GPT-3 using few-shot in-context learning.
- We connect FLAD to the MAB setting and focus on the exploration-exploitation trade-off by designing two algorithms, EXP3-FLAD and UCB1-FLAD along with three reward functions that are both simple and efficient (in space and computational complexity).
- We empirically validate that our methods improve few-shot performance of pretrained language models and show that strategies that employ only exploration or exploitation lead to sub-optimal performance.
- We perform case studies to better understand the dynamics of our reward functions and their interaction with the dynamics of large language model training.
A long history of works have found success when combining auxiliary data with target data [4, 24, 6, 25, 26, 5, 18, 7, 27, 28, 8]. Some works have explored the addition of auxiliary learning objectives to aid the learning of the target task [24, 26, 25, 5, 17]. More similar to our work are
methods that perform auxiliary learning by introducing additional data sources beyond the target data [4, 6, 18, 7, 27, 28, 8]. As opposed to the few-shot setting on which this work focuses, previous works have studied auxiliary learning in settings with large quantities of target data. For example, Chen et al. [18] and Verboven et al. [7] assume access to 10,000 labeled target samples, Ivison et al. [28] and Lin et al. [27] assume access to 1,000s of unlabeled target samples, and Du et al. [6] and Albalak et al. [8] assume access to 100s of labeled target samples. Additionally, many of the previous works that study auxiliary learning have only considered settings with 1-3 auxiliary datasets [6, 18, 7, 8]. For example, Verboven et al. [7] propose a task-weighting method that requires solving a system of equations that becomes underspecified with multiple auxiliary tasks, limiting their method to only a single auxiliary task. Furthermore, Chen et al. [18] experiment with 3 auxiliary tasks because their method requires learning a target-aware classifier for each source task, so the computation scales as  $O(|\mathcal{A}||\mathcal{T}|)$  where  $|\mathcal{A}|$  is the number of auxiliary tasks and  $|\mathcal{T}|$  is the number of target tasks, making it impractical to scale to large numbers of source and target tasks. In this work, we focus on improving auxiliary learning with very few target samples (20-70 samples) by scaling up the number of auxiliary datasets orders of magnitude greater than previous work. In order to scale up the learning process, efficiency is a central concern of this work, unlike prior works.
Data selection studies a similar (but distinct) problem where the goal is to selectively utilize a subset of a single large dataset rather than selecting data from auxiliary datasets. Recent research on data selection has found that intelligent data selection can provide significant improvements to model performance [29, 30, 31, 32].
3 Multi-armed bandits for few-shot learning with auxiliary data
In this section, we first define the few-shot learning with auxiliary data (FLAD) setting. Then, we formulate FLAD as a multi-armed bandits (MAB) problem, shown in Figure 1. Next, we define reward functions that are efficient to compute and appropriate for FLAD. Finally, we describe our adaptations of two popular MAB algorithms: EXP3-FLAD and UCB1-FLAD.
FLAD problem setting. Few-shot learning with auxiliary data (FLAD) fits into the following setting: assume access to a large set of auxiliary datasets  $\mathcal{D}_{\mathcal{A}}$  where, for all  $a \in \mathcal{A}$ ,  $\mathcal{D}_a$  is an individual auxiliary dataset. Given a small quantity of data belonging to a target dataset  $\mathcal{D}_{\mathcal{T}}$ , the goal of FLAD is to find parameters  $\theta$  of a model  $f_{\theta}$  that achieve high performance on the unknown distribution underlying  $\mathcal{D}_{\mathcal{T}}$  while utilizing only the available data,  $\mathcal{D}_{\mathcal{T}} \cup \mathcal{D}_{\mathcal{A}}$ .
Formulating FLAD as MAB. In this work, we adopt the multi-armed bandit (MAB) setting by formulating FLAD as a Markov decision process [33] and defining a learner and environment, illustrated in Figure 1. The learner consists of a policy  $\pi$  defining a selection strategy over all  $\mathcal{D}_a\in \mathcal{D}_A$ . The environment consists of the target dataset  $\mathcal{D}_{\mathcal{T}}$ , auxiliary datasets  $\mathcal{D}_{\mathcal{A}}$ , and model  $f_{\theta}$ . In this formulation the learner interacts with the environment over  $N$  rounds. At each round  $t$  the learner selects one of the environment's  $|\mathcal{A}|$  datasets  $\mathcal{D}_a\in \mathcal{D}_A$ . Next, the environment samples a batch  $\{\mathbf{x},\mathbf{y}\} \sim \mathcal{D}_a$  and calculates the gradient w.r.t.  $\theta$  using a task-appropriate loss function as  $\nabla_{a} = \nabla_{\theta}\mathcal{L}(f_{\theta},\mathbf{x},\mathbf{y})$ . Then, the environment computes the target gradient  $\nabla_{\mathcal{T}} = \nabla_{\theta}\mathcal{L}(f_{\theta},\mathcal{D}_{\mathcal{T}})$  and updates model parameters w.r.t.  $\nabla_{\mathcal{T}} + \nabla_{a}$ . Finally, the learner uses a gradient-based reward  $\mathcal{R}_{a,t}(\nabla_a,\nabla_\tau)$  to update its policy  $\pi$ . See Appendix A and Lattimore & Szeptsvari [34] for further details on multi-armed bandits.
Designing the reward functions. We design the reward function  $\mathcal{R}$  with our desiderata in mind. To ensure that our algorithm adds minimal memory and computational overhead we consider rewards that utilize information intrinsic to the model and the losses being optimized, not an external model or metric (e.g. accuracy or BLEU). In this work we propose three gradient-based reward functions inspired by previous works: gradient alignment [6, 24, 35], gradient magnitude similarity [36, 37], and their aggregation. Formally, at turn  $t$  let  $\nabla_{a}$  be the gradient of the auxiliary batch and  $\nabla_{\mathcal{T}}$  be the target dataset gradient. Gradient alignment is defined as  $\mathcal{R}_{a,t}^{GA} = \frac{\nabla_a \cdot \nabla_\tau}{\|\nabla_a\|_2 \|\nabla_\tau\|_2}$ , i.e. the cosine similarity between the gradients of the sampled auxiliary dataset batch and the whole target dataset. Gradient magnitude similarity is defined as  $\mathcal{R}_{a,t}^{GMS} = \frac{2\|\nabla_a\|_2 \|\nabla_\tau\|_2}{\|\nabla_a\|_2^2 + \|\nabla_\tau\|_2^2}$  so that when the two gradients have equal magnitude, this value is equal to 1 and as the magnitudes differ the value goes to zero. In
addition to the individual reward functions, we also consider an aggregate reward. To ensure that the aggregate is not dominated by either individual reward, we normalize  $\mathcal{R}^{GA} \in [0,1]$ , the same range as  $\mathcal{R}^{GMS}$  and define the aggregate to be their sum:  $\mathcal{R}_{a,t}^{AGG} = \frac{1 + \mathcal{R}_{a,t}^{GA}}{2} + \mathcal{R}_{a,t}^{GMS}$ . We provide further discussion on the design of reward functions in Section 6.
3.2 Adapting the EXP3 algorithm.
EXP3 Background We base our first algorithm, EXP3-FLAD, on the EXP3 algorithm [21] ("Exponential-weight algorithm for Exploration and Exploitation"). EXP3 targets the adversarial MAB setting, which assumes that the reward-generating process is controlled by an adversary who is given access to the learner's policy  $\pi$  and determines the sequence of rewards,  $(R_{a,t})_{t=1}^{N}$ , for each arm prior to play [38]. We consider the adversarial MAB formulation due to the highly non-convex loss landscape of deep neural networks and our use of stochastic gradient descent-based optimization methods. These factors imply that we cannot guarantee our rewards to be stationary, independent, or follow any particular distribution (e.g. Gaussian). Further details on adversarial MAB are included in Appendix A and in [21].
In EXP3-FLAD, the learner selects arms according to a Gibbs distribution based on the empirically determined importance-weighted rewards of arms [39]. To allow for exploration, we mix the Gibbs distribution with a uniform distribution [21]. Formally, let  $\mathcal{E}_t$  be the exploration rate at turn  $t$  and, recalling that  $K = |\mathcal{A}|$  is the number of auxiliary datasets, then  $\pi$  defines the probability of selecting a given arm  $a \in \mathcal{A}$  as the linear combination of Gibbs and uniform distributions  $\pi_t(a) = (1 - K\mathcal{E}_t)\frac{\exp(\mathcal{E}_{t-1}\hat{R}_a)}{\sum_{a'} \exp(\mathcal{E}_{t-1}\hat{R}_{a'})} + \mathcal{E}_t$  where  $\hat{R}_{a,t}$  is the importance weighted reward  $\hat{R}_{a,t} = \hat{R}_{a,t-1} + \frac{R_{a,t}}{\pi_{t-1}(a)}$ . We want the learner to explore more in early training than in later stages, so we use a decaying exploration rate  $\mathcal{E}_t = \min\left\{\frac{1}{K}, \sqrt{\frac{\ln K}{K \cdot t}}\right\}$  as proposed by Seldin et al. [39]. The use of an importance-weighted estimated reward compensates the rewards of actions that are less likely to be chosen, guaranteeing that the expected estimated reward is equal to the actual reward for each action. EXP3-FLAD is designed to be nearly optimal in the worst case, but due to the exploration rate it will select "bad" actions at a rate of  $\mathcal{E}_t$ . The exploration of EXP3-FLAD combined with importance-weighting allows the policy to handle non-stationary reward-generating processes.
3.3 Adapting the UCB1 algorithm.
UCB1 background. While EXP3-FLAD is applicable in unconstrained settings with highly stochastic and non-stationary rewards, it can be outperformed by other algorithms in settings that are constrained. One such algorithm is the upper confidence bound (UCB1) algorithm [22], which was originally designed to be optimal for stationary, normally distributed reward functions. Nevertheless, variants of UCB1 have been demonstrated to be effective in a range of settings, such as those involving non-stationary, sub-Gaussian, or heavy-tailed distributions [40, 41]. The UCB1 algorithm and its variants assign each arm a value called the upper confidence bound based on Hoeffding's inequality [42] and are based on the principle of optimism in the face of uncertainty, meaning that with high probability the upper confidence bound assigned to each arm is an overestimate of the unknown mean reward.
In UCB1-FLAD, the learner greedily selects arms according to their upper confidence bound. UCB1 is originally designed for stationary reward-generating processes, so to accommodate non-stationarity we include an exponential moving average when estimating the mean reward for a given arm. Formally, let  $R_{a,t}$  be the observed reward for arm  $a$  at turn  $t$ , then we calculate the estimated mean reward as  $\hat{R}_a = (1 - \beta)\hat{R}_a + \beta R_{a,t}$  where  $\beta$  is the smoothing factor. Then, we define the upper confidence bound to be  $UCB_{a,t} = \hat{R}_a + \sqrt{\frac{2\ln t}{n_a}}$ . In the original MAB setting all interactions with the environment occur online, but FLAD is a unique situation where the learner can interact with the auxiliary data prior to training. To take advantage of this, rather than initializing estimated rewards with a single mini-batch, we initialize them with larger data quantities to improve the approximation of the true dataset gradients. This is done for each auxiliary dataset by calculating the gradient  $\nabla_{a} = \nabla_{\theta}\mathcal{L}(f_{\theta},\mathbf{x},\mathbf{y})$ , where the number of samples in  $\{\mathbf{x},\mathbf{y}\}$  can be significantly larger than a mini-batch, and can be up to the size of the full dataset. In practice, we use 1,000 examples which is computed in  $\sim 2$  minutes on a single GPU.
Algorithms The EXP3-FLAD and UCB1-FLAD algorithms are visualized in Figure 1. At each turn, both methods will first select an auxiliary dataset  $\mathcal{D}_a$ . EXP3-FLAD first computes the current exploration rate  $\mathcal{E}_t$  and samples  $\mathcal{D}_a$  according to the distribution defined by  $\pi_t(\mathcal{A})$ , while UCB1-FLAD greedily selects  $\mathcal{D}_{a^*}$  corresponding to the arm with largest upper confidence bound,  $a^* = \arg \max_{a\in \mathcal{A}}UCB_{a,t}$ . Next, for both methods, the environment samples a batch from the selected dataset,  $\{\mathbf{x},\mathbf{y}\} \sim \mathcal{D}_a$ , and calculates the gradient  $\nabla_{a} = \nabla_{\theta}\mathcal{L}(f_{\theta},\mathbf{x},\mathbf{y})$ . Let  $G$  be the number of rounds between model updates, then the previous steps will repeat  $G$  times, at which point the environment calculates the gradient of the target dataset  $\nabla_{\theta}\mathcal{L}(f_{\theta},\mathcal{D}_{\mathcal{T}})$  and updates the model w.r.t.  $\nabla_{\mathcal{T}} + \sum_{a}\nabla_{a}$ . Finally, EXP3-FLAD calculates the importance-weighted reward for each auxiliary batch using the observed rewards, while UCB1-FLAD calculates the smoothed estimated mean reward. Pseudocode is found in Appendix B.
Models. For our experiments, we utilize encoder-decoder Transformer models from the T5 family of pre-trained language models [43]. Specifically, we experiment with LM-adapted T5 (T5-LM) and T0. The T5-LM model further trains the T5.1.1 model for 100,000 steps (corresponding to 100B tokens) from the C4 dataset [43] on the prefix language modeling objective [44]. The T0 model was initialized from T5-LM and further trained on a multitask mixture of prompted datasets as described by Sanh et al. [16]. We repeat each experiment with T5-LM XL (hereafter T5-XL) and T0-3B as our base model. Both models use the same architecture with 2.85 billion parameters, and we used model checkpoints from Hugging Face Transformers [45]).
Target datasets. We obtain all datasets from Hugging Face Datasets<sup>1</sup>, and cast them to the text-to-text format by applying prompt templates from the Public Pool of Prompts (P3) [23] that was used to train T0. To evaluate our few-shot methods, we utilize the same held-out datasets as T0, which cover four distinct tasks: sentence completion (COPA [46], HellaSwag [47], Story Cloze [48]), natural language inference (ANLI [49], CB [50], RTE [51]), coreference resolution (WSC [52], Winogrande [53]), and word sense disambiguation (WiC [54]). For each dataset, we randomly sample five few-shot splits from their training data, containing the same number of training examples as previous works, between 20 to 70 [55, 56]. We further divide each split into equal training and validation partitions for true few-shot learning [57](e.g. 10 train and 10 validation samples for HellaSwag). Only ANLI datasets have a publicly available test set, so for all other datasets we evaluate models on the original validation set (not utilized for few-shot training or validation).
Auxiliary datasets. We compare the performance of our methods using two sets of auxiliary data and never include any of the target datasets as part of auxiliary data. First, we use the collection of datasets used for multitask training of T0 (henceforth referred to as T0Mix), including 35 unique datasets covering question answering, sentiment analysis, topic classification, summarization, paraphrase detection and structure-to-text. Second, we utilize all datasets in P3 [23] (which forms a superset of T0Mix) and prevent data leakage by filtering out datasets that overlap with any target dataset, leading to 260 available datasets (list in Appendix H). For each auxiliary dataset, we use at most 10,000 of the dataset's examples.
Baseline methods. We compare our proposed methods with several FLAD and non-FLAD baselines. Target-Only (non-FLAD) directly fine-tunes the base model on the target dataset (i.e. without using auxiliary data). Explore-Only [8] is a commonly used FLAD method which simultaneously trains on auxiliary and target data by mixing auxiliary datasets equally. Originally called Multitask in [8], we call this Explore-Only because it is equivalent to continuously exploring auxiliary data and never exploiting knowledge of its relation to the target data. Exploit-Only computes gradient alignment prior to training (as in UCB1), and multitask-trains the model by mixing auxiliary datasets according to a Gibbs distribution over the alignments (similar to that in EXP3), resulting in an algorithm that exploits the relations determined prior to training, but never exploring. Both explore- and exploit-only mix target and auxiliary data with a ratio of  $M$  times the highest auxiliary sampling probability. For instance, explore-only with  $M = 5$  and  $\mathcal{D}_A = \mathrm{T0Mix}$  has a  $1/35$  probability to sample auxiliary dataset  $\mathcal{D}_a \in \mathcal{D}_A$  and a  $5/35$  probability for the target dataset. Loss-Scaling [6] is a FLAD method similar to EXP3 and UCB1; the main difference being that it scales auxiliary batch losses by their gradient alignment instead of modifying sampling probabilities. Du et al. [6] originally propose to
<sup>1</sup> https://huggingface.co/datasets
Training Method	BASE MODEL Auxiliary Data	T5-XL	T0-3B
T0Mix	P3	T0Mix	P3
Target-Only	52.82	56.44
Loss-Scaling [6] (GA)	53.22	55.19	59.47	60.66
Loss-Scaling [6] (GMS)	55.98	56.40	60.47	60.70
Explore-Only [8]	59.18	60.64	61.17	62.77
Exploit-Only [8]	59.79	60.49	60.87	62.87
EXP3-FLAD (RGA)	61.50	64.07	62.87	65.98
UCB1-FLAD (RGA)	62.01	65.52	62.35	66.29
EXP3-FLAD (RGMS)	61.72	65.57	62.78	65.51
UCB1-FLAD (RGMS)	61.67	65.21	62.85	66.00
EXP3-FLAD (RAGG)	62.05	65.47	62.84	66.84
UCB1-FLAD (RAGG)	62.08	65.63	62.93	66.29