Learning to Balance: Bayesian Meta-Learning for Imbalanced and Out-of-distribution Tasks

1 INTRODUCTION . Despite the success of deep learning in many real-world tasks such as visual recognition and machine translation , such good performances are achievable at the availability of large training data , and many fail to generalize well in small data regimes . To overcome this limitation of conventional deep learning , recently , researchers have explored meta-learning ( Schmidhuber , 1987 ; Thrun & Pratt , 1998 ) approaches , whose goal is to learn a model that generalizes well over distribution of tasks , rather than instances from a single task , in order to utilize the obtained meta-knowledge across tasks to compensate for the lack of training data for each task . However , so far , most existing meta-learning approaches ( Santoro et al. , 2016 ; Vinyals et al. , 2016 ; Snell et al. , 2017 ; Ravi & Larochelle , 2017 ; Finn et al. , 2017 ; Li et al. , 2017 ) have only targeted an artificial scenario where all tasks participating in the multi-class classification problem have equal number of training instances per class . Yet , this is a highly restrictive setting , as in real-world scenarios , tasks that arrive at the model may have different training instances ( task imbalance ) , and within each task , the number of training instances per class may largely vary ( class imbalance ) . Moreover , the new task may come from a distribution that is different from the task distribution the model has been trained on ( out-of-distribution task ) ( See ( a ) of Figure 1 ) . Under such a realistic setting , the meta-knowledge may have a varying degree of utility to each task . Tasks with small number of training data , or close to the tasks trained in meta-training step may want to rely mostly on meta-knowledge obtained over other tasks , whereas tasks that are out-of-distribution or come with more number of training data may obtain better solutions when trained in a task-specific ∗Equal contribution manner . Furthermore , for multi-class classification , we may want to treat the learning for each class differently to handle class imbalance . Thus , to optimally leverage meta-learning under various imbalances , it would be beneficial for the model to task- and class-adaptively decide how much to use from the meta-learner , and how much to learn specifically for each task and class . To this end , we propose a novel Bayesian meta-learning framework , which we refer to as Bayesian Task-Adaptive Meta-Learning ( Bayesian TAML ) , that learns variables to adaptively balance the effect of meta- and task-specific learning . Specifically , we first obtain set-representations for each task , which are learned to convey useful statistics about the task or class distribution , such as mean , variance , and cardinality ( the number of elements in the set ) , and then learn the distribution of three balancing variables a function of the set : 1 ) task-dependent learning rate multiplier , which decides how far away to deviate from the meta-knowledge , when performing task-specific learning . Tasks with higher shots could benefit from taking gradient steps afar , while tasks with few shots may need to stay close to the initial parameter . 2 ) class-dependent learning rate , which decides how much information to use from each class , to automatically handle class imbalance where the number of instances per class can largely vary . 3 ) task-dependent modulator for initial model parameter , which modifies the shared initialization for each task , such that each task can decide how much and what to use from the shared initial model parameter and what to ignore based on its set representation . This is especially useful when handling out-of-distribution task , which may need to ignore some of the meta-knowledge . We validate our model on CIFAR-FS and miniImageNet dataset , as well as a new dataset that consists of heterogeneous datasets , under a scenario where every class in each episode can have any number of shots , that leads to task and class imbalance , and where the dataset at meta-test time is different from that of meta-training time . The experimental results show that our Bayesian TAML significantly improves the performance over the existing approaches under these realistic scenarios . Further analysis of each component reveals that the improvement is due to the effectiveness of the balancing terms for handling task and class imbalance , and out-of-distribution tasks . To summarize , our contribution in this work is threefold : • We consider a novel problem of meta-learning under a realistic task distribution , where the number of instances across classes and tasks could largely vary , or the unseen task at the meta-test time is largely different from the seen tasks . • For effective meta-learning with such imbalances , we propose a Bayesian task-adaptive meta-learning ( Bayesian TAML ) framework that can adaptively adjust the effect of the meta-learner and the task-specific learner , differently for each task and class . • We validate our model on realistic imbalanced few-shot classification tasks with a varying number of shots per task and class and show that it significantly outperforms existing meta-learning models . 2 RELATED WORK . Meta-learning Meta-learning ( Schmidhuber , 1987 ; Thrun & Pratt , 1998 ) is an approach to learn a model to generalize over a distribution of task . The approaches in general can be categorized into either memory-based , metric-based , and optimization-based methods . A memory-based approach ( Santoro et al. , 2016 ) learns to store correct instance and label into the same memory slot and retrieve it later , in a task-generic manner . Metric-based approaches learn a shared metric space ( Vinyals et al. , 2016 ; Snell et al. , 2017 ) . Snell et al . ( 2017 ) defines the distance between the instance and the class prototype , such that the instances are closer to their correct prototypes than to others . As for optimization-based meta-learning , MAML ( Finn et al. , 2017 ) learns a shared initialization parameter that is optimal for any tasks within few gradient steps from the initial parameter . Meta-SGD ( Li et al. , 2017 ) improves upon MAML by learning the learning rate differently for each parameter . For effective learning of a meta-learner , meta-learning approaches adopt the episodic training strategy ( Vinyals et al. , 2016 ) which trains and evaluates a model over a large number of tasks , which are called meta-training and meta-test phase , respectively . However , existing approaches only consider an artificial scenario which samples the classification of classes with exactly the same number of training instances , both within each episode and across episodes . On the other hand , we consider a more challenging scenario where the number of shots per class and task could vary at each episode , and that the task given at meta-test time could be an out-of-distribution task . Task-adaptive meta-learning The goal of learning a single meta-learner that works well for all tasks may be overly ambitious and leads to suboptimal performances for each task . Thus recent approaches adopt task-adaptively modified meta-learning models . Oreshkin et al . ( 2018 ) proposed to learn the temperature scaling parameter to work with the optimal similarity metric . Qiao et al . ( 2018 ) also suggested a model that generates task-specific parameters for the network layers , but it only trains with many-shot classes , and implicitly expects generalization to few-shot cases . Rusu et al . ( 2019 ) proposed a network type task-specific parameter producer , and Lee & Choi ( 2018 ) proposed to differentiate the network weights into task-shared and task-specific weights . Our model also aims to obtain task-specific parameter for each task , but is rather focused on learning how to balance between the meta-learning and task-/class-specific learning . Probabilistic meta-learning Recently , a probabilistic version of MAML has been proposed ( Finn et al. , 2018 ) , where they interpret a task-specific gradient update as a posterior inference process under variational inference framework . Kim et al . ( 2018 ) proposed Bayesian MAML with a similar motivation but with a stein variational inference framework and chaser loss . Gordon et al . ( 2019 ) proposed a probabilistic meta-learning framework where the parameter for a novel task is rapidly estimated under decision theoretic framework , given a set representation of a task . The motivation behind these works is to represent the inherent uncertainty in few-shot classification tasks . Our model also uses Bayesian modeling , but it focuses on leveraging the uncertainties of the meta-learner and the gradient-direction in order to balance between meta- and task- or class-specific learning . 3 LEARNING TO BALANCE . We first introduce notations and briefly recap the model-agnostic meta-learning ( MAML ) by Finn et al . ( 2017 ) . Suppose a task distribution p ( τ ) that randomly generates task τ consisting of a training set Dτ = { Xτ , Yτ } and a test set D̃τ = { X̃τ , Ỹτ } . Then , the goal of MAML is to meta-learn the initial model parameter θ as a meta-knowledge to generalize over the task distribution p ( τ ) , such that we can easily obtain the task-specific predictor θτ in a single ( or a few ) gradient step from the initial θ . Toward this goal , MAML optimizes the following gradient-based meta-learning objective : min θ ∑ τ∼p ( τ ) L ( θ − α∇θL ( θ ; Dτ ) ; D̃τ ) ( 1 ) where α denotes stepsize and L denotes empirical loss such as negative log-likelihood of observations . Note that by meta-learning the initial point θ , the task-specific predictor θτ = θ − α∇θL ( θ ; Dτ ) can minimize the test loss L ( · ; D̃τ ) even with Dτ consisting of only a few samples . We can easily extend the Eq . ( 1 ) , such that we obtain θτ with more than one inner-gradient steps from the initial θ . However , the existing MAML framework has the following limitations that prevent the model from efficiently solving real-world problems involving class/task imbalance and out-of-distribution tasks . 1 . Class imbalance . MAML does not provide any framework to handle class imbalance within each task . Therefore , classes with large number of training instances ( head classes ) may dominate the task-specific learning during the inner-gradient steps , yielding low performance on classes with fewer shots ( tail classes ) . 2 . Task imbalance . The model has a fixed number of inner-gradient steps and stepsize α across all tasks , which prevents the model from adaptively deciding how much to resort to the meta-knowledge or how much to learn from the given dataset , depending on the number of the training examples per task . 3 . Out-of-distribution tasks . The model assumes that the initial model parameter θ will be equally useful for the unseen tasks , but for unseen tasks that are significantly different from the previously seen tasks , the initial parameter may be less useful . 3.1 TASK-ADAPTIVE META-LEARNING ( TAML ) . As shown in Figure 1 for the concepts , we introduce three balancing variables ωτ , γτ , zτ to tackle each problem mentioned above . How to generate these variables will be described in Section 4 . Also , see the experimental section for how to generate the realistic tasks with class and task imbalance . Tackling class imbalance . To handle class imbalance , we vary the learning rate of class-specific gradient for each inner-optimization step . Specifically , for class c = 1 , . . . , C , we introduce a set of class-specific scalars ωτ = ( ωτ1 , . . . , ω τ C ) ∈ [ 0 , 1 ] C , which are multiplied to each of the class-specific gradients∇θL ( θ ; Dτ1 ) , . . . , ∇θL ( θ ; DτC ) , where Dτc is the set of instances and labels for class c. We expect ωτc to be large for tail-classes with small number of training instances , such that they could be considered more in the inner-optimization steps . We generate ωτ with Softmax function , and denote the input to the function as ω̃τ . Tackling task imbalance . To control whether the model parameter for the current task stays close to the initial parameter or deviate far from it , we introduce task-dependent learning-rate multipliers γτ = ( γτ1 , . . . , γ τ L ) ∈ [ 0 , ∞ ) L , such that for each layer l = 1 , . . . , L , the learning rate becomes γτ1α , γ τ 2α , . . . , γ τ Lα 1 . We expect γτ to be large for large tasks , such that they rely more on taskspecific updates , while small tasks use small γτ to benefit more from the meta-knowledge . To amplify the step-size difference between large and small tasks , we generate γτ with an exponential function , and denote the input to the function as γ̃τ . Tackling out-of-distribution tasks . Finally , we introduce zτ which modulates the initial parameter θ for each task . We expect zτ to learn to relocate the initial θ to a new starting point , such that out-of-distribution ( OOD ) tasks can deviate much from the shared initialization θ if the current initialization is suboptimal for the given task . Specifically , we use zτ = 1+ z̃τ for the channel of the convolutional network weights and zτ = z̃τ for the biases , which modify the initial parameter as follows : θ0 ← θ ◦ zτ for the weights and θ0 ← θ + zτ for the biases , where θ0 denotes the new initialization modulated by zτ . We denote this operation as θ ∗ zτ , which is similar to task-dependent modulation of batch normalization parameters ( Oreshkin et al. , 2018 ; Requeima et al. , 2019 ) . A unified framework . Finally , we assemble all these components together into a single unified framework . With a slight abuse of notations , the update rule can be written as follows : θ0 = θ ∗ zτ , ( 2 ) θk = θk−1 − γτ ◦α ◦ C∑ c=1 ωτc∇θk−1L ( θk−1 ; Dτc ) for k = 1 , . . . , K ( 3 ) where α is a multi-dimensional global learning rate vector that is learned ( Li et al. , 2017 ) , and the multiplication operator ◦ is appropriately defined . The last step θK corresponds to the task-specific predictor θτ .

Summary: This paper introduces a mechanism for gradient-based meta-learning models for few-shot classification to be able to adapt to diverse tasks that are imbalanced and heterogeneous. In particular, each encountered task may have varying numbers of shots (task imbalance) and even within each task, different classes may have different numbers of shots (class imbalance). Further, test tasks might come from a different distribution than the training tasks. They propose to handle this scenario by introducing three new types of variables which control different facets of the degree and type of task adaptation, allowing to decide how much to reuse meta-learned knowledge versus new knowledge acquired from the training set of the given task.