Accurate Bayesian Meta-Learning by Accurate Task Posterior Inference
Authors: Michael Volpp
Philipp Dahlinger
Philipp Becker
Christian Daniel
Gerhard Neumann
Venue: ICLR 2023 poster (2023)
Keywords: Bayesian Meta-Learning
Neural Processes
Variational Inference

Abstract: Bayesian meta-learning (BML) enables fitting expressive generative models to small datasets by incorporating inductive priors learned from a set of related tasks. The Neural Process (NP) is a prominent deep neural network-based BML architecture, which has shown remarkable results in recent years. In its standard formulation, the NP encodes epistemic uncertainty in an amortized, factorized, Gaussian variational (VI) approximation to the BML task posterior (TP), using reparametrized gradients. Prior work studies a range of architectural modifications to boost performance, such as attentive computation paths or improved context aggregation schemes, while the influence of the VI scheme remains under-explored. We aim to bridge this gap by introducing GMM-NP, a novel BML model, which builds on recent work that enables highly accurate, full-covariance Gaussian mixture (GMM) TP approximations by combining VI with natural gradients and trust regions. We show that GMM-NP yields tighter evidence lower bounds, which increases the efficiency of marginal likelihood optimization, leading to improved epistemic uncertainty estimation and accuracy. GMM-NP does not require complex architectural modifications, resulting in a powerful, yet conceptually simple BML model, which outperforms the state of the art on a range of challenging experiments, highlighting its applicability to settings where data is scarce.

Published as a conference paper at ICLR 2023
ACCURATE BAYESIAN META-LEARNING BY ACCURATE TASK POSTERIOR INFERENCE
Michael Volpp' 1.2
Philipp Dahlinger'
Philipp Becker
Christian Daniel?
Gerhard Neumann'
Karlsruhe Institute of Technology, Karlsruhe, Germany Bosch Center for Artificial Intelligence, Renningen, Germany
ABSTRACT
Bayesian meta-learning (BML) enables fitting expressive generative models to small datasets by incorporating inductive priors learned from set of related tasks. The Neural Process (NP) is a prominent deep neural network-based BML ar- chitecture, which has shown remarkable results in recent years. In its standard formulation, the NP encodes epistemic uncertainty in an amortized, factorized, Gaussian variational (VI) approximation
the BML task posterior (TP), using reparametrized gradients. Prior work studies a range of architectural modifications to boost performance, such as attentive computation paths or improved context ag- gregation schemes, while the influence of the VI scheme remains under-explored. We aim to bridge this gap by introducing GMM-NP, a novel BML model, which builds on recent work that enables highly accurate, full-covariance Gaussian mix- ture (GMM) TP approximations by combining VI with natural gradients and trust regions. We show that GMM-NP yields tighter evidence lower bounds, which increases the efficiency of marginal likelihood optimization, leading to improved epistemic uncertainty estimation and accuracy. GMM-NP does not require plex architectural modifications, resulting in a powerful, yet conceptually simple BML model, which outperforms the state of the on range of challenging experiments, highlighting its applicability to settings where data is scarce.
INTRODUCTION
Driven by algorithmic advances in the field of deep learning (DL) and the availability of increasingly powerful GPU-assisted hardware, the field of machine learning achieved plethora of impressive
sults in recent years (Parmar et al. 20188) Radford et al. 2019 Mnih et al, (015)). These were enabled to large extent by the availability of huge datasets, which enables training expressive deep ral network (DNN) models. In practice, e.g., in industrial settings, such datasets are unfortunately rarely available, rendering standard DI. approaches futile. Nevertheless, it is often the case that simi- lar tasks arise repeatedly, such that the number of context examples on a novel target task is typically relatively small, but the joint meta-dataset of examples from all tasks accumulated over time can be massive, s.t. powerful inductive biases can be extracted using meta-learning (Hospedales et al.) 2022). While these inductive biases allow restricting predictions to only those compatible with the meta-data, there typically remains epistemic uncertainty due to task ambiguity, as the context data often not informative enough identify the target task exactly. Bayesian meta-learning (BML) aims at accurate quantification of this uncertainty, which is crucial for applications like active learning, Bayesian optimization (Shahriari et al.) 2016), model-based reinforcement learning (Chual ct al. 2018), robotics (Deisenroth et al. 2011) and in safety-critical scenarios.
"Correspondence Mi chael. vo appdde. bosch.com

---

Published as a conference paper at ICLR 2023
Latent space
Predictions
5.0
0
2.5
-1
-2
0.0
-3
-2.5
-4
-
3 -5.0
-1 0
-
0
4
Figure 1: Visualization of our GMM-NP model for d. 2 dimensional latent space, trained on a meta-dataset of sinusoidal functions with varying amplitudes and phases, after having observed a single context example (red cross, right panel) from unseen task (black dots, right panel). Left panel: unnormalized task posterior (TP) distribution (contours) and GMM TP approximation with K = 3 components (ellipses, mixture weights in %). Right panel: corresponding function samples from our model (blue lines), A single context example leaves much task ambiguity, reflected in a highly correlated, multi-modal TP. Our GMM approximation correctly captures this: predictions in accordance with (i) the observed data (all samples pass close to the red context example). and with (ii) the learned inductive biases (all samples are sinusoidal), cf. also Fig. 12 in App. A.5.5
A prominent BML approach is the Neural Process (NP) (Garnelo et al. 2018b) which employs DNN-based conditional latent variable (CLV) model, in which the Bayesian belief about the tar- get task is encoded in factorized Gaussian task posterior (TP) approximation, and inference is amortized over tasks using set encoders (Zaheer et al. (017). This architecture can be optimized effi- ciently using variational inference (VI) with standard, reparametrized gradients (Kingma & Welling, 2014). A range of modifications, such as adding deterministic, attentive, computation paths (Kim) al. 2019), Bayesian encoders (Volpp al.) (02T), have been proposed in recent years to improve predictive performance. Interestingly, the VI scheme with an amortized, factorized Gaus- sian TP. optimized using standard gradients, remains largely unaltered. Yet. it is well known that (i) the factorized Gaussian assumption rarely holds in Bayesian learning (MacKay [003 Wilson Izmailov] 2020), (ii) amortized inference can yield suboptimal posterior approximations (Cre- mer et al., 2018), and (iii) natural gradients are superior to standard gradients for VI in terms of optimization efficiency and robustness (Khan & Nielsen 2018).
Building on these insights and on recent advances in VI (Lin et al. 2020 Arenz et al. 2022 we
propose GMM-NP, a novel NP-based BML algorithm that employs (i) full-covariance Gaussian mixture (GMM) TP approximation, optimized in (ïi) non-amortized fashion, using (iii) robust and efficient trust region natural gradient (TRNG)-V VI. We demonstrate through extensive empirical evaluations and ablations that our approach yields tighter evidence lower bounds, more efficient model optimization, and, thus, markedly improved predictive performance, outperforming the state- of-the-art both in terms of epistemic uncertainty quantification and accuracy. Notably, GMM-NP does not require complex architectural modifications, which shows that accurate TP inference is crucial for accurate BML, an insight we believe will be valuable for future research.
2 RELATED WORK
Multi-task learning aims leverage inductive biases learned on meta-dataset of similar tasks for improved data efficiency on unseen target tasks of similar structure. Notable variants include transfer-learning (Zhuang al., 2020), that refines and combines pre-trained models (Golovin al.) 2017 Krizhevsky et al. 2012)) and meta-learning (Schmidhuber 1987, Thrun & Pratt] 1998 Vilalta & Drissi [0055 Hospeeales
al. 2022), which makes the multi task setting explicit in the model design by formulating fast adapation mechanisms in order to learn how to solve tasks with little context data ("few-shot learning").. ". A plethora of architectures were studied the literature, includ- ing learner networks that adapt model parameters (Bengio et al.) [991] [Schmidhuber] 1992 Ravi & Larochelle 2017), memory- muemory-augmented DNNs (Santoro et al. 2016). early instances of Bayesian meta-models (Edwards & Storkey 2017 Hewitt et al. 2018), and algorithms that that make use of learned measures of task similarity (Koch et al. 2015 Vinyals et al. 2016 Snell et al. 2017).
Arguably the most prominent meta-learning approaches (MAML) and the Neural Process (NP) model families, due to their generality and flexibility. While the original MAMI. (Finn et al. 2017) and Conditional NP (Garnelo et al. 2018a) formulations do
the Model-agnostic Meta-learning
2

---

Published as a conference paper at ICLR 2023
not explicitly model the epistemic uncertainty arising naturally in few-shot settings due to task am- biguity, both model families were extended to fully Bayesian veta-learning (BML) algorithms that
explicitly infer the TP based on CLV formulation (Heskes) 2000 Bakker & Heskes 2003). Im-
portant representatives are Probabilistic MAMI. (Grant et al. 2018 Finn et al. 2018) and Bayesian MAMI. (Kim al., 2018), as well several NP-based BMI approaches that inspire our work. These include the Standard NP (Garnelo et al. 20188)) which was extended by attentive computa- tion paths to avoid underfitting (Kim al. 2019), or by Bayesian set encoders (Zaheer al. 2017 Wagstaff et al. [019 Volpp et al. (2220) for improved handling of task ambiguity. as well as by hier- archical (Wang& & Van Hoof 2020), bootstrapped (Lee et al., 2020). or graph-based (Louizos al.) 2019) latent distributions. While the original NP formulation employs an amortized, reparametrized,
stochastic gradient VI objective (Kingma & Welling, 2014; Rezende et al. 2014), Monte-Carlo
(MC)-based objective functions were also studied (Gordon et al. 2019 Volpp al. 2021..
From a more general perspective, emeerge as a central tool in many areas of probabilisticc machine learning, which require tractable approximations of intractable probability distributions, typically arising as the posterior in Bayesian models (Gelman et al. 2004 Koller & Friedman 2009) Neal
1996 Wilson & Izmailov, 2020). While early approaches (Attias) 2000) allow analytic updates, more complex algorithms employ stochastic gradients w.r.t. the variational parameters (Ranganath al.) [2014] Kingma & Welling 2014 Blundell al. (015). Such approaches are straightforward
to implement and computationalyy efficient for factorized Gaussian variational distributions, but ig- nore the information geometry of the loss landscape. leading to suboptimal convergence rates (Khan & Nielsen, 2018). Natural gradient (NG)-VI (Amari 1998) alleviates this problem and recent work (Hoffman et al.] 2013, Winn & Bishop 2005 Khan & Nielsen 2018 Khan et al. 2018) success- fully applies this idea at scale to complex models, requiring only first-order gradient information (Lin et al., 2019). Further extensions enable NG-VI for structured variational distributions such as mixture models by decomposing the NG update into individual updates per mixture component (Arenz et al. 20188; [in et al. (020) which, in combination with trust region (TR) step size control Abdolmaleki et al.) |015]| Arenz et al. 2022), yields robust and efficient VI algorithms for versatile and highly expressive variational distributions such as Gaussian mixture models (GMMs).
3 PRELIMINARIES
We now briefly recap the TRNG-VI algorithm (Lin et al. 2020 Arenz et al., 2022) as well as the NP model (Garnelo al. 2018b), which form the central building blocks of our GMM-NP model.
3.1 TRUST REGION NATURAL GRADIENT VI WITH GAUSSIAN MIXTURE MODELS
Variational Inference. We consider probability distribution (z) over random variable Rd, which intractable in the sense that we know it only up to some normalization constant Z, i.c., p (z) p(z) / with Z tractable distribution qs (z), parametrized by 6. Variational inference (VI) frames this task as the minimization w.r.t. & of the reverse Kullback-Leiber (KL) divergence (Kullback & Leibler 1951)
p2) de and tractable p(z). We seek to approximate p (z) by
5 (z)
KL [gollp] = (x)
log
+ log Z (L log Z,
Job (z)
where we introduced evidence lower bound (ELBO) L (Ø). As Z independent of Φ, minimizing the KL divergence is equivalent to maximizing the ELBO.
Natural Gradients. standard approach employs stochastic, reparametrized gradients w.r.t. Φ (Kingma & Welling 2014) for optimization. While this is computationally efficient, it ignores the geometry of the statistical manifold defined by the set of probability distributions Gø. which can lead to suboptimal convergence rates (Khan & Niclsen 2018). A more efficient solution is to perform updates in the natural gradient (NG) direction, i.c., the direction of steepest ascent w.r.t. the Fisher information metric (Amari 1998). State-of-the-art approaches estimate the NG from first-order gradients of p (z) by virtue of Stein's lemma (Lin et al. 2019), 9) yielding efficient NG-VI algorithms that scale complex problems (Khan et al. 2018 Lin al. 2020 Arenz et al. 2022).
Trust Regions. Selecting appropriate step sizes for updates in &f can be intricate, which is why Abdolmaleki et al. (001)) propose (zero-order) algorithm that incorporates a trust region constraint
3

---

Published as a conference paper at ICLR 2023
of the form KL Iqol I E, which restricts the updates in distribution space and can be enforced with manageable computational overhead (a scalar, convex optimization problem in the Lagrangean parameter for the constraint). As shown by Arenz al. (2022), such trust regions easily be combined with gradient information, and allow more aggressive updates in comparison to setting the step size directly, while still ensuring robust convergence.
VI with Gaussian Mixture Models. The quality of the approximation depends on the expressive- ness of the distribution family 9Φ. In settings where p corresponds the Bayesian posterior of complex latent variable models (MacKay 2003] Wilson & Izmailov 2020), simple Gaussian ap- proximations do not yield satisfactory results, as p typically is multimodal. In such cases, Gaussian mixture models (GMMs) are an appealing choice. as they provide cheap sampling, evaluation, and marginalization while allowing expressive approximations (Arenz et al.] 2018). However, a naive application of VI is futile because gradients are coupled between GMM components, leading to computationally intractable updates. Fortunately, (Arenz et al. (2018) and et al. (2020) show that updating the components and weights individually is possible, while preventing collapse of the approximation onto a single posterior mode. This leads to two state-of-the-art algorithms for NG-VI with variational GMMs, that differ most notably in the way the step sizes for the updates are trolled: iBayes-GMa GMM (Lin et al. 22000), which directly sets step sizes for the updates, and TRNG- VI Arenz al. 2022). which employs trust regions for more efficient and robust convergence.
3.2
BAYESIAN
WITH NEURAL PROCESSES
The Multi-Task Latent Variable Model. We aim to fit genera- tive model to a meta-dataset De = { Ve.1.N } with inputs It,n E Rd and corresponding evalu- ations ye,n E RR of unknown functions ft. i.c., Yt.n = fe (zt,n) ) + En, where e, denotes (possibly heteroskedastic) noise. Tasks are assumed to share statistical structure as formalized in the multi-task CLV model shown to the right, defining the joint probability distribution
DI.L, consisting of regression tasks
2t.m
3/t.
N
pe (YI:L,I:N. 21.L (IL.I.)
= II f,7
pe (yenk (z)).
where žạ E Rdd denote latent task descriptors and O denotes task-global parameters that capture shared statistical structure. Having observed context data DE {œī I,M, , y.M.M... C D. from target task D., predictions are provided in terms of the marginal predictive distribution
pe (y,,1:N, IIN. D.)
/II pe (Brrn/ z.) pe (z. dz.,
II
with the task posterior (TP) distribution po (z. IIm (o S m: 2. )p (z.) / po (D!).
The Neural Process. In its standard formulation, the Neural Process (NP) (Grnelo et al. 2018b) defines factorized Gaussian likelihood pe (y/ ty| 2) = N' (y/dec (x, 2). diag (o,)), where the decoder deck is a DNN with weights 0, and observation noise variance o.- As the TP is in- tractable for this likelihood choice, NP computes a factorized Gaussian approximation qo (z. (D)) = |Zzzlenc" (Dt). ; diag(enc) (D!)) with deep set encoders (Zaheer et al. 2017 Wagstaff et al., 2019 enc encz, parametrized by Φ. The parameters +b = (0, Φ) are optimized jointly on the meta-data by stochastic gradient ascent on the ELBO Ete Le (T) w.r.t. the approximate log marginal predictive likelihood defined by
log qa (yt,I:1 IE,,IN Df) = log /II po (e,, zx) gs (zelDf) dzt
(4)
> (z Σ log pe ke ze) + log 4s (ze (ze (Dj)
do (zl) (De) Ce (Φ),
(5)
where Dj C Dr, and stochastic gradients w.r.t. & are estimated using the reparametrization trick Kingma& & Welling, 2014). Note that NP amortizes inference (the variational parameters o shared across tasks) and that it re-uses qs (æzl-) to compute the variational distribution qs (z)),,), taking advantage of its deep set encoder, which allows condition it on datasets of arbitrary size.
4

---

Published as a conference paper at ICLR 2023
4 BAYESIAN META-LEARNING WITH GMM TASK POSTERIORS
Motivation. Our work is motivated by the observation that the current fatee-o-the-wee art approach for training NP-based BML models is suboptimal. Concretely, we identify three interrelated issues with the optimization objective Eq.
(II) Expressivity of the Variational Distribution, le is a (i) factorized, (ii) unimodal Gaussian distribution, (ii) amortized over tasks. In effect, this parametrization only allows crude
approximations of the TP distribution (MacKay 2003 Cremer et al. 2018).
(12) Optimization of the Variational Parameters. (i) Naive gradients of Eq. ( ignoring the information geometry of qΦ, with (ii) direct step size control are employed for optimization, yielding brittle convergence suboptimal rates (Khan & Nielsen 2018 Arenz et al. 2022)
(I3) Optimization of the Model Parameters. Due to the suboptimal VI scheme (I1,12), the TP approximation is poor, resulting in a loose ELBO Eq. (5). In effect, optimization w.r.t. the model parameters 0 is inefficient, cf. App. A.1.4 for detailed discussion.
Armed with these insights, develop novel BML model algorithm that is close in spirit to the NP but solves (I1-I3) through TRNG-VI with GMM TP approximations.
Model. Our algorithm builds on the standard multi-task CLV architecture Eq. (2) and retains the likelihood parametrization using a decoder DNN, dec (x.,), z), as this allows for expressive BMI. models. Under this parametrization, the log marginal likelihood for single task reads
log pe ( log / II pe (ye,l ze) p (ze) dæe log (0).
(6)
where Ze (0) is the normalization constant of the TP po Dr) Pe (zz) / Ze (0) with p(zz) = II, Po (Ve, De,,, ze) (zz). In contrast to Eq. (4). we do not condition the left hand side on context set Dj, which yields a tractable integrand p(zz) that does not require further approximation. To tackle (II), we approximate po (z₂ (Dz) by an expressive variational GMM of the form
1Φ. (zx) = Σ WE. (z)(k)
Σ WEAN (zelpeR, Ee,e),
Σ wee =1,
where we train individual GMMs with parameters Φ: = (we,e P.A: E. E fl...., K} for each task f, to not impair approximation quality by introducing inaccuracies through amortization.
Update Equations for the Variational Parameters. ensure efficient and robust optimization of Φe (12), we employ TRNG-VI as proposed by Arenz et al. (2022), with the update equations
1
Eg neww
Re
(8a)
E,k.old
11
1/
keww ne.
"7
He.. + + (ren - Res.eex.oo)
E
(8b)
Week.new X exp Pe.ks
(8c)
where Ry Tr,u, and Pr, are defined as expectations that can be approximated from per-component samples using MC and require most first-order gradients of ő (ze), which are readily available us- ing standard automatic differentiation software (Abadi et al. 2015) Paszke et al., (019). Due to space constraints, we move details to App. A.I.1 The optimal value for the trust region parame- ter 2 0 is defined by scalar convex optimization problem that can be solved efficiently by bracketing search, which also ensures positive definiteness of the new covariance matrix Σek.
Updates for the Model Parameters. To optimize the model parameters 0, we decompose the log marginal likelihood log Ze (0) according to Eq. (I) as
log Po (ye. De,n, zt) + log p (2r)
log Ze (0) = (xge
Σ n
+ KT. (KT. (.) |po (-[De)),
qsh (ze De)
5

---

Published as a conference paper at ICLR 2023
GMM-NP (ours)
MA-NP
BA-NP
BNP
ANP
BANP
LMLHD
MSE
LMLHD
MSE
10
-25
10
10a
-50
10
10a
-50
0
5
10
0
10
15
0
5
10
0
10
15
context size
context
context size
context
(a) Sinusoidal functions.
(b) Mix of affine and sinusoidal functions.
10
2.5
2.5
0.0
0.0
-5
-2.5
-2.5
-10
5
0
5
S
5
0
5
5
(c) GMM-NP (ours).
(d) BA-NP.
(c) ANP.
BANP.
Figure 2: Panels (a), (b): LMLHD and MSE on two synthetic function classes. Panels (c) - (f): func- tion samples of models trained on the affine-sinusoidal class (b), given one context example (red) from sinusoidal instance (black). GMM-NP outperforms the baselines, as it accurately quantifies epistemic uncertainty through diverse samples. BA-NP also shows variability in its samples, but does not achieve competitive performance due to its inaccurate TP approximation. ANP and BANP produce essentially deterministic predictions that fail to give reasonable estimates of the predictive distribution. Cf. also Figs. 11 in App. A.5.4
where the first term on the right hand side is the ELBO w.r.t. log Z, (0), which we denote by L (0). We expect L (0) to be comparably tight, as our inference scheme allows accurate GMM TP approx- imations Ip.- s.t., the KI. term will be small. Consequently, maximization of Zz (0) w.r.t. A can be performed efficiently by maximization of L (0) (I3), cf. also App. A.1.4 App.A...4 As is standard, we use the Adam optimizer (Kingma & Ba 2015) to perform updates in 0, with MC gradient estimates from samples žτ,a ~ qs, (zz): (o x a.n. Ve log Po (ye,nlzeen,, zt,,).
Meta- Training. The goal of any BMI. algorithm is to compute accurate predictions with well- calibrated uncertainty estimates according to Eq. 4 based on samples from the approximate TP qs. (z.(DD) x pe (z(D.), conditioned on a context set DE from a target task. During a meta- training stage on meta-data DI.L, we aim to encode inductive biases in the model parameters O, s.t. small (few-shot) context sets D suffice for accurate predictions. To find versatile solutions that work for variable context set sizes, it is necessary emulate this during meta-training by evaluating gradients for 0 on samples zt,s from approximate TPs informed by a range of context set sizes. Standard NPs achieve this by sampling minibatch of auxiliary subtasks, with a random number of datapoints, from DI.L for each step in the parameters & (cf. Sec. A...2).. Our algorithm uses similar approach: starting from a fixed set of randomly initialized variational GMMs Φe, and a randomly initialized model 0 we iterate through the meta-data in minibatches of auxiliary subtasks, and perform one update step in Φe for all subtasks in the minibatch, according to Eqs. (8), followed by one gradient step in O. Thus, variational and model parameters evolve jointly in similar fashion as for standard NP, resulting in a meta-training stage with comparable computational complexity, cf. App. A.5.6 As this approach retains a fixed set of variational GMMs over the whole course of meta-training (one for each auxiliary subtask), we accordingly sample a fixed set of auxiliary subtasks at the beginning of meta- veca-training. We summarize our algorithm in A.I Alg. D
A.I
Predictions. As our architecture does not amortize inference over tasks and, thus, does not learn set encoder architecture, the variational GMMs learned during meta-training are not required for predictions on test tasks and can be discarded. To make predictions, we fix the model parameters 0 and fit new variational GMM go. DE by iterating Eqs. (8) until convergence. Afterwards, we can cheaply generate arbitrarily many samples Z,,s ~ qs. (z.), and generate corresponding function samples, evaluated at arbitrary input locations æ., by a single forward pass through the decoder DNN to approximate the predictive distribution according to Eq. (4) cf. App. A. Alg. 2
5 EMPIRICAL EVALUATION
Our empirical evaluation aims to study the effect on the predictive performance of (i) our improved TRNG-VI approach well of (ii) expressive variational GMM TP approximations in NP-based
6

---

Published as a conference paper at ICLR 2023
GMM-NP (ours): TRNG-VI foll,
TRNG-VI full, K-1
TRNG-VI, ding.
SGD-VI, diug.
LMLHD
ELBO looseness
LMLHD
ELBO looseness
10
-10
10"
-50
-20
10
-30
-100
0
5
10
5
10
15
10
15
5
10
15
context size
context size
context size
context size
(a) Sinusoidal functions.
(b) Mix of affine and sinusoidal functions.
Figure 3: LMLHD and ELBO looseness over context set size for different versions of our algorithm (blue). Our improved TRNG-VI inference scheme yields tighter ELBOs than standard SGD-VI (orange) and, thus, improved performance (cf. text and App. A.1.4 for details).
BML, in (iii) comparison to the state-of-the-ae on (iv) a range of practically relevant meta-learning tasks. To this end, we evaluate our GMM-NP architecture on a diverse set of BML experiments, and present comparisons to
BML algorithms, namely the original NP with mean context aggregation (MA-NP) (Garnelo et al. 2018b), the NP with Bayesian context aggregation (BA-NP) [Volpp al. (02T), the Attentive NP (ANP) (Kim et al., 2019), well the Bootstrapping (At- lentive) NP (B(A)NP) (Lee et al.) 2020). Tab. 1 in App. A.2 gives an overview of the architectural differences of these algorithms. We move details on data generation to App. A.4.A.4 and on the baseline implementations to App. A.2 For a fair comparison, we employ a fixed experimental protocol for all datasets and models: we first perform a Bayesian hyperparameter search (HPO) to determine optimal algorithm settings, individually for each model-dataset combination. We then retrain the best model with 8 different random seeds and report the median log marginal predictive likelihood (LMLHD) as well as the median mean squared error (MSE), both in dependence of the context set size. To foster reproducibility, we provide further details on our experimental protocol in App. App.[A.3 [A.3 the resulting hyperparameters and architecture sizes in App. A.5.7 and publish our source code Lastly, we include a detailed discussion of limitations and computational resources in App. A.5.6
5.1 SYNTHETIC DATASETS
We first study two synthetic function classes (Finn et al. 2017] 2018) on which predictions can be easily visualized: (i) sinusoidal functions with varying amplitudes and phases, as well as (ii) a mix of these sinusoidal functions with affine functions with varying slopes and intercepts. Fig. 2 shows that our GMM-NP outperforms all baselines by a large margin over the whole range of context sizes, both in terms of LMLHD and MSE. This indicates that GMM-NP's improved TP approx- imation indeed yields improved epistemic uncertainty estimation (higher LMLHD). Interestingly, GMM-NP also shows improved accuracy (lower MSE) and, notably, achieves this without any ad- ditional architectural modifications like parallel deterministic paths with attention modules. This is particularly pleasing, as the results show that such deterministic paths indeed improve accuracy, but degrade epistemic uncertainty estimation massively: (B) ANP performs worst in terms of LMLHD. This is further substantiated by (i) observing that MA-NP and BA-NP, both of which don't employ deterministic paths, are among the best baselines w.r.t. LMLHD, and (ii) by visualizing model pre- dictions (Figs. 2 TO III demonstrating that (B)N compute essentially deterministic function samples that fail to correctly estimate the predictive distribution, while our GMM-NP yields esti- mates uncertainty well through variable samples. BNP does not achieve competitive performance, presumably because the bootstrapping approach does not work well for small context sets.
5.2 ABLATION: TASK POSTERIOR INFERENCE
We now demonstrate that GMM-NP's improved performance can indeed be explained by the im- proved TRNG-VI algorithm with accurate GMM TP approximation. To this end, we compare: (i) BA-NP, i.e., amortized VI with reparameterized gradients and unimodal, factorized Gaussian TP (SGD-VI, diag. K = 1). (ii) our GMM-NP, i.e., non-amortized TRNG- VI and full-covariance GMM TP (TRNG-VI, full, K > 1), as well as two models employing TRNG-VI, but unimodal Gaus- sian TP with (iii) full, and (iv) diagonal covariance. The results are shown in Fig. 3 In addition, compare (v) an architecture with full-covariance GMM TP, but trained with iBayes-GMM (Lin ct al. 2020), i.c., with direct step size control instead of trust regions (Fig. App. A.
https: thub. com/ www.w/.hh////o.... gmm_npp
7

---

Published as a conference paper at ICLR 2023
GMM-NP (ours)
MA-NP
BA-NP
BNP
ANP
BANP
10
REGRET
REGRET
LMLHD
MSE
10
10"
300
10"
10"
10
10
10"
10
2
5
10
15
30
20 context size
30
10
20
30
itcration
iteration
context size
(a) BO on Bramin (2D). (b) BO on Hartmann (3D).
(c) 4D Furuta dynamics prodiction.
Figure 4: Panels (a), (b): simple regret over iteration, when using BML models as Bayesian op- timization (BO) surrogates (further results in App. A.5.2). As BO relies on well-calibrated uncer- tainty predictions, the results demonstrate that GMM-NP provides superior uncertainty estimates. Panel (c): log marginal likelihood (LMLHD) and MSE on one-step ahead predictions of 4D Furuta pendulum dynamics. While GMM-NP generally performs best, BANP also shows strong results.
VI Algorithm. Considering the LMLHD metric, we observe a significant performance boost when keeping the traditional factorized Gaussian approximation, but switching from SGD-VI to TRNG- VI, indicating that the standard SGD-V VI approach indeed suboptimal for BML. To study this further, we estimate the looseness of the ELBO (cf. App. A.3.3), i.e., the median (over tasks) value App.
of the KL divergence KL 99,, (-1D,) Pe (-1D;)] between the true and approximate TPs. We observe that TRNG-VI VI provides ELBOs that are tighter by at least one order of magnitude in comparison to SGD-VI. As discussed above, this allows for more efficient optimization of the model parameters O, explaining the performance gain. Lastly, we find that trust regions yield tighter ELBOs than direct step size control and, consequently, improve predictive performance, cf. Fig. 6
Posterior Expressivity. We now study the effect of increasing the expressiveness of the TP approxi- mation. This discussion is supplemented by Fig. where we visualize the TP and its approximation for d. = 2 dimensional latent space. First, observe tighter ELBOs and improved performance when considering full-covariance (but still unimodal) Gaussian TP approximations, and this effect is particularly pronounced for small context sets. This is intuitive, as small context sizes leave lot of task ambiguity, leading to highly correlated latent dimensions (Fig. [Tp. If we now switch to multimodal TP approximations, i.c.. our full GMM-NP architecture with K 1 components (K optimized by HPO), we observe further increase in performance, as the multimodality of the true TP can be captured more accurately (Figs. 112.. This effect especially pronounced for the affin- sinusoidal mix, but also present for the purely sinusoidal function class. As more complex function classes exhibit stronger task ambiguity, the TP will likely exhibit multimodal, correlated structure over wider ranges of context sizes, s.t. an accurate TP approximation will be even more important.
5.3 BAYESIAN OPTIMIZATION
One important application area for probabilistic regression models is as the surrogate model of Bayesian optimization (BO), global black-box optimization algorithm well-known for its sampling efficiency (Shahriari et al. 2016). BO serves as an interesting experiment to benchmark Bayesian models, it relies on well calibraed uncertainty estimates in order to trade-off exploration against exploitation, which is crucial for efficient optimization. As proposed by [arnelo et al. (2018h), use Thompson sampling (Russo al. 2018) as the BO acquisition function and present results on four function classes: (i) TD functions sampled from Gaussian process (GP) priors with RBF kernels with varying lengthscales and signal variances (Kim al. 2019). and parametrized ver- sions of the global optimization benchmark functions (ii) Forrester (ID) (Forrester et al.) (008). (iii) Branin (2D) (Picheny et al. 2013), and (iv) Hartmann-3 (3D) (Szego
1978) as proposed by (Volpp et al., 2020). In Figs. 4a 4b 7 we report the median simple regret, i.e., the difference of the current incumbent value to the function's minimum, over BO iteration. We observe that our
GMM-NP model represents more powerful BO surrogate compared to the baselines, providing fur- ther evidence that TRNG-VI with GMM TP approximations yields superior epistemic uncertainty estimates. We provide further results in App. A.5.2 Figs. g
5.4 DYNAMICS MODELING
We further investigate a challenging dynamics modeling problem on a function class obtained by simulating Furuta pendulum (Furuta et al. 1992), highly non-linear 4D dynamical system, as
8

---

Published as a conference paper at ICLR 2023
GMM-NP (ours)
BA-NP
ANP
LMLHD
MSE
200
0
6
6
6
10
100
4x 10
8
1
6
6
6
6
3x100
0
3
6
6
6
6
2 10
128 256 context size
392
0
128
256
392
7
2
6
6
6
6
context size
9
3
6
6
(a) Log marginal likelihond (LMLHDD and MSE.
(b) GMM-NP (ours).
(c) ANP.
Figure 5: Results on 2D image completion on MNIST. Panels (b),C) visualize predictions on an unseen task showing the digit "6". The first row shows the context pixels, the remaining rows show five corresponding samples. The results are consistent with earlier observations (e.g.. Fig. P our GMM-NP model shows highly variable samples for small context sets, yielding an accurate estimate of epistemic uncertainty, and contracts properly around the ground truth when more context information is available. ANP yields crisp predictions but massively overfits to the noise, explaining bad LMLHD and MSE scores.
We provide further results in App. A.5.3 Fig. 9
proposed by Volpp et al. (2021). The task is to predict the difference of the next system state Tnext E _dd to the current system state E R4, i.e., we study one-step ahead dynamics predictions + y Aæ Inext x E R4. The function class is generated by simulating L = 64 episodes of N = 64 timesteps each (At 0.1 s), where for each episode we randomly sample the 7 physical parameters of the pendulum (3 lengths, 2 masses, 2 friction coefficients). The results (Fig. Ae) show that GMM-NP outperforms the baselines terms of LMLHD by large margin, demonstrating its applicability to complex dynamics prediction tasks where reliable uncertainty estimates are required, e.g., in robotics applications (Deisenroth et al., 2011). Interestingly, while neither ANP nor BNP can reliable solve this task, BANP performs strongly, reaching GMM-NPs asymptotic performance in terms of LMLHD and yielding even slightly better MSE for small context sets.
5.3 IMAGE COMPLETION
To show that our architecture scales to large meta-datasets, we provide results on 2D image com- pletion experiment on the MNIST database of handwritten digits (LeCun & Cortes 2010), as pro- posed by Garnelo al. (2018b). The task is predict pixel intensities y E at 2D pixel locations E R², given a set of context pixels. To obtain a realistic regression task, we add Gaussian noise to each context pixel. The meta-dataset consists of = 60000 images with N = 784 pixels each. The results (Fig. 5) are consistent with our previous findings: GMM-NP yields markedly improved performance, outperforming the baselines over the whole range of context sizes. The architectures with deterministic paths (BB)NP)) fail at properly estimating epistemic uncertainties, leading to low LMLHD values, i.p., for large context sizes. Figs. Figs.5b|5c 5b explain why this is the case: GMM-NP (and also, to some extent, BA-NP) generate meaningful images of high variability, corresponding to well-calibrated uncertainty estimates. In contrast, AB) produce essentially deterministic sam- ples that overfit the noise in the context data. While these samples might appear less blurry than those of GMM-NP and BA-NP, they represent inferior solutions of the regression problem.
6 CONCLUSION AND OUTLOOK
We proposed GMM-NP, a novel BML algorithm inspired by the NP model architecture. Our approach focuses on accurate task posterior inference, a central algorithmic building block that until now has been treated by amortized inference with set encoders optimized using standard, reparametrized gradients. We demonstrate that this approach leads to suboptimal task posterior ap- proximations and, thus, inefficient optimization of model parameters. We apply modern TRNG-VI techniques that enable expressive variational GMMs, which yields tight ELBOs, efficient optimiza- tion, and markedly improved predictive performance in terms of both epistemic uncertainty esti- mation and accuracy. Despite its simplicity, GMM-NP outperforms the state-of-the-an on range of experiments and demonstrates its applicability in practical settings, i.p.. when meta and context data is scarce. This demonstrates that complex architectural extensions, like Bayesian set encoders deterministic, attentive computation paths not required in fact, we observe that determin- istic modules degrade epistemic uncertainty estimation. Therefore, we hope that our work inspires further research on accurate task posterior inference as this turns suffice for accurate BML.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
We took great care to present fair comparison of our GMM-NP algorithm with the baseline models, with statistically reliable results that can be easily reproduced. In particular, we
clearly state the hyperparameter settings hyperparameter optimization procedure we used (Sec. -.3.I).
clearly state the generating process for the datasets on which we evaluated our algorithm (Sec. (A.4),
concisely define the evaluation metries reported (Sec. A.3.3
made to evaluate these metrics large test and sample sets, well as multiple (8)
random seeds, s.t., our results carry statistical significance (Sec. A.3).
use source code from the original authors for all baselines (Sec. A.2
make the source code for our proposed algorithm available online (Sec. A.2
ETHICS STATEMENT
We do not expect any negative ethical or societal impact of our work. Lp., did not sensi- tive/personal data in our experiments.
ACKNOWLEDGMENTS
This work was performed on the HoreKa supercomputer funded by the Ministry of Science, Re- search and the Arts Baden- WiadenWWrrttmbberggg and by the Federal Ministry of Education and Research. The authors further acknowledge support by the state of Baden- Württemberg through bwHPC.
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg. Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, llya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wat- tenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learn- ing on heterogeneous systems, 2015. URL https: www. wwnsorrisooooo .or.:/oooooooooo/o/o/oooooo
Abbas Abdolmaleki, Rudolf Lioutikov, Jan R Peters, Nuno Lau, Luis Pualo Reis, and Gerhard Neumann. Model-based relative entropy stochastic search. Advances in Neural Information Pro- cessing Systems, 2015.
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 1998.
Oleg Arenz, Gerhard Neumann, and Mingjun Zhong. Efficient gradient-free variational inference using policy search. International Conference on Machine Learning, 2018.
Oleg Arenz, Philipp Dahlinger, Zihan Ye. Michael Volpp, and Gerhard Neumann. A unified per- spective on natural gradient variational inference with Gaussian mixture models. arXiv, 2022.
Hagai Attias A variational Bayesian framework for graphical models. Advances in Neural Infor- mation Processing Systems, 2000.
Bart Bakker and Tom Heskes. Task clustering and gating for Bayesian multitask learning. Journal of Machine Learning Research, 2003.
Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning a synaptic learning rule. International Joint Conference on Neural Networks. 1991.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https://www
wandb. com/
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
Charles Blundell, Julien Corneiise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. International Conference on Machine Learning, 2015.
Ben Cazzolato and Zebb Prime. On the dynamics of the Furuta pendulum. Journal of Control Science and Engineering, 2011.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in handful of trials using probabilistic dynamics models. Advances in Neural Information Processing Systems, 2018.
Chris Cremer, Xuechen Li. and David Duvenaud. Inference suboptimality in variational autoen- coders. International Conference on Machine Learning, 2018.
Mare Peter Deisenroth, Gerhard Neumann, and Jan Peters. A survey on policy search for robotics. Foundations and Trends in Robotics, 2011.
Harrison A. Edwards and Amos Storkey. Towards neural statistician. International Conference on Learning Representations, 2017.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leariig for fast adaptation of deep networks. International Conference on Machine Learning, 2017.
Chelsea Finn, Kelvin Xu, and Sergey Levine. Probabilistic model-agnostic meta-learning. Advances in Neural Information Processing Systems, 2018.
Alexander Forrester, Andras Sobester, and Andy J. Keane. Engineering Design via Surrogate Modelling A Practical Guide. Wiley, 2008

---

Published as a conference paper at ICLR 2023
K. Furuta, M. Yamakita, and S. Kobayashi. Swing-up control of inverted pendulum using pseudo- state feedback. Journal of Systems and Control Engineering, 1992.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and M. Ali Eslami. Conditional neural processes. International Conference Machine Learning, 2018a.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo Jimenez Rezende, S. M. Ali Eslami, and Yee Whye Teh. Neural processes. ICML Workshop on Theoretical Foun- dations and Applications of Deep Generative Models, 2018b.
Andrew Gelman, John B. Carlin, Hal S. Stern, and Donald B. Rubin. Bayesian Data Analysis. Chapman and Hall/CRC, 2nd ed. edition, 2004.
Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and D. Sculley. Google vizier: a service for black-box optimization. International Conference on Knowledge Discovery and Data Mining, 2017.
Jonathan Gordon, John Bronskill, Matthias Bauer, Sebastian Nowozin, and Richard E. Turner. Meta- learning probabilistic inference for prediction. International Conference on Learning Represen- lutions, 2019.
Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas L. Griffiths. Recasting gradient-based meta-learning as hierarchical Bayes. International Conference on Learning Rep- resentations, 2018.
Tom Heskes. Empirical Bayes for learning to leam. International Conference on Machine Learning, 2000.
Luke B. Hewitt, Maxwell Nye, Andreea Gane, Tommi S. Jaakkola, and Joshua B. Tenenbaum. The variational homoencoder: learning to infer high capacity generative models from few examples. Conference on Uncertainty in Artificial Intelligence, 2018.
Matthew D. Hoffman, David M. Blei, Chong Wang, and John W. Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
T. Hospedales, A. Antoniou, P. Micaelli, and A. Storkey. Meta-learning in neural networks: a survey. IEEE Transactions Pattern Analysis and Machine Intelligence, 2022.
Mohammad Emtiyaz. Khan and Didrik Nielsen. Fast yet simple -atuiel-rrdient descent for vari- ational inference in complex models. International Symposium on Information Theory and Its Applications (ISITA), 2018.
Mohammad Emtiyaz Khan. Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivas- tava. Fast and scalable Bayesian deep learning by weight-perturbation Adam. International Conference on Machine Learning, 2018.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Tch. Attentive neural processes. International Conference on Learning Representations, 2019.
Taesup Kim, Jaesik Yoon, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn. Bayesian model-agnostic meta-learning. Advances in Neural Information Processing Systems, 2018.
Diederik P. Kingma and Jimmy Ba. Adam: a method for stochastic optimization. International Conference on Learning Representations, 2015.
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. arXiv:1312.6114 [cs, stat), 2013.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. International Conference Learning Representations, 2014.
12

---

Published as a conference paper at ICLR 2023
Gregory Koch, Richard Zemel, and Ruslan Salakhutdinov. Siamese neural networks for one-shot image recognition. International Conference on Machine Learning, 2015.
D. Koller and N. Friedman. Probabilistic Graphical Models: 2009.
Principles and Techniques. MIT Press,
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in Neural Information Processing Systems, 2012.
Kullback and R. A. Leibler, On information and sufficiency. The Annals of Mathematical Statis- tics, 1951.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010.
Juho Lee, Yoonho Lee, Jungtaek Kim, Eunho Yang, Sung Ju Hwang, and Yee Whye Teh. Boot- strapping neural processes. Advances in Neural Information Processing Systems, 2020.
Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt. Stein's lemma for the reparameterization trick with exponential family mixtures. arXiv, 2019.
Wu Lin, Mark Schmidt, and Mohammad Khan. Handling the positive-definite constraint in the Bayesian learning rule. International Conference Machine Learning. 2020.
Christos Louizos, Xiahan Shi, Klamer Schutte, and M. Welling. The functional neural process. Advances in Neural Information Processing Systems, 2019.
David J. C. MacKay. Information Theory, Inference, and Learning Algorithms. Cambridge Univer- sity Press, 2003.
Volodymyr Mnih, Koray Kavukcuoglu. David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 2015.
Radford M. Neal. Bayesian Learning for Neural Networks. Springer-Veriee Verlag, 1996.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam M. Shazeer, Alexander Ku, and Dustin Tran. Image Transformer. International Conference on Machine Learning, 2018.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Detito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- nigh-performance deep learning library. Advances in Neural Information Processing Systems, 2019.
Victor Picheny, Tobias Wagner, and David Ginsbourger. A benchmark of kriging- based infill criteria for noisy optimization. Structural and Multidisciplinary Optimization, 2013.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners. arXiv, 2019.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. International Conference on Artificial Intelligence and Statistics, 2014.
Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. International Conference Learning Representations, 2017.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. International Conference on Machine Learning, 2014.
Daniel J. Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, and Zheng Wen. A tutorial on Thompson sampling. Foundations and Trendy in Machine Learning, 2018.
13

---

Published as a conference paper at ICLR 2023
Adam Santoro, Sergey Bartunov, Matthew M. Botvinick, Daan Wierstra, and Timothy P. Lillicrap- Meta-learning with memory-augmented neural networks. International Conference on Machine Learning. 2016.
Jürgen Schmidhuber.
Evolutionary principles in self-referential learning. learning how to learn: The mete--meta-mecaa.-hoook Master's thesis, 1987.
Jürgen Schmidhuber. Learning to control fast-weight memories: an alternative to dynamic recurrent networks. Neural Computation, 1992.
Bobak Shahriari, Kevin Swersky, Ziyu Wang, Ryan Adams, and Nando Freitas. Taking the Human Out of the Loop: A Review of Bayesian Optimization. Proceedings of the IEEE, 2016.
Jake Snell, Kevin Swersky, and Richard S. Zemel. Prototypical networks for few-shot learning. Advances in Neural Information Processing Systems, 2017.
G. P. Szego and L. C. W. Dixon. Towards global optimisation. North-Holland Pub. Co., 1978.
Sebastian Thrun and Lorien Pratt. Learning to Learn. Kluwer Academic Publishers, 1998.
Ricardo Vilalta and Youssef Drissi. A perspective view and survey of meta-learning, Artificial Intelligence Review, 2005.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Koray Kavukcuoglu, and Daan Wierstra. Match- ing networks for one shot learning. Advances in Neural Information Processing Systems, 2016.
Michael Volpp, Lukas P. Fröhlich, Kirsten Fischer, Andreas Doerr, Stefan Falkner, Frank Hutter, and Christian Daniel. Meta-learning acquisition functions for transfer learning in bayesian opti- mization. International Conference on Learning Representations, 2020.
Michael Volpp, Fabian Flürenbrock, Lukas Großberger, Christian Daniel, and Gerhard Neumann. Bayesian context aggregation for neural processes. International Conference Learning Repre- sentations, 2021.
Edward Wagstaff, Fabian B. Fuchs, Martin Engelcke, Ingmar Posner, and Michael A. Osborne. On the limitations of representing functions on sets. International Conference on Machine Learning. 2019.
Qi Wang and Herke Van Hoof. Doubly stochastic variational inference for neural processes with hierarchical latent variables. International Conference on Machine Learning. 2020.
Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. Advances in Neural Information Processing Systems, 2020.
Winn and C.M. Bishop. Variational message passing. Journal of Machine Learning Research, 2005.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. Advances Neural Information Processing Systems, 2017.
Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey transfer learning. Proceedings of the IEEE, 2020.
14

---

Published as a conference paper at ICLR 2023
A APPENDIX
This appendix provides further details that supplement the main part of our paper.
A.1
ALGORITHMIC DETAILS
In this section we lay the full set of variational update equations and provide pseudocode for GMM-NP algorithm.
A.1.1
VARIATIONAL UPDATE EQUATIONS
We provide the full set of equations required to compute the TRNG update for the variational pa- rameters Φi (WE,es Me,k, De,k). k K,,, parametrizing GMM TP approximation
9 (zt) Σ we (ze(k) Σ We (zelpe, EEN,,
Σ
wigk
(10)
The TRNG-VI update equations, as proposed by Arenz et al. 2022 read
7 E-1
1
new
E,k,old
Re,
(11a)
1
1
new
new
E-1
KeRooeee, +
(re, - Re..
(11b)
++
new O cxp Pe,ks
(11c)
where Rt Te,k, and Pe, are defined as expectations that can approximated from per-component samples using MC:
R4, Reor (xel)) (ze(k)
Ee (ze - (ee,k,old) VT Ne, (zt)
(12a)
Tt.k B (zllk)
Vzchee (ze)
(12b)
hew (ze) log ae (zelk)) J.
Pe. = Euer (zllk)
(12c)
Here, we defined
he (ze) log (e(e)) + log adc..u (ze(k) log Lebr (zz)
(13)
The optimal value for the Lagrangean parameter that enforces the trust region constraint
KT. [g 1 S e;
(14)
defined by a scalar convex optimization problem that can be solved efficiently by a bracketing search, which also ensures positive definiteness of the new covariance matrix Σ ε
A.1.2 GMM INITIALIZATION
We provide details on the initialization of the variational GMMs before meta-training and testing. As we use the same procedure for each task, we drop task indices avoid clutter. Given number K of components for the GMM task posterior (TP) (z(z) Eq. (7). we use a prior p(z) with K components. To initialize the means μ4;, covariances Σg. and mixture weights uy for k E (1...., K), we use the same simple heuristic as Arenz et al. (002))
Draw the means b4 from a didennnnna standard Normal distribution,
The covariances Σ₂ initialized diagonal matrices (with the diagonal),
The weights are initialized uniformly as Wi 1/K.
A.1.3 ALGGORITHM SUMMARY
We provide pseudocode for the veta-training stage of our GMM-NP algorithm in Alg. | and for the prediction stage in Alg.
15

---

Published as a conference paper at ICLR 2023
Algorithm GMM-NP -
Require: Meta-data D; = (æ1,1.N, W.I.NN,El:.
Sample variably-sized auxiliary tasks Dj { Be }. & : L, cf. Sec. A.3.2
Initialize variational parameters $1.L = W MLLKKK E1.[,1:k)
Initialize model parameters 0 while not converged do
each minibatch tasks C 1 (I......) do
Sample žt,k,, ~ 9Φ, (ze(k) for l E k
Evaluate hg, on æg,k,, and D. for f E I, E 1: K, 8: Eq. 13
Update variational parameters de for f I, Eq.
Sample žt,s ~ ls. (2z) for & E I, EEEEl :::5
Estimate gradient of ELBO Eq. 9 Vol (0)
Σs,n Vo log (ye, De ze,a)
Perform step in 0 using Adam end for end while
return Model parameters
Algorithm GMM-NP (Prediction)
Require:
Context data DE {æe I.M, V. J.M. }. model parameters 0, target inputs æf 2..N. Initialize variational parameters Φ. = [w.,1:K Mx,1:K Σ.,
while not converged do
Sample Z., Qs. (z.lk) for & 1
Evaluate he on ke, and DE for 1: K, 1: S. Eq. 13
Update variational parameters Φ.. Eq. (8)
end while
Sample z, 94_ (z.)
return Predictions y decs (zi,n) 2.), 1N:N,
16

---

Published as a conference paper at ICLR 2023
A.1.4 DISCUSSION CONVERGENCE PROPERTIES
Convergence of the ELBO, Our algorithm inherits the convergence guarantee of the variational Bayes algorithm as discussed, e.g. in Bishop (206). In general, convergence of variational Bayes independent of the concrete optimization strategy used for (Φ. 0): as long as both the E-step (step in Φ) and the M-step (step in 0) increase the ELBO objective (first term in Eq. (9). the algorithm guaranteed to converge to local optimum of the ELBO. While standard, reparametrized, variational Bayes (as employed by the baseline methods studied in Sec. () (Φ, O) are optimized jointly using, e-g., Adam (Kingma & Ba 2015), our method alternates between step in & using PRNG-VI (Arenz et al. 2022) and step in 0 using Adam. Nevertheless, both steps increase the ELBO, so our algorithm will converge.
Convergence of the Marginal Likelihood. As discussed in Sec. 4 our GMM-NP algorithm is designed to improve the convergence behaviour w.r.t the marginal likelihood Eq. (9) in compari- son to existing NP-based BML approaches. Recall that the convergence guarantee of the classical expectation-maximiaaiie
(EM) algorithm w.r.t. the marginal likelihood is lost as soon as the E-step becomes intractable, i.e., as soon the posterior distribution cannot be computed exactly, and, thus, has to be approximated by a variational distribution, cf., e.g.. [iishop (2006). This is the case for most models of reasonable complexity, e.g., for the variational autoencoder (Kingma & Welling 2013) or the NP model family (Garnelo et al. (018b). Our GMM-NP model is no exception here, as we build on the NP model for which the TP distribution cannot be computed analytically. Conver- gence of the marginal likelihood when using the ELBO (first term in Eq. (9) as surrogate objective is guaranteed if the ELBO is tight after the E-step, which is the setting of the aforementioned EM algorithm and only the case for perfect TP approximation, i.e., if (=))loo (z|D") O, cf. also App. A...3 For imperfect approximations, the tightness of the bound is controlled by the 0. A better approximate posterior (=(2) yields tighter ELBO, which in turn brings us closer to the EM setting, i.e., typically improves convergence. Our GMM-NP algorithm builds exactly on this insight: we use an expressive TP approximation by full- covariance GMM and powerful optimizer for & TRRG-VVI, (Arenz et al. 2202)) to obtain tighter BL than existing BMI. approaches in order to achieve optimization of the model parameters in way that efficiently maximizes the marginal likelihood.
variational gap (L(go(z)][pg (I
17

---

Published as a conference paper at ICLR 2023
A.2 BASELINE ALGORITHM
Tab. D gives an overview of the architectural differences of the BML approaches we compared in empirical evaluation (Sec.
Table Comparison of state-of-the-ar approaches for Bayesian meta-learning (TRNGD re- gion natural gradient descent, RSGD = reparametrized stochastic gradient descent, SGD = stochastic gradient descent, SE set encoder, MA mean aggregation, BA Bayesian aggregation, SA self attention, CA cross attention).
TP Approx.
VI Approach
Amortization
Det. Path
GMM-NP (ours)
Full-cov. GMM
TRNGD
none SE + MA
none
MA-NP Garnelo et al.
2018b
Diag. Gaussian
RSGD
none
BA-NP
Volpp et al.
2021
Diag. Gaussian
RSGD
SE BA
none
BNP
Lee et al. 2020)
Non-parametric
SGD
SE + MA SA SE MA SA SE + MA
none
ANP
Kim al. 2019
Diag. Gaussian
RSGD
CA
BANP
Lee et al.
2020
Non-parametric
SGD
CA
To compute our results, we consistently use code by the original authors. We also provide source code for our proposed GMM-NP algorithm:
Source code four our GMM-NP algorithm:
catcp::://ittuu..com//at cot..ccc///..uubggggua....o.a.a. np
MA-NP, ANP:
BA-NP:
BNP, BANP:
Beteas//aaahuub.com/uuoolll////
18

---

Published as a conference paper at ICLR 2023
A.3 EXPERIMENTAL PROTOCOL
To foster reproducibility, we provide details on our experimental protocol.
A.3.1 MODEL HYPERPARAMETERS
To arrive at fair comparison of our GMM-NP model with the baseline approaches, optimize model hyperparameters individually for each model-dataset combination presented in Sec. El Con- cretely, perform Bayesian hyperparameter sweep with 256 trials for each model-datasee com- bination over the parameters detailed below. For the image completion experiment on MNIST, we employ a grid search with fewer trials to keep the computational effort manageable. For hyper- parameters not mentioned below, we consistently use standard settings proposed by the original authors. To implement the hyperparameter search, we the wandb sweep functionality (Biewald, 2020).
Observation Noise Parametrization. As detailed in Sec. 3.2 all compared models (including our GMM-NP) employ Gaussian likelihood of the form
po (y x 2) N (yldect 2), diag (d.)),
(15)
where the mean is computed by a decoder DNN dec& receiving the input location æ and latent sample z. However, different parametrizations of the observation noise variance of are used in the literature. As it is not clear which setting is fairest, we also treat the observation noise parametriza- tion as a hyperparameter. Concretely, for each model-dataset combination, we test the following settings for the observation noise (with individual hyperparameter sweeps) and report the best per- forming one:
1. a o attuue with o n,true being the true noise variance on the data,
2. af a single float value, optimized jointly with 0,
3. of = deeg (ar), i.e., observation noise parametrized by a second decoder network, opti-
mized jointly with decý, but receiving only the input location,
4. of = dec (x, x). i.c., observation noise is parametrized by a second decoder network, op- timized jointly with decg. and also receiving both the input location and the latent sample.
For all compared models, and regardless of the parametrization, we bound the observation noise from below using a softplus transformation s.t. Tn 2 Gn,min = 0.1, as proposed by (Garnelo et al.) 2018 8b
Kim et al. 2019 Lee al. 2020).
DNN Architectures. For all experiments and all baseline models, we use encoder and decoder DNNs with two hidden layers. Likewise, our GMM-NP model uses a decoder DNN with two hidden layers. We optimize the number of hidden units per layer within the bounds {8..... (8....,66 64].
Latent Dimensionalities. For baseline models with parametric latent distributions (all except B(A)NP), we optimize the latent dimension de within the bounds {I...., 64]. As our GMM-NP algorithm employs full covariance matrices, we restrict the bounds for d2 to {1....,8 8) for a fair comparison.
Number of GMM components. For our GMM-NP algorithm, as well as for iBayes-GMM (Ln et al. 2020) used for the comparison in Sec. A.5.I we optimize the number of GMM components within the bounds {1...., ......,10) 10).
Learning Rates and Trust Region Bounds. All algorithms use the Adam optimizer with standard settings to update DNN weights. We optimize the corresponding learning rates a log-uniform scale within the bounds [10- 10-'] We use the same settings to optimize the step size for the GMM updates of the variational parameters of the iBayes-GMii GMM algorithm (Lin et al.) 2020 used for the comparison in Sec. A.5.T As proposed by Arenz et al. (2022), we optimize the Lagrangean parameter 7 of our GMM-NP algorithm using bracketing search the interval [10-3 10-1].
19

---

Published as a conference paper at ICLR 2023
A.3.2 AUXILIARY SUBTASK GENERATION FOR META-TRAINING
We describe the procedure to sample auxiliary subtasks during meta-training in more detail, cf. Sec.
Nomenclature. Recall from Sec. 3.2 that we define meta-task the of all available (noisy) evaluations Dz, l {1,.... L from an unknown function ft and that each meta-task contains N examples. Thus, meta-task Dj is all data a BML algorithm has available to learn about ft during meta-training, A subtask of meta-task Dj is defined as an arbitrary subset of De.
Auxiliary Subtask Sampling. As described in Sec. standard NP meta meta-trainin training samples auxil- lary subtasks from the metadata for each minibatch step order provide the decoder with sam- ples from task posterior approximations informed by a range of context sizes. We use the following standard procedure (Garnelo et al. 2018b) Kim et al., 2019: Lee al.) 2020) sample auxiliary subtasks evaluate the optimization objectives of the baseline approaches (e.g., Eq. 5 for standard NP). Given a minibatch C {I,.... L} of meta-tasks Di, l E L, we first sample auxiliary subtasks Dr with a size Ñ drawn uniformly from Ñ {NNin + ,,,..,,,Nau with Min 2 1 and Nmax S N.
Then, we sample context sets Dj C CD of size M, drawn uniformly from M E {I....,
N. Dj
and Dj are then used in Eq. to compute the ELBO objective for the current minibatch.
As described in Sec, 4 our GMM-NP algorithm uses similar approach: we employ auxiliary subtasks with sizes N drawn uniformly from Ñ E (Nmin: .... Nuax } to evaluate the updates for the variational GMM parameters and the model parameters. Note that our algorithm does not require sample context sets during meta training from the auxiliary subtasks. Furthermore, recall that we train one variational GMM for each auxiliary subtask and retain those GMMs over the whole course of meta training, so we fix set of L auxiliary subtasks at the beginning of meta-training (in contrast standard NPs, which sample new subtasks for each minibatch).
We use the following settings for Nmins Nmax in our experiments: Nmin 1, Nmxx N, except for MNIST image completion where we use N Nmax N/2.
Further, we use L 32L, except for MNIST image completion where we use L 8.
A.3.3 METRICS
For each model-dataset combination, we retrain the best hyperparameter setting determined accord- ing to Sec. A.3.I with 8 different random seeds used for model initialization, and report the median value together with (5%, 95%) percentiles of the metrics computed according the formulae pro- vided below. For all experiments (except the MNIST image completion experiment), we evaluate all metrics on 7. 256 unseen test tasks Di. with D; {yN,I:N, .1,1:N] and N 64, from which we sample context sets D; C Dr. For the image completion experiment we use L = 1024 and N 784 (the number of pixels per image). We report the results in dependence of the context set size.
Log Marginal Predictive Likelihood (LMLHD). For a given task the LMLHD is defined by
Eq. 44. which we restate here for convenience:
log do (yr,1-N DL,I,, D2) Df) logg / II po we ze) (ze (D) dzz.
(16)
Here, use the generic notation q (2y(Dī) to denote the task posterior TP approximation, the concrete form of which depends on the BML model under consideration. As the integral analyt- ically intractable, we resort MC approximation. To this end, sample S = 1024 samples
20

---

Published as a conference paper at ICLR 2023
2t, (ze(Di) in the test set and compute (Volpp et al.
2021
N
log do (yr, Eg,I:N Df) = log
II
po (ye, Dens ze) (z dze
(17)
1==1
N
Rs log
Σ
II
Po (ye, aee, Ze,s)
(18)
5
s=
log S + logsumexp
log pa (yr. / Re,s).
(19)
#=1
where logsumexp denotes the numerically stable implementation of the function log E. exp(z,), available in any scientific computing framework. We then compute the median of this expression over all tasks of the test set.
Mean Squared Error (MSE). We report the MSE w.r.t. the mean prediction. That is, for a given task we again draw S 1024 samples Zt,s ~ q (zelDi) and compute
MSE (YE,IN (...I.,EEE.N
dec (ze, ze,e) ye,
(20)
We then compute the median of this expression over all tasks of the test
ELBO Looseness. For given task f, we define the ELBO looseness the KL-divergence be- tween the approximate and task posteriors. According to Eq. 44. this decomposes as
KL (zc(De) lpe (zx[De)] log 9o (yC,I:N / Dj)
(21)
q (zx
(D(
Σ log po (ye, Fe ze) + log
(22)
q (De)
with log do UUL.NNJa,,IN Df) defined by Eq. 16 The second term the ELBO,
N
(z) (Dj
(2
L0, D2) = E,( (P))
log Po (ye, Le.. zt) + log (z) (2))
(23)
Ln=1
where we made its dependence on both the test set Dr and the context set D; C Dr explicit (in contrast our notation in the main part of this paper). We say the ELBBO tight if its looseness is zero. Then, log qe (ye.1:Nl 11,1.N, D2) L (0, Dz), and optimization of the ELBO w.r.t. 0 is
cquivalent optimization of the LMLHD.
For our ablation study (Sec. 5.2) we estimate the looseness of the ELBO by computing the differ- ence of an importance- -pportance-weighted MC estimate with proposal distribution q (zz (De) of the LMLHD and an MC estimate of the ELBO Eq. (23) with = 1024 samples zt, (D.).
21

---

Published as a conference paper at ICLR 2023
A.4 DATA GENERATION
We provide details on the meta-datasets we use to train the models we compare in Sec. Concretely, provide
the dimension dæ inputs Zt,n E R
the domain C C 2 from which we uniformly sample zt,n,
the dimension dy of targets ye.r E IRdv,
an expression for the function Jt : R2d - IRd,, s.t., ye - Se (xt.,) + En.
the noise standard deviation 0, s.t., ~ ''' o0.o²).
the number of meta-tasks and the number N datapoints for each meta-task.
We denote the uniform distribution on (a, b)d C Rd by U (a, 6)*
Sinusoidal Functions.
d
C 1-5.0,5.0) 5.0]
dy =1
(x) = sin (x - (e).. Ae ~ (0.1,5,..) de U x0.0,
0 = 0.25
L 64, N 16
Mix of Affine and Sinusoidal Functions.
d₂
C = 5.00,
dy =
f (x) a&x + by, at ~
17 (x) Assin - Φb),, A ~ U (0.1,5.0), di T).0,
fe is given either by ft or ft with probability 0.5.
0.25
L 64, N = 16
RBF-GP samples.
dg
C 2.0, 2.0]
dy =1
Je is drawn from a Gaussian process prior with RBF kernel with lengthscale
4 (0.5, 1.0) and signal variance ~ U (0.5, 1.0).
o =01 0.1
L 64, N 16
Forrester ID.
dg =
dy 1
We use the parametrized Forrester function Forrester et al. 2008 as defined on http.haaaassssaaaaoootootttaaaaa hmm
σ 0.25
L. 64, N 16
22

---

Published as a conference paper at ICLR 2023
Branin 2D.
dæ =
dy =1
We use the definition given on http:p//wwww..uuu...ssssrraaaaar
and apply translations Te ~ U (-0.25, 0.25)* to æ. and scale the function values by
U (0.75, 1.25).
0.25 L 64, N = 16
Hartmann 3D.
d₂ =
dy 1
We use the definition on wt 3hhmll
and apply translations Te ~ U (--0.25,0.2 0.25) to æ, and scale the function values by
8g (0.75, 1.25).
0.1 L 64, N 16
4D Furuta Dynamics Prediction.
dæ
dy
We the dynamics equations given in Cazzolato & Prime (201) simulate episodes, starting from the pendulum balancing in the upright position. The input is the current system state x E R.. the target is the difference the next system state Tnext E R i.c.,
y Ax = Enext
E IR4.
Noise is generated by random actions on the joints.
64, N 64
2D MNIST Image Completion.
dy =
d, =1
We use the MNIST handwritten image database (LeCun Cortes 2010). Each image corresponds one task. The input æ is the pixel location, the target is the pixel intensity.
0.25
L 60000, = 784
23

---

Published as a conference paper at ICLR 2023
A.5
FURTHER EXPERIMENTAL RESULTS
We provide further experimental results for the experiments presented in Sec. 5
A.5.1 ABLATION: TRUST REGIONS
In Fig. 6 we compare two methods for step size control for natural gradient VI, namely direct step size control proposed by (Lin et al. 2020) and trust region step size control (Arenz et al. 2022 as used by our GMM-NP algorithm. We observe that trust regions lead to more robust optimization of the variational parameters, and, thus, to tighter ELBOs. This allows more efficient optimization of the model parameters, leading to improved predictive performance.
GMM-NP (ours): TRNG-VI, full, K=1
NG-VI, fall, K+1
ELBO looseness
LMLHD
LMLHD
ELBO looseness
-10
10
10
10
-50
-30
-100
10"
5
10
15
context size
context size
context
context size
(a) Sinusoidal functions.
(b) Mix affine and sinusoidal functions.
Figure 6: Log marginal predictive likelihood (LMLHD) and ELBO looseness context size for our trust region natural gradient VI (TRNG-VII-based (Arenz et al. 2022) GMM-NP algorithm in comparison iBayes-GMM (Lin et al. 2020 that uses direct step size control instead of trust regions (N). Trust regions improve variational optimization, leading to tighter ELBOs, and, consequently, to improved predictive performance.
24

---

Published as a conference paper at ICLR 2023
A.5.2 BAYESIAN OPTIMIZATION EXPERIMENTS
We provide the full set of results for our Bayesian optimzation experiments, cf. Sec. 5.3 Fig. shows the optimization regret for all four evaluated function classes, and Fig. g the corresponding results for LMLHD and MSE.
GMM-NP (ours)
MA-NP
BA-NP
BNP
ANP
BANP
REGRET
REGRET
REGRET
REGRET
10
10
10
10"
10a
10
10"
10
10
15
5
10
15
5
15
10 iteration
iteration
itcration
iteration
(a) ID RBF-GP samples.
(b) ID Forrester.
(c) 2D Branin.
(d) 3D Hartmann.
Figure 7: Simple regret over optimization iteration, when using BML models as Bayesian Opti- mization (BO) surrogates on various function classes, As BO relies on well-calibrated uncertainty predictions, the results demonstrate that GMM-NP provides superior uncertainty estimates.
GMM-NP (ours)
MA-NP
BA-NP
BNP
ANP
BANP
LMLIID
MSE
LMLHD
MSE
10
10
-100
10a
-50
-150
10"
5
10
15
0
10
15
10
15
0
10
context size
context size
context size
context size
(a) ID RBF-GP samples.
(b) ID Forrester functions.
LMLHD
MSE
LMLHD
MSE
10
10"
-25
10
-50
-25
10
0
5
10
15
0
10
15
10
15
0
10
15
context size
context size
context size
context size
(c) 2D Branin functions.
(d) 3D Hartmann-3 functions.
Figure 8: Log marginal predictive likelihood (LMLHD) and mean squared error (MSE) over context size on various function classes. GMM-NP generally performs favorably, showing accurate predic- tions with well-calibrated uncertaintics.
25

---

Published as a conference paper at ICLR 2023
A.5.3
2D IMAGE COMPLETION ON MNIST
Fig.
9
shows the full of predictions on the 2D image completion experiment on MNIST.
2
6
4
6
6
6
6
3
6
2
6
6
$
4
4
4
6
4
6
8
1
6
6
6
6
6
6
6
4
4
6
6
19
4
6
6
O
3
6
3
6
6
6
5 4
6
6
8
0
6
6 6
7
2
6
6
6
3
6
2
6
6
3
6
6
8
6
9
3
6
6
6
6
4
4
6
6
4
6
8
(a) GMM-NP (ours).
(b) BA-NP.
(c) ANP.
(d) BANP.
Figure 9: Predictions on an unseen instance of the MNIST 2D image completion task, showing the
digit "6". The first row of each panel shows the context pixels (ranging from zero pixels in the left column to the full image in the right column). The remaining rows show five samples from the BML models, conditioned on the context pixels shown the first row. The results are consistent with observations from the other experiments (e.g., Fig. T0): our GMM-NP model shows highly variable samples for small context sets, yielding an accurate estimate of epistemic uncertainty, and contracts properly around the ground truth when more context information is available. BA-NP also shows variable samples, albeit of lower quality. ANP and BANP yield crisp predictions but massively overfit the noise, explaining their low LMLHD scores. Note also that BANP does not allow predictions for empty context sets.
26

---

Published as a conference paper at ICLR 2023
A.5.4
VISUALIZATION OF MODEL PREDICTIONS
Figs. I0 and show further visualization of predictions of models trained on the mix of affine and sinusoidal functions.
0.0
-5
0
X
10
2.5
2.5
2
0.0
>
0.0
2.5
-2.
10
5
0
5
5
2.5
0.0
2
-2.5
5
0
5
5
X
-5
x
X
(a) GMM-NP (ours).
(b) BA-NP
(c) ANP.
(d) BANP
Figure 10: Function samples computed by various BML models (columns), trained on a function class consisting of mix of affine and sinusoidal functions (ef. Sec. (.T). when provided with in- creasing amounts of context examples (red crosses, rows) from an unseen sinusoidal representative function. We observe that our GMM-NP model accurately quantifies epistemic uncertainty through the variability of its function samples. BA-NP also shows variable samples, but does not achieve the same predictive performance due its inaccurate approximation of the task posterior distribution. ANP and BANP, both of which employ deterministic computation paths with attention modules, produce essentially deterministic predictions that massively overfit the context data and fail to give a reasonable estimate of the predictive distribution. Therefore, these models have to quantify epis- temic uncertainty through the likelihood noise variance, which ineffective, cf. Fig. Note also that BANP does not provide predictions for empty context sets.
27

---

Published as a conference paper at ICLR 2023
-10
-25
5
-5
-5
-5
x
-5
0
5
-5
0
-5
0
(a) GMM-NP (ours).
(b) BA-NP
(c) ANP.
(d) BANP.
Figure 11: This figure shows the same data Fig. 10 but for each function sample also show band of +1 standard deviation of the observation noise, as computed by the decoder DNN. GMM- NP quantifies epistemic uncertainty correctly through its task posterior approximation, and thus does not have to rely on the decoder DNN to quantify epistemic uncertainty through the observation noise. In contrast, ANP and BANP fail produce variable function samples, and have to make up for that by quantifying epistemic uncertainty through the observation noise, which is ineffective. Note also that BANP does not provide predictions for empty context sets.
28

---

Published as a conference paper at ICLR 2023
A.5.5 VISUALIZATION OF LATENT SPACE STRUCTURE
We provide further visualizations similar Fig. D of the task posterior approximation and corre- sponding function samples of our GMM-NP, when trained on the sinusoidal function class.
29

---

Published as a conference paper at ICLR 2023
5.0
2.5
-1
the:
22
0.0
-
-25
5.00
4
2
0
2
4
(a) A small context set (one single example indicated by the red cross) yields highly correlated, multi-modal task posterior distribution. Our GMM approximation correctly captures this, s.t., amplitudes and phases of the predicted sinusoidal functions accordance with the observed context data point.
5.0
2.5
0.0
-2
27 GE%
-1
-25
5.0
-4
0
2
4
(b) A second example on another instance sinusoidal function class, where the task posterior shows pronounced multimodality, which translates into a bimodal predictive distribution.
2
-12
-144
-2
-18
0 0.0%
31.00%
2
2
(c) Larger context sizes (three examples, red crosses) leave less task ambiguity, resulting in unimodal and nearly isotropic task posterior distribution. Our GMM approximation again correctly approximates this distri- bution, making use of only two of the K 3 mixture components (the mixture weight of the orange component close to zero, so no samples from this component are observed).
Figure 12: Visualization of our GMM-NP model for a de = 2 dimensional latent space, trained on sinusoidal functions with varying amplitudes and phases, ef. Sec. 5.1 Left panels: unnormalized task posterior distribution (contours) and variational GMM approximation with K = 3 components (ellipses, mixture weights in %). Right panels: corresponding samples from our model (blue lines), when having observed a context data set (red crosses), together with unobserved ground truth data (black dots). The visualizations show that (i) the true task posterior distribution can be highly cor- related and multimodal, i.p., for small context sets (panels a.b), (ii) our variational task posterior approximation correctly approximates this distribution, which (iii) leads to expressive predictive distributions that incorporate both the inductive priors learned from the meta-dataset (all samples are sinusoidal in shape) and the additional information contained the context set (all samples pass close to the context data).
30

---

Published as a conference paper at ICLR 2023
A.5.6 RUNTIME COMPARISON
Discussion of Limitations. As the meta-training stage of GMM-NP requires computational effort comparable to standard NP (cf. Sec. T. the only computational overhead of our algorithm occurs at test time, due to the optimization loop required to fit variational GMM to D.. While this can be trivially parallelized for multiple test tasks, it incurs a higher computational burden in comparison to the single forward pass through NP's set encoder (we provide an evaluation of the runtime of our algorithm on the synthetic tasks studied in Sec. belloow We leave detailed examination for future work, but mention two possible remedies: (i) for problems where test data arrives sequentially, expect that few update steps in &h suffice to reach convergence, and (ii) it might be possible to find amortized approximations to Eqs. (8) similar in spirit to standard NP, that retain the advantages of TRNG-VI.
shows the learning curves for meta-training corresponding to the results presented in the main part of this paper. As discussed in Sec. GMM-NP incurs computational cost comparable to the baseline methods.
Meta- Training. Fig. 13
Loss muuta-triiiing Loss (meta-triinngg)
model GMM-NP (ours)
model MA-NP
model
BA-NP
5.0
200
20
2.5
0
5
10
model
BNP
model ANP
model BANP
5.0
5.0
10
2.5
2.5
10
20
Wall clock time (min)
Wall clock time (min)
Wall clock time (min)
(a) Learning curves on the sinusoidal function class.
Loss (mela-training Loss (meta- (osti-training)
model
GMM-NP (ours)
model MA-NP
model BA-NP
500
10
5.0
5
2.5
10
10
10
model BNP
model ANP
model BANP
20
10
5
10
20
Wall clock time (min)
Wall clock time (min)
Wall clock time (min)
(b) Learning curves on the mix of affine and sinusoidal function class.
Figure 13: Learning curves for meta-training on the synthetic datasets, ef. Sec. For each method, we show the learning curves for the seeds used to compute the results presented the main text. For GMM-NP, joint loss for the encoder and decoder parameters (Φ, 0). Note that for GMM-NP, convergence of implies convergence of the variational parameters Φb. As discussed in Sec. GMM-NP incurs computational cost comparable to the baseline methods.
show the loss for the decoder parameters 0. for the other methods we show the
Test-time Adaptation. As discussed in Sec. 4 GMM-NP does not amortize TP inference, i.e., it does not learn a set encoder architecture, but adapts new variational GMMs at test time. Naturally, this incurs higher computational cost in comparison amortized architectures, which compute
31

---

Published as a conference paper at ICLR 2023
predictions on test tasks in single forward pass through their set-encoder - decoder architecture. In Fig. [4] we show the learning curves for fitting variational GMMs (by iterating Eqs. (8) to the test tasks and for the range of context sizes used to compute the results presented in Sec. [. GMM-
NP's TRNG-VI optimization converges in approx. 0.1 s - 1 s per test task (depending on the context
set size).
Loss (adapt)
Loss (adapt)
500
Context size
500
Context size
0.0
0.5
1000
Wall clock time / task (sec)
Step
(a) Sinusoidal functions.
Loss (adapt)
Loss (adapt)
2000
Context size
2000
Context size
16
1000
Wall clock time / task (see)
Step
(b) Mix of affine and sinusoidal functions.
Figure 14: Learning curves for fitting variational GMMs to the test tasks by TRNG-VI (Arenz et al. 2022), used by our GMM-NP (Sec. 4), the synthetic datasets (Sec. 5.1). The quantity labelled "Loss (adapt)" is the expected negative log density of the unnormalized TP under the GMM TP approximation. Note that this is not the loss function optimized by iterating Eqs. (8). but serves as a proxy to judge convergence. We show results in terms of wall clock time per test task (left panels) and in terms of TRNG-VI steps (right panels), for the range of context sizes used to compute the results in the main text. GMM-NP's TRNG-VI optimization converges in approx. 0.1 1 per test task (depending on the context set size).
32

---

Published as a conference paper at ICLR 2023
A.5.7 ANALYSIS OF HPO RESULTS
optimized architectural hyperparameters individually for each model-dataset combination presented in our empirical evaluation, in order to arrive at fair com- parison of our GMM-NP with the baseline methods. In Tab. 2 we provide the resulting settings for the latent dimensionality de. and the number of parameters of the BML models compared in Sec. 1.1 While the number variational parameters during meta-training is naturally comparably high for non-amortizing methods such as GMM-NP, we observe that the expressive GMM-NP TP approximation allows comparably lightweight decoders and small latent dimensions. This is intu- itive, simple TP approximations require (i) large latent dimensions encode relevant information the latent space, together with (ii) expressive decoder architectures to transform the simple latent distribution into expressive predictive distribution. Note further that the variational parameters belonging to different tasks are not coupled in non-amortizing architectures such as ours, which al- lows trivial parallelization of the variational optimization between tasks, explaining why the compu- tational cost is easily managable, cf. Sec. A.5566) Note also that the number of variational parameters one has to store and adapt for GMM-NP make predictions on unseen test tasks is comparably small because the variational GMMs learned during meta-traine meta-training can be discarded as they are not required for predictions at test time.
As discussed in Sees. 5 and A.3
Table 2: Results of our hyperparameter optimization on the sinusoidal function class and on the mix of affine and sinusoidal functions. We provide the settings for the latent dimensionality d, and number of parameters of the BML models compared Sec. 5.1 (i.e., the number of decoder parameters |0] as well the number of encoder variational parameters (g)). If attentive modules are present, their parameters are counted as being part of the encoder. For our GMM-NP, we also provide the number of GMM-components K. Furthermore, inference but learns separate variational GMMs for each subtask generated from the meta-dataset (ef. Secs. and (...2), also provide the total number of variational GMM parameters during meta-training. Note that these variational GMMs are decoupled and can be optimized in parallel. Furthermore they not required for predictions at test time and can be discarded after meta- training.
GMM-NP does not amortize TP-
dz
( (per task)
(Φ] Amneta-training
Sinusoid
GMM-NP (ours)
79872
MA-NP
10
2298
2595
2595
BA-NP
27
7334
7386
7386
BNP
56
12770
19488
19488
ANP
19
4096
10335
10335
BANP
40
28320
28320
Line-Sine
GMM-NP (ours)
2289
59
120832
MA-NP
19
2562
3679
3679
BA-NP
30
11592
11650
11650
BNP
541
11882
18141
18144
ANP
30
4496
14448
14448
BANP
32
4226
18304
18304
33


Reviews:
[{"review_id": "fY84-ypA5Oa", "reviewer": "Reviewer_sb3h", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes GMM-NP, a novel BML algorithm inspired by the NP model architecture, which focuses on accurate task posterior inference. Despite its simplicity, GMM-NP outperforms the state-of-the-art on a range of experiments and demonstrates its applicability in practical settings", "summary_of_the_review": "This paper is well-done for improving the optimization and inference of BNP. This paper can be improved by providing more empirical results and elaborations. ", "strength_and_weaknesses": "Strengths: \n1. identifies the optimization suboptimal and solves by combining VI with natural gradients and trust regions\n2. strong experimental results which covers multiple settings and applications\nWeakness/Questions:\n1. As said in the paper, the additional overhead in the testing time can be parallelized. Can you give some empirical results to support it?\n2. How would your proposed method be used for Bayesian uncertainty learning models, especially those multi-modal, time-varying models for classification/regression?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "CUPGnYErqQ", "reviewer": "Reviewer_FbJt", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method for Bayesian meta-learning that aims to learn a conditional generative model f(x,z) such that variational inference on a latent z can model the posterior predictive distribution given context data from a new task. Specifically, the authors draw connections between their approach and the Neural Process (NP) model, and emphasize how performing non-amortized variational inference via trust-region natural gradient optimization of a GMM distribution on z can yield higher quality posterior predictive distributions than NP-style models, even those with more complex encoders.", "summary_of_the_review": "Overall, I think this is a strong paper, presenting a simple, novel algorithm which achieves impressive results. The work is well presented, and the experimental evaluation is thorough, so I would recommend acceptance.", "strength_and_weaknesses": "Strengths:\n- Overall, the paper was clear and easy to follow. The preliminaries introduced the intuition behind the core topics of the approach well. The description of VI was clear and made the exposition of the core contribution of this paper easy to understand.\n- The approach is simple yet effective, and serves to highlight the utility of more expressive task posteriors in a latent-space meta-learning context.\n- The experiments are thorough and demonstrate the capabilities of the architecture at representing an accurate task posterior both qualitatively, as well as quantitatively via the log-marginal likelihood metric on synthetic task distributions as well as the utility of accurate posterior modeling on downstream tasks such as bayesian optimization, dynamics modeling, and image completion. I also appreciated the ablations measuring ELBO looseness as a function of VI strategy and distribution expressivity to support the author's hypothesis that the quality of variational approximation correlated with downstream performance.\n\nWeaknesses:\n- As stated by the authors, a main limitation is the added computational complexity due to the non-amortized variational inference. While this leads to higher quality task posteriors, it brings   a natural tradeoff of prediction quality vs runtime. It would have been useful to see more details on the computational overhead in terms of wall clock speed -- how many iterations of VI are necessary for convergence on average? How does wall clock runtime compare to algorithms with more complex encoders, such as attention based NPs?\n- The method requires optimizing a separate task posterior not only for each task in the meta-dataset, but also for each context size size considered for each task, each requiring storing the set of GMM parameters. In contrast, an amortized inference approach would share parameters across tasks, but here the parameters would be the weights of an encoding network. How do the number of training parameters compare across methods?\nThe number of latent-dimensions differed between algorithms (and was selected in the hyperparameter optimization). However, this made it difficult to compare different approaches -- perhaps in a higher dimensional latent space, a unimodal distribution can encode similar information as in a multimodal distribution in a lower-dimensional latent space? The paper would be strengthened by either a comparison between methods with a fixed latent dimension, or a more in-depth discussion of any trends in the choice of latent-space dimension that emerge through the hyperparameter optimization.\n\nQuestions:\n- What data was used for hyperparameter optimization?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "AMr7KArTor", "reviewer": "Reviewer_tZEh", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The submission modified the Neural Process model by replacing the Gaussian approximated posterior parameterized by encoder with a Gaussain mixture distribution parameterized by full-covariance matrix. For inference, the proposed method divide the VI into two alternating updating steps between variaitonal and model parameters. For variaitonal parameters, the method uses the existing trust region natural gradient based VI in [Arenz et al. 2022]; for model parameters, it uses the standard VI. The propsoed model and inference is validated on a range of experiments. \n\n", "summary_of_the_review": "From the technical perspective, the methodology in the submission is correct in my eyes (I did not check it carefully) and the experiments are rich and solid. Generally speaking, the submission is a good paper without obvious defects. However, from the model inference perspective, this work simply applied the method in \"A Unified Perspective on Natural Gradient Variational Inference with Gaussian Mixture Models\" [Arenz et al. 2022] to modify the Neural process model; all contributions mentioned in this work, \"full covariance matrix\", \"Gaussian mixture distribution\", \"non-amortized\", \"trust region natural gradient descent\", are inherited from this previous work. Therefore, the incremental contribution dampens my enthusiasm to give a very high score. ", "strength_and_weaknesses": "Strength: The paper is generally easy to follow, and the derivation seems correct to me although I did not check it carefully. Replacing the Gaussian approximated posterior parameterized by an encoder with a Gaussain mixture distribution is an interesting idea. \n\nWeakness: The proposed inference method is inherited from the TRNG-VI in [Arenz et al. 2022] and the contribution is incremental. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7jfM7rwuMzr", "reviewer": "Reviewer_7Kjw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a new way of doing Bayesian Meta Learning (BML) by fitting a full-covariance Gaussian Mixture Model (GMM) to approximate each Task Posterior (TP). TRNG-VI proposed by Arenz et al. (2022) was used to ensure efficient and robust optimization of the variational bound. The training complexity is similar to existing approaches, and does amortize inference over tasks and thus does not learn a set encoder architecture. Experiments on synthetic datasets showed the proposed GMM-NP method outperformed all baselines by a large margin over the whole range of context sizes, both in terms of LMLHD and MSE, in different test scenarios regarding task posterior inference, Bayesian optimization, dynamic modeling and image completion.", "summary_of_the_review": "The GMM-NP approach proposed by this paper achieves much better results than existing approaches for BML, with less requirements in model architectures and objective approximations. The experiments illustrated the effect of each part of the design in a high quality. I would recommend an accept for it, to make people in the BML/NP area aware of the importance of accurate task posterior inference and work along this fundamental direction.", "strength_and_weaknesses": "1. What is the distribution of p(z), the prior used?\n\n2. Is there any convergence guarantee for Alg.1, if $\\phi$ is updated via TRNG and $\\theta$ via Adam?\n\n3. For Figure 1:\n- What is the definition of  z1 and z2 - amplitudes and phases?\n- \"Right panel: corresponding function samples from our model (blue lines)\" - is each line is a sample from q(z)?\n- The paper mentioned \"all samples pass through the red context example\". But from the plot it does not seem so. Is there any guarantee that each sampled functions will certainly pass through the points in the context data points?\n\n4. The paper mentioned \"sample a fixed set of auxiliary subtasks\" during training. What's the difference between tasks and sub-tasks?\n\n5. The paper mentioned the TRNG-VI require at most first-order gradients of $\\tilde{p}(z_l)$. Could you point out which variational update equation needs this gradient info?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster