Title: Meta-Learning with Neural Bandit Scheduler
Venue: NeurIPS 2023
Accepted: yes
Keywords: Meta Learning

Contextual Bandits
Abstract: Meta-learning has been proven an effective learning paradigm for training machine learning models with good generalization ability. Apart from the common practice of uniformly sampling the meta-training tasks, existing methods working on task scheduling strategies are mainly based on pre-defined sampling protocols or the assumed task-model correlations, and greedily make scheduling decisions, which can lead to sub-optimal performance bottlenecks of the meta-model. In this paper, we propose a novel task scheduling framework under Contextual Bandits settings, named BASS, which directly optimizes the task scheduling strategy based on the status of the meta-model. By balancing the exploitation and exploration in meta-learning task scheduling, BASS can help tackle the challenge of limited knowledge about the task distribution during the early stage of meta-training, while simultaneously exploring potential benefits for forthcoming meta-training iterations through an adaptive exploration strategy. Theoretical analysis and extensive experiments are presented to show the effectiveness of our proposed framework.

Meta-Learning with Neural Bandit Scheduler

Abstract

Meta-learning has been proven an effective learning paradigm for training machine learning models with good generalization ability. Apart from the common practice of uniformly sampling the meta-training tasks, existing methods working on task scheduling strategies are mainly based on pre-defined sampling protocols or the assumed task-model correlations, and greedily make scheduling decisions, which can lead to sub-optimal performance bottlenecks of the meta-model. In this paper, we propose a novel task scheduling framework under Contextual Bandits settings, named BASS, which directly optimizes the task scheduling strategy based on the status of the meta-model. By balancing the exploitation and exploration in metalearning task scheduling, BASS can help tackle the challenge of limited knowledge about the task distribution during the early stage of meta-training, while simultaneously exploring potential benefits for forthcoming meta-training iterations through an adaptive exploration strategy. Theoretical analysis and extensive experiments are presented to show the effectiveness of our proposed framework.

1 Introduction

Meta-learning algorithms [19] have been receiving increased attention due to their strong generalization performance across a wide range of tasks [32, 20]. Most existing meta-learning methods often assume a uniform distribution when drawing meta-training tasks, treating each task as equally important [19]. However, this assumption can possibly fail in real-world scenarios. For example, during data collection, candidate training tasks can be subject to noise perturbation, leading to performance bottlenecks in the meta-model if noisy and clean tasks are treated on an equal footing [43, 45]. In addition, some tasks can pose greater challenges for the meta-model to adapt to, necessitating a more flexible allocation of computational resources during the meta-training process. Furthermore, the task distribution may be skewed, with "tail" tasks receiving inadequate attention under uniform sampling. As a result, the task scheduling methods [14, 51, 31, 24, 29] have been proposed for a refined meta-training strategy.

Existing scheduling approaches mainly aim to improve meta-training strategies based on various pre-defined criteria and assumptions, including both non-adaptive and adaptive methods. Among them, non-adaptive methods working on the gradient that updates trainable parameters [31, 13] have proven their superiority over meta-training methods with uniform sampling. But they are unable to adjust the task scheduling strategy adaptively based on the status of the meta-model. On the other hand, adaptive methods aim to schedule the tasks based on Curriculum Learning [14, 52, 50] or adaptively sample the tasks based on the task adaptation difficulty (loss) [51]. However, the existing approaches are all greedy algorithms, which means that they tend to make locally optimal decisions based on the current knowledge (i.e., exploitation only). As the learner only has limited knowledge regarding the task and data distribution at the early stage of meta-training, the greedy

6

Meta-learner:

Task-learner:

Meta-learning:

Learning/adaptation:

①(k-1)

VL6

06

VL1

日(k)

VL3

①(k)

VL5

01

03

05

VL2

VL4

#2

04

Exploitation

vs.

Exploration

strategy can lead to the sub-optimal meta-model due to multiple reasons, such as being misled by the noisy tasks or affected by the skewness of the task distribution. Here, we use Figure 2 to illustrate an example where the greedy approach (only exploitation) may be trapped in sub-optimal solutions.

One intuitive solution for the aforementioned problems is tackling the exploitation-exploration dilemma [4] in meta-training process: balancing the exploitation of current knowledge of selected tasks, and the exploration of under-explored tasks for potential long-term benefits. Therefore, unlike existing approaches with greedy strategies, we propose a novel task scheduling framework under the Contextual Bandits settings [15, 28, 4, 1, 54, 53, 2], named BA ndit Ta S k S cheduler (BASS). In each meta-training iteration, we formulate each candidate meta-training task as an arm under the contextual bandit settings, and the corresponding arm reward will naturally be the meta-model generalization ability score after including this candidate task (arm) into the meta-training process. To achieve this objective, BASS directly learns the mapping from the meta-parameters to the meta-model generalization ability score, instead of depending on the hand-crafted criteria or assumptions. This design also enables us to update meta-parameters and BASS simultaneously within one round of meta-optimization. In particular, instead of greedily scheduling the meta-training tasks solely based on the current knowledge (i.e., exploitation), BASS leverages an additional adaptive exploration module with two different exploration objectives to explore for unrevealed benefits. Our main contributions can be summarized as follows:

[Problem and Proposed Framework] We are the first to formulate the meta-learning task scheduling problem under the Contextual Bandits settings, where we optimize the meta-model w.r.t. chosen task batches in each meta-training iteration. To tackle this problem, we propose a novel banditbased meta-learning task scheduling framework named BASS, which is model-agnostic and can be applied to various meta-learning frameworks. Different from existing works, BASS directly learns the relationship between the meta-model parameters and the meta-model generalization ability. In addition, instead of greedily exploiting the current knowledge as in existing works, BASS utilizes a novel exploration module to adaptively plan for the future meta-training iterations.

[Theoretical Analysis] Under the general meta-learning settings as well as standard assumptions of over-parameterized neural networks and neural bandits, we derive the theoretical performance guarantee for the proposed BASS framework in terms of regret bound.

[Experiments] To demonstrate the effectiveness of BASS, we compare our method against seven strong baselines on three real data sets, with different specifications. In addition, complementary experiments as well as a case study on Ensemble Inference are also provided to better understand the property and behavior of BASS.

2 Related Works

In this section, we briefly review the related works of our proposed BASS framework from the aspects of task scheduling in meta-learning and contextual Bandits.

Task Scheduling in Meta-Learning. By considering each meta-training task to be equally important, many existing works sample the training tasks uniformly from the given task distribution [19, 40]. Then, with fixed sampling strategies, [24, 35] propose to assign the task sampling probability based on the quantity of task information, and [31] utilizes a probabilistic sampling method based on class-pairs. Meanwhile, there are also works try to improve task-specific gradients for the randomly sampled tasks [13, 34]. On the other hand, Curriculum-based approaches [14, 52, 50] schedule the tasks based on the difficulty of task adaptation, and [51] propose to adaptively schedule the tasks based on the assumed correlation between the difficulty of meta-training tasks and the meta-model generalization ability. However, the existing works are all greedy approaches that solely focus on the instant benefits, which can lead to sub-optimal meta-models due to the potential performance bottleneck. Compared with existing works, our proposed BASS directly learns the mapping from the meta-parameters to the generalization loss with no pre-defined criteria. With the adaptive exploration strategy, our proposed BASS helps tackle the insufficient knowledge regarding the task distribution by balancing the exploitation and exploration, as well as focusing on the long-term effects.

Contextual Bandits. Contextual bandits algorithms aim to solve the sequential decision making problem under the online learning settings, such as online recommendation [28, 49, 6]. Assuming that the mapping from arm contexts to rewards is linear, linear upper confidence bound (UCB) algorithms [15, 28, 4] are proposed to solve the exploitation-exploration dilemma. After kernel-based methods [46, 17] being applied to deal with the non-linear setting where the reward mapping is assumed to be a function in Reproducing Kernel Hilbert Space (RKHS), neural bandit algorithms [54, 53, 7, 5, 39, 8] are proposed to utilize neural networks for the reward and confidence bound estimation. In particular, neural bandit algorithms have demonstrated their superior performance over linear and kernel-based algorithms [54], thanks to the representation power of neural networks. Moreover, instead of using non-negative UCB, EE-Net [9, 10] achieves adaptive exploration by adopting an additional neural network to estimate the potential gain of reward estimations. In this paper, we model the task scheduling problem under the Contextual Bandit settings to balance the exploitation and exploration dilemma regarding the meta-task scheduling.

3 Problem Definition and Learning Objective

Under the settings of few-shot supervised meta-learning [19], our goal is to train a meta-model F ( · ; Θ ) with parameters Θ ∈ R p that can generalize well to meta-testing tasks, where p is the number of trainable meta-parameters. The meta-model parameters are initialized as Θ (0) . Note that in this work, the trainable parameters are all represented by column vectors, for the simplicity of notation. Following similar settings in [19, 51], in each training iteration k ∈ [ K ] , we will receive a pool of candidate training tasks Ω ( k ) task = {T k,i } i ∈N where its cardinality | Ω ( k ) task | = N task and the index set N = { 1 , . . . , N task } . Given a task distribution P ( T ) , we draw each candidate task T k,i from P ( T ) to form the candidate pool, i.e., T k,i ∼ P ( T ) . Each task T k,i is also associated with a support data set D s k,i and a query data set D q k,i , where the samples (including their labels) of D s k,i , D q k,i are drawn from the corresponding task data distribution D T k,i of task T k,i .

Meta-training Process. Then, we need to choose a batch of B tasks Ω k = {T k, ̂ i } ̂ i ∈ ̂ N k ⊂ Ω ( k ) task for each meta-training iteration k ∈ [ K ] , where ̂ N k ⊂ N refers to the indices of chosen tasks. With Θ ( k -1) being the meta-model parameters after completing the ( k -1) -th meta-training iteration, for each chosen meta-training task T k, ̂ i , ̂ i ∈ ̂ N k , we run Gradient Descent (GD) for J steps on its support set D s k, ̂ i to obtain the task-specific parameters Θ ( J ) k, ̂ i . This refers to the inner-loop optimization . In this work, we consider a loss function L ( · ; · ) that maps the meta-model parameters and the input sample (or sample set) to the loss value, with the range [0 , 1] (e.g., the MSE loss on normalized user ratings [27]). Denoting Θ (0) k, ̂ i = Θ ( k -1) in meta-training iteration k , each GD step is represented as

where η 1 ∈ R + is the learning rate of GD. For the simplicity of notation, we formulate the above J -iteration inner-loop optimization as an operator I ( · , · ) : T × Θ ↦→ Θ , such that

Meta-trainingwithQk

4.ParameterUpdatefor

Jk,1

BASS&Meta-model

Chosen

自

Tk,2

Task

Tki

Feed

Arm

Meta-model

F(;0(k-1))

Current

Sample

tasks Ωk

Meta-

十

optimization

Meta-

Xki

Xki

S

b

&

GetVf()

&Feedto

into

Network f1()

Exploitation

Reward Est.

Ski=a·fki+eki

SelectBArms

withBenefit

Scores

Candidate

nask

Tasks

BASS

Tk,Ntask!

Meta-parameters

①(k-1)

Current

Arm (task)

Contexts

Exploration

ScoreEst.

eki

Candidate Arm

Exploration

(task) Pool Q(k)k

Networkf2()

Task Dist.

'task

P(J)

1.Arm(Task)ContextDerivation

2.Reward&ExplorationScoreEstimation

3.ArmSelection

Then, we optimize the meta-parameters through the outer-loop optimization with query sets

where η 2 ∈ R + is the learning rate. Here, the trained meta-model parameters Θ ( K ) is expected to minimize the generalization loss E T ∼P ( T ) , x ∼D T [ L ( x ; I ( T , Θ ( K ) ) )] , where we let D T being the associated data distribution of task T .

Evaluating Task Batch Benefits. Recall that in meta-training iteration k ∈ [ K ] , the meta-model will be updated based on the chosen batch of B tasks ( Eq. 3), and we will need to evaluate the benefit in terms of the whole task batch. By the task scheduling problem definition, the learner will select one task batch Ω k ⊂ Ω ( k ) task based on its strategy, and the chosen task batch Ω k will be used for meta-training in this iteration k . Here, we define the reward of the corresponding task batch Ω k as

where Θ ( k -1) [Ω k ] refer to the meta-parameters after adapting Θ ( k -1) to task batch Ω k based on Eq. 3. For simplicity, we let h : R p ↦→ R be the mapping function conditioned on distribution P ( T ) , which maps the trained meta-model parameters Θ ( k -1) [Ω k ] to task batch rewards.

Learning Objective. Up to meta-training iteration k , we have Ω ∗ ( K ) = { Ω ∗ 1 , . . . , Ω ∗ K } being a series of unknown optimal task batches that minimizes the generalization loss. Each optimal batch is Ω ∗ k = {T k,i ∗ } i ∗ ∈N ∗ k ⊂ Ω ( k ) task , k ∈ [ K ] including a total of | Ω ∗ k | = B tasks, with the indices N ∗ k ⊂ N . The corresponding optimal meta-parameters are defined as Θ ( K ) , ∗ = arg inf Θ ∈W [ E T ∼P ( T ) , x ∼D T [ L ( x ; I ( T , Θ ) )] ] , where W ⊂ R p refers to the reachable param. Meanwhile, eter space, which contains the possible trained meta-parameter values, after observing the available candidate tasks { Ω ( k ) task } , k ∈ [ K ] and the randomly initialized meta-parameters Θ (0) denoting our selection of the meta-training task batches as Ω( K ) = { Ω 1 , . . . , Ω K } , we will have the corresponding trained meta-parameters Θ ( K ) . Here, we can define the K -iteration regret as

which measures the difference of the generalization ability between the trained meta-parameters Θ ( K ) and the optimal meta-parameters Θ ( K ) , ∗ , after K meta-training iterations.

4 Proposed Framework: BASS

In this section, we introduce our proposed BASS framework, which simultaneously optimizes the meta-model and the task scheduling strategy on the fly. The pseudo-code is presented in Algo. 1, and the illustration for each meta-training iteration k ∈ [ K ] is shown in Figure 2.

Algorithm 1 BA ndit Ta S k S cheduler (BASS)

1: Input: Task distribution P ( T ) . Iterations K . GD steps J . Number of chosen tasks B . Learning rates for meta-model η 1 , η 2 . Learning rates for BASS η θ 1 , η θ 2 . Exploration coefficient α ∈ (0 , 1] .

2: Output: Trained meta-parameters Θ ( K ) .

3: Initialization: Randomly initialized meta-model parameters Θ (0) , and BASS parameters θ (0) .

4: for each meta-training iteration k ∈ [ K ] do

5: Sample a pool of N task candidate meta-training tasks Ω ( k ) task from the task distribution P ( T ) . ▷ ------Task Scheduling ------

6: for For each candidate task T k,i ∈ Ω ( k ) task do

7: Derive two arm context vectors χ s k,i , χ q k,i [ Eq. 6].

8: Calculate arm benefit score ̂ s k,i = α · f 1 ( · ; θ ( k -1) 1 ) + f 2 ( · ; θ ( k -1) 2 ) with BASS. [ Eq. 7]

10: From Ω ( k ) task , choose the topB tasks Ω k ⊂ Ω ( k ) task with the highest scores { ̂ s k,i } i ∈N . ▷ ------Parameter Updating ------

9: end for

11: for each chosen training task T k, ̂ i ∈ Ω k do

12: Derive the corresponding rewards ˜ r k,i and exploration score ˜ e k,i . [ Eq. 8-9]

14: Update meta-parameters to Θ ( k ) based on chosen arms (tasks) T k, ̂ i ∈ Ω k . [ Remark 2, Eq. 3]

13: end for

15: Update BASS parameters to θ ( k ) with chosen arms . [ Eq. 8-9, Subsec. 4.4]

16: end for

Remark 1 (Task Scheduling with Contextual Bandits) . By the problem definition, there are N batch = ( N task B ) candidate task batches in each meta-training iteration, which can be a large number. In this case, enumerating all possible task batches and estimating their rewards will be time consuming. Therefore, under the Contextual Bandits settings, BASS alternatively considers each candidate task T k,i ∈ Ω ( k ) task as an arm , and directly chooses B arms as the meta-training tasks Ω k ⊂ Ω ( k ) task . As a result, BASS can (1) reduce the arm space size from O ( N batch ) to O ( N task ) , while (2) enjoy the performance guarantee ( Section 5) in terms of regret bound ( Eq. 5).

Section Outline. Here, we will first present our definition of the arm contexts ( Subsec. 4.1), whose formulation is challenging, because our settings are different from conventional Contextual Bandits where arm contexts are readily available from the environment. Then, applying two neural networks f 1 , f 2 for exploitation and exploration respectively, we formulate the arm benefit score ( Subsec. 4.2), which measures the benefit if we include the corresponding arm (task) into the current meta-training process. Next, we define the arm rewards and exploration scores as the labels for training f 1 , f 2 respectively. In particular, to deal with the challenge of achieving exploration in task scheduling, we incorporate the information and the dynamics w.r.t. meta-optimizations, and formulate two separate exploration objectives for a refined exploration strategy ( Subsec. 4.3). Finally, we update BASS with GD ( Subsec. 4.4), and train the meta-model with chosen tasks.

4.1 Formulating Arm Contexts

To encode the information from both the task side and the meta-model side, for each candidate task (i.e., arm) T k,i ∈ Ω ( k ) task in meta-training iteration k ∈ [ K ] , we formulate its arm contexts as the meta-parameters after task adaptations, denoted by

where Θ ( J ) k,i are the task-specific parameters after adapting meta-parameters Θ ( k -1) to task T k,i with inner-loop optimization ( Eqs. 1-2); while Θ ( k -1) [ T k,i ] refer to the meta-parameters after adapting the current Θ ( k -1) to task T k,i , with both inner-loop and outer-loop optimization (as in Eq. 3). In particular, we assign each arm T k,i with two different arm contexts to model the dynamics of meta-parameters w.r.t. inner-loop optimization and outer-loop optimization respectively. For example, the variance of the corresponding data distribution D T k,i can be high. In this case, the support set D s k,i and the query set D q k,i will be considerably different, which tends to make the corresponding

arm contexts χ s k,i , χ q k,i divergent. As a result, the gradient vectors ∇ θ f 1 ( χ s k,i ) , ∇ θ f 1 ( χ q k,i ) will likely be distinct from each other. Alternatively, if the support set D s k,i and the query set D q k,i are not significantly distinct (the distance between χ s k,i and χ q k,i is also likely to be relatively small), these two gradient vectors tend to change dramatically when adapting to T k,i . The reason is possibly that the exploitation model f 1 is not well adapted to this task T k,i . For both scenarios above, it can be beneficial to include more exploration for the task T k,i , and the target is helping f 1 better learn the reward for this task by actively acquiring the knowledge of it. And the two formulated arm contexts can provide important reference for our exploration module.

Remark 2 (Recycling Arm Contexts) . In order to derive the arm contexts ( Eq. 6), the gradients for the outer-loop optimization ∇ Θ L ( D q k, ̂ i ; Θ ( J ) k, ̂ i ) , ̂ i ∈ ̂ N k of the chosen arms T k, ̂ i ∈ Ω k are calculated. As a result, these gradients can be recycled to update the meta-model parameters based on Eq. 3 (line 15, Algo. 1), which helps reduce the computational cost when updating the meta-model.

4.2 Estimating Benefit Scores for Tasks

To determine which arms (tasks) should be included to the meta-training iteration k , we formulate the arm benefit score estimation for each candidate arm T k,i ∈ Ω ( k ) task . The estimated benefit score consists of two parts: (1) the estimated arm reward of choosing this task based on existing knowledge (i.e., exploitation); (2) and the exploration score for the future potential benefit (i.e., exploration). Inspired by recent advances in neural bandits [9], we introduce two separate neural networks, f 1 ( · ; θ 1 ) and f 2 ( · ; θ 2 ) , to estimate the arm reward and exploration score respectively. The exploitation network f 1 ( · ; θ 1 ) aims to learn the mapping h ( · ) from arm contexts (i.e., meta-parameters) to rewards, while the exploration network f 2 ( · ; θ 2 ) aims to learn the uncertainty of reward estimations as the exploration criterion. Different from conventional bandit models, e.g. [9], that works on static arm contexts given by the environment, our design alternatively leverages the evolving information from both the task (arm) side and meta-parameters side, across meta-training iterations. In addition, we consider the dynamics of meta-optimizations for a more comprehensive modeling of the exploration aspect, and the details will be introduced later. Here, given a candidate arm T k,i ∈ Ω ( k ) task , its estimated benefit score ̂ s k,i is formulated as

where α ∈ (0 , 1] is the exploration coefficient to balance exploitation and exploration. Notice that f 2 ( · ; θ 2 ) will take the concatenated gradient of f 1 ( · ; θ 1 ) w.r.t. both arm contexts χ s k,i , χ q k,i as the input, represented by [ ∇ θ f 1 ( χ s k,i ); ∇ θ f 1 ( χ q k,i )] . And the output will be the exploration score estimation ˆ e k,i . To obtain ∇ θ f 1 ( χ q k,i ) , we also calculate f 1 ( χ q k,i ; θ 1 ) and run the back-propagation. Afterwards, we choose the topB arms with the highest estimated benefit scores ̂ s k,i , i ∈ N , as the chosen task batch Ω k ⊂ Ω ( k ) task (line 10, Algo. 1).

Design Intuition. First, recent advances of neural Contextual Bandits [54, 9, 38] have shown that the uncertainty of reward estimations is directly related to the gradients of the estimation model. Therefore, we leverage an exploration module f 2 ( · ; θ 2 ) to directly learn this unknown relationship. Second, since D s k,i , D q k,i are from the same data distribution D T k,i , if these two gradients ∇ θ f 1 ( χ s k,i ) , ∇ θ f 1 ( χ q k,i ) are distinct, the reason can be: (1) the variance of the data distribution D T k,i is high (due to the potentially noisy or difficult task); or (2) the gradients of f 1 ( · ; θ 1 ) tend to change significantly when adapting to task T k,i . In both cases, it can be harder for f 1 ( · ; θ 1 ) to accurately predict the arm reward r k,i , and the meta-model can fail to properly adapt to the task T k,i . In this case, we apply the concatenated gradients w.r.t. both arm contexts as the input of f 2 ( · ; θ 2 ) , in order to provide the information for f 2 ( · ; θ 2 ) to evaluate exploration scores.

Network Architecture and Parameter Initialization. Here, we consider f 1 ( · ; θ 1 ) , f 2 ( · ; θ 2 ) to be two L -layer fully-connected (FC) networks with network width m , while θ = { θ 1 , θ 2 } refer to their trainable parameters. For their randomly initialized parameters θ (0) = { θ (0) 1 , θ (0) 2 } , the weight matrix entries for the first L -1 layers are drawn from the Gaussian distribution N (0 , 2 /m ) , while the entries of the last layer ( L -th layer) are sampled from N (0 , 1 /m ) .

Remark 3 (Reducing Input Complexity) . The input of f 1 ( · ; θ 1 ) is the arm context χ q k,i , whose dimensionality is the number of meta-parameters p . A similar situation also exists for the exploration

network f 2 ( · ; θ 2 ) . Inspired by the idea of learning dense low-dimensional representations with Convolutional Neural Networks (CNNs) (e.g., [44]), we apply the average pooling approach to approximate original inputs for reducing the running time and space complexity in practice. To show its effectiveness, we will apply this approach on BASS for all the experiments in Section 6.

4.3 Formulating Arm Rewards and Exploration Scores

Different from the conventional neural bandit algorithms [53, 54, 9] where the reward is provided by the environment oracle, we need to carefully design the arm rewards to reflect the arm benefit in terms of the meta-model's generalization ability. Analogous to task batch rewards ( Eq. 4), we formulate the single arm reward r k,i = h ( Θ ( k -1) [ T k,i ]) = 1 -E T ∼P ( T ) , x ∼D T [ L ( x ; I ( T , Θ ( k -1) [ T k,i ]) )] for arm T k,i . Since it is impractical to calculate the arm reward by enumerating over P ( T ) , we sample a batch of validation tasks Ω valid k to derive the unbiased reward approximation, denoted by

Here, we adopt the single-step inner-loop optimization [19, 40] to derive I ( T k, ̂ i , Θ ( k -1) [ T k,i ]) in Eq. 8, in order to save the computational cost in practice. Under the few-shot settings, the computation of arm rewards is efficient, since the support set is generally small for inner-loop optimization. The approximation error here can be bounded by the concentration inequality, as the validation tasks Ω valid k are sampled from P ( T ) .

On the other hand, to formulate the exploration score e k,i (i.e., the label for f 2 ( · ; θ 2 ) ), we consider two separate exploration objectives: (1) the prediction uncertainty for the exploitation module f 1 ( · ; θ 1 ) , which is r k,i -f 1 ( χ q k,i ; θ 1 ) ; (2) the validation loss of the meta-model, which represents the difficulty of adapting to T k,i , inspired by the "task difficulty measurer" in Curriculum Learning [52, 50]. As a result, with L k,i = L ( D q k,i ; Θ ( J ) k,i ) being the validation loss of arm T k,i , we formulate the exploration score as e k,i = α · ( r k,i -f 1 ( χ q k,i ; θ ( k -1) 1 ) ) +(1 -α ) · L k,i . Analogously, with the approximated reward ˜ r k,i ( Eq. 8), we calculate the exploration score approximation by

Here, the exploration coefficient α ∈ (0 , 1] (in Eq. 7) is also used to balance our two exploration objectives, which are (1) prediction uncertainty r k,i -f 1 ( · ; θ 1 ) : if the exploitation model f 1 ( · ; θ 1 ) is under-estimating the arm reward, leading to the positive residual r k,i -f 1 ( · ; θ 1 ) , we will have a high exploration score to enhance the exploration for this arm; otherwise, when r k,i -f 1 ( · ; θ 1 ) is negative, it indicates an excessively high estimation, which will alternatively lower the exploration score to compensate for the over-estimation. With a higher α value, our exploration strategy will focus more on the behavior of the exploitation model f 1 ( · ) ; (2) the difficulty of task adaptation (i.e., validation loss) L k,i : if the current meta-model does not generalize well to arm T k,i , the validation loss L k,i will be high, which will also lead to a high exploration score. In this way, our formulation considers two different exploration objectives as well as the dynamics of meta-optimizations (base on concatenated network gradients), for a refined exploration strategy.

4.4 Updating Bandit Scheduler Parameters

After updating the meta-parameters (Line 14, Algo. 1, Remark 2) with tasks Ω k = {T k, ̂ i } ̂ i ∈ ̂ N k , we proceed to update the parameters of BASS (Line 15, Algo. 1). Recall that f 1 ( · ; θ 1 ) tries to learn the reward mapping function h ( · ) , and f 2 ( · ; θ 2 ) aims to learn the exploration score. Given the selected arms Ω k , with η θ 1 , η θ 2 ∈ R + being the learning rates, we apply the GD and quadratic loss to update the parameters of BASS, denoted by θ ( k ) 1 = θ ( k -1) 1 -η θ 1 · ∇ θ 1 ( 1 B ∑ T k, ̂ i ∈ Ω k ∣ ∣ f 1 ( χ q k, ̂ i ; θ ( k -1) 1 ) -˜ r k, ̂ i ∣ ∣ 2 ) , θ ( k ) 2 = θ ( k -1) 2 -η θ 2 · ∇ θ 2 ( 1 B ∑ T k, ̂ i ∈ Ω k ∣ ∣ f 2 ( [ ∇ θ f 1 ( χ s k, ̂ i ); ∇ θ f 1 ( χ q k, ̂ i )]; θ ( k -1) 2 ) -˜ e k, ̂ i ∣ ∣ 2 ) . We refer to Eqs. 8-9 for calculating the approximated arm reward ˜ r k, ̂ i and exploration score ˜ e k, ̂ i .

5 Theoretical Analysis

Recall that in each iteration k ∈ [ K ] , we receive candidate arms (tasks) Ω ( k ) task = {T k,i } i ∈N , and each arm T k,i is associated with two context vectors χ s k,i , χ q k,i . For the sake of analysis, we normalize these

two contexts such that ∥ χ s k,i ∥ 2 = ∥ χ q k,i ∥ 2 = 1 , and set the exploration coefficient α = 1 . Following the existing work [47, 48], we let the meta-model F ( · ; Θ ) be a L F -layer FC network with Gaussian Initialization, with the network width m F . Note that our results can also be generalized to other network architectures, such as CNN and ResNet [22], based on the analysis of over-parameterized neural networks [11, 48, 3]. For the theoretical analysis, we adopt Sigmoid activation for f 1 and ReLU for f 2 , in order to make f 1 Lipschitz smooth under over-parameterization settings. Then, we draw trained parameters of BASS with { θ ( k ) 1 , θ ( k ) 2 } ∼ { ˜ θ ( τ ) 1 , ˜ θ ( τ ) 2 } τ ∈ [ k ] . Here, starting from the randomly initialized parameters { θ (0) 1 , θ (0) 2 } , each parameter pair { ˜ θ ( τ ) 1 , ˜ θ ( τ ) 2 } , τ ∈ [ k ] is separately trained on past arm rewards { r τ ′ , ̂ i } τ ′ ∈ [ τ ] , ̂ i ∈ ̂ N τ ′ and exploration scores { e τ ′ , ̂ i } τ ′ ∈ [ τ ] , ̂ i ∈ ̂ N τ ′ with J θ -iteration GD. Next, similar to existing neural bandit works (e.g., [54, 9, 53]) and the works on meta-model convergence analysis (e.g., [47, 48]), we have the following separateness assumption.

̸

Assumption 5.1 ( ρ -Separateness) . After K meta-training iterations, for every pair of arm contexts χ q k,i , χ q k ′ ,i ′ with k, k ′ ∈ [ K ] such that the corresponding arms T k,i ∈ Ω k ∧ T k ′ ,i ′ ∈ Ω k ′ , if ( k, i ) = ( k ′ , i ′ ) , we have ∥ χ q k,i -χ q k ′ ,i ′ ∥ 2 ≥ ρ where 0 < ρ ≤ O ( 1 L ) .

The assumption above is mild because of two main reasons: (1) since L is manually chosen (e.g., L = 2 ), we can easily satisfy the condition 0 < ρ ≤ O ( 1 L ) as long as no two arm contexts are identical; (2) since the meta-parameters Θ ( k ) , k ∈ [ K ] are constantly changing, the corresponding arm contexts will also be distinct across different meta-training iterations. Additional discussions on this assumption are in Appendix Section B. With standard settings of over-parameterized neural networks [47, 3, 11] and the definition of regret R ( K ) in Eq. 5, we have the following Theorem 5.2. Theorem 5.2. Define δ ∈ (0 , 1) , 0 < ξ 1 , ξ 2 ≤ O (1 /K ) , ξ f = max { ξ 1 , ξ 2 } , 0 < ρ ≤ O (1 /L ) , c ξ > 0 , ξ L = ( c ξ ) L . Suppose the network width m ≥ Ω ( Poly ( K,L,N task , ρ -1 ) log(1 /δ ) ) ; m F ≥ Ω ( Poly ( K,L F , N task ) · log(1 /δ ) ) . Then, let the learning rates be η 1 , η 2 = Θ ( m -1 F Poly ( K,N task ,L F ) ) ; η 1 θ , η 2 θ = Θ ( ρ · m -1 Poly ( K,N task ,L ) ) . J θ = Θ ( Poly ( K,N task ,L ) ρ · δ 2 · log( 1 ξ 1 ) ) . Following Algo. 1, with probability at least 1 -δ , the K -round R ( K ) of BASS could be bounded by

where γ m = O ( 1 m 1 /cγ ) , c γ > 1 , and γ 1 = O (1) with sufficient network width m of BASS.

The proof is presented in Appendix Section D. Here, the first term on the RHS is scaled by the 1 / √ K term, which means the regret bound will shrink along with more iterations K . The second term on the RHS is scaled by 1 / √ m F , which makes it a diminutive term under the over-parameterization settings. Since the network depth L of BASS is a small integer (we apply L = 2 for experiments in Section 6), ξ L will also be a relatively small constant. Meanwhile, γ m will also decrease significantly with increasingly large network width m of BASS. In contrast, with a convex loss function (e.g., L 2 loss or cross-entropy loss) and the same over-parameterization settings, the regret upper bound of the uniform sampling strategy [19, 40] can possibly scale up to 1 for the worst-case scenario, and the upper bound will not decrease with more iterations K (Appendix Lemma D.13). Alternatively, BASS works under the bandit settings, by directly measuring the meta-model performance difference w.r.t. the chosen task batch and the optimal one. With more iterations K , BASS tends to make more accurate scheduling decisions, which makes our regret bound possible. For the existing works, [51] prove that they can improve the optimization landscape with the assumed correlation of task difficulties and meta-model generalization ability. [13] show that their self-paced strategy can improve the model robustness when facing noisy training tasks. Different from previous works, we provide the performance guarantee for the proposed BASS under the neural bandit framework.

6 Experiments

In this section, we compare BASS against seven strong baselines, including: (1) Uniform Sampling; non-adaptive self-paced methods and task schedulers (2) SPL [26], (3) Focal-Loss (FOCAL) [30], (4) DAML [29], (5) GCP [31], (6) PAML [24]; and the adaptive task scheduler (7) ATS [51]. Since GCP is not originally compatible with our problem settings and can only work with classification

problems, we properly adapt it by choosing the tasks with the highest probabilities, and apply it on the classification data sets. ANIL [40] is adopted as the backbone meta-learning framework. Due to page limit, we include the complementary experiments (e.g., parameter study for α , effects of different levels of task skewness), and the configurations to Appendix Section A.

6.1 Real Data Sets with Noisy Meta-Training Tasks

We adopt Drug [36], Mini-ImageNet (M-ImageNet) [42] and CIFAR-100 (CIFAR) [25] data sets under the few-shot learning scenario. Similar to [51], we apply classification accuracy as the evaluation criterion for the Mini-ImageNet and CIFAR-100 data sets, and consider the squared Pearson coefficient for the Drug data set. For each meta-learning iteration k ∈ [ K ] , the learner is given a candidate pool of 10 tasks (i.e., | Ω ( k ) task | = N task = 10 ), and it will need to choose a batch of B = 2 tasks as the training tasks Ω k for this iteration. For the Mini-ImageNet and CIFAR-100 data sets, we consider half of the meta-training tasks are perturbed by the label flipping noise [21], where the chance of a label being flipped is ϵ ∈ [0 , 1] . As the Drug data set stands for a regression problem, we draw the label noise from the Gaussian distribution N (0 , ϵ 2 ) . The experimental settings are under 1-shot or 5-shot, 5-Way (for classification data sets) learning scenario with the noise level ϵ = 0 . 5 . The experiment results are shown in Table 1 and Figure 3. For the average ranking column, we exclude results from the M-ImageNet (1) setting, since BASS and the baselines tend to train a meta-model performing "random guessing". More discussions are in the next paragraph.

Here, BASS can generally outperform the baselines by directly learning the mapping from the metaparameters to the arm rewards, as well as balancing the exploitation and exploration. ATS also achieves good performance as it adaptively learns the correlation between the task adaptation difficulty and task scheduling, which proves that it is necessary to apply the adaptive scheduling strategy instead of staying with a fixed protocol. Meanwhile, BASS can generally train the meta-model more efficiently ( Figure 3), leading to good performances at the early stage of

0

5000

10000

15000

20000

Iterations

0.22

0.24

0.26

0.28

0.30

0.32

0.34

0.36

M-ImageNet (5-shot,

: 0.5)

0

5000

10000

15000

20000

Iterations

0.25

0.30

0.35

0.40

0.45

0.50

0.55

CIFAR (5-shot,

: 0.5)

Uniform

SPL

FOCAL LOSS

DAML

GCP

PAML

ATS

BASS

meta-training. In particular, for the Mini-ImageNet data set under the 1-shot, 5-way settings, all the algorithms fail to train an effective meta-model. Here, under the 5-way classification scenario, all the methods will likely generate a meta-model performing "random guessing" (around 20% accuracy). In this case, utilizing Ensemble Inference techniques [12, 16] can help alleviate this problem, and we include further discussion in the case study ( Subsec. 6.3).

6.2 Effects of the Skewed Task Distribution

The skewed task distribution P ( T ) commonly exists in real-word cases. For instance, consider an animal image classification data set where each class (i.e., task) corresponds to one

kind of animals. In this case, felid classes can be considered as "frequent" tasks in the task distribution due to their large quantity and strong mutual correlations, compared with "tail" tasks like kangaroo classes. In this case, paying insufficient attention to the "tail" classes can impair the

generalization performance of the trained meta-model. Thus, to investigate the effects of when the task distribution P ( T ) is skewed, we randomly choose some tasks from P ( T ) , and assign them with higher sampling probabilities (weights). This corresponds to the situation when P ( T ) is skewed, so that sampling from P ( T ) will likely lead to similar tasks. Here, with the CIFAR-100 data set, we sample 10 tasks and assign them with higher sampling probabilities (weights) (5 tasks with 10% , 5 tasks with 5% ), while the rest of the tasks equally share the remaining 25% probability.

From Table 2, our proposed BASS maintains the best performance due to its adaptive scheduling strategy and the ability of balancing exploitation and exploration. On the other hand, the baselines are unable to adjust their meta-training strategies towards exploration, when the meta-model has already well adapted to the "frequent" tasks. This can possibly lead to the sub-optimal performance of the meta-model. In particular, based on Figure 4, when facing a skewed task distribution, BASS can actively explore the "tail" tasks, after sufficiently exploring the "frequent" tasks. BASS is also able to achieve good performances with fewer meta-training iterations.

6.3 Case Study: BASS-aided Ensemble Inference

From Figure 3, we notice that BASS can train a meta-model that achieves good generalization performance at the early stage of meta-training. One application of this property is using BASS to assist meta-learning models under the Ensemble Inference settings, where separate models are combined to enhance the generalization ability. One renowned ensemble approach is the modelparameter ensemble [41, 33]. With a collection of N E individual models {F ( · ; Θ i ) } i ∈ [ N E ] of the same architecture, the ensemble model will be F E ( · ; Θ E ) , and its parameters Θ E = 1 N E ∑ i ∈ [ N E ] Θ i are the averaged parameters across individual models. Then, the ensemble model F E ( · ; Θ E ) will be applied as the inference model for downstream problems. Here, one natural way of obtaining the individual models F ( · ; Θ i ) , i ∈ [ N E ] is deeming the models from different training iterations as the individual models for ensemble [12, 16]. Here, we conduct experiments using the ensemble techniques with individual models F ( · ; Θ i ) trained by baselines and BASS. We choose the top N E = 10 models with the smallest validation loss across different meta-training iterations as the individual models {F ( · ; Θ i ) } i ∈ [ N E ] for ensemble. The results are shown in Table 3. We label the ensemble version of BASS as "BASS-E", and the non-ensemble version as "BASS-S".

Compared with the non-ensemble settings ( Table 1), we see that the BASS-aided ensemble model can generally perform better. In particular, the ensemble model can improve the meta-model inference performance in significantly difficult cases, such as the Mini-ImageNet under the 1-shot setting ( Subsec. 6.1). As a result, BASS can help generate high-quality ensemble model with the meta-models trained in different iterations. While the ensemble inference technique can also benefit the other baselines, BASS still maintains decent performances.

7 Conclusion

In this paper, we formulate the task scheduling problem in meta-learning under the Contextual Bandits settings, and propose a novel bandit-based task scheduling framework named BASS. It directly optimizes the task sampling strategy based on the status of the meta-model rather than applying fixed task scheduling protocols. Instead of greedily making decisions, BASS can help deal with the insufficient knowledge problem at the early stage of meta-training, as well as plan for the future meta-training iterations with the adaptive exploration strategy. We include both theoretical analyses and a comprehensive set of experiments to demonstrate the effectiveness of our proposed framework as well as its key properties.

Acknowledgments and Disclosure of Funding

This work is supported by National Science Foundation under Award No. IIS-1947203, IIS-2117902, IIS-2137468, IIS-2002540, and Agriculture and Food Research Initiative (AFRI) grant no. 202067021-32799/project accession no.1024178 from the USDA National Institute of Food and Agriculture. The views and conclusions are those of the authors and should not be interpreted as representing the official policies of the funding agencies or the government.

References

Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. Advances in neural information processing systems , 24:2312-2320, 2011.

Shipra Agrawal and Navin Goyal. Thompson sampling for contextual bandits with linear payoffs. In ICML , pages 127-135. PMLR, 2013.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

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

Yikun Ban and Jingrui He. Convolutional neural bandit: Provable algorithm for visual-aware advertising. arXiv preprint arXiv:2107.07438 , 2021.

Yikun Ban and Jingrui He. Local clustering in contextual multi-armed bandits. In Proceedings of the Web Conference 2021 , pages 2335-2346, 2021.

Yikun Ban, Jingrui He, and Curtiss B Cook. Multi-facet contextual bandits: A neural network perspective. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining , pages 35-45, 2021.

Yikun Ban, Yunzhe Qi, Tianxin Wei, and Jingrui He. Neural collaborative filtering bandits via meta learning. arXiv preprint arXiv:2201.13395 , 2022.

Yikun Ban, Yuchen Yan, Arindam Banerjee, and Jingrui He. EE-net: Exploitation-exploration neural networks in contextual bandits. In International Conference on Learning Representations , 2022.

Yikun Ban, Yuchen Yan, Arindam Banerjee, and Jingrui He. Neural exploitation and exploration of contextual bandits. arXiv preprint arXiv:2305.03784 , 2023.

Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in Neural Information Processing Systems , 32:10836-10846, 2019.

Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems , 34:22405-22418, 2021.

Dong Chen, Lingfei Wu, Siliang Tang, Xiao Yun, Bo Long, and Yueting Zhuang. Robust metalearning with sampling noise and label noise via eigen-reptile. arXiv preprint arXiv:2206.01944 , 2022.

Yudong Chen, Xin Wang, Miao Fan, Jizhou Huang, Shengwen Yang, and Wenwu Zhu. Curriculum meta-learning for next poi recommendation. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining , pages 2692-2702, 2021.

Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandits with linear payoff functions. In AISTATS , pages 208-214, 2011.

Xu Chu, Yujie Jin, Wenwu Zhu, Yasha Wang, Xin Wang, Shanghang Zhang, and Hong Mei. Dna: Domain generalization with diversified neural averaging. In International Conference on Machine Learning , pages 4010-4034. PMLR, 2022.

Aniket Anand Deshmukh, Urun Dogan, and Clay Scott. Multi-task learning for contextual bandits. In NeurIPS , pages 4848-4856, 2017.

Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning , pages 1675-1685. PMLR, 2019.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning , pages 1126-1135. PMLR, 2017.

Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. In International Conference on Machine Learning , pages 1568-1577. PMLR, 2018.

B Han, Q Yao, X Yu, G Niu, M Xu, W Hu, I Tsang, and M Sugiyama. Robust training of deep neural networks with extremely noisy labels. In Thirty-fourth Conference on Neural Information Processing Systems (NeurIPS) , volume 2, page 4, 2020.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv:1806.07572 , 2018.

Jean Kaddour, Steindór Sæmundsson, et al. Probabilistic active meta-learning. Advances in Neural Information Processing Systems , 33:20813-20822, 2020.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

MKumar, Benjamin Packer, and Daphne Koller. Self-paced learning for latent variable models. Advances in neural information processing systems , 23, 2010.

Hoyeop Lee, Jinbae Im, Seongwon Jang, Hyunsouk Cho, and Sehee Chung. Melu: Meta-learned user preference estimator for cold-start recommendation. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 1073-1082, 2019.

Lihong Li, Wei Chu, John Langford, and Robert E Schapire. A contextual-bandit approach to personalized news article recommendation. In WWW , pages 661-670, 2010.

Xiaomeng Li, Lequan Yu, Yueming Jin, Chi-Wing Fu, Lei Xing, and Pheng-Ann Heng. Difficulty-aware meta-learning for rare disease diagnosis. In International Conference on Medical Image Computing and Computer-Assisted Intervention , pages 357-366. Springer, 2020.

Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision , pages 2980-2988, 2017.

Chenghao Liu, Zhihao Wang, Doyen Sahoo, Yuan Fang, Kun Zhang, and Steven CH Hoi. Adaptive task sampling for meta-learning. In European Conference on Computer Vision , pages 752-769. Springer, 2020.

Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv:1806.09055 , 2018.

Shiwei Liu, Tianlong Chen, Zahra Atashgahi, Xiaohan Chen, Ghada Sokar, Elena Mocanu, Mykola Pechenizkiy, Zhangyang Wang, and Decebal Constantin Mocanu. Deep ensembling with no overhead for either training or testing: The all-round blessings of dynamic sparsity. arXiv preprint arXiv:2106.14568 , 2021.

Zhuoqun Liu, Yuankun Jiang, Chenglin Li, Wenrui Dai, Junni Zou, and Hongkai Xiong. Adaptive task sampling and variance reduction for gradient-based meta-learning. 2022.

Ricardo Luna Gutierrez and Matteo Leonetti. Information-theoretic task selection for metareinforcement learning. Advances in Neural Information Processing Systems , 33:20532-20542, 2020.

Eric J Martin, Valery R Polyakov, Xiang-Wei Zhu, Li Tian, Prasenjit Mukherjee, and Xin Liu. All-assay-max2 pqsar: activity predictions as accurate as four-concentration ic50s for 8558 novartis assays. Journal of chemical information and modeling , 59(10):4450-4459, 2019.

Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision , pages 1406-1415, 2019.

Yunzhe Qi, Yikun Ban, and Jingrui He. Neural bandit with arm group graph. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 1379-1389, 2022.

Yunzhe Qi, Yikun Ban, and Jingrui He. Graph neural bandits. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 1920-1931, 2023.

Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning or feature reuse? towards understanding the effectiveness of maml. In International Conference on Learning Representations , 2019.

Alexandre Rame, Matthieu Kirchmeyer, Thibaud Rahier, Alain Rakotomamonjy, Patrick Gallinari, and Matthieu Cord. Diverse weight averaging for out-of-distribution generalization. arXiv preprint arXiv:2205.09739 , 2022.

Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. 2016.

Mengye Ren, Wenyuan Zeng, Bin Yang, and Raquel Urtasun. Learning to reweight examples for robust deep learning. In International conference on machine learning , pages 4334-4343. PMLR, 2018.

Yelong Shen, Xiaodong He, Jianfeng Gao, Li Deng, and Grégoire Mesnil. Learning semantic representations using convolutional neural networks for web search. In Proceedings of the 23rd international conference on world wide web , pages 373-374, 2014.

Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Metaweight-net: Learning an explicit mapping for sample weighting. Advances in neural information processing systems , 32, 2019.

Michal Valko, Nathaniel Korda, Rémi Munos, Ilias Flaounas, and Nelo Cristianini. Finite-time analysis of kernelised contextual bandits. arXiv preprint arXiv:1309.6869 , 2013.

Haoxiang Wang, Ruoyu Sun, and Bo Li. Global convergence and generalization bound of gradient-based meta-learning with deep neural nets. arXiv preprint arXiv:2006.14606 , 2020.

Haoxiang Wang, Yite Wang, Ruoyu Sun, and Bo Li. Global convergence of maml and theoryinspired neural architecture search for few-shot learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 9797-9808, June 2022.

Qingyun Wu, Huazheng Wang, Quanquan Gu, and Hongning Wang. Contextual bandits in a collaborative environment. In SIGIR , pages 529-538, 2016.

Tongtong Wu, Xuekai Li, Yuan-Fang Li, Gholamreza Haffari, Guilin Qi, Yujin Zhu, and Guoqiang Xu. Curriculum-meta learning for order-robust continual relation extraction. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 10363-10369, 2021.

Huaxiu Yao, Yu Wang, Ying Wei, Peilin Zhao, Mehrdad Mahdavi, Defu Lian, and Chelsea Finn. Meta-learning with an adaptive task scheduler. Advances in Neural Information Processing Systems , 34:7497-7509, 2021.

Ji Zhang, Jingkuan Song, Yazhou Yao, and Lianli Gao. Curriculum-based meta-learning. In Proceedings of the 29th ACM International Conference on Multimedia , pages 1838-1846, 2021.

Weitong Zhang, Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural thompson sampling. arXiv preprint arXiv:2010.00827 , 2020.

Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural contextual bandits with ucb-based exploration. In International Conference on Machine Learning , pages 11492-11502. PMLR, 2020.

A Appendix: Experiments (Continue)

A.1 Further Details for the Experiment Settings

For the data partitioning, we have the Mini-ImageNet and CIFAR-100 data sets divided into the partitions 64 : 16 : 20 , which correspond to the training set, validation set and the testing set respectively. Each class is corresponding to a task. Then, for the Drug data set, we partition the tasks into 4100 : 76 : 100 representing the training set, validation set and the testing set.

For our BASS, we apply two 2 -layer FC networks for f 1 ( · ; θ 1 ) , f 2 ( · ; θ 2 ) respectively, and set network width m = 200 . For deriving approximated arm rewards, we let | Ω valid k | = 5 . Recall that we apply the approximation approach mentioned in Remark 3 to reduce the space complexity and time complexity in practice for the experiments. Here, we tune the pooling step such that the inputs of f 1 ( · ; θ 1 ) , f 1 ( · ; θ 1 ) are approximately 50 and 20 respectively. For the learning rate, we find the learning rate for BASS with grid search from { 0 . 01 , 0 . 001 , 0 . 0001 } , and choose the learning rates for the meta-model η 1 = 0 . 01 , η 2 = 0 . 001 . The meta-model architecture as well as its learning rates will stay the same for all the baselines and our proposed BASS. For the CIFAR-100 and Mini-ImageNet data sets, we use the the meta-model with four convolutional blocks where the network width of each block is 32, followed by an FC layer as the output layer. For the Drug data set, we apply a meta-model with two FC layers, where the network width is 500. All the experiments are performed on a Linux machine with Intel Xeon CPU, 128GB RAM, and Tesla V100 GPU. Code will be made available at https://github.com/yunzhe0306/Bandit_Task_Scheduler .

A.2 Effect of the Task Noise Magnitude

We conduct the experiments to show the effects of the noise magnitude factor ϵ on the Drug and CIFAR-100 data sets. The experiment results are shown in Table 4.

With increasing noise magnitude ϵ , the performances of the meta-model trained by baselines and our BASS tend to drop, which is intuitive. In particular, for the CIFAR-100 data set, when we increase ϵ , the performance difference between BASS and the other baselines tends to increase. This can be the reason that the greedy baselines with no exploration strategies can be more susceptible to the task noise perturbation, which can lead to the sub-optimal performances of the meta-model.

From the Table 5, we can see that when there is no noise, the overall performance does not differ significantly across different methods. The reason could be that since the meta-learning backbone remains the same for all the methods, the meta-model performance upper bound can be similar for different scheduling algorithms, without the presence of other confounding factors (e.g., noise, task distribution skewness). In the practical application scenarios with noisy data, BASS-guided meta-models tend to perform well in presence of task noise and skewness compared with baselines, as presented by our experiments in the main body.

A.3 Parameter Study for Exploration Coefficient

As in Eq. 7 and Eq. 9, BASS involves an exploration coefficient α to balance the exploitationexploration and the two exploration objectives. Here, we conduct the parameter study for the exploration coefficient α , and include the results with no exploration (i.e., removing f 2 ).

From the results in Table 6, we see that the exploration module can indeed improve the performance of BASS compared with the performance with no exploration. This also fits our initial argument that the greedy algorithm alone can lead to sub-optimal performances of meta-model. By properly choosing the α value, we will be able to achieve a good balance between exploitation and exploration, as well as between the two exploration objectives. Here, setting α ∈ [0 . 5 , 0 . 7] will be good enough to achieve satisfactory performances. Meanwhile, we also note that even with no exploration, our BASS still achieves good performances by directly learning the correlation between the adapted meta-parameter and the generalization score, and refining the scheduling strategy based on the status of the meta-model.

A.4 Running Time Comparison

In Figure 5, we include the running time comparison with baselines. We can see that BASS can achieve significant improvement in terms of the running time, and can take as little as 50% of ATS's running time. The intuition is that our proposed BASS only needs one round of the optimization process to update the meta-model and BASS. On the other hand, from Algorithm 1 of the ATS paper [51], we see that ATS requires two optimization rounds for each meta-training iteration to (1) update the scheduler with the temporal meta-model, and (2) update the actual meta-model respectively. Based on the figure on the RHS, we also see that BASS can achieve a relatively good balance between computational cost and performance.

D (5-shot)

D (1-shot)

M (5-shot)

M (1-shot)

C (5-shot)

C (1-shot)

0

2000

4000

6000

8000

10000

12000

14000

16000

Running time (seconds)

Running time comparison with ATS

BASS

ATS

A.5 Performances with Different Task Skewness Settings

In Table 7, we include the experiments with different levels of skewness. Here, we see that with less skewness levels (the skewness level reduces from Setting 1 to Setting 3), the accuracy of BASS as well as the baselines will continue to improve, while BASS still maintains decent performances.

A.6 Performances with Different Batch Size

With Table 8, we include additional experiments with different batch sizes B , in comparison with the ATS and the uniform sampling approach. Here, we see that with larger B values, the accuracy of BASS as well as the baselines will generally improve.

A.7 Performances with Different Embedding Approaches of Arm Contexts

In Table 9, we include additional experiments with different levels of average pooling, such that after the average pooling, the dimensionality of the pooled vector representation will fall into { 20 , 100 , 500 } . We see that overly small dimensionality of the average-pooled vector representation (e.g., 20) can lead to sub-optimal performance of the BASS framework. In addition, we see that setting the dimensionality to 50 can generally lead to good enough performance.

Here, we also include additional experimental results using MLP to map the original context into the lower dimensional space instead of using our proposed average pooling (Remark 3). Results are shown in Table 10. Here, we use the one-layer MLP with the ReLU activation to embed the original meta-parameters to the low-dimensional vector representations. We can see that the MLPbased method can indeed lead to some performance improvement. But in general, the performance difference between MLP-based embedding and the average-pooling vector representation is subtle. We also note that the MLP-based mapping approach is more time consuming compared with the average pooling approach, since we also need to train the additional embedding layer, which has a considerable number of trainable parameters.

A.8 Additional Experiments on the "DomainNet" data set

In Table 11, we include additional experiments on the new "DomainNet" data set [37]. Within the "real" domain, we filter 100 classes that have at least 600 images. In this way, with each class being a task with 600 images, we will have a total of 100 tasks. Compared with image data sets in our paper (Mini-ImageNet and CIFAR-100), we increase the image resolution of "DomainNet" by resizing its images to 128 × 128 pixels. Following the settings in our paper, we divide tasks into 64 : 16 : 20 portions that correspond to the training set, validation set and the test set respectively. For the few-shot settings, we formulate the problem to be 5-shot, 5-way / 7-way with uniform sampling and ATS as baselines. With a higher image resolution of the "DomainNet" data set, BASS can still maintain the good performance compared with the baselines.

B Appendix: Additional Discussion on the Necessity of Assumption 5.1

We would like to mention that in order to finish the convergence and generalization analysis for the neural Contextual Bandit works (e.g., [54, 2, 9]), the separateness assumption of the arm context is the minimum requirement of the data set. This is because the training data needs to be non-degenerate (i.e., every pairs of samples are distinct) to ensure that the neural network can consistently converge, as indicated by Assumption 2.1 in [3]. Therefore, our Assumption 5.1 regarding the arm separateness aims to ensure that the BASS is able to adequately learn the underlying reward mapping function with sufficient information. Comparing with the existing works, in the convergence analysis works on meta-learning [47, 48], they measure the arm separateness in terms of the minimum eigenvalue λ 0 (with λ 0 > 0 ) of the Neural Tangent Kernel (NTK) [23] matrix, which is comparable with our Euclidean separateness ρ . For existing neural bandit works, Assumption 5.1 in [9] is similar to our separateness assumption. Meanwhile, Assumption 4.2 in [54] and Assumption 3.4 from [53] also imply that no two arms are the same in terms of the minimum NTK matrix eigenvalue λ 0 > 0 .

C Appendix: Limitation

One potential limitation of BASS is that its improvement over baselines may not be significant when dealing with noise-free settings and non-skewed task distributions ( Table 5). Meanwhile, the non-adaptive FOCAL-LOSS [30] tends to achieve a similar performance comparing with the adaptive method ATS [51], while enjoying an advantage in terms of the computational cost. In practical terms, although BASS can generally achieve the decent performance and enjoys a smaller computational cost than ATS, the practitioner still needs to consider whether their task distribution is noisy or skewed in order to strike a good balance between the computational resource needed and the meta-model performance, as BASS can achieve a more significant advantage over baselines given the noisy or skewed task distribution.

D Appendix: Theoretical Analysis

In this section, we present the proof for Theorem 5.2. Here, instead of directly going for the batch setting where we adopt training task batch Ω k for each iteration k ∈ [ K ] ( | Ω k | = | Ω ∗ k | = B ), we first introduce the results of the single-task setting (Subsec. D.1), i.e., | Ω k | = | Ω ∗ k | = 1 . Then, the results will be extended to the batch settings as in Subsec. D.2. Recall that for the meta-model, we first consider it to be a L F -layer fully-connected (FC) network (of width m F for the theoretical analysis (lines 237-239). In particular, we follow the settings in [3] for the Gaussian initialization of weight matrices. For the weight matrix elements in meta-model's first ( L F -1) layers, we draw each of them from the Gaussian distribution N (0 , 2 /m F ) . Then, for the weight matrix elements of the last layer ( L F -th layer), we draw each of them from the Gaussian distribution N (0 , 1) .

D.1 Single-task settings

For the brevity of notation, we denote the scheduler output f ( Θ ( K -1) [ T k,i ]; θ ( k -1) ) = f 1 ( χ q k,i ; θ ( k -1) 1 ) + f 2 ( [ ∇ θ f 1 ( χ s k,i ); ∇ θ f 1 ( χ q k,i )]; θ ( k -1) 2 ) , which corresponds to the definition in Eq. 7. In this case, T ( K ) = {T 1 , . . . , T K } refer to the chosen tasks and T ∗ ( K ) = {T ∗ 1 , . . . , T ∗ K } are the optimal ones. Based on the problem definition, we will have

where the first inequality is due to the arm pulling mechanism, i.e., f ( Θ ( K -1) [ T ∗ K ]; θ ( K -1) ) ≤ f ( Θ ( K -1) [ T K ]; θ ( K -1) ) . Here, f ( · ; ˜ θ ( K -1) ) is defined as the "shadow" bandit model that are trained on optimal tasks {T ∗ 1 , T ∗ 2 , . . . , T ∗ K -1 } and the corresponding meta-model parameters. Here, denote χ K = Θ ( K -1) [ T K ] ∈ R p as the arm context given the arm T K and the meta-model parameter Θ ( K -1) ; similarly, we have χ ∗ K = Θ ( K -1) , ∗ [ T ∗ K ] ∈ R p being the arm context given the arm T ∗ K and the meta-model parameter Θ ( K -1) , ∗ . Thus, for the term I 0 on the RHS, we have

Then, inserting the inequality will lead to

Here, the terms I 1 , I 2 refer to the approximation error for the two bandit models (our possessed model f ( · ; θ ( K -1) ) and the pseudo model f ( · ; ˜ θ ( K -1) ) ). Then, the third term I 3 bounds the output difference when given the same input Θ ( K -1) , ∗ [ T K ] to two separate bandit models, and the final

term I 4 refers to the difference of the meta-model parameters when adapted to the same task with two individual sets of parameters. Here, the terms I 1 , I 2 can be bounded by Lemma D.1, Corollary D.2. Then, the point is to bound the difference term I 4 when given different inputs to the bandit model.

D.1.1 Bounding error terms and assembling the regret bound

[Bounding term I 3 ] For error term I 3 , it focuses on bounding the output difference between two bandit models f ( · ; ˜ θ ( K -1) ) , f ( · ; θ ( K -1) ) given the same input Θ ( K -1) , ∗ [ T ∗ K ] , and we have

With the defined ξ L , applying Lemma D.11 as well as Corollary D.12, we will have

Then, for term I 3 . 2 , we have

Here, for the first term on the RHS, we apply Lemma D.11 as well as Corollary D.12 to bound.

Then, for the second term, with Gaussian initialization of weight matrices, for the over-parameterized FC network f with Lipschitz-smooth activation functions (e.g., Sigmoid), we can have | f ( x ) -f ( x ′ ) | , ∥∇ f ( x ) -∇ f ( x ′ ) ∥ ≤ ξ · ∥ x -x ′ ∥ due to its Lipschitz continuity / smoothness property [47, 18]. Meanwhile, we also have the Lipschitz continuity property for over-parameterized FC network f ′ with ReLU activation [3], such that | f ′ ( x ) -f ′ ( x ′ ) | ≤ ξ ′ · ∥ x -x ′ ∥ . By the Gaussian initialization of BASS's weight matrices and the properties of over-parameterized neural networks [3, 47, 18], we have ξ L = max { ξ, ξ ′ } ≤ O ( c L ξ ) being the Lipschitz constant for our f 1 , f 2 , where c ξ > 1 is a small constant. Applying the conclusion above, we will have

where the last inequality is by Theorem 5 in [3] and the proof of Lemma D.11. With the above results, it will give us

[Bounding term I 4 ] On the other hand, applying the analogous procedure for term I 4 , denoting χ ∗ K = Θ ( K -1) , ∗ [ T ∗ K ] and ¯ χ ∗ K = Θ ( K -1) [ T ∗ K ] for the brevity of notation, we can have

where χ s, ∗ K , χ q, ∗ K respectively represents the support set and query set for task T ∗ K and the metaparameters Θ ( K -1) , ∗ . Similar notation also applies to ¯ χ ∗ K = Θ ( K -1) [ T ∗ K ] . And the inequality is due to the fact that ReLU networks are naturally Lipschitz continuous w.r.t. some coefficient ξ L when they are wide enough [3], as we have discussed above.

[Bounding term I 4 . 1 ] Based on the meta-optimization procedure (inner-loop optimization + outerloop optimization), we have

where Θ ( J ) , ∗ K is the task-specific parameter of T ∗ K after adapting on Θ ( K -1) , ∗ with inner-loop optimization, and the Θ ( J ) K is the similar parameter after adapting on Θ ( K -1) . Here, we simplify the formula by representing the gradient derivation (inner-loop + outer-loop) with the mapping H : T × Θ ↦→ R p , which leads to

Recall that the past arms, including the actual chosen arms {T 1 , T 2 , . . . , T K } as well as the optimal ones {T ∗ 1 , T ∗ 2 , . . . , T ∗ K } are all from the candidate pool where each candidate arm is drawn i.i.d. from the task distribution P ( T ) . Therefore, denoting the bound as ∥ H ( T ∗ K , Θ ( K -1) , ∗ ) -H ( T ∗ K , Θ ( K -1) ) ∥ 2 ≤ S 1 ( K ) , we can have the upper bound as I 4 . 1 ≤ η 2 · ξ L K · S 1 ( K ) .

Then, for the term S 1 ( K ) , by definition we have ∥ H ( T ∗ K , Θ ( K -1) , ∗ ) -H ( T ∗ K , Θ ( K -1) ) ∥ ≤ S 1 ( K ) , applying mean-reduction for the sample loss will further leads to

This inequality essentially bound the gradient difference when given the same input task T ∗ K w.r.t. different sets of model parameters. Based on the conclusion from Lemma 9 of [47] and Lemma

B.3 of [11], we have ∥∇ Θ l f ( x ; Θ K ) ∥ F , ∥∇ Θ l L ( x ; Θ K ) ∥ F ≤ O ( √ m F ) , ∀ l ∈ [ L F ] for any set of parameters within the sphere Θ K ∈ B ( Θ 0 , ω ) where Θ 0 is the center and ω is the corresponding radius (which is a small value). With a total of L F layers for the meta-model and each layer of m F hidden units, this will give us ∥∇ Θ L ( x ; Θ ( J ) , ∗ K ) ∥ 2 , ∥∇ Θ L ( x ; Θ ( J ) K ) ∥ 2 ≤ O ( √ m F L F ) ( Lemma D.15). And this makes S 1 ( K ) ≤ O ( √ m F L F ) . Since we have η 1 , η 2 ≤ O ( 1 m F ) , summarizing the results above, the upper bound can then be derived.

[Bounding term I 4 . 2 ] Next, we proceed to bound I 4 . 2 , which will be

where the inequalities are due to the Lipschitz continuity / smoothness properties of overparameterized FC networks as we discussed above. Here, we notice that the second term on the RHS can be bounded by directly applying the proving procedure of term I 4 . 1 . Then, for the first term on the RHS, we can following a similar procedure as for I 4 . 1 , by

where the last inequality is due to Lemma D.15 and by iterating through K meta-training iterations. Summing up the results above, we will have I 4 . 2 ≤ O ( η 2 ξ 2 L · K √ m F L F ) + O ( η 1 ξ 2 L · KJ √ m F L F ) .

[Summing up the results] Then, combining all the results, we would have

where

Note that the majority of the terms above can be cancelled to O (1) with proper networks width m indicated in Theorem 5.2. With increasingly large network width m , these terms will also become diminutive enough to achieve our regret bound in the main body.

D.2 Extending the result to the batch settings (Proof of Theorem 5.2)

With the results and conclusions from Subsection D.1, we proceed to provide the proof of Theorem 5.2 under the batch settings. Recall that in our original problem formulation and Algorithm 1, we are expected to select a batch of B arms in each meta-training iteration, denoted by { Ω k } k ∈ [ K ] . Note that each of the candidate arms from Ω ( k ) task are drawn i.i.d. from the task distribution P ( T ) . Meantime,

we will have the corresponding optimal arm batches, denoted by { Ω ∗ k } k ∈ [ K ] , which minimizes the loss objective in Eq. 5. Recall that we update the meta-model parameters with

where Θ ( J ) k,i is the task-specific parameter for T k,i after the inner-loop optimization for J steps.

Analogously, for the notation brevity and the sake of analysis, we denote f ( Θ ( K -1) [Ω K ]; θ ( k -1) ) = f 1 ( χ q k ; θ ( k -1) 1 ) + f 2 ( [ ∇ θ f 1 ( χ s k ); ∇ θ f 1 ( χ q k )]; θ ( k -1) 2 ) where we have χ q k := Θ ( K -1) [Ω K ] being the meta-parameters adapted to batch of tasks Ω K , and the batch-specific parameter is defined as χ s k := 1 | Ω K | ∑ T k, ̂ i ∈ Ω K [ Θ ( J ) k, ̂ i ] . Then, the regret under the batch setting can be denoted by

and after applying properties of the arm pulling mechanism, it is equivalent to

where the average value of estimated benefit scores for individual tasks T K,i ∈ Ω K is represented as ̂ f ( Θ ( K -1) [Ω K ]) = 1 | Ω K | · ∑ T K,i ∈ Ω K f ( Θ ( K -1) [ T K,i ]) = 1 | Ω K | · ∑ T K,i ∈ Ω K f ( Θ ( k -1) -η 2 · ∇ Θ L ( D q K,i ; Θ ( J ) K,i ); θ ( K -1) ) , and the inequality is due to the pulling mechanism of BASS. Here, I 5 , I 9 individually correspond to I 1 , I 2 in the single-task setting and can be bounded by Lemma D.3, Corollary D.4. Term I 7 can be upper bounded by I 3 + I 4 from the single-task setting above. Then, for the rest terms I 6 , I 8 , we proceed to bound them separately.

D.2.1 Bounding error terms and assembling the regret bound

We begin with the term I 8 , and then proceed to I 6 . For the chosen batch of tasks Ω K in the round K , we will have f 1 ( Θ ( K -1) [Ω K ]) = f 1 ( Θ ( K -1) -η 2 ·∇ Θ ( 1 | Ω K | ∑ T K,i ∈ Ω K L ( D q K,i ; Θ ( J ) K,i ) ) ; θ ( K -1) 1 ) , In this case, the average value of estimation sampling probabilities for tasks T K,i ∈ Ω K is

[Bounding the f 1 output difference] Next, let us first proceed to bound the output difference with respect to the exploitation module f 1 , where we can transform this term into

Then, applying the Lipschitz continuity property will lead to

Here, by the definition of the outer-loop optimization of first-order meta-learning, we will have an alternative form the inequality, denoted by

For the term in the parentheses on the RHS, substituting the backward operation with the H ( · , · ) mapping function, we have

with | Ω K | = B . Therefore, the f 1 part of error term I 8 could be bounded by f 1 ( Θ ( K -1) [Ω K ]) -̂ f 1 ( Θ ( K -1) [Ω K ]) ≤ η 2 · ξ L · B · S 1 ( K ) . where the upper bound S 1 ( K ) ≤ O ( √ m F L F ) can be found in the procedure bounding term I 4 . 1 .

[Bounding the f 2 output difference] Then, with χ K = Θ ( K -1) [Ω K ] , we proceed to bound the output difference with respect to the exploration module, which is represented by

By adopting the Lipschitz continuity property of f 2 , we will have

where the last inequality is by applying the conclusion when bounding the output difference w.r.t. the exploitation module f 1 . Then, for the second term on the RHS,

where the last inequality is due to Lemma D.15. Summing up all the results above will give us the upper bound for f 2 output difference f 2 ( Θ ( K -1) [Ω K ]) -̂ f 2 ( Θ ( K -1) [Ω K ]) ≤ O ( η 2 · ξ 2 L B · √ m F L F + η 1 · BJ · √ m F L F ) .

[Similar procedure for term I 6 ] Analogously, we can apply the same derivation for the error term I 6 , which leads to

Following a similar procedure as that of term I 8 will give us a similar bound as

where the learning rate η 1 , η 2 ≤ O ( 1 m F ) is a small value. Then, the upper bounds for error terms I 6 , I 8 are given as desired.

[Assembling the results] Then, combining all the results, we would have

where

Similarly, with proper networks width m as in Theorem 5.2, the majority of the terms above can be cancelled to O (1) . With increasingly large network width m under the over-parameterization settings, γ 1 , γ m will also become diminutive.

D.3 Performance Guarantee for the Exploitation and Exploration Modules

In this subsection, we would like to give the performance guarantee for the proposed BASS framework, and the corresponding performance bound can be applied to derive an upper bound for the error terms I 1 , I 2 for the single-task settings and I 5 , I 9 under the batch settings. Up to meta-training iteration k ∈ [ K ] (before updating the meta-parameters and BASS), we denote all the past records received as P k -1 . Before presenting the main lemmas, we first introduce the following operator. Inspired by [3], with two arbitrary vectors ˜ χ , χ such that ∥ ˜ χ ∥ 2 ≤ 1 , ∥ χ ∥ 2 = 1 , we have the following operator

as the concatenation of the two vectors ˜ χ √ 2 , χ 2 and one constant c , where c = √ 3 4 -( ∥ ˜ χ ∥ 2 √ 2 ) 2 ≥ 1 2 . And this operator transforms the transformed vector into unit norm, ∥ ϕ (˜ χ , χ ) ∥ 2 = 1 . The idea of this operator is to make the gradients ∇ θ f 1 ( · ; θ 1 ) of the exploitation model, which is the input of the exploration model f 2 ( · ; θ 2 ) , comply with the normalization requirement and the separateness assumption ( Assumption 5.1). For the sake of analysis, we will adopt this operation in the following proof. Note that this operator is just one possible solution, and our results could be easily generalized to other forms of input gradients under the unit-length and separateness assumption. Similar ideas are also applied in previous works [9]. We begin to bound the single-task settings with the following lemma.

Lemma D.1. For the constants c ′ g > 0 , 0 < ρ ≤ O ( 1 L ) and ξ 1 ∈ (0 , 1) , given the past records P k -1 , we suppose m,η 1 , η 2 satisfy the conditions in Theorem 5.2, and randomly draw the parameter { θ ( k ) 1 , θ ( k ) 2 } ∼ { ˜ θ ( τ ) 1 , ˜ θ ( τ ) 2 } τ ∈ [ k ] . Consider the past records P k -1 up to round k are generated by a fixed policy when witness the candidate arms { Ω ( τ ) task } τ ∈ [ k ] . Then, with probability at least 1 -δ given an arm-reward pair ( T k, ̂ i , r k, ̂ i ) , we have

where

Proof. The proof of this lemma is inspired by Lemma C.1 from [9]. First, we can derive the output upper bound

by triangle inequality and applying the generalization result of FC networks ( Lemma D.5) on f 1 ( · ; θ 1 ) , f 2 ( · ; θ 2 ) .

For the brevity of notation, we use ∇ f 1 ( T k, ̂ i ) to denote ϕ ( [ ∇ θ f 1 ( χ s k, ̂ i ); ∇ θ f 1 ( χ q k, ̂ i )] c ′ g L , χ q k, ̂ i ) and apply ( χ k , r k ) as ( χ q k, ̂ i , r k, ̂ i ) for the following proof. Define the difference sequence as

Since the past rewards and the received arm-reward pairs ( χ τ , r τ ) are generated by the same reward mapping function, we have the expectation

where F τ denotes the filtration given the past records P τ , up to round τ ∈ [ k ] . This also gives the fact that V (1) τ is a martingale difference sequence. Then, after applying the martingale difference sequence over [ k ] , we have

Then, by applying the Azuma-Hoeffding inequality, it leads to

where the first equality is due to the sampling of candidate tasks and the model parameters. Here, the upper bound (i) is derived by applying the conclusions of Lemma D.6 and Lemma D.10, and the inequality (ii) is derived by adopting Lemma D.6 while defining the empirical loss to be

Here, analogous to the trained parameters, we consider the shadow parameters as { θ ( k ) , ∗ 1 , θ ( k ) , ∗ 2 } ∼ { ˜ θ ( τ ) , ∗ 1 , ˜ θ ( τ ) , ∗ 2 } τ ∈ [ k ] . Similarly, each pair { ˜ θ ( τ ) , ∗ 1 , ˜ θ ( τ ) , ∗ 2 } is separately trained on past received rewards of the optimal arm(s) { r τ ′ ,i ∗ } τ ′ ∈ [ τ ] , T τ ′ ,i ∗ ∈ Ω ∗ k and past exploration scores of the optimal arm(s) { e τ ′ ,i ∗ } τ ′ ∈ [ τ ] , T τ ′ ,i ∗ ∈ Ω ∗ k with J θ -iteration GD, starting from the random initialization { θ (0) 1 , θ (0) 2 } .

Corollary D.2. For the constants 0 < ρ ≤ O (1 /L ) and ξ 1 ∈ (0 , 1) , given the past records P k -1 , we suppose m,η 1 , J satisfy the conditions in Theorem 5.2, and randomly draw the parameters { θ ( k ) , ∗ 1 , θ ( k ) , ∗ 2 } ∼ { ˜ θ ( τ ) , ∗ 1 , ˜ θ ( τ ) , ∗ 2 } τ ∈ [ k ] . For the optimal arm T k,i ∗ ∈ Ω k task , consider its union set with the the collection of past optimal arms P ∗ k -1 ∪ {T k,i ∗ , r k,i ∗ } are generated by a fixed policy when witness the candidate arms { Ω ( τ ) task } τ ∈ [ k ] , with P ∗ k -1 being the collection chosen by this policy. Then, with probability at least 1 -δ , we have

where r τ,i ∗ is the corresponding reward generated by the mapping function given an arm χ τ,i ∗ , and

Proof. This corollary is the direct application of Lemma D.1 by following a similar proof procedure. First, suppose the shadow models f 1 ( · ; θ 2 ) , f 2 ( · ; θ 2 ) are trained on the alternative trajectory P ∗ k -1 . Analogous to the proof of Lemma D.1, we can define the following martingale difference sequence with regard to the previous records P ∗ k -1 up to round τ ∈ [ t ] as

Since the records in set P ∗ k -1 are sharing the same reward mapping function, we have the expectation

where F ∗ τ denotes the filtration given the past records P ∗ k -1 . The mean value of V (1) , ∗ τ across different time steps will be

with the expectation of zero. Afterwards, applying the Azuma-Hoeffding inequality, with a constant δ ∈ (0 , 1) , it leads to

To bound the output difference between the shadow model f 1 ( · ; θ ( k -1) , ∗ 1 ) , f 2 ( · ; θ ( k -1) , ∗ 2 ) and the model we trained based on received records f 1 ( · ; θ ( k -1) 1 ) , f 2 ( · ; θ ( k -1) 2 ) , we apply the conclusion from Lemma D.11, which leads to that given arbitrary input vectors x , x ′ , we have

Finally, combining all the results will finish the proof.

We will also be able to have the performance guarantee under the batch settings. Recall that given a batch of chosen tasks Ω k ⊂ Ω ( k ) task , k ∈ [ K ] , we have the meta-parameters adapted to this batch of tasks being Θ ( K -1) [Ω K ] , which we consider as the input for the f 1 ( · ; θ 1 ) model, where the tasks within each collection are sampled from the task distribution. Thus, chosen task batches from different iterations are also independent from each other. Intuitively, we can also define the corresponding reward for arm batch Ω k as r k = h ( Θ ( k -1) [Ω k ]) . Then, we bound the batch settings with the following lemma and corollary.

Lemma D.3. For the constants c ′ g > 0 , ρ ∈ (0 , O ( 1 L )) and ξ 1 ∈ (0 , 1) , given the past records P k -1 , we suppose m,η 1 , η 2 satisfy the conditions in Theorem 5.2, and randomly draw the parameter { θ ( k ) 1 , θ ( k ) 2 } ∼ { ˜ θ ( τ ) 1 , ˜ θ ( τ ) 2 } τ ∈ [ k ] . Consider the past records P k -1 up to round k are generated by a fixed policy when witness the candidate arms { Ω ( τ ) task } τ ∈ [ k ] . Then, with probability at least 1 -δ given the pair of chosen arm batch and the reward (Ω k , r k ) in round k , we have

where

Proof. The proof of this lemma is analogous to the proof of Lemma D.1. First, we can derive the output upper bound

by triangle inequality and applying the generalization result of FC networks ( Lemma D.5) on f 1 ( · ; θ 1 ) , f 2 ( · ; θ 2 ) , where c ′ g > 0 is a positive number to scale the concatenated gradient vector.

For the brevity of notation, we use ∇ f 1 (Ω k ) to denote ϕ ( [ ∇ θ f 1 ( χ s k ); ∇ θ f 1 ( χ q k )] c ′ g L , χ q k ) and apply ( χ k , r k ) as ( χ q k , r k ) for the following proof. Define the difference sequence as

Since the past rewards and the received arm batch-reward pairs ( χ τ , r τ ) are generated by the same reward mapping function, we have the expectation

where F τ denotes the filtration given the past records P τ , up to round τ ∈ [ k ] . This also gives the fact that V (2) τ is a martingale difference sequence. Then, after applying the martingale difference sequence over [ k ] , we have

By the Azuma-Hoeffding inequality, it leads to P [ 1 k ∑ τ ∈ [ k ] V (2) τ -1 k ∑ τ ∈ [ k ] E [ V (2) τ ] ≥ (1 +

2 γ 1 ) √ 2 log(1 /δ ) k ] ≤ δ. As we have discussed, the tasks within each collection are sampled from the task distribution, which makes chosen task batches from different iterations Ω k , k ∈ [ K ] are also independent from each other. Since the expectation of V (2) τ is zero, with the probability at least 1 -δ and an existing set of parameters θ 2 s.t. ∥ θ 2 -θ (0) 2 ∥ ≤ O ( k 3 ρ √ m log m ) , the above inequality implies

where the first equality is due to the sampling of candidate tasks and the model parameters. Here, the upper bound (i) is derived by applying the conclusions of Lemma D.6 and Lemma D.10, and the inequality (ii) is derived by adopting Lemma D.6 while defining the empirical loss to be

Analogously, we consider the shadow parameters as { θ ( k ) , ∗ 1 , θ ( k ) , ∗ 2 } ∼ { ˜ θ ( τ ) , ∗ 1 , ˜ θ ( τ ) , ∗ 2 } τ ∈ [ k ] where each pair { ˜ θ ( τ ) , ∗ 1 , ˜ θ ( τ ) , ∗ 2 } is separately trained on past received rewards of the optimal arm(s) { r τ ′ ,i ∗ } τ ′ ∈ [ τ ] , T τ ′ ,i ∗ ∈ Ω ∗ k and past exploration scores of the optimal arm(s) { e τ ′ ,i ∗ } τ ′ ∈ [ τ ] , T τ ′ ,i ∗ ∈ Ω ∗ k with J θ -iteration GD starting from the random initialization { θ (0) 1 , θ (0) 2 } .

Corollary D.4. For the constants ρ ∈ (0 , O ( 1 L )) and ξ 1 ∈ (0 , 1) , given the past records P k -1 , we suppose m,η 1 , J satisfy the conditions in Theorem 5.2, and randomly draw the parameters { θ ( k ) , ∗ 1 , θ ( k ) , ∗ 2 } ∼ { ˜ θ ( τ ) , ∗ 1 , ˜ θ ( τ ) , ∗ 2 } τ ∈ [ k ] . For the optimal arm batch Ω ∗ k ⊂ Ω k task , consider its union set with the the collection of past optimal arms P ∗ k -1 ∪{ Ω ∗ k , r ∗ k } are generated by a fixed policy when witness the candidate arms { Ω ( τ ) task } τ ∈ [ k ] , with P ∗ k -1 being the collection chosen by this policy. Then, with probability at least 1 -δ , we have

where r τ,i ∗ is the corresponding reward generated by the mapping function given an arm χ τ,i ∗ , and

This corollary is a directly application of Lemma D.3 and can be obtained with a similar proof as in Corollary D.2.

D.4 Ancillary Lemmas

Applying P k -1 as the training data, we have the following properties for the over-parameterized FC network f ( · ; θ ) after GD.

Lemma D.5. For the constants ρ ∈ (0 , O ( 1 L )) and ξ 1 ∈ (0 , 1) , given the past records P k -1 up to time step k , we suppose m,η 1 , J 1 satisfy the conditions in Theorem 5.2. Then, with probability at least 1 -δ , given a sample-label pair ( x , r ) , we have

Proof. The LHS of the inequality could be written as

Here, we could bound the first term on the RHS with Lemma D.7. Applying Lemma D.8 on the second term, and recalling ∥ θ -θ (0) ∥ 2 ≤ ω , would give

Then, applying the conclusion of Lemma D.6 would lead to

Lemma D.6 (Theorem 1 from [3]) . For any 0 < ξ 1 ≤ 1 , 0 < ρ ≤ O ( 1 L ) . Given the past records P k -1 , suppose m,η 1 , J satisfy the conditions in Theorem 5.2, then with probability at least 1 -δ , we could have

L ( θ ) ≤ ξ 1 after J iterations of GD.

In particular, Lemma D.6 above provides the convergence guarantee for f ( · ; θ ) after certain rounds of GD training on the past records P k -1 .

Lemma D.7 (Lemma 4.1 in [11]) . Assume a constant ω such that O ( m -3 / 2 L -3 / 2 [log( TnL 2 /δ )] 3 / 2 ) ≤ ω ≤ O ( L -6 [log m ] -3 / 2 ) and n training samples. With randomly initialized θ (0) , for parameters θ , θ ′ satisfying ∥ θ -θ (0) ∥ , ∥ θ -θ (0) ∥ ≤ ω , we have

with the probability at least 1 -δ .

Lemma D.8. Assume m,η 1 , J satisfy the conditions in Theorem 5.2 and θ (0) is randomly initialized. Then, with probability at least 1 -δ and given an arm ∥ x ∥ 2 = 1 , we have

Proof. The conclusion (1) is a direct application of Lemma 7.1 in [3]. Suppose the parameters of the L -layer FC network are θ = { θ 1 , . . . , θ L } . For conclusion (2), applying Lemma 7.3 in [3], for each layer θ l ∈ { θ 1 , . . . , θ L } , we have

where D is the diagonal matrix corresponding to the activation function. Then, we could have the conclusion that

Lemma D.9 (Theorem 5 in [3]) . Assume m,η 1 , J satisfy the conditions in Theorem 5.2 and θ (0) being randomly initialized. Then, with probability at least 1 -δ , and for all parameter θ such that ∥ θ -θ (0) ∥ 2 ≤ ω , we have

Lemma D.10. Assume m,η 1 satisfy the condition in Theorem 5.2. For notation brevity, suppose the training sample-label pairs are { x τ , r τ } τ ∈ [ k ] . With the probability at least 1 -δ , we have

Proof. With the notation from Lemma 4.3 in [11], set R = k 3 log( m ) δ , ν = R 2 , and ϵ = LR √ 2 νk . Then, considering the loss function to be L ( θ ) := ∑ τ ∈ [ k ] | f ( x τ ; θ ) -r τ | would complete the proof.

Lemma D.11. Consider a randomly initialized L-layer ReLU fully-connected network f ( · ; θ 0 ) . For any 0 < ξ 2 ≤ 1 , 0 < ρ ≤ O ( 1 L ) . Let there be two sets of training samples P k , P ′ k with the unit-length and the ρ -separateness assumption, and let θ be the trained parameter on P k while θ ′ is the trained parameter on P ′ k . Suppose the conditions in Theorem 5.2 are satisfied. Then, with probability at least 1 -δ , we have

when given a new sample x ∈ R d .

Proof. First, based on the conclusion from Theorem 1 from [3] and regarding the t samples, the trained the parameters satisfy ∥ θ -θ 0 ∥ 2 , ∥ θ ′ -θ 0 ∥ 2 ≤ O ( k 3 ρ √ m log( m )) = ω where θ 0 is the randomly initialized parameter. Then, we could have

w.r.t. the conclusion from Theorem 1 and Theorem 5 of [3]. Then, regarding the Lemma 4.1 from [11], we would have

Therefore, the our target could be reformed as

Substituting the ω with its value would complete the proof.

Corollary D.12. Following a similar settings as in Lemma D.11, consider a randomly initialized L-layer fully-connected network f ( · ; θ 0 ) with Sigmoid activation. For any 0 < ξ 2 ≤ 1 , 0 < ρ ≤ O ( 1 L ) . Let there be two sets of training samples P k , P ′ k with the unit-length and the ρ -separateness assumption, and let θ be the trained parameter on P k while θ ′ is the trained parameter on P ′ k . Suppose the conditions in Theorem 5.2 are satisfied. Then, with probability at least 1 -δ , we have

when given a new sample x ∈ R d .

Proof. This corollary is an intuitive extension of Lemma D.11. Since the result from Theorem 1 of [3] also applies to Lipschitz-smooth (i.e., Sigmoid) activation functions, combining the proof of Lemma D.11 and the result from Lemma 7 in [47] will give the conclusion.

D.5 Regret Bound for Uniform Sampling

Lemma D.13 (Regret Bound for the Uniform Sampling Approach) . When applying the uniform sampling as in most meta-learning frameworks, we denote the corresponding sampled task series as [ ]

Θ ( K ) u refer to the meta-parameters trained with uniform sampling. With ∥ Θ ( K ) u -Θ ( K ) , ∗ ∥ 2 ≤ ω , we have the regret bound for the uniform sampling as

Proof. Here, for the simplicity of notation, we denote Θ = I ( T , Θ ) , and neglect the expectation terms. Note that the difference between adapted meta-parameters and the original meta-parameters is

small enough and can be well-bounded. We will then have

where the two sets of meta-parameters are trained with uniformly sampled tasks and the optimal tasks, and ˜ Θ is used to denote the adapted meta-parameters I ( T , Θ ) for simplicity. With any convex loss function (e.g., L 2 loss or cross-entropy loss) under the over-parameterization settings, we will have the generalization loss being almost convex w.r.t. the meta-parameters as in Lemma D.14, which leads to

where ϵ = O ( ω 4 / 3 L 3 F √ m F log( m F )) > 0 , and ∥ Θ ( K ) , ∗ u -Θ ( K ) u ∥ 2 ≤ ω . Here, the first inequality is due to Lemma D.14 and the convexity of the loss function. The third inequality is due to the upper bound for meta-model gradients ( Lemma D.15). The (i) is due to Lemma D.16 and sufficiently small learning rate η 1 ≤ O ( 1 m F ) . Based on Lemma D.15, we will have ∥∇ Θ L ( x ; Θ ( J ) , ∗ K ) ∥ 2 , ∥∇ Θ L ( x ; Θ ( J ) K ) ∥ 2 ≤ O ( √ m F L F ) . Since we have η 1 , η 2 ≤ O ( 1 m ) , starting from randomly initialized Θ (0) , the parameter shift caused by GD can be upper bounded by ∥ Θ ( K ) , ∗ u -Θ ( K ) u ∥ 2 ≤ ω = O ( K · √ L F m F ) . The implication (iii) is because the loss function L ( · ; · ) has the value range [0, 1].

Here, we notice that the RHS of the regret bound in Lemma D.13 has two terms. Although the second term can be reduced to O (1) with sufficiently large meta-model width m F > O ( Poly ( K,L,ρ -1 )) , the first term tends to grow along with more iterations K and the larger meta-model width m F . The reason is that the radius for the parameter shift during meta-training ω can be as large as O ( 1 √ m F ) , which means that it cannot cancel out the effects of gradient norms, which have the order of O ( √ m F ) . In this case, we will not able to include a m F term to the denominator to scale down the regret with m F , and make the upper bound narrower than 1.

Lemma D.14. Given an arbitrary sample x and its label, let ˜ L ( Θ ) = L ( x ; Θ ) . Suppose m F , η 1 , η 2 satisfy the conditions in Theorem 5.2. With probability at least 1 -O ( KL 2 F ) · exp[ -Ω( m F ω 2 / 3 L F )] over randomness of Θ (0) , for all k ∈ [ K ] , and Θ , Θ ′ satisfying ∥ Θ -Θ (0) ∥ 2 ≤ ω and ∥ Θ ′ -Θ (0) ∥ 2 ≤ ω with ω ≤ O ( L -6 F [log m F ] -3 / 2 ) , it holds uniformly that

with ϵ = O ( ω 4 / 3 L 3 F √ log m F )) being a small constant.

proof. This proof follows an analogous approach as the proof of Lemma 4.2 in [11]. Let ∇ F ˜ L ( Θ ′ ) be the derivative of ˜ L with respective to F ( x ; Θ ) . Then, it holds that |∇ F ˜ L ( Θ ′ ) | ≤ O (1) based on

Lemma D.15. Then, by convexity of ˜ L , we have

where (i) is due to the convexity of the loss function L , (ii) is an application of triangle inequality, and (iii) is the application of and Lemma D.16. Finally, denoting ϵ = O ( ω 4 / 3 L 3 F √ m F log m F ) will complete the proof.

Lemma D.15. Suppose m F , η 1 , η 2 satisfy the conditions in Theorem 5.2. With probability at least 1 -O ( KL F ) · exp( -Ω( m F ω 2 / 3 L F )) over the random initialization, Θ satisfying ∥ Θ -Θ (0) ∥ 2 ≤ ω with ω ≤ O ( L -9 / 2 F [log m F ] -3 ) , it holds uniformly that

Proof. This lemma is a direct application of Lemma 9 of [47] and Lemma B.2, B.3 of [11].

Lemma D.16. Suppose m F , η 1 , η 2 satisfy the conditions in Theorem 5.2. With probability at least 1 -O ( KL F ) · exp( -Ω( m F ω 2 / 3 L F )) , for all t ∈ [ T ] , i ∈ [ k ] , Θ , Θ ′ (or Θ , Θ ′ ) satisfying ∥ Θ -Θ (0) ∥ 2 , ∥ Θ ′ -Θ (0) ∥ 2 ≤ ω with ω ≤ O ( L -9 / 2 F [log m F ] -3 ) , it holds uniformly that

Proof. The proof for this lemma directly follows the proof of Lemma 4.1 in [11] and Lemma 7 in [47].

Peer reviews:
comments: SUMMARY:
This paper proposes scheduling tasks for meta-learning using context bandits.

STRENGTHS:
1. The idea of applying context bandits to select tasks for meta-learning is novel
2. The experimental results are strong, especially when the task distribution is skewed

WEAKNESSES:
1. The method is computationally inefficient in two aspects: a) The arm context requires computing both the adapted parameters and the meta-parameters for every task. While some of these gradients can be reused, the computational cost can still be significant as the number of training tasks increases. b) The networks $f_1,f_2$ used to compute the benefit score take parameters as inputs, which is huge in practice. Although, in Remark 3, the authors mention applying the average pooling to reduce the cost, it is unclear how it is done in practice. Moreover, $f_2$ requires computing the gradients of $f_1$ for two different parameter sets.
2. There is no training time comparison between different methods.

QUESTIONS:
1. What’s the definition of $\mathbf{\theta}$ in Equation (7)? Is it $\theta_1$?
2. Why not simply use $\tilde r_{k,i}+\tilde e_{k,i}$ as the benefit score? In that case, there would be no need to learn $f_2$. Moreover, I believe that, with a little more effort, $f_1$ could be removed either.
3. Could you please detail the average pooling process in Remark 3?

LIMITATIONS:
This work does not involve some potential negative social impacts.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

comments: SUMMARY:
This paper considers the problem of task scheduling in meta-learning. Under the gradient-based meta-learning framework, the authors propose BASS, which uses contextual bandits parameterized by neural networks. The tasks in each batch (arms) are selected in an optimistic manner, with the reward estimated using a validation set and its uncertainty estimated using errors in the neural network predictions. A theoretical analysis of the regret is done in the overparameterized neural network setting. Experiments are done on Mini-ImageNet, CIFAR100, and Drug datasets, showing competitiveness with SOTA and outperforming them when there is task imbalance.

STRENGTHS:
1. The authors provide detailed mathematical motivation for their algorithmic choices.
2. The use of contextual bandits for meta-learning is novel, as far as I know.
3. The authors provide theoretical analysis of their algorithm.
4. Experiments are done on a range of datasets, with ablations.
5. BASS is able to lead to statistically significant improvements when there is task imbalance and is competitive otherwise.

WEAKNESSES:
1. There is no discussion of the computational burden of BASS. It would strengthen the work to provide wall clock times or plots of the performance as a function of time.
2. It is not clear how BASS would scale, as it takes the parameters of one neural network as input into another neural network.

QUESTIONS:
1. How does the computation time of BASS compare to the other algorithms?
2. How would BASS scale with the size of the classifier, e.g. if we were to meta-learn an initialization for 100 layers?

LIMITATIONS:
The paper proposes a new meta-learning algorithm, so I think the authors are ok in not including a broader impact section. The authors do adequately discuss limitations in the appendix.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a task scheduling approach in meta-learning under a contextual bandit framework. The proposed methodology, named BASS, treats each meta-learning task as an arm, prioritizing the selection of these arms according to exploration and exploitation scores. These scores are computed by a trainable neural network with the input being the meta-model's status. The authors provide theoretical proof demonstrating the convergence of the regret bound. Empirical results indicate that BASS outperforms baselines in meta-learning tasks, particularly those involving noisy or skewed datasets.\n\nSTRENGTHS:\n**S1. Clarity**\n\nThe authors present their work with high clarity. The notations, a critical aspect in avoiding confusion in this complex subject, are well-defined. Also, the authors have strategically placed numerous remarks in the main section, which significantly assist readers in comprehending the principal concepts.\n\n**S2. Advantages from the use of bandits**\n\nThe application of the contextual bandit in place of the previously used greedy update offers multiple strengths.\n-\tIt eliminates the need for unstable bi-level optimization\n-\tIn naturally manages the balance between exploration and exploitation. The way the extrinsic and intrinsic rewards are formulated is both intriguing and intuitive. Notably, the exploration factor, which takes into account both task uncertainty and difficulty, stands out as a significant advantage of this work.\n-\tAs a consequence of the bandit application, the regret bound can be theoretically limited. The authors have put in considerable effort the substantiate their theoretical analysis in the appendix.\n\nWEAKNESSES:\n**W1. Experimental results**\n\nWhile the proposed method demonstrates significant enhancements over baseline models for noisy or skewed training tasks, it appears to yield marginal or negligible improvement on standard datasets, as acknowledged by the authors. Incorporating additional data regarding computation cost could provide valuable insight, especially when comparing the effect of noise or skewness on different datasets and assessing the increased computation demand posed by BASS.\n\n**Minor comment** \n\nIt is recommended to incorporate the 'Related Works' section into the main body of the manuscript instead of including it in the appendix. Given that all submissions adhere to the same page limit, placing this section in the appendix can disrupt the contextualization of the work within the existing literature.\n\n**Acknowledgment Following Rebuttal**\n\nThe author's rebuttal successfully addressed concerns about the experiment. The inclusion of additional results on average-pooling is particularly noteworthy.\n\nQUESTIONS:\n**Q1.** The explanation provided from Line 179 to Line 185 is somewhat unclear. Could the authors provide further clarification on this matter?\n\n**Q2.** The notations used in Line 108 and Line 326 appear to be inconsistent. Could these be typos?\n\n**Q3.** It's great that BASS employs all parameters of the meta-model to infer its status. However, BASS uses average pooling to reduce millions of parameters to about 50. Can this average-pooled feature truly encompass all necessary information about the status of the meta-model?\n\n**Q4.** In relation to Q3, the use of average-pooling seems to be due to the high dimensionality of the input. Have the authors considered employing other encoding techniques, such as an MLP with smaller hidden layers?\n\nLIMITATIONS:\nThe authors have adequately addressed their limitations regarding the computational cost and the marginal improvement for standard dataset in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed a task scheduling framework BASS based on the status of meta-model based on contextual bandits setting. BASS addressed the performance bottleneck of meta-models by balancing exploitation and exploration and handled the data scarcity in the early stages of meta-training iterations with planning future meta-learning iteration strategies. The experimental and theoretical analysis showed the effectiveness of BASS.\n\nSTRENGTHS:\n1 [Writing & Presentation] This paper is well-written and easy to follow. \n\n2 [Motivation] The previous work about task schedulers aimed to improve meta-training strategies based on the various pre-defined criteria and assumptions, ignoring the global knowledge. It may result in a sub-optimal meta-model affected by noise perturbation or skewed task distributions. Motivated by this limitation, the authors proposed a novel framework to solve the meta-learning task scheduling problem. \n\n3 [Contribution] Different from the existing methods that exploit the current/local knowledge for greedy scheduling, the proposed method leveraged a novel method to adaptively learn the relationship between the meta-model parameters and meta-model generalization ability. \n\n4 [Experiments] The experiments demonstrated the effectiveness of BASS compared with seven strong basslines on three real datasets. \nThe experimental results showed superiority in terms of accuracy and efficiency. Besides, BASS can explore the ‘tail’ task and enjoy a good performance in the ensemble inference setting, which further enhanced the generalization ability of meta-learning models.\n\nWEAKNESSES:\n1 [Algorithm] In Algorithm 1. How about the detailed setting of the initialization? Intuitively, the meta-model parameters would rely on the various data distributions. \n\n2 [ Theory] In Section 4, the $k$-iteration regret is the assumption of Theorem 4.2. However, the detailed setting of regret was not given. For example, the explanation of independence of regret bounds is not well illustrated. \n\n3 [Experiment] In Section 5.1, the experimental analysis was on three datasets, the Drug dataset is a textual dataset, and both ImageNet and CIFAR are visual/image datasets. Since the proposed method focused on the generalization ability, it would be better to verify the performance on various categories of datasets. \n\n4 [Experiment] In Table.1, the performance for 1-shot of BASS increased slightly. Does it mean the meta-model is invalid? If so, what is the main point of BASS leading to this result?\n\nQUESTIONS:\n1 [Refer to Weakness 1] How about the assumption on various datasets? Does the assumption independent of each other? Will the results vary with different initialization? \n\n2 [Refer to Weakness 1] What are the constraints or the policy of batch size B? How will the performance vary with different settings of B?\n \n3 [Refer to Weakness 2] How about the bound of regret? How about the distribution of arms? Whether the performance change with different settings/assumptions of regret/arms?\n\n4 [Refer to Weakness 3] How about the performance on various kinds of datasets? For example, image dataset, textual dataset, vision dataset, or dataset represented by deep features. Furthermore, the data dimension, data sparsity, even the image resolution may also affect the performance.\n\nLIMITATIONS:\nThe target problem is interesting and worth studying. However, the detailed explanation of theory analysis needs to be clarified, and more experiments should be conducted to verify the superiority of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper proposes to adaptively sample the meta-training tasks by optimizing the task scheduling strategy based on the status of the meta-model. The proposed method treats task scheduling as a contextual multi-arm bandit problem with a reward function balancing the exploitation and exploration. The authors provide theoretical analysis about the regret bound of the proposed method and conduct experiments on real data sets to show the effectiveness.

STRENGTHS:
**DISCLAIMER:** I have not checked the proof thoroughly and cannot verify the correctness of the theorems.

* The presentation of the paper is clean and clear.
* It seems to be novel to formulate the curriculum learning as a contextual bandit in the context of meta-learning.

WEAKNESSES:
* The authors claim the proposed method can deal with the data scarcity problem at the early stage of meta-learning. However, I cannot further elaboration anywhere in the paper nor any experimental results supporting this claim.
* Since the proposed method involves extra computation efforts, the quantitative comparison of computation cost is not included in the experiment results.
* If the proposed method outperforms the existing method, especially when the data is noisy and skewed, the authors should include experiments with different levels of noise $\epsilon$ and different levels of skewness to show the trends of improvements.

QUESTIONS:
* How is the K-round regret bound in Theorem 4.2 compared with the existing work?
* How is the performance of the proposed method compared with the existing method when the data is noise-free and unskewed?

LIMITATIONS:
The authors discuss the limitation in the appendix.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06