Title: Improving Few-Shot Generalization by Exploring and Exploiting Auxiliary Data
Venue: NeurIPS 2023
Accepted: yes
Keywords: Few-shot learning

natural language processing

few shot learning

NLP

multi-armed bandit

multi armed bandit
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

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 F ew-shot L earning with A uxiliary D ata (FLAD), a training paradigm that assumes access to auxiliary data during few-shot learning in hopes of improving generalization. Previous works have proposed automated methods for mixing auxiliary and target data, but these methods typically scale linearly (or worse) with the number of auxiliary datasets, limiting their practicality. In this work we relate FLAD to the explore-exploit dilemma that is central to the multi-armed bandit setting and derive algorithms whose computational complexity is independent of the number of auxiliary datasets, allowing us to scale to 100 × more auxiliary datasets than prior methods. We propose two algorithms - EXP3FLAD and UCB1-FLAD - and compare them with prior FLAD methods that either explore or exploit, finding that the combination of exploration and exploitation is crucial. Through extensive experimentation we find that our methods outperform all pre-existing FLAD methods by 4% and lead to the first 3 billion parameter language models that outperform the 175 billion parameter GPT-3. Overall, our work suggests that the discovery of better, more efficient mixing strategies for FLAD may provide a viable path towards substantially improving generalization in few-shot learning. All of our code is available at github.com/alon-albalak/FLAD.

1 Introduction

Few-shot learning is an attractive learning setting for many reasons: it promises efficiency in cost and time, and in some scenarios data is simply not available due to privacy concerns or the nature of the problem. However, few-shot learning is also a challenging setting that requires a delicate balance between learning the structure of the feature and label spaces while preventing overfitting to the limited training samples [1, 2, 3]. One approach to improving the generalizability of models in the few-shot setting is F ew-shot L earning with A uxiliary D ata (FLAD), where additional auxiliary datasets are used to improve generalization on the target few-shot task [4, 5, 6, 7].

However, FLAD methods introduce their own challenges, including increased algorithmic and computational complexity. Specifically, incorporating auxiliary data during training introduces a large space of design choices (e.g. how and when to train on auxiliary data). Manually designing the curriculum for training on large quantities of auxiliary data is not feasible due to the combinatorially large search space, and hand-picking which auxiliary data to use based on heuristics (e.g. from the same domain or task as the target few-shot dataset) can lead to sub-optimal results [8]. Delegating such choices to an algorithm can lead to better solutions, as demonstrated in the transfer learning [8, 9, 10], meta-learning [11, 12], multi-task learning [13, 14, 15, 16], and auxiliary learning literature [4, 17]. However, prior auxiliary learning algorithms often assume that only 1-3 related auxiliary datasets are

3.

Calculate auxiliary and target gradients, update model

Model

5.

Update policy

according to

1.

Select

dataset

according to

Policy

4.

Calculate

Reward

Environment

2.

Sample a batch

available and design algorithms whose computational complexity grows linearly (or worse) with the number of auxiliary datasets [18, 8], motivating the search for more efficient methods as the number of auxiliary datasets grows.

To overcome the challenges of prior works, we desire a FLAD algorithm that (1) makes no assumptions on available auxiliary data a-priori (in-domain, on-task, quality, quantity, etc.), (2) scales well with the number of auxiliary datasets, and (3) adds minimal memory and computational overhead. We design algorithms that satisfy our desiderata by drawing inspiration from the central problem in multi-armed bandit (MAB) settings: the exploration-exploitation trade-off [19, 20]. We relate the set of auxiliary datasets to the arms of a MAB and tailor the classic EXP3 [21] and UCB1 [22] algorithms to fit the FLAD framework by designing three efficient gradient-based reward signals. The combination of our MAB-based algorithms and efficient gradient-based rewards allows us to scale to 100 × more auxiliary datasets than previous methods. Figure 1 provides a basic illustration of how we formulate FLAD as a MAB problem.

To empirically validate our approaches, we focus on few-shot training of language models and utilize P3 [23], a readily available resource with hundreds of auxiliary language datasets. We evaluate our methods on the same held-out tasks as the T0 language model [16] and show that, when using the same collection of auxiliary datasets, our algorithms outperform a directly fine-tuned T0 by 5.6% (EXP3-FLAD) and 5.7% (UCB1-FLAD) absolute. Furthermore, incorporating all available datasets in P3 (i.e. not just those used to train T0) increases the improvement to 9.1% and 9.2%. Finally, we compare models trained with our methods against state-of-the-art few-shot methods, finding that our methods improve performance by >3%, even though one model utilizes a large collection of unlabeled target dataset samples. Furthermore, to the best of our knowledge, our methods lead to the first 3 billion parameter model that improves over 175B GPT-3 using few-shot in-context learning.

In summary, our main contributions are:

We connect FLAD to the MAB setting and focus on the exploration-exploitation trade-off by designing two algorithms, EXP3-FLAD and UCB1-FLAD along with three reward functions that are both simple and efficient (in space and computational complexity).

We empirically validate that our methods improve few-shot performance of pretrained language models and show that strategies that employ only exploration or exploitation lead to sub-optimal performance.

We perform case studies to better understand the dynamics of our reward functions and their interaction with the dynamics of large language model training.

2 Related work

A long history of works have found success when combining auxiliary data with target data [4, 24, 6, 25, 26, 5, 18, 7, 27, 28, 8]. Some works have explored the addition of auxiliary learning objectives to aid the learning of the target task [24, 26, 25, 5, 17]. More similar to our work are

methods that perform auxiliary learning by introducing additional data sources beyond the target data [4, 6, 18, 7, 27, 28, 8]. As opposed to the few-shot setting on which this work focuses, previous works have studied auxiliary learning in settings with large quantities of target data. For example, Chen et al. [18] and Verboven et al. [7] assume access to 10,000 labeled target samples, Ivison et al. [28] and Lin et al. [27] assume access to 1,000s of unlabeled target samples, and Du et al. [6] and Albalak et al. [8] assume access to 100s of labeled target samples. Additionally, many of the previous works that study auxiliary learning have only considered settings with 1-3 auxiliary datasets [6, 18, 7, 8]. For example, Verboven et al. [7] propose a task-weighting method that requires solving a system of equations that becomes underspecified with multiple auxiliary tasks, limiting their method to only a single auxiliary task. Furthermore, Chen et al. [18] experiment with 3 auxiliary tasks because their method requires learning a target-aware classifier for each source task, so the computation scales as O ( |A||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 a similar (but distinct) problem where the goal is to selectively utilize a subset of a single large dataset rather than selecting data from auxiliary datasets. Recent research on data selection has found that intelligent data selection can provide significant improvements to model performance [29, 30, 31, 32].

3 Multi-armed bandits for few-shot learning with auxiliary data

In this section, we first define the few-shot learning with auxiliary data ( FLAD ) setting. Then, we formulate FLAD as a multi-armed bandits ( MAB ) problem, shown in Figure 1. Next, we define reward functions that are efficient to compute and appropriate for FLAD. Finally, we describe our adaptations of two popular MAB algorithms: EXP3-FLAD and UCB1-FLAD.

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 D A where, for all a ∈ A , D a is an individual auxiliary dataset. Given a small quantity of data belonging to a target dataset D T , the goal of FLAD is to find parameters θ of a model f θ that achieve high performance on the unknown distribution underlying D T while utilizing only the available data, D T ∪ D A .

Formulating FLAD as MAB. In this work, we adopt the multi-armed bandit (MAB) setting by formulating FLAD as a Markov decision process [33] and defining a learner and environment, illustrated in Figure 1. The learner consists of a policy π defining a selection strategy over all D a ∈ D A . The environment consists of the target dataset D T , auxiliary datasets D A , and model f θ . In this formulation the learner interacts with the environment over N rounds. At each round t the learner selects one of the environment's |A| datasets D a ∈ D A . Next, the environment samples a batch { x , y } ∼ D a and calculates the gradient w.r.t. θ using a task-appropriate loss function as ∇ a = ∇ θ L ( f θ , x , y ) . Then, the environment computes the target gradient ∇ T = ∇ θ L ( f θ , D T ) , and updates model parameters w.r.t. ∇ T + ∇ a . Finally, the learner uses a gradient-based reward R a,t ( ∇ a , ∇ T ) to update its policy π . See Appendix A 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 we consider rewards that utilize information intrinsic to the model and the losses being optimized, not an external model or metric (e.g. accuracy or BLEU). In this work we propose three gradient-based reward functions inspired by previous works: gradient alignment [6, 24, 35], gradient magnitude similarity [36, 37], and their aggregation. Formally, at turn t let ∇ a be the gradient of the auxiliary batch and ∇ T be the target dataset gradient. Gradient alignment is defined as R GA a,t = ∇ a ·∇ T ∥∇ a ∥ 2 ∥∇ T ∥ 2 , i.e. the cosine similarity between the gradients of the sampled auxiliary dataset batch and the whole target dataset. Gradient magnitude similarity is defined as R GMS a,t = 2 ∥∇ a ∥ 2 ∥∇ T ∥ 2 ∥∇ a ∥ 2 2 + ∥∇ T ∥ 2 2 so that when the two gradients have equal magnitude, this value is equal to 1 and as the magnitudes differ the value goes to zero. In

addition to the individual reward functions, we also consider an aggregate reward. To ensure that the aggregate is not dominated by either individual reward, we normalize R GA ∈ [0 , 1] , the same range as R GMS and define the aggregate to be their sum: R AGG a,t = 1+ R GA a,t 2 + R GMS a,t . We provide further discussion on the design of reward functions in Section 6.

3.2 Adapting the EXP3 algorithm.

EXP3 Background We base our first algorithm, EXP3-FLAD, on the EXP3 algorithm [21] (' Exp onential-weight algorithm for Exp loration and Exp loitation'). EXP3 targets the adversarial MAB setting, which assumes that the reward-generating process is controlled by an adversary who is given access to the learner's policy π and determines the sequence of rewards, ( R a,t ) N t =1 , for each arm prior to play [38]. We consider the adversarial MAB formulation due to the highly non-convex loss landscape of deep neural networks and our use of stochastic gradient descent-based optimization methods. These factors imply that we cannot guarantee our rewards to be stationary, independent, or follow any particular distribution (e.g. Gaussian). Further details on adversarial MAB are included in Appendix A and in [21].

In EXP3-FLAD, the learner selects arms according to a Gibbs distribution based on the empirically determined importance-weighted rewards of arms [39]. To allow for exploration, we mix the Gibbs distribution with a uniform distribution [21]. Formally, let E t be the exploration rate at turn t and, recalling that K = |A| is the number of auxiliary datasets, then π defines the probability of selecting a given arm a ∈ A as the linear combination of Gibbs and uniform distributions π t ( a ) = (1 -K E t ) exp( E t -1 ˆ R a ) ∑ a ′ exp( E t -1 ˆ R a ′ ) + E t where ˆ R a,t is the importance weighted reward ˆ R a,t = ˆ R a,t -1 + R a,t π t -1 ( a ) . We want the learner to explore more in early training than in later stages, so we use a decaying exploration rate E t = min { 1 K , √ ln K K · t } as proposed by Seldin et al. [39]. The use of an importance-weighted estimated reward compensates the rewards of actions that are less likely to be chosen, guaranteeing that the expected estimated reward is equal to the actual reward for each action. EXP3-FLAD is designed to be nearly optimal in the worst case, but due to the exploration rate it will select 'bad' actions at a rate of E t . The exploration of EXP3-FLAD combined with importance-weighting allows the policy to handle non-stationary reward-generating processes.

3.3 Adapting the UCB1 algorithm.

UCB1background. While EXP3-FLAD is applicable in unconstrained settings with highly stochastic and non-stationary rewards, it can be outperformed by other algorithms in settings that are constrained. One such algorithm is the upper confidence bound (UCB1) algorithm [22], which was originally designed to be optimal for stationary, normally distributed reward functions. Nevertheless, variants of UCB1 have been demonstrated to be effective in a range of settings, such as those involving non-stationary, sub-Gaussian, or heavy-tailed distributions [40, 41]. The UCB1 algorithm and its variants assign each arm a value called the upper confidence bound based on Hoeffding's inequality [42] and are based on the principle of optimism in the face of uncertainty , meaning that with high probability the upper confidence bound assigned to each arm is an overestimate of the unknown mean reward.

In UCB1-FLAD, the learner greedily selects arms according to their upper confidence bound. UCB1 is originally designed for stationary reward-generating processes, so to accommodate non-stationarity we include an exponential moving average when estimating the mean reward for a given arm. Formally, let R a,t be the observed reward for arm a at turn t , then we calculate the estimated mean reward as ˆ R a = (1 -β ) ˆ R a + βR a,t where β is the smoothing factor. Then, we define the upper confidence bound to be UCB a,t = ˆ R a + √ 2 ln t n a . In the original MAB setting all interactions with the environment occur online, but FLAD is a unique situation where the learner can interact with the auxiliary data prior to training. To take advantage of this, rather than initializing estimated rewards with a single mini-batch, we initialize them with larger data quantities to improve the approximation of the true dataset gradients. This is done for each auxiliary dataset by calculating the gradient ∇ a = ∇ θ L ( f θ , x , y ) , where the number of samples in { x , y } can be significantly larger than a mini-batch, and can be up to the size of the full dataset. In practice, we use 1,000 examples which is computed in ∼ 2 minutes on a single GPU.

Algorithms The EXP3-FLAD and UCB1-FLAD algorithms are visualized in Figure 1. At each turn, both methods will first select an auxiliary dataset D a . EXP3-FLAD first computes the current exploration rate E t and samples D a according to the distribution defined by π t ( A ) , while UCB1FLAD greedily selects D a ∗ corresponding to the arm with largest upper confidence bound, a ∗ = arg max a ∈A UCB a,t . Next, for both methods, the environment samples a batch from the selected dataset, { x , y } ∼ D a , and calculates the gradient ∇ a = ∇ θ L ( f θ , 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 ∇ θ L ( f θ , D T ) and updates the model w.r.t. ∇ T + ∑ a ∇ a . Finally, EXP3-FLAD calculates the importance-weighted reward for each auxiliary batch using the observed rewards, while UCB1-FLAD calculates the smoothed estimated mean reward. Pseudocode is found in Appendix B.

4 Experimental setup

Models. For our experiments, we utilize encoder-decoder Transformer models from the T5 family of pre-trained language models [43]. Specifically, we experiment with LM-adapted T5 (T5-LM) and T0. The T5-LM model further trains the T5.1.1 model for 100,000 steps (corresponding to 100B tokens) from the C4 dataset [43] on the prefix language modeling objective [44]. The T0 model was initialized from T5-LM and further trained on a multitask mixture of prompted datasets as described by Sanh et al. [16]. We repeat each experiment with T5-LM XL (hereafter T5-XL ) and T0-3B as our base model. Both models use the same architecture with 2.85 billion parameters, and we used model checkpoints from Hugging Face Transformers [45]).

Target datasets. We obtain all datasets from Hugging Face Datasets 1 , and cast them to the text-totext format by applying prompt templates from the Public Pool of Prompts (P3) [23] that was used to train T0. To evaluate our few-shot methods, we utilize the same held-out datasets as T0, which cover four distinct tasks: sentence completion (COPA [46], HellaSwag [47], Story Cloze [48]), natural language inference (ANLI [49], CB [50], RTE [51]), coreference resolution (WSC [52], Winogrande [53]), and word sense disambiguation (WiC [54]). For each dataset, we randomly sample five few-shot splits from their training data, containing the same number of training examples as previous works, between 20 to 70 [55, 56]. We further divide each split into equal training and validation partitions for true few-shot learning [57](e.g. 10 train and 10 validation samples for HellaSwag). Only ANLI datasets have a publicly available test set, so for all other datasets we evaluate models on the original validation set (not utilized for few-shot training or validation).

Auxiliary datasets. We compare the performance of our methods using two sets of auxiliary data and never include any of the target datasets as part of auxiliary data. First, we use the collection of datasets used for multitask training of T0 (henceforth referred to as T0Mix), including 35 unique datasets covering question answering, sentiment analysis, topic classification, summarization, paraphrase detection and structure-to-text. Second, we utilize all datasets in P3 [23] (which forms a superset of T0Mix) and prevent data leakage by filtering out datasets that overlap with any target dataset, leading to 260 available datasets (list in Appendix H). For each auxiliary dataset, we use at most 10,000 of the dataset's examples.

Baseline methods. Wecompare our proposed methods with several FLAD and non-FLAD baselines. Target-Only (non-FLAD) directly fine-tunes the base model on the target dataset (i.e. without using auxiliary data). Explore-Only [8] is a commonly used FLAD method which simultaneously trains on auxiliary and target data by mixing auxiliary datasets equally. Originally called Multitask in [8], we call this Explore-Only because it is equivalent to continuously exploring auxiliary data and never exploiting knowledge of its relation to the target data. Exploit-Only computes gradient alignment prior to training (as in UCB1), and multitask-trains the model by mixing auxiliary datasets according to a Gibbs distribution over the alignments (similar to that in EXP3), resulting in an algorithm that exploits the relations determined prior to training, but never exploring. Both explore- and exploit-only mix target and auxiliary data with a ratio of M times the highest auxiliary sampling probability. For instance, explore-only with M = 5 and D A = T0Mix has a 1 / 35 probability to sample auxiliary dataset D a ∈ D A and a 5 / 35 probability for the target dataset. Loss-Scaling [6] is a FLAD method similar to EXP3 and UCB1; the main difference being that it scales auxiliary batch losses by their gradient alignment instead of modifying sampling probabilities. Du et al. [6] originally propose to

use gradient alignment ( Loss-Scaling ( GA ) ), but we also propose a version that scales losses by gradient magnitude similarity ( Loss-Scaling ( GMS ) ).

Training details. For the target-only baseline, we use learning rates in { 1 e -4 , 3 e -4 } . For all other methods, we always use a learning rate of 1 e -4 . For target-, explore-, and exploit-only baselines we use batch sizes in { 32 , 128 } . For loss-scaling, EXP3-FLAD, and UCB1-FLAD we use mini-batches of 8 samples and let G be in { 4 , 16 } to match the batch size of all methods. For explore- and exploit-only, we use a target dataset mixing ratio of M ∈ { 1 , 5 , 10 } . For all experiments we use the Adafactor optimizer [58] and validation-based early stopping for model checkpoint selection. In preliminary experiments we consider rewards using gradients from various model partitions: the full model, encoder-only, decoder-only, and the weights of the output vocabulary matrix (language modeling head). We find that using the parameters from the language modeling head provides the best performance and contains only 2.3% of the full model parameters, significantly reducing memory consumption. For UCB1-FLAD we found the smoothing factor β = 0 . 9 to work well in preliminary experiments and initialize auxiliary dataset gradient alignment using 1,000 samples from each auxiliary dataset. Additional implementation details can be found in Appendix C

Experiment procedure. The FLAD experiment process involves training a model that is specialized for each target dataset. For each proposed method and baseline, we train and evaluate a model on each of the 11 target datasets. We repeat training and evaluation on 5 random seeds and include the aggregated results in Table 1. Each cell shows the accuracy averaged across all 55 (11 target datasets, 5 random seeds) experiments. This experimental process is performed for each training method on both models and auxiliary datasets. We include the non-aggregated results in Appendix D.

5 Findings and analysis

In Table 1 we compare the empirical results of our MAB-based methods (EXP3-FLAD and UCB1FLAD) and corresponding baselines on 11 target datasets (expanded results in Appendix D. For each base model and auxiliary data combination (each column) EXP3-FLAD and UCB1-FLAD outperform all the baselines. In fact, we find that for every single task our methods always perform equal to or better than the baselines. This demonstrates that our MAB-based methods provide a strong improvement in few-shot generalization over previous FLAD methods. For a fair comparison where each method utilizes equal data, we compare the performance of Target-Only using T0 and T0Mix (56.44) against the proposed FLAD methods and baselines using T5 and T0Mix (left column). From this comparison it becomes clear that Loss-Scaling actually does worse than multitask training followed by direct fine-tuning by 0.5-3.2%. However, we do find that the remaining FLAD methods lead to improvements (between 2.7-5.6% absolute improvement). We find small performance

Natural Language Inference

60

62

64

66

68

70

35

40

45

50

ANLI-R1

30.0

32.5

35.0

37.5

40.0

42.5

ANLI-R2

35.0

37.5

40.0

42.5

45.0

ANLI-R3

70

80

90

CB

60

70

80

90

RTE

60

70

80

WSC

50

60

70

80

Winogrande

70

80

90

100

COPA

75

80

85

90

95

100

StoryCloze

40

50

60

70

80

45

50

55

60

WiC

Average Accuracy

Coference Resolution

Sentence Completion

Word Sense

EXP3-FLAD (T5)

EXP3-FLAD (T0)

UCB1-FLAD (T5)

UCB1-FLAD (T0)

T-Few

DEFT-Few

GPT-3 (175B)

differences between EXP3-FLAD and UCB1-FLAD across the three reward functions. In general, R AGG leads to the best performance, but we perform a two-sided Wilcoxon rank-sum test to check for significance between average scores and find that the other rewards frequently have no significant difference ( p > 0 . 05 ).

The importance of prioritized sampling. Loss-Scaling was originally proposed for use with only a single auxiliary dataset and it was unclear, a priori, how it would cope with larger quantities. Additionally, Du et al. [6] purposefully choose an auxiliary dataset that is related to the target, while in our setting we make no such assumptions. We find that our methods outperform Loss-Scaling methods by 6.3% on average. In Figure 3 (and Figure 5 in Appendix E) we show that, over the course of training, the value of gradient alignments and gradient magnitude similarities for most datasets will converge to 0, leading to very small gradient updates for Loss-Scaling. More importantly, the auxiliary data that is relevant to the target task is seen less frequently for Loss-Scaling than our MAB-based methods. This can be seen by comparing the difference in performance of Loss-Scaling methods when using less (T0Mix) vs. more (P3) auxiliary data. We find that, at best, Loss-Scaling ( GA ) improves 2% when using T5 and, at worst, only 0.2% for Loss-Scaling ( GMS ) with T0. This is compared with the notable improvements of EXP3-FLAD and UCB1-FLAD of 2.6-4% when considering the same data increase from T0Mix to P3.

The importance of exploration and exploitation. Interestingly, we expected that Exploit-Only would outperform the Explore-Only method because it utilizes relational information between the target and auxiliary tasks, but find no statistical difference between the methods (two-sided Wilcoxon rank-sum test gives p > 0 . 05 ). Furthermore, when comparing the ability to leverage additional auxiliary data (i.e. going from T0Mix to all of P3), we find that the improvement for Explore- and Exploit-Only methods is minimal with only 0.7-2% improvement. On the other hand, EXP3-FLAD and UCB1-FLAD show a notable improvement of 2.6-4%, emphasizing the importance of both exploration and exploitation, particularly when dealing with large collections of auxiliary data.

FLAD provides improved generalization over non-FLAD methods. Next, we compare the performance of our best models trained on P3 using R AGG with state-of-the-art few-shot methods: T-Few, DEFT-Few, and GPT-3. T-Few [56] is a variant of the T0-3B model that multi-task pretrains parameter-efficient (IA) 3 modules followed by target-only fine-tuning of the (IA) 3 modules. DEFT-Few [28] is a variant of the T5-XL model that uses retrieved auxiliary data for multi-task training. It first trains a T5-XL model on the 500 nearest neighbor samples from P3 using 1000 unlabeled target dataset samples, and then performs few-shot target-only fine-tuning with the (IA) 3 modules from Liu et al. [56]. Finally, we also compare against the 175 billion parameter variant of

Gradient Alignment

Gradient Magnitude Similarity

Step 0

Step 300

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.000.0

0.2

0.4

0.6

0.8

1.0

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0.0

0.2

0.4

0.6

0.8

1.0

GPT-3 [55], which utilizes in-context learning. We find that, on average, models trained using our FLAD-based methods outperform all other methods and, to the best of our knowledge, our methods lead to the first 3 billion parameter model that outperforms GPT-3 on this dataset mixture (previous smallest models have 11 billion parameters), despite using 62 . 5 times fewer parameters than GPT-3. Additionally, we find that our FLAD-based methods provide robust performance across datasets, achieving the best or second-best performance on 8 / 11 datasets, and never performing worst. The use of task-specific modules lead T-Few and DEFT-Few to significant improvements over target-only fine-tuning, preventing the models from ending up in poor local minima. However, these results demonstrate that with the same data, simultaneously fine-tuning with auxiliary and target data leads to improved few-shot generalization, providing a complementary means of improving performance.

Investigating the Reward-Generating Processes. In Section 3.2, we mention that due to the highly non-convex loss landscape and the use of stochastic gradient descent-based optimization techniques, we cannot ensure that our reward generating process is stationary, independent across auxiliary datasets, or follows a normal distribution. To gain a deeper understanding of our reward-generating processes, we examine the distribution of each reward using 5,000 samples from all 35 auxiliary datasets of T0Mix and 32 samples from a few-shot target dataset, WSC [52]. The resulting histograms at every 100 steps can be found in Appendix E, and Figure 3 shows an abbreviated version. The left side of Figure 3 demonstrates that for R GA , almost every dataset yields a Gaussian reward distribution, with a few multi-modal distributions. Notably, WikiBio [59] (dark orange) exhibits peaks at 0.25 and -0.75. Interestingly, R GA results in polarized rewards across datasets, with minimal distribution density between -0.75 and 0.25. In contrast, the right side of Figure 3 displays more non-Gaussian distributions for R GMS , as well as flatter distributions compared to R GA . Remarkably, we observe that R GA produces more stationary reward distributions, as the distribution for almost every dataset (30/35) converges rapidly towards 0 after only 100 steps. Although most distributions for R GMS also converge towards 0, the convergence occurs at a slower pace, taking nearly 500 steps.

Probing the training dynamics. To better understand the training dynamics of our proposed methods, we perform a case study on T5-XL with T0Mix and R GA and find two datasets where either algorithm improves significantly over the other (full details and figures in Appendix F). First, we study RTE, where UCB1-FLAD outperforms EXP3-FLAD. We calculate the empirical distribution of samples seen from each auxiliary dataset and find that EXP3-FLAD samples nearly uniformly from all datasets while UCB1-FLAD forms a bimodal sampling distribution with peaks at 2.5% and 3.25% (30% relative difference). The uniformity of the EXP3-FLAD distribution is counterintuitive, as we do find that it achieves separation between auxiliary tasks in the cumulative estimated reward (as shown in Figure 7), but this does not lead to separation in the sampling probability space. Additionally we find that even on COPA, where EXP3-FLAD outperforms UCB1-FLAD, EXP3-FLAD still achieves

good separation between cumulative estimated rewards, but has a unimodal sampling distribution, while UCB1-FLAD does not have as clear of a bimodal distribution as in RTE. The difference in empirical sampling distributions is likely due to the difference between the greedy policy of UCB1FLAD and the stochastic policy of EXP3-FLAD. Empirically, we find that EXP3-FLAD very rarely assigns an auxiliary dataset a probability < 1 %, leading to many 'bad' batches over the course of thousands of turns. On the other hand, the optimistic policy of UCB1-FLAD spends much less time exploring and will sample 'bad' batches much less frequently.

The effect of scaling |A| To assess the scalability of our proposed methods, we conduct a study by varying the size of |A| ∈ { 35 , 75 , 125 , 175 , 225 , 260 } . For each value of |A| , we consistently include the 35 datasets from T0Mix 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 T0 base model and R AGG reward. The experiment is repeated with three random seeds. Figure 4 shows the mean across the 11 target datasets, along with the standard deviation between seeds.

We find that both EXP3-FLAD and UCB1FLAD experience a sharp increase from |A| = 35 to 75 . In addition, we observe improvements up to the maximum value of |A| = 260 , ultimately improving accuracy by 2.54 for EXP3FLAD and 3.12 for UCB1-FLAD when transitioning from 35 to 75 datasets, with further increases of 1.54 for EXP3-FLAD and 0.47 for UCB1-FLAD when increasing |A| from 75 to 260.

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, R AGG , we examine four combinations of rewards: low R GA and R GMS , high R GA but low R GMS , low R GA but high R GMS , and high R GA and R GMS . When both rewards are high, we can assume that the auxiliary gradient is useful. However, when one reward is high and the other is low, it is difficult to draw conclusions as a high R GA on its own means the auxiliary gradient will update weights in the right direction, but low R GMS can mean that we significantly overshoot or undershoot the target, where overshooting can be much more detrimental than undershooting. If both R GA and R GMS are small, we know the auxiliary gradient leads us away from the target solution space, but we don't know if its magnitude is much larger or smaller than the target. At the beginning of training, we can't know if the target or auxiliary gradient has larger magnitude, but as training progresses, it becomes significantly more likely that the auxiliary gradient is greater than the target. Thus, when both R GA and R GMS are low, we are likely to be pulled far from our current solution.

This work uses training set-based rewards, but validation set-based rewards are also possible. One downside of validation-based rewards is they calculate validation score frequently, which increases computational complexity. Additionally, we focus on the few-shot setting and use validation-based early stopping. If we use a validation-based reward, then to prevent overfitting we will need to further split the data into 3 partitions: train, early-stopping validation, and reward-validation.

Choice of baselines. With respect to the number of auxiliary datasets |A| and target datasets |T | , our methods and the baselines we compare against have a 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 a baseline whose computation grows linearly w.r.t. |A| ,

50

100

150

200

250

# of Auxiliary Datasets

0.62

0.63

0.64

0.65

0.66

0.67

EXP3-FLAD

UCB1-FLAD

O ( |A||T | ) (e.g. [18]), these experiments would not be feasible without a large amount of hardware: Training such a model with T0Mix would take over 200 GPU-hours (over 8 GPU-days) for a single target dataset , and over 1500 GPU-hours ( over 2 GPU-months ) when using all of P3.

Why we don't include theoretical guarantees. The design of MAB algorithms generally comes with theoretical proofs of regret bounds, but in this work we omit such analysis. Although we could make guarantees on the regret bounds of our algorithms, they would not be meaningful because they would be with respect to the rewards, not the accuracy on a held-out dataset (which is the quantity we actually care about).

How does FLAD relate to few-shot learning and multitask learning? Both few-shot learning and FLAD are concerned with optimizing model performance on a single target task with a limited number of examples from the target task. In few-shot learning, the model is given only the target task data D T and there is no auxiliary data. Effectively, D A is the empty set for few-shot learning. In contrast, for the FLAD setting |D A | > 1 . Based on the findings from this study, we highly recommend that practitioners utilize auxiliary data when it is available.

Multitask learning is concerned with optimizing a model for performance on multiple target datasets simultaneously. This is in direct opposition with the FLAD methods presented here, which aim to optimize a model for a single target task. However, it is possible to extend our MAB-based methods to optimize for multiple target tasks simultaneously by aggregating multiple rewards. We believe this would make for an interesting future study.

Limitations. One of the implicit assumptions in the FLAD setting (made by this work and all prior works) is that there is at least some auxiliary data that will be useful for the target task. However, one of the main distinctions of our methods from prior works in the FLAD setting is that prior works make a strong assumption that all auxiliary data are useful, and thus appropriate auxiliary datasets must be hand-picked by humans. On the other hand, our methods allow for only a small portion of the auxiliary data to be useful - our proposed algorithm explores to find useful auxiliary datasets and then exploits them.

7 Conclusion

Recall the desiderata for our algorithm, expressed in the introduction: our algorithm should (1) make no assumptions on the available auxiliary data a-priori, (2) scale well with the number of auxiliary datasets, and (3) add minimal memory and computational overhead. (1) When designing our algorithm, we purposefully formulate the problem as a multi-armed bandit. MAB algorithms, in general, make no assumptions on the quality of rewards and, in particular, EXP3 even assumes that the auxiliary datasets will play an adversarial role when returning rewards. (2) As previously mentioned, our algorithms have a single-turn computational complexity that is independent of the number of auxiliary datasets. (3) Finally, our method adds minimal computational overhead beyond usual training computations. Every gradient that we utilize for our reward functions are also used to update the model, adding no additional computations. The only computational overhead is to compute gradient alignment (three vector dot products, two scalar square roots, and two scalar multiplications) or magnitude similarity (four vector dot products, two scalar square roots, three scalar multiplications, and one scalar addition). Additionally, our method adds a small amount of memory overhead, used to store gradients between model updates. Our rewards consider only the gradient w.r.t the language modelling head and, in practice, require 0.25Gb per auxiliary gradient to store, slightly increasing the space complexity above standard fine-tuning.

The methods proposed in this work demonstrate the effectiveness of simultaneous training on auxiliary and target datasets in few-shot settings, continuously updating beliefs by exploring and exploiting auxiliary data, and framing FLAD as a MAB problem. We further showed that by satisfying our desiderata, we are able to scale up FLAD to hundreds of auxiliary datasets and outperform traditional few-shot fine-tuning and in-context learning methods. While the presented algorithms satisfy our desiderata, the findings from this study can inform future work to further improve upon these methods in a number of ways, such as improving the reward function and reducing the space complexity.

References

Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. In International Conference on Learning Representations , 2017.

Yaqing Wang, Quanming Yao, James T. Kwok, and Lionel M. Ni. Generalizing from a few examples: A survey on few-shot learning. ACM Comput. Surv. , 53(3), jun 2020.

Archit Parnami and Minwoo Lee. Learning from few examples: A summary of approaches to few-shot learning, 2022.

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.

Reza Esfandiarpoor, Amy Pu, Mohsen Hajabdollahi, and Stephen H. Bach. Extended few-shot learning: Exploiting existing resources for novel tasks, 2020.

Yunshu Du, Wojciech M. Czarnecki, Siddhant M. Jayakumar, Mehrdad Farajtabar, Razvan Pascanu, and Balaji Lakshminarayanan. Adapting auxiliary losses using gradient similarity, 2020.

Sam Verboven, Muhammad Hafeez Chaudhary, Jeroen Berrevoets, Vincent Ginis, and Wouter Verbeke. Hydalearn. Applied Intelligence , Jul 2022.

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.

Tu Vu, Tong Wang, Tsendsuren Munkhdalai, Alessandro Sordoni, Adam Trischler, Andrew Mattarella-Micke, Subhransu Maji, and Mohit Iyyer. Exploring and predicting transferability across NLP tasks. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 7882-7926, Online, November 2020. Association for Computational Linguistics.

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.

Sebastian Thrun and Lorien Pratt, editors. Learning to Learn . Kluwer Academic Publishers, USA, 1998.

Trapit Bansal, Rishikesh Jha, Tsendsuren Munkhdalai, and Andrew McCallum. Self-supervised meta-learning for few-shot natural language classification tasks. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 522-534, Online, November 2020. Association for Computational Linguistics.

Sen Wu, Hongyang R. Zhang, and Christopher Ré. Understanding and improving information transfer in multi-task learning. In International Conference on Learning Representations , 2020.

Armen Aghajanyan, Anchit Gupta, Akshat Shrivastava, Xilun Chen, Luke Zettlemoyer, and Sonal Gupta. Muppet: Massive multi-task representations with pre-finetuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 5799-5811, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.

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 Representations , 2022.

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 Szczechla, 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.

Lucio M. Dery, Paul Michel, Mikhail Khodak, Graham Neubig, and Ameet Talwalkar. AANG : Automating auxiliary learning. In The Eleventh International Conference on Learning Representations , 2023.

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.

William G Macready and David H Wolpert. Bandit problems and the exploration/exploitation tradeoff. IEEE Transactions on evolutionary computation , 2(1):2-22, 1998.

Alex Simpkins, Raymond De Callafon, and Emanuel Todorov. Optimal trade-off between exploration and exploitation. In 2008 American Control Conference , pages 33-38. IEEE, 2008.

Peter Auer, Nicolo Cesa-Bianchi, Yoav Freund, and Robert E Schapire. The nonstochastic multiarmed bandit problem. SIAM journal on computing , 32(1):48-77, 2002.

Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning , 47(2):235-256, 2002.

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.

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.

Lucio M. Dery, Paul Michel, Ameet S. Talwalkar, and Graham Neubig. Should we be pretraining? an argument for end-task aware training as an alternative. ArXiv , abs/2109.07437, 2021.

Aviv Navon, Idan Achituve, Haggai Maron, Gal Chechik, and Ethan Fetaya. Auxiliary learning by implicit differentiation. In International Conference on Learning Representations , 2021.

Bill Yuchen Lin, Kangmin Tan, Chris Miller, Beiwen Tian, and Xiang Ren. Unsupervised cross-task generalization via retrieval augmentation. ArXiv , abs/2204.07937, 2022.

Hamish Ivison, Noah A. Smith, Hannaneh Hajishirzi, and Pradeep Dasigi. Data-efficient finetuning using cross-task nearest neighbors, 2022.

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. In International Conference on Machine Learning , 2021.

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.

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.

Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S. Morcos. Semdedup: Data-efficient learning at web-scale through semantic deduplication, 2023.

RICHARDBELLMAN. A markovian decision process. Journal of Mathematics and Mechanics , 6(5):679-684, 1957.

Tor Lattimore and Csaba Szepesvári. Bandit Algorithms . Cambridge University Press, 2020.

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.

Wufeng Xue, Lei Zhang, Xuanqin Mou, and Alan C. Bovik. Gradient magnitude similarity deviation: A highly efficient perceptual image quality index. IEEE Transactions on Image Processing , 23(2):684-695, 2014.

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.

Peter Auer, Nicolo Cesa-Bianchi, Yoav Freund, and Robert E Schapire. Gambling in a rigged casino: The adversarial multi-armed bandit problem. In Proceedings of IEEE 36th annual foundations of computer science , pages 322-331. IEEE, 1995.

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-01 Jul 2013. PMLR.

Aurélien 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.

Lai Wei and Vaibhav Srivastava. Nonstationary stochastic multiarmed bandits: Ucb policies and minimax regret, 2021.

Wassily Hoeffding. Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association , 58(301):13-30, 1963.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine 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.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 3045-3059, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.

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 Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of-the-art 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.

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 1: 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-398, Montréal, Canada, 7-8 June 2012. Association for Computational Linguistics.

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 Association for Computational Linguistics , pages 4791-4800, Florence, Italy, July 2019. Association for Computational Linguistics.

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.

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 4885-4901, Online, July 2020. Association for Computational Linguistics.

Marie-Catherine de Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In Sinn und Bedeutung , 2019.

Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Joaquin Quiñonero-Candela, Ido Dagan, Bernardo 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.

Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning , 2012.

Keisuke Sakaguchi, Ronan Bras, Chandra Bhagavatula, and Choi Yejin. Winogrande: An adversarial winograd schema challenge at scale. Proceedings of the AAAI Conference on Artificial Intelligence , 34:8732-8740, 04 2020.

Mohammad Taher Pilehvar and Jose Camacho-Collados. 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 1 (Long and Short Papers) , pages 1267-1273, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

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

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.

Haokun Liu, Derek Tam, Muqeeth Mohammed, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

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.

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 Learning Research , pages 4596-4604. PMLR, 10-15 Jul 2018.

Rémi Lebret, David Grangier, and Michael Auli. Generating text from structured data with application to the biography domain. CoRR , abs/1603.07771, 2016.

A Multi-armed bandits

The Multi-Armed Bandit ( MAB ) setting is a problem from machine learning where a learner interacts with an environment over N rounds by following a policy π . At each round t the learner chooses one of the environment's K arms, a ∈ A where K = |A| , after which the environment provides a reward R t . 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 = ∑ N t =1 R t . 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 [34].

Adversarial MAB. The adversarial MAB setting assumes that the reward-generating 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 to the learner's policy π and determines the sequence of rewards, ( R a,t ) N t =1 , for each arm prior to play [38]. At each turn π determines a distribution over actions, p ( A ) , and an action is sampled from the distribution, a ∼ p ( A ) . See Lattimore & Szepesvári [34] for further details.

The EXP3 algorithm. The EXP3 algorithm (' Exp onential-weight algorithm for Exp loration and Exp loitation') targets the adversarial multi-armed bandit problem by choosing arms according to a Gibbs distribution based on the empirically determined importance-weighted rewards of arms [21]. To allow for exploration, EXP3 mixes the Gibbs distribution with a uniform distribution.

Formally, let the exploration rate be γ ∈ (0 , 1] . At round t , π defines the probability of selecting a given arm, a ∈ A , as a linear combination of Gibbs and uniform distributions

where the importance weighted reward ˆ R a,t is calculated as

̸

and R a,t denotes the observed reward. All unplayed arms, a ′ = a have unchanged importance weighted rewards; ˆ R a ′ ,t = ˆ R a ′ ,t -1 .

Algorithmically, EXP3 takes the following steps at each round: First, calculate the sampling distribution p t and sample an arm from the distribution. Then a reward R a,t is observed and the algorithm updates the importance weighted reward ˆ R a,t for the played arm.

Informally, the use of an importance-weighted estimated reward compensates the rewards of actions that are less likely to be chosen, guaranteeing that the expected estimated reward is equal to the actual reward for each action. EXP3 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-weighting allows the policy to handle non-stationary reward-generating processes.

The UCB1 algorithm. While the adversarial setting makes almost no assumptions about the reward-generating 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 reward-generating processes are stationary Gaussian distributions. A common policy used to solve this MAB setting is the Upper Confidence Bound (UCB1) algorithm, which assigns each arm a value called the upper confidence bound based on Hoeffding's inequality [22]. The UCB1 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 a times be ˆ R a and the true mean reward be R a , then

derived from Hoeffding's inequality (following equation 7.1 of Lattimore & Szepesvári [34]), where δ is the confidence level that quantifies the degree of certainty in the arm. In this work we let δ = 1 /t where t is the number of rounds played, shrinking the confidence bound over rounds. Thus, we define the upper confidence bound for arm a at turn t as

Algorithmically, UCB1 takes the following steps at each round. First, the UCB1 policy plays the arm with largest upper confidence bound, a ∗ = arg max a ∈A UCB a,t . Next, a reward R a ∗ ,t is observed and the algorithm updates ˆ R a ∗ (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. n a is small because the arm is not well explored.

B Pseudo-code

We include here pseudo-code for our 2 proposed algorithms. Algorithm 1 contains the pseudo-code for EXP3-FLAD, and Algorithm 2 contains the pseudo-code for UCB1-FLAD.

̸

̸

C Training details

We train all models (FLAD and non-FLAD) on 40Gb A100s.

For all experiments, we use validation-based early stopping, and train for a 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, β , in UCB1-FLAD we ran preliminary experiments using values of { 0 . 99 , 0 . 9 , 0 . 75 , 0 . 5 } and found 0.9 to work well across datasets. All reported scores use β = 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 was ∼ 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 UCB1-FLAD are found on the next page.

Dataset

Anli-r3

CB

COPA

RTE

WiC

Winogrande

WSC

Average

Direct Fine-Tuning

37.6

36.2

35.0

83.2

53.8

51.0

54.2

75.9

51.6

49.6

53.1

52.8

T5-3B

T0Mix

Loss-Scaling (

GA

)

35.7

36.4

35.3

82.5

58.0

51.8

59.0

79.6

49.8

50.4

46.9

53.2

Loss-Scaling (

GMS

)

37.8

37.6

36.0

80.0

76.4

52.6

55.3

85.7

50.6

52.0

51.7

56.0

Exploration-Only

38.1

40.3

36.7

88.6

85.6

51.2

67.6

88.8

51.0

55.5

47.7

59.2

Exploitation-Only

38.8

40.5

38.0

86.1

86.0

51.1

69.4

89.5

52.8

59.2

46.3

59.8

EXP3-FLAD (

R

GA

)

40.6

39.9

36.9

86.1

89.8

52.0

76.7

90.8

50.5

60.3

52.9

61.5

UCB1-FLAD (

R

GA

)

41.8

39.0

38.0

85.4

87.0

52.0

79.1

91.4

49.7

62.7

56.2

62.0

EXP3-FLAD (

R

GMS

)

42.0

40.2

36.6

87.1

87.2

52.4

77.5

90.9

51.1

61.9

51.9

61.7

UCB1-FLAD (

R

GMS

)

41.3

39.7

38.0

82.5

89.8

51.0

76.6

90.5

51.0

62.0

56.0

61.7

EXP3-FLAD (

R

AGG

)

38.6

39.8

39.1

86.8

91.2

51.2

78.8

90.4

50.7

63.0

52.9

62.0

UCB1-FLAD (

R

AGG

)

42.0

41.0

36.6

88.2

86.8

51.0

77.3

90.3

51.1

63.3

55.4

62.1

P3

Loss-Scaling (

GA

)

38.7

39.5

34.8

80.7

64.4

52.7

62.9

80.1

50.3

51.9

51.2

55.2

Loss-Scaling (

GMS

)

39.2

38.7

36.4

85.0

67.8

51.9

62.4

84.8

50.3

51.8

52.1

56.4

Exploration-Only

40.1

37.7

36.0

85.4

83.6

52.1

77.3

89.1

51.5

57.2

57.1

60.6

Exploitation-Only

40.4

37.2

37.3

87.1

84.4

51.0

78.6

90.3

51.3

56.2

51.5

60.5

EXP3-FLAD (

R

GA

)

46.9

38.8

40.2

89.6

88.0

51.5

76.9

91.2

53.4

66.2

61.9

64.1

UCB1-FLAD (

R

GA

)

49.1

38.8

40.1

88.6

88.2

51.6

83.7

90.2

54.3

68.0

68.3

65.5

EXP3-FLAD (

R

GMS

)

46.2

40.6

39.4

88.9

90.4

51.6

85.1

91.3

54.4

65.8

67.5

65.6

UCB1-FLAD (

R

GMS

)

48.1

40.1

39.1

87.5

89.4

52.0

83.7

89.4

51.7

65.8

70.6

65.2

EXP3-FLAD (

R

AGG

)

47.6

40.6

40.6

90.0

90.6

51.4

84.5

91.0

53.2

66.7

64.0

65.5

UCB1-FLAD (

R

AGG

)

47.1

39.0

41.2

86.8

90.4

51.5

85.5

91.1

52.7

66.3

70.6

65.6

Direct Fine-Tuning

40.9

39.1

37.1

79.6

66.4

43.5

67.1

83.2

52.5

54.6

56.7

56.4

T0-3B

T0Mix

Loss-Scaling (

GA

)

41.3

40.0

36.9

81.8

78.0

51.2

76.5

86.9

50.7

54.7

56.2

59.5

Loss-Scaling (

GMS

)

40.5

40.5

37.8

81.1

79.0

52.0

77.0

88.8

52.7

55.0

60.8

60.5

Exploration-Only

44.4

40.3

37.0

82.5

85.6

47.9

77.6

90.1

52.1

58.6

56.9

61.2

Exploitation-Only

42.5

39.3

37.2

84.3

82.8

48.1

79.7

88.8

52.8

57.8

56.3

60.9

EXP3-FLAD (

R

GA

)

46.2

41.5

37.7

83.9

87.6

49.4

80.0

90.1

52.6

63.4

59.0

62.9

UCB1-FLAD (

R

GA

)

43.7

40.8

37.6

86.1

85.4

48.6

80.5

91.3

53.4

63.5

61.0

62.9

EXP3-FLAD (

R

GMS

)

43.4

41.1

38.2

84.6

86.6

49.1

81.0

90.6

53.0

63.1

59.8

62.8

UCB1-FLAD (

R

GMS

)

43.2

41.2

38.7

86.4

86.6

48.4

82.8

91.4

52.2

61.0

59.4

62.8

EXP3-FLAD (

R

AGG

)

43.8

41.6

38.0

83.9

87.8

48.9

81.9

90.7

52.5

62.3

59.8

62.8

UCB1-FLAD (

R

AGG

)

44.0

41.6

38.3

85.4

87.4

48.6

81.1

90.6

53.0

63.1

59.2

62.9

P3

Loss-Scaling (

GA

)

44.0

40.4

38.9

86.4

77.6

51.0

75.1

86.8

51.7

55.6

59.8

60.7

Loss-Scaling (

GMS

)

43.8

38.6

39.3

82.5

79.2

50.6

80.6

89.1

51.6

56.6

56.0

60.7

Exploration-Only

45.4

40.3

38.0

82.5

87.8

50.6

82.2

88.8

52.4

61.8

60.6

62.8

Exploitation-Only

45.5

40.0

38.8

87.5

82.2

49.9

79.6

90.9

52.2

60.1

64.8

62.9

EXP3-FLAD (

R

GA

)

50.4

40.0

41.2

87.9

88.4

49.7

86.1

91.6

52.8

67.5

70.4

66.0

UCB1-FLAD (

R

GA

)

48.2

41.8

41.2

90.0

86.6

50.0

86.1

91.5

53.6

65.6

74.6

66.3

EXP3-FLAD (

R

GMS

)

49.5

40.8

39.5

87.1

89.2

49.4

85.8

91.4

53.9

65.4

68.7

65.5

UCB1-FLAD (

R

GMS

)

48.2

41.8

40.5

89.6

88.0

49.6

83.2

91.6

52.6

66.1

74.6

66.0

EXP3-FLAD (

R

AGG

)

51.1

40.3

39.9

89.6

91.4

49.0

86.5

91.6

52.6

66.4

76.7

66.8

UCB1-FLAD (

R

AGG

)

49.8

39.9

40.8

86.8

88.4

49.6

84.7

91.0

53.2

68.0

76.9

66.3

E Probing the reward generating processes.

Gradient Alignment

Gradient Magnitude Similarity

Step 0

Step 300

Step 200

Step 100

Step 400

Step 500

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

0.0

0.2

0.4

0.6

80

10

1.00

0.75

05'0-

0.25

0.00

0.25

0.50

0.75

1.00

0.0

0.2

0.4

0.6

0.8

10

-1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

100 0.0

0.2

0.4

0.6

0.8

10

-1.00

-0.75

-0.50

0.25

0.00

0.25

0.50

0.75

100

0.0

0.2

0.4

0.6

0.8

1.0

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.000.0

0.2

0.4

0.6

0.8

1.0

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0.0

0.2

0.4

0.6

0.8

10

F EXP3-FLAD and UCB1-FLAD training dynamics

The following 4 pages include a case study on the training dynamics of EXP3-FLAD and UCB1FLAD when training T5-XL using T0Mix as the auxiliary data. First, we find datasets where EXP3-FLAD and UCB1-FLAD improve significantly over the baseline FLAD methods, but also where either EXP3-FLAD or UCB1-FLAD clearly outperforms the other. The two datasets that fulfill our interests are RTE and COPA.

We find that UCB1-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 UCB1-FLAD on COPA, and show their respective training dynamics in Figure 8 (UCB1) 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 a fraction of the total samples seen. For UCB1-FLAD figures, we include charts of the upper confidence index, estimated gradient alignment, and the empirical sampling distribution.

0.0

0.5

1.0

1.5

2.0

2.5

Upper Confidence Bound

amazon_polarity

imdb

kilt_tasks/hotpotqa

qasc

wiki_qa

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

Estimated

Gradient Alignment

0

1600

3200

4800

6400

8000

9600

Number of Auxiliary Batches Seen

0.02

0.03

0.04

0.05

Cumulative Empirical

Sampling Distribution

12000

10000

8000

6000

4000

2000

0

Cumulative

Estimated Reward

amazon_polarity

imdb

kilt_tasks/hotpotqa

qasc

wiki_qa

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

Gradient Alignment

0.00

0.01

0.02

0.03

0.04

0.05

Instantaneous

Sampling Distribution

0

3200

6400

9600

12800

Number of Auxiliary Batches Seen

0.01

0.02

0.03

0.04

0.05

Cumulative Empirical

Sampling Distribution

0.0

0.5

1.0

1.5

2.0

2.5

Upper Confidence Bound

cnn_dailymail/3.0.0

multi_news

qasc

quartz

trec

0.8

0.6

0.4

0.2

0.0

0.2

0.4

0.6

0.8

Estimated

Gradient Alignment

0

3200

6400

9600

12800

16000

Number of Auxiliary Batches Seen

0.015

0.020

0.025

0.030

0.035

0.040

0.045

0.050

Cumulative Empirical

Sampling Distribution

750

500

250

0

250

500

750

Cumulative

Estimated Reward

cnn_dailymail/3.0.0

multi_news

qasc

quartz

trec

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

Gradient Alignment

0.01

0.02

0.03

0.04

0.05

Instantaneous

Sampling Distribution

0

3200

6400

9600

12800

16000

19200

Number of Auxiliary Batches Seen

0.01

0.02

0.03

0.04

0.05

Cumulative Empirical

Sampling Distribution

Table 3: Number of training iterations for T0-3B to converge using a training method (column) and a set of auxiliary datasets (row). The number of iterations to 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-convergence

As we have described in this work, the computational complexity for a 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 it is not possible to 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| on the time-to-model-convergence, we perform a study on the number of iterations to convergence for different FLAD algorithms. Table 3 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 Exploreand 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 T0Mix, at which point it 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 Datasets 2 .

Zaid/quac_expanded, acronym_identification, ade_corpus_v2/Ade_corpus_v2_classification, ade_corpus_v2/Ade_corpus_v2_drug_ade_relation, ade_corpus_v2/Ade_corpus_v2_drug_dosage_relation, adversarial_qa/adversarialQA, adversarial_qa/dbert, adversarial_qa/dbidaf, adversarial_qa/droberta, aeslc, ag_news, ai2_arc/ARC-Challenge, ai2_arc/ARC-Easy, amazon_polarity, ama-zon_reviews_multi/en, amazon_us_reviews/Wireless_v1_00, ambig_qa/light, app_reviews, aqua_rat/raw, art, asset/ratings, asset/simplification, banking77, billsum, bing_coronavirus_query_set, biosses, blbooksgenre/title_genre_classifiction, blended_skill_talk, cbt/CN, cbt/NE, cbt/P, cbt/V, cbt/raw, cc_news, circa, climate_fever, cnn_dailymail/3.0.0, codah/codah, codah/fold_0, codah/fold_1, codah/fold_2, codah/fold_3, codah/fold_4, code_x_glue_tc_text_to_code, com-mon_gen, commonsense_qa, conv_ai, conv_ai_2, conv_ai_3, cord19/metadata, cos_e/v1.0, cos_e/v1.11, cosmos_qa, covid_qa_castorini, craffel/openai_lambada, craigslist_bargains, crows_pairs, dbpedia_14, discofuse/discofuse-sport, discofuse/discofuse-wikipedia, discovery/discovery, docred, dream, drop, duorc/ParaphraseRC, duorc/SelfRC, e2e_nlg_cleaned, ecthr_cases/alleged-violation-prediction, emo, emotion, enriched_web_nlg/en, esnli, ev-idence_infer_treatment/1.1, evidence_infer_treatment/2.0, fever/v1.0, fever/v2.0, finan-cial_phrasebank/sentences_allagree, freebase_qa, generated_reviews_enth, gigaword, glue/ax, glue/cola, glue/mnli, glue/mnli_matched, glue/mnli_mismatched, glue/mrpc, glue/qnli, glue/qqp, glue/rte, glue/sst2, glue/stsb, glue/wnli, google_wellformed_query, great_code, guardian_authorship/cross_genre_1, guardian_authorship/cross_topic_1, guardian_authorship/cross_topic_4, guardian_authorship/cross_topic_7, gutenberg_time, hans, hate_speech18, head_qa/en, health_fact, hlgd, hotpot_qa/distractor, hotpot_qa/fullwiki,

humicroedit/subtask-1, humicroedit/subtask-2, hyperpartisan_news_detection/byarticle, hyperparti-san_news_detection/bypublisher, imdb, jfleg, kelm, kilt_tasks/hotpotqa, kilt_tasks/nq, lama/trex, lambada, liar, limit, math_dataset/algebra__linear_1d, math_dataset/algebra__linear_1d_composed, math_dataset/algebra__linear_2d, math_dataset/algebra__linear_2d_composed, math_qa, mc_taco, mdd/task1_qa, mdd/task2_recs, mdd/task3_qarecs, medal, medical_questions_pairs, meta_woz/dialogues, mocha, movie_rationales, multi_news, multi_nli, multi_x_science_sum, mwsc, narrativeqa, ncbi_disease, neural_code_search/evaluation_dataset, newspop, nlu_evaluation_data, nq_open, numer_sense, onestop_english, openai_humaneval, openbookqa/additional, openbookqa/main, paws-x/en, paws/labeled_final, paws/labeled_swap, paws/unlabeled_final, piqa, poem_sentiment, pubmed_qa/pqa_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, scan/addprim_jump, scan/addprim_turn_left, scan/filler_num0, scan/filler_num1, scan/filler_num2, scan/filler_num3, scan/length, scan/simple, scan/template_around_right, scan/template_jump_around_right, scan/template_opposite_right, scan/template_right, scicite, scientific_papers/arxiv, scientific_papers/pubmed, sciq, scitail/snli_format, scitail/tsv_format, scitldr/Abstract, selqa/answer_selection_analysis, sem_eval_2010_task_8, sem_eval_2014_task_1, sent_comp, sick, sms_spam, snips_built_in_intents, snli, social_i_qa, species_800, squad, squad_adversarial/AddSent, squad_v2, squadshifts/amazon, squadshifts/new_wiki, squadshifts/nyt, sst/default, stsb_multi_mt/en, subjqa/books, subjqa/electronics, subjqa/grocery, subjqa/movies, subjqa/restaurants, subjqa/tripadvisor, super_glue/axb, super_glue/axg, super_glue/boolq, su-per_glue/multirc, super_glue/record, swag/regular, tab_fact/tab_fact, tmu_gfm_dataset, trec, trivia_qa/unfiltered, turk, tweet_eval/emoji, tweet_eval/emotion, tweet_eval/hate, tweet_eval/irony, tweet_eval/offensive, tweet_eval/sentiment, tweet_eval/stance_abortion, tweet_eval/stance_atheism, tweet_eval/stance_climate, tweet_eval/stance_feminist, tweet_eval/stance_hillary, tydiqa/primary_task, tydiqa/secondary_task, web_questions, wiki_bio, wiki_hop/masked, wiki_hop/original, wiki_qa, wiki_split, wino_bias/type1_anti, wino_bias/type1_pro, wino_bias/type2_anti, wino_bias/type2_pro, winograd_wsc/wsc273, winograd_wsc/wsc285, wiqa, xnli/en, xquad/xquad.en, xquad_r/en, xsum, yahoo_answers_qa, yahoo_answers_topics, yelp_polarity, yelp_review_full, zest

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\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. 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.\n\nQUESTIONS:\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? 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?\n\nLIMITATIONS:\n* I do not have much practical experience with RL and it is difficult for me to judge the limitations of this method.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nImportant problem\nImproves the result with a smaller model\nShows nice empirical results\n\nWEAKNESSES:\nThe novelty of the paper is incremental\nThere are some problems in the presentation that cause confusion as explained below.\n\nQUESTIONS:\nOne 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.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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).\n\nWEAKNESSES:\nAs mentioned by the authors, the proposed training has to be performed from scratch for every domain.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nAuthors adequately addressed associated limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: 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.

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.

- 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.

- 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.

- 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|.

QUESTIONS:
Please address the concerns in the Weaknessese.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06