Uncertainty-Aware Meta-Learning for Multimodal Task Distributions

Cesar Almecija, Apoorva Sharma, Young-Jin Park, Navid Azizan

Keywords: Meta-learning, Bayesian inference, neural network linearization, uncertainty estimation, Gaussian Process, NTK

TL;DR: We present a novel meta-learning algorithm that makes probabilistic predictions efficiently, detects out-of-distribution context data, and performs well on heterogeneous, multimodal task distributions.

Abstract: Meta-learning 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 this work, we present UNLIMITD (uncertainty-aware meta-learning for multimodal6 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.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2210.01881/code)