Meta-learning with an Adaptive Task Scheduler

1 Introduction . Meta-learning has emerged in recent years as a popular paradigm to benefit the learning of a new task in a sample-efficient way , by meta-training a meta-model ( e.g. , initializations for model parameters ) from a set of historical tasks ( called meta-training tasks ) . To learn the meta-model during metatraining , the majority of existing meta-learning methods randomly sample meta-training tasks with a uniform probability . The assumption behind such uniform sampling is that all tasks are equally important , which is often not the case . First , some tasks could be noisy . For example , in our experiments on drug discovery , all compounds ( i.e. , examples ) in some target proteins ( i.e. , tasks ) are even labeled with the same bio-activity value due to improper measurement . Second , the number of meta-training tasks is likely limited , so that the distribution over different clusters of tasks is uneven . There are 2 , 523 target proteins in the binding family among all 4 , 276 proteins for meta-training , but only 55 of them belong to the ADME family . To fill the gap , we are motivated to equip the meta-learning framework with a task scheduler that determines which tasks should be used for meta-training in the current iteration . Recently , a few studies have started to consider introducing a task scheduler into meta-learning , by adjusting the class selection strategy for construction of each few-shot classification task [ 17 , 19 ] , ∗H . Yao and Y. Wang contribute equally ; correspondence to : Y. Wei 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . directly using a self-paced regularizer [ 3 ] , or ranking the candidate tasks based on the amount of information associated with them [ 11 ] . While the early success of these methods is a testament to the benefits of introducing a task scheduler , developing a task scheduler that adapts to the progress of the meta-model remains an open challenge . Such a task scheduler could both take into account the complicated learning dynamics of a meta-learning algorithm better than existing manually defined schedulers , and explicitly optimize the generalization capacity to avoid meta-overfitting [ 23 , 30 ] . To address these limitations , in this paper , we propose an Adaptive Task Scheduler ( ATS ) for metalearning . Instead of fixing the scheduler throughout the meta-training process , we design a neural scheduler to predict the probability of each training task being sampled . Concretely , we adopt a bi-level optimization strategy to jointly optimize both the meta-model and the neural scheduler . The meta-model is optimized with the sampled meta-training tasks by the neural scheduler , while the neural scheduler is learned to improve the generalization ability of the meta-model on a set of validation tasks . The neural scheduler considers two meta-model-related factors as its input : 1 ) the loss of the meta-model with respect to a task , and 2 ) the similarity between gradients of the meta-model with respect to the support and query sets of a task , which characterize task difficulty from the perspective of the outcome and the process of learning , respectively . On this account , the neural scheduler avoids the pathology of a poorly generalized meta-model that is corrupted by a limited budget of tasks or detrimental tasks ( e.g. , noisy tasks ) . The main contribution of this paper is an adaptive task scheduler that guides the selection of metatraining tasks for a meta-learning framework . We identify two meta-model-related factors as building blocks of the task scheduler , and theoretically reveal that the scheduler considering these two factors improves the meta-training loss as well as the optimization landscape . Under different settings ( i.e. , meta-learning with noisy or a limited number of tasks ) , we empirically demonstrate the superiority of our proposed scheduler over state-of-the-art schedulers on both an image classification benchmark ( up to 13 % improvement ) and a real-world drug discovery dataset ( up to 18 % improvement ) . The proposed scheduler demonstrates great adaptability , tending to 1 ) sample non-noisy tasks with smaller losses if there are noisy tasks but 2 ) sample difficult tasks with large losses when the budget is limited . 2 Related Work . Meta-learning has emerged as an effective paradigm for learning with small data , by leveraging the knowledge learned from previous tasks . Among the two dominant strands of meta-learning algorithms , we prefer gradient-based [ 4 ] over metric-based [ 26 ] for their general applicability in both classification and regression problems . Much of the research up to now considers all tasks to be equally important , so that a task is randomly sampled in each iteration . Very recently , Jabri et al . [ 8 ] explored unsupervised task generation in meta reinforcement learning according to variations of a reward function , while in [ 11 , 20 ] a task is sampled from existing meta-training tasks with the probability proportional to the amount of information it offers . Complementary to these methods specific to reinforcement learning , a difficulty-aware meta-loss function [ 15 ] and a greedy class-pair based task sampling strategy [ 17 ] have been proposed to attack supervised meta-learning problems . Instead of using these manually defined and fixed sampling strategies , we pursue an automatic task scheduler that learns to predict the sampling probability for each task to directly minimize the generalization error . There has been a large body of literature that is concerned with example sampling , dating back to importance sampling [ 12 ] and AdaBoost [ 5 ] . Similar to AdaBoost where hard examples receive more attention , the strategy of hard example mining [ 25 ] accelerates and stabilizes the SGD optimization of deep neural networks . The difficulty of an example is calibrated by its loss [ 16 ] , the magnitude of its gradient [ 31 ] , or the uncertainty [ 2 ] . On the contrary , self-paced learning [ 13 ] presents the examples in the increasing order of their difficulty , so that deep neural networks do not memorize those noisy examples and generalize poorly [ 1 ] . The order is implemented by a soft weighting scheme , where easy examples with smaller losses have larger weights in the beginning . More self-paced learning variants [ 9 , 28 ] are dedicated to designing the scheme to appropriately model the relationship between the loss and the weight of an example . Until very recently , Jiang et al . [ 10 ] and Ren et al . [ 24 ] proposed to learn the scheme automatically from a clean dataset and by maximizing the performance on a hold-out validation dataset , respectively . Nevertheless , task scheduling poses more challenges than example sampling that these methods are proposed for . 3 Preliminaries and Problem Definition . Assume that we have a task distribution p ( T ) . Each task Ti is associated with a data set Di , which is further split into a support set Dsi and a query set D q i . Gradient-based meta-learning [ 4 ] , which we focus on in this work , learns a well-generalized parameter θ0 ( a.k.a. , meta-model ) of a base predictive learner f from N meta-training tasks { Ti } Ni=1 . The base learner initialized from θ0 adapts to the i-th task by taking k gradient descent steps with respect to its support set ( inner-loop optimization ) , i.e. , θi = θ0 − α∇θL ( Dsi ; θ ) . To evaluate and improve the initialization , we measure the performance of θi on the query set Dqi and use the corresponding loss to optimize the initialization as ( out-loop optimization ) : θ ( k+1 ) 0 = θ ( k ) 0 − β N∑ i=1 L ( Dqi ; θi ) , ( 1 ) where α and β denote the learning rates for task adaptation and initialization update , respectively . After training K time steps , we learn the well-generalized model parameter initializations θ∗0 . During the meta-testing time , θ∗0 can be adapted to each new task Tt by performing a few gradient steps on the corresponding support set , i.e. , θt = θ∗0 − α∇θL ( Dst ; θ ) . In practice , without loading all N tasks into memory , a batch of B tasks { T ( k ) i } B i=1 are sampled for training at the k-th meta-training iteration . Most of existing meta-learning algorithms use the uniform sampling strategy , except for a few recent attempts towards manually defined task schedulers ( e.g. , [ 11 , 17 ] ) . Either the uniform sampling or manually defined task schedulers may be sub-optimal and at a high risk of overfitting . 4 Adaptive Task Scheduler . To address these limitations , we aim to design an adaptive task scheduler ( ATS ) in meta-learning to decide which tasks to use next . Specifically , as illustrated in Figure 1 , we design a neural scheduler to predict the probability of each candidate task being sampled according to the real-time feedback of meta-model-based factors at each meta-training iteration . Based on the probabilities , we sample B tasks to optimize the meta-model and the neural scheduler in an alternating way . In the following subsections , we detail our task scheduling strategy and bi-level optimization process . 4.1 Adaptive Task Scheduling Strategy . The goal for this subsection is to discuss how to select the most informative tasks via ATS . We define the scheduler as g with parameter φ and formulate the sampling probability w ( k ) i of each candidate task Ti at training iteration k as w ( k ) i = g ( Ti , θ ( k ) 0 ; φ ( k ) ) , ( 2 ) where w ( k ) i is conditioned on task Ti and the meta-model θ ( k ) 0 at the current meta-training iteration . To quantify the information covered in Ti and θ ( k ) 0 , we here propose two representative factors : 1 ) the loss L ( Dqi ; θ ( k ) i ) on the query set , where θ ( k ) i is updated by performing a few-gradient steps starting from θ ( k ) 0 ; 2 ) the gradient similarity between the support and target sets with respect to the current meta-model θ ( k ) 0 , i.e. , 〈 ∇ θ ( k ) 0 L ( Dsi ; θ ( k ) 0 ) , ∇θ ( k ) 0 L ( Dqi ; θ ( k ) 0 ) 〉 . Here , we use inner product as an exemplary similarity measurement , other metrics like cosine similarity can also be applied in practice . They are associated with the learning outcome and learning process of the task Ti , respectively . Specifically , the gradient similarity signifies the generalization gap from the support to the query set . A large query loss may represent a true hard task if the gradient similarity is large ; a task with noises in its query set , however , could lead to a large query loss but small gradient similarity . Considering these two factors simultaneously , we reformulate Eqn . ( 2 ) as : w ( k ) i = g ( L ( Dqi ; θ ( k ) i ) , 〈 ∇ θ ( k ) 0 L ( Dsi ; θ ( k ) 0 ) , ∇θ ( k ) 0 L ( Dqi ; θ ( k ) 0 ) 〉 ; φ ( k ) ) . ( 3 ) After obtaining the sampling probability weight w ( k ) i , we directly use it to sample B tasks from the candidate task pool for the current meta-training iteration , where a larger value of w ( k ) i represents higher probability . The out-loop optimization is revised as : θ ( k+1 ) 0 = θ ( k ) 0 − β 1 B B∑ i=1 L ( Dqi ; θ ( k ) i ) . ( 4 )

Summary: Meta-learning is a popular strategy to learn a well-generalized model from different meta-training tasks. Classical meta-learning assumes equal importance for all the meta-training tasks. Recent work has proposed sampling of the tasks based on various factors such as difficulty, amount of information, etc. The paper presents an extension to these ideas by proposing an adaptive task scheduler. In contrast to the prior approaches where the task distribution is hard-crafted and fixed, the authors propose a learnable task scheduler that outputs the sampling probability score of each task in a task pool to directly minimize the generalization error. The inputs to the task scheduler - the query loss and the similarity between the gradients on the query and support set are theoretically grounded and well-motivated. The meta-model and the tasks scheduler are learned using a bi-level optimization; in particular, REINFORCE is used to update the task scheduler. The authors conduct experiments on the mini-imagenet (classification) and drug discovery (regression) datasets to validate the effectiveness of the task scheduler. The results indicate the robustness of the task scheduler on both the datasets under noisy conditions and limited task budget.