Bayesian Context Aggregation for Neural Processes
Authors: Michael Volpp
Fabian Flürenbrock
Lukas Grossberger
Christian Daniel
Gerhard Neumann
Venue: ICLR 2021 Poster (2021)
Keywords: Aggregation Methods
Neural Processes
Latent Variable Models
Meta Learning
Multi-task Learning
Deep Sets

Abstract: Formulating scalable probabilistic regression models with reliable uncertainty estimates has been a long-standing challenge in machine learning research.  Recently, casting probabilistic regression as a multi-task learning problem in terms of conditional latent variable (CLV) models such as the Neural Process (NP) has shown promising results. In this paper, we focus on context aggregation, a central component of such architectures, which fuses information from multiple context data points. So far, this aggregation operation has been treated separately from the inference of a latent representation of the target function in CLV models. Our key contribution is to combine these steps into one holistic mechanism by phrasing context aggregation as a Bayesian inference problem. The resulting Bayesian Aggregation (BA) mechanism enables principled handling of task ambiguity, which is key for efficiently processing context information. We demonstrate on a range of challenging experiments that BA consistently improves upon the performance of traditional mean aggregation while remaining computationally efficient and fully compatible with existing NP-based models.

Published as a conference paper at ICLR 2021
BAYESIAN CONTEXT AGGREGATION FOR NEURAL PROCESSES
Michael Volpp" 1. 1.2.
Fabian Flürenbrock
Lukas Grossherger! Christian Daniel'
Gerhard Neumann 2.3
Bossh Center for Artificial Intelligence, Renningen, Germany 2 "Karlsrhhe Institute of Technology, Karlsruhe, Germany University of Tübingen, Tübingen, Germany
ABSTRACT
Formulating scalable probabilisticc regression models with reliable uncertainty estimates has been a long-standing challenge in machine learning research. Re- cently, casting probabilistic regression multi-task learning problem in terms of conditional latent variable (CLV) models such as the Neural Process (NP) has shown promising results. In this paper, focus context aggregation, central component of such architectures, which fuses information from multiple context data points. So far, this aggregation operation has been treated separately from the inference of a latent representation of the target function in CLV models. Our key contribution is combine these steps into one holistic mechanism by phrasing context aggregation gregation (BA) mechanism enables principled handling of task ambiguity, which is key for efficiently processing context information. We demonstrate on range of challenging experiments that BA consistently improves upon the performance traditional mean aggregation while remaining computationally efficient and fully compatible with existing ba-ba models.
Bayesian inference problem. The resulting Bayesian Ag-
INTRODUCTION
Estimating statistical relationships between physical quantities from measured data of central importance in all branches of science and engineering and devising powerful regression models for this purpose forms a major field of study in statistics and machine learning. When judging representative power, neural networks (NNs) are arguably the most prominent member of the regression toolbox. NNs cope well with large amounts of training data and are computationally efficient at test time. On the downside, standard NN variants do not provide uncertainty estimates over their predictions and tend to overfit on small datasets. Gaussian processes (GPs) may be viewed as complementary to NNs as they provide reliable uncertainty estimates but their cubic (quadratic) scaling with the number of context data points at training (test) time in their basic formulation affects the application tasks with large amounts of data or on high-dimensional problems.
Recently, a lot of interest in the scientific community is drawn to combinations of aspects of NNs and GPs. Indeed, prominent formulation of probabilistic regression is as a multi-task learning problem formalized in terms of amortized inference in conditional latent variable (CLV) models, which results in NN-based architectures which learn a distribution over target functions. Notable variants are given by the Neural Process (NP) (Garnelo et al., 2018b) and the work of Gordon al. (2019), which presents a unifying view on range of related approaches the language of CLV models.
Inspired by this research, study context aggregation, central component of such models, and propose a new, fully Bayesian, aggregation mechanism for CLV-based probabilistic regression models.
"Correspondence Mi chael. vol volppêde. bosch.co....o

---

Published as a conference paper at ICLR 2021
To transform the information contained in the context data into latent representation of the target function, current approaches typically employ a mean aggregator and feed the output of this aggregator into a NN to predict distribution over global latent parameters of the function. Hence, aggregation and latent parameter inference have so far been treated as separate parts of the learning pipeline. Moreover, when using mean aggregator, every context sample is assumed to carry the same amount of information. Yet, in practice, different input locations have different task ambiguity and, therefore, samples should be assigned different importance the aggregation process. In contrast, our Bayesian aggregation mechanism treats context aggregation and latent parameter inference as one holistic mechanism, i.e., the aggregation directly yields the distribution over the latent parameters of the target function. Indeed, we formulate context aggregation as Bayesian inference of latent parameters using Gaussian conditioning in the latent space. Compared to existing methods, the resulting aggregator improves the handling of task ambiguity, as it can assign different variance levels to the context samples. This mechanism improves predictive performance, while it remains conceptually simple and introduces only negligible computational overhead. Moreover, our Bayesian aggregator can also be applied to deterministic model variants like the Conditional NP (CNP) (Gannelo et al.) 2018a).
In summary, our contributions are (i) a novel Bayesian Aggregation (BA) mechanism for context aggregation in NP-based models for probabilistic regression, (ii) its application to existing CLV architectures as well as deterministic variants like the CNP, and (iii) exhaustive experimental evaluation, demonstrating BA's superiority over traditional mean aggregation.
2 RELATED WORK
Prominent approaches to probabilistic regression are Bayesian linear regression and its kernelized counterpart, the Gaussian process (GP) (Rasmussen and Williams 2005). The formal correspondence of GPs with infinite-width Bayesian NNs (BNNs) has been established in Neal (996) and Williams
(999).. A broad range of research aims to overcome the cubic scaling behaviour of GPs with the number of context points, e.g., through sparse GP approximations (Smola and Bartlet 2001]
Lawrence et al.] 2002 Snelson and Ghahramani 2005 (Qunooero--aaneeee and Rasmussen 2005).
by deep kernel learning (Wilson et al. 2016), by approximating the posterior distribution of BNNs (MacKay 1992 Hinton and van Camp 1993; Gal and Ghahramani 2016, Louizos and Welling 2017). or, by adaptive Bayesian linear regression, i.e., by performing inference over the last layer of NN which introduces sparsity through linear combinations of finitely many learned basis functions alaaro-rrreiilla and Figueiras- Vidal, 2010; Hinton and Salakhutdinov 2008 Snoek et al. 2012 Calandra et al., 2016). An sense complementary approach aims to increase the data-efficiency of deep architectures by fully Bayesian treatment of hierarchical latent variable models ("DeepGiPs")
DDamianou and Lawrence 2013).
A parallel line of research studies probabilistic regression in the multi-task setting. Here, the goal is to formulate models which are data-efficient on an unseen target task by training them on data from set of related source tasks. BBrdenet et al. (2013); Yogatama and Mann (0014), and Golovin
ct al. (2017) study multi-task formulations of Ga-based models. More general approaches of this kind employ the meta-learning framework (Schmidhuber 1987 Thrun and Pratt 1998 Vilalta and Drissi 2005), where a model's 's training procedure is formulated in a way which incentivizes it to learn how to solve unseen tasks rapidly with only few context examples ("learning to learn". "few-shot learning" (Fei-Fei et al.) 2006 Lake et al. 2011)). A range of such methods trains a meta-learner to learn how to adjust the parameters of the learner's 's model (Bengio et al. 1991 [chhichhuerr] 1992). an approach which has recently been applied to few-shot image classification (Ravi and Larochelle 2017), or to learning data-efficient optimization algorithms (Hochreiter et al. 2001 Li and Malik
2016 Andrychowicz et al. 2016 Chen et al.) 2017 Perrone et al. 2018 Volpp et al. 2019). Other branches of meta-learning research aim to learn similarity metrics to determine the relevance of context samples for the target task (Koch et al., 20155 Vinyals et al., 22166] Snell et al. 2017 Sung cct al., 2017). ) or explore the application of eeemory-augmenter neural networks for meta-learning (Santoro al. 2016). Finn al. (2017) propose model-agnostic meta-learning (MAML), general framework for fast parameter adaptation in gradient- based learning methods.
A successful formulation of probabilistic regression as few-shot learning problem in multi-task setting is enabled by recent advances in the area of probabilistic learnin methods which allow a quantitative treatment of the uncertainty arising due to task ambiguity, feature particularly
2

---

Published as a conference paper at ICLR 2021
relevant for few-shot learning problems. One line of work specifically studies probabilistic extensions
of MAML (Grant et al. 2018: Ravi and Larochelle| 2017 Rusu et al.] 2018) Finn et al. 2018)
Kim al. 2018). Further important approaches
based amortized inference in multi-task
CLV models (Heskes) 2000 Bakker and Heskes 2003 Kingma and Welling 2013) Rezende et al.
2014 (Shnn 2013), which forms the basis of the Neural Statistician proposed by Edwards and
Storkey ( (017) and of the NP model family (Garnelo et al. 201 8b, Kim et al., 2019 Louizos et al.
2019). Gordon et al. (2019) present a unifying view on many the aforementioned probabilistic architectures. Building on the conditional NPs (CNPs) proposed by Garnelo et al. (2018a), range of NP-based architectures, such Garnelo et al. 2018b) and Kim et al. (2019), consider combinations of deterministic and CLV model architectures. Recently, Gordon et al. (2020) extended CNPs to
include translation equivariance the input space, yielding state-of-the-art predictive performance.
In this paper, we also employ formulation of probabilistic regression in terms of multi-task CLV model. However, while in previous work the context aggregation mechanism (Zaheer et al.. 2017 Wagstaff et al. 2019) was merely viewed as a necessity to consume context sets of variable size, we take inspiration from Becker et al. (2019) and emphasize the fundamental connection of latent parameter inference with context aggregation and, hence, base our model on novel Bayesian aggregation mechanism.
PRELIMINARIES
We present the standard multi-task CLV model which forms the basis for our discussion and present traditional mean context aggregation (MA) and the variational inference (VI) likelihood approximation
as employed by the NP model family (Garneoo al., 2018a Kim al. 2019), as well an alternative Monte Carlo (MC)-based approximation.
Problem Statement. We frame probabilistic regression as a multi-task learning problem. Let F denote family of functions fr Rd - RR4, with some form of shared statistical structure. We assume to have available data sets Dj {( of evaluations ye,t = fe(22,)) + from a subset functions ("tasks") {fe) C F with additive Gaussian noise = ~ N (0, o.). From this data, aim learn the posterior predictive distribution p (yel Ie, Dÿ) over a (set of) ye. given the corresponding (set of) inputs ze as well as a context set D C CD.
Ve.n
The Multi-Task CLV Model. We formalize the multi- task learning problem terms of a CLV model (Heskes) 2000) Gordon et al. 2019) as shown in Fig. ! The model employs task-specific global latent variables ze E RRd., as well as task-innependent latent variable 6, capturing the statistical structure shared between tasks. To learn 8, we split the data into context sets Dj {(x V/ )
Mr
Figure I: Multi-task CLV model with task-specific global latent variables 2; and task- iusk-independent variable 0 describing statistical structure shared between tasks.
and target sets Dj- = {(at m v ) and maximize the posterior predictive likelihood function
L II 8=1
/
M
P (vb.a PLIM,, D2.0) C=1 w.r.t. 0. In what follows, we omit task indices E to avoid clutter.
p D2.0) 0)
II m==
p (vê.ml Ze, ah 0) dae
(I)
Likelihood Approximation. Marginalizing over the task-specific latent variables = is intractable for reasonably complex models, so one has to employ some form of approximation. The NP-family of models (Garnelo al. 2018b Kim al. 2019) uses an approximation of the form
log p (y/m 2, x/, 8) + -log qs Dc
log p (vf.ar ai.M: D 0) 2 E Ego(=
Σ
go (=) D U Dt)
3

---

Published as a conference paper at ICLR 2021
Being derived using variational approach, this approximation utilizes an approximate posterior distribution ge D9) = p(( z|22 D 0). Note, however, that it does not constitute a proper evidence lower bound for the posterior predictive likelihood since the intractable latent posterior p [2] D .) has been replaced by 2o (=l DD) in the nominator of the rightmost term (L.e et al. 2018). An alternative approximation, employed for instance in Gordon et al. 2019). also replaces the intractable latent posterior distribution by an approximate distribution do DD) z p [2] D 0) and uses Monte-Carlo (MC) approximation of the resulting integral based on K latent samples, i.e.
log (Biaal Fi.Ms Dr, 0) log K log Σ
II P (y'ml 21: o 0)
24 ~ qs (z) (3)
Note that both approaches employ approximations qs (2) DC) of the latent posterior distribution (z] DD 0) and, as indicated by the notation, amortize inference in the sense that one single set of parameters is shared between all context data points. This enables efficient inference at test time, as no per-data-pon point optimization loops are required. As is standard in the literature (Garneo et al., 2018b Kim al. 2019), we represent 93 (D)) and p (v |2, n/m 0) by NNs and refer them as the encoder (enc, parameters Φ) and decoder (dec, parameters 0) networks, respectively. These networks set the means and variances of factorized Gaussian distributions, i.e..
9o (=) DF) N (=1 [12, diag (o2)), μ1= encu,, (D)), of enc (D"), (4) p. (m 2, mm 8) N (ym My. diag (oz) Hy dec (z, 2 oy dec (z, (m.am)) (5)
Context Aggregation. The latent variable = is global in the sense that depends on the whole context set Dc Therefore, some form of aggregation mechanism is required to enable the encoder to consume context sets D of variable size. To represent a meaningful operation on sets, such an aggregation mechanism has to be invariant to permutations of the context data points. /Zaheer et al. (2017) characterize possible aggregation mechanisms w.r.t. this permutation invariance condition, resulting in the structure of traditional aggregation mechanisms depicted in Fig. [(a) Each context data tuple (x vi)) is first mapped onto a latent observation rn. = enc,, (I yn) E IRd,, Then, permutation-invariant operation is applied to the set {ra | to obtain aggregated latent observation F. One prominent choice, employed for instance in Garnelo et al. 2018a), Kim et al. (2019), and Gordon et al. (201 is to take the mean, i.c..
n
Subsequently, F is mapped onto the parameters μl= and 0² of the approximate posterior distribution la ( = DD) using additional encoder networks, i.e., = enc,.. (F) and o² = enc,3,o (f). Note that three encoder networks are employed here: (i) enc,, to map from the context pairs Tn. (ii) enc to compute μæ from the aggregated mean and (iii) enc,g. a compute the variance of from F. In what follows, we refer to this aggregation mechanism as mean aggregation (MA) and to the networks enc,,.. and enc.g. collectively as "7-to-z-networks..
4 BAYESIAN CONTEXT AGGREGATION
We propose Bayesian Aggregation (BA), a novel context data aggregation technique for CLV models which avoids the detour via an aggregated latent observation and directly treats the object of interest, namely the latent variable z, as the aggregated quantity. This reflects a central observation for CLV models with global latent variables: context data aggregation and hidden parameter inference are fundamentally the same mechanism. Our key insight is to define probabilistic observation model p(r(z) for which depends on 2. Given a new latent observation r, = ( vr),, we can update p(z) by computing the posterior (z) (r,) p(r, (=)p(=)/p((()))))))))))))))) prr)).. Hence, by formulating context data aggregation as Bayesian inference problem, we aggregate the information contained in De directly into the statistical description of based on first principles.
4.1 BAYESIAN CONTEXT AGGREGATION VIA GAUSSIAN CONDITIONING
BA can easily be implemented using a factorized Gaussian observation model of the form (rr]lz) = N (ral diag((?))))) Tr enc,, (1 vh), of enc,2,6 (1 vi) (7)
4

---

Published as a conference paper at ICLR 2021
Ta
I'
2
Un
NV
N
(a) Traditional mean aggregation (MA).
(b) Our Bayesian aggregation (BA).
Figure Comparison of aggregation mechanisms in CLV models. Dashed lines correspond to learned components of the posterior approximation qu [z] D.). BA avoids the detour via a mean-aggregaaed latent observation /' and aggregates DD directly the statistical description of 2. This allows to
incorporate quantification of the information content of each context tuple (z vi)) as well as of
into the inference in a principled manner, while MA assigns the same weight to each context tuple.
Note that, in contrast to standard variational auto-encoders (VAEs) (Kingma and Welling 2013). we do not learn the mean and variance of a Gaussian distribution, but we learn the latent observation I' (which can be considered as a sample of p(z)) together with the variance o² of this observation. This architecture allows the application of Gaussian conditioning while this is difficult for VAEs. Indeed, we impose a factorized Gaussian prior po (z) A (
Ms,0; diag (o2.)) and arrive at Gaussian aggregation model which allows to derive the parameters of the posterior distribution do (2) [=|2) D*) in closed form
(cf. App. 7.1
o²
(o
B + (o%
Hs Ms.0 of ① Σ (n. - M2,0) @ (0%
Here 9, ④ and 2 denote element-wise inversion, product, and division, respectively, These equations naturally lend themselves to efficient incremental updates as new context data (if, yn) arrives by using the current posterior parameters Hz,old and of ood in place of the prior parameters, i.e.,
mmnew = (n)) (o22)) Nz = M + o z.ne @ (rn Mzzood @ (o².)
(9)
BA employs two encoder networks, enc,, and enc o. mapping context tuples to latent observations and their variances, respectively. In contrast to MA, it does not require 7-to-2-networks, because the set (rn J is aggregated directly into the statistical description of by means of Eq. (8). ef. Fig. 2(b Note that our factorization assumptions avoid the expensive matrix inversions that typically occur in Gaussian conditioning and which are difficult to backpropagate. Using factorized distributions renders BA cheap evaluate with only marginal computational overhead in comparison to MA. Furthermore, we can easily backpropagate through BA to compute gradients to optimize the parameters of the encoder and decoder networks. As the latent space 2 is shaped by the encoder network, the factorization assumptions are valid because the network will find space where these assumptions work well. Note further that BA represents a permutation-invarian operation on D
Discussion. BA includes MA special case. Indeed, Eq. g reduces to the mean aggerggggee latent observation Eq. (6) if we impose a non-informative prior and uniform observation variances o2 1P This observation sheds light on the benefits of Bayesian treatment of aggregation. MA assigns the same weight 1/N to each latent observation Tne independent of the amount of information contained in the corresponding context data tuple (xf, v.), well independent of the uncertainty about the current estimation of 2. Bayesian aggregation remedies both of these limitations: the influence of r. the parameters Mazold and o2 describing the current aggregated state is determined by the relative magnitude of the observation variance of and the latent variance
Note that an extended observation model of the form p (r.l [2]22) A' (ral= pr, diag(o".,)) with pr. given by third encoder output, does not lead to resulting posterior variances would stay unchanged and the posterior mean would read μ. = Ms,o + o3 o distinct encoder outputs computed >n (Fr- Mr. Mzz0) @ (o?.. ). Therefore, expressivity, which is why we set Pr., 0.
expressive aggregation mechanism. Indeed, the
would just subtract from the same inputs, resulting exactly the same
As motivated above, we consider r' as the aggregated quantity of MA and the distribution over described by (t and o², as the aggregated quantity of BA. Note that Fq. (8) does not necessarily generalize μt, and of after nonlinear 2-to-z-networks.
5

---

Published as a conference paper at ICLR 2021
22. cf. Eq. (9.. This emphasizes the central role of the learned observation variances of they allow to quantify the amount of information contained in each latent observation Tn- BA can therefore handle task ambiguity more efficiently than MA, the architecture can learn to assign little weight (by predicting high observation variances of ) to context points (I yi) located in areas with high task ambiguity, i.e., points which could have been generated by many of the functions in F. Conversely, in areas with little task ambiguity. i.e., if (xt Vin) contains a lot of information about the underlying function, BA can induce strong influence on the posterior latent distribution. In contrast, MA has to find ways to propagate such information through the aggregation mechanism by encoding it in the mean- aggn-gggge latent observation
4.2 LIKELIHOOD APPROXIMATION WITH BAYESIAN CONTEXT AGGREGATION
We show that BA is versatile in the sense that it can replace traditional MA in various CLV-based NP architectures as proposed, e.g., in Garnelo et al. (018b) and Gordon et al. 2019), which employ samples from the approximate latent posterior qa 2 DD) to approximate the likelihood (as discussed in Sec. 3), as well as in deterministic variants like the CNP (Garnelo et al. 2018a).
Sampling-Bare Likelihood Approximations. BA is naturally compatible with both the VI and MC likelihood approximations for CLV models. Indeed, BA defines a Gaussian latent distribution from which we can easily obtain samples in order to evaluate Eq. or Eq. using the decoder parametrization Eq.
Bayesian Context A Aggregation for Conditional Neural Processes. BA motivates a novel, alter- native, method to approximate the posterior predictive likelihood Eq. ① resulting in deterministic loss function which can be efficiently optimized for 0 and o in an end-to-end fashion. To this end, we employ Gaussian approximation of the posterior predictive likelihood of the form
p (vieM) Zf.M: D=.0) 8) = N (viaal Mg. E.y)
(10)
This is inspired by GPs which also define a Gaussian likelihood. Maximizing this expression yields the optimal solution By = Buy Ey S,y. with By and Σ,, being the first and second moments of the true posterior predictive distribution. This is a well-known result known as moment matching, a popular variant of deterministic approximate inference used, e.g., in Deisenroth and Rasmussen (2011) and Becker et al. (2019) By and Σ, are functions of the moments μtx and o of the latent posterior p 22) D", which motivates the following decoder parametrization:
My dec e (Ma, x'))
o deco3.0 0 (f,, My diag (o
(1I)
Here, μı= and o are given by the BA Eqs. (8) Note that we define the Gaussian approximation to be factorized w.r.t. individual vin an assumption which simplifies the architecture but could be dropped if a more expressive model was required. This decoder can be interpreted as "moment matching network". computing the moments of y given the moments of z. Indeed, in contrast to decoder networks of CLV-based NP architectures as defined in Eq. ( it operates on the moments μa and of of the latent distribution instead of samples which allows to evaluate this approximation in a deterministic manner. In this sense, the resulting model is akin to the CNP which defines a deterministic, conditional model with decoder operating on the mean-aggregated latent observation F. However, BA-based models trained in this deterministic manner still benefit from BA's ability to accurately quantify latent parameter uncertainty which yields significantly improved predictive likelihoods. In what follows, we refer to this approximation scheme as direct parameter-based (PB) likelihood optimization.
Discussion. The concrete choice of likelihood approximation equivalently, model architecture depends mainly on the intended use-case, Sampling- models are generally more expressive as they can represent complex, i.e., structured, non-Gaussian, posterior predictive distributions. Moreover, they yield true function samples while deterministic models only allow approximate function samples through auto-regressive (AR) sampling schemes. Nevertheless, deterministic models exhibit several computational advantages. They yield direct probabilistic predictions in a single forward pass, while the predictions of sampling-based methods only defined through averages over multiple function samples and hence require multiple forward passes. Likewise, evaluating the MC-based likelihood approximation Eq. 3 during training requires draw multiple

---

Published as a conference paper at ICLR 2021
Table 1: Posterior predictive log- likelihhod functions drawn from GP priors with RBF, weakly periodic, and Matern-5/2 kernels, averaged over context sets with N E {0, 1......., 64) points (table) and in dependence of N (figure). BA consistently outperforms MA, independent of the likelihood approximation, with MC being the most expressive choice. PB represents an efficient, deterministic alternative, while the VI approximation tends to perform worst, in particular for small N.
PB/det
VI
MC
ANP MA Atlention
BA
MA (CNP)
(LP-NP)
RBF GP
1.37 = 0.15
0.91
1.40 = 0.04
0.45 0.12
1.62 2: 0.05
1.07 =07=00 0.05
0.98 + 0.02
Weakly Periodic
0.76 0.02
0.89 0.03
0.07
1,30
0.85
0.04
1.02
Materm-5/2 GP
-0.50 1: 0.07
0.68 4: 0.01
-0.79 -: 0.01
-1.09 2 0.10
-0.33 = 0.01
-0.90 2 0.15
0.25 1: 0.02
Weskly Periodic
Maten 5/2
1op pG we .P)
Aggreaatoo
-0.5
BA
Loss
PRiles
-1.5
Vi-imspined
MC
25
0 25
50
0
25
50
N
N
(K) latent samples z. While the VI likelihood approximation Eq. (2) can be optimized on a single function sample per training step through stochastic gradient descent (Bishop) (006), it has the disadvantage that it requires to feed target sets DD' through the encoder which can impede the training for small context sets D as discussed in detail in App. 7.2
5 EXPERIMENTS
We present experiments to compare the performances of BA and of MA in NP-based models. To pro- vide a complete picture. we evaluate all combinations of likelihood approximations (PB/deterministic Eq. (1O), VI Eq. (2). MC Eq.
(3) and aggregation methods (BA Eq. 18. MA Eq. (6), resulting in six different model architectures, ef. Fig. 4 App. 7.5.2] Two of these architectures correspond to existing members of the NP family: MA + deterministic is equivalent the CNP (Gamelo et al., 2018a), and MA + VI corresponds to the Latent-Path NP (LP-NP) (Garnelo et al., 2018b), i.c., the NP without deterministic path. We further evaluate the Attentive Neural Process (ANP) (Kim et al.) 2019), which employs hybrid approach, combining LP-NP with cross-attentic a parallel deterministic path as well as an -P-architecture using MA with self-attentive (SA) encoder network. Note that BA can also be used in hybrid models like ANP or in combination with 7.4| we discuss NP-based regression in relation to
mechanism in
SA, an idea we leave for future research. App. 7.4 other methods for (scalable) probabilistic regression.
The performance of Na-based models depends heavily on the encoder and decoder network archi- tectures as well the latent space dimensionality dz. To assess the influence of the aggregation mechanism independently from all other confounding factors, we consistently optimize the encoder and decoder network architectures. the latent-space dimensionality dz, well as the learning rate of the Adam optimizer (Kingma and Ba. 2015), independently for all model architectures and for all experiments using the Optuna (Akiba | 2019) framework, ef. App. 7.5.3 If not stated differently, we report performance in terms of the mean posterior predictive likelikld over 256 test tasks with 256 data points each, conditioned on context sets containing N E (0, ,,..., } data points methods (VI, MC, ANP), we report the joint log- likelioed over
(cf. App. 7.5.4). For sampling-bast the test sets using a Monte-Carlo approximation with 25 latent samples, cf. App. 77..4 We average the resulting log- likeehodd values over 10 training runs with different random seeds and report 95% confidence intervals. We publish source code to reproduce the experimental results online
GP Samples. We evaluate the architectures on synthetic functions drawn from GP priors with different kernels (RBF, weakly periodic, Matern-5/2). as proposed by Gordon et al. (2020). cf. App. 7.5.I We generate a new batch of functions for each training epoch. The results
(Tab. I show that BA consistently outperforms MA, independent of the model architecture. In-
code from http:///tt.tatttetttao.eee/////r https: http::/www...........aaaaaaaaoa.
For ANP, we use
com

---

Published as a conference paper at ICLR 2021
Table 3: Posterior predictive log- lellikelihood on ID and 3D quadratic functions with limited numbers of training tasks, averaged over context sets with E {0, ,,...,20) data points. BA outperforms MA by considerable margins in this regime of little training data.
PBAde
VI
MC
ANP MA Attention
BA
MA (CNP)
MA (LP-NP)
BA
MA
Qundratic ID. 64
1.42 1 0.20
0.47 0.25
1.48 0.05
-0.32 = 0.55
1.71 E 0.23
1.27 3: 0.06
0.69
Quadraric 3D, L 128
2.46 + 0.12
2.73 0.10
2.53 + 0.07
3.45 + 0.12
1.79 + 0.07
2.14 + 0:05
3.08 + D.02
NA+
MM-dot
RAME
MA+iT
MA+MCC
AMP
Figure 3: Predictions two instances (dashed lines) of the ID quadratic function class, given N = 3 context data points (circles). We show mean and standard deviation predictions (solid line, shaded area), and 10 function samples (AR samples for deterministic methods). Cf. also App. 7.6
terestingly, despite employing factorized Gaussian approximation, our deterministic PB approx- imation performs at least on-par with the traditional VI approximation which tends to perform particularly poorly for small context sets, reflecting the intricacies discussed Sec. 4.2 As expected, the MC approximation yields the best results in terms of pre- dictive performance, as is more expressive than the deterministic approaches and does not share the prob- lems of the VI approach. As shown in Tab, and Tab. | App. 7.6 our proposed PB likelihood approximation much cheaper to evaluate compared to both sampling based approaches which require multiple forward passes per prediction. We further observe that BA tends to require smaller encoder and decoder networks as it is more efficient at propagating context information to the latent state as discussed Sec. [4.] The hybrid ANP approach competitive only on the Matern-5/2 function class. Yet, we refer the reader to Tab. T0 App. (7.6. demonstrating that the attention mechanism greatly improves performance in terms of MSE.
Table 2: Relative evaluation runtimes and #parameters of the optimized network ar- chitectures on RBF GP. Also cf. Tab.
PHoter
VI
BA
MA (CNP)
BA
MA (LP-N)
BA
MA.
Runnime +
1.4
18
25
12
27
Parameters
73k.
96L
63L
776
122
153k
Quadratic Functions. We further seek to study the performance of BA with very limited amounts of training data. To this end, we consider two quadratic function classes, each parametrized by three real parameters from which we generate limited numbers I. of training tasks. The first function class defined on one-dimensional domain, i.e.. E R, and we choose L 64, while the second function class, as proposed by Perrone et al. 2018). defined on E R with L = 128, cf. App. 7.5.1 As shown in Tab. 3 BA again consistently outperforms MA, often by considerably large margins, underlining the efficiency of our Bayesian approach to aggregation in the regime of little training data. On the ID task, all likelihood approximations perform approximately on-par in combination with BA, while MC outperforms both on the more complex 3D task. Fig. 3 comparee prediction qualities.
Dynamics a Furuta Pendulum. We study BA on realistic dataset given by the simulated dynamics of rotary inverted pendulum, better known the Furuta pendulum (Furuta et al. 1992), which is a highly non-linear dynamical system, consisting of an actuated arm rotating in the horizontal plane with attached pendulum rotating freely in the vertical plane, parametrized by two masses, three lengths, and two damping constants. The regression task is defined as the one-step-ahead prediction of the four-dimensional system state with a step-size of At 0.1 s, as detailed in App. [ The results (Tab. 4 show that BA improves predictive performance also on complex, non-synthetic regression tasks with nigher-dimension with our previous findings regarding the likelihood approximations, with MC being strongest in terms of predictive likelihood, followed by our efficient deterministic alternative PB.
input- and output spaces. Further, they are consistent
2D Image Completion. We consider 2D image completion experiment where the inputs are pixel locations images showing handwritten digits, and we regress onto the corresponding pixel intensities cf. App. [7. Interestingly, we found that architectures without deterministic paths were not able to solve this task reliably which is why we only report results for deterministic models.
8

---

Published as a conference paper at ICLR 2021
Table 4: Posterior predictive log- legllikelioood the dynamics of Furuta pendulum, averaged over context sets with N E {0, 1,...,200 20) state transitions. BA performs favorably on this weal-world task.
PB/da.
MC
ANP Attention
BA
MA (CNP)
MA (LP-NP)
MA
Furuta Dynamics 7.50 + 0.27
7.06 0.12
7.32 0.18
5.57 + 0.21
8.25 + 0.33
7.35 4: 0.24
4.74 0.16
Table 6: Comparison of the posterior predictive log-likelihooe lileliiieiood of our BA with traditional MA, combined with a self-attention (SA) mechanism in the encoder (BA does not use SA mechanism), using the PB and MC likelihood approximations. We provide results for Laplace SA (L-SA), dot- product SA (DP-SA), and mulihead SA (MH-SA) and repeat the results for BA and MA without SA ("no SA"). While L-SA and DP-SA do not increase predictive performance compared to MA without SA, MH-SA results in significant improvements. Nevertheless, vanilla BA still performs better or at least on par,par while being computationally more efficient.
BA SA
MA PM
HA MC to
MA
1-5A
XII-SA
SA
1-5A
DP-SA
MII-A
Weakly Partedia GP
1.13 0.08
0.25 0.003
0.38 00007
0.21 0.00
1.13
1,50
0.01
00:0
0.87 0113
1.29 0..03
Matere-A -22----- E2"
-0.50 0.077
-0.05 0.01
-1.05-.00 8.000
-0.76 1.000
-0.61 1..0
-0.33 1.000
-0.90 0.033
-040 0002
-0.45 1.0.0
-0.39 (110
Ocoabatic
1.42 0.20
0:47 00:35
0:15 10337
0.17 0.24
1:09 0:31
1.71 0:23
1:27 1.05
0.0
1.32 0.31
1.60 0.012
Fonen 7:50 0:27
7.05 00.32
0.13
7.030
7.00 0:446
1.211
7.00
1..3
7.07 0.14
8.39 1:220
As shown in Tab. 5 BA improves performance in comparison to MA by large margin. This highlights that BA's ability to quantify the information content of a context tuple is particularly beneficial on this task, as, e.g., pixels the middle area of the images typically convey more information about the identity of the digit than pixels located near the borders.
Table 5: Predictive log-likelihood on 2D image completion task on MNIST, averaged over N E (0, 1, 392] context pixels.
PRider
ANP
BA
MA (CNP)
MA + Attemion
2D Image Complction
2.75 a 0.20
2.05 = 0.36
1.02 A 0000
Self-attentive Encoders. Another interesting baseline for BA is MA. combined with self-attention (SA) mechanism in the encoder. Indeed, similar to BA, SA yields non-uniform weights for the latent observations ru, where given weight computed from some form of pairwise spatial relationship with all other latent observations in the context set (cf. App. 7.3 for a detailed discussion). As BA's weight for Tn only depends on (Ii, Mn) itself, BA is computationally more efficient: SA scales like (NN2) in the number N of context tuples while BA scales like ONN), and, furthermore, SA does not allow for efficient incremental updates while this is possible for BA, cf. Eq. (9). Tab, 6 shows comparison of BA with MA in combination with various different SA mechanisms in the encoder. We emphasize that we compare against BA in its vanilla form, i.c., BA does not use SA in the encoder. The results show that Laplace SA and dot-product SA do not improve predictive performance compared to vanilla MA, while multihead SA yields significantly better results. Nevertheless, vanilla BA still performs better at least on-par and is computationally more efficient. While being out of the scope of this work, according to these results, a combination of BA with SA seems promising if computational disadvantages can be accepted in favour of increased predictive performance, cf. App.
6 CONCLUSION AND OUTLOOK
We proposed a novel Bayesian Aggregation (BA) method for NP-based models, combining context aggregation and hidden parameter inference in one holistic mechanism which enables efficient handling of task ambiguity. BA is conceptually simple, compatible with existing NP-based model architectures, and consistently improves performance compared to traditional mean aggregation. It introduces only marginal computational overhead, simplifies the architectures in comparison to existing CLV models (no n-to-k-nttworks). and tends to require less complex encoder and decoder network architectures. Our experiments further demonstrate that the VI likelihood approximation traditionally used to train NP-based models should be abandoned favor a MC-based approach, and that our proposed PB likelihood approximation represents efficient deterministic alternative with strong predictive performance. We believe that a range of existing models, e.g., the ANP or NPs with self-attentive encoders, can benefit from BA, especially when reliable quantification of uncertainty is crucial. Also, more complex Bayesian aggregation models are conceivable, opening interesting avenues for future research.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGMENTS
We thank Philipp Becker, Stefan Falkner, and the anonymous reviewers for valuable remarks and discussions which greatly improved this paper.
REFERENCES
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A Next-generation Hyperparameter Optimization Framework. 2019.
Marcin Andrychowicz, Misha Denil, Sergio Gomez Colmenarejo, Matthew W. Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning to Learn by Gradient Descent by Gradient Descent. Advances in Neural Information Processing Systems, 2016.
Bart Bakker and Tom Heskes. Task Clustering and Gating for Bayesian Multitask Learning. Journal of Machine Learning Research, 2003.
Rémi Bardenet, Mátyás Brendel, Balázs Kégl, and Michèle Sebag. Collaborative Hyperparameter Tuning. International Conference on Machine Learning. 2013.
Philipp Becker, Harit Pandya, Gregor H. W. Gebhardt, Cheng Zhao, C. James Taylor, and Gerhard Neumann. Recurrent Kalman Networks: Factorized Inference in Spaces. International Conference Machine Learning, 2019.
Deep Feature
Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning Synaptic Learning Rule. International Joint Conference on Neural Networks, 1991.
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
R. Calandra, J. Peters, C. E. Rasmussen, and M. P. Deisenroth. Manifold Gaussian Processes for Regression. International Joint Conference on Neural Networks, 2016.
Benjamin Seth Cazzolato and Zebb Prime. On the Dynamics of the Furuta Pendulum. Journal of Control Science and Engineering, 2011.
Yutian Chen, Matthew W. Hoffman, Sergio Gómez Colmenarejo, Misha Denil, Timothy P. Lillicrap, Matt Botvinick, and Nando de Freitas. Learning Learn without Gradient Descent by Gradient Descent. International Conference on Machine Learning, 2017.
Andreas Damianou and Neil Lawrence. Deep Gaussian Processes. International Conference Artificial Intelligence and Statistics, 2013.
MP. Deisenroth and CE. Rasmussen. PILCO: A Model-Based and Data-Efficient Approach to Policy Search. International Conference on Machine Learning, 2011.
Harrison A. Edwards and Amos J. Storkey. Towards a Neural Statistician. International Conference on Learning Representations, 2017.
Li Fei-Fei, R. Fergus, and P. Perona. One-shot Learning of Object Categories. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2006.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model- Aggoo-Ag Meta-Learning for Fast Adaptation of Deep Networks. International Conference on Machine Learning, 2017.
Chelsea Finn, Kelvin Xu, and Sergey Levine. Probabilistic Model-Agnostic Meta-Leaming, Advances in Neural Information Processing Systems, 2018.
K. Furuta, M. Yamakita, and S. Kobayashi. Swing-up Control of Inverted Pendulum Using Pseudo- State Feedhack. Journal of Systems and Control Engineering, 1992.
Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. International Conference on Machine Learning, 2016.
10

---

Published as a conference paper at ICLR 2021
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Tch, Danilo Rezende, and M. Ali Eslami. Conditional Neural Processes. International Conference Machine Learning, 2018a.
Marta Gamelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo Jimenez Rezende, S. M. Ali Eslami, and Yee Whye Teh. Neural Processes. arXiv abs/1807.01622, 2018b.
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 Representations, 2019.
Jonathan Gordon, Wessel P. Bruinsma, Andrew Y. K. Foong, James Requeima, Yann Dubois, and Richard F. Turner. Convolutional Conditional Neural Processes. International Conference on Learning Representations, 2020.
Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas L. Griffiths. Recasting Gradient-Based Meta-Learning as Hierarchical Bayes. International Conference on Learning Representations, 2018.
Tom Heskes. Empirical Bayes for Learning to Learn International Conference Machine Learning, 2000.
Geoffrey E. Hinton and Russ R. Salakhutdinov. Using Deep Belief Nets to Learn Covariance Kernels for Gaussian Processes. Advances in Neural Information Processing Systems, 2008
Geoffrey E. Hinton and Drew van Camp. Keeping the Neural Networks Simple by Minimizing the Description ength of the Weights. Annual Conference on Computational Learning Theory, 1993.
Sepp Hochreiter, A. Steven Younger, and Peter R. Conwell. Learning to Learn Using Gradient Descent. International Conference on Artificial Neural Networks, 2001.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, S. M. Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive Neural Processes. International Conference on Learning Representations, 2019.
Taesup Kim, Jaesik Yoon, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn. Bayesian Model-Agnosti Meta-Learning, Advances in Neural Information Processing Systems, 2018.
Diederik P. Kingma and Jimmy Ba. Adam: Method for Stochastic Optimization. International Conference on Learning Representations, 2015.
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. International Conference on Learning Representations, 2013.
Gregory Koch, Richard Zemel, and Ruslan Salakhutdinov. Siamese Neural Networks for One-shot
Image Recognition. International Conference on Machine Learning, 2015.
Brenden M. Lake, Ruslan Salakhutdinov, Jason Gross, and Joshua B. Tenenbaum. One Shot Learning of Simple Visual Concepts. Cognitive Science, 2011.
Neil Lawrence, Matthias Seeger, and Ralf Herbrich. Fast Sparse Gaussian Process Methods: The Informative Vector Machine. Advances in Neural Information Processing Systems, 2002.
M. Lazaro-Gredde and A. R. Figueiras- Vidal. Marginalized Neural Network Mixtures for Large- Scale Regression, JEEE Transactions on Neural Networks, 2010.
Tuan Anh Le, Hyunjik Kim, and Marta Gamelo. Empirical Evaluation of Neural Process Objectives. Third Workshop Bayesian Deep Learning, 2018.
Yann LeCun and Corinna Cortes. MNIST Handwritten Digit Database. 2010.

---

Published as a conference paper at ICLR 2021
Ke Li and Jitendra Malik. Learning to Optimize. International Conference Machine Learning. 2016.
Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: Novel Bandit-Based Approach to Hyperparameter Optimization. Journal of Machine Learning Research, 2017.
Christos Louizos and Max Welling. Multiplicative Normalizing Flows for Variational Bayesian Neural Networks. International Conference Machine Learning. 2017.
Christos Louizos, Xiahan Shi, Klamer Schutte, and M. Welling. The Functional Neural Process. Advances in Neural Information Processing Systems, 2019.
David C. MacKay. A Practical Bayesian Framework for Backpropagation Networks. Neural Computation, 1992.
Radford M. Neal. Bayesian Learning for Neural Networks. Springer-
Springer-Verlag. Verlag, 1996.
Valerio Perrone, Rodolphe Jenatton, Matthis Seeger, and Cedric Archambeau. Scalable Hyperpa- rameter Transfer Learning. Advances in Neural Information Processing Systems, 2018.
Joaquin Quinoneeo-Candelaa and Carl Edward Rasmussen. A Unifying View of Sparse Approximate Gaussian Process Regression. Journal of Machine Learning Research, 2005.
Carl Edward Rasmussen and Christopher K. Williams. Gaussian Processes for Machine Learning. The MIT Press, 2005.
Sachin Ravi and Hugo Larochelle. Optimization as Model for Few-Shot Learning. International Conference on 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.
Andrei A. Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osindero, and Raia Hadsell. Meta-Learning with Latent Embedding Optimization. International Conference on Learning Representations, 2018.
Adam Santoro, Sergey Bartunov, Matthew M. Botvinick, Daan Wierstra, and Timothy P. Lillicrap. Meta-Learming with Memory-Augmented Neural Networks. International Conference on Machine Learning, 2016.
Jürgen Schmidhuber. Evolutionary Principles in Self- Referentiil Learning. On Learning how to Learn: The Meeta-Meta-Meta..-HHlee Diploma Thesis, Technische Universitat München, Germany, 1987.
Jürgen Schmidhuber. Learning to Control Fast-Weight Memories: An Alternative to Dynamic Recurrent Networks. Neural Computation, 1992.
Alex J. Smola and Peter I., Bartlett. Sparse Greedy Gaussian Process Regression. Advances in Neural Information Processing Systems, 2001.
Jake Snell, Kevin Swersky, and Richard S. 7emel. Prototypical Networks for Few-shot Learning. Advances in Neural Information Processing Systems, 2017.
Edward Snelson and Zoubin Ghahramani. Sparse Gaussian Processes Using Pseudo-Inputs. Interna- tional Conference on Neural Information Processing Systems, 2005.
Jasper Snoek, Hugo I Larochelle, and Ryan P. Adams. Practical Bayesian Optimization of Machine Learning Algorithms. Advances in Neural Information Processing Systems, 2012.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning Structured Output Representation using
Deep Conditional Generative Models. Advances Neural Information Processing Systems, 2015.
12

---

Published as a conference paper at ICLR 2021
Flood Sung, Yongxin Yang. Li Zhang, Tao Xiang, Philip H. S. Torr, and Timothy M. Hospedales. Learning to Compare: Relation Network for Few-Shot Learning. Conference on Computer Vision and Pattern Recognition, 2017.
Scbastian Thrun and Lorien Pratt. Learning to Learn. Kluwer Academic Publishers, 1998.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. Advances in Neural Information Processing Systems, 2017.
Ricardo Vilalta and Youssef Drissi. A Perspective View and Survey 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 Optimization. International Conference Learning Representations, 2019.
Edward Wagstaff, Fabian B. Fuchs. Martin Engelcke, Ingmar Posner, and Michael A. Osbore.. On the Limitations of Representing Functions on Sets. International Conference on Machine Learning, 2019.
Christopher K. I. Williams. Computing with Infinite Networks. Advances Neural Information Processing Systems, 1996.
Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P. Xing. Deep Kernel Learning. International Conference on Artificial Intelligence and Statistics, 2016.
Dani Yogatama and Gideon Mann. Efficient Transfer Learning Method for Automatic Hyperparameter Tuning. International Conference Artificial Intelligence and Statistics, 2014.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alexander J. Smola. Deep Sets, Advances in Neural Information Processing Systems, 2017.
13

---

Published as a conference paper at ICLR 2021
7 APPENDIX
We present the derivation of the Bayesian aggregation update equations (Eqs. 8 detail. To foster reproducibility, we describe all experimental settings well the hyperparameter optimization procedure used to obtain the results reported in Sec. and publish the source code We further provide additional experimental results and visualizations of the predictions of
( in more
online & the compared architectures.
7.1 DERIVATION OF THE BAYESIAN AGGREGATION UPDATE EQUATIONS
We derive the full Bayesian aggregation update equations without making any factorization assump- tions. We start from Gaussian observation model of the form
p (ral I 2) N (rxl z, Mr.), In = encr.d (2 vū), Er, = ence, (z Bi))
(12)
where In and Σr, are learned by the encoder network. If we impose Gaussian prior in the latent space, ie.,
p (2) = N (#) Ms,0; Σz,0),
(13)
we arrive Gaussian aggregation model which allows to derive the parameters of the posterior distribution, i.e., of
do (= D = (z)
(14)
in closed form using standard Gaussian conditioning (Bishop 2006
N
Σg (E2.0)
+
(15a)
n=
Ne = 12,0 M Eg Σ (2rr,) (rn M2,0)
(15b)
As the latent space is shaped by the encoder network, it will find space where the following factorization assumptions work well (given dg is large enough);
Er, diag (o?,), of enco?.o (2 vi), Σg,0 diag (02.0)
(16)
This yields a factorized posterior, i.e..
(=1 = N (=1 diag (o²)),
(17)
with
N
0 =
(02,0)
e
(18a)
N
(4 = P2.0 + o²
(rn (2,0) 2 (o?.).
(18b)
n=1
Here O and @ denote element-wise inversion, product, and division, respectively. This is the result Eq. (8) from the main part this paper.
7.2 DISCUSSION OF LIKELIHOOD APPROXIMATION
To highlight the limitations of the VI approximation, note that decoder networks of models employing the PB or the MC likelihood approximation are provided with the same context information training and test time: the latent variable (which is passed samples z (for MC) or in the form of parameters Hs, of describing the latent distribution (for PB)) is in both cases conditioned only on the context D In contrast, the variational approximation Eq. (2), the expectation is w.r.t. conditioned on the union of the context set D and the target set D As Dt is not available at test time, this introduces mismatch between how the model trained
the decoder in the form of latent
14

---

Published as a conference paper at ICLR 2021
and how it is used at test time. Indeed, the decoder is trained on samples from 96 (2) D U D*) but evaluated on samples from qe (z) D*). This is not a serious problem when the model is evaluated on context sets with sizes large enough to allow accurate approximations of the true latent posterior distribution. Small context sets, however, usually contain too little information to infer 2 reliably. Consequently, the distributions do (=l DD) and 2o (=l De U 'D') typically differ significantly in this regime. Hence, incentivizing the decoder to yield meaningful predictions on small context sets requires intricate and potentially expensive additional sampling procedures choose suitable target sets D' during training. As a corner case, we point out that it is not possible to train the decoder on samples from the latent prior, because the right hand side of Eq. vanishes for D D Ø.
7.3 SELF-ATTENTIEE ENCODER ARCHITECTURES
Kim et al. (2099) propose to use attention-mechanisms to improve the quality of NP-based regression. In general, given a set of key-value pairs {(In: yn)] I Zn E R Vn E R and query x² E Rd an attention mechanism A produces a weighted sum of the values, with the weights being computed from the keys and the query:
N w (Inn x²) yn-
A Mn)
(19)
There are several types of attention mechanisms proposed in the literature (Vaswani et al., 2017).
each defining specific form of the weights. Laplace atlention adjusts the weights according to the spatial distance of keys and query:
Wt. (In x") x exp (-1lamn - x x²Il1).
(20)
Similarly, dot- dot-product attention computes
WDP (In, x exp ( (V/x//Vd,)))
(21)
more complex mechanism multihead attention, which employs a set of 3H learned linear mappings { | k- (CX)" {cx] A==1" where H is a hyperparameter. For each h, these mappings are applied to keys, values, and queries, respectively. Subsequently, dot-product attention is applied to the set of transformed key-value pairs and the transformed query. The resulting H values are then again combined by further learned linear mapping C to obtain the final result.
Self-attention (SA) is defined by setting the of queries equal the set of keys. Therefore, SA produces again a set of N weighted values. Combining SA with an NP-encoder, i.e., applying SA to the set (fx , '' 1== of inputs In and corresponding latent observations r, (where we also consider a possible nonlinear transformation fa of the inputs) and subsequently applying MA yields an interesting baseline for our proposed BA. Indeed, similar to BA, SA computes weighted sum of the latent observations Tn. Note, however, that SA weighs each latent observation according to some form of spatial relationship of the corresponding input with all other latent observations in the context set. In contrast, BA's weight for a given latent observation is based only on features computed from the context tuple corresponding to this very latent observation and allows incorporate an estimation of the amount of information contained in the context tuple into the aggregation (ef. Sec. This leads to several computational advantages of BA over SA: (i) SA scales quadratically in the number N of context tuples, as it has to be evaluated on all 2 pairs of context tuples. In contrast, BA scales linearly with N. (ii) BA allows for efficient incremental updates when context data arrives sequentially (cf. Eq. 9 while using SA does not provide this possibility: it requires to store and encode the whole context set D at once and to subsequently aggregate the whole set of resulting (SA-weighted) latent observations.
The results in Tab. 6 Sec. showw that multihead SA leads to significant improvements in predictive performance compared to vanilla MA. Therefore, a combination of BA with self-attentive encoders seems promising in situations where computational disadvantages can be accepted in favour of increased predictive performance. Note that BA relies on a second encoder output of, (in addition to the latent observation r.)) which assesses the information content in each context tuple (Ins Bn). As each SA- rn is informed by the other latent observations in the context set, obviously, one would have to also process the set of o² in a manner consistent with the SA-weighting. We leave such combination of SA and BA for future research.
15

---

Published as a conference paper at ICLR 2021
Table 7: Comparison of the predictive log-likelihood of NP-based architectures with two simple GP- based baselines. (i) Vanilla GP (optimizes the hyperparameters individually on each target task and ignores the source data) (ii) Multi-task GP (optimizes one set of hyperparameters on all source tasks and uses them without further adaptation on the target tasks). Both GP implementations use RBF. kernels. As in the main text, we average performance over context sets with sizes N E {O.....64] for RBF GP and N E {0, .... 203 for the other experiments. Multi-task GP constitutes the optimal model (assuming it fits the hyperparameters perfectly) for the RBF GP experiment, which explains its superior performance. On the Quadratic 1D experiment, Multi-task GP still performs better than the other methods as this function class shows a relatively low degree of variability, In contrast, on more complex experiments like Quadratic 3D and the Furuta dynamics, none of the GP variants is able to produce meaningful results given the small budget of at most 20 context points, while NP-based methods produce predictions of high quality they incorporate the source data more efficiently.
NPs with MCsloss
GP
BA
MA
Vanilla
Multi-task
RBF GP
1.62 + 0.05
1.07 = 0.05
1.96
2.99
Quadratic ID, L 64
1.71
1.27 0.06
-1.56
2.11
Quadratic 3D, T. - 128
-1.79 0.07
+.14
-472.76
-173.78
Furuta Dynamics
8.25 + 0.33
7.55 0.24
-6.16
-2.47
7.4 NEURAL PROCESS-BASED MODELS IN THE CONTEXT OF SCALABLE PROBABILISTIC REGRESSION
We discuss in more detail how NP-based models relate to other existing methods for (scalable) probabilistic regression, such (multi-task) GPs (Rasmussen and Williams 2005; Bardenet et al.)
2013 Yogatama and Mann 2014
Golovin et al. 2017). Bayesian neural networks (BNNs) (MacKay Gal and Ghabramani 2016). and DeepGPs (Damianou and Lawrence 2013).
1992
NPs are motivated in Garnelo et al. 22018a b). Kim et al. (011)), as well as in our Sec. as models which combine the computational efficiency of neural networks with well-calibrated uncertainty estimates (like those of GPs). Indeed, NPs scale linearly in the number N of context and M of target data points, i.e., like O(N + M). while GPs scale like O(N3 + M²). Furthermore, NPs are shown to exhibit well-calibrated uncertainty estimates. In this sense, NPs can be counted members of the family of scalable probabilistic regression methods.
A central aspect of NP training which distinguishes NPs from range of standard methods is that they are trained in a multi-task fashion (cf. Sec. B). This means that NPs rely on data from a set of related source tasks from which they automatically learn powerful priors and the ability to adapt quickly to unseen target tasks. This multi-task training procedure of NPs scales linearly in the number L of source tasks, which makes it possible train these architectures large amounts of source data. Applying GPs in such a multi-task setting can be challenging, especially for large numbers of source tasks. Similarly, BNNs as well DeepGPs are in their vanilla forms specifically designed for the single-task setting. Therefore, GPs, BNNs, and DeepGPss are not directly applicable in the NP multi-task setting, which is why they are typically not considered as baselines for Na-based models, as discussed in (Kim et al. 2019).
The experiments presented in Garnelo al. 201 Sa and Kim al. 2019) focus mainly on evaluating NPs in the context of few-shot probabilistic regression, i.e., on demonstrating the data- efficiency of NPs on the target task after training on data from a range of source tasks. In contrast, the application of NPs in situations with large (> 1000) numbers of context/target points per task has to the best of our knowledge not yet been investigated in detail in the literature. Furthermore, it has not been studied how to apply NPs in situations where only single very few source tasks are available. The focus of our paper is a clear-cut comparison of the performance of our BA with traditional MA in the context of NP-based models. Therefore, we also consider experiments similar to those presented in (Garnelo et al. 2018a] 1b] Kim et al., 2019) and leave further comparisons with existing methods for (multi-task) probabilistic regressions for future work.
Nevertheless, to illustrate this discussion, we provide two simple GP-based baseline methods: (i) a vanilla GP, which optimizes the hyperparameters on each target task individually and does not use
16

---

Published as a conference paper at ICLR 2021
the source data, and (ii) naive but easily interpretable example of multi-task GP, which optimizes one set of hyperparameters on all source tasks and uses it for predictions on the target tasks without further adaptation. The results in Tab. Tab. show that those GP- GP-based models can only compete with NPs on function classes where either the inductive bias as given by the kernel functions fits the data well (RBF GP), ID). On more complex function classes, NPs produce predictions of much better quality, as they incorporate the source data more efficiently.
function classes which exhibit relatively low degree of variablity (Quadratic
7.5 EXPERIMENTAL DETAILS
We provide details about the data sets well as about the experimental setup used in our experiments
in Sec. 5
7.5.1 DATA GENER ATION
In our experiments, we use several classes of functions to evaluate the architectures under considera- tion. To generate training data from these function classes, we sample I. random tasks (as described in Sec.5), and Nooe random input locations z for each task. For each minibatch of training tasks, we uniformly sample a context set size N E {mmiin.... mmax and use a random subset of N data points from each task as context sets D The remaining M = Noot N data points used the target sets D' (cf. App. 7.5.3 for the special case of the VI likelihood approximation). Tab. 8 provides details about the data generation process.
GP Samples. We sample one-dimensional functions f R + IR from GP priors with three different stationary kernel functions as proposed by Gordon et al. 2020).
A radial basis functions (RBF) kernel with lenghtscale | - 1.0:
KRBF (r) exp (-0.5r").
(22)
A weakly periodic kernel:
kwp (r) = exp (-2 sin (0.5r) 0.125r
(23)
A Matern-5/2 kernel with lengthscale [ 0.25:
v5r
5r
v5r
k/ (r)
1
exp
(24)
0.25
0.25
0.25
Quadratic Functions. We consider two classes of quadratic functions. The first class fQ.ID :R R R is defined on a one-dimensional domain and parametrized by three parameters a. b. R:
QQ.I (z) a² (x b)² cc.
(25)
The second class FQ.3D; R3 - defined on three-dimensional domain and also parametrized by three parameters a, b, E R:
GQ.3D (4 T2, 0.5a { + 23 x3) + b (x11 x2 + xz) + 3c.
(26)
This function class was proposed in Perrone et (2018).
For both function classes we add Gaussian noise with standard deviation σn to the evaluations, cf. Tab.
Furuta Pendulum Dynamics. We consider function class obtained by integrating the non-linear
equations of motion governing the dynamics of a Furuta pendulum (Furuta et al.) 1992 Cazzolato and Prime 2011) for time span of At 0.1 s. More concretely, we consider the mapping
+ (l) + + (t) + (1),
(27)
17

---

Published as a conference paper at ICLR 2021
Table 8: Input spaces and parameters used to generate data for training and testing the architectures
discussed in the main part of this paper. U (a. b) denotes the uniform distribution on the interval [a, b]. and, likewise U {a, + n denotes the uniform distribution the {a, a + I..... n).
Symbol
Description
Value/Sampling distribution
GP Samples
T
Input
U(-2.0,+2.0)
Niot
Number of data points per task
128 (3....,50 50)
(Amins--. Rmax
Context set sizes during training
1D Quadratic Functions
Input
U (11.0.++..0)
a
Parameter
U ((-0.5.+1.3
b
Parameter
U
C
Parameter
U(-1..+
On
Noise standard deviation
0.01
Noot
Number of data points per task
128
(Aminy.. Rmmax
Context set sizes during training
Uf0, 20,
3D Quadratic Functions
21, T2, T3
Inputs
U (-1.0, +1.0)
a, b, c
Parameters
(+0.1,+10.0)
On Noot {Amiu.... Rmmax
Noise standard deviation
0.01
Number of data points per task
128
Context set sizes during training
Furuta Dynamics
arm, ppndd
Input angles
(0.0, 2T rad)
Barns Opend
Input angular velocities
/0.5s)
Marm
Mass arm
U U6.00 100-110-"gg...0 6.0 10 kg)
Moend
Mass pendulum
10-
lam
Length arm
Larm
Distance joint arm - mass arm
11.00-10-1m,300 10-
L Leaad
Distance joint pend. - mass pend.
U1.00 11.0-10-'m,33 10-²
barm
Damping constant arm
U2.0 10- Nmms, 2 Nms)
ppend
Damping constant pendulum
(5.6- 15 1mms, 5.6 10-3 1mm
rrarm
Action noise standard dev. arm
0.5 Nm
OT Tr.pend
Action noise standard dev. pend.
0.5 Nm
Niot
Number of data points per task
256
(Amin Rmaax
Context set sizes during training
2D Image Completion MNIST
F1, .72
Input pixel locations Number of data points per task
U 27] (scaled to [0, 1)
Ntot
28 28
(Aeiiu.... Tamax
Context set sizes during training
Uf0. ........8............. 28 22.2822))
where + = 0amm (t), 9pend (t), 0am (t) 8pend (t) denotes the four-dimensional vector describ- ing the dynamical state the Furuta pendulum. The Furuta pendulum parametrized by seven parameters (two masses, three lengths, two damping constants) as detailed in Tab. g During training, we provide = 64 tasks, corresponding to 64 different parameter configurations. We consider the free system and generate noise by applying random torques at each integration time step (Atruler 0.001 s) to the joints the arm and pendulum drawn from Gaussian distributions with standard deviations a pend, arm: respectively.
18

---

Published as a conference paper at ICLR 2021
MLP
14
MLP
Py
Bayesian aggregation
MLP
MLP
(a) BA PB.
MLP
Py
Mcan aggregation
(2 2
MLP
MT.P
(b) MA det. (CNP).
MLP
MLP
P9
Bayesian aggregation
(Sry) Ser
MLP
MLP
(c) BA VI
VI MC.
ML.P
ML.P
Mean aggregation
MLP
MLP
MLP
(d) MA + (I MA + MC.
Figure 4: Model architectures used for experiments in Sec. For the ANP architecture refer the reader to Kim et al. (2019). Orange rectangles denote MI.Ps. Blue rectangles denote aggregation operations. Variables green rectangles given by the incoming nodes. To arrive at fair comparison, we optimize all MLP architectures, the latent space dimensionality de, as well the Adam learning rate, individually for all model architectures and all experiments, cf. App. 7.5.3
sampled from normal distributions with parameters
2D Image Completion. For this task, we use the MNIST database of 28 X 28 images of handwritten digits (LeCun and Cortes 2010), and define 2D functions mapping pixel locations I1, $2 {O.... 27} (scaled to the unit square) to the corresponding pixel intensities (scaled the unit interval), cf. Tab. [ One training task corresponds to one image drawn randomly from the training set (consisting of 60000 images) and for evaluation we use subset of the test set (consisting 10000 images).
{0, ..... 255}
7.5.2 MODEL ARCUITECTURES
We provide the detailed architectures used for the experiments in Sec. in Fig. For ANP multihead cross attention and refer the reader to Kim et al. (2019) for details about the architecture.
7.5.3 HYPERPARAMETERS AND HYPERPARAMETER OPTIMIZATION
To arrive fair comparison of our BA with MA, it is imperative to use optimal model architectures for each aggregation method and likelihood approximation under consideration. Therefore, we optimize the number of hidden layers and the number hidden units per layer of each encoder and decoder MLP (as shown in Fig. individually for each model architecture and each experiment. For the ANP, we also optimize the multihead attention MLPs. We further optimize the latent space dimensionality d, and the learning rate of the Adam optimizer. For this hyperparameter optimization. we use the Optuna framework (Akiba et al. 2019) with TPE Sampler and Hyperband pruner (Li et al. minibatch size of 16. Further,
2017). We consistently
10 latent samples to evaluate
19

---

Published as a conference paper at ICLR 2021
the MC likelihood approximation during training. To evaluate the VI likelihood approximation, sample target set sizes between Mot and N in each training epoch, cf. Tab.
7.5.4 EVALUATION PROCEDURE
To evaluate the performance of the various model architectures we generate 256 unseen test tasks with target sets Dj consisting of M = 256 data points each and compute the average posterior given context sets Di of size N.
predictive log-likelihood M Ei= log (Vi T.... Di. 0
Depending on the architecture, we approximate the posterior predictive log-likelihood according to:
For BA PB likelihood approximation:
1 1
log v/.m Mz,
(28)
.
For MA + deterministic loss (= CNP):
M Σ 8=
11
log (v am Fe, 0)
(29)
T M
For architectures employing sampling-bas based likelihood approximations (VI, MC-LL) we report the joint log- liglliieldood over all data points in test set, i.e.
/
log
90 (zd D2)
II p m
xlim ze, 0) dzz
(30)
M
s
Md
log
II
yi. x2 Zt,0: 0) 8)
(31)
L M
t=1
mmm
1
1 1 T.M
M Σ
Ze.0 (im 2))
log S +
logsumexp
log (vé, =
(32)
M
s=1
where 2e,s 2o (=) Dz). We employ S 25 latent samples.
To compute the log- ligllikel values given in tables, we additionally average over various context set sizes N as detailed in the main part of this paper.
We report the mean posterior predictive log-likelihood computed in this way w.r.t 10 training runs with different random seeds together with 95% confidence intervals
20

---

Published as a conference paper at ICLR 2021
Table 9: Relative evaluation runtimes and numbers of parameters of the optimized network archi- tectures on the GP tasks. The deterministic methods (PB, det.) are much more efficient regarding evaluation runtime, as they require only on forward pass per prediction, while the sampling- based approaches (VI, MC) require multiple forward passes (each corresponding to one latent sample) to compute their predictions. We use 5 25 latent samples, described in App. [75.44 Furthermore, BA tends to require less complex encoder and decoder network architectures compared to MA, because it represents a more efficient mechanism to propagate information from the context set to latent state.
PB/dt. MA (CNP)
VI
MC
BA
(LP-NP)
RBF GP
Runtime
1.4
Parameters
72k
96k
63k
77k
122k
153k
Weakly Periodic GP
Rumtime
1.4
11
15
#Parameters
51k
87k
48k
72k
87K
89k
Matern-5/2 GP
Runtime
1.1
6,5
11
19
Parameters
100k
Table 10: Posterior predictive mean squared error (MSE) on all experiments presented in this paper. We average over the same context set sizes used compute the posterior predictive log- likelihood, cf. Sec. and again use S = 25 latent samples to compute the mean prediction of sampling based methods. Our BA consistently improves predictive performance compared to MA not only in terms of likelihood (as shown in Sec. 5) but also in terms of MSE. Furthermore, while ANP tends to perform poorly in terms of likelihood (ef. Sec.
(5. it's MSE is improved greatly by the attention mechanism.
PENder
VI
MC
ANP
BA
MA (CNP)
BA
MA (LP-NP)
BA
MA
MA + Atiantian
RBF GP
0. 0.0623 1: 0.0009
0.0887
0.0000
0.0736 2: 0.0005
0.8388 0,0036
0.06U7 1. 0000
0,0741
0.0012
0.0G50 &: 0009
Weakly Periodic GP
0. 0.00779
0007
0.0766
0.00L
0.0879
1: 0.0017
0.3326
0G18
0.0677
1 0008
0,1832
0.009
0.0592 1 0009
Maserm
0.2452 0088
0,3921
00005
0.3702
& 0.01000
0.1077
0.2321
0.0019
0.5105
0.1438
0.980 0012
Quadratics
0.0065 190.5
0.1613
0.1757
0.1473
0.0107
0.0082
0.1330
00077
Quadratics
195.4:1 1,5
253.1
278.1 40.5
196.8 2.6
200.
0.0092
0.1980
0.0088
(12000
0.0154
20 Image Complction
D 0348 1.0010
0.0117 (10006
0.0215 0O0S
7.6 ADDITIONAL EXPERIMENTAL RESULTS
We provide additional experimental results accompanying the experiments presented in Sec.
Results for relative evaluation runtimes and numbers of parameters of the optimized network
architectures on the full GP suite of experiments, Tab. g
The posterior predictive mean squared error all experiments, cf. Tab. 10
The context-size dependent results for the predictive posterior log- likelihoodd for the ID and 3D Quadratic experiments, the Furuta dynamics experiment, as well as the 2D image completion experiment, cf. Fig. 5
More detailed plots of the predictions on one-dimensional experiments (ID Quadratics (Figs. 6 RBF-G, (Figs. g 9), Weakly Periodic GP (Figs. I0] ITI. and Matern-5/2 GP
(Figs.
12
21

---

Published as a conference paper at ICLR 2021
Quadratic
Quadratic
Furui
Aggregaror
log an
BA
Less
Vi-irsprod
MCC
20
20
10
20
200
Figure 5: Posterior predictive log-likelihood in dependence of the context set size N for the 1D) and 3D Quadratic experiments, the Furuta dynamics experiment as well as the 2D image completion experiment.
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
(a) BA + PB
(b) MA + det. (CNP)
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
(c) BA + VI
(d) MA + VI (LP-NNPP
(c) ANP
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
( BA MC-LI
(g) MA + MC-LL
Figure Predictions on two instances (dashed lines) of the ID quadratic function class, given N context data points (circles). We plot mean and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
22

---

Published as a conference paper at ICLR 2021
(a) BA PB
(b) MA + det. (CNP)
1.00
(c) BA VI
(d) MA (LP-NN)
(e) ANP
0.00
0.25
0.30
0.75
1.00
0.00
0.25
0.50
0.75
1.00
(f) BA + MC-LI
(g) MA + MC-LL
Figure Predictions on two instances (dashed lines) of the ID quadratic function class, given N 19 context data points (circles). We plot mean and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
23

---

Published as a conference paper at ICLR 2021
(a) BA PB
(b) MA + det. (CNP)
-2
2
(c) BA VI
(d) MA + (L
(e) ANP
(f) BA + MC-LI
(g) MA + MC-LL
Figure 8: Predictions on two instances (dashed lines) of the RBF GP function class, given = 20 context data points (eircles). We plot mean and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
24

---

Published as a conference paper at ICLR 2021
(a) BA PB
(b) MA + det. (CNP)
-2
2
-2
(c) BA VI
(d) MA + (L
(e) ANP
(f) BA + MC-LI
(g) MA + MC-LL
Figure 9: Predictions on two instances (dashed lines) of the RBF GP function class, given = 60 context data points (eircles). We plot mean and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
25

---

Published as a conference paper at ICLR 2021
(a) BA PB
(b) MA + det. (CNP)
(c) BA VI
(d) MA (LPP-NP)
(e) ANP
(f) BA + MC-LI
(g) MA + MC-LL
Figure 10: Predictions on two instances (dashed lines) of the Weakly Periodic GP function class, given N 20 context data points (circles). We plot and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
26

---

Published as a conference paper at ICLR 2021
(a) BA PB
(b) MA + det. (CNP)
(c) BA VI
(d) MA (LP-PN)
(e) ANP
(f) BA + MC-II
(g) MA + MC-LL
Figure 11: Predictions on two instances (dashed lines) of the Weakly Periodic GP function class, given N 60 context data points (circles). We plot and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
27

---

Published as a conference paper at ICLR 2021
(a) BA PB
(b) MA + det. (CNP)
(c) BA VI
(d) MA (L
(e) ANP
(f) BA + MC-LI
(g) MA + MC-LL
Figure 12: Predictions on two instances (dashed lines) of the Matern-5/ 5/2 GP function class, given N 20 context data points (circles). We plot mean and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
28

---

Published as a conference paper at ICLR 2021
(a) BA PB
(b) MA + det. (CNP)
(c) BA VI
(d) MA + (L
(e) ANP
(f) BA + MC-LI
(g) MA + MC-LL
Figure 13: Predictions on two instances (dashed lines) of the Matern-5/ GP function class, given N 60 context data points (circles). We plot mean and standard deviation (solid line, shaded area) predictions together with 10 function samples (for deterministic methods we employ AR sampling).
29


Reviews:
[{"review_id": "rbzGHvGMoiQ", "reviewer": "AnonReviewer3", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors make two contributions to neural process-like CLV models. First, they replace the somewhat adhoc variational-like approach to learning the amortized latent variable distribution with a monte carlo based approximation. Second, they replace the step of context aggregation with direct latent variable inference over z. \n\nOverall, in my opinion these modifications make the neural process model significantly cleaner from a Bayesian perspective and is quite nice. In the context of neural processes, I have very little criticism for the authors' methods. Everything makes sense, and the MC approach in equation (3) seems cleaner to me than the somewhat ad hoc \"VI like\" approach in equation (2).\n\nThe biggest difficulty I have is determining how to evaluate the authors' clear improvements to neural processes in the broader context of scalable probabilistic regression, an area to which the authors claim membership. To start with, the authors clearly demonstrate the value of both Bayesian context aggregation and a MC based likelihood approximation scheme on precisely the same types of problems that existing neural processes papers (e.g., Garnelo et al., 2018) have considered (with the notable exception that the 2D image completion task considers only MNIST as a target dataset). In this respect, it's difficult to fault the experimental evaluation.\n\nHowever, this paper and many neural process papers are written in the context of \"Formulating scalable probabilistic regression models with reliable uncertainty estimates.\" Surely, at some point, this should involve a comparison of these approaches to existing techniques for probabilistic regression, whether that be deep Gaussian processes, dropout based approaches, Bayesian neural networks, or other approaches. I don't mean to imply here that the authors are unaware of this large body of literature -- indeed, the authors have a decent if incomplete overview of techniques in this area (notably missing work on deep GPs).\n\nRather, it just seems surprising to me that the discussion of the relevant probabilistic regression literature ends at \"well, it exists.\" What I would like to see is a discussion of where the authors' impressive improvements to neural processes leave the model family in this broader context. How close or far off is the family on performance for standard benchmark regression tasks? Are there settings in which we can leverage the fully NN based nature of neural processes to achieve probabilistic regression in settings where the inductive biases of kernel methods are poor, like in computer vision or natural language processing? The relatively toy nature and limited dimensionality of the problems considered suggests that there is still significant progress to be made before such a comparison would be reasonable or even possible.\n\nTo summarize, in the context of neural processes I feel the paper makes good methodological contributions in presenting a much cleaner and more natural (from a Bayesian perspective) version of the model that has more of the flavor of standard amortized inference for latent variable models. Within the very narrow context of neural process papers, I therefore have very little to complain about. However, from a broader scientific perspective I would feel that the paper would be significantly strengthened by a fair evaluation to the rest of this literature, whether empirical or simply in discussion, regardless of how the authors' approach fares in comparison.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Z5i8d9My61", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper builds upon previous lines of research on multi-task learning problem, such as conditional latent variable models including the Neural Process. As shown by the extensive Related Work section, this seems to be an active research direction. This makes it difficult for me to judge originality and significance, but it is well-written and clear.\n\nSpecific comments\n- the approximate posterior distribution q_\\phi is often referred to as \"the posterior distribution\". I would keep \"approximate\" here. \n- p2: \"correspondence of GPs with infinite Bayesian NNs (BNNs)\", what is meant by \"infinite BNN\"? is it infinite-width BNN? please specify.\n- p2: \"adaptive BLR\" please describe the acronym.\n- p6: the Gaussian approximation of the posterior predictive likelihood (10) is said to be \"inspired by GPs which also define a Gaussian likelihood\". This is also essentially what is done by Synthetic Likelihood (a Nature paper by SN Wood, 2010) which is I think more closely related to the proposed approach than GPs.\n- p6, one line below: define PB acronym the first time it is used (not three lines after).\n\nTypos\n- top p2: \"does not introducing noticeable computational overhead\".\n- p6: \"conditional model with an decoder operating on\".\n- The references list could be tidied: some first names are abbreviated, some not.", "extra_scores": {}, "extra_text": {}}, {"review_id": "8Zf4uJ1IKqb", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of the Paper:\n  \n        This paper describes Bayesian context aggregation for neural processes. These models are useful to address regression problems in which a set of related tasks are available for inference with associated context information in the form of extra data. These models assume that there is a task-specific global latent variable and task-independent latent variable. They are learned via approximate maximum posterior likelihood, in which the latent variables specific for each tasks are marginalzied out. For this, an approximation to the posterior distribution of these variables is need. This requires conditioning to the context dataset which is challenging. In the past, a latent representation is used and the context data set is aggregated as the mean of the latent representation. In this paper a Bayesian way of aggregating context information is proposed. This is based on using Bayes rule and a Gaussian generative model for the latent representations. The proposed method also leads to a new way of training CLV models which is based on moment matching. The method is validated on several synthetic an real-world experiments showing improvements over mean aggregation.\n\nDetailed Comments:\n\nI believe that this is a relevant paper. Context aggregation is a difficult problem that is required to address the learning tasks described in the paper. Previous solution look limited and the proposed method seems natural and a more effective method of aggregating this information. The paper is well written and the proposed method is sound. The experiments are also convincing and exhaustive. I believe that this is a relevant paper for the conference.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "3PwHp4elcTP", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors present the Bayesian Aggregation (BA) mechanism in the context of Neural Processes (NPs) for aggregating the context information into the latent variable z in the form of posterior updates to z. The authors show that this improves predictive performance (in terms of likelihood) compared to mean aggregation MA that it replaces on various regression tasks with varying input-output dimensionality.\n\nStrengths:\n1. The idea is simple and leads to a notable improvement compared to MA in terms of likelihood\n2. The background and method is presented very clearly.\n3. The evaluation is done on a wide variety of tasks, ranging from standard 1D regression of GP samples to pendulum trajectory prediction tasks.\n\nWeaknesses:\n1. The evaluation is missing an important baseline model, which are (A)NP models that have self-attention in the encoder for processing the contexts (c.f. model figure in ANP paper (Kim et al., 2019b)). Contrary to the NP/CNP baselines that are compared against in the paper, the ANP with self-attention in the encoder does not give uniform weights to each context point - the self-attention allows the model to assign varying importance to the different context points (despite using mean-aggregation after the self-attention), which is presented as a key motivation for the BA mechanism introduced in the paper. Hence for the experiments, I strongly suggest comparing against CNP/NP/ANP with self-attention in the deterministic/latent/latent path of the encoder. For completeness, if would be nice to also compare against models that have both deterministic and latent paths, since BA can also be applied to these models. At the same time, I understand that BA would be more interpretable for showing which observations have little/high effect on z compared to the approach of using self-attention in the encoder, but it would still be very informative for the reader to be able to compare the two approaches. Also these two approaches can be combined to have self-attention in the encoder + BA, which might also yield improved performance.\n2. The claim that “BA includes MA as a special case” doesn’t seem to be true. Using a “non-informative prior and uniform observation variances” leads to constant sigma_z and mu_z being linearly proportional to mean(r_n) (i.e. sum_n r_n / N), which is not quite the same as MA - MA allows sigma_z and mu_z to be non-linear functions of mean(r_n), hence is strictly more expressive than this special case.\n3. In Equation (7), it seems as though the context points (x_n,y_n) only affects r_n via the variance, which seems unnecessarily limiting. Why not have the mean also depend on r_n? e.g. p(r_n|z) = N(r_n| z + mu_{r_n}, diag(sigma_{r_n}^2) where mu_{r_n} is also computed as a function of (x_n,y_n)? This will still give a closed-form posterior p(z|r_{1:N}) since the mean of p(r_n|z) is still linear in z, creating a model that’s strictly more expressive with very similar efficiency. It would be informative to see how this changes the experimental results.\n4. I’m guessing the VI objective was used to train the ANP. Given the clear advantage of training with the MC objective, shouldn’t the ANP also be trained with MC?\n5. The latent variable models were not evaluated on 2D image completion tasks because “architectures without deterministic paths were not able to solve this task”. Why not then add a deterministic path to these latent variable models to allow them to train?\n\nOther points\n- In the text, it says that the model is also compared against ANP to show that BA can compete with SOTA. This is arguably incorrect since ConvCNP models are SOTA among models of the NP family, showing a significant improvement over ANP. Hence to achieve the goal mentioned in the text, it would make sense to compare with ConvCNP models as part of the evaluation against other deterministic NPs.\n\nOverall the paper is presented very clearly with a simple yet effective idea tested on a wide variety of tasks. However it’s missing an important baseline that uses self-attention in the encoder, along with several other baselines that would be informative to compare against. I am willing to increase my score should these results be included in the revised version of the paper.\n\n=================\n\nScore raised to 6 after inclusion of MA + SA results in rebuttal.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors present a Bayesian approach for context aggregation in neural processes based models. The article is well written, and provides a nice and comprehensive framework. The reviewers raised some issues regarding the lack of comparisons to proper baselines. The authors provided additional comparisons in the revised version. The comparisons were found satisfactory by some some reviewers, who increased their scores. Based on the revised version, I recommend acceptance.
Decision: Accept (Poster)