Improving Few-Shot Generalization by Exploring and Exploiting Auxiliary Data
Authors: Alon Albalak
Colin Raffel
William Yang Wang
Venue: NeurIPS 2023 poster (2023)
Keywords: Few-shot learning
natural language processing
few shot learning
NLP
multi-armed bandit
multi armed bandit
TL;DR: This work focuses on few-shot learning with auxiliary data to improve model generalization, presenting two methods inspired by multi-armed bandits.

Abstract: 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 100x 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.

Improving Few-Shot Generalization by Exploring and Exploiting Auxiliary Data
Alon Albalak University of California, Santa Barbara al alon_ aonaalaalakuucsh edu
Colin Raffel University of Toronto Vector Institute craffel Serffel@mmail.com
William Yang Wang University of California, Santa Barbara willi am@cs ucsb edu
Abstract
Few-shot learning is valuable in many real-world applications, but learning generalizable model without overfitting to the few labeled datapoints is challenging. In this work, focus Few-shot Learning with Auxiliary Data (FLAD), 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 FL.D 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 100x more auxiliary datasets than prior methods. We propose two algorithms EXP3- FLAD and UCBI-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 pre-existing FL.D methods by 4% and lead to the first 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 viable path towards substantially improving generalization in few-shot learning. All our code available at wiigghb.coo/allooo http://loooallaaa/AAADDD
Introduction
Few-shot learning is an attractive learning setting for many reasons: 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 also challenging setting that requires delicate balance between learning the structure of the feature and label spaces while preventing overfitting to the limited training samples [T 3. One approach to improving the gencralizability 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 7.
However, FLAD methods introduce their own challenges, including increased algorithmic and computational complexity. Specifically, incorporating auxiliary data during training introduces 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 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 X Delegating such choices to algorithm can lead to better solutions, demonstrated in the transfer learning {8] 10] meta-lcarning [1] multi-task learning [13 [4 [6), and auxiliary learning literature 44 [7]
However, prior auxiliary learning algorithms often assume that only 1-3 related auxiliary datasets
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Learner
Environment
Select dalaset DD according to
Policy
DA
2. Sample batch
D
D
(x,y) ~ D.
Model
5. Update pollcy according Rey
Calculate Reward
V
fg
Rat
Calculate auxilarr and larget gradients, updata model
Dr
Figure 1: Overview of lew-shot learning with auxiliary data (FLAD) as a multi-armed bandit problem. On the left is the learner which defines a policy that determines which auxiliary dataset to sample from. On the right is the environment that includes the set of auxiliary datasets DA. target dataset Dy. and the model fe. At each turn the following five steps take place, further described in Section [.1]] 1. The learner selects an auxiliary dataset Da according to its policy я. 2. The environment samples batch yx.y) Da 3. The model fe calculates gradients for the sampled batch (V and the target dataset (VT), then updates the parameters 0. 4. A reward Ra, is calculated based V. and V 5. The learner updates based on R.,,-
available and design algorithms whose computational complexity grows linearly (or worse) with the number auxiliary datasets [8] 8) motivating the search for more efficient methods as the number of auxiliary datasets grows.
To overcome the challenges prior works, we desire a FLAD algorithm that (1) makes no assump- tions 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-exploitatoo trade-off 19 200 We relate the set of auxiliary datasets to the arms of MAB and tailor the classic EXP3 [11] and UCB1 [22]
algorithms to fit the FLAD framework by designing three efficient gradient-based reward signals. The combination of our LAB-based algorithms and efficient gradient-based rewards allows to scale to 100x more auxiliary datasets than previous methods. Figure provides a basic illustration of how we formulate FLAD MAB problem.
To empirically validate our approaches, we focus on few-shot training of language models and utilize P3 [23] readily available resource with hundreds of auxiliary language datasets. We evaluate our methods on the same held-out tasks as the TO language model [16) and show that, when using the same collection of auxiliary datasets, our algorithms outperform directly fine-tuned TO by 5.6% (EXP3-FLAD) and 5.7% (UCBI-FI --LLDD absolute. Furthermore, incorporating all available datasets in P3 (i.e. not just those used to train TO) increases the improvement to 9, and 9.2%. Finally, compare models trained with our methods against state-of-the-art few-shot methods, finding that methods improve performance by >3%, unlabeled target dataset samples. Furthermore, to the best of our knowledge, our methods lead the first 3 billion parameter model that improves 175B GPT-3 using few-shot in in-context learning.
though one model utilizes large collection of
In summary, our main contributions are:
We connect FLAD to the MAB setting and focus on the exploration-exploiae trade-off by designing two algorithms, EXP3-FLAD and UCBI-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 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.
2
Related work
A long history of works have found success when combining auxiliary data with target data [
24 25 26
28 8). Some works have explored the addition auxiliary learning
objectives to aid the learning of the target task [24 26 25 5 [7]. More similar to work are

---

methods that perform auxiliary learning by introducing additional data sources beyond the target 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 al. [8] and Verboven et 7 assume access to 10,000 laheled target samples, Ivison et
data 4 6 18
al. [28) and Lin et al.
[27] assume access to 1, 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 [ I8 7 8 For example, Verboven et al. 7 propose a task-weigtting method that requires solving system of equations that becomes underspecified with multiple auxiliary tasks, limiting their method to only a single auxiliary task. Furthermore, Chen et al. [8] experiment with 3 auxiliary tasks because their method requires learning a target-aware classifier for each source task, so the computation scales as ((IA[]T) where |A] is the number of auxiliary tasks and [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 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-armes bandits (MAB) problem, shown in Figure [] 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 UCBI-FLAD.
3.1 Setup
FLAD problem setting. Few-shot learning with auxiliary data (FLAD) fits into the following setting: assume access to a large set of auxiliary datasets DA where, for all E A, D, is an individual auxiliary dataset. Given small quantity of data belonging a target dataset Dy. the goal of FLAD is to find parameters 0 of a model fo that achieve high performance on the unknown distribution underlying Df while utilizing only the available data, Dy UDA.
Formulating FLAD as MAB. In this work, we adopt the multi-armes bandit (MAB) setting by formulating FLAD as a Markov decision process (3) and defining learner and environment, illustrated in Figure [ The learner consists of a policy π defining a selection strategy over all Da E DA. The environment consists of the target dataset 'D,, auxiliary datasets DA and model fg. In this formulation the learner interacts with the environment over N rounds. At each round the learner selects one of the environment's 's |A| datasets D, E DA. Next, the environment samples a batch (x. ~ Da and calculates the gradient w.r.t. using a task-appropriate loss function as V. VoL(fo, y). Then, the environment computes the target gradient V, VeC(fo, Dr), and updates model parameters w.r.t. V + Va- Finally, the learner uses a gradient- based reward Ra, (... V to update its policy æ. See Appendix Aland and Lattimore & Szepesvári [34] for further details on multi-armed bandits.
Designing the reward functions. We design the reward function R with our desiderata in mind. To ensure that our algorithm adds minimal memory and computational overhead consider rewards that utilize information intrinsic to the model and the losses being optimized, not an external model or metric (e.g. accuracy BLEU). In this work we propose three gradient-based reward functions inspired by previous works: gradient alignment (4]4][[5) gradient magnitude similarity [36] 37], and their aggregation. Formally, at turn t let V. be the gradient of the auxiliary batch and V be the target dataset gradient. Gradient alignment is defined as R.G.A.
V 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 RCMS 2/| >
V..VT
|V 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
3

---

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 RGA E [o, 1], the same range
as RGMS and define the aggregate to be their sum: RAGG 1+R
+ RGMS, We provide further
2
discussion on the design of reward functions in Section
3.2 Adapting the EXP3 algorithm.
EXP3 Background We base our first algorithm, EXP3-FLAD, on the EXP3 algorithm 21
("Exponential-weighe algorithm for Exploration and Exploitation"), EXP3 targets the adversarial MAB setting, which assumes that the geward-generating process is controlled by an adversary who is given access to the learner's s policy и and determines the sequence of rewards, (Ra.)) for each arm prior to play 38 We consider the adversarial MAB formulation due 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 Jnnd in 221.
In EXP3-FLAD, the learner selects arms according to a Gibbs distribution based on the empirically determined importance -peertance-weighted weighted rewards of arms [9]. To allow for exploration, we mix the Gibbs distribution with uniform distribution [21] Formally, let & be the exploration rate at turn t and, recalling that K |A| is the number of auxiliary datasets, then иT defines the probability of selecting a given arm E A as the linear combination of Gibbs and uniform distributions
(a) = =11 -E E.' exp(E,-1R1,))
exp(Er_1R.)) 1hk) + & where Ra, is the importance weighted reward Ro, =
Ra,t-1 + xi i(a) We want the learner to explore more in early training than later stages, so we
Ru.
use a decaying exploration rate & = min K: K- as proposed by Seldin et al. 33 The use of an supprraacc-weighhee 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 &, The exploration of EXP3-FLAD combined with amportance-weighting allows the policy to handle non-stationary reward wwwrr-gennerating processes.
3.3 Adapting the UCB1 algorithm.
UCB1 background. While EXP3-FLAD applicable in unconstrained settings with highly stochas- tic and statstattonnry rewards, it can be outperformed by other algorithms in settings that are con- strained. One such algorithm is the upper confidence bound (UCB1) algorithm [22], which was originally designed to be optimal for stationary, normally distributed reward functions. Neverthe- less, variants of UCB1 have been demonstrated to be effective a range of settings, such as those involving non-stationary, Gau-Gausiinn, or heavy-tailed distributions 40) 41). The UCB1 1 algorithm and its variants assign each arm a value called the upper confidence bound based on Hoeffding's inequality [4] 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 UCBI-FLAD, the learner greedily selects arms according to their upper confidence bound. UCB1 is originally designed for stationary reward- genur--enertiing processes, so to accommodate non-stationarity we include an exponential moving average when estimating the mean reward for a given arm. Formally, let RA,: be the observed reward for arm at turn i, then we calculate the estimated mean reward as Rr (1 fR. + 3Ra., where 3 is the smoothing factor. Then, we define the upper confidence bound be UCBa, R. + / /2lmt In the original MAB setting all interactions with na
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 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 V. VoL(Jo, x, y), where the number of samples in (x, can be significantly larger than mini-batch, and can be up to the size of the full dataset. practice, we use 1,000 examples which is computed in 2 minutes single GPU,

---

Algorithms The EXP3-FLAD and UCBI-FLAD algorithms are visualized in Figure I At each turn, both methods will first select an auxiliary dataset D,, EXP3-FLAD first computes the current exploration rate & and samples Da according to the distribution defined by Ty (A), while UCB1- FLAD greedily selects DD-- corresponding to the arm with largest upper confidence bound, a" arg m a.t Next, for both methods, the environment samples batch from the selected dataset, {x.y} ~ Da. and calculates the gradient V. = VeL(fe, x, 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 Vo DT) and updates the model w.r.t. V E. Va. Finally, EXP3-FLAD calculates the sepportance-weightee reward for each auxiliary batch using the observed rewards, while UCB1-FLAD calculates the smoothed estimated mean reward. Pseudocode is found in Appendix B
4 Experimental setup
Models. For our experiments, we utilize encoder-decoder Transformer models from the T5 family of pre-trained language models [43). Specifically, experiment with L.M-ddapeed T5 (T5-LM) and TO. The T5-LM model further trains the T5.1.1 model for 100,000 steps (corresponding to 100B tokens) from the C4 dataset
(3] the prefix language modeling objective + 'The TO model was initialized from T5-LM and further trained on multitask mixture of prompted datasets as described by Sanh et al. [6]. We repeat each experiment with T5-LM XL (hereafter T5-XL) and TO-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
Target datasets. We obtain all datasets from Hugging Face Dataset: 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 TO. To evaluate our few-shot methods, we utilize the same held-out datasets as TO, which
cover four distinct tasks: sentence completion (COPA 46). HellaSwag 47). Story Cloze [8]] ).
natural language inference (ANLI [49] CB (50), RTE 511 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 and validation partitions for true few-shot learning (e.g. 10 train and 10 validation samples for HellaSwag). Only ANLI datasets have publicly available test set, for all other datasets evaluate models on the original validation set (not utilized for few-shot training or validation).
[55 [6). We further divide each split into equal training
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, use the collection of datasets used for multitask training of TO (henceforth referred to as TOMix), including 35 unique datasets covering question answering, sentiment analysis, topic classification, summarization, para- phrase detection and structure-to-texx
Second, we utilize all datasets in P3 [23] (which forms a superset of TOMix) and prevent data leakage by filtering out datasets that overlap with any target dataset, leading to 260 available datasets (list in Appendix(HF) (H). For each auxiliary dataset, we use at most 10,000 of the dataset's is 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 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 equivalent 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 UCB 1), and multitask-trains the model by mixing auxiliary datasets according to 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 oxploit-only mix target and auxiliary data with a ratio of AM times the highest auxiliary sampling probability. For instance, explore-only with M = 5 and DA TOMix has a 1/35 probability to sample auxiliary dataset DD E D 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
'anps//hugginnftaee.oodass coldttttsss
5

---

BASE MODEL
T5-XL
T0-3B
Training Method
Auxiliary Data
TOMiX
P3
TOMix
P3
Target-Onle
52.82
56.44
Loss-Scaling (G.)
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 (R G4)
61.50
64.07
62.87
65.98
UCBI-FLAD (RGA)
62.01
65.52
62.35
66.29
EXP3-FLAD (RGMS)
61.72
65.57
62.78
65.51
UCBI-FLAD (RGMS)
61.67
65.21
62.85
66.00
EXP3-FLAD (R AGG)
62.05
65.47
62.84
66.84
UCBI-FLAD (R. (RAGG)
62.08
65.63
62.93
66.29
Table 1: Main results. Each cell contains the score of training 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 (G))), but we also propose version that scales losses by gradient magnitude similarity (Loss-Scaling (GMS)).
Training details. For the target-only baseline, we use learning rates in {le-4, 3e-4J. For all other methods, we always use a learning rate of le-4. For target-, explore-, and exploit-only baselines we use batch sizes in {32, 128}. For loss scalinc. 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 target dataset mixing ratio of M {1, 5, 10). For all experiments we the Adafactor optimizer (58) and validation-based early stopping for model checkpoint selection. In preliminary experiments
and UCB1-FLAD we use mini-batches
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 I-FL-D found the smoothing factor 3 0.9 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 model that is specialized for each target dataset. For each proposed method and baseline, we train and evaluate model on each of the 11 target datasets. We repeat training and evaluation on random seeds and include the aggregated results in Table [ Each cell shows the accuracy averaged across all 55 (11 target datasets, 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 [ (w 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 UCB I-LLAD 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 TO and TOMix (56.44) against the proposed FLAD methods and baselines using T5 and TOMix (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 find that the remaining FLAD methods lead to improvements (between 2.7-5.6% absolute improvement). We find small performance

---

Natural Language Inference
Average Accuracy
ANUI-R
ANU-R2
ANLIR3
CB
RTE
70
30
45.0
90
40.0
42.5
66
37.5
40,9
64
40
70
3255
35
60
30.0
35.0
Coference Resolution
Sentence Completion
Word Sense
WSC
Winegrande
COPA
StoryCloze
HellaSwag
WIC
100
70
55
60
50
60
45
50
70
40
EXP3-TLAD (15)
EXPS-FLAD (T()
UCB1-FLAD (T5]
UCB1-FLAD (TO)
1-few
DEFT-Few
GPT-3 (1758)
Figure 2: Comparison of state-of-the-art few-shot methods with FLAD methods trained on P3 using RAGG 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 UCBI-FLAD across the three reward functions. In general, R.GGG 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 Loss-Scaling originally proposed for with only a single auxiliary dataset and was unclear, a priori, how it would cope with larger quantities. Additionally, Du et al. (6] purposefully choose 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 ( Figure 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 O, leading to very small gradient updates for Loss- Loss-Scaling. More importantly, the auxiliary data that is relevant to the targel 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 (TOMix) 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 TO. This is compared with the notable improvements of EXP3-FLAD and UCBI-FLAD of 2.6-4% when considering the same data increase from TOMix 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 TOMix 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 UCBI-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 FOA-FLAD methods. Next, we compare the performance of our best models trained on P3 using R AGG with sfate-of-the-art few-shot methods: T-Few, DEFT-Few, and GPT-3. T-Few (56) is a variant of the TO-3B model that multi-task pre- trains sarameter-efficien (IA)3 modules followed by target-only fine-tuning of the (IA)² modules. DEFT-Few (2 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)² modules from Liu et al. (56). Finally, also compare against the 175 billion parameter variant of
7

---

Gradient Alignment R
Gradient Magnitude Similarity RCMS
Step
44
Slep 300
6/0 0.5 175 100
Figure 3: Reward distributions of R and RGMS prior to training (step 0) and after 300 gradient updates for the T5-XL model with TOMix the auxiliary dataset and WSC [52] 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 3lin in Appendix H
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, the best of our knowledge, methods lead to the first 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 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 complementary means of improving performance.
Investigating the Sewward-Genertting Processes. In Section 3.2 we mention that due to the highly non-convex loss landscape and the of stochastic gradient descent-based optimization techniques, we cannot ensure that our reward generating process is stationary, independent across auxiliary datasets, follows a normal distribution. To gain deeper understanding of our reward-generatin processes, we examine the distribution of each reward using 5,000 samples from all 35 auxiliary datasets of TOMix and 32 samples from 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 RE almost every dataset yields Gaussian reward distribution, with a Tew multi-modal distributions. Notably, WikiBio [59] (dark orange) exhibits peaks 0.25 and -0.75. Interestingly. RGA 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- non-Gaussian distributions for RGMS well as flatter distributions compared to we observe that R RGA produces more stationary reward distributions, as the distribution for almost every dataset (30/35) converges rapidly towards after only 100 steps. Although most distributions for RGSS also converge towards O, the convergence occurs at a slower pace, taking nearly 500 steps.
Remarkably,
Probing the training dynamics. To better understand the training dynamics of our proposed methods, we perform a case study on T5-XL with TOMix and 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-FB -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 UCB (CBI-FIU -LLAD forms a bimodal sampling distribution with peaks at 2.5% and 3.25% (30% relative difference). The uniformity of the EXP3-FLAD distribution 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 UCBI-FLAR EXP3-FLAD still achieves
8

---

good separation between cumulative estimated rewards, but has unimodal sampling distribution, while UCBI-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 UCBI- FLAD and the stochastic policy of EXP3-FLAD. Empirically, we find that EXP3-FLAD very rarely assigns an auxiliary dataset probability < 1%, leading to many "bad" batches over the course of thousands of turns. On the other hand, the optimistic policy of UCBI I-FLAD spends much less time exploring and will sample "bad" batches much less frequently.
The effect of scaling |A] To assess the scal- ability of our proposed methods, we con- duct study by varying the size of |A| E {35, 75, 125, 175, 225, 260). For each value of [A], we consistently include the 35 datasets from TOMix and randomly select additional auxiliary datasets from P3 until we reach the desired [A]. The study is performed on the same 11 target datasets as the main study, using the TO base model and RAGG reward. The experiment is re- peated with three random seeds. Figure shows the mean across the 11 target datasets, along with the standard deviation between seeds.
5.66
Accuracy
0.64
6,63
EXP3-FLAD UCB1-FLAD
We find that both EXP3-FLAD and UCB1- FLAD experience sharp increase from |A| = 35 to 75. In addition, we observe improvements up to the maximum value of |A| = 260, ulti- mately improving accuracy by 2.54 for EXP3- FLAD and 3.12 for UCBI-FLAD when tran- sitioning from 35 to 75 datasets, with further increases 1.54 for EXP3-FLAD and 0.47 for UCB1-FLAD when increasing |A| from 75 to 260.
100 # of Auxillary Datasets
200
Figure 4: The effect of scaling |A| on target task performance. Each line represents mean score across 11 datasets and three random seeds, with shaded regions falling between one standard devia- tion of the mean.
6 Discussion
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, RAGG we examine four combinations of rewards: low RGA and R RCMS, high RGA but low RGMS, low RGA but high RGMS and high RGA and RGMS. When both rewards are high, we can assume that the auxiliary gradient useful. However, when one reward is high and the other is low, it is difficult to draw conclusions as high R on its own means the auxiliary gradient will update weights in the right direction, but low R RCMS can mean that we significantly overshoot or undershoot the target, where overshooting can be much more detrimental than undershooting. If both R.G and R RGMS 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 '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 R&A and R RGMS are low, we are likely to be pulled far from our current solution.
This work uses training bet-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-batt based early stopping. If we use a validation- reward, then to prevent overfitting we will need to further split the data into 3 partitions: train, early-stopping validation, and veward-validation.
Choice of baselines. With respect to the number of auxiliary datasets |A| and target datasets [T], our methods and the baselines we compare against have computational complexity of O((T), independent of [A]. For our model and these baselines, the models we train require ~ 6 GPU-hours per target dataset. If we were to consider baseline whose computation grows linearly w.r.t [A],
9

---

O(IAIT) (e.g. (8)), these experiments would not be feasible without large amount of hardware: Training such model with TOMix 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 single target task with limited number of examples from the target task. In few-shot learning, the model is given only the target task data Dy and there is no auxiliary data. Effectively, DA is the empty set for few-shot learning. In contrast, for the FLAD setting [DAl > 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 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, 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.
7 Conclusion
Recall the desiderata for our algorithm, expressed in the introduction: our algorithm should (1) make no assumptions on the available auxíliary 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 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 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 fhot 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 number of ways, such improving the reward function and reducing the space complexity.
10

---

Acknowledgements
This material is based on work that 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 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.
References
[1] Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. In Interna- tional Conference Learning Representations, 2017.
[2] Yaqing Wang, Quanming Yao, James T. Kwok, and Lionel M. Ni. Generalizing from a few examples: A survey 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, Wojeiech M. Czamecki, Siddhant M. Jayakumar, Mehrdad Farajtabar, Razvan Pascanu, and Balaji Lakshminarayanan. Adapting auxiliary losses using gradient similarity, 2020.
[7] Sam Verhoven, 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, Tsendsuren Munkhdalai, Alessandro Sordoni, Adam Trischler, Andrew Mattarella-Micke, Subhransu Maji, and Mohit lyyer. 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 Natural Language Processing (EMNLP)), pages 522-534, Online, November 2020. Association for Computational Linguistics.

---

[13] Sen Wu, Hongyang R. Zhang, and Christopher Ré. Understanding and improving information transfer in multi-task learning. In International Conference on Learning Representations, 2020.
[14] Armen Aghajanyan, Anchit Gupta, Akshat Shrivastava, Xilun Chen, Luke Zettlemoee, and Sonal Gupta. Muppet: Massive multi-task representations with pre-finetuning. In Proceedings of the 202/ Conference on Empirical Methods in Natural Language Processing, pages 5799-581 I, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
[15] Vamsi Aribandi, Yi Tay, Tal Schuster, Jinfeng Rao, Huaixiu Steven Zheng, Sanket Vaibhav Mehta, Honglei Zhuang. Vinh Q. Tran, Dara Bahri, Jianmo Ni, Jai Gupta, Kai Hui, Sebastian Ruder, and Donald Metzler. Ext5: Towards extreme multi-task scaling for transfer learning. In International Conference on Learning Reprexentations, 2022.
[16] Victor Sanh, Albert Webson, Colin Raffel, Stephen Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Arun Raja, Manan Dey, M Saiful Bari, Canwen Xu, Urmish Thakker, Shanya Sharma Sharma, Eliza Szezechla, Taewoon Kim, Gunjan Chhablani, Nihal Nayak, Debajyoti Datta, Jonathan Chang, Mike Tian-Jian Jiang, Han Wang. Matteo Manica, Sheng Shen, Zheng Xin Yong. Harshit Pandey, Rachel Bawden, Thomas Wang, Trishala Neeraj, Jos Rozen, Abheesht Sharma, Andrea Santilli, Thibault Fevry, Jason Alan Fries, Ryan Teehan, Teven Le Scao, Stella Biderman, Leo Gao, Thomas Wolf, and Alexander M Rush. Multitask prompted training enables zero-shot task generalization. In International Conference on Learning Representations, 2022.
[17] Lucio M. Dery, Paul Michel, Mikhail Khodak, Graham Neubig, and Ameet Talwalkar. AANG: Automating auxiliary learning. In The Eleventh International Conference on Learning Repre- sentations, 2023.
[18] Shuxiao Chen, Koby Crammer, Hangfeng He, Dan Roth, and Weijie J Su. Weighted training for cross-task learning. In International Conference on Learning Representations, 2022.
[19] William G Macready and David H Wolpert. Bandit problems and the exploration/exploitttio tradeoff. IEEE Trunsuctions on evolutionary computation, 2(1):2-22, 1998.
[20] Alex Simpkins, Raymond De Callafon, and Emanuel Todorov. Optimal trade-off between exploration and exploitation. 2008 American Control Conference, pages 33-38. IEEE, 2008
[21] Peter Auer, Nicolo Cesa-Bianchi, Yoav Freund, and Robert E Schapire. The nonstochastic multiarmed bandit problem. SIAM journal computing, 32(1):48-771 2002.
[22] Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2): 2002.
[23] Stephen Bach, Victor Sanh, Zheng Xin Yong, Albert Webson, Colin Raffel, Nihal V. Nayak, Abheesht Sharma, Taewoon Kim, M Saiful Bari, Thibault Fevry, Zaid Alyafeai, Manan Dey, Andrea Santilli, Zhiqing Sun, Srulik Ben-david, Canwen Xu, Gunjan Chhablani, Han Wang, Jason Fries, Maged Al-shaibani, Shanya Sharma, Urmish Thakker, Khalid Almubarak, Xiangru Tang, Dragomir Radev, Mike Tian-jian Jiang, and Alexander Rush. PromptSource: An integrated development environment and repository for natural language prompts. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics: System Demonstrations, pages 93-104, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[24] Xingyu Lin, Harjatin Baweja, George Kantor, and David Held. Adaptive auxiliary task weighting for reinforcement learning. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc. E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[25] Lucio M. Dery, Paul Michel, Ameet S. Talwalkar, and Graham Neubig. Should we be pre- training? an argument for end-task aware training as an alternative. ArXiv, abs/2109.07437, 2021.
[26] Aviv Navon, Idan Achituve, Haggai Maron, Gal Chechik, and Ethan Fetaya. Auxiliary learning by implicit differentiation. In International Conference on Learning Representations, 2021.
12

---

[27] Bill Yuchen Lin, Kangmin Tan, Chris Miller, Beiwen Tian, and Xiang Ren. Unsupervised cross-task generalization via retrieval augmentation. ArXiv, abs/2204.07937, 2022.
[28] Hamish Ivison, Noah A. Smith, Hannaneh Hajishirzi, and Pradeep Dasigi. Data-efficient finetuning using cross-task nearest neighbors, 2022.
[29] Sören Mindermann, Muhammed Razzak, Winnie Xu, Andreas Kirsch, Mrinank Sharma, Adrien Morisot, Aidan N. Gomez, Sebastian Farquhar, Janina Brauner, and Yarin Gal. Prioritized training on points that are learnable, worth learning, and not yet learned. International Conference Machine Learning. 2021.
[30] Shoaib Ahmed Siddiqui, Nitarshan Rajkumar, Tegan Maharaj, David Krueger, and Sara Hooker. Metadata archaeology: Unearthing data subsets by leveraging training dynamics. ArXiv, abs/2209.10015, 2022.
[31] Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari S. Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[32] Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S. Morcos. Semdedup: Data-efficient learning at web-scale through semantic deduplication, 2023.
[33] RICHARD BELLMAN. A markovian decision process. Journal of Mathematics and Mechanics, 6(5):79-688) 6844, 1957.
[34] Tor Lattimore and Csaba Szepesvári. Bandit Algorithmss. Cambridge University Press, 2020.
[35] Zirui Wang, Yulia Tsvetkov, Orhan Firat, and Yuan Cao. Gradient vaccine: Investigating and improving multi-task optimization in massively multilingual models. In International Conference on Learning Representations, 2021.
[36] Wufeng Xue, Lei Zhang, Xuanqin Mou, and Alan C. Bovik. Gradient magnitude similarity deviation: A highly efficient perceptual image quality index. IEEE Transactions Image Processing, 23(2):684-695, 2014.
[37] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. arXiv preprint arXiv:2001.06782, 2020.
[38] Peter Auer, Nicolo Desa-Bianchh, Yoav Freund, and Robert E Schapire. Gambling in rigged casino: The adversarial multi-armed bandit problem. In Proceedings of IEEE 36th annual foundations of computer science, pages 322-331. IEEE, 1995.
[39] Yevgeny Seldin, Csaba Szepesvári, Peter Auer, and Yasin Abbasi- Yadkori. Evaluation and analysis of the performance of the exp3 algorithm in stochastic environments. In Marc Peter Deisenroth, Csaba Szepesvári, and Jan Peters, editors, Proceedings of the Tenth European Workshop on Reinforcement Learning, volume 24 of Proceedings of Machine Learning Research, pages 103-116, Edinburgh, Scotland, 30 Jun-O1 Jul 2013. PMLR.
[40] Aurelien Garivier and Eric Moulines. On upper-confidence bound policies for switching bandit problems. In Jyrki Kivinen, Csaba Szepesvári, Esko Ukkonen, and Thomas Zeugmann, editors, Algorithmic Learning Theory, pages 174-188, Berlin, Heidelberg, 2011. Springer Berlin Heidelberg.
[41] Lai Wei and Vaibhav Srivastava. Nonstationary stochastic multiarmed bandits: Ucb policies and minimax regret, 2021.
[42] Wassily Hoeffding. Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association, (8(30)):13333) 1963.
[43] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine I Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res,, 21(1), jun 2020.
13

---

[44] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for -effcittteef prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Lan- guage Processing, pages 3045-3059, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
[45] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jemite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of-the-ar natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online, October 2020. Association for Computational Linguistics.
[46] Andrew Gordon, Zornitsa Kozareva, and Melissa Roemmele. SemEval-2012 task 7: Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In *SEM 2012: The First Joint Conference on Lexical and Computational Semantics Volume I: Proceedings of the main conference and the shared task, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation (SemEval 2012), pages 394-399 Montréal, Canada, 7-8 June 2012. Association for Computational Linguisties.
[47] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Associa- tion for Computational Linguistics, pages 17900, Florence, Italy, July 2019. Association for Computational Linguistics.
[48] Rishi Sharma, James Allen, Omid Bakhshandeh, and Nasrin Mostafazadeh. Tackling the story ending biases in the story cloze test. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 752-757, Melbourne, Australia, July 2018. Association for Computational Linguistics.
[49] Yixin Nie, Adina Williams, Emily Dinan, Mohit Bansal, Jason Weston, and Douwe Kiela. Adversarial NLI: A new benchmark for natural language understanding. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 1185-4900 Online, July 2020. Association for Computational Linguistics.
[50]
Marie-Cathering de Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In Sinn und Bedeutung, 2019.
[51] Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Joaquin Quinonero-Candela, Ido Dagan, Bernaroo Magnini, and Florence d' Alché Buc, editors, Machine Learning Challenges. Evaluating Predictive Uncertainty, Visual Object Classification and Recognising Tectual Entailment, pages 177-190, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg.
[52] Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Thirleenth international conference on the principles of knowledge representation and reasoning, 2012.
[53] Keisuke Sakaguchi, Ronan Bras, Chandra Bhagavatula, and Choi Yejin. Winogrande: An adversarial winograd schema challenge scale. Proceedings of the AAAI Conference on Artificial Intelligence, 34:8732-877 07473, 04 2020.
[54] Mohammad Taher Pilehvar and Jose Camacho Colladoo.. WiC: the word-in-context dataset for evaluating context- sensitive meaning representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 1267-1273, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[55] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott
14

---

Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[56] Haokun Liu, Derek Tam, Muqeeth Mohammed, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin Raffel. Few-shot purameter-efficient fine- muee-tning is better and cheaper than in-context learning. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[57] Ethan Perez, Douwe Kiela, and Kyunghyun Cho. True few-shot learning with language models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems; 2021.
[58] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Leurning Research, pages 4596-4604. PMLR, 10-15 Jul 2018.
[59] Rémi Lebret, David Grangier, and Michael Auli. Generating text from structured data with application to the biography domain. CoRR, abs/1603.07771, 2016.
15

---

A Multi-armed bandits
The Multi-Armed Bandit (MAB) setting is problem from machine learning where learner interacts with an environment over N rounds by following a policy т. At each round & the learner chooses one of the environment's K arms, reward R₁. Rewards for unplayed arms are not observed. The goal of the learner is to adopt a policy that selects actions that lead to the largest cumulative reward over N rounds, R > 1t=1 Rs. In this work we assume a finite K and that the underlying reward distribution of each arm may have a variety of properties (e.g. stochasticity or stationarity) depending on the exact scenario, leading to different optimal policies
A where K = |A,, after which the environment provides a
Adversarial MAB. The adversarial MAB setting assumes that the reward-generald process is controlled by an adversary. This assumption allows for modelling non-stationary and highly stochastic reward signals. We will later show why our FLAD formulation fits into this setting. Under this setting, it is assumed that an adversary is given access the learner's policy æ and determines the sequence of rewards, () for each arm prior to play (38]. At each turn и determines a distribution p(A). See Lattimore &
over actions, p(AA), and an action is sampled from the distribution, Szepesvári (34] for further details.
The EXP3 algorithm. The EXP3 algorithm ("Exponential-weigtt algorithm for Exploration and Exploitation") targets the adversarial multi-armed bandit problem by choosing arms according to Gibbs distribution based on the empirically determined importance-weighen rewards of arms [2I]. To allow for exploration, EXP3 mixes the Gibbs distribution with a uniform distribution.
Formally, let the exploration rate be y E (0, 1]. At round t, π defines the probability of selecting a given arm, A, as a linear combination of Gibbs and uniform distributions
(a) = (1 - -) Ixp()Ra.t-1/K))
7
Ea exp(y Rao
where the importance weighted reward Ra.t is calculated as
Ra,t
Ret Ro,c-1 Pt-1(a)
(2)
and Ra,, denotes the observed reward. All unplayed arms, a' F have unchanged importance weighted rewards; Ra Re',t-1.
Algorithmically, EXP3 takes the following steps at each round: First, calculate the sampling distribu- tion Pt and sample an arm from the distribution. Then a reward Re, is observed and the algorithm updates the importance weighted reward Ru, for the played arm.
Informally, the use of an apportance--eighhee estimated reward compensates the rewards of actions that are less likely to be chosen, guaranteeing that the expected estimated reward equal to the actual reward for each action. EXP3 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 > K. The exploration of EXP3 combined with importance inportance-weighting allows the policy to handle non non-stationary reward- seauur--gneraingg processes.
The UCB1 algorithm. While the adversarial setting makes almost no assumptions about the reward-gent generrting process and therefore maintains its performance guarantees under almost any circumstances, it can be outperformed in settings that are constrained. In this section we assume that the wwwrrdggnnerating processes are stationary Gaussian distributions. A common policy used to solve this MAB setting is the Upper Confidence Bound (UCBI) algorithm, which assigns each arm value called the upper confidence bound based on Hoeffding' 's inequality [22]. The UCBI algorithm is 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.
Formally, let the estimated mean reward of arm a after being played n. times be RS and the true mean reward be R,, then
2 2m1(1/8)
P Ra Ru
<8
4g (0.1)
nh
16

---

derived from Hoeffding's s inequality (following equation 7.1 of Lattimore & Szepesvári [4]). where ò is the confidence level that quantifies the degree of certainty the arm. In this work let 6 1/t where is the number of rounds played, shrinking the confidence bound over rounds. Thus, we define the upper confidence bound for arm a at turn as
00,
n = 0 otherwise
UCBa,t
(3)
Roo +
2
Algorithmically, UCBI takes the following steps at each round. First, the UCB1 policy plays the arm with largest upper confidence bound, a" arg max.5. U.B... Next, reward Ra is observed and the algorithm updates R.- (the estimated mean reward for a and the upper confidence bounds for all a. Informally, this algorithm suggests that the learner should play arms more often if they either 1. have large expected reward, R, or 2. na small because the arm is not well explored.
B Pseudo-code
We include here pseudo-code for our proposed algorithms. Algorithm I contains the pseudo-code for EXP3-FLAD, and Algorithm contains the pseudo-code for UCBI-FLAD.
Algorithm EXP3-FLAD
Require: Dy: Auxiliary and target datasets
Require: Jo: Parameterized model
Require: Gradient accumulation steps
1: Initialize: = |A]: &o = k:
Va E A: V. = 0, R. =
2: for = 1,2...., ......N. do
E, = min R: K
K
4: Ya A: (T))= + (1 K&E) Dn exp(Er-111.)) E.
exp(E, Ran)
5: Sample a (AA) and batch {x. y) D.
V. - V. VoC()xx,y)
7: ift (mod G) = 0 then
8: VT + VeL(fo Dr)
Update model parameters W.I.L. E. V. for all (a V. & 0)
11:
R, Rs
xie)
12: V. +
13:
end for
end
15:
end for
17

---

Algorithm 2 UCBI-FLAD
Require: D,, Dy: Auxiliary and target datasets
Require: Ja: Parameterized model
Require: G: Gradient accumulation steps
Require: 8: Smoothing factor
I: Initialize:
Ya A:
R_ con(VoL(o,,T)) VoL(6,
for 1.2.......
3: a" = argmax Ra + V 2ln aEA
4: Sample batch (x, y) Da-
V.- V.. + VoL(Jo, x,
6: Pu- + Nu- -m+++1
7:
ift (mod G) = 0 then
8:
VT + VoL(f, DT)
Update model parameters V E, V. for all {a V. # 0}
11:
R₂ (1 R. RRat
V. 0
13:
end for
14:
end if
15:
end for
C Training details
We train all models (FLAD and non-FLAD) 40Gb A100s.
For experiments, we use validation- based early stopping, and train for maximum of 10,000 gradient update steps. In practice, we find that early-stopping leads to significantly fewer than 10,000 updates, usually between 50-150 for direct fine-tuning, and 1-2,000 for other methods.
For the smoothing factor, 8, in UCBI-FLAD we ran preliminary experiments using values of (0.99, 0.9, 0.75, 0.5) and found 0.9 work well across datasets. All reported scores 3 0.9.
In preliminary experiments we consider rewards using gradients from multiple model partitions: the full model, encoder-only, decoder-only, and language modelling head (token classifier). We find that using the parameters from the LM head provides best performance, followed by the decoder-only, encoder-only, and full model gradients. The differential from best to worst method ~ 3% relative performance. Recall that with a gradient accumulation factor of G, our algorithms need to store at most G + 1 gradients at any time. So not only does using the LM head provide performance improvements, but also saves memory. For the models we use, the LM head contains only 2.3% of the full model parameters.
D Full results
The full results of experiments on target-only fine-tuning, explore-only, exploit-only, EXP3-FLAD, and UCBI-FLAR found on the next page.
18

---

Table 2: Detailed results from the main experiment including direct fine-tuning, exploration-only. exploitation-only baselines and our proposed methods, EXP3-FLAD and UCBI-FLAD.
TOMiK
TOMix
P3
UCBL PLADD (22 (2.
A (Te CoMa cooo
AAAAAA B8B1--1 (R 5
Sos--ine (R 1ec22
Lose Loss
Exyas
Coa---a (R Savss
na FLAM aoca aoptaen (R aecc
Loss axcce Fite- dumn
EXPx FLADD 1RA4A
ACBI-FLADD (R Sack UCBI-FT
EXP2 FLAD
SEBI-FLA (R Only
Loss Low: Scalnn (GA)
EXPR A (RRAAA UICRI SEAAEE aaio
BOBI-FLAA (R CoUS
PRCA SoPMIRE (R GMrS
CEBI-EELA (R
Loss Scaling (G2A) Loss Scaling CCi (C Eoportttinn Oaly
AAAAAA UR
Espleitation
Sepkuaaai Only
Loss Scatin (GAA)
Expleitatio Onin
Explone ChC
Scaling
Exspeitatinn Only
Direct FeelFin
Fine- duiigg
(R CoNEs
Veco
Loo Lo
Only
111
285
56
2.3
555
079
055
855
25
73
22
3
3
600
199
22
10
0
20,
300
2.5
0.0
30
40,0
38.8
38
35.1
330
811
XT1
000
0109
E1
210
911
9119
215
1119
800
C15
519
019
0.60
909
1009
338
338
115
0113
228
219
205
3
X00
570
098
6%
E33
099
9:23
022
699
1:25
CIF
1'6
F6E
10P
200
099
4200
999
1766
080
399
105
200
200
97
976
1:2
0000
575
F30
58
6139
F9R
915
190
ET3
575
118
18
9.6
0'06
688
988
08
1:28
H:50
009
120
$99
172
555
1/
11
000
213
160
988
2.66
0.99
1
18
99
9/9
000
1:90
90
68
p706
2.99
0.90
1.98
888
799
16
075
8
100
3
0P
0700
0.08
7.
0TH
0.0
6.0
970
0/13
9.59
6:3
69
9'8
1'99
6.2
0'3
219
532
F
0'28
911
9115
515
0/15
1'2
615
228
0/12
2115
015
728
023
003
110
211
115
1.98
788
198
199
9/6
228
2
9/05
1.5
1'88
6188
88
1818
0/00
11
5.0
F39
5:5
5'88
88
1588
630
6-2
399
0.7
000
110
116
310
910
F16
5T5
916
$180
$.98
9
110
906
C16
100
598
1100
6.8
0%
833
1:16
016
F88
E16
700
C10
1'8
8'8
108
P00
500
606
-16
1000
284
2.3
976
913
9/15
215
0'00
225
2:22
0'5
F33
073
123
205
273
273
C15
P39
535
$15
205
F15
205
011
115
550
300
25
31
000
XX6
0.
0.50
299
199
950
C00
109
X19
999
19
972
019
139
539
998
0'5
99
59
859
0389
299
295
818
615
E39
0'0
0729
619
137
609
263
555
023
966
91h
090
704
$99
909
0.95
65
8'8
F66
8'65
0'19
0'08
EW9
6%9
200
2.9
297
0'19
7'0
5'5
E39
611
11
125
215
6'2
099
615
295
49
115
113
699
133
5,0
990
5:0
-90
090
S79
223
209
20
2
220
620
619
2:1
68
9:59
55
59
590
[19
199
1709
0729
C1
E19
0229
S10
276
0.5
235
19

---

E Probing the reward generating processes.
Gradient Alignment RGA
Gradient Magnitude Similarity RGMS
Step
65
46
Slep 100
-1.00
-179
6.6
128
660
033
850
6n
-1.0
-152
-335
6.6
050
6.7
150
14
14
Slep 300
44
64
Step 400
-200
4404
6.25
420
94
Step 500
Figure 5: Reward distributions of RGA and GMS prior training and every 100 gradient updates thereafter. We probe the reward distributions using the T5-XL model with the TOMix auxiliary dataset and WSC 52 as the target dataset.
20

---

EXP3-FLAD and UCBI-FLAD training dynamics
The following 4 pages include a case study on the training dynamics of EXP3-FLAD and UCBI- FLAD when training T5-XL using TOMix as the auxiliary data. First, we find datasets where EXP3-FLAD and UCBI-FLAD improve significantly over the baseline FLAD methods, but also where either EXP3-FLAD or UCB1- (CB1-FLAD clearly outperforms the other. The two datasets that fulfill our interests are RTE and COPA.
We find that UCBI-FLAN --FLAD outperforms EXP3-FLAD on RTE, and show their respective training dynamics in Figure 6 (UCB1) and Figure 7 (EXP3).
We find that EXP3-FLAD outperforms UCBI-FLAD on COPA, and show their respective training dynamics in Figure (UCBI) and Figure 9 (EXP3).
We include details and takeaways in the caption for each figure. For EXP3-FLAD figures, we include charts of the cumulative estimated reward, empirical gradient alignment, instantaneous sampling distribution determined by the policy, and the empirical sampling distribution determined by the total number of samples seen per dataset as fraction of the total samples seen. For UCBI-FI, figures, we include charts of the upper confidence index, estimated gradient alignment, and the empirical sampling distribution.
21

---

Upper Confidence Bound
amazon_polarity
2.5
imdb
kilt_ hilitttsks/ooooaqa
qasc
1.5
wiki_ Liki_q
1.0
0.0
Estimated Gradient Alignment
0.75
0.50
0.00
-0.25
-0.50
Cumulative Empirical
Sampling Distribution
0.04
0.03
0.02
1600
3200
4800
6400
8000
9600
Number of Auxiliary Batches Seen
Figure 6: Training dynamics of UCBI-FLAD, case study using RTE as target dataset and TOMix as auxiliary data, where UCB I-FLAD outperforms EXP3-FLAD. Colored lines sample of auxiliary datasets with interesting properties, the remaining datasets are shown in grey. We find that even though wiki qa's estimated gradient alignment falls below 0 (middle), UCB1 does abandon sampling from it in the future. finding that between 3200 and 4800 batches, it becomes the dataset with largest upper confidence bound (top). Similarly. we see that UCB1 alternates between wiki qa. amazon_ amazon_polarity, and qasc as the datasets with higher gradient alignment and upper confidence bounds. kilt tasks/hotpotqa has very high gradient alignment prior training, but UCBI samples very infrequently from it, due to it'ls lower upper confidence bound. This is failure case for transfer learning- baased based methods. Interestingly, UCBI I never estimates imdb to have negative gradient, and gradually samples from it more and more frequently over the course of training.
22

---

Cumulative Estimated Reward
-2000
44000
amazon_ portrrtt imdb
8000
kilt patta
qasc
-10000
wiki_qp
-12000
Gradient Alignment
1.00
0.25
0.00
0.550
-0.75
-1.00
Sampling Distribution Sampling ling Distribution
0.05
Instantaneour
0.04
0.01
0.00
Cumulative Empirical
3200
6400
9600
12800
Number of Auxiliary Batches Seen
Figure 7: Training dynamics of EXP3-FLAD, a case study using RTE as target dataset and TOMix as auxiliary data, where UCBI-FLAD outperforms EXP3-FLAD. Colored lines are sample of auxiliary datasets with interesting properties, the remaining datasets are shown in grey. We find that the gradient alignment signal is particularly noisy for EXP3-FLAD, possibly leading to it's slightly worse performance on RTE. All five highlighted auxiliary datasets have high instantaneous sampling probability, but over the course of training, the empirical sampling distribution is very condensed across the full set of auxiliary datasets, unlike UCBI which is able to find better separation.
23

---

Upper Confidence Bound
2.5
cnn donndailymaiil3...
multi_news news
qasc
1.5
quartz
trec
1.0
0.5
0.0
0.8
0.6
Estimated Gradient Alignment
0.4
0.2
0.0
-0.4
-0.6
0.050
Cumulative Empirical
Sampling Distribution
0.045
0.035
0.000
0.025
0.020
3200
6400
9600
12800
16000
Number of Auxiliary Batches Seen
Figure 8: Training dynamics of UCBI-FLAD, case study using COPA as target dataset and TOMix auxiliary data, where EXP3-FLAD outperforms UCBI-FLAD. Colored lines are sample of auxiliary datasets with interesting properties, the remaining datasets are shown in grey. We find that although qasc and quartz start with very high gradient alignment, they very quickly fall to negative alignment (middle figure, green and yellow). In the end, we find that the algorithm samples much more from qase than from quartz (bottom figure). Interestingly, we find that although both cnn daillyaailllllllll and multi news start off with very negative gradient alignment, they quickly become the most aligned with the target task (middle figure, blue and red). We find that the three auxiliary datasets with highest upper confidence index (top figure) and largest sampling percent (bottom figure) are enn dailymail, multi news, and trec even though these considered dissimilar to the target prior to training.
24

---

Cumulative Estimated Reward
750
250
cnn hin_dallymail33...
multi _eew
-250
qasc
quartz
7700
trec
Gradient Alignment
1.00
0.75
0.50
0.00
-0.25
00.7
Sampling Distribution
Instantaneous
0.03
0.02
0.01
Cumulative Empirical
Sampling Distribution
0.01
6400
9600
12800
16000
19200
Number of Auxiliary Batches Seen
Figure 9: Training dynamics of EXP3-FLAD, case study using COPA as target dataset and TOMix as auxiliary data, where EXP3-FLAD outperforms UCBI-FLAD. Colored lines are sample of auxiliary datasets with interesting properties, the remaining datasets are shown in grey. This is an impressive example of the importance- weighted estimated reward. We see that cnn con_dailymal and multi_ news both start with very negative alignment, but EXP3 quickly updates it's estimated reward once their alignment becomes positive. Similar to RTE, we see that EXP3 never makes large separations in the empirical sampling distribution, possibly reason why UCBI outperforms EXP3 overall. Compared to RTE, we find that gradient alignments are much less variable, with a maximum alignment close to 0.5 and minimum alignment close to -0.5. Whereas in RTE, alignments regularly reach close to 1.0 and -1.0.
25

---

Explore-Only
Exploit-Only
EXP3-FLAD (RCA)
UCBI-FLAD (R.A)
35
(TOMix)
570.9
549.1
769.1
700.0
A
260
(P3)
863.6
692.7
832.7
794.5
% increase
51.3%
26.2%
8.3%
13.5%
Table 3: Number of training iterations for TO-3B to converge using a training method (column) and set of auxiliary datasets (row). The number of iterations convergence is averaged across 11 target datasets and 5 seeds, leading to 55 experiments aggregated per cell.
G Effect of scaling [A] on time-to-convergencce
As we have described in this work, the computational complexity for single turn of our methods are independent of the number of auxiliary datasets. However, it is unclear whether the computation complexity of the multi-armed bandits are dependent on the number of auxiliary datasets through their exploration rates. Thus, the computational complexity of an individual training run may be influenced by the number of auxiliary datasets (,A),, but is not possible characterize this relation explicitly as it relates to the complex and stochastic process of training and large language model.
To better understand the empirical effects of increasing |A| the fime--o-model--coneegeeee perform a study on the number of iterations to convergence for different FLAD algorithms. Table B shows that all methods require longer training to converge when increasing from |A| 35 to 260. We find that, compared with baseline methods, our MAB-based methods require more steps for the smaller set of auxiliary datasets, but the number of additional steps required to train our methods only increases modestly (~ 10%) when increasing |A| by a factor of nearly 10. In contrast, the Explore- and Exploit-Only methods do not scale nearly as well when increasing the number of auxiliary datasets. Notably, the Explore-Only method requires over 50% more training iterations for P3 than for TOMix, at which point takes longer to converge than either of the MAB-based methods.
H Auxiliary Datasets
Here we include the full list of auxiliary datasets from P3 (23) used to train models for the ANLI target tasks. Other target datasets have slightly different auxiliary datasets due to test set decontamination, but are generally the same. Datasets are listed by their name as found in HuggingFace Dataset
Zaid/quac_ _ecelade_epaddedd,, acronym identification, ade corpus v2/Ade_corpus ade_ wttppcoorpu///eeee v2/Ade corpus v2 drug add_relation, ade_ corpus_ citiss_2/Adeeeeeee corpus v2 drug dosage_ relation, adversarial Satprraal__aaaddeerrrraaAQ/ adversarial _ adversarial qa/dbidaf, adversarial Baaersariallgg/drbbbretttt
classification,
aeslc, ag_ news,
ai2 arc/ARC-Fasy, rww.ees//Wirelessss v1 v1_00,
amazon_polarine app_reviews, lin__coronavirus.cmmm soronavirus_guery_.ee
ama-
zon zon_reviews_ geviiess_mlltien.. amazon _usooonnns____
roua_rat/raw, art, asset/ratings, asset/simplification, banking77, billsum, bing_ biosses, blbooksgenre/title__eeee Laarieggnnneeeeeeeeeeeeeeee classifiction, blended _kill _tlk, cb/CN, cb/NE, cbuP, cbt/V, cbt/raw, cc ec_news, circa, climate_ enn dailymail/3.00. 0.0.0, codah/codah, codah/fold O, codah/fold_1, codah/fold_2 codah/fold_3, codah/fold 4, code_x_glux sttc_tt_____e_ cec_teet_to___ee cooeoeoooooooooooeeeoeeeeeeeeeee com- mon gen, commonsense __a, conv ai, conv ai conv ai 3, cord19/metadata, cos cos_e/vl.0_ e/v1.0, cos_eht los_e/v1.11, cosmos_ee covid castorini, craffel/openai lambada, craigslist_barge crows rrors_ppirr, dbpedia 14, discoouseedtiseofuss--spp cry/discovery, docred, dream, drop, duorc/ParaphraseeC. duorc/SelIRC,, e2e nlg cleaned, ecthr www.ee/wallleeeddoolttooo-ooooee idence_infer idence_infee__treammm//LLLL evidence_infes diadenee_infer_rraamee2222.00
discov-
emo, emotion, enriched, web web_nigg/en, esnli, fever/v1.0. fever/v2.0, finan- cial hol__phraaecann//eettencs. aallggeee freebase_ge generated reviews_entle
gigaword, glue/ax, glue/cola, glue/mnli, glue/mnli hiee/mmni_mmtchhdd, glue/mnli cawialiimmmmmthheed... glue/mrpe, glue/qnli, glue/qqp. glue/rte, glue/sst2, glue/stsb, glue/wnli, google_ wellformed, _uerr, great_code, guardian_authonse www.aaooshpp/rrrosgggaooaaa genooe__
guardian_authorsee wawaaaoorshppoorosooooooooo _tppi__, guardian authorship/cross htcpipp//ooosssocic___pe 4, guardian authorship/cross topic _, gutenberg time, hans, hate satee_specch8, head qa/en, head_ga/en,
health fact, hlgd, hotpot_ gappppt_addsittctttrr hotpot gathttlwwwllliiiii....
Panps:/huuggingface.oottaa
26

---

humicroedid/subtask-l humicroedit/ssbiuskkk hyperpartisan, news. detection/byarticle, hyperparti- san news detection/bypubliheer, imdb, jfleg, kelm, kilt tasks/hotpotqa, kilt tasks/nq, lama/trex, lambada, liar, limit, math aaah__ataset/algebre hatpr//llleba__iearr__aaa linear_ 1d, dataasaeaaaaebbaa__iin Baneebrr___iearr_____ lere__d_ccmm id_composed, math, _ _iieear 2d, math math_ __taaaet/algebba__iiin linear 2d_ _d_composed, math_qq _a, me_taco, mdd/task1 _aa, _qa, mdd/task2_res _eccs, mdd/task3_ qareess medal, coadical_questtonn
_e_____ssseee_eie__e mwse. search/evaluation_daae _attaatt newspop, wew.eeeellllndleeeeeee openbookkaaaddiitonnl, open- bookqa/main, paws-x/en, paws/labeled final, paws/labeled, _wapp paws/unlabeled final, piqa. semm_ssntmmen, pubmed pareeddd_apppaaaeeeeeeee labeled, qa_srl, qa_zre, qasc, qed, quac, quail, quarel, quartz, quora, quoref, race/all, race/high, race/middle, riddle sense, ropes, rotten tomatoes, samsum, acan/addprim_jumg jump, scan/addprim Saaprri_ttt___ttt turn_left, scan/filler_mung
meta_woz/diallog mocha, moovie_rational multi_news, multi_nli, narrativeqa, ncbi bibease, neural _owellloooeeeeeeeeeeeeeeeeeeeeee nq_open, sense, ceestop_englisee
labeeee
haa/filler_numl ocan/fillernnmm2 (can/fillernnumi scan/length, scan/simple, scan/template sareee_rrruudddddee httpheempppteeoppostteeeeeee sright, scan/template_rig rrigtt scicite, seiq, scitail/snli format, scitail/tsv _format
scan/template_jump caomp_around_rre
scientific crcentificppapes//ariii scitldr/Abstract, selqa/answer_selen cetrrrsseeeettooseeeeeee annallsss, sem__val__ee Leveellllll0___ 1100_ass___ 8, sem_ 2014_task_1, sent comp, sick, sms_spo spam, snips_buill built_in _niit_in_intents, intents, snli, social_i_qa, species_ 800, squad, aalaaaadvvaraaaddaoo
squad_v2, v2, squadshifts/amazzn. squadshifts/new_witte squadshifts/ny, sst/default, stsb multi mt/en, subjqa/books, subjqu/clectronies. subjqa/grocery, subjqa/mmovies. subjqu/restaurants, subjqa/tripadvisor, super_ glue//xb//xxb.... super glue/axg, super glue/boolg, su- per glue/multire, super glue/record, swag/regular. tab fact/tab fact, tmu tmu_ggfm dataset, trec. Livia_gg/unfilteredd turk, tweet_eet seee___el//eooji, tweet_ vrcet_eval/eootion, tweet_eval/heee tweet sereeetrevll/rony... weeet_eallooffensive tweet_eval/sentimen leree__llltttteeeeeeee ffeminist, tweet_ prce__eeal/stnnce hillary, ty- tweet ureet_eval/sttnce_ehll cliimete, tweet_e
tweet evwl/stance_attiiieee
diqa/primary_siee _taak, tydiqa/secondary _takk, web_ questions, wiki_ wiki hop/maskedd
wino siino_baasttpel__pp viki_hop/original wiki_qa. wiki_split, wino catio_iia/ttpee_eae winograd hinograd_wsc/wseed wsc/wsc285, sino_bias/tppe__ntti bino_bias/type2.11
wiqa, xnli/en, xquad/xquad.en xquad_r/en xsum, yahoo_anse aamioo_asswer_gpr qa, aahoo_ansswerscoo topics,
selp_polarity, yelp repiereview.... full, zest
27


Reviews:
[{"review_id": "SgqIq7M2yf", "reviewer": "Reviewer_EHmY", "summary": "This paper focuses on the problem of few-shot learning with auxiliary data (FLAD) problem. The authors formulate FLAD as a multi-armed bandit problem, and employ exploration-exploitation algorithms to tackle it. Experiment results show that the proposed method outperforms exploration or exploitation-only approaches by significant improvement in accuracy.", "questions": "Please address the concerns in the Weaknessese.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- This paper models the FLAD problem with the multi-armed bandit framework, which makes a difference from previous works. The proposed method scales well to the number of auxiliary datasets, and adds a limited amount of computational and memory overhead.\n\n- The proposed approach shows a significant accuracy improvement over other exploration-only and exploitation-only baselines.", "weaknesses": "-The novelty is limited since the proposed multi-armed bandit method is a existing technique.\n\n- The authors should clarify the difference in definitions and settings between FLAD and other similar problems such as general few-shot learning and multi-task learning.\n\n- More experiment results had better be provided. For example, in order to demonstrate the scalability of the proposed method, it would be better to provide what the changes of performance be like with different values of |\\mathcal{A}| or |\\mathcal{D}_T|.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their useful feedback, which points out some details of the paper that required clarification and for the suggested studies that can strengthen the paper. Our responses are below:\n\n>The novelty is limited since the proposed multi-armed bandit method is a existing technique.\n\nWe respectfully disagree with the claim that our work lacks novelty. While our work shows that the MAB framework is a natural fit for FLAD, to the best of our knowledge, MABs have not previously been used in the context of FLAD. The novelty of our work lies in how we adapt and apply multi-armed bandits to FLAD. Our approach results in significantly improved computational complexity compared with prior FLAD methods, allowing us to scale to 100x more auxiliary datasets than previous methods, leading to a 4% accuracy improvement. These results mark a leap in the efficiency and effectiveness of few-shot learning, which can hardly be considered incremental. Thus, while we have built upon existing techniques from multi-armed bandits, we believe the outcome is a transformative contribution to the field of few-shot learning.\n\n>The authors should clarify the difference in definitions and settings between FLAD and other similar problems such as general few-shot learning and multi-task learning.\n\nThank you for your valuable question. We elaborate on the differences below.\n\nAs a reminder, we use the following notation (section 3.1 lines 99-103):\n- $\\mathcal{D}_{T}$ is the target task dataset.\n- $\\mathcal{D}_{A}$ is the set of all auxiliary datasets.\n- $\\mathcal{D}_{a}$ is a single auxiliary dataset.\n\n**Few-shot Learning vs. FLAD**: 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}\\_{T}$, and there is no auxiliary data used (so $\\mathcal{D}\\_{A}$ is effectively the empty set for few-shot learning). In contrast, in the FLAD setting $|\\mathcal{D}\\_{A}|$ > 1, and in our settings specifically we use T0Mix and P3 which have $|\\mathcal{D}\\_{A}| = 35$ and $|\\mathcal{D}\\_{a}| = 260$ auxiliary datasets respectively. This information all comes from section 4 under the paragraph “Auxiliary datasets'' on lines 206-213.\n\n**Multitask Learning vs. FLAD**: Multitask learning is concerned with optimizing a model for performance on multiple target datasets simultaneously. This is in direct opposition with FLAD which aims to optimize a model for a single target task. Formally, FLAD considers the scenario where $\\mathcal{D}\\_{A}$ is a set of multiple datasets and $\\mathcal{D}\\_{T}$ is a single dataset. On the other hand, multitask learning would consider a scenario where $\\mathcal{D}\\_{A}$ is effectively the empty set, and $\\mathcal{D}\\_{T}$ is the concatenation of multiple datasets where one attempts to simultaneously optimize for performance all datasets.\n\nWe hope this clarifies the differences between few-shot learning, FLAD, and multitask learning. We have made these distinctions clearer in the updated draft.\n\n>More experiment results had better be provided. For example, in order to demonstrate the scalability of the proposed method, it would be better to provide what the changes of performance be like with different values of $|\\mathcal{A}|$ or $|\\mathcal{D}_T|$.\n\nWe thank you for your insightful suggestion of additional ablation studies.\n\nFirst, we would like to reiterate that the main experiments in this work do include results for two different sizes of  $\\mathcal{A}$ (specifically, 35 and 260) to demonstrate the scalability of our proposed methods. On lines 273-277, we discuss how our methods improve between 2.6-4% when increasing $|\\mathcal{A}|$ from 35 to 260, while the Loss-Scaling baseline only improves between 0.2-2%. Additionally, in lines 281-285, we show how the Explore-Only and Exploit-Only baselines only improve between 0.7-2%, further showing that our methods demonstrate improved scalability compared with existing FLAD baseline methods.\n\nWe agree that the additional ablation studies you suggested will provide for interesting additional analysis. We therefore performed an ablation on values of $|\\mathcal{A}|$ in {35, 75, 125, 175, 225, 260} (all taken from P3) for all target tasks using T0 as a base model and repeat the experiments with 3 random seeds. The results are in Figure 1 of the pdf attached to the global rebuttal. We found that the performance for both EXP3-FLAD and UCB1-FLAD jumps significantly from $|\\mathcal{A}|$ = 35 to 75, followed by a more steady increase in performance all the way up to $|\\mathcal{A}| = 260$. The results show an increase of 2.54 accuracy for EXP3-FLAD and 3.12 for UCB1-FLAD when increasing from 35 to 75 auxiliary datasets, with an additional increase of 1.54 for EXP3-FLAD and 0.47 for UCB1-FLAD when increasing from 75 to 260.\n\nTo your point about varying $|\\mathcal{D}_T|$, our main experiments do not vary the size of the target dataset in order to maintain fair comparisons with past work that uses the same standard dataset sizes that we did. However, we see that there is value in better understanding the relation of our method with varying target dataset sizes. Unfortunately, we were unable to perform experiments varying $|\\mathcal{T}|$ in time for the rebuttal due to computational resources, but will do our best to include them during the discussion period, and will include them in the final draft.\n\nWe hope this information is helpful. Please let us know if you have any further questions."}, {"review_id": "b0zYW51BLo", "reviewer": "Reviewer_yskH", "summary": "In the context of Language Models, this paper studies the use of Auxiliary Data during training to gain in generalization. The Few Shot settings is also considered here where proposed tasks only have limited set of labeled data. Compared to previously proposed approaches using auxiliary data, here it is consider to exploit more auxiliary dataset than before. In this context, then it is crucial to select what are the good auxiliary datasets to be used (otherwise performance could degrade) with a reasonable complexity. \n\nA multi-armed bandit (MAB) approach is proposed here to iteratively refine a selection policy among these auxiliary datasets. Two MAB algorithms have been studied (EXP3 and UCB1) in conjunction with different reward functions (based on gradient alignment, gradient magnitude alignment of a combination of both). Experiments are reported on specific training of T5-XL and T0 model. Significant gains with respect to previously proposed approaches are reported. Moreover the complexity associated to the proposed method allows to use high number of auxiliary datasets which helps to boost the performance further. \n\nDiscussions on the various aspect of the method are also provided.", "questions": "1. Author mentioned that complexity is independent of the number of auxiliary datasets (l. 370-371). Is it really true? When considering even higher number of auxiliary datasets, it could need more iterations to reach a good policy $\\pi$ among datasets (need more exploration). So aren't there really an independence there? Or practically since rewards distribution tend to converge quite fast (see fig 4 in annex), impact should be limited?", "limitations": "Authors adequately addressed associated limitations.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is clearly written. The adaptation of MAB here allows to derive a performant policy to select auxiliary datasets for training. \n\nSpecifically since model updates and policy updates are intrinsically performed, the overall complexity is not armed by the number of auxiliary dataset. \n\nSignificant gain are reported with respect to previous works, and relevant discussions are provided upon the importance of rewarding functions used and associated algorithms (e.g. benefits of having exploration and exploitation phases thanks to the bandit).", "weaknesses": "As mentioned by the authors, the proposed training has to be performed from scratch for every domain. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the overall supportive review of our work, and for the insightful question about complexity which led us to a very interesting analysis that we now include in the updated draft.\n\n>As mentioned by the authors, the proposed training has to be performed from scratch for every domain.\n\nIt’s true that the training needs to be performed from scratch on each target dataset, but this is true of most few-shot methods, and the computational costs of the proposed methods are fairly cheap. We found an average training time of 6 hours on a single 40Gb GPU.\n\n>Author mentioned that complexity is independent of the number of auxiliary datasets (l. 370-371). Is it really true? When considering even higher number of auxiliary datasets, it could need more iterations to reach a good policy ($\\pi$) among datasets (need more exploration). So aren't there really an independence there? Or practically since rewards distribution tend to converge quite fast (see fig 4 in annex), impact should be limited?\n\nThe computational complexity of a single turn of our method is independent of the number of auxiliary datasets (as you mentioned in Strengths, the policy updates are intrinsically performed). However, you are correct that the computational complexity of the multi-armed bandits are dependent on the number of auxiliary datasets due to their exploration rates (line 146 for Exp3 and 168 for UCB1). Thus, the computational complexity of an entire training run will be influenced by the number of auxiliary datasets, but is challenging to characterize exactly as it relates to the convergence of training a large language model.\n\nTo better understand the empirical effects of increasing the number of auxiliary datasets on the time to model convergence, we performed additional analysis of the convergence time for different algorithms. Indeed, we do find that the number of auxiliary datasets requires a longer training, but we find the same to be true for any multitask fine-tuning (see the table below). We find that, compared with baseline methods, our methods require more steps for the smaller auxiliary dataset, but the number of additional steps required to train our MAB-based methods on P3 only increases modestly when increasing the number of auxiliary datasets by a factor of nearly 10 (in contrast with the the explore- and exploit-only baselines, where the number of steps to convergence increases significantly).\n\n*The table below shows the number of training iterations for T0-3B to converge using the training method across the top row, and the auxiliary dataset on the left column. Iterations are averaged over 11 target datasets and 5 seeds leading to 55 experiments averaged per cell.*\n\n| Average # iterations | Explore-Only | Exploit-Only | EXP3-FLAD (GA) | UCB1-FLAD (GA) |\n| ---- | ---- | ---- | ---- | ---- |\n| T0Mix (35 auxiliary datasets) | 570.9 | 549.1 | 769.1 | 700.0 |\n| P3 (260 auxiliary datasets ) | 863.6 | 692.7 | 832.7 | 794.5 |\n| % increase | 51.3% | 26.2 | 8.3% | 13.5% |\n\nWe will clarify the statement on lines 370-371 to say “our algorithms have a single-turn computational complexity that is independent of the number of auxiliary datasets”.\n\nWe appreciate your thoughtful question! It has led us to some very interesting analysis that we will include in the discussion section of the revised paper.\n"}, {"review_id": "cVoqMRK8t0", "reviewer": "Reviewer_jiND", "summary": "The paper addresses the problem of using auxiliary  data for the few-shot learning. The main novelty of the proposed approach is to formulate the selection of the auxiliary dataset from a pool of datasets as a multi-arm bandit problem (balancing exploration and exploitation) and adapting EXP3 and UCB1 to that goal. This allows to scale the number axillary datasets and improve generalization in the few-shot learning. The method is evaluated in application to NLP.", "questions": "One of the main claims in the paper is that the proposed approach almost doesn't add to computational cost.  The  description in line 172-174: \" This is done for each auxiliary dataset by calculating the gradient ∇a = ∇θL(fθ, x, y), where the number of samples in {x, y} is significantly larger than a mini-batch,  and can be up to the size of the full dataset.\"  doesn't seem to go along with that clain, as this computation is large and it has to be done in every round. Do I miss something here?\n\nA more minor thing that needs to be taken care of is in line 219. Exploit-only is suppose to be an opposite to Explore-only. So it's confusing how Exploit-only is an extension of Explore-Only. The text later explains the idea, but it needs amendment. ", "limitations": "I did not really see any clear discussion of limitations. While the selection of the datasets are done using the proposed algorithms of balancing exploration and exploitation, there is still an implicit assumption that the datasets are useful. I believe that in language tasks this is less of a problem but in general it could be the case. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Important problem\nImproves the result with a smaller model\nShows nice empirical results", "weaknesses": "The novelty of the paper is incremental\nThere are some problems in the presentation that cause confusion as explained below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your very insightful questions and for the opportunity to respond with the clarifications below:\n\n>The novelty of the paper is incremental\n\nWe would like to clarify what we believe is an important misunderstanding. While the building blocks of our algorithms may be derived from existing algorithms, we emphasize that the ultimate impact of research often lies in how the components are combined and applied, and the resulting effectiveness. Our work presents a novel application of the multi-armed bandit framework, and demonstrates a significant advancement of methods in the Few-shot Learning with Auxiliary Data (FLAD) setting.\n\nOur novel algorithms (EXP3-FLAD and UCB1-FLAD) are shown to solve an important problem within FLAD, namely how to scale to the ever-increasing number of available datasets. Through our adaptation of MAB to the FLAD setting, we show that we can scale to 100x more auxiliary datasets than previous methods, leading to a 4% accuracy improvement. These results mark a leap in the efficiency and effectiveness of few-shot learning, which can hardly be considered incremental. Thus, while we have built upon existing techniques from multi-armed bandits, we believe the outcome is a transformative contribution to the field of few-shot learning.\n\nAdditionally, we would like to emphasize that it was not obvious, a priori, that our proposed methods would work. In particular, recall from lines 202-204 that we divide each few-shot target dataset into equal training and validation partitions, where the training partition is used to calculate rewards and update model parameters, while the validation partition is used for early stopping and selecting the best model checkpoint. One particularly novel and unexpected finding is that even with only 10-35 training samples, our reward functions are able to create a meaningful signal for the MAB algorithms.\n\n>One of the main claims in the paper is that the proposed approach almost doesn't add to computational cost. The description in line 172-174: \" This is done for each auxiliary dataset by calculating the gradient $\\nabla a=\\nabla L(f_{\\theta},x,y)$, where the number of samples in {x, y} is significantly larger than a mini-batch, and can be up to the size of the full dataset.\" doesn't seem to go along with that clain, as this computation is large and it has to be done in every round. Do I miss something here?\n\nWe would like to clarify that the computation you are referring to is **not** performed every round. As described in lines 170-171, this computation is used to initialize the estimated reward for each arm prior to training, thus it only occurs once. To provide you with a more detailed answer, in traditional UCB1 the algorithm starts by playing each arm for one turn (to initialize rewards), and then continues accordingly. In our work, we take advantage of the fact that we can obtain a better estimate of the initial reward by using a larger number of samples prior to training (rather than a single batch as would be done if we strictly followed UCB1). In practice, the auxiliary gradients computed prior to training can be reused for each target dataset, so the cost is amortized across target datasets. Furthermore, we found that computing gradients for 1000 samples is very fast on a single GPU (~2 minutes per dataset). We will add this clarification to the paper to avoid additional confusion.\n\n>A more minor thing that needs to be taken care of is in line 219. Exploit-only is suppose to be an opposite to Explore-only. So it's confusing how Exploit-only is an extension of Explore-Only. The text later explains the idea, but it needs amendment.\n\nThank you for pointing out the confusion about Explore-Only vs. Exploit-Only. We’ve altered the text from line 219 to say: “Exploit-Only computes gradient alignment prior to training…” at which point it continues as originally written.\n\n>I did not really see any clear discussion of limitations. While the selection of the datasets are done using the proposed algorithms of balancing exploration and exploitation, there is still an implicit assumption that the datasets are useful. I believe that in language tasks this is less of a problem but in general it could be the case.\n\nYou are correct that one of the assumptions for all methods in the Few-shot Learning with Auxiliary Data (FLAD) setting (all prior works included) is that there is at least *some* auxiliary data which will be useful for the target task.\n\nHowever, 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 must be useful, and thus the auxiliary datasets should be hand-picked by humans. On the other hand, our methods allow for only a small portion of the auxiliary data to be useful, and we allow the proposed algorithm to explore and find those auxiliary datasets which will be most useful, then exploit them.\n\nOf course, as you mention, because our experiments are on language tasks we cannot make guarantees about any other domains/modalities (e.g. vision, time-series data). However, we believe that our findings suggest that in any domain where multi-task learning is an option, the presented methods can be beneficial.\n\nThank you for mentioning this limitation. We have added your point, and our thoughts, to the discussion section."}, {"review_id": "r3fd4QIWfh", "reviewer": "Reviewer_5GTn", "summary": "The paper presents a new approach to Few-shot Learning with Auxiliary Data (FLAD), a method that uses additional data to improve model generalization. Traditional methods for integrating auxiliary and target data have limitations, particularly in terms of computational scale with an increasing number of auxiliary datasets. This study introduces two new algorithms, EXP3-FLAD and UCB1-FLAD, that relate FLAD to the explore-exploit problem in multi-armed bandit scenarios, thereby achieving computational complexity that is independent of the quantity of auxiliary datasets. This paper shows that a combined exploration and exploitation strategy is vital, with the new methods surpassing all existing FLAD methods by 4%. The research also yielded the first 3 billion parameter language models that outperform the 175 billion parameters GPT-3.", "questions": "* Why is the policy configured to select a singular specific auxiliary dataset? Could there not be more value in a more flexible approach, allowing the policy to identify and utilize the most appropriate samples across the entirety of the auxiliary data? Especially considering the successes achieved with cross-task transfer for T5 and T0, it seems that an exploration of this approach could be beneficial.\n\n* Figure 2 the HellaSwag task, where all methods show inferior performance, with the exception of GPT-3. Could there be certain features inherent to this task that align better with the GPT-3 model, or are there particular aspects of the task that the other models struggle with? Do you have some intuition?\n\n* Do you have updated results against GPT-4?", "limitations": "* I do not have much practical experience with RL and it is difficult for me to judge the limitations of this method.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This study introduces a novel approach to enhance few-shot learning in language models by leveraging auxiliary data. The innovative use of additional datasets significantly improves the generalization capabilities of the models in the few-shot learning scenario.\n\n* The incorporation of an exploration/exploitation framework, often seen in multi-armed bandit scenarios, adds an interesting aspect to the methodology. It ensures the computational complexity is independent of the number of auxiliary datasets, thus addressing the significant issue of scalability often seen in previous methodologies. \n\n* The experimental results displayed in Table 1 and Figure 2 are good and offer strong support for the efficacy of this approach.\n\n* One of the most exciting features of this methodology is its universal applicability; it can be effectively utilized with any language model. It provides a path to significantly elevate the performance of relatively smaller models, leveraging them to outperform even the more sizable models, as demonstrated by the 3 billion parameter language model surpassing the 175 billion parameters GPT-3.", "weaknesses": "* Figure 2 presents an interesting picture of the model's performance, revealing an interesting dependency on the auxiliary data for some tasks, while seeming almost indifferent between using T0 or T5 for others. The interpretation of these results is not straightforward and questions about the behavior exhibited by the model across different tasks. \n\n* In terms of its responsiveness to auxiliary data and the use of T0 or T5, the performance differs significantly from task to task. This variance potentially suggests that certain tasks might be more intrinsically aligned with the structure or features of the auxiliary data, or that some tasks may respond more favorably to the model architectures embodied by T0 and T5, but there is no ablation or discussion about this aspect.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful questions and overall appreciation of our work, we share an excitement about the proposed methodology.\n\n>Figure 2 presents an interesting picture of the model's performance, revealing an interesting dependency on the auxiliary data for some tasks, while seeming almost indifferent between using T0 or T5 for others. ...\n\nWe’d like to clarify that all four of our models in Figure 2 were trained with the same auxiliary data, P3. The difference between each model is the algorithm used in FLAD training (EXP3-FLAD or UCB1-FLAD), and the base model used to initialize FLAD training (T5 or T0).\n\nFigure 2 presents a very clear separation in performance for two datasets: anli-r1 and wsc (T0-based models have significantly better performance than T5). The fact that we *only* see such a separation in 2/11 datasets was a surprising result for us, as we expected that the multitask-trained T0 would always be a better starting point for our FLAD training. The exact reason why starting with T0 only improves performance on these two datasets is difficult to determine. Because the models in Figure 2 were all trained on P3, but only T0 has been multitask trained explicitly on T0Mix, these results suggest anli-r1 and wsc benefit from the extensive multitask training of T0. Conversely, this suggests that for the remaining 9/11 target datasets, our FLAD methods influence few-shot generalization ability more than general-purpose multitask training.\n\nThank you for pointing out this interesting direction of analysis, we will add this surprising result to the analysis of the paper.\n\n>In terms of its responsiveness to auxiliary data and the use of T0 or T5, the performance differs significantly from task to task. ...\n\nWe agree that the variance across tasks is likely coming from the available auxiliary data, but should note that there is no architectural difference between T5 and T0, so we don’t believe that architectural differences could cause the variations in performance.\n\nFigure 3 demonstrates how certain auxiliary tasks are more aligned with the features of the WSC target task, and In particular we show that at step 300 there is a clear differentiation between those auxiliary datasets which we believe to be aligned with the target task (Figure 4 in the appendix expands this figure to every 100 training steps).\n\nAdditionally, Section F in the Appendix shows four figures containing in-depth information on the alignment between target and auxiliary tasks. Specifically, consider Figure 5 in the appendix. We look at the training dynamics of the target task RTE, which is a binary classification task. The figure shows that amazon_polarity, imdb, wiki_qa, qasc are the auxiliary datasets which are most aligned with RTE, and thus most sampled datasets. Interestingly, these tasks are all binary classification, or multiple choice datasets formulated as binary classification. On the other hand, we see that kilt_tasks/hotpotqa, which is a generation task, is the least sampled dataset. These examples suggest that the structure of the auxiliary datasets are, at least partially, a predictor of usefulness for target tasks.\n\nWe will add these interesting findings into the discussion section of the main paper.\n\n>Why is the policy configured to select a singular specific auxiliary dataset? Could there not be more value in a more flexible approach, allowing the policy to identify and utilize the most appropriate samples across the entirety of the auxiliary data? ...\n\nFirst, there is definitely value in applying approaches that identify and utilize the most appropriate samples across the full set of auxiliary data! In fact, there is an excellent direction of study called data selection that focuses on exactly that question (we briefly discuss this area and provide references in the related works section on lines 89-92).\n\nSecond, we actually compare against DEFT-Few (Figure 2) which does something similar to what you are suggesting. In DEFT-Few they train their model on the 500 nearest neighbor auxiliary samples to the target dataset (more detail on lines 290-293). We found that our method outperforms DEFT-Few by 2.7% absolute accuracy under the same settings (T0 as the base model, P3 as the auxiliary data).\n\nFinally, our method calculates the policy online (at every turn) and there are some downsides to identifying the most appropriate samples across the entirety of the auxiliary data at every turn. Mainly, the method would need to reindex the entirety of the auxiliary data at every turn (or every N turns), making it very costly and inefficient. In contrast, to update our policy, we only need to know the reward from the current turn, making it very efficient.\n\n>Figure 2 the HellaSwag task, ...\n\nThank you for pointing this out, it’s a known phenomena that is difficult to give an exact answer to because the exact details of GPT-3 training is unknown. Sanh et al. [1] and Wei et al. [2] also found that their models (T0 and FLAN) underperform GPT-3 on HellaSwag. They suggest that because this task can be formulated as finishing an incomplete sentence, it is ideally formulated for large language models such as GPT-3. Furthermore, we note that HellaSwag was created through adversarial filtering [3], leading to a dataset with adversarially challenging samples and may adversely affect smaller models more than large models.\n\n>Do you have updated results against GPT-4?\n\nWe assume the performance of GPT-4 would be very strong in this setting, however, in this work we focus on methods of achieving high performance from much smaller base models.\n\n[1] - Sanh et al. Multitask Prompted Training Enables Zero-Shot Task Generalization, 2023. https://arxiv.org/abs/2110.08207\n\n[2] - Wei et al. Finetuned Language Models are Zero-Shot Learners, 2022. https://arxiv.org/abs/2109.01652\n\n[3] Zellers et al. HellaSwag: Can a Machine Really Finish Your Sentence?, 2019. https://arxiv.org/abs/1905.07830"}]

Author rebuttal:
We would like to express our gratitude to the reviewers for their valuable and detailed feedback. We appreciate the time and effort they took to review our draft and provide us with their insightful comments. We carefully consider all of the suggestions and incorporate them into our final version. We are confident that the final product is improved as a result of their feedback.

See the attached pdf for results of a study on adjusting the number of auxiliary datasets, $|\mathcal{A}|\in$ {35, 75, 125, 175, 225, 260}.

Meta-review: The paper tackles the problem of few shot learning with auxiliary datasets (FLAD) by relating it to the multi-armed bandit (MAB) problem. The authors make use of know MAB algorithms in order to build a solution for the FLAD problem. There is a consensus in the reviews that the problem tackled is important, and the experiments show a convincing picture regarding the effectiveness of the proposed method. The main concerns raised were about (1) the novelty of the work, (2) the clarity of manuscript. Regarding (1), the authors indeed did not invent a new algorithm, rather used existing ones as building blocks towards their solution. Given that other papers already tackled the FLAD problem and did not make use MAB in order to solve it, and that the proposed solution here is agreed to perform well, I do not find this to be a major limitation of the paper. For weakness (2), it seems that all the major problems pointed out in the reviews are resolved in the rebuttal. It seems that the required work needed to integrate these clarifications to the paper is not major and can be done towards a camera ready version. Given this, I recommend accepting the paper and urge the authors to carefully go over the discussions here and insert the appropriate changes to their paper.
Decision: Accept (poster)