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: The paper proposes a Bayesian approach for meta learning in settings were the tasks might be OOD or have imbalanced class distribution. The proposed approach has 3 task-specific balancing variables with a prior and an inference network. Using an amortized inference scheme, the model unifies the meta-learning objective loss with the lower bound of probabilistic model marginal likelihood.