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: This paper proposed Adaptive Task Scheduler (ATS) for meta-learning, which trains a neural network to adaptively select tasks. A bi-level optimization strategy is used to optimize the the meta-model and neural scheduler. Experiments are conducted on miniImagenet  and drug compounds, which beats the baseline methods.