MARS: M ETA - LEARNING AS SCORE MATCHING IN THE FUNCTION SPACE
A BSTRACT
Meta-learning aims to extract useful inductive biases from a set of related datasets.
In Bayesian meta-learning this is typically achieved by constructing a prior distribution over neural network parameters.
However, specifying families of computationally viable prior distributions over the high-dimensional neural network parameters is difficult.
As a result, existing approaches resort to meta-learning restrictive diagonal Gaussian priors, severely limiting their expressiveness and performance.
To circumvent these issues, we approach meta-learning through the lens of functional Bayesian neural network inference which views the prior as a stochastic process and performs inference in the function space.
Specifically, we view the meta-training tasks as samples from the data-generating process and formalize meta-learning as empirically estimating the law of this stochastic process.
Our approach can seamlessly acquire and represent complex prior knowledge by meta-learning the score function of the data-generating process marginals instead of parameter space priors.
In a comprehensive benchmark, we demonstrate that our method achieves state-of-the-art performance in terms of predictive accuracy and substantial improvements in the quality of uncertainty estimates
I NTRODUCTION
Using data from related tasks is of key importance for sample efficiency.
Meta-learning attempts to extract prior knowledge (i.e., inductive bias) about the unknown data generation process from these related tasks and embed it into the learner so that it generalizes better to new learning tasks (Thrun &
Pratt, 1998; Vanschoren, 2018).
Many meta-learning approaches try to amortize or re-learn the entire inference process (e.g., Santoro et al., 2016; Mishra et al., 2018; Garnelo et al., 2018) or significant parts of it (e.g., Finn et al., 2017; Yoon et al., 2018).
As a result, they require large amounts of meta-training data and are prone to meta-overfitting (Qin et al., 2018; Rothfuss et al., 2021a).
The Bayesian framework provides a sound and statistically optimal method for inference by combining prior knowledge about the data-generating process with new empirical evidence in the form of a dataset.
In this work, we adopt the Bayesian framework for inference at the task level and only focus on meta-learning informative Bayesian priors.
Previous approaches (Amit & Meir, 2018; Rothfuss et al., 2021a) meta-learn Bayesian Neural Network (BNN) prior distributions from a set of related datasets; by meta-learning the prior distribution and applying regularization at the meta-level, they facilitate positive transfer from only a handful of meta-training tasks.
However,
BNNs lack a parametric family of (meta-)learnable priors over the high-dimensional space of neural network (NN) parameters that is both computationally viable and, simultaneously, flexible enough to account for the over-parametrization of NNs.
In practice, both approaches use a Gaussian family of priors with a diagonal covariance matrix, which is too restrictive to accurately match the complex probabilistic structure of the data-generating process.
To address these shortcomings, we take a new approach to formulating the meta-learning problem and represent prior knowledge in a novel way.
We build on recent advances in functional approximate inference for BNNs that perform Bayesian inference in the function space rather than in the parameter space of neural networks (Wang et al., 2018; Sun et al., 2019).
When viewing the BNN prior and posterior as stochastic processes, the perfect Bayesian prior is the (true) data-generating ∗ Equal contribution.
stochastic process itself.
Hence, we view the meta-training datasets as samples from the meta-datagenerating process and interpret meta-learning as empirically estimating the law of this stochastic process.
More specifically, we meta-learn the score function of its marginal distributions, which can then directly be used as a source of prior knowledge when performing approximate functional BNN inference on a new target task.
This ultimately allows us to use flexible neural network models for learning the score and overcome the issues of meta-learning BNN priors in the parameter space.
In our experiments, we demonstrate that our proposed approach, called Meta-learning via Attentionbased Regularised Score estimation (MARS) , consistently outperforms previous meta-learners in predictive accuracy and yields significant improvements in the quality of uncertainty estimates.
Notably, MARS enables positive transfer from only a handful of tasks while maintaining reliable uncertainty estimates.
This promises fruitful future applications to domains like molecular biology or medicine, where meta-training data is scarce and reasoning about epistemic uncertainty is crucial.
R ELATED W ORK
Meta-Learning.
Common approaches in meta-learning amortize the entire inference process (Santoro et al., 2016; Mishra et al., 2018; Ravi & Beatson, 2018; Garnelo et al., 2018), learn a good neural network initialization (Finn et al., 2017;
Rothfuss et al., 2019; Nichol et al., 2018;
Kim et al., 2018) or a shared embedding space (Baxter, 2000; Vinyals et al., 2016; Snell et al., 2017).
Although these approaches can meta-learn complex inference patterns, they require a large amount of metatraining data and often perform poorly in settings where data is scarce.
Another line of work uses a hierarchical Bayesian approach to meta-learn priors over the NN parameters (Pentina & Lampert, 2014; Amit & Meir, 2018;
Rothfuss et al., 2021a).
Such methods perform much better on small data.
However, they suffer from the lack of expressive families of priors for the high-dimensional and complex parameter space of NNs, making too restrictive assumptions to represent complex inductive biases.
Our approach overcomes these issues by viewing the problem in the function space and directly learning the score, which can easily be represented by a NN instead of a prior distribution.
Also related to our stochastic process approach are methods that meta-learn Gaussian Process (GP) priors (Fortuin et al., 2019;
Rothfuss et al., 2021b; 2022).
However, the GP assumption is quite limiting, while MARS can, in principle, match the marginals of any data-generating process.
Score estimation
We use score estimation as a central element of our meta-learning method.
In particular, we use a parametric approach to score matching and employ an extended version of the score matching objective of Hyv ¨ arinen & Dayan (2005).
For high-dimensional problems, Song et al.
; Pang et al.
propose randomly sliced variations of the score matching loss.
Alternatively, there is a body of work on nonparametric score estimation (Canu & Smola, 2006; Liu et al., 2016; Shi et al., 2018; Engl et al., 1996; Zhou et al., 2020).
Among those, the Spectral Stein
Gradient Estimator (Shi et al., 2018) has been used for estimating the stochastic process marginals for functional BNN inference in a setting where the stochastic prior is user-defined and allows for generating arbitrarily many samples (Sun et al., 2019).
Such estimators make it much harder to add explicit dependence on the measurement sets and prevent meta-overfitting via regularization, making them less suited to our problem setting.
B ACKGROUND
Bayesian Neural Networks.
Consider a regression task with data D = (cid:0) X D , y D (cid:1) that consists of m i.i.d.
noisy function evaluations y j = f ( x j )
+ ϵ j of an unknown function f : X (cid:55)→ Y .
Here, X D = { x j } mj =1 ∈
X m denotes training inputs and y D
= {
y j } mj =1 ∈ Y m the corresponding noisy function values.
Let h θ : X → Y be a function parametrized by a NN with weights θ ∈ Θ .
For regression, where Y ⊆ R , we can use h θ to define a conditional distribution over the noisy observations p ( y | x , θ ) =
N ( y | h θ ( x ) , σ 2 ) , with flexible mean represented by h θ and observation noise variance σ 2 .
Given a prior distribution p ( θ ) over the model parameters θ , Bayes’ theorem yields a posterior distribution p ( θ | X D , y D )
∝ p ( y D | X D , θ ) p ( θ ) , where p ( y D | X D , θ ) =
(cid:81) mj
=1 p ( y j | x j , θ ) .
For an unseen test point x ∗ , we compute the predictive distribution, which is defined as p ( y ∗ | x ∗ , X D , y D ) =
(cid:82) p ( y ∗ | x ∗ , θ ) p ( θ | X D , y D ) dθ and obtained by marginalizing out the posterior over parameters θ .
BNN inference in the function space.
Posterior inference for BNNs is difficult due to the high-dimensional parameter space Θ and the over-parameterized nature of the NN mapping h θ ( x ) .
An alternative approach views BNN inference in the function space, i.e., the space of regression functions h : X (cid:55)→ Y , yielding the Bayes rule p ( h | X D , y D )
∝ p ( y D | X D , h ) p ( h ) (Wang et al., 2019; Sun et al., 2019).
Here, p ( h ) is a stochastic process prior with index set X , taking values in Y .
Stochastic processes can be understood as infinite-dimensional random vectors, and, thus, are hard to work with computationally.
However, given finite measurement sets X :=
[ x 1 , ..., x k ] ∈ X k , k ∈ N , the stochastic process can be characterized by its corresponding marginal distributions of function values ρ ( h X ) :
=
ρ ( h ( x 1 ) , ...h ( x k ))
(cf.
Kolmogorov Extension Theorem, Oksendal (2013)).
Thus, we can break down the functional posterior into a more tractable form by re-phrasing it in terms of posterior marginals based on measurement sets X : p ( h X | X , X D , y D )
∝ p ( y D | h X D ) p ( h X ) .
This functional posterior can be tractably approximated to achieve functional BNN inference (Sun et al., 2019; Wang et al., 2019).
We briefly describe how this can be done via functional Stein
Variational Gradient Descent (fSVGD, Wang et al., 2019).
The procedure is explained in more detail in Appendix A.1.
The fSVGD algorithm approximates the posterior using a set of L NN parameter particles { θ 1 , ..., θ L } .
To optimize posterior approximation, fSVGD iteratively re-samples measurement sets X from a measurement distribution ν , supported on X , computes SVGD updates (Liu & Wang, 2016) in the function space, and projects them back into the parameter space, in order to update the parameter-space particle configuration approximating the posterior.
To achieve this, fSVGD uses the functional posterior score, i.e., the gradient of the log-density ∇ h
X ln p
[Equation]
=
[Equation]
y D |
[Equation]
X D )
+
[Equation]
X ln p ( h X ) to guide the particles towards areas of high posterior probability.
Formally, for all l = 1 , ..., L the particle updates are computed as
[Equation]
where K = [ k ( h X θ l , h X θ i )] li is the kernel matrix between the function values in the measurement points based on a positive semi-definite kernel function k ( · , · ) .
A key insight that will later draw upon in the paper is that such functional approximate inference techniques only use the prior scores and, in principle, do not require a full prior density that integrates to 1.
M ETA -L EARNING AS S CORE E STIMATION IN THE F UNCTION S PACE 4. P ROBLEM S TATEMENT : M ETA -L EARNING
Meta-learning extracts prior knowledge (i.e., inductive bias) from a set of related learning tasks to accelerate inference on a new and unseen learning task.
The meta-learner is given n datasets
D 1 , ..., D n , where each dataset D i =
( X D i , y D i ) consists of m i noisy function evaluations
y i,t = f i ( x i,t )
+ ϵ corresponding to a function
f i : X (cid:55)→ Y ⊆ R and additive σ sub-Gaussian noise ϵ .
In short, we write X D i =
( x i, 1 , ..., x i,m i ) ⊤ for the matrix of function inputs and y D i
=
( y i, 1 , ..., y i,m i ) ⊤ for the vector of corresponding observations.
Following previous work (e.g., Baxter, 2000; Pentina & Lampert, 2014; Rothfuss et al., 2021a), we assume that the functions f i ∼ T are sampled i.i.d.
from a task distribution T , which can be thought of as a stochastic process p ( f ) that governs the random function f : X (cid:55)→ Y .
Our goal is to extract knowledge from the observed datasets, which can then be used as a form of prior for learning a new, unknown target function f ∗ ∼ T from a corresponding dataset D ∗ .
By performing Bayesian inference with a meta-learned prior that is attuned to the task distribution T , we hope to obtain posterior predictions that generalize better.
4. S HORTCOMINGS OF M ETA -L EARNING P RIORS IN THE P ARAMETER S PACE
Previous work phrases meta-learning as hierarchical Bayesian problems and tries to meta-learn a prior distribution p ( θ ) on the NN weight space that resembles p ( f ) (e.g., Amit & Meir, 2018;
Rothfuss et al., 2021a).
This approach suffers from the following central issues: Due to the overparameterized nature of neural networks, a large set of parameters θ correspond to exactly the same mapping h θ ( · ) .
This makes specifying a good prior distribution p ( θ ) typically very difficult.
At the same time, there is a lack of sufficiently rich parametric families of distributions over the highdimensional parameter space Θ that are also computationally viable.
So far, only the very restrictive
Gaussian family of priors with diagonal covariance matrices has been employed.
As a result, the meta-learned posterior lacks the necessary flexibility to match the complex probabilistic structure of the data-generative process accurately.
4. M ETA -L EARNING AS S CORE E STIMATION ON THE D ATA -G ENERATING P ROCESS
Aiming to address these issues, we acquire prior knowledge in a data-driven manner with a new perspective.
We develop a novel approach to meta-learning which hinges upon three key ideas:
First, we view BNN inference in the function space (see Sec. 3), i.e., as posterior inference p ( h | X D , y D )
∝ p ( y D | X D , h ) p ( h ) over neural network mappings h θ : X (cid:55)→ Y instead of parameters θ .
From this viewpoint, the prior, which is the target of our meta-learning problem, is p ( h ) , a stochastic process on the function space.
This alleviates the aforementioned problem of over-parametrization, which arises when considering priors on the parameter space.
Second, we ask ourselves what is a desirable prior: from a Bayesian perspective, the best possible prior is the stochastic process of the task generating distribution T itself, i.e., p ( h ) =
p ( f ) .
Hence, we would want to meta-learn a prior that matches p ( f ) as closely as possible.
Since the meta-training datasets D 1 , ..., D n constitute noisy observations of the function draws f 1 , ...., f n ∼ p ( f ) , we can use them to estimate the stochastic process marginals p ( f X ) .
Crucially, we tractably represent and estimate the stochastic process p ( f ) through its finite marginals p ( f X ) in measurement sets X .
Our third insight is that all popular approximate inference methods for BNNs only use the prior score, i.e., the gradient of the log-prior, and not the prior distribution itself.
In addition, parametrizing and estimating the prior score is computationally easier than the prior distribution itself since it does not have to integrate to 1, allowing for much more flexible neural network representations.
For this reason, our meta-learning approach directly forms an estimate of the scores s ( f X , X ) =
∇
f
X ln p ( f X ) of the data-generating process marginals.
In summary, our high-level approach is to meta-learn / estimate the prior score ∇ h X ln p ( h X ) that matches data-generating stochastic process from meta-training data D 1 , ..., D n .
At meta-test time, the meta-learned prior marginals are used in the approximate functional BNN inference on a target dataset D ∗ to infuse the acquired prior knowledge into the posterior predictions.
In the following section, we discuss in more detail how to implement this general approach as a concrete, computationally feasible algorithm with strong empirical performance.
T HE MARS META - LEARNING ALGORITHM
In Section 4, we have motivated and sketched the idea of meta-learning via score estimation of the data-generating process marginals.
We now discuss particular challenges and outline the design choices for translating this idea into a practical meta-learning algorithm.
Score estimation of stochastic process marginals has previously been used in the context of existing functional BNN approaches (Wang et al., 2019; Sun et al., 2019).
However, they presume a known stochastic process prior with oracle access to arbitrarily many sampled functions that can be evaluated in arbitrary locations of the domain X .
Compared to such a stochastic process prior with oracle access, we face two key challenges which make the score estimation problem much harder:
We only have access to a finite number of datasets {D i } ni =1 , each corresponding to one function
f
i drawn from the data-generating stochastic process.
This means, we have so estimate the marginal scores ∇ h X ln p ( h X ) from only n function samples without over-fitting.
Per function f
i , we only have a limited number of m i noisy function evaluations y D
i in input locations
X D i which are given to us exogenously, and we have no control over them.
However, to perform score estimation for the marginal p ( h X ) distributions, for each function, we need a vector of function values f X i evaluated in the measurements points X .
In the following subsections, we present our approach to estimating the stochastic prior scores and discuss how to solve these challenges.
In addition, we discuss how to perform functional approximate BNN inference with the meta-learned score estimates.
5. P ARAMETRIC S CORE M ATCHING FOR S TOCHASTIC P ROCESS M ARGINALS
Performing functional approximate inference with fSVGD updates as in (1) requires estimation of the prior marginal scores ∇ h X ln p ( h X ) for arbitrary measurement sets X .
A key property of stochastic process marginal scores is their permutation equivariance: if we permute the order of the points in the measurement set, the rows of the score permute in the same manner.
Similarly, the measurement sets can be of different sizes, implying score vectors/matrices of different dimensionalities.
Aiming to embed these properties into a parametric model of the prior marginal scores, we use a transformer encoder architecture (Vaswani et al., 2017), that takes as an input a measurement set
X ∈ R dim ( X ) × k of k points and corresponding query function values h X ∈ R dim ( Y )
× k and performs attention over the second dimension, i.e., the k columns corresponding to measurement points.
Our score network model is denoted as s ϕ ( h X , X ) with trainable parameters ϕ
and outputs an estimate of the score, i.e., s ϕ
( h X , X )
(cid:55)→ ˜ ∇ h X p ( h X ) , a matrix of the same size as the query function values.
The model is permutation equivariant w.r.t.
the columns corresponding to measurement points and supports inputs/outputs for varying measurement set sizes k .
We schematically illustrate the score network architecture in Figure 4 and provide details in Appendix A.3.
We train the score network with a modified version of the score matching loss (Hyv¨arinen & Dayan, 2005; Song et al., 2020).
This loss minimizes the Fisher divergence between the data-generating process marginals and our score network, based on samples f X 1 , ..., f X n from the data-generating process marginals.
In contrast to the standard score matching loss, our modified loss takes the dependence of stochastic process marginals on their measurement sets into account.
It uses an expectation over randomly sampled measurement sets, similar to the functional approximate inference in Section 3.
[Equation]
As noted previously, our meta-training data only provides noisy function evaluations for a limited, exogenously given, set of input locations per task f i .
The input locations X D i per task might even differ in number and be non-overlapping.
Thus, we do not have access to the function values f X 1 , ..., f X n for arbitrary measurement sets X which we require for our score matching loss in (2).
In the following section, we discuss how to resolve this problem.
5. I NTERPOLATING THE D ATASETS ACROSS X
In our meta-learning setup (see Section 4.1), we only receive noisy function evaluations y D i in m i locations
X D i for each meta-training task.
However, to minimize the score matching loss in (2), we need the corresponding function values in arbitrary measurement locations in the domain X .
Hence, we interpolate each dataset D i by a regression model.
Importantly, the regression model should be able to quantify the epistemic uncertainty that stems from having observations at only a finite subset of points X D i ⊂ X .
While a range of method could be used, we employ a Gaussian
Process (GP) with Mat´ern − 5 / 2 kernel or a Bayesian neural network (BNN) based on Monte Carlo dropout (Gal & Ghahramani, 2016) for this purpose.
We fit the respective Bayesian model independently on each dataset D i which gives us a posterior p ( f X i | X , X D i , y D i ) over function values f
X i for arbitrary measurement sets X .
In Appx.
A.4 we provide more details on how we fit each model.
While we could simply use posterior mean values for computing the score matching loss in 2, this would not reflect the epistemic interpolation uncertainty.
Instead, for each task, we sample function values ˜ f
X i ∼ p ( f X i | X , X D
i , y D i ) from the corresponding posterior and use them as an input to the score matching loss in (2).
Correspondingly, the score matching loss modifies to
[Equation]
The further a measurement point x is away from the closest point in X D i , the higher is the (epistemic) uncertainty of the corresponding function value f i ( x ) .
By sampling from the posterior, in expectation, the loss in (3) effectively propagates the interpolation uncertainty into the score matching procedure, preventing over-confident score estimates for areas of the domain X with scarce data.
This is illustrated in Appx.
D, Figure 6.
5. P REVENTING M ETA -O VERFITTING OF THE P RIOR S CORE N ETWORK
The final challenge we need to address is over-fitting to the meta-training tasks (Yin et al., 2020;
Rothfuss et al., 2021a).
As we only have meta-training data {D i } n i =1 , corresponding to n functions
Algorithm
Input: datasets D 1 , ..., D n , measurement point distribution ν , step size η
Initialize score network parameters ϕ for i = 1 , ..., n do fit GP or BNN on D i , obtain posterior p ( f i |
X D i , y D i ) while not converged do
X iid ∼ ν // sample measurement set for i = 1 , ..., n do ˜ f
X i ∼ p ( f X i | X , X D
i , y D i ) //
sample function values from posterior marginal ˆ L ( ϕ ) ← 1 n (cid:80) ni
=1 (cid:16) tr ( ∇ ˜ f
X i s ϕ
(˜ f X i , X ))
+
12 || s ϕ (˜ f
X i , X ) || 22 (cid:17) // score matching loss ϕ ← ϕ + η ∇ ϕ
ˆ L ( ϕ ) // score network gradient update
Output: trained score network s ϕ f
i drawn from the data-generating stochastic process, i.e., we only have n samples for estimation of marginal scores ∇ h X ln p ( h X ) , making us prone to overfitting the prior score network.
If we performed functional BNN inference with such an overfitted score estimate, the BNN’s posterior predictions are likely to be over-confident and too biased towards the functions seen during meta-training.
To counteract the tendency to overfit, we regularize the score network via spectral normalization (Miyato et al., 2018) of the linear layers in the transformer encoder blocks.
Spectral normalization controls the Lipschitz constant of the neural network layers by dividing their weight matrix W by its spectral norm
|| W || , i.e., re-parametrizing the weights as ˜ W :
= W / || W || .
Hence, by applying spectral normalization, we bias our score network towards smoother score estimates corresponding to marginal distributions p ( h X ) with higher entropy.
Empirically we find that spectral normalization effectively combats meta-overfitting and prevents the estimated prior score from inducing over-confident posterior predictions when employed in functional BNNs.
In our experiments, we examined several other regularization methods, e.g., gradient penalties (Gulrajani et al., 2017) or a spectral regularization penalty in the loss, but found spectral normalization to work the best.
5. T HE FULL MARS ALGORITHM
We now summarize our meta-learning algorithm MARS.
It consists of two stages:
Stage 1: Meta-Learning the Prior Score Network.
After initializing the parameters of the score network, we fit a GP or BNN to each of the n meta-training datasets D i .
Then, we train the score network by stochastic gradient descent on the modified score matching loss ˜
L ( ϕ ) in (3).
In each iteration, we first sample a measurement set X i.i.d.
from the measurement distribution ν = U ( ˜ X ) , chosen as uniform distribution over the hypercube ˜ X ⊆
X conservatively covering the data in X .
Based on the measurement set, we sample a vector of functions values ˜ f
X i from the corresponding
GP or BNN posterior marginals p ( f X i | X , X D i , y D i ) , i = 1 , ..., n .
Based on these samples, we form an empirical estimate ˆ L ( ϕ ) of the score matching loss and perform a gradient update step on the score network parameters ϕ .
This is repeated till convergence and summarized in Alg. 1.
Depending on whether we use a GP or BNN as interpolator, we refer to our method as MARS-GP or MARS-BNN.
Stage 2: Functional BNN Inference with the Prior Score Network.
When concerned with a target learning task with a training dataset D ∗ =
( X D∗ , y D∗ ) , we can infuse the meta-learned inductive bias into the BNN inference by using the prior score network s ϕ from Stage 1 for the approximate inference.
In particular, we can either perform functional VI (Sun et al., 2019) or fSVGD (Wang et al., 2019), using the score network s ϕ ( h X , X ) predictions as swap-in replacement for the marginal scores ∇ h X ln p ( h X ) of a user-specified stochastic process prior.
In our experiments, we use fSVGD (see Section 3) as functional approximate inference method.
The resulting fSVGD BNN inference procedure with our meta-learned priors score network is summarized in Alg. 2 in Appx.
A.
E XPERIMENTS
We provide a detailed benchmark comparison with existing meta-learning methods, demonstrating that MARS:
(i) achieves state-of-the-art performance in terms of predictive accuracy, (ii) yields
4 x   f ( x ) true fun pred mean train data pred std (a) Vanilla BNN
4 x   f ( x ) (b) fBNN with GP prior    4 x   f
( x ) (c) MARS (ours)
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
well-calibrated uncertainty estimates.
6. E XPERIMENT S ETUP Environments.
We consider five realistic meta-learning environments for regression .
The first environment corresponds to data of different calibration sessions of the Swiss Free Electron Laser ( SwissFEL ) (Milne et al., 2017).
Here, a task requires predicting the laser’s beam intensity based on undulators parameters.
We also consider electronic health measurements ( PhysioNet ) from intensive care patients (Silva et al., 2012), in particular, the Glasgow Coma Scale ( GCS ) and the hematocrit value ( HCT ).
Here, each task correspond to a patient.
We also use the Intel Berkeley Lab temperature sensor dataset ( Berkeley-Sensor )
(Madden, 2004), requiring auto-regressive prediction of temperature measurements of sensors in different locations of the building.
Finally, the ArgusControl environment requires predicting the total variation of a robot from its target position based on its PID controller parameters (Rothfuss et al., 2022).
Here, tasks correspond to different step sizes between the source and target location.
See Appendix B for details.
Baselines.
As non-meta-learning baselines, we use a Vanilla GP with RBF kernel and a Vanilla BNN with a zero-mean, spherical Gaussian prior and SVGD inference (Liu & Wang, 2016).
To compare to previous meta-learners, we report results for model agnostic meta-learning ( MAML )
(Finn et al., 2017), Bayesian MAML ( BMAML ) (Yoon et al., 2018) and neural processes ( NPs ) (Garnelo et al., 2018).
In addition, we include PACOH (Rothfuss et al., 2021a) which performs hierarchical Bayes inference to meta-learn a GP prior ( PACOH-GP ) or prior over NN parameters ( PACOH-NN ).
6. E MPIRICAL B ENCHMARK S TUDY
Qualitative illustration.
Fig.
1 illustrates the posterior predictions of a) a BNN with Gaussian Prior in the parameter space, b) a functional BNN with GP prior (Wang et al., 2019) and c) a fBNN trained with the meta-learned MARS-GP marginal scores.
For the meta-training we use n = 20 tasks with each m = 8 data points, generated from a student-t process with sinusoidal mean function 2 x + 5 sin(2 x ) .
Similar to previous work (e.g., Fortuin et al., 2022) we observe that weight-space BNN priors provide poor inductive bias and yield grossly over-confident uncertainty estimates.
Approaching the BNN inference problem in the function space (see Fig 1b) results in much better uncertainty estimates.
In this case, however, the uninformative GP prior does not result in accurate mean predictions.
Finally, the meta-learned scores in the function space provide useful inductive biased towards the linear+sinusoidal data-generating pattern while yielding tight but reliable confidence intervals.
[Table] [Table]
MARS provides accurate predictions.
We perform a comprehensive benchmark study with the environments and baselines introduced in Sec. 6.1.
Table 1 reports RMSE on unseen meta-test tasks.
Both MARS-GP and MARS-BNN yield substantial improvements in the RMSE compared to the non-meta-learning baselines and significantly outperforms all other meta-learning baselines in most environments.
This shows that MARS can acquire valuable inductive biases from data reliably.
We hypothesize that MARS performs substantially better than PACOH-NN, which meta-learns restrictive diagonal Gaussian priors over Θ since it has much more flexibility in expressing prior knowledge about the data-generating process.
In the majority of environments, MARS also outperforms
PACOH-GP which meta-learns a GP prior.
In contrast to Gaussian marginals, MARS can meta-learn any stochastic process marginal, and, thus has much more flexibility to express inductive bias.
MARS yields well-calibrated uncertainty estimates.
We hypothesize that by performing meta-learning in the function space, we avoid the pitfalls of NN over-parametrization, which often lead to over-confidence.
To investigate the quality of uncertainty estimates, we compute the calibration error , which measures how much the predicted confidence regions deviate from the actual frequencies of test data in the respective regions.
We list the results in Table 2 1 .
While, in most cases, still better than the baselines, the uncertainty estimates MARS-BNN are on average worse than those of MARS-GP.
This is not surprising as the MC-dropout uncertainty estimates are typically not as reliable as those of a GP.
MARS-GP consistently yields the best-calibrated uncertainty estimates and, for some environments, reduces the calibration error by more than 30 % compared to the next best baseline.
This provides further evidence for the soundness and efficacy of our function space approach to meta-learning.
Finally, the superior calibration performance suggests that the spectral regularization, together with the posterior uncertainty sampling in Sec. 5, effectively prevents meta-overfitting and accounts for epistemic uncertainty.
6. A BLATION S TUDY
We empirically investigate the algorithm components introduced in Sec. 5 and provide supporting empirical evidence for our design decisions.
First, we perform a quantitative ablation study where we vary/remove components of our algorithm.
We consider MARS-GP with 1) nonparametric score estimator SSGE (Shi et al., 2018) instead of the parametric score network + score matching from
Sec. 5.1, 2) without spectral regularization, and 3) without GP posterior sampling, i.e., using the
GP posterior mean instead of samples for the score matching.
Table 3 and 4 report the quantitative results of this ablation experiment.
In the following, we discuss the three aspects separately:
Parametric score matching outperforms nonparametric score estimation.
In Sec.
5.1 we have introduced a parametric score network which we train with the score matching loss.
Alternatively, for each measurement set X , we could apply the nonparametric score estimator SSGE (Shi et al., 2018) to the function values ˜ f X 1 , ..., ˜ f
X n , sampled from the posteriors, and directly use the resulting
[Figure]
[Table]
[Figure]
MARS w/o regularization overfit and underestimate the variance, whereas spectral normalization biases MARS towards higher entropy.
score estimates during the fSVGD posterior inference in (1).
This produces score estimates ’ad-hoc’ and does not require us to train an explicit score network.
In contrast, our score network is a global function that explicitly considers the measurement set and thus can exploit similarity structure across
X , which SSGE cannot.
To SSGE we also cannot simply add inductive bias towards higher entropy as we do through spectral normalization.
Overall, the experiment results in Table 3 and 4 suggest that instantiating our general approach with SSGE performs worse than MARS.
We can also observe this visually in Fig.
2a, where we plot the score estimate of MARS and SSGE for a GP marginal where the true score is known.
While the MARS score network slightly overestimates the variance of true generative-process marginal, SSGE implicitly underestimates the prior variance, leading to over-confident predictions.
Finally, Table 7 in Appx.
D quantitatively shows that the MARS score network provides the most accurate estimates compared to a variety of nonparametric estimators.
Spectral normalization prevents meta-overfitting.
In Sec.
5.3, we added spectral normalization to our score network.
Here, we empirically investigate what happens when we remove the spectral normalization from MARS.
Figure 2b and 2c illustrate how our score network over-fits and under-estimates the true data-generating variance when we do not use spectral regularization.
Quantitatively, we observe substantial increases in calibration errors and a consistent worsening of the predictive accuracy.
Overall, this highlights the empirical importance of spectral normalization for preventing meta-overfitting and biasing the score estimates toward higher entropy.
Accounting for epistemic uncertainty of the GP interpolators is crucial.
In MARS, we interpolate each dataset with a GP or BNN and, when performing score estimation, use samples from the posterior marginals to account for the epistemic uncertainty of the interpolation.
Here, we empirically study what happens if we ignore the epistemic uncertainty and take each GP’s mean predictions.
When doing so, we observe detrimental effects on the RMSE and calibration error in the majority of the environments in Table 3 and 4.
This affirms that propagating the epistemic uncertainty of the interpolation into the score estimates is a critical component of MARS.
C ONCLUSION
We have introduced a novel meta-learning approach in the function space that estimates the score of the data-generating process marginals from a set of related datasets.
When facing a new learning task, we use the meta-learned score network as a source of prior for functional approximate BNN inference.
By representing inductive bias as the score of a stochastic process, our approach is versatile and can seamlessly acquire/represent complex prior knowledge.
Empirically, this translates into strong performance when compared to previous meta-learning methods.
The substantial improvements of MARS in terms of the quality of uncertainty estimates open up many potential extensions toward interactive machine learning where exploration based on epistemic uncertainty is vital.
A CKNOWLEDGMENTS
This research was supported by the European Research Council (ERC) under the European Union’s
Horizon 2020 research and innovation program grant agreement
no. 815943 and the Swiss National
Science Foundation under NCCR Automation, grant agreement 51NF40 180545.
Jonas Rothfuss was supported by an Apple Scholars in AI/ML fellowship.
We thank Alex H¨agele, Parnian Kassraie,
Lars Lorch and Danica J. Sutherland for their valuable feedback.
A MARS IMPLEMENTATION DETAILS
This section focuses on a detailed explanation of the proposed framework.
We first elaborate on the fSVGD inference process and then provide information regarding the network architecture and the hyperparameter configuration.
A. F SVGD INFERENCE
To perform approximate BNN inference with the meta-learned score network, we use functional
Stein Variational Gradient Descent (fSVGD) (Wang et al., 2018).
In the following, we explain our fSVGD implementation and how it interplays with the score network in our MARS framework.
In our context, the goal of fSVGD is to approximate the posterior p ( h θ | X D , y D )
∝ p ( y D | X D , h θ ) p ( h θ ) over neural network mappings h θ .
Recall from Section 3, that X D = { x j } mj =1 ∈
X m are the training inputs, y D
= { y j } mj =1 ∈ Y m the corresponding noisy function values for a new, unseen learning tasks.
The fSVGD algorithm approximates the posterior using a set of L NN parameter particles { θ 1 , ..., θ L } where each θ l corresponds to the weights and biases of a neural network.
The particles (i.e., weights and biases) are initialized based on the scheme of
Steinwart (2019) (see Appendix A.2 for details).
To make the BNN inference in the function space tractable, in each iteration, fSVGD samples a measurement set X from a measurement distribution ν and performs its particle updates based on the posterior marginals p ( h X | X , X D , y D )
∝ p ( y D | h X D ) p ( h X ) corresponding to X .
We use a uniform distribution over the bounded domain X as measurement distribution ν and sample the L measurement points i.i.d.
from it.
For each NN particle, we compute the NN function values in the measurement points, i.e. h X θ l =
( h lθ ( x 1 ) , ..., h lθ ( x k )) , l = 1 , ..., L and the corresponding posterior marginal scores: ∇
[Equation]
Here, we use a Gaussian likelihood p ( y D | h X D ) =
(cid:81) kj =1 N ( y D
i ; h lθ ( x j ) , σ 2 ) where σ 2 is the likelihood variance.
Unlike in Wang et al.
(2018) where the stochastic process prior is exogenously given, and its marginals p ( h X ) are approximated as multivariate Gaussians, we use our meta-learned marginal prior scores.
In particular, we use the score networks s ϕ ( h X θ l , X ) output as a swap-in for the prior score ∇ h X θl
ln ˆ p ( h X θ l ) .
Finally, based on the score in (4) and the function values h X θ l , we can compute the SVGD updates in the function space, and project them back into the parameter space via the NN Jacobian ∇ θ l h X θ l :
[Equation]
Here, γ is the SVGD step size and K =
[ k ( h X θ l , h
X θ i )]
li is the kernel matrix between the function values in the measurement points based on a kernel function k ( · , · ) :
Y k ×Y k (cid:55)→ R .
We use the RBF kernel k ( h , h ′ ) =
exp (cid:0) ( −|| h − h ′
|| 2 ) / (2 ℓ k ) (cid:1) where ℓ k is the bandwidth hyper-parameter.
The particle update in (5) completes one iteration of fSVGD.
We list the full procedure in Algorithm 2.
fSVGD hyperparameter selection.
For the fBNN training using fSVGD, among other parameters, we need to choose the step size γ , kernel bandwidth ℓ k , and likelihood standard deviation σ , as we found these three to have the most substantial impact on training dynamics.
We fix the number of particles to L = 10 and perform 10000 fSVGD update steps.
We always standardize both the input and output data based on the meta-training data’s empirical mean and standard deviation.
For the NNs, we use three hidden layers, each of size 32 and with leaky ReLU activations.
To initialize the NN weights, we use He initialization with a uniform distribution (He et al., 2015) and the bias initializer of Steinwart (2019) (see Section A.2 for details).
Generally, we choose the
Input: SVGD kernel function k ( · , · ; ℓ k ) , bandwidth ℓ k , step size γ
Input: dataset D ∗ =
( X D∗ , y D∗ ) for target task, trained score network s ϕ ( · , · )
Initialize set of BNN particles { θ 1 , ..., θ L } while not converged do
X iid ∼ ν // sample measurement set for l = 1 , ..., L do h X θ l ← ( h lθ ( x 1 ) , ..., h lθ ( x k ))
where
X =
( x 1 , ..., x k ) // compute NN function values in X ∇
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
// posterior marginal score
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
K
[Figure]
(cid:19) //
fSVGD update
Output: Set of BNN particles { θ 1 , ..., θ L } that approximate the BNN posterior process
[Figure]
a) The underlying data generating process and corresponding fBNN prediction after training using the fSVGD algorithm for 2000 iterations with b) constant bias initializer, and c) Steinwart bias initializer.
With the Steinwart bias initializer, we get the desirable non-linear behavior of the fSVGD BNN much faster.
kernel bandwidth ℓ k via a random hyper-parameter search over the values range of [0 . 1 , 10] .
When comparing to the original fSVGD implementation with SSGE, we fix the SSGE lengthscale to 0 .for comparability.
Note that we also experimented with using the median heuristic (as proposed by
Shi et al.
; however, we observed this to yield inferior performance.
A. B IAS INITIALIZATION
When initializing the biases to zeros or small positive constants, we find that the learned neural network maps behave like linear functions further away from zero and that it takes many SVGD iterations for them to assume non-linear behavior at the boundaries of the domain.
To address this issue, we use the bias initialization scheme of Steinwart (2019), which initializes the biases in such a way that the kinks of the leaky ReLU functions are more evenly distributed across the domain and less concentrated around zero.
More specifically, we initialize each bias b i as
[Equation]
where w i is uniformly sampled from a sphere by taking w i
=
a i || a
i || , a i ∼ U (0 , 1) .
x ∗ i ’s are sampled uniformly: x ∗ i ∼ U ( min x , max x ) , where min x and max x are the minimum and maximum points in the input domain respectively.
By using Steinwart initialization, the neural networks show much more non-linear behavior after initialization.
Compared to constant bias initialization, learning non-linear functional relationships happens much more quickly.
We showcase this in Figure 3: Figure 3a displays a simple data generation process of sinusoids of varying amplitude, frequency, phase shift, and slope.
Figure 3b displays corresponding BNN predictions after 2000 iterations with a constant bias initializer (where the constant equals 0.01), and Figure 3c shows the result of the same training dynamics, using Steinwart’s bias initializer instead of the constant initializer.
The BNNs with the Steinwart initialization assume desirable non-linear behavior much earlier during training, thus, significantly speeding up training.
However, if trained for a large enough number of iterations, the performance of the two networks with different bias initializations becomes much less distinct.
For further details on the implications of the Steinwart initializer on the training dynamics, we refer to Steinwart (2019).
A. S CORE ESTIMATION NETWORK
We now give an overview of the score estimation network.
We start by providing motivation for the required architecture, detailing the permutation equivariance properties of the network.
We construct the proposed architecture step-by-step, giving the architectural details and mentioning additional architectural designs we experimented with.
Finally, we comment on the optimization method and acknowledge the libraries we used in our implementations.
Incorporating task invariances.
The proposed network is permutation equivariant across the k dimension: reordering the measurement inputs would result in reordering the network’s prediction in the same manner: Formally, for any permutation π of the measurement set indices we have that s ϕ
( h X π (1: d ) , X π (1: d ) )
i,j (cid:55)→ ˜ ∇ h X 1: d p ( h X 1: d ) π ( i ) ,π ( j ) .
To impose permutation equivariance, we use the self-attention mechanism (Vaswani et al., 2017).
Permutation equivariance is obtained by concatenating inputs (measurement points) and the corresponding functional evaluations (i.e., the concatenation constructs an object corresponding to Transformer tokens ), which are then embedded and inputted to the attention mechanism.
The architecture is displayed in Figure 4.
Constructing the network architecture.
The core of our model is composed of two identical blocks, each consisting of two residual layers, the first one applying multi-head self-attention and the second one position-wise feed-forward neural network, similar to the vanilla Transformer encoder (Vaswani et al., 2017).
Since the multi-head attention is permutation equivariant over the measurement point (i.e., token) dimension, the representation is permutation equivariant at all times (Lee et al., 2019).
Finally, to minimize training time, we select attention embedding dimensions proportional to the data dimension of the environments; higher-dimensional environments (e.g., SwissFEL, detailed in Appendix B) correspond to the higher number of attention parameters/embedding dimensions and lower environments (e.g., the Sinusoid and Berkeley environment) to lower embeddings.
Furthermore, we tune the step size and report the chosen configurations under https://tinyurl.com/376wp8xe .
We describe the corresponding implementation details in the following subsection.
Constructing the score network.
We train the score network on the input/output pair concatenations, which are then embedded onto a higher-dimensional space.
After this, we perform selfattention.
Specifically, we fix the initialization scale of self-attention weights to 2.0.
As mentioned, the size of the model and embeddings varies across meta-learning environments.
Next follows a position-wise feed-forward neural network, for which we use Exponential Linear-Unit (ELU) as the activation, as implemented in Haiku’s vanilla Transformer encoder (Hennigan et al., 2020).
Afterward, we apply the self-attention mechanism again, following a position-wise feed-forward neural network, after which we apply the final linear layer.
Variants of the attention mechanism.
We experimented with multiple variations of the attention mechanism, all being permutation equivariant in the k dimension, similar to the work of Lorch et al.
We mention two other architectural designs: in the first design, rather than using the embeddings of the concatenation of input-output points as keys, queries, and values (as performed in the vanilla Transformer encoder), we experimented with using the embeddings of the inputs as the keys and queries and the embeddings of functional outputs.
However, this design yields varying performance: on several low-dimensional tasks, the performance was slightly better, whereas performance on tasks with higher-dimensional inputs was substantially worse.
In another attempt, we experimented with learning different embeddings for inputs and functional outputs, which we then concatenated and used as keys, queries, and values.
The difference in performance in this method was marginal, and we decided against it for simplicity.
Optimizer setup and employed libraries.
Finally, across all experiments, we use gradient clipping of the prior score together with the ADAM optimizer (Kingma & Ba, 2014), with the default values set in Jax (Bradbury et al., 2018) and Haiku (Hennigan et al., 2020).
For the gradient clipping, we use values of 1., 10., or 100., depending on the task and the underlying properties of the data generating process.
For example, for the first experiment in Appendix D, larger clipping values (50 or 100) performed better for the heavier-tailed Student’t- t process, whereas clipping at 10.
resulted in a good performance of the GP task.
We train the score network for 20000 iterations.
For the GP and Student- t process implementations, we use GP-Jax (Pinder & Dodd, 2022), scikit-learn (Pedregosa et al., 2011) and TensorFlow Distributions packages (Dillon et al., 2017).
A. I NTERPOLATING THE DATASETS ACROSS X
Interpolation with Gaussian Processes.
In GP regression, each data point corresponds to a feature-target tuple
z i,j = ( x i,j , y i,j
) ∈ R d × R .
For the i -th dataset, we write D i =
( X i , y i ) , where X i = ( x i, 1 , . . .
, x i,m i ) ⊤ and y i
=
( y i, 1 , . . .
, y i,m i ) ⊤ .
GPs are a Bayesian method in which the prior P ( h ) =
GP ( h | m ( x ) , k
( x, x ′ )) is specified by a positive definite kernel k : X × X
→ R and a mean function m : X → R .
In this section, we assume a zero mean GP and omit writing the dependence on m ( · ) .
As the GP kernel, we use the Mat ´ ern covariance function
[Equation]
with degree ν , lengthscale ℓ , and Γ( · ) representing the gamma function, and K v the modified Bessel function of the second kind.
We use fix the degree of the Mat´ern kernel to ν = 5 / 2 and choose the lengthscale ℓ via 4-fold cross-validation (CV) from 10 log-uniformly spaced points in [0 . 001 , 10] .
We select the lengthscale that maximizes the 4-fold CV log marginal likelihood, averaged across the n tasks.
In Alg. 3, we give the full procedure of selecting the lengthscale and fitting GPs to the meta-training tasks.
Input: n datasets {D i } ni =1 , where D i = { ( x ij , y ij ) } m i
{GP ( f | k ℓ j ) i } pj =1 , { score i,j } pj =1 = CV 4-fold ( D i , P ) //
cross-validation on D i j ∗ = arg max j =1 ,...,p 1 n (cid:80) ni
=1 score i,j // selecting the optimal kernel parameters
Output: n GP posteriors {GP ( f | k ( ℓ,ν ) j ∗ ) } ni =
Interpolation with Bayesian Neural Networks.
For MARS-BNN, in order to interpolate the datasets, we use the Monte-Carlo dropout (MC-dropout) approach by Gal & Ghahramani (2016) which trains neural networks with dropout and also uses dropout at inference/test time to generate random forward passes through the network.
Consider a neural network with L layers.
For each layer l , let us denote as M l the (random) weight matrix of dimensions K l × K l − 1 .
MC-dropout samples these weight matrices by randomly dropping rows:
[Equation]
given some dropout probabilities p l and trained weight matrices W l as variational parameters.
We write ω
= { W l } Ll =1 for the set variational parameters.
The binary variable z l,j = 0 corresponds to neuron j in layer l being dropped.
We can view these random weights M l as samples from some variational distribution with parameters ω .
For every dataset D i we fit the respective MC-dropout BNNs to the dataset D i , resulting in a set of variational parameters ω i .
This yields the posterior distribution ˜ p ( f X i | X , ω i ) over function values f
X i for arbitrary measurement sets X .
Sampling from ˜ p
( f X i | X , ω i ) corresponds to a random (i.e., with randomly dropped neurons) forward pass with inputs X .
Analogously to (3) case, the score matching loss modifies to
[Equation]
In order to approximate the inner expectation in Eqn.
(7), we sample one realisation of ˜ p
( f X i | X , ω i ) per iteration by performing a random forward pass through the BNN.
We fix the dropout probabilities of all layers p 1 = ...
=
p L = p , across all datasets to a value p , which is a tunable hyperparameter.
In the experiments, we use an fully connected neural network with 3 hidden layers of size 32 each.
As the activation, we use the leaky-relu.
To train each network is we minimize the MSE and with the
Adam optimizer with weight decay.. Weights and bias initialization are set to default initializations by Haiku (Hennigan et al., 2020), i.e., truncated normal distiribution for the weights and constant initialization for the biases.
We fix the number of epochs to 100 .
The hyper-parameters are selected from the following:
• learning rate: { 1 e − 2 , 1 e − 3 , 1 e − 4 , 5 e − 3 , 5 e − 4 } • weight decay: { 0 ., 1 e − 2 , 1 e − 3 , 1 e − 4 , 5 e − 3 , 5 e − 4 } • batch size: { 1, 2, 4, 8, 16 } • dropout: { 0.1, 0.25, 0.5, 0.7, 0.8 }
Note that for Physionet dataset, as the datasets are varying and some have less than 16 data points, if the selected batch size is 16 and the dataset contains less than 16 points, we select the full dataset.
A. C HOOSING THE MEASUREMENT DISTRIBUTION ν
The measurement distribution ν should be supported on relevant parts of the domain X from which we may see queries at test time.
In our experiments, we choose ν = U ( ˜ X ) as uniform distribution over the hypercube ˜ X
⊆ X ⊆ R d which conservatively covers the data.
In particular for each dimension k = 1 , ..., d we compute the minimum and maximum value that occurs in the metatraining data, i.e.,
[Equation]
and expand the respective ranges by 20% on each side
[Equation]
We construct the hypercube
[Table]
B M ETA -L EARNING E NVIRONMENTS
In the following, we provide details on the meta-learning environments used in our experiments in
Section 6.
We list the number of tasks and samples in each environment in Table 5.
[Table]
B. S INUSOIDS (S YNTHETIC E NVIRONMENT )
The sinusoid environment corresponds to a simple 1-dimensional regression problem with a sinusoidal structure.
It is used for visualization purposes in Figure 1 and Figure 3a.
Each task of the sinusoid environment corresponds to a parametric function
[Equation]
yielding a sum of affine and a sinusoid function.
Tasks differ in the function parameters ( a, b, c, β ) that are sampled from the task environment T as follows:
[Equation]
Figure 3a displays functions f a,b,c,β with parameters sampled according to (9).
To draw training samples from each task, we uniformly sample x from U ( − 5 , 5) and add Gaussian noise with standard deviation 0 .
1 to the function values f ( x ) :
[Equation] [Equation]
B. S WISS FEL
Free-electron lasers (FELs) accelerate electrons to a very high speed to generate shortly pulsed laser beams with wavelengths in the X-ray spectrum.
The X-ray pulses from the accelerator can map nanometer-scale structures, thus facilitating molecular biology and material science experiments.
The accelerator and the electron beam line of an FEL consist of multiple magnets and other adjustable components, which have several parameters that experts adjust in order to maximize the pulse energy (Kirschner et al., 2019a).
Due to different operational modes, parameter drift, and changing (latent) conditions, the laser’s pulse energy function, in response to its parameters, changes across time.
As a result, optimizing the laser’s parameters is a recurrent task.
Meta-learning setup.
The meta-learning environment represents different parameter optimization runs (i.e., tasks) on SwissFEL, an 800-meter-long free electron laser located in Switzerland (Milne et al., 2017).
The input space is 12-dimensional and corresponds to the laser parameters, whereas the regression target corresponds to the one-dimensional pulse energy.
We refer to Kirschner et al.
(2019b) for details on the individual parameters.
Each optimization run consists of roughlydata points generated with online optimization methods, yielding non-i.i.d.
data, which becomes successively less diverse throughout the optimization.
To avoid issues with highly dependent data points, we take the first 400 data points per run and split them into training and test subsets of size 200.
As we have a total of 9 runs (tasks) available, we use 5 of them for meta-training and the remaining 4 for meta-testing.
B. P HYSIO N ET
In the context of the Physionet competition 2012, Silva et al.
(2012) have published an open-access dataset of patient stays in the intensive care unit (ICU).
The dataset consists of measurements taken during the patient stays, where up to 37 clinical variables are measured over the span of 48 hours, yielding a time series of measurements.
The intended task for the competition was the binary classification of patient mortality.
However, the dataset is also often used for time series prediction methods due to a large number of missing values (around 80 % across all features).
Meta-learning setup.
To set up the meta-learning environment, we treat each patient as a separate task and the different clinical variables as different environments.
Out of the 37 variables, we picked the Glasgow coma scale (GCS) and hematocrit value (HCT) as environments for our study since
[Table]
On the full dataset, MARS performance is less competitive due to the strong auto-correlation of the data which is not taken into account in the BNN likelihood.
On the partial dataset, which has less dependency among the data points, MARS outperforms all other meta-learning methods.
they are among the most frequently measured variables in this dataset.
From the dataset, we remove all patients where less than four measurements of CGS (and HCT, respectively) are available.
From the remaining patients, we used 100 patients for meta-training and 500 patients for meta-validation and meta-testing.
Since the number of available measurements differs across patients, the number of training points m i ranges between 4 and 24.
B. B ERKELEY -S ENSOR
The Berkeley dataset contains data from 46 temperature sensors deployed in different locations at the Intel Research lab in Berkeley (Madden, 2004).
The temperature measurements are taken over four days and sampled at 10-minute intervals.
Each task corresponds to one of the 46 sensors and requires auto-regressive prediction, particularly predicting the subsequent temperature measurement given the last ten measurements.
Meta-learning setup.
The Berkeley environment, as used in Rothfuss et al.
(2021a), uses 36 sensors (tasks) with data for the first two days for meta-training and the last 10 for meta-testing.
The meta-training and meta-testing are separated temporally and spatially since the data is non-i.i.d.
Data are abundant, and the measurements are taken at very close intervals.
Thus, the features and the train/context data points are strongly correlated, violating the i.i.d.
assumption that underlies our factorized Gaussian likelihood and Bayes rule in Section 3, causing the BNN to over-weight the empirical evidence and making over-confident predictions.
To alleviate this problem, we subsample the data.
In particular, we randomly select 10 out of the 36 training tasks, and instead of using all measurements, we randomly sample 30 of them.
This has two effects: First, it makes the data less dependent/correlated and thus more compatible with our Bayesian formulation.
Second, we increase the epistemic uncertainty by using less data, making the calibration metrics more meaningful.
The results reported in Section 6 correspond to the sub-sampled data.
For completeness, we also report the results for the full dataset as in Rothfuss et al.
(2021a)
in Table 6.
Other meta-learning baselines, such as MAML or MLAP, perform better than MARS on the full dataset since they do not explicitly use the Bayes rule with i.i.d.
assumption or weight the empirical evidence less.
Note that MARS performs worse due to the Bayesian inference at meta-test time rather than our meta-learning approach.
Accounting for the strong auto-correlation of the data in the likelihood would most likely resolve the issue.
On the sub-sampled environment, MARS again performs best.
B. A RGUS -C ONTROL
The final environment we use in our experiments is a robot case study.
In particular, it aims at tuning the controller of the Argus linear motion system by Schneeberger Linear Technology.
The goal is to choose the controller parameters so that the position error is minimal.
In our setup, each task is a regression problem where the goal is to predict the total variation (TV) of the robot’s position error signal when controlled by a PID controller in simulation.
The regression features are the three PID controller gain parameter parameters.
Meta-learning setup.
Overall, the environment consists of 24 tasks, of which 20 are used for meta-training and the remaining 4 for meta-testing.
Each task corresponds to a different step size for the robot to move, ranging from 10 µm to 10 mm .
At different scales, the robot behaves differently in response to the controller parameters, resulting in different target functions.
This presents a good environment for transferring similarities across different scales while leaving enough flexibility in the prior to adjust to the target function at a step size.
C E XPERIMENTAL METHODOLOGY
In the following, we describe our experimental methodology used in Section 6.
C. O VERVIEW OF THE META - TRAINING AND META - TESTING PHASES
Evaluating the performance of a meta-learner consists of two phases, meta-training and meta-testing .
The latter phase, meta-testing , can be further split into target training and target testing .
In particular, for MARS the phases consist of the following:
• Meta-training: The meta-training datasets D ni =1 are used to train the score estimator network (see Algorithm 1).
• Target training: Equipped with knowledge about the underlying data-generation process,
i
e., the score network, we perform BNN inference on a new target task with a corresponding context dataset D ∗ .
In particular, we run fSVGD with the score network as a swap-in for the marginal scores of the stochastic process prior (see Algorithm 2).
As a result, we obtain a set of NN particles that approximates the BNN posterior in the function space.
• Target testing: Finally, we evaluate the approximate posterior predictions on a test set D † corresponding to the same target task.
In particular, we compute the residual mean squared error (RMSE) and the calibration error as performance metrics.
We describe the evaluation metrics in more detail in Section C.2.
The target training and testing are performed independently with the meta-learned score network for each test task.
The metrics are reported as averages over the test tasks.
The entire meta-training and meta-testing procedure are repeated for five random seeds that influence the score network initialization, the sampling-based estimates in Algorithm 1 as well as the initialization of the BNN particles for target training.
The reported averages and standard deviations are based on the results obtained for different seeds.
C. E VALUATION METRICS
During target-testing , we evaluate the posterior predictions on a held-out test dataset D † .
Among the methods employed in Section 6, MARS, PACOH-NN, NPs, MLAP, Vanilla BNN and Vanilla
GP yield probabilistic predictions ˆ p ( y † | x † , D ∗ ) for the test points x † ∈ D † .
For instance, in the case of MARS, PACOH-NN, and Vanilla BNN where the posterior is approximated by a set of
NN particles { θ 1 , ..., θ L } and we use a Gaussian likelihood, the predictive distribution is an equally weighted mixture of Gaussians:
[Equation]
[Equation]
The respective mean prediction corresponds to the expectation of ˆ p , that is ˆ y
=
ˆ E ( y ∗ | x ∗ , D ∗ ) .
In the case of MAML, only the mean prediction is available.
[Table]
2020) with either an IMQ or RBF kernel.
Evaluating prediction accuracy (RMSE)
Based on the mean predictions, we compute the rootmean-squared error (RMSE)
[Equation]
which quantifies how accurate the mean predictions are.
Evaluating uncertainty calibration (Calibration error)
In addition to the prediction accuracy, we also assess the quality of the uncertainty estimates.
For this purpose, we use the concept of calibration, i.e., a probabilistic predictor is calibrated if the predicted probabilities are consistent with the observed frequencies on unseen test data.
We use a regression calibration error similar to
Kuleshov et al.
(2018) in order to quantify how much the predicted probabilities deviate from the empirical frequencies.
Let us denote the predictor’s cumulative density function (CDF) as
[Equation]
=
[Equation]
y −∞ ˆ p (˜ y | x , D ∗ ) d ˜ y , where ˆ p (˜ y | x , D ∗ ) d ˜ y is the predictive posterior distribution.
For confidence levels 0 ≤ q h < . . .
<
q H ≤ 1 , we compute the corresponding empirical frequency
[Equation]
based on some test dataset D † .
If the predictions are calibrated, we expect that ˆ q h → q h as m → ∞ .
Following Kuleshov et al.
(2018), we define the calibration error metric as a function of the residuals ˆ q h − q h :
[Equation]
(cid:88)
[Equation]
Note that we report the average of absolute residuals | ˆ q h − q h | , rather than reporting the average of squared residuals | ˆ q h − q h | 2 , as done by Kuleshov et al.
This is done to preserve the units and keep the calibration error easier to interpret.
In our experiments, we compute the empirical frequency using M = 20 equally spaced confidence levels between 0 and 1.
C. H YPER -P ARAMETER S ELECTION
For each of the meta-environments and algorithms, we ran a separate hyper-parameter search to select the hyper-parameters.
In particular, we use 30 randomly sampled hyperparameter configurations across five randomly selected seeds and select the best-performing one in terms of either RMSE or
[Table]
) with either an IMQ or RBF kernel.
calibration error.
For the reported results, we provide the chosen hyperparameters and detailed evaluation results in https://tinyurl.com/376wp8xe .
The code scripts for reproducing the experimental results are provided in our repository 2 .
D F URTHER E XPERIMENT R ESULTS
We provide further empirical evidence for the proposed method.
In particular, we showcase that:
Our parametric score matching approach performs favorably to many nonparametric score estimators.
Regularization via spectral normalization does not hinder the flexibility of the network.
Sampling from the GP posteriors when training the score network successfully incorporates epistemic uncertainty in areas of the domain where meta-training data is scarce.
D. P ARAMETRIC VS N ONPARAMETRIC S CORE E STIMATION M ETHODS
We start by comparing the performance of MARS to that of several nonparametric score estimators, described in Zhou et al.
Nonparametric score estimators.
There are several theoretically motivated nonparametric score estimation methods with well-understood properties and straightforward implementations.
Their simplicity and flexibility make them a popular choice (Coyle et al., 2020;
Wang et al., 2019;
Deledalle et al., 2014).
KEF (Canu & Smola, 2006) performs regularized score matching inside a kernel exponential family; this can further be efficiently approximated by the Nystr¨om method (NKEF)
(Sutherland et al., 2018).
Liu et al.
Chwialkowski et al.
(2016) propose a Stein estimator, and Shi et al.
(2018) propose Spectral Stein Gradient Estimator (SSGE) by expanding the score function in terms of the spectral eigenbasis.
Together with iterative methods (the ν -method and Landweber iteration (Engl et al., 1996)), these approaches can be naturally unified through a regularized, nonparametric regression framework (Zhou et al., 2020), in which the conjugate-gradient version of the KEF estimator is also proposed (KEF-CG).
In our evaluations, we reports results for curl-free sccore estimators with both an inverse multiquatric (IMQ) or RBF kernel.
MARS vs nonparametric estimators.
To compare the performance of MARS to nonparametric methods, we consider score estimation of marginal distributions of Gaussian Processes (GP) (Rasmussen, 2003) and Student’s- t processes (TP) (Shah et al., 2014).
In both cases, as the stochastic
10 x 1.
x l n p ( x )
MARS True    20 x 1.
x l n p ( x ) process mean function, we use the sinusoidal mean function 2 x + 5 sin (2 x ) , and for the covariance kernel, we choose the radial basis function kernel, with the lengthscale parameter set to ℓ = 1 .
We use Tensorflow Probability’s (Dillon et al., 2017) implementation of both GP and TP and set all other parameters to their default values.
We perform four experiments, the first two on estimating the marginal scores of a GP and the last two on the TP.
In all experiments, we sample one measurement set X containing either two ( X = { x 1 , x 2 } ) or three points ( X = { x 1 , x 2 , x 3 } ) of dimension one, i.e., x i ∈ R , ∀ i .
For all experiments, the x i follow a uniform distribution: for the GP experiment, x i ∼ U ([ − 5 , 5]) , whereas for the TP experiment, x i ∼ U ([ − 1 , 1]) , ∀ i .
In the experiments with measurement sets consisting of two points, the score network and the nonparametric estimators are trained on 50 samples from the corresponding two -dimensional marginal distributions.
For the measurement sets of size 3, we train the score network and nonparametric estimators on 200 samples from the corresponding three -dimensional marginal distribution.
When evaluating the performance of MARS, we take the average performance over five different seeds after training the network for 2000 iterations in the two -dimensional marginal case and 5000 iterations in the three -dimensional case.
We measure the quality of the score estimates via the RMSE and cosine similarity between the estimates and true scores.
The corresponding evaluation results are reported in Tables 7-8.
The RMSE measures deviations of the estimated scores in both direction and magnitude of the gradients.
For the SVGD particle estimation, it is more important that the gradients in the vector field point in the correct direction than having the correct magnitude.
Thus, we also report the cosine similarity, which only quantifies how well the directions in the estimated vector field match the directions in the true vector field while neglecting errors in the score magnitude.
As we can see in Table 7 and 8, MARS consistently outperforms the nonparametric score estimators.
Among the nonparametric methods, SSGE performs best on average.
Hence, we use SSGE in our ablation study in Section 6.3.
D. F LEXIBILITY OF THE REGULARIZED NETWORK
Recall from section 5.3 that, in order to prevent overfitting, we perform spectral normalization of the weights of the linear layers by re-parameterizing the weights by ˜ W := W / ||
W || .
To investigate whether this hinders the flexibility of the network’s outputs, we visually examine the network’s predictions on the following two tasks, in which we estimate the scores of one-dimensional Laplace and Student’s- t distribution.
In order to estimate a score ∇ x log p ( x ) for some unknown distribution p ( x ) (rather than marginal distribution scores ∇ f X ln p ( f X ) for some unknown stochastic process p ( f ) ), we use 45 samples
X = { x 1 , ..., x 45 } from: (i) one-dimensional Student’s- t distribution with ν = 5 degrees of freedom, location parameter µ = 0 , and scale parameter σ = 1 , and (ii) one-dimensional standard
Laplace distribution, i.e., with location parameter µ = 0 and scale parameter σ = 1 .
To perform distributional score estimation, we need to make a slight modification to the original score network architecture; rather than concatenating input-output pairs (i.e., concatenating x i and f ( x i ) , as shown in Fig. 4), we use only the distributional samples as inputs.
To be more precise, at every iteration, we randomly select (without replacement) k
= 8 inputs { x ′ 1 , ..., x ′ k } ⊂ X , which we input to the score network.
We train the network using the regularized score matching loss, i.e., perform spectral normalization of the layers.
The network is trained for 2000 iterations, using a learning rate of 0 .
For this experiment, we set the attention embedding dimension to 32 , key size to 16 , and use 8 attention heads.
The remaining hyperparameters are set according to Sec.
A.3.
The respective predictions are displayed in Figure 5.
We observe that the score network approximates the score functions well and that the regularization does not hinder flexibility too much.
D. I NCORPORATING UNCERTAINTY THROUGH GP INTERPOLATION
In addition to the ablation study performed in Section 6.3, we visually investigate the implications of sampling functions from the GP posteriors during score matching.
In particular, we do so using the sinusoid environment, detailed in Appendix B.1.
Experiment setup.
In this experiment, we sample 10 functions from the environment and evaluate each function at five randomly selected inputs in the [ − 5 , − 2] ∪ [2 , 5] range.
We investigate whether the proposed GP interpolation method promotes uncertainty in the [ − 2 , 2] part of the domain.
In order to do so, let us recall Algorithm 1, in which we learn the trained score network s ϕ .
The algorithm first fits a GP to each of the n datasets D 1 , ..., D n .
Then, at every step, the algorithm samples a measurement set X iid ∼ ν , and then, for every dataset D i (i.e., for every collection of input/output pairs X D i , y D i ), the algorithm samples function values from the GP posterior marginal:
˜ f
X i ∼ p ( f X i | X , X D
i , y D i ) .
We compare this approach to training the score network s ϕ on the mean predictions of the posterior marginal p ( f X i | X , X D i , y D i ) , i.e., where ˜ f
X i = µ p ( f X i | X , X D
i , y D i ) .
In order to distinguish between the two approaches, we denote the score network trained on the posterior marginal mean predictions with s µϕ , and use ˜ f
X µ i := µ p ( f X i | X , X D
i , y D i ) as a shorthand notation for the mean predictions of the posterior marginals.
Observe that, for measurement sets X close to X D i , the samples ˜
f
X i and ˜ f
X µ i will be very similar, whereas when X is far from X D i , the variability of the samples ˜ f
X i will be much larger.
In the following, we showcase that this variability successfully incorporates uncertainty about the areas of the domain X with little or no data.
In order to compare the two approaches, we visually compare the overall framework when the fSVGD network is trained with s µϕ (the score network trained on ˜
f X µ i , the GP mean predictions), and when it is trained with s ϕ (the score network trained on ˜ f
X i , the samples from the GP posterior marginal), as performed in the original MARS algorithm.
Empirical findings.
We plot the corresponding fSVGD-BNN (fBNN) predictions, trained using the two score networks s µϕ and s ϕ .
Both fBNNs are fitted to four points, where inputs x 1 , ..., x 4 are sampled uniformly from [ − 5 , 5] (i.e., the whole input domain), and their functional evaluations are obtained according to the sinusoid environment.
The results are visualized in Figure 6.
The first two plots on the left in Figure 6 correspond to posterior predictions of two randomly selected GPs fitted to the meta-tasks, where no task contains information in the [ − 2 , 2] input range.
The middle plot corresponds to the fBNN network predictions, where the network was trained using the fSVGD algorithm and the score network s µϕ .
The last plot corresponds to the fBNN network predictions, where fBNN is trained using the fSVGD algorithm and the score network s ϕ .
We observe a clear difference between the two approaches:
MARS (trained using s ϕ ) successfully incorporates the epistemic uncertainty in the [ − 2 , 2] part of the input domain into the fBNN posterior, yielding less confident predictions in the area where no data was available during meta-training.
In contrast, when we use GP posterior means instead of samples when fitting the score network, the resulting BNN predictions entirely ignore the epistemic uncertainty that arises due to the fact that we don’t know the function values in [ − 2 , 2] .
This may lead to overconfident posterior predictions.
D. C OMPARISON OF I NTERPOLATION M ETHODS
So far, we have experimented with using GPs (MARS-GP) and MC-dropout BNNs as interpolators (MARS-BNN) and sample random random function values from the corresponding posteriors when training our score network.
In the ablation study (performed in Sec. 6.3), we also experimented with
[Table]
Both the deterministic NN and the NN ensemble interpolators are trained similarly to the MCdropout BNNs, as described in Section A.4.
For the NN ensembles, we choose the number of of ensemble members as a hyper-parameter from { 5 , 10 , 20 } .
In case of the deterministic NN, we do not perform any posterior sampling in (3) and just use the NN’s predictions as function values for the score estimation.
In case of the ensemble, we randomly sample the function values corresponding to one ore multiple ensemble member.
The number of function value samples from the ensemble is chosen as hyper-parameter and selected from { 1 , 2 , 3 , 4 } .
The results for employing MARS with different interpolators are reported in Table 9-10.
Not surprisingly, the not accounting for epistemic uncertainty (i.e., GP mean and deterministic NN) performs much worse in the majority of environments, in particular, w.r.t.
the calibration error.
Among the neural network based interpolators, the
MC-dropout BNNs consistently perform the best.
This is why suggest their usage and evaluate them in the main part of the paper.
[Table]