Meta-learning with an Adaptive Task Scheduler
Year: 2022

Abstract:
To benefit the learning of a new task, meta-learning has been proposed to transfer a well-generalized meta-model learned from various meta-training tasks. Existing meta-learning algorithms randomly sample meta-training tasks with a uniform probability, under the assumption that tasks are of equal importance. However, it is likely that tasks are detrimental with noise or imbalanced given a limited number of meta-training tasks. To prevent the meta-model from being corrupted by such detrimental tasks or dominated by tasks in the majority, in this paper, we propose an adaptive task scheduler (ATS) for the meta-training process. In ATS, for the first time, we design a neural scheduler to decide which meta-training tasks to use next by predicting the probability being sampled for each candidate task, and train the scheduler to optimize the generalization capacity of the metamodel to unseen tasks. We identify two meta-model-related factors as the input of the neural scheduler, which characterize the difficulty of a candidate task to the meta-model. Theoretically, we show that a scheduler taking the two factors into account improves the meta-training loss and also the optimization landscape. Under the setting of meta-learning with noise and limited budgets, ATS improves the performance on both miniImageNet and a real-world drug discovery benchmark by up to 13% and 18%, respectively, compared to state-of-the-art task schedulers.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Preliminaries and Problem Definition
text: 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.

heading: 4 Adaptive Task Scheduler
text: 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.

heading: 4.1 Adaptive Task Scheduling Strategy
text: 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)

heading: 4.2 Bi-level Optimization with Gradient Approximation
text: In this subsection, we aim to discuss how to jointly learn the parameters of neural scheduler φ and the meta-model θ0 during the meta-training process. Instead of directly using the meta-training tasks to optimize both meta-model and the neural scheduler, ATS aims to optimize the loss on a disparate validation set with Nv tasks, i.e., {Tv}Nvv=1, where the performance on the validation set could be regarded as the reward or fitness. Specifically, ATS searches the optimal parameter φ∗ of the neural scheduler by minimizing the average loss 1
Nv
∑Nv v=1 Lval(Tv; θ ∗ 0(φ)) over the validation tasks, where
the optimal parameter θ∗0 is obtained by optimizing the meta-training loss 1B ∑B i=1 Ltr(Ti; θ0, φ) over the sampled tasks. Formally, the bi-level optimization process is formulated as:
min φ
1
Nv Nv∑ v=1 Lval(Tv; θ∗0(φ)), where θ∗0(φ) = argmin θ0 1 B B∑ i=1 Ltr(Ti; θ0, φ) (5)
It is computational expensive to optimize the inner loop in Eqn. (5) directly. Inspired by the differentiate hyperparameter optimization [18], we propose a strategy to approximate θ∗0(φ) by performing one gradient step starting from the current meta-model θ(k)0 as:
θ∗0(φ (k)) ≈ θ̃(k+1)0 (φ (k)) = θ (k) 0 − β∇θ(k)0
1
B B∑ i=1 Ltr(Ti; θ(k)i , φ (k)).
s.t. θ (k) i = θ (k) 0 − α∇θL(D s i ; θ),
(6)
where the task-specific parameter θ(k)i is adapted with a few gradient steps starting from θ (k) 0 . As such, for each meta-training iteration, the bi-level optimization process in Eqn. (5) is revised as:
φ(k+1) ← min φ(k)
1
Nv Nv∑ v=1 Lval(Tv; θ̃(k+1)0 (φ (k))),
s.t., θ̃ (k+1) 0 (φ (k)) = θ (k) 0 − β∇θ(k)0
1
B B∑ i=1 Ltr(Ti; θ(k)0 , φ (k)).
(7)
We then focus on optimizing the neural scheduler in the outer loop of Eqn. (7). In ATS, tasks are sampled from the candidate task pool according to the sampling probabilities w(k)i . It is intractable to directly optimizing the validation loss Lval, since the sampling process is non-differentiable. We thereby use a policy gradient method to optimize φ(k), where REINFORCE [29] is adopted. We also
Algorithm 1 Meta-training Process with ATS
Require: learning rates α, β, task distribution p(T ), batch size B, candidate task pool size Npool 1: Initialize the meta-model θ0 and the neural scheduler φ 2: for each training iteration k do 3: Randomly select Npool tasks and construct the candidate task pool 4: for each task Ti in the candidate task pool do 5: Compute two factors L(Dqi ; θ (k) i ) and 〈 ∇ θ (k) 0 L(Dsi ; θ (k) 0 ),∇θ(k)0 L(Dqi ; θ (k) 0 ) 〉
by using the support set Dsi and the query set Dqi
6: Calculate the sampling probability w(k)i by Eqn. (2) 7: end for 8: Sample B tasks from the candidate task pool via the sampling probabilities W(k) 9: Calculate the training loss and obtain a temporal meta-model via 1-step gradient descent
10: Sample Nv validation tasks 11: Calculate the accuracy (reward) R(k)i using the temporal meta-model θ̃ (k+1) 0 12: Update the neural scheduler by Eqn. (8) and get φ(k+1) 13: Sample another B tasks {T ′ i }Bi=1 by using the updated task scheduler φ(k+1) 14: Update the meta-model as: θ(k+1)0 (φ (k)) = θ
(k) 0 − β∇θ(k)0 1 B
∑B i=1 Ltr(T ′ i ; θ (k) 0 , φ (k+1))
15: end for
equip a baseline function for REINFORCE to reduce the gradient variance. Regarding the accuracy of each sampled validation task Ti as the reward R(k)i , we define the optimization process as:
φ(k+1) ← φ(k) − γ∇φ(k) logP (W (k);φ(k))(
1
Nv Nv∑ i=1 R (k) i − b), (8)
where W(k) = {w(k)i } B i=1 and the baseline b is defined as the moving average of validation accuracies. After updating the parameter of neural scheduler φ, we use it to resample B tasks from the candidate task pool and update the meta-model from θ(k)0 to θ (k+1) 0 via Eqn. (7). The whole optimization algorithm of ATS is illustrated in Alg. 1.

heading: 5 Theoretical Analysis
text: In this section, we would extend the theoretical analysis in [6] to our problem of task scheduling, to theoretically study how the neural scheduler g improves the meta-training loss as well as the optimization landscape. Without loss of generality, here we consider a weighted version of ATS with hard sampling, where the meta-model θ0 is updated by solving the meta-training loss weighted by the task sampling probability wi = g(Ti, θ0;φ) over all candidate tasks in the task pool, i.e.,
θ∗0 = argmin θ0 Npool∑ i=1 wiL(Dqi ; θi), θi = θ0 − α∇θL(D s i ; θ). (9)
Denote the meta-training loss without and with the task scheduler as L(θ0) = 1Npool ∑Npool i=1 L(D q i ; θi)
and Lw(θ0) = ∑Npool i=1 wiL(D q i ; θi), respectively. Then we have the following result:
Proposition 1. Suppose that w=[w1,· · · ,wNpool ] denotes the random variable for sampling probabilities, Lθ0 = [L(D q 1; θ0), · · · ,L(D q Npool ; θ0)] denotes the random variable for the loss using the meta-model, and ∇θ0 = [〈∇θ0L(Ds1; θ0),∇θ0L(D q 1; θ0)〉 , · · · , 〈 ∇θ0L(DsNpool ; θ0),∇θ0L(D q Npool ; θ0) 〉 ] denotes the random variable for the inner product between gradients of the support and query sets with respect to the meta-model. Then the following equation connecting the meta-learning losses with and without the task scheduler holds:
Lw(θ0) = L(θ0) + Cov(Lθ0 ,w)− αCov(∇θ0 ,w). (10)
From Proposition 1, we conclude that the task scheduler improves the meta-training loss, as long as the sampling probability w is negatively correlated with the loss but positively correlated with the gradient similarity between the support and the query set. Specifically, if the loss L(Dqi ; θi) is large
as a result of a quite challenging or noisy task Ti, the sampling probability wi is expected to be small. Moreover, a large value of wi is anticipated, when a large inner product between the gradients of the support and the query set with respect to the meta-model 〈∇θ0L(Dsi ; θ0),∇θ0L(D q i ; θ0)〉 signifies that the generalization gap from the support set Dsi to the query set D q i is small.
Consistent with [6], the optimal meta-model is assumed to also minimize the covariance Cov(Lθ0 ,w) and maximize the covariance Cov(∇θ0 ,w), i.e., θ∗0 = argminL(θ0) = argmin[Cov(Lθ0 ,w) − αCov(∇θ0 ,w)]. Under this assumption, the task scheduler does not change the global minimum, i.e., θ∗0 = argminL(θ0) = argminLw(θ0), while modifies the optimization landscape as the following. Proposition 2. With the sampling probability defined as
w∗i = e− [ L(Dqi ;θ ∗ 0 )−α〈∇θ0L(Dsi ;θ∗0 ),∇θ0L(Dqi ;θ∗0 )〉 ] ∑B i=1 e − [ L(Dqi ;θ∗0 )−α〈∇θ0L(Dsi ;θ∗0 ),∇θ0L(Dqi ;θ∗0 )〉
] , (11) the following hold:
∀θ0 : Cov(Lθ0 − α∇θ0 , e −(Lθ∗0 −α∇θ∗0 ) ) ≥ 0, Lw(θ0)− Lw(θ∗0) ≥ L(θ0)− L(θ∗0), ∀θ0 : Cov(Lθ0 − α∇θ0 , e −(Lθ∗0 −α∇θ∗0 ) ) ≤ −Var(Lθ∗0 − α∇θ∗0 ), L w(θ0)− Lw(θ∗0) ≤ L(θ0)− L(θ∗0).
Proposition 2 sheds light on how the optimization landscape is improved by an ideal task scheduler: 1) for those parameters θ0 that are far from the optimal meta-model θ∗0 (i.e., Cov(Lθ0 − α∇θ0 , e −(Lθ∗0 −α∇θ∗0 ) ) ≥ 0), the gradients towards the direction of θ∗0 become overall steeper for speed-up; 2) for those parameters θ0 that are within the variance of the optimum (i.e., Cov(Lθ0 − α∇θ0 , e −(Lθ∗0 −α∇θ∗0 ) ) ≤ −Var(Lθ∗0 − α∇θ∗0 )), the minima tends to be flat with better generalization ability [7, 14]. Though the optimal meta-model θ∗0 remains unknown and the ideal task scheduler with the sampling probabilities in (11) is inaccessible, we learn a neural scheduler to dynamically accommodate the changes of both the loss Lθ0 and the gradient similarity ∇θ0 . Detailed proofs of Proposition 1 and 2 are provided in Appendix A.

heading: 6 Experiments
text: In this section, we empirically demonstrate the effectiveness of the proposed ATS through comprehensive experiments on both regression and classification problems. Specifically, two challenging settings are studied: meta-learning with noise and limited budgets.
Dataset Description and Model Structure. We conduct comprehensive experiments on two datasets. First, we use miniImagenet as the classification dataset, where we apply the conventional N-way, Kshot setting to create tasks [4]. For miniImagenet, we use the standard model with four convolutional blocks, where each block contains 32 neurons. We report accuracy with 95% confidence interval over all meta-testing tasks. The second dataset aims to predict the activity of drug compounds [21], where each task as an assay covers drug compounds for one target protein. There are 4,276 assays in total, and we split 4,100 / 76 / 100 tasks for meta-training / validation / testing, respectively. We use two fully connected layers in the drug activity prediction as the base model, where each layer contains 500 neurons. For each assay, the performance is measured by the square of Pearson coefficient (R2) between the predicted values and the actual values. Follow [21], we report the mean and medium R2 as well as the number of assays with R2 > 0.3, all of which are considered as reliable metrics in the pharmacology domain. We provide more detailed descriptions of datasets in Appendix B.1.
In terms of the neural scheduler φ, we separately encode the loss on the query set and the gradient similarity by two bi-directional LSTM network. The percentage of training iterations are also fed into the neural scheduler to indicate the progress of meta-training. Finally, all encoded information are concatenated and feed into a two-layer MLP for predicting the sampling probability (More detail descriptions are provided in Appendix B.2).
Baselines. We compare the proposed ATS with the following two categories of baselines. The first category contains easy-to-implement example sampling methods that can be adapted for task scheduling, including focal loss (FocalLoss) [16] and self-paced learning loss (SPL) [13]. The second category covers state-of-the-art task schedulers for meta-learning that are non-adaptive, which includes DAML [15], GCP [17], and PAML [11]. Note that GCP is a class-driven task scheduler
and thus it can not be applied to regression problems such as drug activity prediction here. We also slightly revise the loss of DAML so that it can be applied to both classification and regression problems. For all baselines and ATS, we use the same base model and adopt ANIL [22] as the backbone meta-learning algorithm.

heading: 6.1 Meta-learning with Noise
text: Experimental Setup. We first apply ATS on meta-learning with noisy tasks, where each noisy task is constructed by only adding noises on the labels of the support set. Therefore, each noisy task contains a noisy support set and a clean query set. In this way, adapting the meta-model on the support set causes inaccurate task-specific parameters and leads to negative impacts on the meta-training process. Specifically, for miniImagenet, we apply the symmetry flipping on the labels of the support set [27]. The default ratio of noisy tasks is set as 0.6. For the drug activity prediction, we sample the label noise from a Gaussian distribution η ∗ N (0, 1), where the noise scalar η is used to control the noise level. Note that, empirically we find that the effect of adding noise on the drug activity prediction is not as great as adding noise on the miniImagenet, and thus we add noise to all assays and use the scalar η to control the noise ratio. By default, we set the noise scalar as 4 during the meta-training process. Besides, since ATS uses the clean validation task, for fair comparison, all other baselines are also fine-tuned on the validation tasks. Detailed experimental setups and hyperparameters are provided in Appendix C.
Results. Table 1 reports the overall results of ATS and other baselines. Our key observations are: (1) The performance of non-adaptive task schedulers (i.e., GCP, PAML, DAML) is similar to the uniform sampling, indicating that manually designing the task scheduler may not explain the complex dynamics of meta-learning, being sub-optimal. (2) ATS outperforms traditional example sampling methods and nonadaptive task schedulers, demonstrating its effectiveness on improving the robustness of meta-learning algorithms by adaptively adjusting the scheduler policy based on the realtime feedback of meta-model-related factors. The findings are further strengthened by the distribution comparison of sampling weights between clean and noisy tasks in Figure 2, where ATS pushes most noisy tasks to small weights (i.e., less contributions in meta-training).
Ablation Study. We further conduct an ablation study under the noisy task setting and investigate five ablation models detailed as follows. • Rank by Sim/Loss: In the first ablation model, we heuristically determine and select tasks by ranking
tasks according to a simple combination of the loss and the gradient similarity, i.e., Sim/Loss. We assume that tasks with higher gradient similarity but smaller losses should be prioritized.
• Random φ: In Random φ, we remove both meta-model-related factors, where the model parameter φ of the neural scheduler is randomly initialized at each meta-training iteration.
• φ+Loss or φ+Sim: The third (φ+Loss) and forth (φ+Sim) ablation models remove the gradient similarity between the support and query sets, as well as the loss of the query set, respectively.
• Reweighting: Instead of selecting tasks from the candidate pool, in the last ablation model, we direct reweigh all tasks in a meta batch, where the weights are learned via the neural scheduler.
We list the results of all ablation models in Table 2, where ATS (φ+Loss+Sim) is also reported for comparison. The results indicate that (1) simply selecting tasks according to the ratio Sim/Loss significantly underperforms ATS since the contribution of each metric is hard to manually define. Besides, the contribution of each metric evolves as training proceeds, which has been ignored in such a simple combination but modeled in the neural scheduler with the percentage of training iterations as input; (2) the superiority of φ+Loss and φ+Sim over Random φ shows the effectiveness of both the query loss and the gradient similarity; (3) the performance gap between Reweight+Loss+Sim and ATS is potentially caused by the number of effective tasks, where more candidate tasks are considered by ATS; (4) including all meta-model-related factors (i.e., ATS) achieves the best performance, coinciding with our theoretical findings in Section 5.
Effect of Noise Ratio. We analyze the performance of ATS with respect to the noise ratio and show the results of miniImagenet and drug activity prediction in Table 3. The performances of the uniform sampling and the best non-adaptive scheduler (BNS) are also reported for comparison. We summarize the key findings: (1) ATS consistently outperforms the uniform sampling and the best non-adaptive task scheduler, indicating its effectiveness of adaptively sampling tasks to guide the meta-training process; (2) With the increase of the noise ratio, ATS achieves more significant improvements. In particular, the results of ATS in miniImagenet are almost stable even with a very large noise ratio, suggesting that involving an adaptive task scheduler does improve the robustness of the model.
Analysis of the Meta-model-related Factors. To analyze our motivations about designing the metamodel-related factors, we randomly select 1000 meta-training tasks and visualize the correlation between the sampling weight wki and each factor in Figures 3a-3d. In these figures, we rank the sampling weight wki and normalize the ranking to [0, 1], where a larger normalized ranking value is associated with a larger sampling weight. Then, we split these tasks into 20 bins according to the rank of sampling weights. For tasks within each bin, we show the mean and standard error of their query losses and gradient similarities. According to these figures, we find that tasks with larger losses are
associated with smaller sampling weights, verifying our assumption that noisy tasks (large query loss) tend to have smaller sampling weight (Figures 3a, 3c). Our motivation is further strengthened by the fact that tasks with more similar support and query sets have larger sampling weights (Figures 3b, 3d).

heading: 6.2 Meta-learning with Limited Budgets
text: Experimental Setup. We further analyze the effectiveness of ATS under the meta-learning setting with limited budgets. Follow [4], In few-shot classification problem, each training episode is a few-shot task by subsampling classes as well as data points and two episodes that share the same classes are considered to be the same task. Thus, we treat the budgets in meta-learning as the number of meta-training tasks. In miniImagenet, the original number of meta-training classes is 64, corresponding to more than 7 million 5-way combinations. Thus, we control the budgets by reducing the number of meta-training classes to 16, resulting in 4,368 combinations. For drug activity prediction, since it only has 4,100 tasks, we do not reduce the number of tasks and use the full dataset for meta-training. We provide more discussions about the setup in Appendix D.
Results. We report the performance on miniImagenet and drug activity prediction in Table 4. Aligning with the meta-learning with noise setting, ATS consistently outperforms other baselines with limited budgets. Besides, compared with the results in miniImagenet, ATS achieves more significant improvement in the drug activity prediction problem under this setting. This is what we expected – as we have discussed in Section 1, the drug dataset contains noisy and imbalanced tasks.
The effectiveness of ATS is further strengthened by the ablation study under the limited budgets setting, where we report the results in Table 5. Similar to the findings in the noisy task setting, the ablation study further verifies the contributions of the proposed two meta-model-related factors on the meta-training process.
Analysis of the Meta-model-related Factors. Under the limited budgets setting, we analyze the correlation between the sampling weight and the meta-model-related factors in Figure 4a-4d. From these figures, we can see that the gradient similarity indeed reflects the task difficulty (Figure 4b, 4d), where larger similarities correspond to more useful tasks (i.e., larger weights). Interestingly, the
correlation between the query loss and the sampling weight under the limited budgets setting is opposite to the correlation under the noisy setting. It is expected since tasks with larger query losses are associated with “more difficult” tasks under this setting, which may contain more valuable information that further benefits the meta-training process.
Effect of the Budgets. In addition, we analyze the effect of budgets by changing the number of meta-training tasks in miniImagenet. The results of uniform sampling, GCP (best non-adaptive task scheduler), ATS under 1-shot scenario are illustrated in Table 6. We observe that our model achieves the best performance in all scenarios. In addition, compared with uniform sampling, ATS achieves more significant improvements with less budgets, indicating its effectiveness on improving the meta-training efficiency.

heading: 7 Conclusion and Discussion
text: This paper proposes a new adaptive task sampling strategy (ATS) to improve the meta-training process. Specifically, we design a neural scheduler with two meta-model-related factors. At each meta-training iteration, the neural scheduler predicts the probability of each meta-training task being sampled according to the received factors from each candidate task. The meta-model and the neural scheduler are optimized in a bi-level optimization framework. Our experiments demonstrate the effectiveness of the proposed ATS under the settings of meta-learning with noise and limited budgets.
One limitation in this paper is that we only consider how to adaptively schedule tasks during the meta-training process. In the future, it would be meaningful to investigate how to incorporate the task scheduler with the sample scheduler within each task. Another limitation is that using ATS is more computationally expensive than random sampling since we alternatively learn the neural scheduler and the meta-model. It would be interesting to explore how to reduce the computational cost, e.g., compressing the neural scheduler.

heading: Acknowledgement
text: This work was supported in part by JPMorgan Chase & Co. Any views or opinions expressed herein are solely those of the authors listed, and may differ from the views and opinions expressed by JPMorgan Chase & Co. or its affiliates. This material is not a product of the Research Department of J.P. Morgan Securities LLC. This material should not be construed as an individual recommendation for any particular client and is not intended as a recommendation of particular securities, financial instruments or strategies for a particular client. This material does not constitute a solicitation or offer in any jurisdiction. This work was also supported in part by the start-up grant from City University of Hong Kong (9610512). Besides, Y. Wei would like to acknowledge the support from the Tencent AI Lab Rhino-Bird Gift Fund. D. Lian was supported by grants from the National Natural Science Foundation of China # 62022077.

Review summaries:
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. 

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. 

In this paper, the authors propose an adaptive task scheduler to pick the set of tasks that a meta-learner should solve at a given iteration which will improve its generalizability/transferability to downstream tasks rather than picking a set of tasks at random. Instead of choosing a metric/heuristic based approach, the authors propose to use a second network that learns-to-learn the right set of tasks to sample and the overall training progresses via an interleaved (bi-level) optimization procedure. Authors provide empirical evidence based on two datasets that their proposed scheduler ATS performs better than similar existing approaches. The paper also contains some theoretical justifications on why ATS performs better than random sampling.