# Uncertainty-Aware Meta-Learning for Multimodal Task Distributions

Cesar Almecija1, Apoorva Sharma2 & Navid Azizan1
1 Massachusetts Institute of Technology
2 NVIDIA Research
<EMAIL_ADDRESS><EMAIL_ADDRESS>
This work was completed prior to Apoorva starting at NVIDIA Research.

###### 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. 111An
implementation is available at https://github.com/azizanlab/unlimitd

##  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, UnLiMiTD-R, and UnLiMiTD-F) that reflect a trade-off
between learning a rich prior distribution over the weights and maintaining
the full expressivity of the network; we show that UnLiMiTD-F strikes a
balance between the two, making it the most appealing variant. Finally, we
demonstrate that (1) our method allows for efficient probabilistic predictions
on in-distribution tasks, that compare favorably to, and in most cases
outperform, the existing baselines, (2) it is effective in detecting context
data from OoD tasks at test time, and that (3) both these findings continue to
hold in the multimodal task-distribution setting.

The rest of the paper is organized as follows. Section 2 formalizes the
problem. Section 3 presents background information on the linearization of
neural networks and Bayesian linear regression. We detail our approach and its
three variants in Section 4. We discuss related work in detail in Section 5.
Finally, we present our experimental results concerning the performance of
UnLiMiTD in Section 6 and conclude in Section 7.

Figure 1:  The true task distribution $p(f)$ can be multimodal, i.e.,
containing multiple clusters of tasks (e.g., lines and sines). Our approach
UnLiMiTD fits $p(f)$ with a parametric, tuneable distribution
$\tilde{p}_{\xi}(f)$ yielded by Bayesian linear regression on a linearized
neural network.

##  2 Problem statement

A task $\mathcal{T}^{i}$ consists of a function $f_{i}$ from which data is
drawn. At test time, the prediction steps are broken down into (1) adaptation,
that is identifying $f_{i}$ using $K$ context datapoints
$({\bm{X}}^{i},{\bm{Y}}^{i})$ from the task, and (2) inference, that is making
predictions for $f_{i}$ on the query inputs ${\bm{X}}^{i}_{*}$. Later the
predictions can be compared with the query ground-truths ${\bm{Y}}^{i}_{*}$ to
estimate the quality of the prediction, for example in terms of mean squared
error (MSE). The meta-training consists in learning valuable features from a
cluster of tasks, which is a set of similar tasks (e.g., sines with different
phases and amplitudes but same frequency), so that at test time the
predictions can be accurate on tasks from the same cluster. We take a
probabilistic, functional perspective and represent a cluster by $p(f)$, a
theoretical distribution over the function space that describes the
probability of a task belonging to the cluster. Learning $p(f)$ is appealing,
as it allows for performing OoD detection in addition to making predictions.
Adaptation amounts to computing the conditional distribution given test
context data, and one can obtain an uncertainty metric by evaluating the
negative log-likelihood (NLL) of the context data under $p(f)$.

Thus, our goal is to construct a parametric, learnable functional distribution
$\tilde{p}_{\xi}(f)$ that approaches the theoretical distribution $p(f)$, with
a structure that allows tractable conditioning and likelihood computation,
even in deep learning contexts. In practice, however, we are not given $p(f)$,
but only a meta-training dataset $\mathcal{D}$ that we assume is sampled from
$p(f)$:
$\mathcal{D}=\\{(\widetilde{{\bm{X}}}^{i},\widetilde{{\bm{Y}}}^{i})\\}_{i=1}^{N}$,
where $N$ is the number of tasks available during training, and
$(\widetilde{{\bm{X}}}^{i},\widetilde{{\bm{Y}}}^{i})\sim\mathcal{T}^{i}$ is
the entire pool of data from which we can draw subsets of context data
$({\bm{X}}^{i},{\bm{Y}}^{i})$. Consequently, in the meta-training phase, we
aim to optimize $\tilde{p}_{\xi}(f)$ to capture properties of $p(f)$, using
only the samples in $\mathcal{D}$.

Once we have $\tilde{p}_{\xi}(f)$, we can evaluate it both in terms of how it
performs for few-shot learning (by comparing the predictions with the ground
truths in terms of MSE), as well as for OoD detection (by measuring how well
the NLL of context data serves to classify in-distribution tasks against OoD
tasks, measured via the AUC-ROC score).

##  3 Background

###  3.1 Bayesian linear regression and Gaussian Processes

Efficient Bayesian meta-learning requires a tractable inference process at
test time. In general, this is only possible analytically in a few cases. One
of them is the Bayesian linear regression with Gaussian noise and a Gaussian
prior on the weights. Viewing it from a nonparametric, functional approach,
this model is equivalent to a Gaussian process (GP) (Rasmussen & Williams,
2005).

Let ${\bm{X}}=({\bm{x}}_{1},\dots,{\bm{x}}_{K})\in\mathbb{R}^{{N_{x}}\times
K}$ be a batch of $K$ ${N_{x}}$-dimensional inputs, and let
${\bm{y}}=({\bm{y}}_{1},\dots,{\bm{y}}_{K})\in\mathbb{R}^{{N_{y}}K}$ be a
vectorized batch of ${N_{y}}$-dimensional outputs. In the Bayesian linear
regression model, these quantities are related according to
${\bm{y}}=\phi({\bm{X}})^{\top}\hat{{\bm{\theta}}}+\varepsilon\in\mathbb{R}^{{N_{y}}K}$
where $\hat{{\bm{\theta}}}\in\mathbb{R}^{P}$ are the weights of the model, and
the inputs are mapped via $\phi:\mathbb{R}^{{N_{x}}\times
K}\rightarrow\mathbb{R}^{P\times{N_{y}}K}$. Notice how this is a
generalization of the usual one-dimensional linear regression (${N_{y}}=1$).

If we assume a Gaussian prior on the weights
$\hat{{\bm{\theta}}}\sim\mathcal{N}({\bm{\mu}},{\bm{\Sigma}})$ and a Gaussian
noise $\varepsilon\sim\mathcal{N}(\bm{0},{\bm{\Sigma}}_{\varepsilon})$ with
${\bm{\Sigma}}_{\varepsilon}=\sigma_{\varepsilon}^{2}{\bm{I}}$, then the model
describes a multivariate Gaussian distribution on ${\bm{y}}$ for any
${\bm{X}}$. Equivalently, this means that this model describes a GP
distribution over functions, with mean and covariance function (or kernel)

$\displaystyle\begin{split}{\bm{\mu}}_{\text{prior}}({\bm{x}}_{t})&=\phi({\bm{x}}_{t})^{\top}{\bm{\mu}},\\\
\text{cov}_{\text{prior}}({\bm{x}}_{t_{1}},{\bm{x}}_{t_{2}})&=\phi({\bm{x}}_{t_{1}})^{\top}{\bm{\Sigma}}\phi({\bm{x}}_{t_{2}})+{\bm{\Sigma}}_{\varepsilon}=:k_{\bm{\Sigma}}({\bm{x}}_{t_{1}},{\bm{x}}_{t_{2}})+{\bm{\Sigma}}_{\varepsilon}.\end{split}$
(1)

This GP enables tractable computation of the likelihood of any batch of data
$({\bm{X}},{\bm{Y}})$ given this distribution over functions. The structure of
this distribution is governed by the feature map $\phi$ and the prior over the
weights, specified by ${\bm{\mu}}$ and ${\bm{\Sigma}}$.

This distribution can also easily be conditioned to perform inference. Given a
batch of data $({\bm{X}},{\bm{Y}})$, the posterior predictive distribution is
also a GP, with an updated mean and covariance function

$\displaystyle\begin{split}{\bm{\mu}}_{\text{post}}({\bm{x}}_{t_{*}})&=k_{\bm{\Sigma}}({\bm{x}}_{t_{*}},{\bm{X}})\left(k_{\bm{\Sigma}}({\bm{X}},{\bm{X}})+{\bm{\Sigma}}_{\varepsilon}\right)^{-1}{\bm{Y}},\\\
\text{cov}_{\text{post}}({\bm{x}}_{{t_{1}}_{*}},{\bm{x}}_{{t_{2}}_{*}})&=k_{\bm{\Sigma}}({\bm{x}}_{{t_{1}}_{*}},{\bm{x}}_{{t_{2}}_{*}})-k_{\bm{\Sigma}}({\bm{x}}_{{t_{1}}_{*}},{\bm{X}})\left(k_{\bm{\Sigma}}({\bm{X}},{\bm{X}})+{\bm{\Sigma}}_{\varepsilon}\right)^{-1}k_{\bm{\Sigma}}({\bm{X}},{\bm{x}}_{{t_{2}}_{*}}).\end{split}$
(2)

Here, ${\bm{\mu}}_{\text{post}}({\bm{X}}_{*})$ represents our model’s adapted
predictions for the test data, which we can compare to ${\bm{Y}}_{*}$ to
evaluate the quality of our predictions, for example, via mean squared error
(assuming that test data is clean, following Rasmussen & Williams (2005)). The
diagonal of $\text{cov}_{\text{post}}({\bm{X}}_{*},{\bm{X}}_{*})$ can be
interpreted as a per-input level of confidence that captures the ambiguity in
making predictions with only a limited amount of context data.

###  3.2 The linearization of a neural network yields an expressive linear
regression model

As discussed, the choice of feature map $\phi$ plays an important role in
specifying a linear regression model. In the deep learning context, recent
work has demonstrated that the linear model obtained when linearizing a deep
neural network with respect to its weights at initialization, wherein the
Jacobian of the network operates as the feature map, can well approximate the
training behavior of wide nonlinear deep neural networks (Jacot et al., 2018;
Azizan et al., 2021; Liu et al., 2020; Neal, 1996; Lee et al., 2018).

Let $f$ be a neural network
$f:\left({\bm{\theta}},{\bm{x}}_{t}\right)\mapsto{\bm{y}}_{t}$, where
${\bm{\theta}}\in\mathbb{R}^{P}$ are the parameters of the model,
${\bm{x}}\in\mathbb{R}^{{N_{x}}}$ is an input and
${\bm{y}}\in\mathbb{R}^{{N_{y}}}$ an output. The linearized network (w.r.t.
the parameters) around ${\bm{\theta}}_{0}$ is

$f({\bm{\theta}},{\bm{x}}_{t})-f({\bm{\theta}}_{0},{\bm{x}}_{t})\approx{\bm{J}}_{\bm{\theta}}(f)({\bm{\theta}}_{0},{\bm{x}}_{t})({\bm{\theta}}-{\bm{\theta}}_{0}),$

where
${\bm{J}}_{\bm{\theta}}(f)(\cdot,\cdot):\mathbb{R}^{P}\times\mathbb{R}^{N_{x}}\rightarrow\mathbb{R}^{{N_{y}}\times
P}$ is the Jacobian of the network (w.r.t. the parameters).

In the case where the model accepts a batch of $K$ inputs
${\bm{X}}=({\bm{x}}_{1},\dots,{\bm{x}}_{K})$ and returns
${\bm{Y}}=({\bm{y}}_{1},\dots,{\bm{y}}_{K})$, we generalize $f$ to
$g:\mathbb{R}^{P}\times\mathbb{R}^{{N_{x}}\times
K}\rightarrow\mathbb{R}^{{N_{y}}\times K}$, with
${\bm{Y}}=g({\bm{\theta}},{\bm{X}})$. Consequently, we generalize the
linearization:

$g({\bm{\theta}},{\bm{X}})-g({\bm{\theta}}_{0},{\bm{X}})\approx{\bm{J}}({\bm{\theta}}_{0},{\bm{X}})({\bm{\theta}}-{\bm{\theta}}_{0}),$

where ${\bm{J}}(\cdot,\cdot):\mathbb{R}^{P}\times\mathbb{R}^{{N_{x}}\times
K}\rightarrow\mathbb{R}^{{N_{y}}K\times P}$ is a shorthand for
${\bm{J}}_{\bm{\theta}}(g)(\cdot,\cdot)$. Note that we have implicitly
vectorized the outputs, and throughout the work, we will interchange the
matrices $\mathbb{R}^{{N_{y}}\times K}$ and the vectorized matrices
$\mathbb{R}^{{N_{y}}K}$.

This linearization can be viewed as the ${N_{y}}K$-dimensional linear
regression

${\bm{z}}=\phi_{{\bm{\theta}}_{0}}({\bm{X}})^{\top}\hat{{\bm{\theta}}}\in\mathbb{R}^{{N_{y}}K},$
(3)

where the feature map
$\phi_{{\bm{\theta}}_{0}}(\cdot):\mathbb{R}^{{N_{x}}\times
K}\rightarrow\mathbb{R}^{P\times{N_{y}}K}$ is the transposed Jacobian
${\bm{J}}({\bm{\theta}}_{0},\cdot)^{\top}$. The parameters of this linear
regression $\hat{{\bm{\theta}}}=\left({\bm{\theta}}-{\bm{\theta}}_{0}\right)$
are the correction to the parameters chosen as the linearization point.
Equivalently, this can be seen as a kernel regression with the kernel
$k_{{\bm{\theta}}_{0}}({\bm{X}}_{1},{\bm{X}}_{2})={\bm{J}}({\bm{\theta}}_{0},{\bm{X}}_{1}){\bm{J}}({\bm{\theta}}_{0},{\bm{X}}_{2})^{\top}$,
which is commonly referred to as the Neural Tangent Kernel (NTK) of the
network. Note that the NTK depends on the linearization point
${\bm{\theta}}_{0}$. Building on these ideas, Maddox et al. (2021) show that
the NTK obtained via linearizing a DNN after it has been trained on a task
yields a GP that is well-suited for adaptation and fine-tuning to new, similar
tasks. Furthermore, they show that networks trained on similar tasks tend to
have similar Jacobians, suggesting that neural network linearization can yield
an effective model for multi-task contexts such as meta-learning. In this
work, we leverage these insights to construct our parametric functional
distribution $\tilde{p}_{\xi}(f)$ via linearizing a neural network model.

##  4 Our approach: UnLiMiTD

In this section, we describe our meta-learning algorithm UnLiMiTD and the
construction of a parametric functional distribution $\tilde{p}_{\xi}(f)$ that
can model the true underlying distribution over tasks $p(f)$. First, we focus
on the single cluster case, where a Gaussian process structure on
$\tilde{p}_{\xi}(f)$ can effectively model the true distribution of tasks, and
detail how we can leverage meta-training data $\mathcal{D}$ from a single
cluster of tasks to train the parameters $\xi$ of our model. Next, we will
generalize our approach to the multimodal setting, with more than one cluster
of tasks. Here, we construct $\tilde{p}_{\xi}(f)$ as a mixture of GPs and
develop a training approach that can automatically identify the clusters
present in the training dataset without requiring the meta-training dataset to
contain any additional structure such as cluster labels.

###  4.1 Tractably structuring the prior predictive distribution over
functions via a Gaussian distribution over the weights

In our approach, we choose $\tilde{p}_{\xi}(f)$ to be the GP distribution over
functions that arises from a Gaussian prior on the weights of the
linearization of a neural network (equation 3). Consider a particular task
$\mathcal{T}^{i}$ and a batch of $K$ context data
$({\bm{X}}^{i},{\bm{Y}}^{i})$. The resulting prior predictive distribution,
derived from equation 1 after evaluating on the context inputs, is
${\bm{Y}}|{\bm{X}}^{i}\sim\mathcal{N}({\bm{\mu}}_{{\bm{Y}}\mid{\bm{X}}^{i}},{\bm{\Sigma}}_{{\bm{Y}}\mid{\bm{X}}^{i}})$,
where

${\bm{\mu}}_{{\bm{Y}}\mid{\bm{X}}^{i}}={\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i}){\bm{\mu}},\quad{\bm{\Sigma}}_{{\bm{Y}}\mid{\bm{X}}^{i}}={\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i}){\bm{\Sigma}}{\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i})^{\top}+{\bm{\Sigma}}_{\varepsilon}.$
(4)

In this setup, the parameters $\xi$ of $\tilde{p}_{\xi}(f)$ that we wish to
optimize are the linearization point ${\bm{\theta}}_{0}$, and the parameters
of the prior over the weights $({\bm{\mu}},{\bm{\Sigma}})$. Given this
Gaussian prior, it is straightforward to compute the joint NLL of the context
labels ${\bm{Y}}^{i}$,

$\displaystyle\mathrm{NLL}({\bm{X}}^{i},{\bm{Y}}^{i})=\frac{1}{2}\left(\left\|{\bm{Y}}^{i}-{\bm{\mu}}_{{\bm{Y}}\mid{\bm{X}}^{i}}\right\|^{2}_{{\bm{\Sigma}}_{{\bm{Y}}\mid{\bm{X}}^{i}}^{-1}}+\log\det{\bm{\Sigma}}_{{\bm{Y}}\mid{\bm{X}}^{i}}+{N_{y}}K\log
2\pi\right).$ (5)

The NLL (a) serves as a loss function quantifying the quality of $\xi$ during
training and (b) serves as an uncertainty signal at test time to evaluate
whether context data $({\bm{X}}^{i},{\bm{Y}}^{i})$ is OoD. Given this model,
adaptation is tractable as we can condition this GP on the context data
analytically. In addition, we can efficiently make probabilistic predictions
by evaluating the mean and covariance of the resulting posterior predictive
distribution on the query inputs, using equation 2.

####  4.1.1 Parameterizing the prior covariance over the weights

When working with deep neural networks, the number of weights $P$ can surpass
$10^{6}$. While it remains tractable to deal with ${\bm{\theta}}_{0}$ and
${\bm{\mu}}$, whose memory footprint grows linearly with $P$, it can quickly
become intractable to make computations with (let alone store) a dense prior
covariance matrix over the weights ${\bm{\Sigma}}\in\mathbb{R}^{P\times P}$.
Thus, we must impose some structural assumptions on the prior covariance to
scale to deep neural network models.

Imposing a unit covariance. One simple way to tackle this issue would be to
remove ${\bm{\Sigma}}$ from the learnable parameters $\xi$, i.e., fixing it to
be the identity ${\bm{\Sigma}}={\bm{I}}_{P}$. In this case,
$\xi=({\bm{\theta}}_{0},{\bm{\mu}})$. This computational benefit comes at the
cost of model expressivity, as we lose a degree of freedom in how we can
optimize our learned prior distribution $\tilde{p}_{\xi}(f)$. In particular,
we are unable to choose a prior over the weights of our model that captures
correlations between elements of the feature map.

Learning a low-dimensional representation of the covariance. An alternative is
to learn a low-rank representation of ${\bm{\Sigma}}$, allowing for a
learnable weight-space prior covariance that can encode correlations.
Specifically, we consider a covariance of the form
${\bm{\Sigma}}={\bm{Q}}^{\top}\text{diag}({\bm{s}}^{2}){\bm{Q}}$, where
${\bm{Q}}$ is a fixed projection matrix on an $s$-dimensional subspace of
$\mathbb{R}^{P}$, while ${\bm{s}}^{2}$ is learnable. In this case, the
parameters that are learned are $\xi=({\bm{\theta}}_{0},{\bm{\mu}},{\bm{s}})$.
We define ${\bm{S}}:=\text{diag}({\bm{s}}^{2})$. The computation of the
covariance of the prior predictive (equation 4) could then be broken down into
two steps:

$\left\\{\begin{array}[]{l}A:={\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i}){\bm{Q}}^{\top}\\\
{\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i}){\bm{\Sigma}}{\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i})^{\top}=A{\bm{S}}A^{\top}\end{array}\right.$

which requires a memory footprint of $O(P(s+{N_{y}}K))$, if we include the
storage of the Jacobian. Because ${N_{y}}K\ll P$ in typical deep learning
contexts, it suffices that $s\ll P$ so that it becomes tractable to deal with
this new representation of the covariance.

A trade-off between feature-map expressiveness and learning a rich prior over
the weights. Note that even if a low-dimensional representation of
${\bm{\Sigma}}$ enriches the prior distribution over the weights, it also
restrains the expressiveness of the feature map in the kernel by projecting
the $P$-dimensional features ${\bm{J}}({\bm{\theta}}_{0},{\bm{X}})$ on a
subspace of size $s\ll P$ via ${\bm{Q}}$. This presents a trade-off: we can
use the full feature map, but limit the weight-space prior covariance to be
the identity matrix by keeping ${\bm{\Sigma}}={\bm{I}}$ (case UnLiMiTD-I).
Alternatively, we could learn a low-rank representation of ${\bm{\Sigma}}$ by
randomly choosing $s$ orthogonal directions in $\mathbb{R}^{P}$, with the risk
that they could limit the expressiveness of the feature map if the directions
are not relevant to the problem that is considered (case UnLiMiTD-R). As a
compromise between these two cases, we can choose the projection matrix more
intelligently and project to the most impactful subspace of the full feature
map — in this way, we can reap the benefits of a tuneable prior covariance
while minimizing the useful features that the projection drops. To select this
subspace, we construct this projection map by choosing the top $s$
eigenvectors of the Fisher information matrix (FIM) evaluated on the training
dataset $\mathcal{D}$ (case UnLiMiTD-F). Recent work has shown that the FIM
for deep neural networks tends to have rapid spectral decay (Sharma et al.,
2021), which suggests that keeping only a few of the top eigenvectors of the
FIM is enough to encode an expressive task-tailored prior. See Appendix A.1
for more details.

####  4.1.2 Generalizing the structure to a mixture of Gaussians

When learning on multiple clusters of tasks, $p(f)$ can become non-unimodal,
and thus cannot be accurately described by a single GP. Instead, we can
capture this multimodality by structuring $\tilde{p}_{\xi}(f)$ as a mixture of
Gaussian processes.

Building a more general structure. We assume that at train time, a task
$\mathcal{T}^{i}$ comes from any cluster
$\left\\{\mathcal{C}_{j}\right\\}_{j=1}^{j=\alpha}$ with equal probability.
Thus, we choose to construct $\tilde{p}_{\xi}(f)$ as an equal-weighted mixture
of $\alpha$ Gaussian processes.

For each element of the mixture, the structure is similar to the single
cluster case, where the parameters of the cluster’s weight-space prior are
given by $({\bm{\mu}}_{j},{\bm{\Sigma}}_{j})$. We choose to have both the
projection matrix ${\bm{Q}}$ and the linearization point ${\bm{\theta}}_{0}$
(and hence, the feature map $\phi(\cdot)={\bm{J}}({\bm{\theta}}_{0},\cdot)$)
shared across the clusters. This yields improved computational efficiency, as
we can compute the projected features once, simultaneously, for all clusters.
This yields the parameters
$\xi_{\alpha}=({\bm{\theta}}_{0},{\bm{Q}},({\bm{\mu}}_{1},{\bm{s}}_{1}),\ldots,({\bm{\mu}}_{\alpha},{\bm{s}}_{\alpha}))$.

This can be viewed as a mixture of linear regression models, with a common
feature map but separate, independent prior distributions over the weights for
each cluster. These separate distributions are encoded using the low-
dimensional representations ${\bm{S}}_{j}$ for each ${\bm{\Sigma}}_{j}$.
Notice how this is a generalization of the single cluster case, for when
$\alpha=1$, $\tilde{p}_{\xi}(f)$ becomes a Gaussian and
$\xi_{\alpha}=\xi$222In theory, it is possible to drop ${\bm{Q}}$ and extend
the identity covariance case to the multi-cluster setting; however, this leads
to each cluster having an identical covariance function, and thus is not
effective at modeling heterogeneous behaviors among clusters..

Prediction and likelihood computation. The NLL of a batch of inputs under this
mixture model can be computed as

$\mathrm{NLL}_{\text{mixt}}({\bm{X}}^{i},{\bm{Y}}^{i})=\log\alpha-\log\operatorname{sum}\exp(-\mathrm{NLL}_{1}({\bm{X}}^{i},{\bm{Y}}^{i}),\ldots,-\mathrm{NLL}_{\alpha}({\bm{X}}^{i},{\bm{Y}}^{i})),$
(6)

where $\mathrm{NLL}_{j}({\bm{X}}^{i},{\bm{Y}}^{i})$ is the NLL with respect to
each individual Gaussian, as computed in equation 5, and
$\log\operatorname{sum}\exp$ computes the logarithm of the sum of the
exponential of the arguments, taking care to avoid underflow issues.

To make exact predictions, we would require conditioning this mixture model.
As this is not directly tractable, we propose to first infer the cluster from
which a task comes from, by identifying the Gaussian $\mathcal{G}_{j_{0}}$
that yields the highest likelihood for the context data
$\left({\bm{X}}^{i},{\bm{Y}}^{i}\right)$. Then, we can adapt by conditioning
$\mathcal{G}_{j_{0}}$ with the context data and finally infer by evaluating
the resulting posterior distribution on the queried inputs ${\bm{X}}^{i}_{*}$.

###  4.2 Meta-training the Parametric Task Distribution

The key to our meta-learning approach is to estimate the quality of
$\tilde{p}_{\xi}(f)$ via the NLL of context data from training tasks, and use
its gradients to update the parameters of the distribution $\xi$. Optimizing
this loss over tasks in the dataset draws $\tilde{p}_{\xi}(f)$ closer to the
empirical distribution present in the dataset, and hence towards the true
distribution $p(f)$.

We present three versions of UnLiMiTD, depending on the choice of structure of
the prior covariance over the weights (see Section 4.1.1 for more details).
UnLiMiTD-I (Algorithm 1) is the meta-training with the fixed identity prior
covariance. UnLiMiTD-R and UnLiMiTD-F (Algorithm 2) learn a low-dimensional
representation of that prior covariance, either with random projections or
with FIM-based projections.

Algorithm 1 UnLiMiTD-I: meta-training with identity prior covariance

1:Initialize ${\bm{\theta}}_{0}$, ${\bm{\mu}}$.

2:for all epoch do

3: Sample $n$ tasks
$\\{\mathcal{T}^{i},({\bm{X}}^{i},{\bm{Y}}^{i})\\}_{i=1}^{i=n}$

4: for all $\mathcal{T}^{i},({\bm{X}}^{i},{\bm{Y}}^{i})$ do

5:
$NLL_{i}\leftarrow\textsc{GaussNLL}({\bm{Y}}^{i};{\bm{J}}{\bm{\mu}},~{}{\bm{J}}{\bm{J}}^{\top}+{\bm{\Sigma}}_{\varepsilon})$
$\triangleright$ ${\bm{J}}={\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i})$

6: end for

7: Update ${\bm{\theta}}_{0}$, ${\bm{\mu}}$ with
$\nabla_{{\bm{\theta}}_{0}\cup{\bm{\mu}}}\sum_{i}NLL_{i}$

8:end for

Algorithm 2 UnLiMiTD-R and UnLiMiTD-F: meta-training with a learnt covariance

1:if using random projections then

2: Find random projection ${\bm{Q}}$

3: Initialize ${\bm{\theta}}_{0}$, ${\bm{\mu}}$, ${\bm{s}}$

4:else if using FIM-based projections then

5: Find intermediate ${\bm{\theta}}_{0}$, ${\bm{\mu}}$ with UnLiMiTD-I
$\triangleright$ see Alg. 1

6: Find ${\bm{Q}}$ via FIMProj(s); initialize ${\bm{s}}$. $\triangleright$ see
Alg. 3

7:end if

8:for all epoch do

9: Sample $n$ tasks
$\\{\mathcal{T}^{i},({\bm{X}}^{i},{\bm{Y}}^{i})\\}_{i=1}^{i=n}$

10: for all $\mathcal{T}^{i},({\bm{X}}^{i},{\bm{Y}}^{i})$ do

11:
$NLL_{i}\leftarrow\textsc{GaussNLL}({\bm{Y}}^{i};{\bm{J}}{\bm{\mu}},~{}{\bm{J}}{\bm{Q}}^{\top}\text{diag}({\bm{s}}^{2}){\bm{Q}}{\bm{J}}^{\top}+{\bm{\Sigma}}_{\varepsilon})$
$\triangleright$ ${\bm{J}}={\bm{J}}({\bm{\theta}}_{0},{\bm{X}}^{i})$

12: end for

13: Update ${\bm{\theta}}_{0}$, ${\bm{\mu}}$, ${\bm{s}}$ with
$\nabla_{{\bm{\theta}}_{0}\cup{\bm{\mu}}\cup{\bm{s}}}\sum_{i}NLL_{i}$

14:end for

Computing the likelihood. In the algorithms, the function
GaussNLL(${\bm{Y}}^{i}$; $m$, $K$) stands for NLL of ${\bm{Y}}^{i}$ under the
Gaussian $\mathcal{N}(m,K)$ (see equation 5). In the mixture case, we instead
use MixtNLL, which wraps equation 6 and calls GaussNLL for the individual NLL
computations (see discussion in Section 4.1.2). In this case, ${\bm{\mu}}$
becomes $\\{{\bm{\mu}}_{j}\\}_{j=1}^{j=\alpha}$ and ${\bm{s}}$ becomes
$\\{{\bm{s}}_{j}\\}_{j=1}^{j=\alpha}$ when applicable.

Finding the FIM-based projections. The FIM-based projection matrix aims to
identify the elements of $\phi={\bm{J}}({\bm{\theta}}_{0},{\bm{X}})$ that are
most relevant for the problem (see Section 4.1.1 and Appendix A.1). However,
this feature map evolves during training, because it is
${\bm{\theta}}_{0}$-dependent. How do we ensure that the directions we choose
for ${\bm{Q}}$ remain relevant during training? We leverage results from Fort
et al. (2020), stating that the NTK (the kernel associated with the Jacobian
feature map, see Section 3.2) changes significantly at the beginning of
training and that its evolution slows down as training goes on. This suggests
that as a heuristic, we can compute the FIM-based directions after partial
training, as they are unlikely to deviate much after the initial training. For
this reason, UnLiMiTD-F (Algorithm 2) first calls UnLiMiTD-I (Algorithm 1)
before computing the FIM-based ${\bm{Q}}$ that yields intermediate parameters
${\bm{\theta}}_{0}$ and ${\bm{\mu}}$. Then the usual training takes place with
the learning of ${\bm{s}}$ in addition to ${\bm{\theta}}_{0}$ and
${\bm{\mu}}$.

##  5 Related work

Bayesian inference with linearized DNNs. Bayesian inference with neural
networks is often intractable because the posterior predictive has rarely a
closed-form expression. Whereas UnLiMiTD linearizes the network to allow for
practical Bayesian inference, existing work has used other approximations to
tractably express the posterior. For example, it has been shown that in the
infinite-width approximation, the posterior predictive of a Bayesian neural
network behaves like a GP (Neal, 1996; Lee et al., 2018). This analysis can in
some cases yield a good approximation to the Bayesian posterior of a DNN
(Garriga-Alonso et al., 2019). It is also common to use Laplace’s method to
approximate the posterior predictive by a Gaussian distribution and allow
practical use of the Bayesian framework for neural networks. This
approximation relies in particular on the computation of the Hessian of the
network: this is in general intractable, and most approaches use the so-called
Gauss-Newton approximation of the Hessian instead (Ritter et al., 2018).
Recently, it has been shown that the Laplace method using the Gauss-Newton
approximation is equivalent to working with a certain linearized version of
the network and its resulting posterior GP (Immer et al., 2021).

Bayesian inference is applied in a wide range of subjects. For example, recent
advances in transfer learning have been possible thanks to Bayesian inference
with linearized neural networks. Maddox et al. (2021) have linearized pre-
trained networks and performed domain adaptation by conditioning the prior
predictive with data from the new task: the posterior predictive is then used
to make predictions. Our approach leverages a similar adaption method and
demonstrates how the prior distribution can be learned in a meta-learning
setup.

Meta-learning. MAML is a meta-learning algorithm that uses as adaptation a few
steps of gradient descent (Finn et al., 2017). It has the benefit of being
model-agnostic (it can be used on any model for which we can compute gradients
w.r.t. the weights), whereas UnLiMiTD requires the model to be a
differentiable regressor. MAML has been further generalized to probabilistic
meta-learning models such as PLATIPUS or BaMAML (Yoon et al., 2018; Finn et
al., 2018), where the simple gradient descent step is augmented to perform
approximate Bayesian inference. These approaches, like ours, learn (during
meta-training) and make use of (at test-time) a prior distribution on the
weights. In contrast, however, UnLiMiTD uses exact Bayesian inference at test-
time. MAML has also been improved for multimodal meta-learning via MMAML
(Vuorio et al., 2019; Abdollahzadeh et al., 2021). Similarly to our method,
they add a step to identify the cluster from which the task comes from (Vuorio
et al., 2019). OoD detection in meta-learning has been studied by Jeong & Kim
(2020), who build upon MAML to perform OoD detection in the classification
setting, to identify unseen classes during training. Iwata & Kumagai (2022)
also implemented OoD detection for classification, by learning a Gaussian
mixture model on a latent space. UnLiMiTD extends these ideas to the
regression task, aiming to identify when test data is drawn from an unfamiliar
function.

ALPaCA is a Bayesian meta-learning algorithm for neural networks, where only
the last layer is Bayesian (Harrison et al., 2018). Such framework yields an
exact linear regression that uses as feature map the activations right before
the last layer. Our work is a generalization of ALPaCA, in the sense that
UnLiMiTD restricted to the last layer matches ALPaCA’s approach. More on this
link between the methods is discussed in Appendix A.2.

##  6 Results and discussion

We wish to evaluate four key aspects of UnLiMiTD. (1) At test time, how do the
probabilistic predictions compare to baselines? (2) How well does the
detection of context data from OoD tasks perform? (3) How do these results
hold in the multimodal setting? (4) Which approach performs better between (a)
the identity covariance (UnLiMiTD-I), (b) the low-dimensional covariance with
random directions (UnLiMiTD-R) and the compromise (c) using FIM-based
directions (UnLiMiTD-F) (see trade-off in Section 4.1.1)? That is, what is
best between learning a rich prior distribution over the weights, keeping a
full feature map, and a compromise between the two?

We consider a cluster of sine tasks, one of linear tasks and one of quadratic
tasks, regression problems inspired from Vuorio et al. (2019). Details on the
problems can be found in Appendix A.4.

Unimodal meta-learning: The meta-learned prior accurately fits the tasks.
First, we investigate the performance of UnLiMiTD on a unimodal task
distribution consisting of sinusoids of varying amplitude and phase, using the
single GP structure for $\tilde{p}_{\xi}(f)$. We compare the performance
between UnLiMiTD-I, UnLiMiTD-R and UnLiMiTD-F. We also compare the results
between training with an infinite amount of available sine tasks (infinite
task dataset), and with a finite amount of available tasks (finite task
dataset). More training details can be found in Appendix A.5. Examples of
predictions at the test time are available in Figure 2, along with confidence
levels.

Figure 2:  Example of predictions for a varying number of context inputs $K$,
after meta-training with UnLiMiTD-F. Top: UnLiMiTD-F, infinite task dataset.
Bottom: UnLiMiTD-F, finite task dataset. The standard deviation is from the
posterior predictive distribution. Note how the uncertainty levels are
coherent with the actual prediction error. Also, note how uncertainty
decreases when there is more context data. Notice how UnLiMiTD-F recovers the
shape of the sine even with a low number of context inputs. Finally, note how
UnLiMiTD-F is able to reconstruct the sine even when trained on fewer tasks
(bottom). More comprehensive plots available in Figure 6.

In both OoD detection and quality of predictions, UnLiMiTD-R and UnLiMiTD-F
perform better than UnLiMiTD-I (Figure 3), and this is reflected in the
quality of the learned prior $\tilde{p}_{\xi}(f)$ in each case (see Appendix
A.3.1). With respect to the trade-off mentioned in Section 4.1.1, we find that
for small networks, a rich prior over the weights matters more than the full
expressiveness of the feature map, making both UnLiMiTD-R and UnLiMiTD-F
appealing. However, after running further experiments on a deep-learning
image-domain problem, this conclusion does not hold for deep networks (see
Appendix A.3.3), where keeping an expressive feature map is important
(UnLiMiTD-I and UnLiMiTD-F are appealing in that case). Thus, UnLiMiTD-F is
the variant that we retain, for it allows similar or better performances than
the other variants in all situations.

Note how UnLiMiTD-F outperforms MAML: it achieves much better generalization
when decreasing the number of context samples $K$ (Figure 3). Indeed,
UnLiMiTD-F trained with a finite task dataset performs better than MAML with
an infinite task dataset: it is able to capture better the common features of
the tasks with a smaller task dataset.

(a) Examples of context data from in-dist. and OoD tasks

(b) AUC for OoD detection

(c) MSE on predictions

Figure 3: Unimodal case: Performance of UnLiMiTD for OoD detection and
inference, as a function of the number of context datapoints $K$. The training
dataset consists of sinusoids, while OoD tasks are lines and quadratic tasks.
We compare different variants (UnLiMiTD-I, UnLiMiTD-R and UnLiMiTD-F), and
against MAML for predictions. We also compare training with a finite and
infinite task dataset. Note how UnLiMiTD-R and UnLiMiTD-F have efficient OoD
detection and outperform MAML in predictions. Also, note how MAML trained with
an infinite task dataset performs worse than UnLiMiTD-R and UnLiMiTD-F trained
on a finite task dataset.

Multimodal meta-learning: Comparing the mixture model against a single GP.
Next, we consider a multimodal task distribution with training data consisting
of sinusoids as well as lines with varying slopes. Here, we compare the
performance between choosing the mixture structure or the single GP structure
(see discussion in Section 4.1.2): in both cases, we use UnLiMiTD-F. More
training details can be found in Appendix A.6).

Both the OoD detection and the prediction performances are better with the
mixture structure than with the single GP structure (Figure 4), indicating
that the mixture model is a useful structure for $\tilde{p}_{\xi}(f)$. This is
reflected in the quality of the learned priors (see Appendix A.3.2 for
qualitative results including samples from the learned priors). Note how the
single GP structure still performs better than both MAML and MMAML for
prediction, especially in the low-data regime. This demonstrates the strength
of our probabilistic approach for multimodal meta-learning: even if the
probabilistic assumptions are not optimal, the predictions are still accurate
and can beat baselines.

(a) Examples of context data from in-dist. and OoD tasks

(b) AUC for OoD detection

(c) MSE on predictions

Figure 4: Multimodal case: Performance of UnLiMiTD for OoD detection and
inference, as a function of the number of context datapoints $K$. The training
dataset includes both sines and lines, while OoD tasks are quadratic
functions. We compare the different variants (UnLiMiTD-F with a single GP or a
mixture model), and against MAML/MMAML for predictions. Note how both versions
of UnLiMiTD-F yield better predictions than the baselines. In particular, even
with a single GP, UnLiMiTD-F outperforms the baselines.

##  7 Conclusion

We propose UnLiMiTD, a novel meta-learning algorithm that models the
underlying task distribution using a parametric and tuneable distribution,
leveraging Bayesian inference with linearized neural networks. We compare
three variants, and show that among these, the Fisher-based parameterization,
UnLiMiTD-F, effectively balances scalability and expressivity, even for deep
learning applications. We have demonstrated that (1) our approach makes
efficient probabilistic predictions on in-distribution tasks, which compare
favorably to, and often outperform, baselines, (2) it allows for effective
detection of context data from OoD tasks, and (3) that both these findings
continue to hold in the multimodal task-distribution setting.

There are several avenues for future work. One direction entails understanding
how the performance of UnLiMiTD-F is impacted if the FIM-based directions are
computed too early in the training and the NTK changes significantly
afterwards. One could also generalize our approach to non-Gaussian
likelihoods, making UnLiMiTD effective for classification tasks. Finally,
further research can push the limits of multimodal meta-learning, e.g., by
implementing non-parametric Bayesian methods to automatically infer an optimal
number of clusters, thereby eliminating a hyperparameter of the current
approach.

#### Acknowledgements

The authors acknowledge the MIT SuperCloud (Reuther et al., 2018) and Lincoln
Laboratory Supercomputing Center for providing HPC resources that have
contributed to the research results reported within this paper. The authors
would like to thank MISTI MIT-France for supporting this research. C.A.
further acknowledges support from Mines Paris Foundation. N.A. acknowledges
support from the Edgerton Career Development Professorship.

## References

  * Abdollahzadeh et al. (2021) Milad Abdollahzadeh, Touba Malekzadeh, and Ngai-Man Man Cheung.  Revisit multimodal meta-learning through the lens of multi-task learning.  _Advances in Neural Information Processing Systems_ , 34:14632–14644, 2021.
  * Azizan et al. (2021) Navid Azizan, Sahin Lale, and Babak Hassibi.  Stochastic mirror descent on overparameterized nonlinear models.  _IEEE Transactions on Neural Networks and Learning Systems_ , 2021\.
  * Farajtabar et al. (2020) Mehrdad Farajtabar, Navid Azizan, Alex Mott, and Ang Li.  Orthogonal gradient descent for continual learning.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 3762–3773. PMLR, 2020.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-agnostic meta-learning for fast adaptation of deep networks.  In _International Conference on Machine Learning_ , pp. 1126–1135. PMLR, 2017.
  * Finn et al. (2018) Chelsea Finn, Kelvin Xu, and Sergey Levine.  Probabilistic model-agnostic meta-learning.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Fort et al. (2020) Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli.  Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel.  _Advances in Neural Information Processing Systems_ , 33:5850–5861, 2020.
  * Gao et al. (2022) Ning Gao, Hanna Ziesche, Ngo Anh Vien, Michael Volpp, and Gerhard Neumann.  What matters for meta-learning vision regression tasks?  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 14776–14786, 2022.
  * Garriga-Alonso et al. (2019) Adrià Garriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison.  Deep convolutional networks as shallow gaussian processes.  In _International Conference on Learning Representations_ , 2019.
  * Harrison et al. (2018) James Harrison, Apoorva Sharma, and Marco Pavone.  Meta-learning priors for efficient online bayesian regression.  In _International Workshop on the Algorithmic Foundations of Robotics_ , pp. 318–337. Springer, 2018.
  * Immer et al. (2021) Alexander Immer, Maciej Korzepa, and Matthias Bauer.  Improving predictions of bayesian neural nets via local linearization.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 703–711. PMLR, 2021.
  * Iwata & Kumagai (2022) Tomoharu Iwata and Atsutoshi Kumagai.  Meta-learning for out-of-distribution detection via density estimation in latent space.  _arXiv preprint arXiv:2206.09543_ , 2022.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Jeong & Kim (2020) Taewon Jeong and Heeyoung Kim.  Ood-maml: Meta-learning for few-shot out-of-distribution detection and classification.  _Advances in Neural Information Processing Systems_ , 33:3907–3916, 2020.
  * Lee et al. (2018) Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein.  Deep neural networks as gaussian processes.  In _International Conference on Learning Representations_ , 2018.
  * Liu et al. (2020) Chaoyue Liu, Libin Zhu, and Misha Belkin.  On the linearity of large non-linear models: when and why the tangent kernel is constant.  _Advances in Neural Information Processing Systems_ , 33:15954–15964, 2020.
  * Maddox et al. (2021) Wesley Maddox, Shuai Tang, Pablo Moreno, Andrew Gordon Wilson, and Andreas Damianou.  Fast adaptation with linearized neural networks.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 2737–2745. PMLR, 2021.
  * Neal (1996) Radford M Neal.  Priors for infinite networks.  In _Bayesian Learning for Neural Networks_ , pp. 29–53. Springer, 1996.
  * Rasmussen & Williams (2005) Carl Edward Rasmussen and Christopher K. I. Williams.  _Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning)_.  The MIT Press, 2005.  ISBN 026218253X.
  * Reuther et al. (2018) Albert Reuther, Jeremy Kepner, Chansup Byun, Siddharth Samsi, William Arcand, David Bestor, Bill Bergeron, Vijay Gadepally, Michael Houle, Matthew Hubbell, Michael Jones, Anna Klein, Lauren Milechin, Julia Mullen, Andrew Prout, Antonio Rosa, Charles Yee, and Peter Michaleas.  Interactive supercomputing on 40,000 cores for machine learning and data analysis.  In _2018 IEEE High Performance Extreme Computing Conference (HPEC)_ , pp. 1–6. IEEE, 2018.
  * Ritter et al. (2018) Hippolyt Ritter, Aleksandar Botev, and David Barber.  A scalable laplace approximation for neural networks.  In _6th International Conference on Learning Representations, ICLR 2018-Conference Track Proceedings_ , volume 6. International Conference on Representation Learning, 2018.
  * Sagun et al. (2017) Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou.  Empirical analysis of the hessian of over-parametrized neural networks.  _arXiv preprint arXiv:1706.04454_ , 2017.
  * Sharma et al. (2021) Apoorva Sharma, Navid Azizan, and Marco Pavone.  Sketching curvature for efficient out-of-distribution detection for deep neural networks.  In _Uncertainty in Artificial Intelligence_ , pp. 1958–1967. PMLR, 2021.
  * Tropp et al. (2017) Joel A Tropp, Alp Yurtsever, Madeleine Udell, and Volkan Cevher.  Practical sketching algorithms for low-rank matrix approximation.  _SIAM Journal on Matrix Analysis and Applications_ , 38(4):1454–1485, 2017.
  * Vuorio et al. (2019) Risto Vuorio, Shao-Hua Sun, Hexiang Hu, and Joseph J Lim.  Multimodal model-agnostic meta-learning via task-aware modulation.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Yoon et al. (2018) Jaesik Yoon, Taesup Kim, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn.  Bayesian model-agnostic meta-learning.  _Advances in Neural Information Processing Systems_ , 31, 2018.

##  Appendix A Appendix

###  A.1 A tractable way of finding the perturbation directions in weight
space that impact the most the predictions of an entire dataset

Deep neural networks have a large number of parameters, making the feature map
$\phi_{{\bm{\theta}}_{0}}$ high-dimensional. However, recent work has shown
that only a small subspace of the weight space is impactful. For example, to
perform continual learning, Farajtabar et al. (2020) leverage the fact that it
is sufficient to update the parameters orthogonally to a few directions only
to avoid catastrophic forgetting. Sagun et al. (2017) have shown that the
Hessian of a deep neural network can be summarized in a few number of
directions, due to rapid spectral decay. This encourages finding a method to
extract these meaningful directions of the weight space.

####  A.1.1 Link with the Fisher Information Matrix

We define these main directions as the ones that have the most impact on the
predictions of a whole dataset. To find them, we first find a way to quantify
the influence of an infinitesimal weight perturbation. Using the second-order
approximation of that quantity, we then describe in the deep-learning context
a tractable way to find the directions.

##### Setting

We take the same setting as Section 3.2, and we describe a method to quantify
the influence of a parameter perturbation $\tilde{{\bm{\theta}}}$ on the
predictions of a dataset of tasks $\mathcal{D}$. To do so, we leverage a
probabilistic interpretation of the model: we assume a Gaussian pdf over the
observations for given inputs and parameters
$p_{{\bm{\theta}}}({\bm{Y}}|{\bm{X}})\sim\mathcal{N}(g({\bm{\theta}},{\bm{X}}),~{}{\bm{\Sigma}}_{\varepsilon})$,
where the covariance of the noise ${\bm{\Sigma}}_{\varepsilon}$ is diagonal
${\bm{\Sigma}}_{\varepsilon}=\sigma_{\varepsilon}^{2}{\bm{I}}$ (just as in
Section 3.1).

Perturbation of the prediction of a batch of inputs. Before quantifying the
influence of the perturbation on the predictions of the whole task dataset, we
do it for the prediction of a batch of inputs $g({\bm{\theta}}_{0},{\bm{X}})$.

We borrow the method from Sharma et al. (2021): we quantify the influence of a
parameter perturbation by computing the Kullback-Leibler divergence between
$p_{{\bm{\theta}}_{0}}({\bm{Y}}|{\bm{X}})$ and
$p_{{\bm{\theta}}_{0}+\tilde{{\bm{\theta}}}}({\bm{Y}}|{\bm{X}})$. The
expansion is:

$\delta({\bm{\theta}}_{0},{\bm{X}})(\tilde{{\bm{\theta}}}):=D_{\mathrm{KL}}(p_{{\bm{\theta}}_{0}}({\bm{Y}}|{\bm{X}})\|p_{{\bm{\theta}}_{0}+\tilde{{\bm{\theta}}}}({\bm{Y}}|{\bm{X}}))\approx\tilde{{\bm{\theta}}}^{\top}{\bm{F}}({\bm{\theta}}_{0},{\bm{X}})\tilde{{\bm{\theta}}}+o(\|\tilde{{\bm{\theta}}}\|^{2})$
(7)

where
${\bm{F}}({\bm{\theta}}_{0},{\bm{X}}):={\bm{J}}({\bm{\theta}}_{0},{\bm{X}})^{\top}{\bm{\Sigma}}_{\varepsilon}^{-1}{\bm{J}}({\bm{\theta}}_{0},{\bm{X}})=\sigma_{\varepsilon}^{-2}{\bm{J}}({\bm{\theta}}_{0},{\bm{X}})^{\top}{\bm{J}}({\bm{\theta}}_{0},{\bm{X}})\in\mathbb{R}^{P\times
P}$ is the empirical Fisher Information Matrix (FIM) of the batch of inputs
${\bm{X}}$, computed on the parameters ${\bm{\theta}}_{0}$.

Generalization: perturbation of the prediction of a dataset. Now we can define
the influence of a parameter perturbation on the whole training dataset
$\mathcal{D}$, by generalizing the previous definition:

$\delta({\bm{\theta}}_{0},\mathcal{D})(\tilde{{\bm{\theta}}}):=\dfrac{1}{N}\sum_{i=1}^{N}\delta({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})(\tilde{{\bm{\theta}}})$

Using equation 7, this quantity verifies:

$\delta({\bm{\theta}}_{0},\mathcal{D})(\tilde{{\bm{\theta}}})\approx\tilde{{\bm{\theta}}}^{\top}\left(\dfrac{1}{N}\sum_{i=1}^{N}{\bm{F}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})\right)\tilde{{\bm{\theta}}}+o(\|\tilde{{\bm{\theta}}}\|^{2})$
(8)

which gives a natural definition for the FIM of the whole dataset by analogy
with equation 7:

${\bm{F}}({\bm{\theta}}_{0},\mathcal{D}):=\dfrac{1}{N}\sum_{i=1}^{N}{\bm{F}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})=\dfrac{1}{N\sigma_{\varepsilon}^{2}}\sum_{i=1}^{N}{\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})^{\top}{\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})\in\mathbb{R}^{P\times
P}$

The expansion in equation 8 shows that the FIM of the dataset is a second-
order approximation describing the influence of a parameter perturbation over
the entire dataset. In particular, the eigenvectors of
${\bm{F}}({\bm{\theta}}_{0},\mathcal{D})$ with the highest eigenvalues are the
directions that impact the most the predictions.

####  A.1.2 Computing the top eigenvectors of the Fisher Information Matrix of
the dataset in a deep learning context

Naively computing the top eigenspace of
${\bm{F}}({\bm{\theta}}_{0},\mathcal{D})$ requires processing a $P\times P$
matrix, which is intractable in the deep-learning context (where $P$ can
surpass $10^{6}$). Instead, we decide to use the method used by Sharma et al.
(2021), which leverages a low-rank-approximation-based technique (namely
matrix sketching) developed by Tropp et al. (2017).

##### Sketching the FIM of the dataset

The key idea behind this technique is to build two small random sketches of
the FIM,
$({\bm{Y}},{\bm{W}}):=\mathcal{S}\left({\bm{F}}({\bm{\theta}}_{0},\mathcal{D})\right)$,
that with high probability contain enough information to reconstruct the top
$s$ eigenvectors of ${\bm{F}}({\bm{\theta}}_{0},\mathcal{D})$. The linearity
of $\mathcal{S}$ simplifies the sketching process by breaking down the
computation into individual sketches:

$({\bm{Y}},{\bm{W}}):=\mathcal{S}({\bm{F}}({\bm{\theta}}_{0},\mathcal{D}))=\dfrac{1}{N\sigma_{\varepsilon}^{2}}\sum_{i=1}^{N}\mathcal{S}\left({\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})^{\top}{\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})\right)=:\dfrac{1}{N\sigma_{\varepsilon}^{2}}\sum_{i=1}^{N}({\bm{Y}}^{i},{\bm{W}}^{i})$

In particular, the sketch $({\bm{Y}},{\bm{W}})$ can be updated in-place and
does not require to store all the individual sketches. Given a sketch budget
$k+l$ (Sharma et al. (2021) recommends choosing $k:=2s+1$ and $l:=4s+3$) and
two random normal matrices $\bm{\Omega}\in\mathbb{R}^{k\times P}$ and
$\bm{\Psi}\in\mathbb{R}^{l\times P}$, the random individual sketches
$({\bm{W}}^{i},{\bm{Y}}^{i})$ are defined as:

$\left\\{\begin{array}[]{rcll}{\bm{Y}}^{i}&:=&((\bm{\Omega}{\bm{J}}_{i}^{\top}){\bm{J}}_{i})^{\top}&\in\mathbb{R}^{P\times
k}\\\
{\bm{W}}^{i}&:=&(\bm{\Psi}{\bm{J}}_{i}^{\top}){\bm{J}}_{i}&\in\mathbb{R}^{l\times
P}\end{array}\right.$

where ${\bm{J}}_{i}:={\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})$.

##### Sketch-based computation of the top eigenspace

Once the sketches are computed, the function FixedRankSymApprox by Tropp et
al. (2017) computes the first eigenvectors and eigenvalues of the FIM of the
dataset. Overall, the memory footprint to find the sketches and the top
eigenspace is $O(P(s+{N_{y}}M))$, where $s$ is the number of queried
eigenvectors and $M$ is the size of $\widetilde{{\bm{X}}}^{i}$. As long as
$s\ll P$, this computation is tractable, given that ${N_{y}}M\ll P$ is usual
deep learning contexts. Algorithm 3 summarizes the process that yields the
FIM-based projections via sketching. We drop the scaling coefficient
$\sigma_{\varepsilon}^{-2}$ as it doesn’t affect the computation, given that
we only want orthogonal eigenvectors and eigenvalues.

1:$s$ (desired size of the subspace)

2:$k\leftarrow 2s+1$

3:$l\leftarrow 4s+3$

4:Draw $\bm{\Omega}\in\mathbb{R}^{k\times P},\bm{\Psi}\in\mathbb{R}^{l\times
P}$, two random normal matrices

5:Initialize ${\bm{Y}}=0\in\mathbb{R}^{P\times
k},{\bm{W}}=0\in\mathbb{R}^{l\times P}$

6:for all training task $\mathcal{T}^{i}$ do

7:
${\bm{J}}_{i}\leftarrow{\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})$

8:
${\bm{Y}}\leftarrow{\bm{Y}}+1/N((\bm{\Omega}{\bm{J}}_{i}^{\top}){\bm{J}}_{i})^{\top}$

9: ${\bm{W}}\leftarrow{\bm{W}}+1/N(\bm{\Psi}{\bm{J}}_{i}^{\top}){\bm{J}}_{i}$

10:end for

Algorithm 3 Computing the FIM-based projections

###  A.2 UnLiMiTD as a generalization of ALPaCA

##### Restraining the linearization to the last layer

Remember the linear regression of equation 3, that we obtained by linearizing
the network with all its layers. Let’s separate the parameters of the network
${\bm{\theta}}$ into two: the parameters of all the layers but the last one
$\bm{\lambda}$, and the parameters of the last layer $\bm{\rho}$:
${\bm{\theta}}=\bm{\lambda}\cup\bm{\rho}$.

We assume that the last layer is dense with biases: we will note
$\bm{\rho}^{w}$ the weight matrix and $\bm{\rho}^{b}$ the biases of this last
layer. ${N_{\psi}}$ will stand as the dimension of the activations right
before the last layer: in particular,
$\bm{\rho}^{w}\in\mathbb{R}^{{N_{y}}\times{N_{\psi}}}$ and
$\bm{\rho}^{b}\in\mathbb{R}^{{N_{y}}}$. $P^{\prime}$ will stand as the size of
$\bm{\rho}$: in our case, $P^{\prime}={N_{\psi}}\times{N_{y}}+{N_{y}}$. We
implicitly vectorize $\bm{\rho}$, such that:

$\bm{\rho}=\begin{pmatrix}\operatorname{vec}{\bm{\rho}^{w}}\\\
\bm{\rho}^{b}\end{pmatrix}\in\mathbb{R}^{{N_{\psi}}{N_{y}}+{N_{y}}}=\mathbb{R}^{P^{\prime}}$

We now restrain the linear regression to the last layer as follows:

${\bm{y}}={\bm{J}}^{\prime}(\bm{\rho}_{0},\psi_{\bm{\lambda}}({\bm{X}}))(\bm{\rho}-\bm{\rho}_{0})+\varepsilon$
(9)

where:

  * •

$\psi_{\bm{\lambda}}(\cdot):\mathbb{R}^{{N_{x}}\times
K}\rightarrow\mathbb{R}^{{N_{\psi}}\times K}$ stands for the function that
maps the inputs and the activations right before the last layer;

  * •

${\bm{J}}^{\prime}(\cdot,\cdot):\mathbb{R}^{P^{\prime}}\times\mathbb{R}^{{N_{\psi}}\times
K}\rightarrow\mathbb{R}^{{N_{y}}K\times P^{\prime}}$ stands for the jacobian
of the last layer with respect to the parameters $\bm{\rho}$. We can write the
jacobian in closed-form due to the linearity of the last layer:

${\bm{J}}^{\prime}(\bm{\rho}_{0},\psi_{\bm{\lambda}}({\bm{X}}))=\begin{pmatrix}\psi_{\bm{\lambda}}({\bm{X}})^{\top}\otimes{\bm{I}}_{{N_{y}}}&{\bm{I}}_{{N_{y}}}\end{pmatrix}\in\mathbb{R}^{{N_{y}}K\times
P^{\prime}}$

Note that ${\bm{J}}^{\prime}(\bm{\rho}_{0},\psi_{\bm{\lambda}}(x))$ does not
depend on the linearization point $\bm{\rho}_{0}$ (could be expected, given
the linearity of the last layer).

  * •

$\bm{\rho}-\bm{\rho}_{0}\in\mathbb{R}^{P^{\prime}}$ is the correction to the
last parameters. We will note the correction to the weights as
$\hat{\bm{\rho}}^{w}:=\bm{\rho}^{w}-\bm{\rho}_{0}^{w}$ and the correction to
the bias as $\hat{\bm{\rho}}^{b}:=\bm{\rho}^{b}-\bm{\rho}_{0}^{b}$.

Using Kronecker’s product identities, the linear regression in equation 9 can
be rewritten:

${\bm{y}}=\hat{\bm{\rho}}^{w}\psi_{\bm{\lambda}}({\bm{X}})+\hat{\bm{\rho}}^{b}+\varepsilon$
(10)

Also, as a side note, another way of getting this linear regression (equation
10) is to rewrite the initial linearization of Section 3.2, but with respect
to $\bm{\rho}$ only:

$\begin{array}[]{rcl}f(\bm{\lambda}\cup\bm{\rho},{\bm{x}}_{t})-f(\bm{\lambda}\cup\bm{\rho}_{0},{\bm{x}}_{t})&=&\bm{\rho}^{w}\psi_{\bm{\lambda}}(x)+\bm{\rho}^{b}-(\bm{\rho}_{0}^{w}\psi_{\bm{\lambda}}(x)+\bm{\rho}_{0}^{b})\\\
&=&(\bm{\rho}^{w}-\bm{\rho}_{0}^{w})\psi_{\bm{\lambda}}(x)+(\bm{\rho}^{b}-\bm{\rho}_{0}^{b})\\\
&=&\hat{\bm{\rho}}^{w}\psi_{\bm{\lambda}}(x)+\hat{\bm{\rho}}^{b}\end{array}$

Doing it this way has the benefit to show that the linearization is exact when
restricted to the last layer. For non-linear neural networks, the
linearization is always an approximation.

##### Adapting UnLiMiTD to the last layer

Just like what we did in the general case, we use the GP theory to make
Bayesian inference on this linear regression. The new parameters of the GP
$\tilde{p}_{\xi}(f)$ are now $\xi=(\bm{\lambda},{\bm{\mu}},{\bm{\Sigma}})$,
where ${\bm{\mu}}$ and ${\bm{\Sigma}}$ are the parameters of the Gaussian
prior over $\hat{\bm{\rho}}$. Note how $\bm{\rho}_{0}$ have disappeared from
$\xi$: contrary to the general case, the linearization point is not optimized,
as it does not impact the computation. Also note that $\bm{\lambda}$ has
replaced $\bm{\rho}_{0}$, as it parameterizes the feature map.

##### Comparison with ALPaCA

In the ALPaCA setting, the linear regression is:

${\bm{y}}_{a}=\hat{\bm{\rho}}^{w}\psi_{\bm{\lambda}}({\bm{X}})+\varepsilon$
(11)

Note that $\hat{\bm{\rho}}^{w}$ still plays the same role in the linear
regression, but is not any correction anymore. Substracting 11 \- 10 yields:

${\bm{y}}-{\bm{y}}_{a}=\hat{\bm{\rho}}^{b}$

UnLiMiTD restricted to the last layer is closely related to ALPaCA, in that
they both perform a linear regression with the same kernel: the only
difference lies in the additional bias term that is not considered in ALPaCA.
Thus, we can think of UnLiMiTD as a generalization of ALPaCA to all the layers
of the network.

###  A.3 Additional results

####  A.3.1 Additional results (single-cluster case)

##### Quality of the priors

To qualitatively analyze the prior after meta-training on the sines, we plot
the mean functions and the covariance functions of the resulting GP, that is
$\tilde{p}_{\xi}(f)$ (Figure 5). All the trainings yield a similar mean for
$\tilde{p}_{\xi}(f)$ (that is a cosine with amplitude 1.5 and offset 1)
(Figure 5(a)), which is close to the theoretical mean of $p(f)$ (a cosine with
amplitude 2.5 and offset 1). The covariance function of UnLiMiTD-R and
UnLiMiTD-F (Figure 5(c)) resembles what we would expect for $p(f)$ (e.g.,
periodicity, negative correlation between 0 and $\pi$, etc.), but it is not
the case for UnLiMiTD-I (Figure 5(b)). This empirical analysis confirms the
quantitative comparison in terms of OoD detection and prediction performance
carried in Section 6.

(a) Mean $\mathbb{E}(f)$ (UnLiMiTD-I, UnLiMiTD-R, UnLiMiTD-F)

(b) Covariance $\text{cov}(f(0),f(x))$ (UnLiMiTD-I)

(c) Covariance $\text{cov}(f(0),f(x))$ (UnLiMiTD-R, UnLiMiTD-F)

Figure 5: Mean (5(a)) and covariance (5(b), 5(c)) functions of
$\tilde{p}_{\xi}(f)$ after different meta-trainings on the sine cluster
(UnLiMiTD-I, UnLiMiTD-R and UnLiMiTD-F, with a finite or infinite dataset).
Note the scale, and how UnLiMiTD-I has a less accurate covariance function
that UnLiMiTD-R and UnLiMiTD-F.

##### Examples of predictions

Figure 6 summarizes the predictions of the model meta-trained on sines, for a
varying number of context inputs $K$, breaking down all the different cases of
training.

(a) UnLiMiTD-I (inf)

(b) UnLiMiTD-R (inf)

(c) UnLiMiTD-F (inf)

(d) UnLiMiTD-I (fin)

(e) UnLiMiTD-R (fin)

(f) UnLiMiTD-F (fin)

Figure 6: Example of predictions for a varying number of context inputs $K$,
after different meta-trainings on the sine cluster (UnLiMiTD-I, UnLiMiTD-R and
UnLiMiTD-F, in the infinite and finite case). Standard deviation is from the
posterior predictive distribution. Note how UnLiMiTD-R and UnLiMiTD-F perform
better than UnLiMiTD-I when it comes to reconstructing the sine with a smaller
amount of context inputs.

####  A.3.2 Additional results (multi-cluster case)

##### Quality of the priors

Figure 7 and Figure 8 show the mean and covariance functions of the GP (when
training with a single GP), and the mean and covariance functions of the GPs
composing the mixture(when training with a mixture of GPs). We note that in
the case of the mixture, both of the Gaussians composing the mixture have
correctly captured the common features shared by each of the clusters
(respectively the linear and the sine cluster). For instance, the mean of the
line cluster is the zero function, which matches Figure 7(b), and the
correlation between $x=1$ and the other inputs is correctly rendered for
linear tasks (Figure 8(b)).

When learning with a single GP, the learnt mean and covariance do not match
any of the two clusters. For example, the mean has an intermediate offset
between the offset of the sine cluster and the line cluster. This empirical
analysis comforts the conclusions of Section 6: the mixture model yields
better results than the single GP case.

(a) Mean function $\mathbb{E}(f)$ of the GP (UnLiMiTD-F with a single GP).

(b) Mean function $\mathbb{E}(f)$ of the first GP of the mixture (UnLiMiTD-F
with a mixture of GPs).

(c) Mean function $\mathbb{E}(f)$ of the second GP of the mixture (UnLiMiTD-F
with a mixture of GPs).

Figure 7: Mean functions of the GP (when learning with a GP) / the GPs
composing the mixture (when learning a mixture of GPs), after training on both
the sine and line cluster with UnLiMiTD-F. Note how the mean of the single GP
is intermediate between the ones of the mixture.

(a) Covariance function $\text{cov}(f(0),f(x))$ of the GP (UnLiMiTD-F with a
single GP).

(b) Covariance function $\text{cov}(f(1),f(x))$ of the first GP of the mixture
(UnLiMiTD-F with a mixture of GPs)

(c) Covariance function $\text{cov}(f(0),f(x))$ of the second GP of the
mixture (UnLiMiTD-F with a mixture of GPs)

Figure 8: Covariance functions of the GP (when learning with a GP) / the GPs
composing the mixture (when learning a mixture of GPs), after training on both
the sine and line cluster with UnLiMiTD-F. Note how the covariance of the
single GP is not accurate for any of the two clusters.

####  A.3.3 UnLiMiTD yields effective predictions on large-scale vision
problems

We consider a regression-vision meta-learning problem from Gao et al. (2022),
Shapenet1D, aiming to predict object orientations in space. In this problem,
each task consists of a different object of which we want to predict the
orientation. For each task, the context data consists of some images of the
same object, but with different orientations; the query inputs are other
images of the same object, with unknown orientations. Details on the problems
and the datasets can be found in appendix A.4.

We train a deep learning model on Shapenet1D, and we compare the performances
on the test set between UnLiMiTD-I, UnLiMiTD-R and UnLiMiTD-F (Figure 9). More
training and test details can be found in the Appendix A.7.

Both UnLiMiTD-I and UnLiMiTD-F yield better performances than MAML, achieving
low angle errors. UnLiMiTD-R however gives poor results, worse than that of
MAML.

In terms of the trade-off of Section 4.1.1, our conclusion from small networks
does not scale up to deep models. Here, the loss of valuable features is
perceptible (that is what happens with UnLiMiTD-R, when the randomness of the
directions may drop such features), and not learning a rich prior over the
weights is not burdensome (UnLiMiTD-I). However, UnLiMiTD-F plays a role of
compromise, by learning the prior covariance while keeping the few important
features of the jacobian, as it gives comparable results to UnLiMiTD-I.

Figure 9: Left: Examples of vision tasks. Right: angle error as a function of
the number of context datapoints $K$, after different meta-trainings on
Shapenet1D (UnLiMiTD-I, UnLiMiTD-R and UnLiMiTD-F vs MAML). Note how
UnLiMiTD-I and UnLiMiTD-F make better predictions than MAML. Also note that
contrary to the small network case, UnLiMiTD-R performs worse than MAML and
than the two other variants.

###  A.4 Details on the regression problems

####  A.4.1 Simple regression problems

Our simple regression problems are inspired by Vuorio et al. (2019)’s work.
They consists of three clusters of different types of tasks, and varying
offset. The first cluster consists of sines with constant frequency and
offset, but with a varying amplitude and phase:

$\left\\{x\mapsto A\sin(x+\varphi)+1|A\in[0.1,5],\varphi\in[0,\pi]\right\\}$

The second cluster consists of lines with a varying slope and no offset:

$\left\\{x\mapsto ax|a\in[-1,1]\right\\}$

The last cluster consists of quadratic functions, with a varying quadratic
coefficient and phase, and with a constant offset:

$\left\\{x\mapsto
a(x-\varphi)^{2}+0.5|a\in[-0.2,0.2],\varphi\in[-2,2]\right\\}$

In all these clusters, we add an artificial Gaussian noise on the context
observations $\mathcal{N}(0,~{}0.05)$. The query datapoints remain noiseless,
to remain coherent with the assumption from Section 3.1 borrowed from
Rasmussen & Williams (2005).

####  A.4.2 Vision problem

We consider the meta-learning vision, regression problem recently created by
Gao et al. (2022) (namely Shapenet1D), that consists in estimating objects
orientation from images. The objects have a vertical angular degree of
freedom, and Figure 10 shows an example of such objects in different
positions.

(a) Angle: 4°

(b) Angle: 26°

(c) Angle: 78°

(d) Angle: 130°

(e) Angle: 202°

(f) Angle: 243°

(g) Angle: 261°

(h) Angle: 285°

Figure 10: Examples of images and ground-truth angles from the Shapenet1D
dataset (Gao et al., 2022)

We use the same training and test datasets as Gao et al. (2022), with no kind
of augmentation (in particular, no artificial noise on the ground-truth
angles). In particular, we use the same intra-category (IC) evaluation dataset
(that is, objects from the same categories as the objects used for training)
and cross-category (CC) evaluation dataset (that is, objects from different
categories as the ones used for training).

###  A.5 Training and test details for the single cluster case

####  A.5.1 Training details of UnLiMiTD (single cluster case)

The model is a neural network with $2$ hidden layers, $40$ neurons each, with
a ReLU non-linearity. In our single-cluster experiment, the cluster is the
sine cluster (see Appendix A.4).

In the case where there are an infinite number of available sine tasks during
training (ie $N=+\infty$), the training is performed with $n=24$ tasks per
epoch, and at each epoch the context inputs are randomly drawn from $[-5,5]$
(which means that $M=+\infty$). In the case where we restrict the available
tasks to a finite number, we randomly choose $N=10$ tasks and $M=50$ context
datapoints per task before training (they are shared among all the “finite”
trainings) and perform the trainings with $n=6$ tasks per epoch.

For all the trainings, the number of context inputs seen during training is
$K=10$.

For all trainings, we train UnLiMiTD on $60,000$ epochs. When dealing with the
UnLiMiTD-F case, we allow half of the epochs ($30,000$) to the training before
finding the intermediate ${\bm{\theta}}_{0}$ and ${\bm{\mu}}$ (see Algorithm
2), and the other half after.

In the UnLiMiTD-F case with infinite available tasks, a finite number of tasks
is needed to compute the FIM: thus we build an artificial finite dataset of
$N=100$ and $M=P$ (arbitrary, but chosen so that
${\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i}){\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})^{\top}$
gets a chance to be full-rank ie contain as much information as possible),
that is used only for that computation.

In the UnLiMiTD-F and the UnLiMiTD-R cases, the subspace size is $s=10$.

For all trainings, the meta-optimizer is Adam with an initial learning rate of
$0.001$. The noise $\sigma_{\varepsilon}=0.05$, equal to the noise added to
the context data. We compute the NLL using Rasmussen & Williams (2005)’s
implementation.

####  A.5.2 Training details of MAML (single cluster case)

We train MAML baselines to compare our results with that of MAML. A large part
of these hyperparameters is directly inspired from Finn et al. (2017)’s work.

The model is a neural network with $2$ hidden layers, $40$ neurons each, with
a ReLU non-linearity. In our single-cluster experiment, the cluster is the
sine cluster (see Appendix A.4).

In the case where there are an infinite number of available sine tasks during
training (ie $N=+\infty$), the training is performed with $n=24$ tasks per
epoch, and at each epoch the context and query inputs are randomly drawn from
$[-5,5]$. In the case where we restrict the available tasks to a finite
number, we randomly choose $N=10$ tasks and $M=50$ datapoints per task (used
for both context and query batches) before training (they are shared among all
the “finite” trainings) and perform the trainings with $n=6$ tasks per epoch.

For all trainings, the number of context datapoints is $K=10$, and the number
of query datapoints is $L=10$. We meta-train for $70,000$ epochs.

For all trainings, the meta-optimizer is Adam with an initial learning rate of
$0.001$. The inner learning-rate is kept constant, at $0.001$. The number of
inner updates is $5$ during training, and $10$ at test time.

####  A.5.3 Test details (single-cluster case)

For the OoD detection evaluation, we plot the AUC as a function of the number
of the context inputs $K$. The AUC is computed using the NLL of the context
inputs wrt to $\tilde{p}_{\xi}(f)$ (our uncertainty metric). The true-
positives are the OoD tasks (lines and quadratic tasks) flagged as such; the
false-positives are the in-distribution tasks (sines) flagged as OoD.

For the predictions, we plot the average and ci95 of the MSE on 1,000 tasks
(100 queried inputs each). In the UnLiMiTD-R case, we also compute the average
and ci95 on 5 different random projections trainings.

###  A.6 Training and test details for the multi cluster case

####  A.6.1 Training details of UnLiMiTD (multi-cluster case)

The model is a neural network with $2$ hidden layers, $40$ neurons each, with
a ReLU non-linearity. In our multi-cluster experiment, $\alpha=2$: the
clusters consist of the sine cluster and the linear cluster (see Appendix
A.4).

For all the trainings, the training is performed with $n=24$ tasks per epoch
(with an infinite number of available sine tasks and linear tasks during
training ie $N=+\infty$), and at each epoch the context inputs are randomly
drawn from $[-5,5]$ (which means that $M=+\infty$). In accordance with our
equal probability assumption from Section 4.1.2, at each epoch $n/2=12$ tasks
come from the sine cluster and $n/2=12$ tasks come from the linear cluster.

For all trainings, the number of context inputs seen during training is
$K=10$.

For all trainings, we train UnLiMiTD algorithm on $60,000$ epochs: because we
deal with the UnLiMiTD-F case, we allow half of the epochs ($30,000$) to the
training before finding the intermediate ${\bm{\theta}}_{0}$ and
$\\{{\bm{\mu}}_{j}\\}_{j=1}^{j=\alpha}$ (see Algorithm 2), and the other half
after.

In the UnLiMiTD-F case, a finite number of tasks is needed to compute the FIM:
thus we build an artificial finite dataset of $N=100$ and $M=P$ (arbitrary,
but chosen so that
${\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i}){\bm{J}}({\bm{\theta}}_{0},\widetilde{{\bm{X}}}^{i})^{\top}$
gets a chance to be full-rank ie contain as much information as possible),
that is used only for that computation.

For all trainings, the subspace size is $s=10$.

For all trainings, the meta-optimizer is Adam with an initial learning rate of
$0.001$. The noise $\sigma_{\varepsilon}=0.05$, equal to the noise added to
the context data. We compute the NLL using Rasmussen & Williams (2005)’s
implementation.

When training with Mixt, we make sure to initialize ${\bm{s}}_{1}$ and
${\bm{s}}_{2}$ randomly with $\mathcal{N}(\bm{0},~{}0.5{\bm{I}})$, so that the
meta-learning can effectively differentiate the two clusters. Also, in the
Mixt case, the mean ${\bm{\mu}}$ is unique before computing the FIM. Once we
have found the projection directions, we initialize
$({\bm{\mu}}_{1},{\bm{\mu}}_{2})$ with the intermediate ${\bm{\mu}}$, thus
yielding two Gaussians.

####  A.6.2 Training details of MAML (multi-cluster case)

We reimplement and train a MAML baseline to compare our results with that of
MAML. A large part of these hyperparameters is directly inspired from Finn et
al. (2017)’s work.

The model is a neural network with $2$ hidden layers, $40$ neurons each, with
a ReLU non-linearity. In our multi-cluster experiment, $\alpha=2$: the
clusters consist of the sine cluster and the linear cluster (see Appendix
A.4).

The training is performed with $n=24$ tasks per epoch (with an infinite number
of available sine and linear tasks during training ie $N=+\infty$) and at each
epoch the context and query inputs are randomly drawn from $[-5,5]$. In
accordance with our equal probability assumption from Section 4.1.2, at each
epoch $n/2=12$ tasks come from the sine cluster and $n/2=12$ tasks come from
the linear cluster.

For all trainings, the number of context datapoints is $K=10$, and the number
of query datapoints is $L=10$. We meta-train for $70,000$ epochs.

For all trainings, the meta-optimizer is Adam with an initial learning rate of
$0.001$. The inner learning-rate is kept constant, at $0.001$. The number of
inner updates is $5$ during training, and $10$ at test time.

####  A.6.3 Training details of MMAML (multi-cluster case)

We train a MMAML by using Vuorio et al. (2019)’s code, using the best setting
mentionned in their paper (FiLM). We adapt their code to train it on our sine
and linear clusters: we add an offset to their sine cluster (via their
parameter bias), change the phase from $\sin(x-\varphi)$ to $\sin(x+\varphi)$
and specify via the arguments the slope range ($[-1,1]$) and the y-intercept
range ($[0,0]$, because it does not vary in our case) of the line tasks. We
also change the number of context inputs that we set to $K=10$, to remain
coherent with the rest of the trainings. The rest of the hyperparameters are
kept identical to the command specified in the repository of MMAML. Finally,
at test time we make the query ground-truths noiseless, to remain coherent
with the rest of the test conditions.

####  A.6.4 Test details (multi-cluster case)

For the OoD detection evaluation, we plot the AUC as a function of the number
of the context inputs $K$. The AUC is computed using the NLL of the context
inputs wrt to $\tilde{p}_{\xi}(f)$ (our uncertainty metric). The true-
positives are the OoD tasks (quadratic tasks) flagged as such; the false-
positives are the in-distribution tasks (lines and sines) flagged as OoD.

For the predictions, we plot the average and ci95 of the MSE on 1,000 tasks
(100 queried inputs each): half of them are sines and half of them are lines.

###  A.7 Training and test details for the vision problem

The model is a deep neural network, identical to the one used by Gao et al.
(2022) except for the last layer: instead of doing a one-dimensional
regression (where the output stands for an angle prediction), we perform a
two-dimensional regression (where the output stands for a cosine and sine
prediction). This choice is motivated by the fact that the Gaussian noise
assumed in Section 3.1 cannot capture the complexity of an angle error (e.g.,
predicting $361^{\circ}$ should yield a low error when compared to the ground-
truth angle $0^{\circ}$), but better renders the MSE that can be applied on
cosine and sine.

####  A.7.1 Training details of UnLiMiTD (vision case)

For all the trainings, there are $n=10$ tasks per epoch and $K=15$ context
inputs per task. When dealing with the UnLiMiTD-F case, we allow half of the
epochs ($5,000$) to the training before finding the intermediate
${\bm{\theta}}_{0}$ and ${\bm{\mu}}$ (see Algorithm 2), and the other half
after.

In the UnLiMiTD-F and UnLiMiTD-R cases, the subspace size is $s=100$.

For all the trainings, the meta-optimizer is Adam with an initial learning
rate of $0.001$. The noise is $\sigma_{\varepsilon}=0.01$. We compute the NLL
using Rasmussen & Williams (2005)’s implementation.

####  A.7.2 Training details of MAML (vision case)

We train a MAML baseline to compare our results with that of MAML. A large
part of these hyperparameters is directly inspired from Gao et al. (2022).

The training is performed with $n=10$ tasks per epoch. The number of context
datapoints is $K=15$, and the number of query datapoints is $L=10$. We meta-
train for $50,000$ epochs.

For all trainings, the meta-optimizer is Adam with an initial learning rate of
$0.0005$. The inner learning-rate is kept constant, at $0.002$. The number of
inner updates is $5$ during training, and $20$ a test time.

####  A.7.3 Test details (vision problem)

At test time, we wrap our model with the $\arctan$ function to convert the
predictions angle predictions. Then, we compare the angle predictions with the
ground-truth angles. To do so, we use the following error from Gao et al.
(2022) to compare two angles:

$\mathcal{E}(\beta,\beta^{*})=\min\\{\mathcal{E}_{\beta^{+},\beta^{*}},\mathcal{E}_{\beta,\beta^{*}},\mathcal{E}_{\beta^{-},\beta^{*}}\\}$

where $\mathcal{E}_{\beta^{\pm},\beta^{*}}=|y\pm 360-y^{*}|$ and
$\mathcal{E}_{\beta,\beta^{*}}=|y-y^{*}|$.

When plotting the performance (Figure 9), we plot the average and ci95 on 100
tasks (15 queried inputs each). For UnLiMiTD-R, the average and ci95 are
computed on 5 different random projection trainings.