UNCERTAINTY-AWARE META-LEARNING FOR MULTI-MODAL TASK DISTRIBUTIONS

Abstract

Meta-learning or learning to learn is a popular approach for learning new tasks with limited data (i.e., few-shot learning) by leveraging the commonalities among different tasks. However, meta-learned models can perform poorly when context data is limited, or when data is drawn from an out-of-distribution (OoD) task. Especially in safety-critical settings, this necessitates an uncertainty-aware approach to meta-learning. In addition, the often multimodal nature of task distributions can pose unique challenges to meta-learning methods. In this work, we present UNLIMITD (uncertainty-aware meta-learning for multimodal task distributions), a novel method for meta-learning that (1) makes probabilistic predictions on in-distribution tasks efficiently, (2) is capable of detecting OoD context data at test time, and (3) performs on heterogeneous, multimodal task distributions. To achieve this goal, we take a probabilistic perspective and train a parametric, tuneable distribution over tasks on the meta-dataset. We construct this distribution by performing Bayesian inference on a linearized neural network, leveraging Gaussian process theory. We demonstrate that UNLIMITD's predictions compare favorably to, and outperform in most cases, the standard baselines, especially in the low-data regime. Furthermore, we show that UNLIMITD is effective in detecting data from OoD tasks. Finally, we confirm that both of these findings continue to hold in the multimodal task-distribution setting.

1. INTRODUCTION

Learning to learn is essential in human intelligence but is still a wide area of research in machine learning. Meta-learning has emerged as a popular approach to enable models to perform well on new tasks using limited data. It involves first a meta-training process, when the model learns valuable features from a set of tasks. Then, at test time, using only few datapoints from a new, unseen task, the model (1) adapts to this new task (i.e., performs few-shot learning with context data), and then (2) infers by making predictions on new, unseen query inputs from the same task. A popular baseline for meta-learning, which has attracted a large amount of attention, is Model-Agnostic Meta-Learning (MAML) (Finn et al., 2017) , in which the adaptation process consists of fine-tuning the parameters of the model via gradient descent. However, meta-learning methods can often struggle in several ways when deployed in challenging real-world scenarios. First, when context data is too limited to fully identify the test-time task, accurate prediction can be challenging. As these predictions can be untrustworthy, this necessitates the development of meta-learning methods that can express uncertainty during adaptation (Yoon et al., 2018; Harrison et al., 2018) . In addition, meta-learning models may not successfully adapt to "unusual" tasks, i.e., when test-time context data is drawn from an out-of-distribution (OoD) task not well represented in the training dataset (Jeong & Kim, 2020; Iwata & Kumagai, 2022) . Finally, special care has to be taken when learning tasks that have a large degree of heterogeneity. An important example is the case of tasks with a multimodal distribution, i.e., when there are no common features shared across all the tasks, but the tasks can be broken down into subsets (modes) in a way that the ones from the same subset share common features (Vuorio et al., 2019) . Our contributions. We present UNLIMITD (uncertainty-aware meta-learning for multimodal task distributions), a novel meta-learning method that leverages probabilistic tools to address the aforementioned issues. Specifically, UNLIMITD models the true distribution of tasks with a learnable distribution constructed over a linearized neural network and uses analytic Bayesian inference to perform uncertainty-aware adaption. We present three variants (namely, UNLIMITD-I, UNLIM-

