SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Main results. Each cell contains the score of training a base model (top row) with auxiliary data (second row) using the specified training method (left column), averaged across 11 target datasets on 5 random seeds (each cell is the average of 55 experiments). Target-Only does not utilize auxiliary data. Bolded scores are those with highest mean for a given base model and auxiliary dataset (column-wise), underlined scores are those where a Wilcoxon rank-sum test fails to find significant difference from the highest score  $(p > 0.05)$ . Expanded results are found in Appendix D.
use gradient alignment (Loss-Scaling  $(GA))$ , but we also propose a version that scales losses by gradient magnitude similarity (Loss-Scaling  $(GMS))$ .
Training details. For the target-only baseline, we use learning rates in  $\{1e - 4,3e - 4\}$ . For all other methods, we always use a learning rate of  $1e - 4$ . For target-, explore-, and exploit-only baselines we use batch sizes in  $\{32,128\}$ . For loss-scaling, EXP3-FLAD, and UCB1-FLAD we use mini-batches of 8 samples and let  $G$  be in  $\{4,16\}$  to match the batch size of all methods. For explore- and exploit-only, we use a target dataset mixing ratio of  $M\in \{1,5,10\}$ . For all experiments we use the Adafactor optimizer [58] and validation-based early stopping for model checkpoint selection. In preliminary experiments we consider rewards using gradients from various model partitions: the full model, encoder-only, decoder-only, and the weights of the output vocabulary matrix (language modeling head). We find that using the parameters from the language modeling head provides the best performance and contains only  $2.3\%$  of the full model parameters, significantly reducing memory consumption. For UCB1-FLAD we found the smoothing factor  $\beta = 0.9$  to work well in preliminary experiments and initialize auxiliary dataset gradient alignment using 1,000 samples from each auxiliary dataset. Additional implementation details can be found in Appendix C
Experiment procedure. The FLAD experiment process involves training a model that is specialized for each target dataset. For each proposed method and baseline, we train and evaluate a model on each of the 11 target datasets. We repeat training and evaluation on 5 random seeds and include the aggregated results in Table 1. Each cell shows the accuracy averaged across all 55 (11 target datasets, 5 random seeds) experiments. This experimental process is performed for each training method on both models and auxiliary datasets. We include the non-aggregated results in Appendix D.
5 Findings and analysis
In Table 1 we compare the empirical results of our MAB-based methods (EXP3-FLAD and UCB1-FLAD) and corresponding baselines on 11 target datasets (expanded results in Appendix D. For each base model and auxiliary data combination (each column) EXP3-FLAD and UCB1-FLAD outperform all the baselines. In fact, we find that for every single task our methods always perform equal to or better than the baselines. This demonstrates that our MAB-based methods provide a strong improvement in few-shot generalization over previous FLAD methods. For a fair comparison where each method utilizes equal data, we compare the performance of Target-Only using T0 and T0Mix (56.44) against the proposed FLAD methods and baselines using T5 and T0Mix (left column). From this comparison it becomes clear that Loss-Scaling actually does worse than multitask training followed by direct fine-tuning by  $0.5 - 3.2\%$ . However, we do find that the remaining FLAD methods lead to improvements (between  $2.7 - 5.6\%$  absolute improvement). We find small performance
Figure 2: Comparison of state-of-the-art few-shot methods with FLAD methods trained on P3 using  $\mathcal{R}^{\mathrm{AGG}}$ . T-Few scores are from [56]. DEFT-Few scores are from [28]. GPT-3 scores are from [55] and utilize few-shot in-context learning. All models utilize the same number of few-shot examples and (other than GPT-3) have 3B parameters.
differences between EXP3-FLAD and UCB1-FLAD across the three reward functions. In general,  $\mathcal{R}^{AGG}$  leads to the best performance, but we perform a two-sided Wilcoxon rank-sum test to check for significance between average scores and find that the other rewards frequently have no significant difference  $(p > 0.05)$ .
The importance of prioritized sampling. Loss-Scaling was originally proposed for use with only a single auxiliary dataset and it was unclear, a priori, how it would cope with larger quantities. Additionally, Du et al. [6] purposefully choose an auxiliary dataset that is related to the target, while in our setting we make no such assumptions. We find that our methods outperform Loss-Scaling methods by  $6.3\%$  on average. In Figure 3 (and Figure 5 in Appendix E) we show that, over the course of training, the value of gradient alignments and gradient magnitude similarities for most datasets will converge to 0, leading to very small gradient updates for Loss-Scaling. More importantly, the auxiliary data that is relevant to the target task is seen less frequently for Loss-Scaling than our MAB-based methods. This can be seen by comparing the difference in performance of Loss-Scaling methods when using less (T0Mix) vs. more (P3) auxiliary data. We find that, at best, Loss-Scaling  $(GA)$  improves  $2\%$  when using T5 and, at worst, only  $0.2\%$  for Loss-Scaling  $(GMS)$  with T0. This is compared with the notable improvements of EXP3-FLAD and UCB1-FLAD of  $2.6 - 4\%$  when considering the same data increase from T0Mix to P3.
The importance of exploration and exploitation. Interestingly, we expected that Exploit-Only would outperform the Explore-Only method because it utilizes relational information between the target and auxiliary tasks, but find no statistical difference between the methods (two-sided Wilcoxon rank-sum test gives  $p > 0.05$ ). Furthermore, when comparing the ability to leverage additional auxiliary data (i.e. going from T0Mix to all of P3), we find that the improvement for Explore- and Exploit-Only methods is minimal with only  $0.7 - 2\%$  improvement. On the other hand, EXP3-FLAD and UCB1-FLAD show a notable improvement of  $2.6 - 4\%$ , emphasizing the importance of both exploration and exploitation, particularly when dealing with large collections of auxiliary data.
FLAD provides improved generalization over non-FLAD methods. Next, we compare the performance of our best models trained on P3 using  $\mathcal{R}^{AGG}$  with state-of-the-art few-shot methods: T-Few, DEFT-Few, and GPT-3. T-Few [56] is a variant of the T0-3B model that multi-task pre-trains parameter-efficient (IA) $^3$  modules followed by target-only fine-tuning of the (IA) $^3$  modules. DEFT-Few [28] is a variant of the T5-XL model that uses retrieved auxiliary data for multi-task training. It first trains a T5-XL model on the 500 nearest neighbor samples from P3 using 1000 unlabeled target dataset samples, and then performs few-shot target-only fine-tuning with the (IA) $^3$  modules from Liu et al. [56]. Finally, we also compare against the 175 billion parameter variant of
Figure 3: Reward distributions of  $R^{GA}$  and  $R^{GMS}$  prior to training (step 0) and after 300 gradient updates for the T5-XL model with T0Mix as the auxiliary dataset and WSC [52] as the target dataset. Each quadrant shows the histograms of reward distributions for all 35 auxiliary datasets. By step 300 most auxiliary datasets provide 0 reward, while only the few remaining "beneficial" datasets provide positive rewards. Results from every 100 gradient updates are shown in Figure 5 in Appendix E.
GPT-3 [55], which utilizes in-context learning. We find that, on average, models trained using our FLAD-based methods outperform all other methods and, to the best of our knowledge, our methods lead to the first 3 billion parameter model that outperforms GPT-3 on this dataset mixture (previous smallest models have 11 billion parameters), despite using 62.5 times fewer parameters than GPT-3. Additionally, we find that our FLAD-based methods provide robust performance across datasets, achieving the best or second-best performance on  $8/11$  datasets, and never performing worst. The use of task-specific modules lead T-Few and DEFT-Few to significant improvements over target-only fine-tuning, preventing the models from ending up in poor local minima. However, these results demonstrate that with the same data, simultaneously fine-tuning with auxiliary and target data leads to improved few-shot generalization, providing a complementary means of improving performance.
Investigating the Reward-Generating Processes. In Section 3.2, we mention that due to the highly non-convex loss landscape and the use of stochastic gradient descent-based optimization techniques, we cannot ensure that our reward generating process is stationary, independent across auxiliary datasets, or follows a normal distribution. To gain a deeper understanding of our reward-generating processes, we examine the distribution of each reward using 5,000 samples from all 35 auxiliary datasets of T0Mix and 32 samples from a few-shot target dataset, WSC [52]. The resulting histograms at every 100 steps can be found in Appendix E, and Figure 3 shows an abbreviated version. The left side of Figure 3 demonstrates that for  $\mathcal{R}^{GA}$ , almost every dataset yields a Gaussian reward distribution, with a few multi-modal distributions. Notably, WikiBio [59] (dark orange) exhibits peaks at 0.25 and -0.75. Interestingly,  $\mathcal{R}^{GA}$  results in polarized rewards across datasets, with minimal distribution density between -0.75 and 0.25. In contrast, the right side of Figure 3 displays more non-Gaussian distributions for  $\mathcal{R}^{GMS}$ , as well as flatter distributions compared to  $\mathcal{R}^{GA}$ . Remarkably, we observe that  $\mathcal{R}^{GA}$  produces more stationary reward distributions, as the distribution for almost every dataset (30/35) converges rapidly towards 0 after only 100 steps. Although most distributions for  $\mathcal{R}^{GMS}$  also converge towards 0, the convergence occurs at a slower pace, taking nearly 500 steps.
Probing the training dynamics. To better understand the training dynamics of our proposed methods, we perform a case study on T5-XL with T0Mix and  $\mathcal{R}^{GA}$  and find two datasets where either algorithm improves significantly over the other (full details and figures in Appendix F). First, we study RTE, where UCB1-FLAD outperforms EXP3-FLAD. We calculate the empirical distribution of samples seen from each auxiliary dataset and find that EXP3-FLAD samples nearly uniformly from all datasets while UCB1-FLAD forms a bimodal sampling distribution with peaks at  $2.5\%$  and  $3.25\%$  ( $30\%$  relative difference). The uniformity of the EXP3-FLAD distribution is counterintuitive, as we do find that it achieves separation between auxiliary tasks in the cumulative estimated reward (as shown in Figure 7), but this does not lead to separation in the sampling probability space. Additionally we find that even on COPA, where EXP3-FLAD outperforms UCB1-FLAD, EXP3-FLAD still achieves
good separation between cumulative estimated rewards, but has a unimodal sampling distribution, while UCB1-FLAD does not have as clear of a bimodal distribution as in RTE. The difference in empirical sampling distributions is likely due to the difference between the greedy policy of UCB1-FLAD and the stochastic policy of EXP3-FLAD. Empirically, we find that EXP3-FLAD very rarely assigns an auxiliary dataset a probability  $< 1\%$ , leading to many "bad" batches over the course of thousands of turns. On the other hand, the optimistic policy of UCB1-FLAD spends much less time exploring and will sample "bad" batches much less frequently.
The effect of scaling  $|\mathcal{A}|$  To assess the scalability of our proposed methods, we conduct a study by varying the size of  $|\mathcal{A}| \in \{35, 75, 125, 175, 225, 260\}$ . For each value of  $|\mathcal{A}|$ , we consistently include the 35 datasets from T0Mix and randomly select additional auxiliary datasets from P3 until we reach the desired  $|\mathcal{A}|$ . The study is performed on the same 11 target datasets as the main study, using the T0 base model and  $\mathcal{R}^{AGG}$  reward. The experiment is repeated with three random seeds. Figure 4 shows the mean across the 11 target datasets, along with the standard deviation between seeds.
We find that both EXP3-FLAD and UCB1-FLAD experience a sharp increase from  $|\mathcal{A}| = 35$  to 75. In addition, we observe improvements up to the maximum value of  $|\mathcal{A}| = 260$ , ultimately improving accuracy by 2.54 for EXP3-FLAD and 3.12 for UCB1-FLAD when transitioning from 35 to 75 datasets, with further increases of 1.54 for EXP3-FLAD and 0.47 for UCB1-FLAD when increasing  $|\mathcal{A}|$  from 75 to 260.
Figure 4: The effect of scaling  $|\mathcal{A}|$  on target task performance. Each line represents mean score across 11 datasets and three random seeds, with shaded regions falling between one standard deviation of the mean.
Discussion on reward functions. In FLAD we want to prioritize training on auxiliary datasets with similar solution spaces as the target task without overfitting to the few-shot target data, and our reward functions are designed to serve this goal. To better understand the reward signal of our aggregate reward,  $\mathcal{R}^{AGG}$ , we examine four combinations of rewards: low  $\mathcal{R}^{GA}$  and  $\mathcal{R}^{GMS}$ , high  $\mathcal{R}^{GA}$  but low  $\mathcal{R}^{GMS}$ , low  $\mathcal{R}^{GA}$  but high  $\mathcal{R}^{GMS}$ , and high  $\mathcal{R}^{GA}$  and  $\mathcal{R}^{GMS}$ . When both rewards are high, we can assume that the auxiliary gradient is useful. However, when one reward is high and the other is low, it is difficult to draw conclusions as a high  $\mathcal{R}^{GA}$  on its own means the auxiliary gradient will update weights in the right direction, but low  $\mathcal{R}^{GMS}$  can mean that we significantly overshoot or undershoot the target, where overshooting can be much more detrimental than undershooting. If both  $\mathcal{R}^{GA}$  and  $\mathcal{R}^{GMS}$  are small, we know the auxiliary gradient leads us away from the target solution space, but we don't know if its magnitude is much larger or smaller than the target. At the beginning of training, we can't know if the target or auxiliary gradient has larger magnitude, but as training progresses, it becomes significantly more likely that the auxiliary gradient is greater than the target. Thus, when both  $\mathcal{R}^{GA}$  and  $\mathcal{R}^{GMS}$  are low, we are likely to be pulled far from our current solution.
This work uses training set-based rewards, but validation set-based rewards are also possible. One downside of validation-based rewards is they calculate validation score frequently, which increases computational complexity. Additionally, we focus on the few-shot setting and use validation-based early stopping. If we use a validation-based reward, then to prevent overfitting we will need to further split the data into 3 partitions: train, early-stopping validation, and reward-validation.
Choice of baselines. With respect to the number of auxiliary datasets  $|\mathcal{A}|$  and target datasets  $|\mathcal{T}|$ , our methods and the baselines we compare against have a computational complexity of  $O(|\mathcal{T}|)$ , independent of  $|\mathcal{A}|$ . For our model and these baselines, the models we train require  $\sim 6$  GPU-hours per target dataset. If we were to consider a baseline whose computation grows linearly w.r.t.  $|\mathcal{A}|$ ,
$O(|\mathcal{A}||\mathcal{T}|)$  (e.g. [18]), these experiments would not be feasible without a large amount of hardware: Training such a model with T0Mix would take over 200 GPU-hours (over 8 GPU-days) for a single target dataset, and over 1500 GPU-hours (over 2 GPU-months) when using all of P3.
Why we don't include theoretical guarantees. The design of MAB algorithms generally comes with theoretical proofs of regret bounds, but in this work we omit such analysis. Although we could make guarantees on the regret bounds of our algorithms, they would not be meaningful because they would be with respect to the rewards, not the accuracy on a held-out dataset (which is the quantity we actually care about).
How does FLAD relate to few-shot learning and multitask learning? Both few-shot learning and FLAD are concerned with optimizing model performance on a single target task with a limited number of examples from the target task. In few-shot learning, the model is given only the target task data  $\mathcal{D}_{\mathcal{T}}$  and there is no auxiliary data. Effectively,  $\mathcal{D}_{\mathcal{A}}$  is the empty set for few-shot learning. In contrast, for the FLAD setting  $|\mathcal{D}_{\mathcal{A}}| > 1$ . Based on the findings from this study, we highly recommend that practitioners utilize auxiliary data when it is available.
Multitask learning is concerned with optimizing a model for performance on multiple target datasets simultaneously. This is in direct opposition with the FLAD methods presented here, which aim to optimize a model for a single target task. However, it is possible to extend our MAB-based methods to optimize for multiple target tasks simultaneously by aggregating multiple rewards. We believe this would make for an interesting future study.
Limitations. One of the implicit assumptions in the FLAD setting (made by this work and all prior works) is that there is at least some auxiliary data that will be useful for the target task. However, one of the main distinctions of our methods from prior works in the FLAD setting is that prior works make a strong assumption that all auxiliary data are useful, and thus appropriate auxiliary datasets must be hand-picked by humans. On the other hand, our methods allow for only a small portion of the auxiliary data to be useful – our proposed algorithm explores to find useful auxiliary datasets and then exploits them.
Recall the desiderata for our algorithm, expressed in the introduction: our algorithm should (1) make no assumptions on the available auxiliary data a-priori, (2) scale well with the number of auxiliary datasets, and (3) add minimal memory and computational overhead. (1) When designing our algorithm, we purposefully formulate the problem as a multi-armed bandit. MAB algorithms, in general, make no assumptions on the quality of rewards and, in particular, EXP3 even assumes that the auxiliary datasets will play an adversarial role when returning rewards. (2) As previously mentioned, our algorithms have a single-turn computational complexity that is independent of the number of auxiliary datasets. (3) Finally, our method adds minimal computational overhead beyond usual training computations. Every gradient that we utilize for our reward functions are also used to update the model, adding no additional computations. The only computational overhead is to compute gradient alignment (three vector dot products, two scalar square roots, and two scalar multiplications) or magnitude similarity (four vector dot products, two scalar square roots, three scalar multiplications, and one scalar addition). Additionally, our method adds a small amount of memory overhead, used to store gradients between model updates. Our rewards consider only the gradient w.r.t the language modelling head and, in practice, require 0.25Gb per auxiliary gradient to store, slightly increasing the space complexity above standard fine-tuning.
The methods proposed in this work demonstrate the effectiveness of simultaneous training on auxiliary and target datasets in few-shot settings, continuously updating beliefs by exploring and exploiting auxiliary data, and framing FLAD as a MAB problem. We further showed that by satisfying our desiderata, we are able to scale up FLAD to hundreds of auxiliary datasets and outperform traditional few-shot fine-tuning and in-context learning methods. While the presented algorithms satisfy our desiderata, the findings from this study can inform future work to further improve upon these methods in a number of ways, such as improving the reward function and reducing the space complexity.
This material is based on work that is partially funded by an unrestricted gift from Google. This work was supported by the National Science Foundation award #2048122. The views expressed are those of the authors and do not reflect the official policy or position of the US government. Additionally, we thank the Robert N. Noyce Trust for their generous gift to the University of California via the Noyce Initiative.
We would also like to thank Dheeraj Baby for insightful advising on multi-armed bandit settings as well as Lucio Dery and Jiachen Li for invaluable feedback on early drafts of this work.
[1] Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. In International Conference on Learning Representations, 2017.
[2] Yaqing Wang, Quanming Yao, James T. Kwok, and Lionel M. Ni. Generalizing from a few examples: A survey on few-shot learning. ACM Comput. Surv., 53(3), jun 2020.
[3] Archit Parnami and Minwoo Lee. Learning from few examples: A summary of approaches to few-shot learning, 2022.
[4] Pengcheng Wu and Thomas G. Dietterich. Improving svm accuracy by training on auxiliary data sources. In Proceedings of the Twenty-First International Conference on Machine Learning, ICML '04, page 110, New York, NY, USA, 2004. Association for Computing Machinery.
[5] Reza Esfandiarpoor, Amy Pu, Mohsen Hajabdollahi, and Stephen H. Bach. Extended few-shot learning: Exploiting existing resources for novel tasks, 2020.
[6] Yunshu Du, Wojciech M. Czarnecki, Siddhant M. Jayakumar, Mehrdad Farajtabar, Razvan Pascanu, and Balaji Lakshminarayanan. Adapting auxiliary losses using gradient similarity, 2020.
[7] Sam Verboven, Muhammad Hafeez Chaudhary, Jeroen Berrevoets, Vincent Ginis, and Wouter Verbeke. Hydalearn. Applied Intelligence, Jul 2022.
[8] Alon Albalak, Yi-Lin Tuan, Pegah Jandaghi, Connor Pryor, Luke Yoffe, Deepak Ramachandran, Lise Getoor, Jay Pujara, and William Yang Wang. FETA: A benchmark for few-sample task transfer in open-domain dialogue. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 10936-10953, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
[9] Tu Vu, Tong Wang, Tsendvarsen Munkhdalai, Alessandro Sordoni, Adam Trischler, Andrew Mattarella-Micke, Subhransu Maji, and Mohit Iyyer. Exploring and predicting transferability across NLP tasks. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7882-7926, Online, November 2020. Association for Computational Linguistics.
[10] Yada Pruksachatkun, Jason Phang, Haokun Liu, Phu Mon Htut, Xiaoyi Zhang, Richard Yuanzhe Pang, Clara Vania, Katharina Kann, and Samuel R. Bowman. Intermediate-task transfer learning with pretrained language models: When and why does it work? In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 5231-5247, Online, July 2020. Association for Computational Linguistics.
[11] Sebastian Thrun and Lorien Pratt, editors. Learning to Learn. Kluwer Academic Publishers, USA, 1998.
[12] Trapit Bansal, Rishikesh Jha, Tsendsuren Munkhdalai, and Andrew McCallum. Self-supervised meta-learning for few-shot natural language classification tasks. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 522-534, Online, November 2020. Association for Computational Linguistics.